- Refactor command structure for better organization - Improve WebSocket client frame handling - Add response handler improvements - Update queue, requeue, and status commands - Add security module for CLI authentication
206 lines
8.2 KiB
Zig
206 lines
8.2 KiB
Zig
const std = @import("std");
|
|
const security = @import("security.zig");
|
|
|
|
pub const Config = struct {
|
|
worker_host: []const u8,
|
|
worker_user: []const u8,
|
|
worker_base: []const u8,
|
|
worker_port: u16,
|
|
api_key: []const u8,
|
|
|
|
// Default resource requests
|
|
default_cpu: u8,
|
|
default_memory: u8,
|
|
default_gpu: u8,
|
|
default_gpu_memory: ?[]const u8,
|
|
|
|
// CLI behavior defaults
|
|
default_dry_run: bool,
|
|
default_validate: bool,
|
|
default_json: bool,
|
|
default_priority: u8,
|
|
|
|
pub fn validate(self: Config) !void {
|
|
// Validate host
|
|
if (self.worker_host.len == 0) {
|
|
return error.EmptyHost;
|
|
}
|
|
|
|
// Validate port range
|
|
if (self.worker_port == 0 or self.worker_port > 65535) {
|
|
return error.InvalidPort;
|
|
}
|
|
|
|
// Validate API key presence
|
|
if (self.api_key.len == 0) {
|
|
return error.EmptyAPIKey;
|
|
}
|
|
|
|
// Validate base path
|
|
if (self.worker_base.len == 0) {
|
|
return error.EmptyBasePath;
|
|
}
|
|
}
|
|
|
|
pub fn load(allocator: std.mem.Allocator) !Config {
|
|
const home = std.posix.getenv("HOME") orelse return error.NoHomeDir;
|
|
const config_path = try std.fmt.allocPrint(allocator, "{s}/.ml/config.toml", .{home});
|
|
defer allocator.free(config_path);
|
|
|
|
const file = std.fs.openFileAbsolute(config_path, .{}) catch |err| {
|
|
if (err == error.FileNotFound) {
|
|
std.debug.print("Config file not found. Run 'ml init' first.\n", .{});
|
|
return error.ConfigNotFound;
|
|
}
|
|
return err;
|
|
};
|
|
defer file.close();
|
|
|
|
// Load config with environment variable overrides
|
|
var config = try loadFromFile(allocator, file);
|
|
|
|
// Apply environment variable overrides (FETCH_ML_CLI_* to match TUI)
|
|
if (std.posix.getenv("FETCH_ML_CLI_HOST")) |host| {
|
|
config.worker_host = try allocator.dupe(u8, host);
|
|
}
|
|
if (std.posix.getenv("FETCH_ML_CLI_USER")) |user| {
|
|
config.worker_user = try allocator.dupe(u8, user);
|
|
}
|
|
if (std.posix.getenv("FETCH_ML_CLI_BASE")) |base| {
|
|
config.worker_base = try allocator.dupe(u8, base);
|
|
}
|
|
if (std.posix.getenv("FETCH_ML_CLI_PORT")) |port_str| {
|
|
config.worker_port = try std.fmt.parseInt(u16, port_str, 10);
|
|
}
|
|
if (std.posix.getenv("FETCH_ML_CLI_API_KEY")) |api_key| {
|
|
config.api_key = try allocator.dupe(u8, api_key);
|
|
}
|
|
|
|
// Try to get API key from keychain if not in config or env
|
|
if (config.api_key.len == 0) {
|
|
if (try security.SecureStorage.retrieveApiKey(allocator)) |keychain_key| {
|
|
config.api_key = keychain_key;
|
|
}
|
|
}
|
|
|
|
try config.validate();
|
|
return config;
|
|
}
|
|
|
|
fn loadFromFile(allocator: std.mem.Allocator, file: std.fs.File) !Config {
|
|
const content = try file.readToEndAlloc(allocator, 1024 * 1024);
|
|
defer allocator.free(content);
|
|
|
|
// Simple TOML parser - parse key=value pairs
|
|
var config = Config{
|
|
.worker_host = "",
|
|
.worker_user = "",
|
|
.worker_base = "",
|
|
.worker_port = 22,
|
|
.api_key = "",
|
|
.default_cpu = 2,
|
|
.default_memory = 8,
|
|
.default_gpu = 0,
|
|
.default_gpu_memory = null,
|
|
.default_dry_run = false,
|
|
.default_validate = false,
|
|
.default_json = false,
|
|
.default_priority = 5,
|
|
};
|
|
|
|
var lines = std.mem.splitScalar(u8, content, '\n');
|
|
while (lines.next()) |line| {
|
|
const trimmed = std.mem.trim(u8, line, " \t\r");
|
|
if (trimmed.len == 0 or trimmed[0] == '#') continue;
|
|
|
|
var parts = std.mem.splitScalar(u8, trimmed, '=');
|
|
const key = std.mem.trim(u8, parts.next() orelse continue, " \t");
|
|
const value_raw = std.mem.trim(u8, parts.next() orelse continue, " \t");
|
|
|
|
// Remove quotes
|
|
const value = if (value_raw.len >= 2 and value_raw[0] == '"' and value_raw[value_raw.len - 1] == '"')
|
|
value_raw[1 .. value_raw.len - 1]
|
|
else
|
|
value_raw;
|
|
|
|
if (std.mem.eql(u8, key, "worker_host")) {
|
|
config.worker_host = try allocator.dupe(u8, value);
|
|
} else if (std.mem.eql(u8, key, "worker_user")) {
|
|
config.worker_user = try allocator.dupe(u8, value);
|
|
} else if (std.mem.eql(u8, key, "worker_base")) {
|
|
config.worker_base = try allocator.dupe(u8, value);
|
|
} else if (std.mem.eql(u8, key, "worker_port")) {
|
|
config.worker_port = try std.fmt.parseInt(u16, value, 10);
|
|
} else if (std.mem.eql(u8, key, "api_key")) {
|
|
config.api_key = try allocator.dupe(u8, value);
|
|
} else if (std.mem.eql(u8, key, "default_cpu")) {
|
|
config.default_cpu = try std.fmt.parseInt(u8, value, 10);
|
|
} else if (std.mem.eql(u8, key, "default_memory")) {
|
|
config.default_memory = try std.fmt.parseInt(u8, value, 10);
|
|
} else if (std.mem.eql(u8, key, "default_gpu")) {
|
|
config.default_gpu = try std.fmt.parseInt(u8, value, 10);
|
|
} else if (std.mem.eql(u8, key, "default_gpu_memory")) {
|
|
if (value.len > 0) {
|
|
config.default_gpu_memory = try allocator.dupe(u8, value);
|
|
}
|
|
} else if (std.mem.eql(u8, key, "default_dry_run")) {
|
|
config.default_dry_run = std.mem.eql(u8, value, "true");
|
|
} else if (std.mem.eql(u8, key, "default_validate")) {
|
|
config.default_validate = std.mem.eql(u8, value, "true");
|
|
} else if (std.mem.eql(u8, key, "default_json")) {
|
|
config.default_json = std.mem.eql(u8, value, "true");
|
|
} else if (std.mem.eql(u8, key, "default_priority")) {
|
|
config.default_priority = try std.fmt.parseInt(u8, value, 10);
|
|
}
|
|
}
|
|
|
|
return config;
|
|
}
|
|
|
|
pub fn save(self: Config, allocator: std.mem.Allocator) !void {
|
|
const home = std.posix.getenv("HOME") orelse return error.NoHomeDir;
|
|
|
|
// Create .ml directory
|
|
const ml_dir = try std.fmt.allocPrint(allocator, "{s}/.ml", .{home});
|
|
defer allocator.free(ml_dir);
|
|
|
|
std.fs.makeDirAbsolute(ml_dir) catch |err| {
|
|
if (err != error.PathAlreadyExists) return err;
|
|
};
|
|
|
|
const config_path = try std.fmt.allocPrint(allocator, "{s}/config.toml", .{ml_dir});
|
|
defer allocator.free(config_path);
|
|
|
|
const file = try std.fs.createFileAbsolute(config_path, .{});
|
|
defer file.close();
|
|
|
|
const writer = file.writer();
|
|
try writer.print("worker_host = \"{s}\"\n", .{self.worker_host});
|
|
try writer.print("worker_user = \"{s}\"\n", .{self.worker_user});
|
|
try writer.print("worker_base = \"{s}\"\n", .{self.worker_base});
|
|
try writer.print("worker_port = {d}\n", .{self.worker_port});
|
|
try writer.print("api_key = \"{s}\"\n", .{self.api_key});
|
|
try writer.print("\n# Default resource requests\n", .{});
|
|
try writer.print("default_cpu = {d}\n", .{self.default_cpu});
|
|
try writer.print("default_memory = {d}\n", .{self.default_memory});
|
|
try writer.print("default_gpu = {d}\n", .{self.default_gpu});
|
|
if (self.default_gpu_memory) |gpu_mem| {
|
|
try writer.print("default_gpu_memory = \"{s}\"\n", .{gpu_mem});
|
|
}
|
|
try writer.print("\n# CLI behavior defaults\n", .{});
|
|
try writer.print("default_dry_run = {s}\n", .{if (self.default_dry_run) "true" else "false"});
|
|
try writer.print("default_validate = {s}\n", .{if (self.default_validate) "true" else "false"});
|
|
try writer.print("default_json = {s}\n", .{if (self.default_json) "true" else "false"});
|
|
try writer.print("default_priority = {d}\n", .{self.default_priority});
|
|
}
|
|
|
|
pub fn deinit(self: *Config, allocator: std.mem.Allocator) void {
|
|
allocator.free(self.worker_host);
|
|
allocator.free(self.worker_user);
|
|
allocator.free(self.worker_base);
|
|
allocator.free(self.api_key);
|
|
if (self.default_gpu_memory) |gpu_mem| {
|
|
allocator.free(gpu_mem);
|
|
}
|
|
}
|
|
};
|