ad3be36a6d
New commands and modules: - Add workers.zig command for worker management and status - Add scheduler_client.zig for scheduler hub communication - Add pii.zig utility for PII detection and redaction in logs/outputs Improvements to existing commands: - groups.zig: enhanced group management with capability metadata - jupyter/mod.zig: improved Jupyter workspace lifecycle handling - tasks.zig: better task status reporting and cancellation support Networking and sync improvements: - ws/client.zig: WebSocket client enhancements for hub protocol - sync_manager.zig: improved sync with scheduler state and conflict resolution - uuid.zig: optimized UUID generation for macOS and Linux Database utilities: - sqlite_embedded.zig: embedded SQLite for CLI-local state caching
1598 lines
57 KiB
Zig
1598 lines
57 KiB
Zig
const std = @import("std");
|
|
const crypto = @import("crypto");
|
|
const io = @import("io");
|
|
const log = @import("log");
|
|
const protocol = @import("../protocol.zig");
|
|
const resolve = @import("resolve.zig");
|
|
const handshake = @import("handshake.zig");
|
|
const frame = @import("frame.zig");
|
|
const response = @import("response.zig");
|
|
const response_handlers = @import("response_handlers.zig");
|
|
const opcode = @import("opcode.zig");
|
|
const utils = @import("utils.zig");
|
|
const tls = @import("tls.zig");
|
|
|
|
/// Transport abstraction for WebSocket connections
|
|
/// Supports both raw TCP and TLS-wrapped connections
|
|
pub const Transport = union(enum) {
|
|
tcp: std.net.Stream,
|
|
tls: *tls.TlsStream,
|
|
|
|
pub fn read(self: Transport, buffer: []u8) !usize {
|
|
return switch (self) {
|
|
.tcp => |s| s.read(buffer),
|
|
.tls => |s| s.read(buffer),
|
|
};
|
|
}
|
|
|
|
pub fn write(self: Transport, buffer: []const u8) !void {
|
|
return switch (self) {
|
|
.tcp => |s| s.writeAll(buffer),
|
|
.tls => |s| s.write(buffer),
|
|
};
|
|
}
|
|
|
|
pub fn close(self: *Transport) void {
|
|
switch (self.*) {
|
|
.tcp => |*s| s.close(),
|
|
.tls => |s| s.close(),
|
|
}
|
|
self.* = undefined;
|
|
}
|
|
};
|
|
const MessageBuilder = struct {
|
|
buffer: []u8,
|
|
offset: usize,
|
|
allocator: std.mem.Allocator,
|
|
|
|
pub fn init(allocator: std.mem.Allocator, total_len: usize) !MessageBuilder {
|
|
const buffer = try allocator.alloc(u8, total_len);
|
|
return MessageBuilder{
|
|
.buffer = buffer,
|
|
.offset = 0,
|
|
.allocator = allocator,
|
|
};
|
|
}
|
|
|
|
pub fn deinit(self: *MessageBuilder) void {
|
|
self.allocator.free(self.buffer);
|
|
}
|
|
|
|
pub fn writeOpcode(self: *MessageBuilder, op: opcode.Opcode) void {
|
|
self.buffer[self.offset] = @intFromEnum(op);
|
|
self.offset += 1;
|
|
}
|
|
|
|
pub fn writeBytes(self: *MessageBuilder, data: []const u8) void {
|
|
@memcpy(self.buffer[self.offset .. self.offset + data.len], data);
|
|
self.offset += data.len;
|
|
}
|
|
|
|
pub fn writeU8(self: *MessageBuilder, value: u8) void {
|
|
self.buffer[self.offset] = value;
|
|
self.offset += 1;
|
|
}
|
|
|
|
pub fn writeU16(self: *MessageBuilder, value: u16) void {
|
|
std.mem.writeInt(u16, self.buffer[self.offset .. self.offset + 2][0..2], value, .big);
|
|
self.offset += 2;
|
|
}
|
|
|
|
pub fn writeU32(self: *MessageBuilder, value: u32) void {
|
|
std.mem.writeInt(u32, self.buffer[self.offset .. self.offset + 4][0..4], value, .big);
|
|
self.offset += 4;
|
|
}
|
|
|
|
pub fn writeU64(self: *MessageBuilder, value: u64) void {
|
|
std.mem.writeInt(u64, self.buffer[self.offset .. self.offset + 8][0..8], value, .big);
|
|
self.offset += 8;
|
|
}
|
|
|
|
pub fn writeStringU8(self: *MessageBuilder, str: []const u8) void {
|
|
self.writeU8(@intCast(str.len));
|
|
if (str.len > 0) {
|
|
self.writeBytes(str);
|
|
}
|
|
}
|
|
|
|
pub fn writeStringU16(self: *MessageBuilder, str: []const u8) void {
|
|
self.writeU16(@intCast(str.len));
|
|
if (str.len > 0) {
|
|
self.writeBytes(str);
|
|
}
|
|
}
|
|
|
|
pub fn send(self: *MessageBuilder, transport: Transport) !void {
|
|
try frame.sendWebSocketFrame(transport, self.buffer);
|
|
}
|
|
};
|
|
|
|
/// WebSocket client for binary protocol communication
|
|
pub const Client = struct {
|
|
allocator: std.mem.Allocator,
|
|
transport: Transport,
|
|
host: []const u8,
|
|
port: u16,
|
|
is_tls: bool = false,
|
|
connected: bool = false,
|
|
|
|
pub fn formatPrewarmFromStatusRoot(allocator: std.mem.Allocator, root: std.json.ObjectMap) !?[]u8 {
|
|
return response.formatPrewarmFromStatusRoot(allocator, root);
|
|
}
|
|
|
|
pub fn connect(allocator: std.mem.Allocator, url: []const u8, api_key: []const u8) !Client {
|
|
// Detect TLS
|
|
const is_tls = std.mem.startsWith(u8, url, "wss://");
|
|
|
|
// Parse URL (simplified - assumes ws://host:port/path or wss://host:port/path)
|
|
const host_start = std.mem.indexOf(u8, url, "//") orelse return error.InvalidURL;
|
|
const host_port_start = host_start + 2;
|
|
const path_start = std.mem.indexOfPos(u8, url, host_port_start, "/") orelse url.len;
|
|
const colon_pos = std.mem.indexOfPos(u8, url, host_port_start, ":");
|
|
|
|
const host_end = blk: {
|
|
if (colon_pos) |pos| {
|
|
if (pos < path_start) break :blk pos;
|
|
}
|
|
break :blk path_start;
|
|
};
|
|
const host = url[host_port_start..host_end];
|
|
|
|
var port: u16 = if (is_tls) 9101 else 9100; // default ports
|
|
if (colon_pos) |pos| {
|
|
if (pos < path_start) {
|
|
const port_start = pos + 1;
|
|
const port_end = std.mem.indexOfPos(u8, url, port_start, "/") orelse url.len;
|
|
const port_str = url[port_start..port_end];
|
|
port = try std.fmt.parseInt(u16, port_str, 10);
|
|
}
|
|
}
|
|
|
|
// Connect to server
|
|
const tcp_stream = try std.net.tcpConnectToAddress(try resolve.resolveHostAddress(allocator, host, port));
|
|
|
|
// Setup transport - raw TCP or TLS
|
|
var transport: Transport = undefined;
|
|
if (is_tls) {
|
|
// Allocate TLS stream on heap
|
|
const tls_stream = try allocator.create(tls.TlsStream);
|
|
errdefer allocator.destroy(tls_stream);
|
|
|
|
// Initialize TLS stream with handshake
|
|
tls_stream.* = tls.TlsStream.init(allocator, tcp_stream, host) catch |err| {
|
|
allocator.destroy(tls_stream);
|
|
tcp_stream.close();
|
|
if (err == error.TlsLibraryRequired) {
|
|
std.log.warn("TLS (wss://) support requires external TLS library integration. Falling back to ws://", .{});
|
|
return error.TLSNotSupported;
|
|
}
|
|
return err;
|
|
};
|
|
transport = Transport{ .tls = tls_stream };
|
|
} else {
|
|
transport = Transport{ .tcp = tcp_stream };
|
|
}
|
|
|
|
// Perform WebSocket handshake over the transport
|
|
try handshake.handshake(allocator, transport, host, url, api_key);
|
|
|
|
return Client{
|
|
.allocator = allocator,
|
|
.transport = transport,
|
|
.host = try allocator.dupe(u8, host),
|
|
.port = port,
|
|
.is_tls = is_tls,
|
|
.connected = true,
|
|
};
|
|
}
|
|
|
|
/// Connect to WebSocket server with retry logic
|
|
pub fn connectWithRetry(allocator: std.mem.Allocator, url: []const u8, api_key: []const u8, max_retries: u32) !Client {
|
|
var retry_count: u32 = 0;
|
|
var last_error: anyerror = error.ConnectionFailed;
|
|
|
|
while (retry_count < max_retries) {
|
|
const client = connect(allocator, url, api_key) catch |err| {
|
|
last_error = err;
|
|
retry_count += 1;
|
|
|
|
if (retry_count < max_retries) {
|
|
const delay_ms = @min(1000 * retry_count, 5000); // Exponential backoff, max 5s
|
|
log.warn("Connection failed (attempt {d}/{d}), retrying in {d}s...\n", .{ retry_count, max_retries, delay_ms / 1000 });
|
|
std.Thread.sleep(@as(u64, delay_ms) * std.time.ns_per_ms);
|
|
}
|
|
continue;
|
|
};
|
|
|
|
if (retry_count > 0) {
|
|
log.success("Connected successfully after {d} attempts\n", .{retry_count + 1});
|
|
}
|
|
return client;
|
|
}
|
|
|
|
return last_error;
|
|
}
|
|
|
|
/// Disconnect from WebSocket server (closes transport only)
|
|
pub fn disconnect(self: *Client) void {
|
|
self.transport.close();
|
|
}
|
|
|
|
/// Fully close client - disconnects transport and frees host memory
|
|
pub fn close(self: *Client) void {
|
|
if (self.connected) {
|
|
self.disconnect();
|
|
self.connected = false;
|
|
}
|
|
if (self.host.len > 0) {
|
|
self.allocator.free(self.host);
|
|
}
|
|
}
|
|
|
|
// Validation helpers
|
|
fn validateApiKeyHash(api_key_hash: []const u8) error{InvalidApiKeyHash}!void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
}
|
|
|
|
fn validateCommitId(commit_id: []const u8) error{InvalidCommitId}!void {
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
}
|
|
|
|
fn validateJobName(job_name: []const u8) error{JobNameTooLong}!void {
|
|
if (job_name.len == 0 or job_name.len > 255) return error.JobNameTooLong;
|
|
}
|
|
|
|
fn getStream(self: *Client) error{NotConnected}!Transport {
|
|
// Return a copy of the transport - both TCP and TLS support read/write
|
|
return self.transport;
|
|
}
|
|
|
|
pub fn sendValidateRequestCommit(self: *Client, api_key_hash: []const u8, commit_id: []const u8) !void {
|
|
const stream = try self.getStream();
|
|
try validateApiKeyHash(api_key_hash);
|
|
try validateCommitId(commit_id);
|
|
|
|
var builder = try MessageBuilder.init(self.allocator, 1 + 16 + 1 + 1 + 20);
|
|
defer builder.deinit();
|
|
|
|
builder.writeOpcode(opcode.Opcode.validate_request);
|
|
builder.writeBytes(api_key_hash);
|
|
builder.writeU8(@intFromEnum(opcode.ValidateTargetType.commit_id));
|
|
builder.writeU8(20);
|
|
builder.writeBytes(commit_id);
|
|
try builder.send(stream);
|
|
}
|
|
|
|
pub fn sendQueryJobByCommit(self: *Client, job_name: []const u8, commit_id: []const u8, api_key_hash: []const u8) !void {
|
|
const stream = try self.getStream();
|
|
try validateApiKeyHash(api_key_hash);
|
|
try validateCommitId(commit_id);
|
|
try validateJobName(job_name);
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [job_name_len: u8] [job_name: var] [commit_id: 20 bytes]
|
|
const total_len = 1 + 16 + 1 + job_name.len + 20;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.query_job);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
|
|
try frame.sendWebSocketFrame(stream, buffer);
|
|
}
|
|
|
|
pub fn sendListJupyterPackages(self: *Client, name: []const u8, api_key_hash: []const u8) !void {
|
|
const stream = try self.getStream();
|
|
try validateApiKeyHash(api_key_hash);
|
|
if (name.len > 255) return error.NameTooLong;
|
|
|
|
var builder = try MessageBuilder.init(self.allocator, 1 + 16 + 1 + name.len);
|
|
defer builder.deinit();
|
|
|
|
builder.writeOpcode(opcode.list_jupyter_packages);
|
|
builder.writeBytes(api_key_hash);
|
|
builder.writeStringU8(name);
|
|
try builder.send(stream);
|
|
}
|
|
|
|
pub fn sendSetRunNarrative(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
patch_json: []const u8,
|
|
api_key_hash: []const u8,
|
|
) !void {
|
|
const stream = self.stream orelse return error.NotConnected;
|
|
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (job_name.len == 0 or job_name.len > 255) return error.JobNameTooLong;
|
|
if (patch_json.len == 0 or patch_json.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash:16]
|
|
// [job_name_len:1][job_name]
|
|
// [patch_len:2][patch_json]
|
|
const total_len = 1 + 16 + 1 + job_name.len + 2 + patch_json.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.set_run_narrative);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @as(u8, @intCast(job_name.len));
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @as(u16, @intCast(patch_json.len)), .big);
|
|
offset += 2;
|
|
@memcpy(buffer[offset .. offset + patch_json.len], patch_json);
|
|
|
|
try frame.sendWebSocketFrame(stream, buffer);
|
|
}
|
|
|
|
pub fn sendSetRunPrivacy(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
patch_json: []const u8,
|
|
api_key_hash: []const u8,
|
|
) !void {
|
|
const stream = self.stream orelse return error.NotConnected;
|
|
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (job_name.len == 0 or job_name.len > 255) return error.JobNameTooLong;
|
|
if (patch_json.len == 0 or patch_json.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash:16]
|
|
// [job_name_len:1][job_name]
|
|
// [patch_len:2][patch_json]
|
|
const total_len = 1 + 16 + 1 + job_name.len + 2 + patch_json.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.set_run_privacy);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @as(u8, @intCast(job_name.len));
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @as(u16, @intCast(patch_json.len)), .big);
|
|
offset += 2;
|
|
@memcpy(buffer[offset .. offset + patch_json.len], patch_json);
|
|
|
|
try frame.sendWebSocketFrame(stream, buffer);
|
|
}
|
|
|
|
pub fn sendAnnotateRun(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
author: []const u8,
|
|
note: []const u8,
|
|
api_key_hash: []const u8,
|
|
) !void {
|
|
const stream = self.stream orelse return error.NotConnected;
|
|
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (job_name.len == 0 or job_name.len > 255) return error.JobNameTooLong;
|
|
if (author.len > 255) return error.PayloadTooLarge;
|
|
if (note.len == 0 or note.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash:16]
|
|
// [job_name_len:1][job_name]
|
|
// [author_len:1][author]
|
|
// [note_len:2][note]
|
|
const total_len = 1 + 16 + 1 + job_name.len + 1 + author.len + 2 + note.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.annotate_run);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @as(u8, @intCast(job_name.len));
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = @as(u8, @intCast(author.len));
|
|
offset += 1;
|
|
if (author.len > 0) {
|
|
@memcpy(buffer[offset .. offset + author.len], author);
|
|
}
|
|
offset += author.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @as(u16, @intCast(note.len)), .big);
|
|
offset += 2;
|
|
@memcpy(buffer[offset .. offset + note.len], note);
|
|
|
|
try frame.sendWebSocketFrame(stream, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithArgsNoteAndResources(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
args: []const u8,
|
|
note: []const u8,
|
|
force: bool,
|
|
cpu: u8,
|
|
memory_gb: u8,
|
|
gpu: u8,
|
|
gpu_memory: ?[]const u8,
|
|
) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
if (args.len > 0xFFFF) return error.PayloadTooLarge;
|
|
if (note.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
const gpu_mem = gpu_memory orelse "";
|
|
if (gpu_mem.len > 255) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash]
|
|
// [commit_id]
|
|
// [priority]
|
|
// [job_name_len][job_name]
|
|
// [args_len:2][args]
|
|
// [note_len:2][note]
|
|
// [force:1]
|
|
// [cpu][memory_gb][gpu][gpu_mem_len][gpu_mem]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 2 + args.len + 2 + note.len + 1 + 4 + gpu_mem.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job_with_note);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = @intCast((args.len >> 8) & 0xFF);
|
|
buffer[offset + 1] = @intCast(args.len & 0xFF);
|
|
offset += 2;
|
|
|
|
if (args.len > 0) {
|
|
@memcpy(buffer[offset .. offset + args.len], args);
|
|
offset += args.len;
|
|
}
|
|
|
|
buffer[offset] = @intCast((note.len >> 8) & 0xFF);
|
|
buffer[offset + 1] = @intCast(note.len & 0xFF);
|
|
offset += 2;
|
|
|
|
if (note.len > 0) {
|
|
@memcpy(buffer[offset .. offset + note.len], note);
|
|
offset += note.len;
|
|
}
|
|
|
|
// Force flag
|
|
buffer[offset] = if (force) 0x01 else 0x00;
|
|
offset += 1;
|
|
|
|
buffer[offset] = cpu;
|
|
buffer[offset + 1] = memory_gb;
|
|
buffer[offset + 2] = gpu;
|
|
buffer[offset + 3] = @intCast(gpu_mem.len);
|
|
offset += 4;
|
|
|
|
if (gpu_mem.len > 0) {
|
|
@memcpy(buffer[offset .. offset + gpu_mem.len], gpu_mem);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithArgsAndResources(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
args: []const u8,
|
|
force: bool,
|
|
cpu: u8,
|
|
memory_gb: u8,
|
|
gpu: u8,
|
|
gpu_memory: ?[]const u8,
|
|
) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
if (args.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
const gpu_mem = gpu_memory orelse "";
|
|
if (gpu_mem.len > 255) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash]
|
|
// [commit_id]
|
|
// [priority]
|
|
// [job_name_len][job_name]
|
|
// [args_len:2][args]
|
|
// [force:1]
|
|
// [cpu][memory_gb][gpu][gpu_mem_len][gpu_mem]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 2 + args.len + 1 + 4 + gpu_mem.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job_with_args);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = @intCast((args.len >> 8) & 0xFF);
|
|
buffer[offset + 1] = @intCast(args.len & 0xFF);
|
|
offset += 2;
|
|
|
|
if (args.len > 0) {
|
|
@memcpy(buffer[offset .. offset + args.len], args);
|
|
offset += args.len;
|
|
}
|
|
|
|
// Force flag
|
|
buffer[offset] = if (force) 0x01 else 0x00;
|
|
offset += 1;
|
|
|
|
buffer[offset] = cpu;
|
|
buffer[offset + 1] = memory_gb;
|
|
buffer[offset + 2] = gpu;
|
|
buffer[offset + 3] = @intCast(gpu_mem.len);
|
|
offset += 4;
|
|
|
|
if (gpu_mem.len > 0) {
|
|
@memcpy(buffer[offset .. offset + gpu_mem.len], gpu_mem);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithSnapshotAndResources(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
snapshot_id: []const u8,
|
|
snapshot_sha256: []const u8,
|
|
cpu: u8,
|
|
memory_gb: u8,
|
|
gpu: u8,
|
|
gpu_memory: ?[]const u8,
|
|
) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
if (snapshot_id.len == 0 or snapshot_id.len > 255) return error.PayloadTooLarge;
|
|
if (snapshot_sha256.len == 0 or snapshot_sha256.len > 255) return error.PayloadTooLarge;
|
|
|
|
const gpu_mem = gpu_memory orelse "";
|
|
if (gpu_mem.len > 255) return error.PayloadTooLarge;
|
|
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 1 + snapshot_id.len + 1 + snapshot_sha256.len + 4 + gpu_mem.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job_with_snapshot);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = @intCast(snapshot_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + snapshot_id.len], snapshot_id);
|
|
offset += snapshot_id.len;
|
|
|
|
buffer[offset] = @intCast(snapshot_sha256.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + snapshot_sha256.len], snapshot_sha256);
|
|
offset += snapshot_sha256.len;
|
|
|
|
buffer[offset] = cpu;
|
|
buffer[offset + 1] = memory_gb;
|
|
buffer[offset + 2] = gpu;
|
|
buffer[offset + 3] = @intCast(gpu_mem.len);
|
|
offset += 4;
|
|
|
|
if (gpu_mem.len > 0) {
|
|
@memcpy(buffer[offset .. offset + gpu_mem.len], gpu_mem);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendValidateRequestTask(self: *Client, api_key_hash: []const u8, task_id: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (task_id.len == 0 or task_id.len > 255) return error.PayloadTooLarge;
|
|
|
|
const total_len = 1 + 16 + 1 + 1 + task_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.validate_request);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
buffer[offset] = @intFromEnum(opcode.ValidateTargetType.task_id);
|
|
offset += 1;
|
|
buffer[offset] = @intCast(task_id.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + task_id.len], task_id);
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJob(self: *Client, job_name: []const u8, commit_id: []const u8, priority: u8, api_key_hash: []const u8) !void {
|
|
// Validate input lengths
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [commit_id: 20 bytes] [priority: u8] [job_name_len: u8] [job_name: var]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset..], job_name);
|
|
|
|
// Send as WebSocket binary frame
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithResources(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
cpu: u8,
|
|
memory_gb: u8,
|
|
gpu: u8,
|
|
gpu_memory: ?[]const u8,
|
|
) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
|
|
const gpu_mem = gpu_memory orelse "";
|
|
if (gpu_mem.len > 255) return error.PayloadTooLarge;
|
|
|
|
// Tail encoding: [cpu:1][memory_gb:1][gpu:1][gpu_mem_len:1][gpu_mem:var]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 4 + gpu_mem.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = cpu;
|
|
buffer[offset + 1] = memory_gb;
|
|
buffer[offset + 2] = gpu;
|
|
buffer[offset + 3] = @intCast(gpu_mem.len);
|
|
offset += 4;
|
|
|
|
if (gpu_mem.len > 0) {
|
|
@memcpy(buffer[offset .. offset + gpu_mem.len], gpu_mem);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithTracking(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
tracking_json: []const u8,
|
|
) !void {
|
|
// Validate input lengths
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
if (tracking_json.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8]
|
|
// [api_key_hash: 16]
|
|
// [commit_id: 20]
|
|
// [priority: u8]
|
|
// [job_name_len: u8]
|
|
// [job_name: var]
|
|
// [tracking_json_len: u16]
|
|
// [tracking_json: var]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 2 + tracking_json.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job_with_tracking);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
// tracking_json length (big-endian)
|
|
buffer[offset] = @intCast((tracking_json.len >> 8) & 0xFF);
|
|
buffer[offset + 1] = @intCast(tracking_json.len & 0xFF);
|
|
offset += 2;
|
|
|
|
if (tracking_json.len > 0) {
|
|
@memcpy(buffer[offset .. offset + tracking_json.len], tracking_json);
|
|
}
|
|
|
|
// Single WebSocket frame for throughput
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueueJobWithTrackingAndResources(
|
|
self: *Client,
|
|
job_name: []const u8,
|
|
commit_id: []const u8,
|
|
priority: u8,
|
|
api_key_hash: []const u8,
|
|
tracking_json: []const u8,
|
|
cpu: u8,
|
|
memory_gb: u8,
|
|
gpu: u8,
|
|
gpu_memory: ?[]const u8,
|
|
) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (job_name.len > 255) return error.JobNameTooLong;
|
|
if (tracking_json.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
const gpu_mem = gpu_memory orelse "";
|
|
if (gpu_mem.len > 255) return error.PayloadTooLarge;
|
|
|
|
// [opcode]
|
|
// [api_key_hash]
|
|
// [commit_id]
|
|
// [priority]
|
|
// [job_name_len][job_name]
|
|
// [tracking_json_len:2][tracking_json]
|
|
// [cpu][memory_gb][gpu][gpu_mem_len][gpu_mem]
|
|
const total_len = 1 + 16 + 20 + 1 + 1 + job_name.len + 2 + tracking_json.len + 4 + gpu_mem.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.queue_job_with_tracking);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
buffer[offset] = priority;
|
|
offset += 1;
|
|
|
|
buffer[offset] = @intCast(job_name.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + job_name.len], job_name);
|
|
offset += job_name.len;
|
|
|
|
buffer[offset] = @intCast((tracking_json.len >> 8) & 0xFF);
|
|
buffer[offset + 1] = @intCast(tracking_json.len & 0xFF);
|
|
offset += 2;
|
|
|
|
if (tracking_json.len > 0) {
|
|
@memcpy(buffer[offset .. offset + tracking_json.len], tracking_json);
|
|
offset += tracking_json.len;
|
|
}
|
|
|
|
buffer[offset] = cpu;
|
|
buffer[offset + 1] = memory_gb;
|
|
buffer[offset + 2] = gpu;
|
|
buffer[offset + 3] = @intCast(gpu_mem.len);
|
|
offset += 4;
|
|
|
|
if (gpu_mem.len > 0) {
|
|
@memcpy(buffer[offset .. offset + gpu_mem.len], gpu_mem);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendCancelJob(self: *Client, job_name: []const u8, api_key_hash: []const u8) !void {
|
|
try validateApiKeyHash(api_key_hash);
|
|
try validateJobName(job_name);
|
|
|
|
var builder = try MessageBuilder.init(self.allocator, 1 + 16 + 1 + job_name.len);
|
|
defer builder.deinit();
|
|
|
|
builder.writeOpcode(opcode.cancel_job);
|
|
builder.writeBytes(api_key_hash);
|
|
builder.writeStringU8(job_name);
|
|
try builder.send(self.transport);
|
|
}
|
|
|
|
pub fn sendPrune(self: *Client, api_key_hash: []const u8, prune_type: u8, value: u32) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [prune_type: u8] [value: u4]
|
|
const total_len = 1 + 16 + 1 + 4;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.prune);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = prune_type;
|
|
offset += 1;
|
|
|
|
// Store value in big-endian format
|
|
buffer[offset] = @intCast((value >> 24) & 0xFF);
|
|
buffer[offset + 1] = @intCast((value >> 16) & 0xFF);
|
|
buffer[offset + 2] = @intCast((value >> 8) & 0xFF);
|
|
buffer[offset + 3] = @intCast(value & 0xFF);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendSyncRun(self: *Client, sync_json: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (sync_json.len > 0xFFFF) return error.PayloadTooLarge;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [json_len: u16] [json: var]
|
|
const total_len = 1 + 16 + 2 + sync_json.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.sync_run);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @intCast(sync_json.len), .big);
|
|
offset += 2;
|
|
|
|
if (sync_json.len > 0) {
|
|
@memcpy(buffer[offset .. offset + sync_json.len], sync_json);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendRerunRequest(self: *Client, run_id: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (run_id.len > 255) return error.PayloadTooLarge;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [run_id_len: u8] [run_id: var]
|
|
const total_len = 1 + 16 + 1 + run_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.rerun_request);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(run_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + run_id.len], run_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendQueryRunInfo(self: *Client, run_id: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (run_id.len > 255) return error.PayloadTooLarge;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [run_id_len: u8] [run_id: var]
|
|
const total_len = 1 + 16 + 1 + run_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.query_run_info);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(run_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + run_id.len], run_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendStatusRequest(self: *Client, api_key_hash: []const u8) !void {
|
|
try validateApiKeyHash(api_key_hash);
|
|
|
|
var builder = try MessageBuilder.init(self.allocator, 1 + 16);
|
|
defer builder.deinit();
|
|
|
|
builder.writeOpcode(opcode.status_request);
|
|
builder.writeBytes(api_key_hash);
|
|
try builder.send(self.transport);
|
|
}
|
|
|
|
pub fn receiveMessage(self: *Client, allocator: std.mem.Allocator) ![]u8 {
|
|
return frame.receiveBinaryMessage(self.transport, allocator);
|
|
}
|
|
|
|
/// Send raw binary message (alias for consistency with tasks.zig)
|
|
pub fn send(self: *Client, data: []const u8) !void {
|
|
return frame.sendWebSocketFrame(self.transport, data);
|
|
}
|
|
|
|
/// Read response (alias for receiveMessage for consistency)
|
|
pub fn read(self: *Client, allocator: std.mem.Allocator) ![]u8 {
|
|
return self.receiveMessage(allocator);
|
|
}
|
|
|
|
/// Receive and handle response with automatic display
|
|
pub fn receiveAndHandleResponse(self: *Client, allocator: std.mem.Allocator, operation: []const u8) !void {
|
|
const message = try self.receiveMessage(allocator);
|
|
defer allocator.free(message);
|
|
|
|
const packet = protocol.ResponsePacket.deserialize(message, allocator) catch {
|
|
// Fallback: treat as plain response.
|
|
std.debug.print("Server response: {s}\n", .{message});
|
|
return;
|
|
};
|
|
defer packet.deinit(allocator);
|
|
|
|
try response_handlers.handleResponsePacket(self, packet, operation);
|
|
}
|
|
|
|
pub fn receiveAndHandleStatusResponse(self: *Client, allocator: std.mem.Allocator, user_context: anytype, options: anytype) !void {
|
|
return response_handlers.receiveAndHandleStatusResponse(self, allocator, user_context, options);
|
|
}
|
|
|
|
pub fn receiveAndHandleCancelResponse(self: *Client, allocator: std.mem.Allocator, user_context: anytype, job_name: []const u8, options: anytype) !void {
|
|
return response_handlers.receiveAndHandleCancelResponse(self, allocator, user_context, job_name, options);
|
|
}
|
|
|
|
pub fn handleResponsePacket(self: *Client, packet: protocol.ResponsePacket, operation: []const u8) !void {
|
|
return response_handlers.handleResponsePacket(self, packet, operation);
|
|
}
|
|
|
|
fn convertServerError(self: *Client, server_error: protocol.ErrorCode) anyerror {
|
|
_ = self;
|
|
return switch (server_error) {
|
|
.authentication_failed => error.AuthenticationFailed,
|
|
.permission_denied => error.PermissionDenied,
|
|
.resource_not_found => error.JobNotFound,
|
|
.resource_already_exists => error.ResourceExists,
|
|
.timeout => error.RequestTimeout,
|
|
.server_overloaded, .service_unavailable => error.ServerUnreachable,
|
|
.invalid_request => error.InvalidArguments,
|
|
.job_not_found => error.JobNotFound,
|
|
.job_already_running => error.JobAlreadyRunning,
|
|
.job_failed_to_start, .job_execution_failed => error.CommandFailed,
|
|
.job_cancelled => error.JobCancelled,
|
|
else => error.ServerError,
|
|
};
|
|
}
|
|
|
|
pub fn sendCrashReport(self: *Client, api_key_hash: []const u8, error_type: []const u8, error_message: []const u8, command: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][error_type_len:2][error_type][error_message_len:2][error_message][command_len:2][command]
|
|
const total_len = 1 + 16 + 2 + error_type.len + 2 + error_message.len + 2 + command.len;
|
|
const message = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(message);
|
|
|
|
var offset: usize = 0;
|
|
|
|
// opcode
|
|
message[offset] = @intFromEnum(opcode.crash_report);
|
|
offset += 1;
|
|
|
|
// API key hash
|
|
@memcpy(message[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
// Error type length and data
|
|
std.mem.writeInt(u16, message[offset .. offset + 2][0..2], @intCast(error_type.len), .big);
|
|
offset += 2;
|
|
@memcpy(message[offset .. offset + error_type.len], error_type);
|
|
offset += error_type.len;
|
|
|
|
// Error message length and data
|
|
std.mem.writeInt(u16, message[offset .. offset + 2][0..2], @intCast(error_message.len), .big);
|
|
offset += 2;
|
|
@memcpy(message[offset .. offset + error_message.len], error_message);
|
|
offset += error_message.len;
|
|
|
|
// Command length and data
|
|
std.mem.writeInt(u16, message[offset .. offset + 2][0..2], @intCast(command.len), .big);
|
|
offset += 2;
|
|
@memcpy(message[offset .. offset + command.len], command);
|
|
|
|
// Send WebSocket frame over transport
|
|
try frame.sendWebSocketFrame(self.transport, message);
|
|
}
|
|
|
|
// Dataset management methods
|
|
pub fn sendDatasetList(self: *Client, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message: [opcode: u8] [api_key_hash: 16 bytes]
|
|
const total_len = 1 + 16;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
buffer[0] = @intFromEnum(opcode.dataset_list);
|
|
@memcpy(buffer[1..17], api_key_hash);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendDatasetRegister(self: *Client, name: []const u8, url: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (name.len > 255) return error.NameTooLong;
|
|
if (url.len > 1023) return error.URLTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [name_len: u8] [name: var] [url_len: u16] [url: var]
|
|
const total_len = 1 + 16 + 1 + name.len + 2 + url.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.dataset_register);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + name.len], name);
|
|
offset += name.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @intCast(url.len), .big);
|
|
offset += 2;
|
|
|
|
@memcpy(buffer[offset .. offset + url.len], url);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
// Jupyter management methods
|
|
pub fn sendStartJupyter(self: *Client, name: []const u8, workspace: []const u8, password: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (name.len > 255) return error.NameTooLong;
|
|
if (workspace.len > 65535) return error.WorkspacePathTooLong;
|
|
if (password.len > 255) return error.PasswordTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode:1][api_key_hash:16][name_len:1][name:var][workspace_len:2][workspace:var][password_len:1][password:var]
|
|
const total_len = 1 + 16 + 1 + name.len + 2 + workspace.len + 1 + password.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.start_jupyter);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + name.len], name);
|
|
offset += name.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @intCast(workspace.len), .big);
|
|
offset += 2;
|
|
@memcpy(buffer[offset .. offset + workspace.len], workspace);
|
|
offset += workspace.len;
|
|
|
|
buffer[offset] = @intCast(password.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + password.len], password);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendStopJupyter(self: *Client, service_id: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (service_id.len > 255) return error.InvalidServiceId;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][service_id_len:1][service_id:var]
|
|
const total_len = 1 + 16 + 1 + service_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.stop_jupyter);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(service_id.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + service_id.len], service_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendRemoveJupyter(self: *Client, service_id: []const u8, api_key_hash: []const u8, purge: bool) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (service_id.len > 255) return error.InvalidServiceId;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][service_id_len:1][service_id:var][purge:1]
|
|
const total_len = 1 + 16 + 1 + service_id.len + 1;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.remove_jupyter);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(service_id.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + service_id.len], service_id);
|
|
offset += service_id.len;
|
|
|
|
buffer[offset] = if (purge) 0x01 else 0x00;
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendRestoreJupyter(self: *Client, name: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (name.len > 255) return error.NameTooLong;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][name_len:1][name:var]
|
|
const total_len = 1 + 16 + 1 + name.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.restore_jupyter);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + name.len], name);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendListJupyter(self: *Client, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16]
|
|
const total_len = 1 + 16;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
buffer[0] = @intFromEnum(opcode.list_jupyter);
|
|
@memcpy(buffer[1..17], api_key_hash);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendDatasetInfo(self: *Client, name: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (name.len > 255) return error.NameTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [name_len: u8] [name: var]
|
|
const total_len = 1 + 16 + 1 + name.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.dataset_info);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset..], name);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendDatasetSearch(self: *Client, term: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message: [opcode: u8] [api_key_hash: 16 bytes] [term_len: u8] [term: var]
|
|
const total_len = 1 + 16 + 1 + term.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.dataset_search);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(term.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset..], term);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendLogMetric(self: *Client, api_key_hash: []const u8, commit_id: []const u8, name: []const u8, value: f64, step: u32) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
if (name.len > 255) return error.NameTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [commit_id: 20 bytes] [step: u32] [value: f64] [name_len: u8] [name: var]
|
|
const total_len = 1 + 16 + 20 + 4 + 8 + 1 + name.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.log_metric);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
std.mem.writeInt(u32, buffer[offset .. offset + 4][0..4], step, .big);
|
|
offset += 4;
|
|
|
|
std.mem.writeInt(u64, buffer[offset .. offset + 8][0..8], @as(u64, @bitCast(value)), .big);
|
|
offset += 8;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset..], name);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendGetExperiment(self: *Client, api_key_hash: []const u8, commit_id: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (commit_id.len != 20) return error.InvalidCommitId;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [commit_id: 20 bytes]
|
|
const total_len = 1 + 16 + 20;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.get_experiment);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
@memcpy(buffer[offset .. offset + 20], commit_id);
|
|
offset += 20;
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendCreateExperiment(self: *Client, api_key_hash: []const u8, name: []const u8, description: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (name.len == 0 or name.len > 255) return error.NameTooLong;
|
|
if (description.len > 1023) return error.DescriptionTooLong;
|
|
|
|
// Build binary message:
|
|
// [opcode: u8] [api_key_hash: 16 bytes] [name_len: u8] [name: var] [desc_len: u16] [description: var]
|
|
const total_len = 1 + 16 + 1 + name.len + 2 + description.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.create_experiment);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(name.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + name.len], name);
|
|
offset += name.len;
|
|
|
|
std.mem.writeInt(u16, buffer[offset .. offset + 2][0..2], @intCast(description.len), .big);
|
|
offset += 2;
|
|
if (description.len > 0) {
|
|
@memcpy(buffer[offset .. offset + description.len], description);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendListExperiments(self: *Client, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
|
|
// Build binary message: [opcode: u8] [api_key_hash: 16 bytes]
|
|
const total_len = 1 + 16;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
buffer[0] = @intFromEnum(opcode.list_experiments);
|
|
@memcpy(buffer[1..17], api_key_hash);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendGetExperimentByID(self: *Client, api_key_hash: []const u8, experiment_id: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (experiment_id.len == 0 or experiment_id.len > 255) return error.InvalidExperimentId;
|
|
|
|
// Build binary message: [opcode: u8] [api_key_hash: 16 bytes] [exp_id_len: u8] [experiment_id: var]
|
|
const total_len = 1 + 16 + 1 + experiment_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.get_experiment);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(experiment_id.len);
|
|
offset += 1;
|
|
@memcpy(buffer[offset .. offset + experiment_id.len], experiment_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
// Logs and debug methods
|
|
pub fn sendGetLogs(self: *Client, target_id: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (target_id.len == 0 or target_id.len > 255) return error.InvalidTargetId;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][target_id_len:1][target_id:var]
|
|
const total_len = 1 + 16 + 1 + target_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.get_logs);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(target_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + target_id.len], target_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendStreamLogs(self: *Client, target_id: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (target_id.len == 0 or target_id.len > 255) return error.InvalidTargetId;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][target_id_len:1][target_id:var]
|
|
const total_len = 1 + 16 + 1 + target_id.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.stream_logs);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(target_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + target_id.len], target_id);
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
pub fn sendAttachDebug(self: *Client, target_id: []const u8, debug_type: []const u8, api_key_hash: []const u8) !void {
|
|
if (api_key_hash.len != 16) return error.InvalidApiKeyHash;
|
|
if (target_id.len == 0 or target_id.len > 255) return error.InvalidTargetId;
|
|
if (debug_type.len > 255) return error.InvalidDebugType;
|
|
|
|
// Build binary message: [opcode:1][api_key_hash:16][target_id_len:1][target_id:var][debug_type:var]
|
|
const total_len = 1 + 16 + 1 + target_id.len + debug_type.len;
|
|
var buffer = try self.allocator.alloc(u8, total_len);
|
|
defer self.allocator.free(buffer);
|
|
|
|
var offset: usize = 0;
|
|
buffer[offset] = @intFromEnum(opcode.attach_debug);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + 16], api_key_hash);
|
|
offset += 16;
|
|
|
|
buffer[offset] = @intCast(target_id.len);
|
|
offset += 1;
|
|
|
|
@memcpy(buffer[offset .. offset + target_id.len], target_id);
|
|
offset += target_id.len;
|
|
|
|
if (debug_type.len > 0) {
|
|
@memcpy(buffer[offset .. offset + debug_type.len], debug_type);
|
|
}
|
|
|
|
try frame.sendWebSocketFrame(self.transport, buffer);
|
|
}
|
|
|
|
/// Receive and handle dataset response
|
|
pub fn receiveAndHandleDatasetResponse(self: *Client, allocator: std.mem.Allocator) ![]const u8 {
|
|
const message = try self.receiveMessage(allocator);
|
|
defer allocator.free(message);
|
|
|
|
const packet = protocol.ResponsePacket.deserialize(message, allocator) catch {
|
|
// Fallback: treat as plain response.
|
|
return allocator.dupe(u8, message);
|
|
};
|
|
defer packet.deinit(allocator);
|
|
|
|
switch (packet.packet_type) {
|
|
.data => {
|
|
if (packet.data_payload) |payload| {
|
|
return allocator.dupe(u8, payload);
|
|
}
|
|
return allocator.dupe(u8, "");
|
|
},
|
|
.success => {
|
|
if (packet.success_message) |msg| {
|
|
return allocator.dupe(u8, msg);
|
|
}
|
|
return allocator.dupe(u8, "");
|
|
},
|
|
.error_packet => {
|
|
// Print details and raise appropriate CLI error.
|
|
_ = response_handlers.handleResponsePacket(self, packet, "Dataset") catch {};
|
|
return self.convertServerError(packet.error_code.?);
|
|
},
|
|
else => {
|
|
// Unexpected packet type.
|
|
return error.UnexpectedResponse;
|
|
},
|
|
}
|
|
}
|
|
};
|