From cba3ad2721604c72ab8163609d7075287c611f56 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 21 Jan 2026 08:16:11 +0700 Subject: [PATCH] chore: checkpoint --- packages/relay/src/live-relay.e2e.test.ts | 172 ++++++ packages/relay/wrangler.toml | 19 + .../server/src/client/daemon-client-v2.ts | 33 ++ .../codex-cli-full-access-ipc.test.ts | 92 ++++ .../codex-mcp-agent-commands.e2e.test.ts | 121 +++++ .../codex-mcp-agent.commands.test.ts | 38 ++ .../codex-mcp-agent.resume-errors.test.ts | 19 + .../server/agent/providers/codex-mcp-agent.ts | 295 +++++++++- .../daemon-restart-resume.e2e.test.ts | 244 +++++++++ packages/server/src/server/session.ts | 82 +++ .../server/src/server/websocket-server.ts | 508 +----------------- .../src/server/websocket-session-bridge.ts | 434 +++++++++++++++ packages/server/src/shared/messages.ts | 28 + 13 files changed, 1597 insertions(+), 488 deletions(-) create mode 100644 packages/relay/src/live-relay.e2e.test.ts create mode 100644 packages/relay/wrangler.toml create mode 100644 packages/server/src/server/agent/providers/codex-cli-full-access-ipc.test.ts create mode 100644 packages/server/src/server/agent/providers/codex-mcp-agent-commands.e2e.test.ts create mode 100644 packages/server/src/server/agent/providers/codex-mcp-agent.commands.test.ts create mode 100644 packages/server/src/server/agent/providers/codex-mcp-agent.resume-errors.test.ts create mode 100644 packages/server/src/server/daemon-e2e/daemon-restart-resume.e2e.test.ts create mode 100644 packages/server/src/server/websocket-session-bridge.ts diff --git a/packages/relay/src/live-relay.e2e.test.ts b/packages/relay/src/live-relay.e2e.test.ts new file mode 100644 index 000000000..1dfd11b2d --- /dev/null +++ b/packages/relay/src/live-relay.e2e.test.ts @@ -0,0 +1,172 @@ +import { describe, it, expect } from "vitest"; +import { WebSocket } from "ws"; +import { + generateKeyPair, + exportPublicKey, + importPublicKey, + deriveSharedKey, + encrypt, + decrypt, +} from "./crypto.js"; + +const RELAY_BASE_URL = "wss://relay.paseo.sh"; + +async function withRetry( + fn: () => Promise, + options: { retries: number; delayMs: number } +): Promise { + let lastError: unknown; + for (let attempt = 0; attempt <= options.retries; attempt++) { + try { + return await fn(); + } catch (error) { + lastError = error; + if (attempt < options.retries) { + await new Promise((r) => setTimeout(r, options.delayMs)); + } + } + } + throw lastError instanceof Error ? lastError : new Error(String(lastError)); +} + +describe("Live relay (relay.paseo.sh) E2E", () => { + it("bridges encrypted traffic end-to-end", { timeout: 45_000 }, async () => { + await withRetry( + async () => { + const sessionId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`; + const serverUrl = `${RELAY_BASE_URL}/ws?session=${encodeURIComponent(sessionId)}&role=server`; + const clientUrl = `${RELAY_BASE_URL}/ws?session=${encodeURIComponent(sessionId)}&role=client`; + + // === Key setup === + const daemonKeyPair = await generateKeyPair(); + const daemonPubKeyB64 = await exportPublicKey(daemonKeyPair.publicKey); + + const clientKeyPair = await generateKeyPair(); + const clientPubKeyB64 = await exportPublicKey(clientKeyPair.publicKey); + + const daemonPubKeyOnClient = await importPublicKey(daemonPubKeyB64); + const clientSharedKey = await deriveSharedKey( + clientKeyPair.privateKey, + daemonPubKeyOnClient + ); + + // === Connect === + const daemonWs = new WebSocket(serverUrl); + const clientWs = new WebSocket(clientUrl); + + const waitOpen = (ws: WebSocket, label: string) => + new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error(`Timed out opening ${label} websocket`)), + 10_000 + ); + ws.once("open", () => { + clearTimeout(timeout); + resolve(); + }); + ws.once("error", (err) => { + clearTimeout(timeout); + reject(err); + }); + }); + + try { + await Promise.all([waitOpen(daemonWs, "server"), waitOpen(clientWs, "client")]); + + // === Handshake === + // Client sends hello with its public key (not encrypted). + clientWs.send(JSON.stringify({ type: "hello", key: clientPubKeyB64 })); + + const daemonReceivedHello = await new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for hello")), + 10_000 + ); + daemonWs.once("message", (data) => { + clearTimeout(timeout); + resolve(data.toString()); + }); + }); + + const hello = JSON.parse(daemonReceivedHello) as { + type: string; + key?: string; + }; + expect(hello.type).toBe("hello"); + expect(typeof hello.key).toBe("string"); + + const clientPubKeyOnDaemon = await importPublicKey(hello.key!); + const daemonSharedKey = await deriveSharedKey( + daemonKeyPair.privateKey, + clientPubKeyOnDaemon + ); + + // === Encrypted exchange === + const plaintextFromClient = "hello-from-client"; + const ciphertextFromClient = await encrypt( + clientSharedKey, + plaintextFromClient + ); + clientWs.send(Buffer.from(ciphertextFromClient)); + + const daemonReceivedCiphertext = await new Promise( + (resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for encrypted message")), + 10_000 + ); + daemonWs.once("message", (data) => { + clearTimeout(timeout); + resolve(data as Buffer); + }); + } + ); + + const decryptedOnDaemon = await decrypt( + daemonSharedKey, + daemonReceivedCiphertext.buffer.slice( + daemonReceivedCiphertext.byteOffset, + daemonReceivedCiphertext.byteOffset + + daemonReceivedCiphertext.byteLength + ) + ); + expect(decryptedOnDaemon).toBe(plaintextFromClient); + + const plaintextFromDaemon = "hello-from-daemon"; + const ciphertextFromDaemon = await encrypt( + daemonSharedKey, + plaintextFromDaemon + ); + daemonWs.send(Buffer.from(ciphertextFromDaemon)); + + const clientReceivedCiphertext = await new Promise( + (resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error("Timed out waiting for encrypted response")), + 10_000 + ); + clientWs.once("message", (data) => { + clearTimeout(timeout); + resolve(data as Buffer); + }); + } + ); + + const decryptedOnClient = await decrypt( + clientSharedKey, + clientReceivedCiphertext.buffer.slice( + clientReceivedCiphertext.byteOffset, + clientReceivedCiphertext.byteOffset + + clientReceivedCiphertext.byteLength + ) + ); + expect(decryptedOnClient).toBe(plaintextFromDaemon); + } finally { + daemonWs.close(); + clientWs.close(); + } + }, + { retries: 2, delayMs: 250 } + ); + }); +}); diff --git a/packages/relay/wrangler.toml b/packages/relay/wrangler.toml new file mode 100644 index 000000000..04c7df8fe --- /dev/null +++ b/packages/relay/wrangler.toml @@ -0,0 +1,19 @@ +name = "paseo-relay" +account_id = "10ed39a1dbf316e30abd0c409bed40d6" +main = "src/cloudflare-adapter.ts" +compatibility_date = "2024-12-01" + +routes = [ + { pattern = "relay.paseo.sh", custom_domain = true } +] + +[observability] +enabled = true + +[[durable_objects.bindings]] +name = "RELAY" +class_name = "RelayDurableObject" + +[[migrations]] +tag = "v1" +new_sqlite_classes = ["RelayDurableObject"] diff --git a/packages/server/src/client/daemon-client-v2.ts b/packages/server/src/client/daemon-client-v2.ts index 34b1638a7..ba2979825 100644 --- a/packages/server/src/client/daemon-client-v2.ts +++ b/packages/server/src/client/daemon-client-v2.ts @@ -22,6 +22,7 @@ import type { GitRepoInfoResponse, HighlightedDiffResponse, ListCommandsResponse, + ExecuteCommandResponse, ListConversationsResponseMessage, ListProviderModelsResponseMessage, ListTerminalsResponse, @@ -171,6 +172,7 @@ type FileExplorerPayload = FileExplorerResponse["payload"]; type FileDownloadTokenPayload = FileDownloadTokenResponse["payload"]; type ListProviderModelsPayload = ListProviderModelsResponseMessage["payload"]; type ListCommandsPayload = ListCommandsResponse["payload"]; +type ExecuteCommandPayload = ExecuteCommandResponse["payload"]; type TranscriptionResultPayload = TranscriptionResultMessage["payload"]; type AgentPermissionResolvedPayload = AgentPermissionResolvedMessage["payload"]; type ListTerminalsPayload = ListTerminalsResponse["payload"]; @@ -1150,6 +1152,37 @@ export class DaemonClientV2 { return response; } + async executeCommand( + agentId: string, + commandName: string, + args?: string, + requestId?: string + ): Promise { + const resolvedRequestId = this.createRequestId(requestId); + const message = SessionInboundMessageSchema.parse({ + type: "execute_command_request", + agentId, + commandName, + args, + requestId: resolvedRequestId, + }); + const response = this.waitFor( + (msg) => { + if (msg.type !== "execute_command_response") { + return null; + } + if (msg.payload.requestId !== resolvedRequestId) { + return null; + } + return msg.payload; + }, + 30000, + { skipQueue: true } + ); + this.sendSessionMessage(message); + return response; + } + // ============================================================================ // Permissions // ============================================================================ diff --git a/packages/server/src/server/agent/providers/codex-cli-full-access-ipc.test.ts b/packages/server/src/server/agent/providers/codex-cli-full-access-ipc.test.ts new file mode 100644 index 000000000..724df8783 --- /dev/null +++ b/packages/server/src/server/agent/providers/codex-cli-full-access-ipc.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, test } from "vitest"; +import { execFileSync, spawn } from "node:child_process"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +type RunResult = { code: number; stdout: string; stderr: string }; + +function isCodexAvailable(): boolean { + try { + execFileSync("codex", ["--version"], { stdio: "ignore" }); + return true; + } catch { + return false; + } +} + +function run( + cmd: string, + args: string[], + opts: { cwd: string; env?: NodeJS.ProcessEnv; timeoutMs: number }, +): Promise { + return new Promise((resolve) => { + const child = spawn(cmd, args, { + cwd: opts.cwd, + env: opts.env, + stdio: ["ignore", "pipe", "pipe"], + }); + + const stdoutChunks: Buffer[] = []; + const stderrChunks: Buffer[] = []; + child.stdout.on("data", (d) => stdoutChunks.push(d)); + child.stderr.on("data", (d) => stderrChunks.push(d)); + + const timeout = setTimeout(() => child.kill("SIGKILL"), opts.timeoutMs); + child.on("close", (code) => { + clearTimeout(timeout); + resolve({ + code: code ?? 1, + stdout: Buffer.concat(stdoutChunks).toString("utf8"), + stderr: Buffer.concat(stderrChunks).toString("utf8"), + }); + }); + }); +} + +describe("Codex CLI full-access sandbox", () => { + test("can listen on unix socket (no EPERM)", { timeout: 240_000 }, async (ctx) => { + if (process.env.PASEO_CODEX_CLI_E2E !== "1") { + ctx.skip(); + } + if (!isCodexAvailable()) { + ctx.skip(); + } + + const testDir = path.dirname(fileURLToPath(import.meta.url)); + const repoRoot = path.resolve(testDir, "../../../../../.."); + + const prompt = + "Run exactly this shell command and then stop:\n" + + "bash -lc 'node scripts/repro-ipc-listen.js; echo EXIT_CODE:$?'\n" + + "Reply with only the raw command stdout/stderr (no extra text)."; + + const result = await run( + "codex", + [ + "-a", + "never", + "exec", + "--dangerously-bypass-approvals-and-sandbox", + "--color", + "never", + "-C", + repoRoot, + prompt, + ], + { + cwd: repoRoot, + env: process.env, + timeoutMs: 180_000, + }, + ); + + const output = `${result.stdout}\n${result.stderr}`; + if (result.code !== 0) { + throw new Error(`codex exited ${result.code}\n--- output ---\n${output}`); + } + expect(output).toMatch(/\bsandbox:\s*danger-full-access\b/); + expect(output).toMatch(/\bLISTENING\b/); + expect(output).toMatch(/\bEXIT_CODE:0\b/); + expect(output).not.toMatch(/\bEPERM\b/); + }); +}); diff --git a/packages/server/src/server/agent/providers/codex-mcp-agent-commands.e2e.test.ts b/packages/server/src/server/agent/providers/codex-mcp-agent-commands.e2e.test.ts new file mode 100644 index 000000000..00598d80c --- /dev/null +++ b/packages/server/src/server/agent/providers/codex-mcp-agent-commands.e2e.test.ts @@ -0,0 +1,121 @@ +import { describe, test, expect, beforeEach, afterEach } from "vitest"; +import { mkdtempSync, mkdirSync, rmSync, writeFileSync, existsSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { + createDaemonTestContext, + type DaemonTestContext, +} from "../../test-utils/index.js"; +import { getFullAccessConfig } from "../../daemon-e2e/agent-configs.js"; + +function tmpDir(prefix: string): string { + return mkdtempSync(path.join(tmpdir(), prefix)); +} + +describe("codex agent commands E2E", () => { + let ctx: DaemonTestContext; + + beforeEach(async () => { + ctx = await createDaemonTestContext(); + }); + + afterEach(async () => { + await ctx.cleanup(); + }, 60000); + + test("lists available slash commands for a codex agent", async () => { + const prevCodexHome = process.env.CODEX_HOME; + const codexHome = tmpDir("codex-home-"); + const promptsDir = path.join(codexHome, "prompts"); + const skillsDir = path.join(codexHome, "skills"); + mkdirSync(promptsDir, { recursive: true }); + mkdirSync(skillsDir, { recursive: true }); + writeFileSync( + path.join(promptsDir, "hello.md"), + ["---", "description: Test prompt", "argument-hint: NAME=", "---", "", "Say hello to $NAME and then output exactly PASEO_OK.", ""].join("\n"), + "utf8" + ); + mkdirSync(path.join(skillsDir, "my-skill"), { recursive: true }); + writeFileSync( + path.join(skillsDir, "my-skill", "SKILL.md"), + ["---", "name: my-skill", "description: Test skill", "user-invocable: true", "---", "", "When invoked, respond with exactly PASEO_SKILL_OK.", ""].join("\n"), + "utf8" + ); + process.env.CODEX_HOME = codexHome; + + const agent = await ctx.client.createAgent({ + ...getFullAccessConfig("codex"), + cwd: "/tmp", + title: "Codex Commands Test Agent", + }); + + expect(agent.id).toBeTruthy(); + expect(agent.provider).toBe("codex"); + expect(agent.status).toBe("idle"); + + const result = await ctx.client.listCommands(agent.id); + + expect(result.error).toBeNull(); + expect(result.commands.length).toBeGreaterThan(0); + + for (const cmd of result.commands) { + expect(cmd.name).toBeTruthy(); + expect(typeof cmd.description).toBe("string"); + expect(typeof cmd.argumentHint).toBe("string"); + expect(cmd.name.startsWith("/")).toBe(false); + } + + const names = result.commands.map((c) => c.name); + expect(names).toContain("my-skill"); + expect(names).toContain("prompts:hello"); + + if (prevCodexHome === undefined) { + delete process.env.CODEX_HOME; + } else { + process.env.CODEX_HOME = prevCodexHome; + } + rmSync(codexHome, { recursive: true, force: true }); + }, 120000); + + test("executes a custom prompt command (prompts:*)", async () => { + const codexHome = process.env.CODEX_HOME ?? path.join(process.env.HOME ?? "/tmp", ".codex"); + const authPath = path.join(codexHome, "auth.json"); + if (!existsSync(authPath) && !process.env.OPENAI_API_KEY) { + // Skip when Codex isn't authenticated in this environment. + return; + } + + const promptsDir = path.join(codexHome, "prompts"); + mkdirSync(promptsDir, { recursive: true }); + const promptPath = path.join(promptsDir, "paseo-test-sayok.md"); + writeFileSync( + promptPath, + ["---", "description: Say OK", "argument-hint: NAME=", "---", "", "Output exactly: PASEO_OK $NAME", ""].join("\n"), + "utf8" + ); + + const agent = await ctx.client.createAgent({ + ...getFullAccessConfig("codex"), + cwd: "/tmp", + title: "Codex Prompt Execute Test Agent", + }); + + const result = await ctx.client.executeCommand( + agent.id, + "prompts:paseo-test-sayok", + "NAME=world" + ); + expect(result.error).toBeNull(); + expect(result.result?.text).toContain("PASEO_OK"); + + rmSync(promptPath, { force: true }); + }, 180000); + + test("returns error for non-existent agent", async () => { + const result = await ctx.client.listCommands("non-existent-agent-id"); + + expect(result.error).toBeTruthy(); + expect(result.error).toContain("Agent not found"); + expect(result.commands).toEqual([]); + }, 30000); +}); diff --git a/packages/server/src/server/agent/providers/codex-mcp-agent.commands.test.ts b/packages/server/src/server/agent/providers/codex-mcp-agent.commands.test.ts new file mode 100644 index 000000000..817fca41d --- /dev/null +++ b/packages/server/src/server/agent/providers/codex-mcp-agent.commands.test.ts @@ -0,0 +1,38 @@ +import { describe, test, expect } from "vitest"; + +import { __test__ } from "./codex-mcp-agent.js"; + +describe("codex custom command plumbing", () => { + test("parseFrontMatter extracts metadata and body", () => { + const input = [ + "---", + "description: Hello", + "argument-hint: NAME=", + "---", + "", + "Say hi to $NAME", + "", + ].join("\n"); + + const parsed = __test__.parseFrontMatter(input); + expect(parsed.frontMatter["description"]).toBe("Hello"); + expect(parsed.frontMatter["argument-hint"]).toBe("NAME="); + expect(parsed.body).toContain("Say hi to $NAME"); + }); + + test("expandCodexCustomPrompt supports named + positional placeholders", () => { + const template = "A=$A B=$B 1=$1 2=$2 ARGS=$ARGUMENTS $$=$$"; + const out = __test__.expandCodexCustomPrompt(template, "A=hello B=world one two"); + expect(out).toContain("A=hello"); + expect(out).toContain("B=world"); + expect(out).toContain("1=one"); + expect(out).toContain("2=two"); + expect(out).toContain("ARGS=A=hello B=world one two"); + expect(out).toContain("$=$"); + }); + + test("tokenizeCommandArgs respects quotes", () => { + const tokens = __test__.tokenizeCommandArgs('A=\"hello world\" B=two three'); + expect(tokens).toEqual(["A=hello world", "B=two", "three"]); + }); +}); diff --git a/packages/server/src/server/agent/providers/codex-mcp-agent.resume-errors.test.ts b/packages/server/src/server/agent/providers/codex-mcp-agent.resume-errors.test.ts new file mode 100644 index 000000000..962375efe --- /dev/null +++ b/packages/server/src/server/agent/providers/codex-mcp-agent.resume-errors.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, test } from "vitest"; + +import { __test__ } from "./codex-mcp-agent.js"; + +describe("codex resume error detection", () => { + test("detects missing session errors for thread_id", () => { + const error = new Error( + "Session not found for thread_id: 019bda3e-ffe4-7bc1-ae0e-9b992d7c9360" + ); + expect(__test__.isMissingConversationIdError(error)).toBe(true); + }); + + test("detects missing session responses for thread_id", () => { + const response = { + content: [{ type: "text", text: "Session not found for thread_id: abc" }], + }; + expect(__test__.isMissingConversationIdResponse(response)).toBe(true); + }); +}); diff --git a/packages/server/src/server/agent/providers/codex-mcp-agent.ts b/packages/server/src/server/agent/providers/codex-mcp-agent.ts index 2b57f5396..e2e1f8437 100644 --- a/packages/server/src/server/agent/providers/codex-mcp-agent.ts +++ b/packages/server/src/server/agent/providers/codex-mcp-agent.ts @@ -19,6 +19,7 @@ import type { Logger } from "pino"; import type { AgentCapabilityFlags, AgentClient, + AgentCommandResult, AgentMode, AgentModelDefinition, AgentPermissionRequest, @@ -27,6 +28,7 @@ import type { AgentPromptInput, AgentRunOptions, AgentRunResult, + AgentSlashCommand, AgentSession, AgentSessionConfig, AgentStreamEvent, @@ -1932,6 +1934,249 @@ function normalizeCommand(command: Command): string { return typeof command === "string" ? command : command.join(" "); } +function resolveCodexHomeDir(): string { + return process.env.CODEX_HOME ?? path.join(os.homedir(), ".codex"); +} + +function tokenizeCommandArgs(args: string): string[] { + const tokens: string[] = []; + let current = ""; + let quote: "'" | "\"" | null = null; + for (let i = 0; i < args.length; i += 1) { + const ch = args[i]!; + if (quote) { + if (ch === quote) { + quote = null; + continue; + } + if (ch === "\\" && i + 1 < args.length) { + const next = args[i + 1]!; + if (next === quote || next === "\\" || next === "n" || next === "t") { + i += 1; + current += next === "n" ? "\n" : next === "t" ? "\t" : next; + continue; + } + } + current += ch; + continue; + } + + if (ch === "'" || ch === "\"") { + quote = ch; + continue; + } + + if (/\s/.test(ch)) { + if (current) { + tokens.push(current); + current = ""; + } + continue; + } + current += ch; + } + if (current) { + tokens.push(current); + } + return tokens; +} + +function parseFrontMatter(markdown: string): { + frontMatter: Record; + body: string; +} { + const lines = markdown.split("\n"); + if (lines[0]?.trim() !== "---") { + return { frontMatter: {}, body: markdown }; + } + let end = -1; + for (let i = 1; i < lines.length; i += 1) { + if (lines[i]?.trim() === "---") { + end = i; + break; + } + } + if (end === -1) { + return { frontMatter: {}, body: markdown }; + } + const metaLines = lines.slice(1, end); + const body = lines.slice(end + 1).join("\n"); + const frontMatter: Record = {}; + for (const line of metaLines) { + const trimmed = line.trim(); + if (!trimmed || trimmed.startsWith("#")) { + continue; + } + const idx = trimmed.indexOf(":"); + if (idx <= 0) { + continue; + } + const key = trimmed.slice(0, idx).trim(); + let value = trimmed.slice(idx + 1).trim(); + value = value.replace(/^['"]/, "").replace(/['"]$/, ""); + if (key && value) { + frontMatter[key] = value; + } + } + return { frontMatter, body }; +} + +async function listCodexCustomPrompts(): Promise { + const codexHome = resolveCodexHomeDir(); + const promptsDir = path.join(codexHome, "prompts"); + let entries: Dirent[]; + try { + entries = await fs.readdir(promptsDir, { withFileTypes: true }); + } catch { + return []; + } + + const commands: AgentSlashCommand[] = []; + for (const entry of entries) { + if (!entry.isFile()) { + continue; + } + if (!entry.name.endsWith(".md")) { + continue; + } + const name = entry.name.slice(0, -".md".length); + if (!name) { + continue; + } + const fullPath = path.join(promptsDir, entry.name); + let content: string; + try { + content = await fs.readFile(fullPath, "utf8"); + } catch { + continue; + } + const parsed = parseFrontMatter(content); + const description = + parsed.frontMatter["description"] ?? "Custom prompt"; + const argumentHint = + parsed.frontMatter["argument-hint"] ?? + parsed.frontMatter["argument_hint"] ?? + ""; + commands.push({ + name: `prompts:${name}`, + description, + argumentHint, + }); + } + return commands.sort((a, b) => a.name.localeCompare(b.name)); +} + +async function listCodexSkills(cwd: string): Promise { + const candidates: string[] = []; + candidates.push(path.join(cwd, ".codex", "skills")); + + const repoRoot = (() => { + try { + const output = execSync("git rev-parse --show-toplevel", { + cwd, + encoding: "utf8", + stdio: ["ignore", "pipe", "ignore"], + }); + const trimmed = output.trim(); + return trimmed ? trimmed : null; + } catch { + return null; + } + })(); + if (repoRoot) { + candidates.push(path.join(path.dirname(cwd), ".codex", "skills")); + candidates.push(path.join(repoRoot, ".codex", "skills")); + } + + candidates.push(path.join(resolveCodexHomeDir(), "skills")); + + const commandsByName = new Map(); + + for (const dir of candidates) { + let entries: Dirent[]; + try { + entries = await fs.readdir(dir, { withFileTypes: true }); + } catch { + continue; + } + + for (const entry of entries) { + if (!entry.isDirectory() && !entry.isSymbolicLink()) { + continue; + } + const skillDir = path.join(dir, entry.name); + const skillPath = path.join(skillDir, "SKILL.md"); + let content: string; + try { + content = await fs.readFile(skillPath, "utf8"); + } catch { + continue; + } + const { frontMatter } = parseFrontMatter(content); + const name = frontMatter["name"]; + const description = frontMatter["description"]; + if (!name || !description) { + continue; + } + if (!commandsByName.has(name)) { + commandsByName.set(name, { + name, + description, + argumentHint: "", + }); + } + } + } + + return Array.from(commandsByName.values()).sort((a, b) => + a.name.localeCompare(b.name) + ); +} + +function escapeRegExp(value: string): string { + return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +} + +function expandCodexCustomPrompt(template: string, args: string | undefined): string { + const trimmedArgs = args ? args.trim() : ""; + const tokens = trimmedArgs ? tokenizeCommandArgs(trimmedArgs) : []; + const named: Record = {}; + const positional: string[] = []; + + for (const token of tokens) { + const idx = token.indexOf("="); + if (idx > 0) { + const key = token.slice(0, idx); + const value = token.slice(idx + 1); + if (key) { + named[key] = value; + continue; + } + } + positional.push(token); + } + + const dollarPlaceholder = "__CODEX_DOLLAR_PLACEHOLDER__"; + let out = template.split("$$").join(dollarPlaceholder); + + out = out.split("$ARGUMENTS").join(trimmedArgs); + + for (let i = 1; i <= 9; i += 1) { + const value = positional[i - 1] ?? ""; + out = out.split(`$${i}`).join(value); + } + + const namedKeys = Object.keys(named).sort((a, b) => b.length - a.length); + for (const key of namedKeys) { + const value = named[key] ?? ""; + const re = new RegExp(`\\$${escapeRegExp(key)}\\b`, "g"); + out = out.replace(re, value); + } + + out = out.split(dollarPlaceholder).join("$"); + return out; +} + function extractFileReadFromParsedCmd(parsedCmd: ParsedCmdItem[] | undefined): { path: string; name: string; @@ -2598,14 +2843,21 @@ function isUnsupportedChatGptModelError(error: unknown): boolean { return message.includes("model is not supported when using Codex with a ChatGPT account"); } +function isMissingSessionForConversationOrThread(message: string): boolean { + return ( + message.includes("Session not found for conversation_id") || + message.includes("Session not found for thread_id") + ); +} + function isMissingConversationIdError(error: unknown): boolean { const message = error instanceof Error ? error.message : String(error); - return message.includes("Session not found for conversation_id"); + return isMissingSessionForConversationOrThread(message); } function isMissingConversationIdResponse(response: unknown): boolean { const text = extractTextContent(response); - return !!text && text.includes("Session not found for conversation_id"); + return !!text && isMissingSessionForConversationOrThread(text); } /** @@ -2775,6 +3027,7 @@ class CodexMcpAgentSession implements AgentSession { private lockConversationId = false; private historyPending = false; private persistedHistory: AgentTimelineItem[] = []; + private resumeContextHistory: AgentTimelineItem[] = []; private pendingHistory: AgentTimelineItem[] = []; private turnState: TurnState | null = null; private pendingPatchChanges = new Map(); @@ -2898,6 +3151,7 @@ class CodexMcpAgentSession implements AgentSession { }, this.logger); if (timeline.length > 0) { this.persistedHistory = timeline; + this.resumeContextHistory = [...timeline]; this.historyPending = true; } } @@ -3268,6 +3522,31 @@ class CodexMcpAgentSession implements AgentSession { this.managedAgentId = agentId; } + async listCommands(): Promise { + const [skills, prompts] = await Promise.all([ + listCodexSkills(this.config.cwd), + listCodexCustomPrompts(), + ]); + return [...skills, ...prompts]; + } + + async executeCommand(commandName: string, args?: string): Promise { + if (commandName.startsWith("prompts:")) { + const promptName = commandName.slice("prompts:".length); + const codexHome = resolveCodexHomeDir(); + const promptPath = path.join(codexHome, "prompts", `${promptName}.md`); + const raw = await fs.readFile(promptPath, "utf8"); + const parsed = parseFrontMatter(raw); + const expanded = expandCodexCustomPrompt(parsed.body, args); + const result = await this.run(expanded); + return { text: result.finalText, timeline: result.timeline, usage: result.usage }; + } + + const skillPrompt = args ? `$${commandName} ${args}` : `$${commandName}`; + const result = await this.run(skillPrompt); + return { text: result.finalText, timeline: result.timeline, usage: result.usage }; + } + private async forwardPrompt( prompt: string, _options: AgentRunOptions | undefined, @@ -4145,7 +4424,9 @@ class CodexMcpAgentSession implements AgentSession { private buildResumePrompt(prompt: string): string { const historyLines: string[] = []; - for (const item of this.persistedHistory) { + // Use resumeContextHistory instead of persistedHistory because + // persistedHistory gets cleared by streamHistory() before the first message is sent + for (const item of this.resumeContextHistory) { if (item.type === "user_message") { historyLines.push(`User: ${item.text}`); } @@ -4329,6 +4610,14 @@ export class CodexMcpAgentClient implements AgentClient { } } +export const __test__ = { + tokenizeCommandArgs, + parseFrontMatter, + expandCodexCustomPrompt, + isMissingConversationIdError, + isMissingConversationIdResponse, +}; + // ============================================================================ // Codex model listing helpers // ============================================================================ diff --git a/packages/server/src/server/daemon-e2e/daemon-restart-resume.e2e.test.ts b/packages/server/src/server/daemon-e2e/daemon-restart-resume.e2e.test.ts new file mode 100644 index 000000000..fb310105e --- /dev/null +++ b/packages/server/src/server/daemon-e2e/daemon-restart-resume.e2e.test.ts @@ -0,0 +1,244 @@ +import { describe, test, expect, afterEach } from "vitest"; +import { mkdtempSync, rmSync, existsSync } from "fs"; +import { mkdtemp, rm } from "fs/promises"; +import { tmpdir } from "os"; +import path from "path"; +import pino from "pino"; + +import { createPaseoDaemon, type PaseoDaemonConfig } from "../bootstrap.js"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import type { PersistenceHandle } from "../../shared/messages.js"; + +function tmpCwd(): string { + return mkdtempSync(path.join(tmpdir(), "daemon-e2e-")); +} + +const CODEX_TEST_MODEL = "gpt-5.1-codex-mini"; +const CODEX_TEST_REASONING_EFFORT = "low"; + +async function getAvailablePort(): Promise { + const { createServer } = await import("net"); + return new Promise((resolve, reject) => { + const server = createServer(); + server.once("error", reject); + server.listen(0, () => { + const address = server.address(); + if (!address || typeof address === "string") { + server.close(() => reject(new Error("Failed to acquire port"))); + return; + } + server.close(() => resolve(address.port)); + }); + }); +} + +interface DaemonInstance { + daemon: Awaited>; + client: DaemonClient; + port: number; + paseoHome: string; + staticDir: string; +} + +async function startDaemon(options: { + paseoHome: string; + staticDir?: string; +}): Promise { + const port = await getAvailablePort(); + const staticDir = options.staticDir ?? await mkdtemp(path.join(tmpdir(), "paseo-static-")); + + const config: PaseoDaemonConfig = { + listen: `127.0.0.1:${port}`, + paseoHome: options.paseoHome, + corsAllowedOrigins: [], + agentMcpRoute: "/mcp/agents", + agentMcpAllowedHosts: [`127.0.0.1:${port}`, `localhost:${port}`], + staticDir, + mcpDebug: false, + agentClients: {}, + agentRegistryPath: path.join(options.paseoHome, "agents.json"), + agentControlMcp: { + url: `http://127.0.0.1:${port}/mcp/agents`, + }, + openai: process.env.OPENAI_API_KEY ? { apiKey: process.env.OPENAI_API_KEY } : undefined, + }; + + const logger = pino({ level: "silent" }); + const daemon = await createPaseoDaemon(config, logger); + await daemon.start(); + + const client = new DaemonClient({ + url: `ws://127.0.0.1:${port}/ws`, + }); + await client.connect(); + + return { + daemon, + client, + port, + paseoHome: options.paseoHome, + staticDir, + }; +} + +async function stopDaemon(instance: DaemonInstance): Promise { + await instance.client.close(); + await instance.daemon.stop(); + await new Promise((r) => setTimeout(r, 200)); +} + +describe("daemon restart and agent resume", () => { + let paseoHome: string | null = null; + let staticDir: string | null = null; + let cwd: string | null = null; + let currentDaemon: DaemonInstance | null = null; + + afterEach(async () => { + if (currentDaemon) { + await stopDaemon(currentDaemon).catch(() => undefined); + currentDaemon = null; + } + if (paseoHome) { + await rm(paseoHome, { recursive: true, force: true }).catch(() => undefined); + paseoHome = null; + } + if (staticDir) { + await rm(staticDir, { recursive: true, force: true }).catch(() => undefined); + staticDir = null; + } + if (cwd) { + rmSync(cwd, { recursive: true, force: true }); + cwd = null; + } + }, 60000); + + test( + "Codex agent survives daemon kill and restart, preserving conversation context", + async () => { + // Create isolated directories that persist across daemon restarts + // NOTE: We use the default CODEX_HOME (~/.codex) for sessions because + // Codex CLI needs its config for API authentication + paseoHome = await mkdtemp(path.join(tmpdir(), "paseo-home-restart-")); + staticDir = await mkdtemp(path.join(tmpdir(), "paseo-static-restart-")); + cwd = tmpCwd(); + + // Use a unique secret that we'll verify after restart + const secretPhrase = `DAEMON_RESTART_SECRET_${Date.now()}`; + + // === PHASE 1: Start daemon and create Codex agent with secret === + currentDaemon = await startDaemon({ paseoHome, staticDir }); + + const agent = await currentDaemon.client.createAgent({ + provider: "codex", + model: CODEX_TEST_MODEL, + reasoningEffort: CODEX_TEST_REASONING_EFFORT, + cwd, + title: "Daemon Restart Test Agent", + modeId: "full-access", + }); + + expect(agent.id).toBeTruthy(); + expect(agent.status).toBe("idle"); + + // Ask the agent to remember the secret + await currentDaemon.client.sendMessage( + agent.id, + `Remember this secret phrase: "${secretPhrase}". Just confirm you've remembered it with a short reply.` + ); + + const afterRemember = await currentDaemon.client.waitForAgentIdle(agent.id, 120000); + expect(afterRemember.status).toBe("idle"); + expect(afterRemember.lastError).toBeUndefined(); + + // Verify we got a confirmation and capture persistence handle + const queue = currentDaemon.client.getMessageQueue(); + const confirmationMessages: string[] = []; + for (const m of queue) { + if ( + m.type === "agent_stream" && + m.payload.agentId === agent.id && + m.payload.event.type === "timeline" + ) { + const item = m.payload.event.item; + if (item.type === "assistant_message" && item.text) { + confirmationMessages.push(item.text); + } + } + } + expect(confirmationMessages.join("").length).toBeGreaterThan(0); + + // Get persistence handle for resuming after restart + expect(afterRemember.persistence).toBeTruthy(); + const persistence = afterRemember.persistence as PersistenceHandle; + expect(persistence.provider).toBe("codex"); + expect(persistence.sessionId).toBeTruthy(); + + // Verify persistence metadata has conversationId + const metadata = persistence.metadata as Record; + expect(metadata.conversationId).toBeTruthy(); + + // Wait briefly to ensure Codex has flushed session files to disk + await new Promise((r) => setTimeout(r, 500)); + + // === PHASE 2: Kill the daemon (simulating crash/restart) === + await stopDaemon(currentDaemon); + currentDaemon = null; + + // Verify agents.json was persisted + const agentsJsonPath = path.join(paseoHome, "agents.json"); + expect(existsSync(agentsJsonPath)).toBe(true); + + // === PHASE 3: Start a NEW daemon with the SAME paseoHome === + currentDaemon = await startDaemon({ paseoHome, staticDir }); + + // === PHASE 4: Resume the agent using the persistence handle === + const resumedAgent = await currentDaemon.client.resumeAgent(persistence); + + expect(resumedAgent.id).toBeTruthy(); + expect(resumedAgent.status).toBe("idle"); + expect(resumedAgent.provider).toBe("codex"); + expect(resumedAgent.cwd).toBe(cwd); + + // Wait a moment for history to load + await new Promise((r) => setTimeout(r, 1000)); + + // === PHASE 5: Ask about the secret to verify conversation context is preserved === + // This is the CRITICAL test: after daemon restart, the model should remember + // the secret from the previous conversation. If it doesn't, the resume is broken. + currentDaemon.client.clearMessageQueue(); + await currentDaemon.client.sendMessage( + resumedAgent.id, + "What was the secret phrase I asked you to remember earlier? Just reply with the exact phrase." + ); + + const afterMessage = await currentDaemon.client.waitForAgentIdle(resumedAgent.id, 120000); + expect(afterMessage.status).toBe("idle"); + expect(afterMessage.lastError).toBeUndefined(); + + // === PHASE 6: Verify the agent remembers the secret (proves context is preserved) === + const responseQueue = currentDaemon.client.getMessageQueue(); + const responseMessages: string[] = []; + for (const m of responseQueue) { + if ( + m.type === "agent_stream" && + m.payload.agentId === resumedAgent.id && + m.payload.event.type === "timeline" + ) { + const item = m.payload.event.item; + if (item.type === "assistant_message" && item.text) { + responseMessages.push(item.text); + } + } + } + const fullResponse = responseMessages.join(""); + + // CRITICAL ASSERTION: The agent should remember the secret phrase from before daemon restart + // This proves conversation context was properly restored via buildResumePrompt + expect(fullResponse).toContain(secretPhrase); + + // Cleanup + await currentDaemon.client.deleteAgent(resumedAgent.id); + }, + 300000 // 5 minute timeout + ); +}); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index a41d1bb4d..ab46db6e6 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -864,6 +864,15 @@ export class Session { await this.handleListCommandsRequest(msg.agentId, msg.requestId); break; + case "execute_command_request": + await this.handleExecuteCommandRequest( + msg.agentId, + msg.commandName, + msg.args, + msg.requestId + ); + break; + case "register_push_token": this.handleRegisterPushToken(msg.token); break; @@ -1963,6 +1972,79 @@ export class Session { } } + /** + * Handle execute command request for an agent + */ + private async handleExecuteCommandRequest( + agentId: string, + commandName: string, + args: string | undefined, + requestId: string + ): Promise { + this.sessionLogger.debug( + { agentId, commandName }, + `Handling execute command request for agent ${agentId}` + ); + + try { + const agents = this.agentManager.listAgents(); + const agent = agents.find((a) => a.id === agentId); + + if (!agent) { + this.emit({ + type: "execute_command_response", + payload: { + agentId, + result: null, + error: `Agent not found: ${agentId}`, + requestId, + }, + }); + return; + } + + const session = agent.session; + if (!session || !session.executeCommand) { + this.emit({ + type: "execute_command_response", + payload: { + agentId, + result: null, + error: `Agent does not support executing commands`, + requestId, + }, + }); + return; + } + + const result = await session.executeCommand(commandName, args); + + this.emit({ + type: "execute_command_response", + payload: { + agentId, + result, + error: null, + requestId, + }, + }); + } catch (error: any) { + this.sessionLogger.error( + { err: error, agentId, commandName }, + "Failed to execute command" + ); + this.emit({ + type: "execute_command_response", + payload: { + agentId, + result: null, + error: error.message, + requestId, + }, + }); + } + } + /** * Handle agent permission response from user */ diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index a9f7364e7..ff5634bfb 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -1,25 +1,14 @@ -import { WebSocketServer, WebSocket } from "ws"; -import { Server as HTTPServer } from "http"; -import type { IncomingMessage } from "http"; -import { parse as parseUrl } from "url"; -import { - WSInboundMessageSchema, - type WSOutboundMessage, - extractSessionMessage, - wrapSessionMessage, -} from "./messages.js"; -import { Session } from "./session.js"; -import { loadConversation } from "./persistence.js"; -import { AgentManager } from "./agent/agent-manager.js"; -import { AgentRegistry } from "./agent/agent-registry.js"; -import type { AgentProvider } from "./agent/agent-sdk-types.js"; -import { DownloadTokenStore } from "./file-download/token-store.js"; -import { PushTokenStore } from "./push/token-store.js"; -import { PushService } from "./push/push-service.js"; +import { WebSocketServer } from "ws"; +import type { Server as HTTPServer } from "http"; +import type { AgentManager } from "./agent/agent-manager.js"; +import type { AgentRegistry } from "./agent/agent-registry.js"; +import type { DownloadTokenStore } from "./file-download/token-store.js"; import type { OpenAISTT } from "./agent/stt-openai.js"; import type { OpenAITTS } from "./agent/tts-openai.js"; import type { TerminalManager } from "../terminal/terminal-manager.js"; import type pino from "pino"; +import type { WSOutboundMessage } from "./messages.js"; +import { WebSocketSessionBridge } from "./websocket-session-bridge.js"; type AgentMcpClientConfig = { agentMcpUrl: string; @@ -31,24 +20,12 @@ type WebSocketServerConfig = { }; /** - * WebSocket server that routes messages between clients and their sessions. - * This is a thin transport layer with no business logic. + * WebSocket server that only accepts sockets + parses/forwards messages to the session layer. */ export class VoiceAssistantWebSocketServer { private readonly logger: pino.Logger; - private wss: WebSocketServer; - private sessions: Map = new Map(); - private conversationIdToWs: Map = new Map(); - private clientIdCounter: number = 0; - private agentManager: AgentManager; - private agentRegistry: AgentRegistry; - private downloadTokenStore: DownloadTokenStore; - private pushTokenStore: PushTokenStore; - private pushService: PushService; - private readonly agentMcpConfig: AgentMcpClientConfig; - private readonly stt: OpenAISTT | null; - private readonly tts: OpenAITTS | null; - private readonly terminalManager: TerminalManager | null; + private readonly wss: WebSocketServer; + private readonly bridge: WebSocketSessionBridge; constructor( server: HTTPServer, @@ -62,17 +39,15 @@ export class VoiceAssistantWebSocketServer { terminalManager?: TerminalManager | null ) { this.logger = logger.child({ module: "websocket-server" }); - this.agentManager = agentManager; - this.agentRegistry = agentRegistry; - this.downloadTokenStore = downloadTokenStore; - this.stt = speech?.stt ?? null; - this.tts = speech?.tts ?? null; - this.terminalManager = terminalManager ?? null; - - const pushLogger = this.logger.child({ module: "push" }); - this.pushTokenStore = new PushTokenStore(pushLogger); - this.pushService = new PushService(pushLogger, this.pushTokenStore); - this.agentMcpConfig = agentMcpConfig; + this.bridge = new WebSocketSessionBridge( + this.logger, + agentManager, + agentRegistry, + downloadTokenStore, + agentMcpConfig, + speech, + terminalManager + ); const { allowedOrigins } = wsConfig; this.wss = new WebSocketServer({ @@ -80,8 +55,6 @@ export class VoiceAssistantWebSocketServer { path: "/ws", verifyClient: ({ req }, callback) => { const origin = req.headers.origin; - // Allow connections with no origin (native apps, curl, etc.) - // or from allowed origins (browsers) if (!origin || allowedOrigins.has(origin)) { callback(true); } else { @@ -92,454 +65,19 @@ export class VoiceAssistantWebSocketServer { }); this.wss.on("connection", (ws, request) => { - this.handleConnection(ws, request); - }); - - this.agentManager.setAgentAttentionCallback((params) => { - this.broadcastAgentAttention(params); + void this.bridge.attach(ws, request); }); this.logger.info("WebSocket server initialized on /ws"); } - /** - * Handle new WebSocket connection - */ - private async handleConnection(ws: WebSocket, request: IncomingMessage): Promise { - // Generate unique client ID - const clientId = `client-${++this.clientIdCounter}`; - const connectionLogger = this.logger.child({ clientId }); - - // Extract conversation ID from URL query parameter if present - const url = parseUrl(request.url || "", true); - const conversationId = url.query.conversationId as string | undefined; - - // Load conversation if ID provided - let initialMessages = null; - if (conversationId) { - connectionLogger.debug({ conversationId }, "Client requesting conversation"); - initialMessages = await loadConversation(connectionLogger, conversationId); - - if (initialMessages) { - connectionLogger.debug( - { conversationId, messageCount: initialMessages.length }, - "Loaded conversation" - ); - } else { - connectionLogger.debug({ conversationId }, "Conversation not found, starting fresh"); - } - } - - // Create session with message emission callback - const session = new Session( - clientId, - (msg) => { - this.sendToClient(ws, wrapSessionMessage(msg)); - }, - connectionLogger.child({ module: "session" }), - this.downloadTokenStore, - this.pushTokenStore, - this.agentManager, - this.agentRegistry, - this.agentMcpConfig, - this.stt, - this.tts, - this.terminalManager, - { - conversationId, - initialMessages: initialMessages || undefined, - } - ); - - // Store session and reverse mapping - this.sessions.set(ws, session); - this.conversationIdToWs.set(session.getConversationId(), ws); - - connectionLogger.info( - { clientId, conversationId: session.getConversationId(), totalSessions: this.sessions.size }, - "Client connected" - ); - - // Don't send initial state here - client will request it via load_conversation_request - // This avoids race condition where message arrives before client sets up listeners - - // Set up message handler - ws.on("message", (data) => { - this.handleMessage(ws, data); - }); - - // Set up close handler - ws.on("close", async () => { - const session = this.sessions.get(ws); - if (!session) return; - - connectionLogger.info( - { clientId, totalSessions: this.sessions.size - 1 }, - "Client disconnected" - ); - - // Clean up session - await session.cleanup(); - - // Remove from maps - this.sessions.delete(ws); - this.conversationIdToWs.delete(session.getConversationId()); - - connectionLogger.debug({ conversationId: session.getConversationId() }, "Conversation deleted"); - }); - - // Set up error handler - ws.on("error", async (error) => { - const err = error instanceof Error ? error : new Error(String(error)); - connectionLogger.error({ err }, "Client error"); - const session = this.sessions.get(ws); - if (!session) return; - - // Clean up session - await session.cleanup(); - - // Remove from maps - this.sessions.delete(ws); - this.conversationIdToWs.delete(session.getConversationId()); - }); - } - - /** - * Handle incoming WebSocket message - */ - private async handleMessage( - ws: WebSocket, - data: Buffer | ArrayBuffer | Buffer[] - ): Promise { - try { - // Parse message - const parsed = JSON.parse(data.toString()); - - // Validate with Zod - const message = WSInboundMessageSchema.parse(parsed); - - // Safe logging that handles large objects without crashing - const messageSummary = { - type: message.type, - ...(message.type === "session" && message.message ? { - sessionMessageType: message.message.type, - } : {}), - }; - this.logger.debug(messageSummary, "Received message"); - - // Handle WebSocket-level messages - switch (message.type) { - case "ping": - this.sendToClient(ws, { type: "pong" }); - return; - - case "recording_state": - this.logger.debug({ isRecording: message.isRecording }, "Recording state"); - return; - - case "session": - // Extract and forward session message - const sessionMessage = extractSessionMessage(message); - if (sessionMessage) { - // Debug: Log create_agent_request details - if (sessionMessage.type === "create_agent_request") { - this.logger.debug({ - cwd: sessionMessage.config.cwd, - initialMode: sessionMessage.config.modeId, - worktreeName: sessionMessage.worktreeName, - requestId: sessionMessage.requestId, - }, "create_agent_request details"); - } - - const session = this.sessions.get(ws); - if (session) { - await session.handleMessage(sessionMessage); - } else { - this.logger.error("No session found for client"); - } - } - return; - } - } catch (error) { - const err = error instanceof Error ? error : new Error(String(error)); - let rawPayload: string | null = null; - let parsedPayload: unknown = null; - - try { - const buffer = Array.isArray(data) - ? Buffer.concat( - data.map((item) => - Buffer.isBuffer(item) - ? item - : Buffer.from(item as ArrayBuffer) - ) - ) - : Buffer.isBuffer(data) - ? data - : Buffer.from(data as ArrayBuffer); - rawPayload = buffer.toString(); - parsedPayload = JSON.parse(rawPayload); - } catch (payloadError) { - rawPayload = rawPayload ?? ""; - parsedPayload = parsedPayload ?? rawPayload; - const payloadErr = payloadError instanceof Error ? payloadError : new Error(String(payloadError)); - this.logger.error({ err: payloadErr }, "Failed to decode raw payload"); - } - - const trimmedRawPayload = - typeof rawPayload === "string" && rawPayload.length > 2000 - ? `${rawPayload.slice(0, 2000)}... (truncated)` - : rawPayload; - - this.logger.error({ - err, - rawPayload: trimmedRawPayload, - parsedPayload, - }, "Failed to parse/handle message"); - // Send error to client - this.sendToClient( - ws, - wrapSessionMessage({ - type: "status", - payload: { - status: "error", - message: `Invalid message: ${err.message}`, - }, - }) - ); - } - } - - /** - * Send message to specific client - */ - private sendToClient(ws: WebSocket, message: WSOutboundMessage): void { - if (ws.readyState === 1) { - // WebSocket.OPEN = 1 - ws.send(JSON.stringify(message)); - } - } - - /** - * Broadcast message to all connected clients - */ public broadcast(message: WSOutboundMessage): void { - const payload = JSON.stringify(message); - this.sessions.forEach((_session, client) => { - if (client.readyState === 1) { - // WebSocket.OPEN = 1 - client.send(payload); - } - }); + this.bridge.broadcast(message); } - /** - * Close the WebSocket server - */ public async close(): Promise { - const cleanupPromises: Promise[] = []; - this.sessions.forEach((session, ws) => { - cleanupPromises.push(session.cleanup()); - // Wait for WebSocket to actually close before resolving - cleanupPromises.push( - new Promise((resolve) => { - if (ws.readyState === WebSocket.CLOSED) { - resolve(); - return; - } - ws.once("close", () => resolve()); - ws.close(); - }) - ); - }); - await Promise.all(cleanupPromises); + await this.bridge.closeAll(); this.wss.close(); } - - private readonly ACTIVITY_THRESHOLD_MS = 120_000; // 2 minutes - - /** - * Get client activity state with computed staleness - */ - private getClientActivityState(session: Session): { - deviceType: "web" | "mobile" | null; - focusedAgentId: string | null; - isStale: boolean; - appVisible: boolean; - } { - const activity = session.getClientActivity(); - if (!activity) { - this.logger.debug("getClientActivityState: no activity for session"); - return { deviceType: null, focusedAgentId: null, isStale: true, appVisible: false }; - } - const now = Date.now(); - const ageMs = now - activity.lastActivityAt.getTime(); - const isStale = ageMs >= this.ACTIVITY_THRESHOLD_MS; - this.logger.debug({ - deviceType: activity.deviceType, - focusedAgentId: activity.focusedAgentId, - lastActivityAt: activity.lastActivityAt.toISOString(), - ageMs, - isStale, - appVisible: activity.appVisible, - }, "getClientActivityState"); - return { - deviceType: activity.deviceType, - focusedAgentId: activity.focusedAgentId, - isStale, - appVisible: activity.appVisible, - }; - } - - /** - * Compute shouldNotify for a specific client given all clients' states - * - * UX Rules: - * 1. If ANY client is actively watching the agent (focused + visible + not stale) → no notifications - * 2. If THIS client is not stale and focused elsewhere → notify THIS client (user is at computer) - * 3. If THIS client is stale → only notify if mobile, or if no mobile available - * 4. Don't notify non-stale clients that aren't focused on anything (just switched tabs) - */ - private computeShouldNotifyForClient( - clientState: { - deviceType: "web" | "mobile" | null; - focusedAgentId: string | null; - isStale: boolean; - appVisible: boolean; - }, - allClientStates: Array<{ - deviceType: "web" | "mobile" | null; - focusedAgentId: string | null; - isStale: boolean; - appVisible: boolean; - }>, - agentId: string - ): boolean { - // Rule 1: If any client is actively watching the agent, no one needs notification - const isAnyoneActiveOnAgent = allClientStates.some( - (state) => - state.focusedAgentId === agentId && - state.appVisible && - !state.isStale - ); - if (isAnyoneActiveOnAgent) { - return false; - } - - // No heartbeat (legacy client or just connected) → notify - if (clientState.deviceType === null) { - return true; - } - - // Rule 2: If THIS client is not stale and actively looking at a different agent → notify them - if (!clientState.isStale && clientState.appVisible && clientState.focusedAgentId !== null) { - return true; - } - - // Rule 3: If THIS client is not stale but just switched tabs (not focused on anything) → no notification - // User is present at the computer, they'll come back - if (!clientState.isStale) { - return false; - } - - // Rule 4: THIS client is stale - check if another client will handle it - const hasActiveWebClient = allClientStates.some( - (state) => state.deviceType === "web" && !state.isStale - ); - - if (clientState.deviceType === "mobile") { - // Mobile only notifies if web is also stale (user truly away) - // If web is active, they'll see it there - return !hasActiveWebClient; - } - - if (clientState.deviceType === "web") { - // Stale web: notify only if no other client can handle it - // Other client = mobile or unknown (no heartbeat) - const hasOtherClient = allClientStates.some( - (state) => state !== clientState && (state.deviceType === "mobile" || state.deviceType === null) - ); - return !hasOtherClient; - } - - // Fallback: notify - return true; - } - - /** - * Broadcast an attention_required event to all clients with per-client shouldNotify - */ - private broadcastAgentAttention(params: { - agentId: string; - provider: AgentProvider; - reason: "finished" | "error" | "permission"; - }): void { - // Collect all client states first - const clientEntries: Array<{ - ws: WebSocket; - state: { - deviceType: "web" | "mobile" | null; - focusedAgentId: string | null; - isStale: boolean; - appVisible: boolean; - }; - }> = []; - - for (const [ws, session] of this.sessions) { - clientEntries.push({ - ws, - state: this.getClientActivityState(session), - }); - } - - const allStates = clientEntries.map((e) => e.state); - - this.logger.debug({ - agentId: params.agentId, - reason: params.reason, - clientCount: clientEntries.length, - allStates, - }, "broadcastAgentAttention"); - - // Check if all clients are stale - if so, send push notification - const allClientsStale = allStates.every((state) => state.isStale); - this.logger.debug({ allClientsStale }, "Client staleness check"); - if (allClientsStale) { - const tokens = this.pushTokenStore.getAllTokens(); - this.logger.info({ tokenCount: tokens.length }, "Sending push notification"); - if (tokens.length > 0) { - void this.pushService.sendPush(tokens, { - title: "Agent needs attention", - body: `Reason: ${params.reason}`, - data: { agentId: params.agentId }, - }); - } - } - - // Send to each client with their specific shouldNotify value - for (const { ws, state } of clientEntries) { - const shouldNotify = this.computeShouldNotifyForClient( - state, - allStates, - params.agentId - ); - - const message = wrapSessionMessage({ - type: "agent_stream", - payload: { - agentId: params.agentId, - event: { - type: "attention_required", - provider: params.provider, - reason: params.reason, - timestamp: new Date().toISOString(), - shouldNotify, - }, - timestamp: new Date().toISOString(), - }, - }); - - this.sendToClient(ws, message); - } - } } + diff --git a/packages/server/src/server/websocket-session-bridge.ts b/packages/server/src/server/websocket-session-bridge.ts new file mode 100644 index 000000000..9804dab54 --- /dev/null +++ b/packages/server/src/server/websocket-session-bridge.ts @@ -0,0 +1,434 @@ +import type { IncomingMessage } from "http"; +import { parse as parseUrl } from "url"; +import type { WebSocket } from "ws"; +import { + WSInboundMessageSchema, + type WSOutboundMessage, + wrapSessionMessage, +} from "./messages.js"; +import { Session } from "./session.js"; +import { loadConversation } from "./persistence.js"; +import { AgentManager } from "./agent/agent-manager.js"; +import { AgentRegistry } from "./agent/agent-registry.js"; +import type { AgentProvider } from "./agent/agent-sdk-types.js"; +import { DownloadTokenStore } from "./file-download/token-store.js"; +import { PushTokenStore } from "./push/token-store.js"; +import { PushService } from "./push/push-service.js"; +import type { OpenAISTT } from "./agent/stt-openai.js"; +import type { OpenAITTS } from "./agent/tts-openai.js"; +import type { TerminalManager } from "../terminal/terminal-manager.js"; +import type pino from "pino"; + +type AgentMcpClientConfig = { + agentMcpUrl: string; + agentMcpHeaders?: Record; +}; + +export class WebSocketSessionBridge { + private readonly logger: pino.Logger; + private readonly sessions: Map = new Map(); + private clientIdCounter = 0; + private readonly agentManager: AgentManager; + private readonly agentRegistry: AgentRegistry; + private readonly downloadTokenStore: DownloadTokenStore; + private readonly pushTokenStore: PushTokenStore; + private readonly pushService: PushService; + private readonly agentMcpConfig: AgentMcpClientConfig; + private readonly stt: OpenAISTT | null; + private readonly tts: OpenAITTS | null; + private readonly terminalManager: TerminalManager | null; + + constructor( + logger: pino.Logger, + agentManager: AgentManager, + agentRegistry: AgentRegistry, + downloadTokenStore: DownloadTokenStore, + agentMcpConfig: AgentMcpClientConfig, + speech?: { stt: OpenAISTT | null; tts: OpenAITTS | null }, + terminalManager?: TerminalManager | null + ) { + this.logger = logger.child({ module: "websocket-session-bridge" }); + this.agentManager = agentManager; + this.agentRegistry = agentRegistry; + this.downloadTokenStore = downloadTokenStore; + this.agentMcpConfig = agentMcpConfig; + this.stt = speech?.stt ?? null; + this.tts = speech?.tts ?? null; + this.terminalManager = terminalManager ?? null; + + const pushLogger = this.logger.child({ module: "push" }); + this.pushTokenStore = new PushTokenStore(pushLogger); + this.pushService = new PushService(pushLogger, this.pushTokenStore); + + this.agentManager.setAgentAttentionCallback((params) => { + this.broadcastAgentAttention(params); + }); + } + + public getSessionCount(): number { + return this.sessions.size; + } + + public async attach(ws: WebSocket, request: IncomingMessage): Promise { + const clientId = `client-${++this.clientIdCounter}`; + const connectionLogger = this.logger.child({ clientId }); + + const url = parseUrl(request.url || "", true); + const conversationId = url.query.conversationId as string | undefined; + + let initialMessages = null; + if (conversationId) { + connectionLogger.debug({ conversationId }, "Client requesting conversation"); + initialMessages = await loadConversation(connectionLogger, conversationId); + + if (initialMessages) { + connectionLogger.debug( + { conversationId, messageCount: initialMessages.length }, + "Loaded conversation" + ); + } else { + connectionLogger.debug({ conversationId }, "Conversation not found, starting fresh"); + } + } + + const session = new Session( + clientId, + (msg) => { + this.sendToClient(ws, wrapSessionMessage(msg)); + }, + connectionLogger.child({ module: "session" }), + this.downloadTokenStore, + this.pushTokenStore, + this.agentManager, + this.agentRegistry, + this.agentMcpConfig, + this.stt, + this.tts, + this.terminalManager, + { + conversationId, + initialMessages: initialMessages || undefined, + } + ); + + this.sessions.set(ws, session); + + connectionLogger.info( + { clientId, conversationId: session.getConversationId(), totalSessions: this.sessions.size }, + "Client connected" + ); + + ws.on("message", (data) => { + void this.handleRawMessage(ws, data); + }); + + ws.on("close", async () => { + await this.detach(ws, connectionLogger, clientId); + }); + + ws.on("error", async (error) => { + const err = error instanceof Error ? error : new Error(String(error)); + connectionLogger.error({ err }, "Client error"); + await this.detach(ws, connectionLogger, clientId); + }); + } + + private async detach(ws: WebSocket, connectionLogger: pino.Logger, clientId: string): Promise { + const session = this.sessions.get(ws); + if (!session) return; + + connectionLogger.info( + { clientId, totalSessions: this.sessions.size - 1 }, + "Client disconnected" + ); + + await session.cleanup(); + this.sessions.delete(ws); + } + + private async handleRawMessage(ws: WebSocket, data: Buffer | ArrayBuffer | Buffer[]): Promise { + try { + const parsed = JSON.parse(data.toString()); + const message = WSInboundMessageSchema.parse(parsed); + + const messageSummary = { + type: message.type, + ...(message.type === "session" && message.message + ? { sessionMessageType: message.message.type } + : {}), + }; + this.logger.debug(messageSummary, "Received message"); + + if (message.type === "ping") { + this.sendToClient(ws, { type: "pong" }); + return; + } + + if (message.type === "recording_state") { + this.logger.debug({ isRecording: message.isRecording }, "Recording state"); + return; + } + + const session = this.sessions.get(ws); + if (!session) { + this.logger.error("No session found for client"); + return; + } + + if (message.type === "session") { + if (message.message.type === "create_agent_request") { + this.logger.debug( + { + cwd: message.message.config.cwd, + initialMode: message.message.config.modeId, + worktreeName: message.message.worktreeName, + requestId: message.message.requestId, + }, + "create_agent_request details" + ); + } + await session.handleMessage(message.message); + } + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + let rawPayload: string | null = null; + let parsedPayload: unknown = null; + + try { + const buffer = Array.isArray(data) + ? Buffer.concat( + data.map((item) => (Buffer.isBuffer(item) ? item : Buffer.from(item as ArrayBuffer))) + ) + : Buffer.isBuffer(data) + ? data + : Buffer.from(data as ArrayBuffer); + rawPayload = buffer.toString(); + parsedPayload = JSON.parse(rawPayload); + } catch (payloadError) { + rawPayload = rawPayload ?? ""; + parsedPayload = parsedPayload ?? rawPayload; + const payloadErr = + payloadError instanceof Error ? payloadError : new Error(String(payloadError)); + this.logger.error({ err: payloadErr }, "Failed to decode raw payload"); + } + + const trimmedRawPayload = + typeof rawPayload === "string" && rawPayload.length > 2000 + ? `${rawPayload.slice(0, 2000)}... (truncated)` + : rawPayload; + + this.logger.error( + { + err, + rawPayload: trimmedRawPayload, + parsedPayload, + }, + "Failed to parse/handle message" + ); + + this.sendToClient( + ws, + wrapSessionMessage({ + type: "status", + payload: { + status: "error", + message: `Invalid message: ${err.message}`, + }, + }) + ); + } + } + + private sendToClient(ws: WebSocket, message: WSOutboundMessage): void { + // WebSocket.OPEN = 1 + if (ws.readyState === 1) { + ws.send(JSON.stringify(message)); + } + } + + public broadcast(message: WSOutboundMessage): void { + const payload = JSON.stringify(message); + for (const ws of this.sessions.keys()) { + if (ws.readyState === 1) { + ws.send(payload); + } + } + } + + public async closeAll(): Promise { + const cleanupPromises: Promise[] = []; + for (const [ws, session] of this.sessions) { + cleanupPromises.push(session.cleanup()); + cleanupPromises.push( + new Promise((resolve) => { + // WebSocket.CLOSED = 3 + if (ws.readyState === 3) { + resolve(); + return; + } + ws.once("close", () => resolve()); + ws.close(); + }) + ); + } + await Promise.all(cleanupPromises); + this.sessions.clear(); + } + + private readonly ACTIVITY_THRESHOLD_MS = 120_000; + + private getClientActivityState(session: Session): { + deviceType: "web" | "mobile" | null; + focusedAgentId: string | null; + isStale: boolean; + appVisible: boolean; + } { + const activity = session.getClientActivity(); + if (!activity) { + this.logger.debug("getClientActivityState: no activity for session"); + return { deviceType: null, focusedAgentId: null, isStale: true, appVisible: false }; + } + const now = Date.now(); + const ageMs = now - activity.lastActivityAt.getTime(); + const isStale = ageMs >= this.ACTIVITY_THRESHOLD_MS; + this.logger.debug( + { + deviceType: activity.deviceType, + focusedAgentId: activity.focusedAgentId, + lastActivityAt: activity.lastActivityAt.toISOString(), + ageMs, + isStale, + appVisible: activity.appVisible, + }, + "getClientActivityState" + ); + return { + deviceType: activity.deviceType, + focusedAgentId: activity.focusedAgentId, + isStale, + appVisible: activity.appVisible, + }; + } + + private computeShouldNotifyForClient( + clientState: { + deviceType: "web" | "mobile" | null; + focusedAgentId: string | null; + isStale: boolean; + appVisible: boolean; + }, + allClientStates: Array<{ + deviceType: "web" | "mobile" | null; + focusedAgentId: string | null; + isStale: boolean; + appVisible: boolean; + }>, + agentId: string + ): boolean { + const isAnyoneActiveOnAgent = allClientStates.some( + (state) => state.focusedAgentId === agentId && state.appVisible && !state.isStale + ); + if (isAnyoneActiveOnAgent) { + return false; + } + + if (clientState.deviceType === null) { + return true; + } + + if (!clientState.isStale && clientState.appVisible && clientState.focusedAgentId !== null) { + return true; + } + + if (!clientState.isStale) { + return false; + } + + const hasActiveWebClient = allClientStates.some( + (state) => state.deviceType === "web" && !state.isStale + ); + + if (clientState.deviceType === "mobile") { + return !hasActiveWebClient; + } + + if (clientState.deviceType === "web") { + const hasOtherClient = allClientStates.some( + (state) => state !== clientState && (state.deviceType === "mobile" || state.deviceType === null) + ); + return !hasOtherClient; + } + + return true; + } + + private broadcastAgentAttention(params: { + agentId: string; + provider: AgentProvider; + reason: "finished" | "error" | "permission"; + }): void { + const clientEntries: Array<{ + ws: WebSocket; + state: { + deviceType: "web" | "mobile" | null; + focusedAgentId: string | null; + isStale: boolean; + appVisible: boolean; + }; + }> = []; + + for (const [ws, session] of this.sessions) { + clientEntries.push({ + ws, + state: this.getClientActivityState(session), + }); + } + + const allStates = clientEntries.map((e) => e.state); + + this.logger.debug( + { + agentId: params.agentId, + reason: params.reason, + clientCount: clientEntries.length, + allStates, + }, + "broadcastAgentAttention" + ); + + const allClientsStale = allStates.every((state) => state.isStale); + this.logger.debug({ allClientsStale }, "Client staleness check"); + if (allClientsStale) { + const tokens = this.pushTokenStore.getAllTokens(); + this.logger.info({ tokenCount: tokens.length }, "Sending push notification"); + if (tokens.length > 0) { + void this.pushService.sendPush(tokens, { + title: "Agent needs attention", + body: `Reason: ${params.reason}`, + data: { agentId: params.agentId }, + }); + } + } + + for (const { ws, state } of clientEntries) { + const shouldNotify = this.computeShouldNotifyForClient(state, allStates, params.agentId); + + const message = wrapSessionMessage({ + type: "agent_stream", + payload: { + agentId: params.agentId, + event: { + type: "attention_required", + provider: params.provider, + reason: params.reason, + timestamp: new Date().toISOString(), + shouldNotify, + }, + timestamp: new Date().toISOString(), + }, + }); + + this.sendToClient(ws, message); + } + } +} diff --git a/packages/server/src/shared/messages.ts b/packages/server/src/shared/messages.ts index 88a681b37..fc9dfa449 100644 --- a/packages/server/src/shared/messages.ts +++ b/packages/server/src/shared/messages.ts @@ -518,6 +518,14 @@ export const ListCommandsRequestSchema = z.object({ requestId: z.string(), }); +export const ExecuteCommandRequestSchema = z.object({ + type: z.literal("execute_command_request"), + agentId: z.string(), + commandName: z.string(), + args: z.string().optional(), + requestId: z.string(), +}); + export const RegisterPushTokenMessageSchema = z.object({ type: z.literal("register_push_token"), token: z.string(), @@ -604,6 +612,7 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ ClearAgentAttentionMessageSchema, ClientHeartbeatMessageSchema, ListCommandsRequestSchema, + ExecuteCommandRequestSchema, RegisterPushTokenMessageSchema, ListTerminalsRequestSchema, CreateTerminalRequestSchema, @@ -937,6 +946,22 @@ export const ListCommandsResponseSchema = z.object({ }), }); +const AgentCommandResultSchema = z.object({ + text: z.string(), + timeline: z.array(AgentTimelineItemPayloadSchema), + usage: AgentUsageSchema.optional(), +}); + +export const ExecuteCommandResponseSchema = z.object({ + type: z.literal("execute_command_response"), + payload: z.object({ + agentId: z.string(), + result: AgentCommandResultSchema.nullable(), + error: z.string().nullable(), + requestId: z.string(), + }), +}); + // ============================================================================ // Terminal Outbound Messages // ============================================================================ @@ -1037,6 +1062,7 @@ export const SessionOutboundMessageSchema = z.discriminatedUnion("type", [ GitRepoInfoResponseSchema, ListProviderModelsResponseMessageSchema, ListCommandsResponseSchema, + ExecuteCommandResponseSchema, ListTerminalsResponseSchema, CreateTerminalResponseSchema, SubscribeTerminalResponseSchema, @@ -1104,6 +1130,8 @@ export type ClearAgentAttentionMessage = z.infer; export type ListCommandsRequest = z.infer; export type ListCommandsResponse = z.infer; +export type ExecuteCommandRequest = z.infer; +export type ExecuteCommandResponse = z.infer; export type RegisterPushTokenMessage = z.infer; // Terminal message types