zoe-agent 0.3.1

package/dist/adapters/server/ws-handlers.js

@@ -0,0 +1,523 @@
+/**
+ * Zoe Server — WebSocket Protocol Handlers
+ *
+ * All handler functions, safeSend helper, and active connections registry.
+ * Extracted from websocket.ts for single-responsibility.
+ */
+import * as crypto from "crypto";
+import { authMiddleware, hasScope } from "./auth.js";
+import { hashKey } from "./session-store.js";
+import { handleWsGetSettings, handleWsUpdateSettings } from "./settings-handlers.js";
+// ── Active connections registry ──────────────────────────────────────
+const activeConnections = new Map();
+// ── Pending tool approvals ───────────────────────────────────────────
+const pendingApprovals = new Map();
+/**
+ * Get the number of currently active WebSocket connections.
+ */
+export function getActiveConnectionCount() {
+    return activeConnections.size;
+}
+/**
+ * Get all connected WS clients (excluding the given one).
+ * Used by settings broadcast to notify other connections of changes.
+ */
+export function getOtherClients(excludeWs) {
+    const clients = [];
+    for (const [ws, state] of activeConnections) {
+        if (ws !== excludeWs) {
+            clients.push({ ws, state });
+        }
+    }
+    return clients;
+}
+// ── Send helper ──────────────────────────────────────────────────────
+export function safeSend(ws, message) {
+    try {
+        if (ws.readyState === 1 /* OPEN */) {
+            ws.send(JSON.stringify(message));
+        }
+    }
+    catch {
+        // Connection may have closed
+    }
+}
+// ── Protocol handler ─────────────────────────────────────────────────
+export function handleConnection(ws, req, ctx) {
+    // Auth check — the token should have been validated during upgrade,
+    // but verify again for safety
+    const key = authMiddleware(req);
+    if (!key) {
+        safeSend(ws, {
+            type: "error",
+            code: "UNAUTHORIZED",
+            retryable: false,
+            message: "Authentication required",
+        });
+        ws.close(4001, "Unauthorized");
+        return;
+    }
+    const state = {
+        sessionId: null,
+        currentAbortController: null,
+        activeProvider: null,
+        activeModel: null,
+        maxPermissionLevel: ctx.maxPermissionLevel,
+        apiKeyHash: hashKey(key.key),
+        apiKey: key,
+    };
+    activeConnections.set(ws, state);
+    // ── Message dispatch ───────────────────────────────────────────────
+    ws.on("message", (data) => {
+        let msg;
+        try {
+            msg = JSON.parse(data.toString("utf-8"));
+        }
+        catch {
+            safeSend(ws, {
+                type: "error",
+                code: "INVALID_MESSAGE",
+                retryable: false,
+                message: "Invalid JSON message",
+            });
+            return;
+        }
+        switch (msg.type) {
+            case "chat":
+                handleChat(ws, msg, state, ctx);
+                break;
+            case "abort":
+                handleAbort(ws, msg, state);
+                break;
+            case "tool_approval_response":
+                handleToolApprovalResponse(ws, msg);
+                break;
+            case "resume":
+                void handleResume(ws, msg, state, ctx);
+                break;
+            case "reconnect":
+                void handleReconnect(ws, msg, state, ctx);
+                break;
+            case "switch_provider":
+                handleSwitchProvider(ws, msg, state);
+                break;
+            case "list_models":
+                handleListModels(ws, ctx);
+                break;
+            case "list_skills":
+                handleListSkills(ws, ctx);
+                break;
+            case "ping":
+                safeSend(ws, {
+                    type: "pong",
+                    serverTime: new Date().toISOString(),
+                });
+                break;
+            case "get_settings":
+                handleWsSettingsMessage(ws, msg, state, ctx, (sCtx) => handleWsGetSettings(msg, ws, state, sCtx));
+                break;
+            case "update_settings":
+                handleWsSettingsMessage(ws, msg, state, ctx, (sCtx) => void handleWsUpdateSettings(msg, ws, state, sCtx));
+                break;
+            case "list_providers":
+                handleWsSettingsMessage(ws, msg, state, ctx, (sCtx) => handleWsListProviders(msg, ws, state, sCtx));
+                break;
+            case "set_provider":
+                handleWsSettingsMessage(ws, msg, state, ctx, (sCtx) => void handleWsSetProvider(msg, ws, state, sCtx));
+                break;
+            case "remove_provider":
+                handleWsSettingsMessage(ws, msg, state, ctx, (sCtx) => void handleWsRemoveProvider(msg, ws, state, sCtx));
+                break;
+            default:
+                safeSend(ws, {
+                    type: "error",
+                    code: "UNKNOWN_MESSAGE_TYPE",
+                    retryable: false,
+                    message: `Unknown message type: ${msg.type}`,
+                });
+        }
+    });
+    // ── Close ──────────────────────────────────────────────────────────
+    ws.on("close", () => {
+        // Abort any in-flight stream
+        if (state.currentAbortController) {
+            state.currentAbortController.abort();
+            state.currentAbortController = null;
+        }
+        activeConnections.delete(ws);
+    });
+    // ── Error ──────────────────────────────────────────────────────────
+    ws.on("error", (err) => {
+        console.error("[ws] Connection error:", err.message);
+        if (state.currentAbortController) {
+            state.currentAbortController.abort();
+            state.currentAbortController = null;
+        }
+        activeConnections.delete(ws);
+    });
+}
+// ── Chat handler ─────────────────────────────────────────────────────
+function handleChat(ws, msg, state, ctx) {
+    const serverMsgId = crypto.randomUUID();
+    // Acknowledge
+    safeSend(ws, {
+        type: "ack",
+        clientMsgId: msg.id,
+        serverMsgId,
+        timestamp: new Date().toISOString(),
+    });
+    // Create session if needed
+    if (!state.sessionId && msg.sessionId) {
+        state.sessionId = msg.sessionId;
+    }
+    // Set up abort controller
+    const abortController = new AbortController();
+    state.currentAbortController = abortController;
+    // Resolve options with connection-level overrides
+    const provider = msg.options?.provider ?? state.activeProvider ?? undefined;
+    const model = msg.options?.model ?? state.activeModel ?? undefined;
+    // Resolve permission level with server ceiling
+    let effectivePermissionLevel = msg.options?.permissionLevel ?? state.permissionLevel;
+    if (effectivePermissionLevel && state.maxPermissionLevel) {
+        const levels = ["strict", "moderate", "permissive"];
+        const maxIdx = levels.indexOf(state.maxPermissionLevel);
+        const reqIdx = levels.indexOf(effectivePermissionLevel);
+        // QA-009: Unknown levels (-1) are capped to the server ceiling
+        if (reqIdx === -1 || reqIdx > maxIdx) {
+            effectivePermissionLevel = state.maxPermissionLevel;
+        }
+    }
+    // Stream text
+    try {
+        ctx.streamText({
+            message: msg.message,
+            model,
+            provider,
+            tools: msg.options?.tools,
+            maxSteps: msg.options?.maxSteps ?? 10,
+            skills: msg.options?.skills,
+            sessionId: state.sessionId ?? undefined,
+            permissionLevel: effectivePermissionLevel,
+            approveTool: createServerApproveTool(ws),
+            signal: abortController.signal,
+            onText: (delta) => {
+                safeSend(ws, {
+                    type: "text",
+                    delta,
+                    serverMsgId,
+                });
+            },
+            onToolCall: (info) => {
+                safeSend(ws, {
+                    type: "tool_call",
+                    callId: info.callId,
+                    name: info.name,
+                    args: info.args,
+                });
+            },
+            onToolResult: (info) => {
+                safeSend(ws, {
+                    type: "tool_result",
+                    callId: info.callId,
+                    output: info.output,
+                    success: info.success,
+                });
+            },
+            onStep: (step) => {
+                // Estimate progress — we don't know totalSteps ahead of time
+                safeSend(ws, {
+                    type: "progress",
+                    step: 0,
+                    totalSteps: 0,
+                    percentage: 0,
+                    activity: step.content ?? step.type,
+                });
+            },
+            onError: (error) => {
+                safeSend(ws, {
+                    type: "error",
+                    code: error.code || "STREAM_ERROR",
+                    retryable: error.code === "PROVIDER_ERROR",
+                    message: error.message,
+                    provider: error.provider,
+                    tool: error.tool,
+                });
+            },
+            onDone: (result) => {
+                safeSend(ws, {
+                    type: "done",
+                    serverMsgId,
+                    usage: result.usage,
+                    finishReason: result.finishReason,
+                });
+                // Add assistant message to session
+                if (state.sessionId) {
+                    const assistantMsg = {
+                        id: serverMsgId,
+                        role: "assistant",
+                        content: result.text,
+                        timestamp: Date.now(),
+                    };
+                    ctx.sessionManager.addMessage(state.sessionId, assistantMsg);
+                }
+                state.currentAbortController = null;
+            },
+        });
+    }
+    catch (err) {
+        const message = err instanceof Error ? err.message : "Stream failed";
+        safeSend(ws, {
+            type: "error",
+            code: "STREAM_ERROR",
+            retryable: false,
+            message,
+        });
+        state.currentAbortController = null;
+    }
+}
+// ── Abort handler ────────────────────────────────────────────────────
+function handleAbort(ws, _msg, state) {
+    if (state.currentAbortController) {
+        state.currentAbortController.abort();
+        state.currentAbortController = null;
+        safeSend(ws, {
+            type: "error",
+            code: "ABORTED",
+            retryable: false,
+            message: "Request aborted by client",
+        });
+    }
+}
+// ── Tool approval handler ────────────────────────────────────────────
+const APPROVAL_TIMEOUT_MS = 30_000; // 30 seconds
+/**
+ * Create an `approveTool` callback for the server adapter.
+ * Sends a `tool_approval_request` to the client and waits for a
+ * `tool_approval_response`. Falls back to auto-deny on timeout.
+ */
+export function createServerApproveTool(ws) {
+    return async (call) => {
+        const callId = crypto.randomUUID();
+        safeSend(ws, {
+            type: "tool_approval_request",
+            callId,
+            name: call.name,
+            args: call.args,
+        });
+        return new Promise((resolve) => {
+            const timer = setTimeout(() => {
+                pendingApprovals.delete(callId);
+                resolve(false); // Timeout → deny
+            }, APPROVAL_TIMEOUT_MS);
+            pendingApprovals.set(callId, { resolve, timer, ws, toolName: call.name, createdAt: Date.now() });
+        });
+    };
+}
+function handleToolApprovalResponse(ws, msg) {
+    const pending = pendingApprovals.get(msg.callId);
+    if (!pending)
+        return;
+    // QA-001: Only the originating connection may resolve the approval
+    if (pending.ws !== ws)
+        return;
+    // Defense-in-depth: verify the tool name matches the pending request
+    if (msg.name !== pending.toolName)
+        return;
+    // Reject expired approvals (defense-in-depth, timer should have fired)
+    if (Date.now() - pending.createdAt > APPROVAL_TIMEOUT_MS) {
+        clearTimeout(pending.timer);
+        pendingApprovals.delete(msg.callId);
+        pending.resolve(false);
+        return;
+    }
+    clearTimeout(pending.timer);
+    pendingApprovals.delete(msg.callId);
+    pending.resolve(msg.approved);
+}
+// ── Resume handler ───────────────────────────────────────────────────
+async function handleResume(ws, msg, state, ctx) {
+    const session = await ctx.sessionManager.getSession(msg.sessionId, state.apiKeyHash);
+    if (!session) {
+        safeSend(ws, {
+            type: "error",
+            code: "SESSION_NOT_FOUND",
+            retryable: false,
+            message: `Session ${msg.sessionId} not found or expired`,
+        });
+        return;
+    }
+    state.sessionId = msg.sessionId;
+    safeSend(ws, {
+        type: "session_resumed",
+        sessionId: msg.sessionId,
+        messages: session.messages,
+    });
+}
+// ── Reconnect handler ────────────────────────────────────────────────
+async function handleReconnect(ws, msg, state, ctx) {
+    const session = await ctx.sessionManager.getSession(msg.sessionId, state.apiKeyHash);
+    if (!session) {
+        safeSend(ws, {
+            type: "error",
+            code: "SESSION_NOT_FOUND",
+            retryable: false,
+            message: `Session ${msg.sessionId} not found or expired`,
+        });
+        return;
+    }
+    state.sessionId = msg.sessionId;
+    // Replay messages — optionally only those after lastSeenId
+    let messages = session.messages;
+    if (msg.lastSeenId) {
+        const lastIndex = messages.findIndex((m) => m.id === msg.lastSeenId);
+        if (lastIndex !== -1) {
+            messages = messages.slice(lastIndex + 1);
+        }
+    }
+    safeSend(ws, {
+        type: "replay",
+        messages,
+        currentStatus: "ready",
+    });
+}
+// ── Switch provider handler ──────────────────────────────────────────
+function handleSwitchProvider(ws, msg, state) {
+    state.activeProvider = msg.provider;
+    if (msg.model) {
+        state.activeModel = msg.model;
+    }
+    safeSend(ws, {
+        type: "ack",
+        clientMsgId: "",
+        serverMsgId: crypto.randomUUID(),
+        timestamp: new Date().toISOString(),
+    });
+}
+// ── List models handler ──────────────────────────────────────────────
+function handleListModels(ws, ctx) {
+    safeSend(ws, {
+        type: "models_list",
+        models: ctx.listModels(),
+    });
+}
+// ── List skills handler ──────────────────────────────────────────────
+function handleListSkills(ws, ctx) {
+    safeSend(ws, {
+        type: "skills_list",
+        skills: ctx.listSkills(),
+    });
+}
+// ── Settings dispatch helper ────────────────────────────────────────────
+function handleWsSettingsMessage(ws, _msg, _state, ctx, fn) {
+    const sCtx = ctx.settingsHandlerContext;
+    if (!sCtx) {
+        safeSend(ws, {
+            type: "error",
+            code: "SERVICE_UNAVAILABLE",
+            retryable: false,
+            message: "Settings not configured",
+        });
+        return;
+    }
+    fn(sCtx);
+}
+// ── WS scope helper ─────────────────────────────────────────────────────
+function requireWsScope(state, scope) {
+    return !!state.apiKey && hasScope(state.apiKey, scope);
+}
+// ── WS Settings: list providers ─────────────────────────────────────────
+function handleWsListProviders(msg, ws, state, ctx) {
+    if (!requireWsScope(state, "agent:read")) {
+        safeSend(ws, { type: "providers_list", id: msg.id, providers: {}, error: { code: "FORBIDDEN", message: "Requires agent:read scope" } });
+        return;
+    }
+    const providers = {};
+    for (const pType of ["openai", "anthropic", "glm", "openai-compatible"]) {
+        const prefix = `providers.${pType === "openai-compatible" ? "openai-compat" : pType}`;
+        const apiKeyVal = ctx.settingsManager.get(`${prefix}.apiKey`).value;
+        if (apiKeyVal != null) {
+            providers[pType] = { type: pType };
+            const model = ctx.settingsManager.get(`${prefix}.model`).value;
+            if (model)
+                providers[pType].model = model;
+            if (pType === "openai-compatible") {
+                const baseUrl = ctx.settingsManager.get(`${prefix}.baseUrl`).value;
+                if (baseUrl)
+                    providers[pType].baseUrl = baseUrl;
+            }
+        }
+    }
+    safeSend(ws, { type: "providers_list", id: msg.id, providers });
+}
+// ── WS Settings: set provider ───────────────────────────────────────────
+async function handleWsSetProvider(msg, ws, state, ctx) {
+    if (!requireWsScope(state, "admin")) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "FORBIDDEN", message: "Requires admin scope" } });
+        return;
+    }
+    const { type: providerType, apiKey, baseUrl, model } = msg.provider;
+    const VALID_PROVIDER_TYPES = new Set(["openai", "anthropic", "glm", "openai-compatible"]);
+    if (!providerType || !VALID_PROVIDER_TYPES.has(providerType)) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "VALIDATION_ERROR", message: "Invalid provider type" } });
+        return;
+    }
+    const prefix = `providers.${providerType === "openai-compatible" ? "openai-compat" : providerType}`;
+    try {
+        if (apiKey)
+            await ctx.settingsManager.set(`${prefix}.apiKey`, apiKey);
+        if (model)
+            await ctx.settingsManager.set(`${prefix}.model`, model);
+        if (baseUrl && providerType === "openai-compatible")
+            await ctx.settingsManager.set(`${prefix}.baseUrl`, baseUrl);
+    }
+    catch (e) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "SET_ERROR", message: e.message } });
+        return;
+    }
+    safeSend(ws, { type: "settings_updated", id: msg.id, applied: { [providerType]: true }, requiresRestart: false, restartAffected: [] });
+}
+// ── WS Settings: remove provider ────────────────────────────────────────
+async function handleWsRemoveProvider(msg, ws, state, ctx) {
+    if (!requireWsScope(state, "admin")) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "FORBIDDEN", message: "Requires admin scope" } });
+        return;
+    }
+    const VALID_PROVIDER_TYPES = new Set(["openai", "anthropic", "glm", "openai-compatible"]);
+    if (!msg.providerType || !VALID_PROVIDER_TYPES.has(msg.providerType)) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "VALIDATION_ERROR", message: "Invalid provider type" } });
+        return;
+    }
+    const prefix = `providers.${msg.providerType === "openai-compatible" ? "openai-compat" : msg.providerType}`;
+    try {
+        await ctx.settingsManager.reset(`${prefix}.apiKey`);
+        try {
+            await ctx.settingsManager.reset(`${prefix}.model`);
+        }
+        catch { /* may not exist */ }
+        try {
+            await ctx.settingsManager.reset(`${prefix}.baseUrl`);
+        }
+        catch { /* may not exist */ }
+    }
+    catch (e) {
+        safeSend(ws, { type: "settings_updated", id: msg.id, error: { code: "RESET_ERROR", message: e.message } });
+        return;
+    }
+    safeSend(ws, { type: "settings_updated", id: msg.id, applied: { removed: msg.providerType }, requiresRestart: false, restartAffected: [] });
+}
+// ── Active connections accessor (for closeWebSocket) ──────────────────
+/**
+ * Close all active connections and clear the registry.
+ * Used by closeWebSocket() during shutdown.
+ */
+export function closeAllConnections() {
+    for (const [ws] of activeConnections) {
+        try {
+            ws.close(1001, "Server shutting down");
+        }
+        catch {
+            // Ignore errors during shutdown
+        }
+    }
+    activeConnections.clear();
+}