From c3c8103563c120c6c9a487d0cecd3a0a6a3dacd6 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 21 Feb 2026 15:33:29 +0700 Subject: [PATCH] refactor: migrate from clientSessionKey to clientId identity --- docs/claude-provider-invariants-test-plan.md | 81 +++ .../app/src/components/agent-input-area.tsx | 4 + .../app/src/components/tool-call-details.tsx | 42 +- packages/app/src/contexts/session-context.tsx | 52 +- packages/app/src/hooks/use-all-agents-list.ts | 5 +- .../app/src/hooks/use-sidebar-agents-list.ts | 5 +- packages/app/src/runtime/host-runtime.test.ts | 26 +- packages/app/src/runtime/host-runtime.ts | 40 +- packages/app/src/types/stream.test.ts | 79 +++ ...-session-key.test.ts => client-id.test.ts} | 24 +- packages/app/src/utils/client-id.ts | 54 ++ packages/app/src/utils/client-session-key.ts | 54 -- packages/app/src/utils/daemon-endpoints.ts | 1 - .../src/utils/test-daemon-connection.test.ts | 61 +-- .../app/src/utils/test-daemon-connection.ts | 84 ++- .../app/src/utils/tool-call-display.test.ts | 23 +- packages/app/src/utils/tool-call-icon.test.ts | 51 ++ packages/app/src/utils/tool-call-icon.ts | 8 +- .../{client-session-key.ts => client-id.ts} | 24 +- packages/cli/src/utils/client.ts | 21 +- packages/relay/src/cloudflare-adapter.test.ts | 40 +- packages/relay/src/cloudflare-adapter.ts | 132 ++--- packages/relay/src/e2e.test.ts | 26 +- packages/relay/src/live-relay.e2e.test.ts | 10 +- packages/relay/src/types.ts | 6 +- .../server/src/client/daemon-client.test.ts | 117 +++-- packages/server/src/client/daemon-client.ts | 129 ++--- .../src/server/agent/activity-curator.test.ts | 43 +- .../src/server/agent/agent-manager.test.ts | 258 ++++++++++ .../server/src/server/agent/agent-manager.ts | 83 ++- .../src/server/agent/agent-sdk-types.ts | 11 + ...agent.interrupt-restart-regression.test.ts | 463 ++++++++++++++++- .../claude-agent.sub-agent-sidechain.test.ts | 351 +++++++++++++ .../agent/providers/claude-agent.test.ts | 16 +- .../server/agent/providers/claude-agent.ts | 479 ++++++++++++++++-- .../src/server/daemon-client.e2e.test.ts | 58 +-- .../src/server/daemon-e2e/checkout-debug.ts | 6 +- .../daemon-e2e/relay-transport.e2e.test.ts | 45 +- ...hile-running-stuck-claude.real.e2e.test.ts | 41 +- ...end-while-running-stuck-test-utils.test.ts | 57 +++ .../send-while-running-stuck-test-utils.ts | 12 +- .../send-while-running-stuck.real.e2e.test.ts | 43 +- .../ui-action-stress.real.e2e.test.ts | 12 +- .../daemon-e2e/wait-for-idle.e2e.test.ts | 60 +++ .../server/src/server/relay-transport.test.ts | 2 +- packages/server/src/server/relay-transport.ts | 61 ++- packages/server/src/server/session.ts | 21 +- .../src/server/test-utils/daemon-client.ts | 10 +- .../websocket-server.relay-reconnect.test.ts | 348 ++++++++----- .../server/src/server/websocket-server.ts | 398 ++++++++++----- .../src/shared/daemon-endpoints.test.ts | 16 +- .../server/src/shared/daemon-endpoints.ts | 31 +- .../shared/messages.stream-parsing.test.ts | 42 ++ .../shared/messages.tool-call-schema.test.ts | 37 ++ .../src/shared/tool-call-display.test.ts | 23 +- .../server/src/shared/tool-call-display.ts | 4 + 56 files changed, 3289 insertions(+), 941 deletions(-) create mode 100644 docs/claude-provider-invariants-test-plan.md rename packages/app/src/utils/{client-session-key.test.ts => client-id.test.ts} (57%) create mode 100644 packages/app/src/utils/client-id.ts delete mode 100644 packages/app/src/utils/client-session-key.ts create mode 100644 packages/app/src/utils/tool-call-icon.test.ts rename packages/cli/src/utils/{client-session-key.ts => client-id.ts} (58%) create mode 100644 packages/server/src/server/agent/providers/claude-agent.sub-agent-sidechain.test.ts create mode 100644 packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.test.ts diff --git a/docs/claude-provider-invariants-test-plan.md b/docs/claude-provider-invariants-test-plan.md new file mode 100644 index 000000000..91af7902c --- /dev/null +++ b/docs/claude-provider-invariants-test-plan.md @@ -0,0 +1,81 @@ +# Claude Provider Invariants and Regression Test Plan + +## Goal +Lock down behavior for Claude provider stream handling without changing provider interface contracts. + +## Non-Negotiable Contract Invariants + +- `I1` Provider interface remains unchanged. + - `stream(prompt) -> AsyncGenerator` stays the only provider surface for run streaming. + - Reference: `packages/server/src/server/agent/agent-sdk-types.ts:337` + +- `I2` One prompt submission produces exactly one terminal turn event. + - Terminal set: `turn_completed | turn_failed | turn_canceled`. + - No ambiguous completion path. + +- `I3` Wait semantics remain status-based, not message-correlation-based. + - `waitForFinish` resolves on lifecycle outcomes (`idle|error|permission|timeout`). + - References: + - `packages/server/src/server/session.ts:5366` + - `packages/server/src/shared/messages.ts:1523` + +- `I4` App chat/input semantics remain status-transition-based. + - Input processing and queue flush continue to rely on `running` transitions and `updatedAt` ordering. + - References: + - `packages/app/src/components/agent-input-area.tsx:222` + - `packages/app/src/components/agent-input-area.tsx:300` + - `packages/app/src/contexts/session-context.tsx:485` + +- `I5` New prompt must never be satisfied by stale pre-prompt assistant/result events. + - Old queued events (including system-triggered activity) must not be misattributed as response to latest prompt. + +## Existing Coverage (Keep) + +- Wait semantics after immediate/rapid sends: + - `packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts:33` + +- Send-while-running recovery under real providers: + - `packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts:18` + - `packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts:18` + +- Claude stale old-turn protection (interrupt-failure path): + - `packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts:217` + +- App stream head/tail flush on terminal events: + - `packages/app/src/types/stream-event.test.ts:115` + +## Missing Coverage (Add) + +### M1: Task Notification Interleaving Regression (Claude Unit) + +- Add mocked SDK sequence where a queued system/task-notification-related burst causes old assistant/result events to appear before new prompt output. +- Assert latest `session.stream(newPrompt)` does not emit stale assistant text as its reply. +- Expected: `I5` enforced. + +### M2: Stale Result Preemption Regression (Claude Unit) + +- Add mocked SDK sequence with old pending `result` available right after pushing a new prompt. +- Assert provider ignores/contains stale terminal for previous activity and does not complete new prompt prematurely. +- Expected: `I2` + `I5` enforced. + +### M3: Wait Final Snapshot Coherence (Daemon E2E/Integration) + +- Add/strengthen test to ensure `waitForFinish` never returns `status: idle` while `final.status === running`. +- Remove need for retry loop workaround in stress scenarios. +- Current stress workaround location: + - `packages/server/src/server/daemon-e2e/ui-action-stress.real.e2e.test.ts:387` +- Expected: `I3` hard guarantee. + +## Implementation Notes (Provider-Internal Only) + +- Turn events are valid internal lifecycle markers. +- Do not introduce user-message to assistant-message correlation as a new cross-layer contract. +- Do not modify app/CLI/public daemon wait interfaces. +- Keep all changes scoped to Claude provider internals and regression tests. + +## Definition of Done + +- Invariants `I1-I5` are unchanged and explicitly covered. +- New tests for `M1-M3` are added and passing. +- Existing tests listed above remain green. +- `npm run typecheck` passes. diff --git a/packages/app/src/components/agent-input-area.tsx b/packages/app/src/components/agent-input-area.tsx index 27d2d9cda..2102b82c7 100644 --- a/packages/app/src/components/agent-input-area.tsx +++ b/packages/app/src/components/agent-input-area.tsx @@ -227,6 +227,10 @@ export function AgentInputArea({ useEffect(() => { const previousUpdatedAt = latestAgentUpdatedAtRef.current if (agentUpdatedAtMs < previousUpdatedAt) { + if (isProcessing && !isAgentRunning) { + prevIsAgentRunningRef.current = false + setIsProcessing(false) + } return } diff --git a/packages/app/src/components/tool-call-details.tsx b/packages/app/src/components/tool-call-details.tsx index 49a1a7a37..27181166a 100644 --- a/packages/app/src/components/tool-call-details.tsx +++ b/packages/app/src/components/tool-call-details.tsx @@ -48,7 +48,8 @@ export function ToolCallDetailsContent({ (detail?.type === "shell" || detail?.type === "edit" || detail?.type === "write" || - detail?.type === "read"); + detail?.type === "read" || + detail?.type === "sub_agent"); const codeBlockStyle = isFullBleed ? styles.fullBleedBlock : styles.diffContainer; if (detail?.type === "shell") { @@ -126,6 +127,45 @@ export function ToolCallDetailsContent({ ); + } else if (detail?.type === "sub_agent") { + const activityLog = detail.log.replace(/^\n+/, ""); + const hasLog = activityLog.length > 0; + const fallbackHeader = + detail.subAgentType && detail.description + ? `${detail.subAgentType}: ${detail.description}` + : detail.subAgentType ?? detail.description ?? "Sub-agent activity"; + sections.push( + + + + + + + {hasLog ? activityLog : fallbackHeader} + + + + + + + ); } else if (detail?.type === "edit") { sections.push( { const current = prev.get(agent.id); if (current && agent.updatedAt.getTime() < current.updatedAt.getTime()) { @@ -469,15 +467,53 @@ function SessionProviderInternal({ setAgentLastActivity(agent.id, agent.lastActivityAt); setPendingPermissions(serverId, (prev) => { - const next = new Map(prev); - for (const [key, pending] of Array.from(next.entries())) { + const existingKeysForAgent: string[] = []; + for (const [key, pending] of prev.entries()) { if (pending.agentId === agent.id) { - next.delete(key); + existingKeysForAgent.push(key); } } - for (const request of agent.pendingPermissions) { - const key = derivePendingPermissionKey(agent.id, request); - next.set(key, { key, agentId: agent.id, request }); + + const nextEntries = agent.pendingPermissions.map((request) => ({ + key: derivePendingPermissionKey(agent.id, request), + agentId: agent.id, + request, + })); + + let changed = existingKeysForAgent.length !== nextEntries.length; + if (!changed) { + const existingKeySet = new Set(existingKeysForAgent); + for (const entry of nextEntries) { + const existing = prev.get(entry.key); + if (!existingKeySet.has(entry.key) || !existing) { + changed = true; + break; + } + + const currentRequest = existing.request; + if ( + currentRequest.id !== entry.request.id || + currentRequest.kind !== entry.request.kind || + currentRequest.name !== entry.request.name || + currentRequest.title !== entry.request.title || + currentRequest.description !== entry.request.description + ) { + changed = true; + break; + } + } + } + + if (!changed) { + return prev; + } + + const next = new Map(prev); + for (const key of existingKeysForAgent) { + next.delete(key); + } + for (const entry of nextEntries) { + next.set(entry.key, entry); } return next; }); diff --git a/packages/app/src/hooks/use-all-agents-list.ts b/packages/app/src/hooks/use-all-agents-list.ts index af1ba8ae5..f533453d4 100644 --- a/packages/app/src/hooks/use-all-agents-list.ts +++ b/packages/app/src/hooks/use-all-agents-list.ts @@ -47,10 +47,9 @@ export function useAllAgentsList(options?: { : null; }, [options?.serverId]); - const session = useSessionStore((state) => - serverId ? state.sessions[serverId] : undefined + const liveAgents = useSessionStore((state) => + serverId ? state.sessions[serverId]?.agents ?? null : null ); - const liveAgents = session?.agents ?? null; const { snapshot } = useHostRuntimeSession(serverId ?? ""); const refreshAll = useCallback(() => { diff --git a/packages/app/src/hooks/use-sidebar-agents-list.ts b/packages/app/src/hooks/use-sidebar-agents-list.ts index 10fd2d488..c3777a6d4 100644 --- a/packages/app/src/hooks/use-sidebar-agents-list.ts +++ b/packages/app/src/hooks/use-sidebar-agents-list.ts @@ -168,10 +168,9 @@ export function useSidebarAgentsList(options?: { [options?.selectedProjectFilterKeys] ); - const session = useSessionStore((state) => - serverId ? state.sessions[serverId] : undefined + const liveAgents = useSessionStore((state) => + serverId ? state.sessions[serverId]?.agents ?? null : null ); - const liveAgents = session?.agents ?? null; const { snapshot } = useHostRuntimeSession(serverId ?? ""); const { entries, projectFilterOptions, hasAnyData, hasMoreEntries } = useMemo(() => { diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index d7bda5bf1..ec4739096 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/packages/app/src/runtime/host-runtime.test.ts @@ -134,7 +134,7 @@ function makeDeps( } return value; }, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }; } @@ -169,7 +169,7 @@ describe("HostRuntimeController", () => { measureLatency: async () => { throw new Error("probe unavailable"); }, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }; const controller = new HostRuntimeController({ host, @@ -188,7 +188,7 @@ describe("HostRuntimeController", () => { expect(controller.getSnapshot().agentDirectoryStatus).toBe("initial_loading"); }); - it("passes resolved client session key into created active clients", async () => { + it("passes resolved client id into created active clients", async () => { const host = makeHost({ connections: [ { @@ -198,23 +198,23 @@ describe("HostRuntimeController", () => { }, ], }); - const seenSessionKeys: string[] = []; + const seenClientIds: string[] = []; const fakeClient = new FakeDaemonClient(); const controller = new HostRuntimeController({ host, deps: { - createClient: ({ clientSessionKey }) => { - seenSessionKeys.push(clientSessionKey); + createClient: ({ clientId }) => { + seenClientIds.push(clientId); return fakeClient as unknown as DaemonClient; }, measureLatency: async () => 10, - getClientSessionKey: async () => "clsk_runtime_stable", + getClientId: async () => "cid_runtime_stable", }, }); await controller.start({ autoProbe: false }); - expect(seenSessionKeys).toEqual(["clsk_runtime_stable"]); + expect(seenClientIds).toEqual(["cid_runtime_stable"]); expect(controller.getSnapshot().connectionStatus).toBe("online"); }); @@ -535,7 +535,7 @@ describe("HostRuntimeController", () => { return client as unknown as DaemonClient; }, measureLatency: async () => 10, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }; const controller = new HostRuntimeController({ host, @@ -621,7 +621,7 @@ describe("HostRuntimeController", () => { } throw new Error("unexpected probe call"); }, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }, }); @@ -685,7 +685,7 @@ describe("HostRuntimeController", () => { } return 10; }, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }, }); @@ -730,7 +730,7 @@ describe("HostRuntimeStore", () => { deps: { createClient: () => fakeClient as unknown as DaemonClient, measureLatency: async () => 5, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }, }); @@ -772,7 +772,7 @@ describe("HostRuntimeStore", () => { measureLatency: async () => { throw new Error("probe unavailable"); }, - getClientSessionKey: async () => "clsk_test_runtime", + getClientId: async () => "cid_test_runtime", }, }); diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index 136f1256b..0b84fc904 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -10,7 +10,7 @@ import { buildRelayWebSocketUrl, } from "@/utils/daemon-endpoints"; import { measureConnectionLatency } from "@/utils/test-daemon-connection"; -import { getOrCreateClientSessionKey } from "@/utils/client-session-key"; +import { getOrCreateClientId } from "@/utils/client-id"; import { selectBestConnection, type ConnectionCandidate, @@ -118,14 +118,14 @@ export type HostRuntimeControllerDeps = { createClient: (input: { host: HostProfile; connection: HostConnection; - clientSessionKey: string; + clientId: string; runtimeGeneration: number; }) => DaemonClient; measureLatency: (input: { host: HostProfile; connection: HostConnection; }) => Promise; - getClientSessionKey: () => Promise; + getClientId: () => Promise; }; export type HostRuntimeStartOptions = { @@ -374,11 +374,12 @@ function selectInitialConnectionId(host: HostProfile): string | null { function createDefaultDeps(): HostRuntimeControllerDeps { return { - createClient: ({ host, connection, clientSessionKey, runtimeGeneration }) => { + createClient: ({ host, connection, clientId, runtimeGeneration }) => { const tauriTransportFactory = createTauriWebSocketTransportFactory(); const base = { suppressSendErrors: true, - clientSessionKey, + clientId, + clientType: "mobile" as const, runtimeGeneration, ...(tauriTransportFactory ? { transportFactory: tauriTransportFactory } @@ -387,7 +388,7 @@ function createDefaultDeps(): HostRuntimeControllerDeps { if (connection.type === "direct") { return new DaemonClient({ ...base, - url: buildDaemonWebSocketUrl(connection.endpoint, { clientSessionKey }), + url: buildDaemonWebSocketUrl(connection.endpoint), }); } return new DaemonClient({ @@ -395,7 +396,6 @@ function createDefaultDeps(): HostRuntimeControllerDeps { url: buildRelayWebSocketUrl({ endpoint: connection.relayEndpoint, serverId: host.serverId, - clientSessionKey, }), e2ee: { enabled: true, @@ -405,7 +405,7 @@ function createDefaultDeps(): HostRuntimeControllerDeps { }, measureLatency: ({ host, connection }) => measureConnectionLatency(connection, { serverId: host.serverId }), - getClientSessionKey: () => getOrCreateClientSessionKey(), + getClientId: () => getOrCreateClientId(), }; } @@ -421,8 +421,8 @@ export class HostRuntimeController { private started = false; private switchCandidateConnectionId: string | null = null; private switchCandidateHitCount = 0; - private clientSessionKeyPromise: Promise | null = null; - private clientSessionKeyHash: string | null = null; + private clientIdPromise: Promise | null = null; + private clientIdHash: string | null = null; private switchRequestVersion = 0; private probeRequestVersion = 0; @@ -740,7 +740,7 @@ export class HostRuntimeController { : toReasonCode(reason); console.info("[HostRuntimeTransition]", { serverId: this.host.serverId, - clientSessionKeyHash: this.clientSessionKeyHash, + clientIdHash: this.clientIdHash, from: input.from, to: input.to, event: event.type, @@ -782,9 +782,9 @@ export class HostRuntimeController { } const requestVersion = ++this.switchRequestVersion; - let clientSessionKey: string; + let clientId: string; try { - clientSessionKey = await this.resolveClientSessionKey(); + clientId = await this.resolveClientId(); } catch (error) { if (!this.isCurrentSwitchRequest(requestVersion)) { return; @@ -792,7 +792,7 @@ export class HostRuntimeController { const message = toErrorMessage(error); this.applyConnectionEvent({ type: "connect_failed", - message: `Failed to resolve client session key: ${message}`, + message: `Failed to resolve client id: ${message}`, }); this.updateSnapshot({ ...toSnapshotConnectionPatch(this.connectionMachineState), @@ -827,7 +827,7 @@ export class HostRuntimeController { const client = this.deps.createClient({ host: this.host, connection, - clientSessionKey, + clientId, runtimeGeneration: nextGeneration, }); if (!this.isCurrentSwitchRequest(requestVersion)) { @@ -910,14 +910,14 @@ export class HostRuntimeController { } } - private resolveClientSessionKey(): Promise { - if (!this.clientSessionKeyPromise) { - this.clientSessionKeyPromise = this.deps.getClientSessionKey().then((value) => { - this.clientSessionKeyHash = hashForLog(value); + private resolveClientId(): Promise { + if (!this.clientIdPromise) { + this.clientIdPromise = this.deps.getClientId().then((value) => { + this.clientIdHash = hashForLog(value); return value; }); } - return this.clientSessionKeyPromise; + return this.clientIdPromise; } } diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index 9ff8b1b42..a04ee626e 100644 --- a/packages/app/src/types/stream.test.ts +++ b/packages/app/src/types/stream.test.ts @@ -159,6 +159,85 @@ describe("stream reducer canonical tool calls", () => { }); }); + it("keeps sub_agent detail through lifecycle updates for the same callId", () => { + const callId = "task-sub-agent-1"; + const updates = [ + { + event: canonicalToolTimeline({ + provider: "claude", + callId, + name: "Task", + status: "running", + detail: { + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository structure", + log: "[Read] README.md\n[Bash] ls", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + { + index: 2, + toolName: "Bash", + summary: "ls", + }, + ], + }, + }), + timestamp: new Date("2025-01-01T10:12:00Z"), + }, + { + event: canonicalToolTimeline({ + provider: "claude", + callId, + name: "Task", + status: "completed", + input: null, + output: { ok: true }, + }), + timestamp: new Date("2025-01-01T10:12:01Z"), + }, + ]; + + const state = hydrateStreamState(updates); + const tools = state.filter(isAgentToolCallItem); + + assert.strictEqual(tools.length, 1); + assert.strictEqual(tools[0].payload.data.status, "completed"); + assert.deepStrictEqual(tools[0].payload.data.detail, { + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository structure", + log: "[Read] README.md\n[Bash] ls", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + { + index: 2, + toolName: "Bash", + summary: "ls", + }, + ], + }); + + const display = buildToolCallDisplayModel({ + name: tools[0].payload.data.name, + status: tools[0].payload.data.status, + error: tools[0].payload.data.error, + detail: tools[0].payload.data.detail, + }); + assert.deepStrictEqual(display, { + displayName: "Explore", + summary: "Inspect repository structure", + }); + }); + it("exposes shell summary from running input before completion", () => { const callId = "running-summary-shell"; const state = hydrateStreamState([ diff --git a/packages/app/src/utils/client-session-key.test.ts b/packages/app/src/utils/client-id.test.ts similarity index 57% rename from packages/app/src/utils/client-session-key.test.ts rename to packages/app/src/utils/client-id.test.ts index e3c8247e0..8b177eb43 100644 --- a/packages/app/src/utils/client-session-key.test.ts +++ b/packages/app/src/utils/client-id.test.ts @@ -9,35 +9,35 @@ vi.mock("@react-native-async-storage/async-storage", () => ({ default: asyncStorageMock, })); -describe("client-session-key", () => { +describe("client-id", () => { beforeEach(() => { vi.resetModules(); asyncStorageMock.getItem.mockReset(); asyncStorageMock.setItem.mockReset(); }); - it("returns stored client session key when present", async () => { - asyncStorageMock.getItem.mockResolvedValue("clsk_existing"); - const mod = await import("./client-session-key"); + it("returns stored client id when present", async () => { + asyncStorageMock.getItem.mockResolvedValue("cid_existing"); + const mod = await import("./client-id"); - const key = await mod.getOrCreateClientSessionKey(); - expect(key).toBe("clsk_existing"); + const key = await mod.getOrCreateClientId(); + expect(key).toBe("cid_existing"); expect(asyncStorageMock.getItem).toHaveBeenCalledTimes(1); expect(asyncStorageMock.setItem).not.toHaveBeenCalled(); }); - it("creates and persists a client session key when missing", async () => { + it("creates and persists a client id when missing", async () => { asyncStorageMock.getItem.mockResolvedValue(null); asyncStorageMock.setItem.mockResolvedValue(); vi.spyOn(globalThis.crypto, "randomUUID").mockReturnValue("12345678-1234-1234-1234-1234567890ab"); - const mod = await import("./client-session-key"); - const key = await mod.getOrCreateClientSessionKey(); + const mod = await import("./client-id"); + const key = await mod.getOrCreateClientId(); - expect(key).toBe("clsk_123456781234123412341234567890ab"); + expect(key).toBe("cid_123456781234123412341234567890ab"); expect(asyncStorageMock.setItem).toHaveBeenCalledWith( - "@paseo:client-session-key-v1", - "clsk_123456781234123412341234567890ab" + "@paseo:client-id-v1", + "cid_123456781234123412341234567890ab" ); }); }); diff --git a/packages/app/src/utils/client-id.ts b/packages/app/src/utils/client-id.ts new file mode 100644 index 000000000..19f307a91 --- /dev/null +++ b/packages/app/src/utils/client-id.ts @@ -0,0 +1,54 @@ +import AsyncStorage from "@react-native-async-storage/async-storage"; + +const CLIENT_ID_STORAGE_KEY = "@paseo:client-id-v1"; + +let cachedClientId: string | null = null; +let inFlightClientId: Promise | null = null; + +function generateClientId(): string { + const randomUuid = (() => { + const cryptoObj = globalThis.crypto as { randomUUID?: () => string } | undefined; + if (cryptoObj && typeof cryptoObj.randomUUID === "function") { + return cryptoObj.randomUUID().replace(/-/g, ""); + } + return `${Date.now().toString(36)}${Math.random().toString(36).slice(2)}`; + })(); + return `cid_${randomUuid}`; +} + +function normalizeStoredClientId(value: unknown): string | null { + if (typeof value !== "string") { + return null; + } + const trimmed = value.trim(); + return trimmed.length > 0 ? trimmed : null; +} + +export async function getOrCreateClientId(): Promise { + if (cachedClientId) { + return cachedClientId; + } + if (inFlightClientId) { + return inFlightClientId; + } + + inFlightClientId = (async () => { + const storedValue = await AsyncStorage.getItem(CLIENT_ID_STORAGE_KEY); + const existing = normalizeStoredClientId(storedValue); + if (existing) { + cachedClientId = existing; + return existing; + } + + const nextValue = generateClientId(); + await AsyncStorage.setItem(CLIENT_ID_STORAGE_KEY, nextValue); + cachedClientId = nextValue; + return nextValue; + })(); + + try { + return await inFlightClientId; + } finally { + inFlightClientId = null; + } +} diff --git a/packages/app/src/utils/client-session-key.ts b/packages/app/src/utils/client-session-key.ts deleted file mode 100644 index 8b7542fba..000000000 --- a/packages/app/src/utils/client-session-key.ts +++ /dev/null @@ -1,54 +0,0 @@ -import AsyncStorage from "@react-native-async-storage/async-storage"; - -const CLIENT_SESSION_KEY_STORAGE_KEY = "@paseo:client-session-key-v1"; - -let cachedClientSessionKey: string | null = null; -let inFlightClientSessionKey: Promise | null = null; - -function generateClientSessionKey(): string { - const randomUuid = (() => { - const cryptoObj = globalThis.crypto as { randomUUID?: () => string } | undefined; - if (cryptoObj && typeof cryptoObj.randomUUID === "function") { - return cryptoObj.randomUUID().replace(/-/g, ""); - } - return `${Date.now().toString(36)}${Math.random().toString(36).slice(2)}`; - })(); - return `clsk_${randomUuid}`; -} - -function normalizeStoredClientSessionKey(value: unknown): string | null { - if (typeof value !== "string") { - return null; - } - const trimmed = value.trim(); - return trimmed.length > 0 ? trimmed : null; -} - -export async function getOrCreateClientSessionKey(): Promise { - if (cachedClientSessionKey) { - return cachedClientSessionKey; - } - if (inFlightClientSessionKey) { - return inFlightClientSessionKey; - } - - inFlightClientSessionKey = (async () => { - const storedValue = await AsyncStorage.getItem(CLIENT_SESSION_KEY_STORAGE_KEY); - const existing = normalizeStoredClientSessionKey(storedValue); - if (existing) { - cachedClientSessionKey = existing; - return existing; - } - - const nextValue = generateClientSessionKey(); - await AsyncStorage.setItem(CLIENT_SESSION_KEY_STORAGE_KEY, nextValue); - cachedClientSessionKey = nextValue; - return nextValue; - })(); - - try { - return await inFlightClientSessionKey; - } finally { - inFlightClientSessionKey = null; - } -} diff --git a/packages/app/src/utils/daemon-endpoints.ts b/packages/app/src/utils/daemon-endpoints.ts index 33cddc5bd..8a0c76e29 100644 --- a/packages/app/src/utils/daemon-endpoints.ts +++ b/packages/app/src/utils/daemon-endpoints.ts @@ -33,7 +33,6 @@ export function decodeOfferFragmentPayload(encoded: string): unknown { export function buildRelayWebSocketUrl(params: { endpoint: string; serverId: string; - clientSessionKey?: string; }): string { return buildSharedRelayWebSocketUrl({ ...params, role: "client" }); } diff --git a/packages/app/src/utils/test-daemon-connection.test.ts b/packages/app/src/utils/test-daemon-connection.test.ts index 27111529d..6a67b1443 100644 --- a/packages/app/src/utils/test-daemon-connection.test.ts +++ b/packages/app/src/utils/test-daemon-connection.test.ts @@ -1,21 +1,19 @@ import { describe, expect, it, vi, beforeEach } from "vitest"; const daemonClientMock = vi.hoisted(() => { - const createdConfigs: Array<{ clientSessionKey?: string; url?: string }> = []; + const createdConfigs: Array<{ clientId?: string; url?: string }> = []; class MockDaemonClient { - private statusHandlers = new Set< - ( - message: { - type: "status"; - payload: { status: string; serverId: string; hostname: string | null }; - } - ) => void - >(); - public lastError: string | null = null; + private lastWelcome = { + type: "welcome" as const, + serverId: "srv_probe_test", + hostname: "probe-host" as string | null, + version: "0.0.0", + resumed: false, + }; - constructor(config: { clientSessionKey?: string; url?: string }) { + constructor(config: { clientId?: string; url?: string }) { createdConfigs.push(config); } @@ -23,35 +21,16 @@ const daemonClientMock = vi.hoisted(() => { return () => undefined; } - on( - event: "status", - handler: ( - message: { - type: "status"; - payload: { status: string; serverId: string; hostname: string | null }; - } - ) => void - ): () => void { - if (event === "status") { - this.statusHandlers.add(handler); - } - return () => { - this.statusHandlers.delete(handler); - }; + on(): () => void { + return () => undefined; } async connect(): Promise { - const message = { - type: "status" as const, - payload: { - status: "server_info", - serverId: "srv_probe_test", - hostname: "probe-host", - }, - }; - for (const handler of this.statusHandlers) { - handler(message); - } + return; + } + + getLastWelcomeMessage() { + return this.lastWelcome; } async ping(): Promise<{ rttMs: number }> { @@ -78,7 +57,7 @@ describe("test-daemon-connection probe client identity", () => { daemonClientMock.createdConfigs.length = 0; }); - it("uses isolated probe clientSessionKey values for direct latency probes", async () => { + it("uses isolated probe clientId values for direct latency probes", async () => { const mod = await import("./test-daemon-connection"); await mod.measureConnectionLatency({ @@ -93,8 +72,8 @@ describe("test-daemon-connection probe client identity", () => { }); const [first, second] = daemonClientMock.createdConfigs; - expect(first?.clientSessionKey).toMatch(/^clsk_probe_/); - expect(second?.clientSessionKey).toMatch(/^clsk_probe_/); - expect(first?.clientSessionKey).not.toBe(second?.clientSessionKey); + expect(first?.clientId).toMatch(/^cid_probe_/); + expect(second?.clientId).toMatch(/^cid_probe_/); + expect(first?.clientId).not.toBe(second?.clientId); }); }); diff --git a/packages/app/src/utils/test-daemon-connection.ts b/packages/app/src/utils/test-daemon-connection.ts index 4c9edfbf3..324b4681e 100644 --- a/packages/app/src/utils/test-daemon-connection.ts +++ b/packages/app/src/utils/test-daemon-connection.ts @@ -1,10 +1,9 @@ import { DaemonClient } from "@server/client/daemon-client"; import type { DaemonClientConfig } from "@server/client/daemon-client"; -import { parseServerInfoStatusPayload } from "@server/shared/messages"; import type { HostConnection } from "@/contexts/daemon-registry-context"; import { buildDaemonWebSocketUrl, buildRelayWebSocketUrl } from "./daemon-endpoints"; import { createTauriWebSocketTransportFactory } from "./tauri-daemon-transport"; -function createProbeClientSessionKey(): string { +function createProbeClientId(): string { const randomUuid = (() => { const cryptoObj = globalThis.crypto as { randomUUID?: () => string } | undefined; if (cryptoObj && typeof cryptoObj.randomUUID === "function") { @@ -12,7 +11,7 @@ function createProbeClientSessionKey(): string { } return `${Date.now().toString(36)}${Math.random().toString(36).slice(2)}`; })(); - return `clsk_probe_${randomUuid}`; + return `cid_probe_${randomUuid}`; } function normalizeNonEmptyString(value: unknown): string | null { @@ -55,10 +54,11 @@ async function buildClientConfig( connection: HostConnection, serverId?: string ): Promise { - const clientSessionKey = createProbeClientSessionKey(); + const clientId = createProbeClientId(); const tauriTransportFactory = createTauriWebSocketTransportFactory(); const base = { - clientSessionKey, + clientId, + clientType: "mobile" as const, suppressSendErrors: true, ...(tauriTransportFactory ? { transportFactory: tauriTransportFactory } : {}), }; @@ -66,7 +66,7 @@ async function buildClientConfig( if (connection.type === "direct") { return { ...base, - url: buildDaemonWebSocketUrl(connection.endpoint, { clientSessionKey }), + url: buildDaemonWebSocketUrl(connection.endpoint), }; } @@ -79,7 +79,6 @@ async function buildClientConfig( url: buildRelayWebSocketUrl({ endpoint: connection.relayEndpoint, serverId, - clientSessionKey, }), e2ee: { enabled: true, daemonPublicKeyB64: connection.daemonPublicKeyB64 }, }; @@ -92,35 +91,9 @@ function connectAndProbe( const client = new DaemonClient(config); return new Promise<{ client: DaemonClient; serverId: string; hostname: string | null }>((resolve, reject) => { - let cleanedUp = false; - let unsubscribe: (() => void) | null = null; - let unsubscribeStatus: (() => void) | null = null; - let serverId: string | null = null; - let hostname: string | null = null; - - const cleanup = () => { - if (cleanedUp) return; - cleanedUp = true; - clearTimeout(timer); - unsubscribe?.(); - unsubscribeStatus?.(); - }; - - const maybeFinishOk = () => { - if (!serverId) return; - cleanup(); - resolve({ client, serverId, hostname }); - }; - - const finishErr = (error: Error) => { - if (cleanedUp) return; - cleanup(); - client.close().catch(() => undefined); - reject(error); - }; - const timer = setTimeout(() => { - finishErr( + void client.close().catch(() => undefined); + reject( new DaemonConnectionTestError("Connection timed out", { reason: "Connection timed out", lastError: client.lastError ?? null, @@ -128,25 +101,32 @@ function connectAndProbe( ); }, timeoutMs); - unsubscribe = client.subscribeConnectionStatus((state) => { - if (state.status === "disconnected") { - const reason = normalizeNonEmptyString(state.reason); - const lastError = normalizeNonEmptyString(client.lastError); - const message = pickBestReason(reason, lastError); - finishErr(new DaemonConnectionTestError(message, { reason, lastError })); + void client.connect().then(() => { + clearTimeout(timer); + const welcome = client.getLastWelcomeMessage(); + if (!welcome) { + void client.close().catch(() => undefined); + reject( + new DaemonConnectionTestError("Missing welcome message", { + reason: "Missing welcome message", + lastError: client.lastError ?? null, + }) + ); + return; } + resolve({ + client, + serverId: welcome.serverId, + hostname: welcome.hostname, + }); + }).catch((error) => { + clearTimeout(timer); + const reason = normalizeNonEmptyString(error instanceof Error ? error.message : String(error)); + const lastError = normalizeNonEmptyString(client.lastError); + const message = pickBestReason(reason, lastError); + void client.close().catch(() => undefined); + reject(new DaemonConnectionTestError(message, { reason, lastError })); }); - - unsubscribeStatus = client.on("status", (message) => { - if (message.type !== "status") return; - const payload = parseServerInfoStatusPayload(message.payload); - if (!payload) return; - serverId = payload.serverId; - hostname = payload.hostname; - maybeFinishOk(); - }); - - void client.connect().catch(() => undefined); }); } diff --git a/packages/app/src/utils/tool-call-display.test.ts b/packages/app/src/utils/tool-call-display.test.ts index 04ea25cc2..e7e08f47e 100644 --- a/packages/app/src/utils/tool-call-display.test.ts +++ b/packages/app/src/utils/tool-call-display.test.ts @@ -38,24 +38,29 @@ describe("tool-call-display", () => { }); }); - it("uses metadata summary for task tool calls", () => { + it("uses sub-agent detail for task label and description", () => { const display = buildToolCallDisplayModel({ name: "task", status: "running", error: null, detail: { - type: "unknown", - input: null, - output: null, - }, - metadata: { - subAgentActivity: "Running tests", + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository structure", + log: "[Read] README.md", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + ], }, }); expect(display).toEqual({ - displayName: "Task", - summary: "Running tests", + displayName: "Explore", + summary: "Inspect repository structure", }); }); diff --git a/packages/app/src/utils/tool-call-icon.test.ts b/packages/app/src/utils/tool-call-icon.test.ts new file mode 100644 index 000000000..a3e997c1e --- /dev/null +++ b/packages/app/src/utils/tool-call-icon.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it, vi } from "vitest"; + +const iconMocks = vi.hoisted(() => ({ + Bot: () => null, + Brain: () => null, + Eye: () => null, + MicVocal: () => null, + Pencil: () => null, + Search: () => null, + SquareTerminal: () => null, + Wrench: () => null, +})); + +vi.mock("lucide-react-native", () => iconMocks); + +import { resolveToolCallIcon } from "./tool-call-icon"; + +describe("tool-call-icon", () => { + it("uses robot icon for task sub-agent details", () => { + const icon = resolveToolCallIcon("Task", { + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository", + log: "[Read] README.md", + actions: [{ index: 1, toolName: "Read", summary: "README.md" }], + }); + + expect(icon).toBe(iconMocks.Bot); + }); + + it("uses robot icon for task calls without canonical detail", () => { + const icon = resolveToolCallIcon("Task", { + type: "unknown", + input: null, + output: null, + }); + + expect(icon).toBe(iconMocks.Bot); + expect(resolveToolCallIcon("Task")).toBe(iconMocks.Bot); + }); + + it("keeps thinking icon override for unknown detail", () => { + const icon = resolveToolCallIcon("thinking", { + type: "unknown", + input: null, + output: null, + }); + + expect(icon).toBe(iconMocks.Brain); + }); +}); diff --git a/packages/app/src/utils/tool-call-icon.ts b/packages/app/src/utils/tool-call-icon.ts index c0087b9fc..10ba09bc0 100644 --- a/packages/app/src/utils/tool-call-icon.ts +++ b/packages/app/src/utils/tool-call-icon.ts @@ -20,6 +20,7 @@ const TOOL_DETAIL_ICONS: Record = write: Pencil, search: Search, worktree_setup: SquareTerminal, + sub_agent: Bot, unknown: Wrench, }; @@ -33,13 +34,12 @@ export function resolveToolCallIcon(toolName: string, detail?: ToolCallDetail): if (lowerName === "speak") { return MicVocal; } + if (lowerName === "task") { + return Bot; + } if (detail) { return TOOL_DETAIL_ICONS[detail.type]; } - - if (lowerName === "task") { - return Bot; - } return Wrench; } diff --git a/packages/cli/src/utils/client-session-key.ts b/packages/cli/src/utils/client-id.ts similarity index 58% rename from packages/cli/src/utils/client-session-key.ts rename to packages/cli/src/utils/client-id.ts index 0387d03fe..af21b7760 100644 --- a/packages/cli/src/utils/client-session-key.ts +++ b/packages/cli/src/utils/client-id.ts @@ -5,31 +5,31 @@ import { homedir } from "node:os"; const CLIENT_SESSION_KEY_FILE = join( process.env.PASEO_HOME ?? join(homedir(), ".paseo"), - "cli-client-session-key" + "cli-client-id" ); -let cachedClientSessionKey: string | null = null; +let cachedClientId: string | null = null; -function normalizeClientSessionKey(value: string): string | null { +function normalizeClientId(value: string): string | null { const trimmed = value.trim(); return trimmed.length > 0 ? trimmed : null; } -function generateClientSessionKey(): string { - return `clsk_${randomUUID().replace(/-/g, "")}`; +function generateClientId(): string { + return `cid_${randomUUID().replace(/-/g, "")}`; } -export async function getOrCreateCliClientSessionKey(): Promise { - if (cachedClientSessionKey) { - return cachedClientSessionKey; +export async function getOrCreateCliClientId(): Promise { + if (cachedClientId) { + return cachedClientId; } try { - const existing = normalizeClientSessionKey( + const existing = normalizeClientId( await readFile(CLIENT_SESSION_KEY_FILE, "utf8") ); if (existing) { - cachedClientSessionKey = existing; + cachedClientId = existing; return existing; } } catch (error) { @@ -39,9 +39,9 @@ export async function getOrCreateCliClientSessionKey(): Promise { } } - const nextValue = generateClientSessionKey(); + const nextValue = generateClientId(); await mkdir(dirname(CLIENT_SESSION_KEY_FILE), { recursive: true }); await writeFile(CLIENT_SESSION_KEY_FILE, nextValue, { mode: 0o600 }); - cachedClientSessionKey = nextValue; + cachedClientId = nextValue; return nextValue; } diff --git a/packages/cli/src/utils/client.ts b/packages/cli/src/utils/client.ts index 9129f2709..ddd667df4 100644 --- a/packages/cli/src/utils/client.ts +++ b/packages/cli/src/utils/client.ts @@ -1,6 +1,6 @@ import { DaemonClient } from '@getpaseo/server' import WebSocket from 'ws' -import { getOrCreateCliClientSessionKey } from './client-session-key.js' +import { getOrCreateCliClientId } from './client-id.js' export interface ConnectOptions { host?: string @@ -40,15 +40,18 @@ function createNodeWebSocketFactory() { export async function connectToDaemon(options?: ConnectOptions): Promise { const host = getDaemonHost(options) const timeout = options?.timeout ?? DEFAULT_TIMEOUT - const clientSessionKey = await getOrCreateCliClientSessionKey() - const encodedSessionKey = encodeURIComponent(clientSessionKey) - const url = `ws://${host}/ws?clientSessionKey=${encodedSessionKey}` + const clientId = await getOrCreateCliClientId() + const url = `ws://${host}/ws` - const client = new DaemonClient({ - url, - webSocketFactory: createNodeWebSocketFactory(), - reconnect: { enabled: false }, - }) + const client = new DaemonClient( + { + url, + clientId, + clientType: 'cli', + webSocketFactory: createNodeWebSocketFactory(), + reconnect: { enabled: false }, + } as unknown as ConstructorParameters[0] + ) // Connect with timeout const connectPromise = client.connect() diff --git a/packages/relay/src/cloudflare-adapter.test.ts b/packages/relay/src/cloudflare-adapter.test.ts index ff1993510..c537fc83c 100644 --- a/packages/relay/src/cloudflare-adapter.test.ts +++ b/packages/relay/src/cloudflare-adapter.test.ts @@ -69,7 +69,7 @@ async function withMockWebSocketPair( } describe("RelayDurableObject versioning", () => { - it("accepts legacy v1 client sockets without clientId", async () => { + it("accepts legacy v1 client sockets without connectionId", async () => { const { state } = createMockState(); await withMockWebSocketPair(async () => { const relay = new RelayDurableObject(state as any); @@ -83,13 +83,21 @@ describe("RelayDurableObject versioning", () => { }); }); - it("rejects v2 client sockets when clientId is missing", async () => { + it("assigns a connectionId when v2 client connects without one", async () => { const { state } = createMockState(); - const relay = new RelayDurableObject(state as any); - const req = new Request("https://relay.test/ws?role=client&serverId=srv_test&v=2"); - const response = await relay.fetch(req); - expect(response.status).toBe(400); - await expect(response.text()).resolves.toBe("Missing clientId parameter"); + await withMockWebSocketPair(async ({ serverWs }) => { + const relay = new RelayDurableObject(state as any); + const req = new Request("https://relay.test/ws?role=client&serverId=srv_test&v=2", { + headers: { Upgrade: "websocket" }, + }); + await relay.fetch(req).catch(() => undefined); + expect(state.acceptWebSocket).toHaveBeenCalled(); + const attachment = serverWs.deserializeAttachment(); + expect(attachment).toMatchObject({ + role: "client", + connectionId: expect.stringMatching(/^conn_/), + }); + }); }); }); @@ -110,7 +118,7 @@ describe("RelayDurableObject control nudge/reset behavior", () => { setTagSockets(`server:${clientId}`, []); const relay = new RelayDurableObject(state as any); - (relay as any).nudgeOrResetControlForClient(clientId); + (relay as any).nudgeOrResetControlForConnection(clientId); vi.advanceTimersByTime(15_000); @@ -124,7 +132,7 @@ describe("RelayDurableObject control nudge/reset behavior", () => { const control = createMockSocket(); const client = createMockSocket({ role: "client", - clientId, + connectionId: clientId, serverId: "srv_test", createdAt: Date.now(), }); @@ -136,7 +144,7 @@ describe("RelayDurableObject control nudge/reset behavior", () => { setTagSockets(`server:${clientId}`, []); const relay = new RelayDurableObject(state as any); - (relay as any).nudgeOrResetControlForClient(clientId); + (relay as any).nudgeOrResetControlForConnection(clientId); vi.advanceTimersByTime(10_000); expect(control.send).toHaveBeenCalledTimes(1); @@ -145,11 +153,11 @@ describe("RelayDurableObject control nudge/reset behavior", () => { expect(control.close).toHaveBeenCalledWith(1011, "Control unresponsive"); }); - it("does not replace existing client sockets for the same clientId", async () => { + it("does not replace existing client sockets for the same connectionId", async () => { const existingClient = createMockSocket({ version: "2", role: "client", - clientId: "clt_same_session", + connectionId: "clt_same_session", serverId: "srv_test", createdAt: Date.now(), }); @@ -160,7 +168,7 @@ describe("RelayDurableObject control nudge/reset behavior", () => { await withMockWebSocketPair(async () => { const relay = new RelayDurableObject(state as any); const req = new Request( - "https://relay.test/ws?role=client&serverId=srv_test&clientId=clt_same_session&v=2", + "https://relay.test/ws?role=client&serverId=srv_test&connectionId=clt_same_session&v=2", { headers: { Upgrade: "websocket", @@ -178,14 +186,14 @@ describe("RelayDurableObject control nudge/reset behavior", () => { const disconnectedClient = createMockSocket({ version: "2", role: "client", - clientId, + connectionId: clientId, serverId: "srv_test", createdAt: Date.now(), }); const stillConnectedClient = createMockSocket({ version: "2", role: "client", - clientId, + connectionId: clientId, serverId: "srv_test", createdAt: Date.now(), }); @@ -208,7 +216,7 @@ describe("RelayDurableObject control nudge/reset behavior", () => { expect(serverData.close).not.toHaveBeenCalled(); expect(control.send).not.toHaveBeenCalledWith( - JSON.stringify({ type: "client_disconnected", clientId }) + JSON.stringify({ type: "disconnected", connectionId: clientId }) ); }); }); diff --git a/packages/relay/src/cloudflare-adapter.ts b/packages/relay/src/cloudflare-adapter.ts index 6c23dfa3e..f2ec4363f 100644 --- a/packages/relay/src/cloudflare-adapter.ts +++ b/packages/relay/src/cloudflare-adapter.ts @@ -74,9 +74,9 @@ interface DurableObjectStub { * - role=client: app/client socket * * v2 WebSockets connect in three shapes: - * - role=server (no clientId): daemon control socket (one per serverId) - * - role=server&clientId=...: daemon per-client data socket (one per clientId) - * - role=client&clientId=...: app/client socket (many per clientId) + * - role=server (no connectionId): daemon control socket (one per serverId) + * - role=server&connectionId=...: daemon per-connection data socket (one per connectionId) + * - role=client&connectionId=...: app/client socket (many per connectionId) */ interface CFResponseInit extends ResponseInit { webSocket?: WebSocket; @@ -84,7 +84,7 @@ interface CFResponseInit extends ResponseInit { export class RelayDurableObject { private state: DurableObjectState; - private pendingClientFrames = new Map>(); + private pendingFrames = new Map>(); constructor(state: DurableObjectState) { this.state = state; @@ -110,42 +110,42 @@ export class RelayDurableObject { } as CFResponseInit); } - private hasServerDataSocket(clientId: string): boolean { + private hasServerDataSocket(connectionId: string): boolean { try { - return this.state.getWebSockets(`server:${clientId}`).length > 0; + return this.state.getWebSockets(`server:${connectionId}`).length > 0; } catch { return false; } } - private hasClientSocket(clientId: string): boolean { + private hasClientSocket(connectionId: string): boolean { try { - return this.state.getWebSockets(`client:${clientId}`).length > 0; + return this.state.getWebSockets(`client:${connectionId}`).length > 0; } catch { return false; } } - private nudgeOrResetControlForClient(clientId: string): void { + private nudgeOrResetControlForConnection(connectionId: string): void { // If the daemon's control WS becomes half-open, the DO can't reliably detect it via ws.send errors // (Cloudflare may accept writes even if the other side is no longer reading). // - // Instead, observe whether the daemon reacts by opening the per-client server-data socket. + // Instead, observe whether the daemon reacts by opening the per-connection server-data socket. // If it doesn't, nudge with a sync message; if still no reaction, force-close the control // socket(s) so the daemon reconnects. const initialDelayMs = 10_000; const secondDelayMs = 5_000; setTimeout(() => { - if (!this.hasClientSocket(clientId)) return; - if (this.hasServerDataSocket(clientId)) return; + if (!this.hasClientSocket(connectionId)) return; + if (this.hasServerDataSocket(connectionId)) return; // First nudge: send a full sync list. - this.notifyControls({ type: "sync", clientIds: this.listConnectedClientIds() }); + this.notifyControls({ type: "sync", connectionIds: this.listConnectedConnectionIds() }); setTimeout(() => { - if (!this.hasClientSocket(clientId)) return; - if (this.hasServerDataSocket(clientId)) return; + if (!this.hasClientSocket(connectionId)) return; + if (this.hasServerDataSocket(connectionId)) return; // Still nothing: assume control is stuck and force a reconnect. for (const ws of this.state.getWebSockets("server-control")) { @@ -159,38 +159,38 @@ export class RelayDurableObject { }, initialDelayMs); } - private bufferClientFrame(clientId: string, message: string | ArrayBuffer): void { - const existing = this.pendingClientFrames.get(clientId) ?? []; + private bufferFrame(connectionId: string, message: string | ArrayBuffer): void { + const existing = this.pendingFrames.get(connectionId) ?? []; existing.push(message); // Prevent unbounded memory growth if a daemon never connects. if (existing.length > 200) { existing.splice(0, existing.length - 200); } - this.pendingClientFrames.set(clientId, existing); + this.pendingFrames.set(connectionId, existing); } - private flushClientFrames(clientId: string, serverWs: WebSocket): void { - const frames = this.pendingClientFrames.get(clientId); + private flushFrames(connectionId: string, serverWs: WebSocket): void { + const frames = this.pendingFrames.get(connectionId); if (!frames || frames.length === 0) return; - this.pendingClientFrames.delete(clientId); + this.pendingFrames.delete(connectionId); for (const frame of frames) { try { serverWs.send(frame); } catch { // If we can't flush, re-buffer and let the daemon re-establish. - this.bufferClientFrame(clientId, frame); + this.bufferFrame(connectionId, frame); break; } } } - private listConnectedClientIds(): string[] { + private listConnectedConnectionIds(): string[] { const out = new Set(); for (const ws of this.state.getWebSockets("client")) { try { const attachment = (ws as WebSocketWithAttachment).deserializeAttachment() as RelaySessionAttachment | null; - if (attachment?.role === "client" && typeof attachment.clientId === "string" && attachment.clientId) { - out.add(attachment.clientId); + if (attachment?.role === "client" && typeof attachment.connectionId === "string" && attachment.connectionId) { + out.add(attachment.connectionId); } } catch { // ignore @@ -230,7 +230,7 @@ export class RelayDurableObject { serverId, role, version: LEGACY_RELAY_VERSION, - clientId: null, + connectionId: null, createdAt: Date.now(), }; (server as WebSocketWithAttachment).serializeAttachment(attachment); @@ -244,30 +244,30 @@ export class RelayDurableObject { request: Request, role: ConnectionRole, serverId: string, - clientId: string + connectionId: string ): Response { - // Clients must provide a clientId so the daemon can create an independent - // E2EE channel per client connection. - if (role === "client" && !clientId) { - return new Response("Missing clientId parameter", { status: 400 }); - } - const upgradeError = this.requireWebSocketUpgrade(request); if (upgradeError) return upgradeError; - const isServerControl = role === "server" && !clientId; - const isServerData = role === "server" && !!clientId; + // If a client didn't provide a connectionId, the relay assigns one for routing. + const resolvedConnectionId = + role === "client" && !connectionId + ? `conn_${crypto.randomUUID().replace(/-/g, "").slice(0, 16)}` + : connectionId; + + const isServerControl = role === "server" && !resolvedConnectionId; + const isServerData = role === "server" && !!resolvedConnectionId; // Close any existing server-side connection with the same identity. // - server-control: single per serverId - // - server-data: single per clientId - // - client: many sockets per clientId are allowed + // - server-data: single per connectionId + // - client: many sockets per connectionId are allowed if (isServerControl) { for (const ws of this.state.getWebSockets("server-control")) { ws.close(1008, "Replaced by new connection"); } } else if (isServerData) { - for (const ws of this.state.getWebSockets(`server:${clientId}`)) { + for (const ws of this.state.getWebSockets(`server:${resolvedConnectionId}`)) { ws.close(1008, "Replaced by new connection"); } } @@ -276,11 +276,11 @@ export class RelayDurableObject { const tags: string[] = []; if (role === "client") { - tags.push("client", `client:${clientId}`); + tags.push("client", `client:${resolvedConnectionId}`); } else if (isServerControl) { tags.push("server-control"); } else { - tags.push("server", `server:${clientId}`); + tags.push("server", `server:${resolvedConnectionId}`); } this.state.acceptWebSocket(server, tags); @@ -289,31 +289,31 @@ export class RelayDurableObject { serverId, role, version: CURRENT_RELAY_VERSION, - clientId: clientId || null, + connectionId: resolvedConnectionId || null, createdAt: Date.now(), }; (server as WebSocketWithAttachment).serializeAttachment(attachment); console.log( - `[Relay DO] v2:${role}${isServerControl ? "(control)" : ""}${isServerData ? `(data:${clientId})` : role === "client" ? `(${clientId})` : ""} connected to session ${serverId}` + `[Relay DO] v2:${role}${isServerControl ? "(control)" : ""}${isServerData ? `(data:${resolvedConnectionId})` : role === "client" ? `(${resolvedConnectionId})` : ""} connected to session ${serverId}` ); if (role === "client") { - this.notifyControls({ type: "client_connected", clientId }); - this.nudgeOrResetControlForClient(clientId); + this.notifyControls({ type: "connected", connectionId: resolvedConnectionId }); + this.nudgeOrResetControlForConnection(resolvedConnectionId); } if (isServerControl) { - // Send current client list so the daemon can attach existing clients. + // Send current connection list so the daemon can attach existing connections. try { - server.send(JSON.stringify({ type: "sync", clientIds: this.listConnectedClientIds() })); + server.send(JSON.stringify({ type: "sync", connectionIds: this.listConnectedConnectionIds() })); } catch { // ignore } } - if (isServerData && clientId) { - this.flushClientFrames(clientId, server); + if (isServerData && resolvedConnectionId) { + this.flushFrames(resolvedConnectionId, server); } return this.asSwitchingProtocolsResponse(client); @@ -323,8 +323,8 @@ export class RelayDurableObject { const url = new URL(request.url); const role = url.searchParams.get("role") as ConnectionRole | null; const serverId = url.searchParams.get("serverId"); - const clientIdRaw = url.searchParams.get("clientId"); - const clientId = typeof clientIdRaw === "string" ? clientIdRaw.trim() : ""; + const connectionIdRaw = url.searchParams.get("connectionId"); + const connectionId = typeof connectionIdRaw === "string" ? connectionIdRaw.trim() : ""; const version = resolveRelayVersion(url.searchParams.get("v")); if (!role || (role !== "server" && role !== "client")) { @@ -343,7 +343,7 @@ export class RelayDurableObject { return this.fetchV1(request, role, serverId); } - return this.fetchV2(request, role, serverId, clientId); + return this.fetchV2(request, role, serverId, connectionId); } /** @@ -371,8 +371,8 @@ export class RelayDurableObject { return; } - const { role, clientId } = attachment; - if (!clientId) { + const { role, connectionId } = attachment; + if (!connectionId) { // Control channel: support simple app-level keepalive. if (typeof message === "string") { try { @@ -392,28 +392,28 @@ export class RelayDurableObject { } if (role === "client") { - const servers = this.state.getWebSockets(`server:${clientId}`); + const servers = this.state.getWebSockets(`server:${connectionId}`); if (servers.length === 0) { - this.bufferClientFrame(clientId, message); + this.bufferFrame(connectionId, message); return; } for (const target of servers) { try { target.send(message); } catch (error) { - console.error(`[Relay DO] Failed to forward client->server(${clientId}):`, error); + console.error(`[Relay DO] Failed to forward client->server(${connectionId}):`, error); } } return; } // server data socket -> client - const targets = this.state.getWebSockets(`client:${clientId}`); + const targets = this.state.getWebSockets(`client:${connectionId}`); for (const target of targets) { try { target.send(message); } catch (error) { - console.error(`[Relay DO] Failed to forward server->client(${clientId}):`, error); + console.error(`[Relay DO] Failed to forward server->client(${connectionId}):`, error); } } } @@ -432,37 +432,37 @@ export class RelayDurableObject { const version = attachment.version ?? LEGACY_RELAY_VERSION; console.log( - `[Relay DO] v${version}:${attachment.role}${attachment.clientId ? `(${attachment.clientId})` : ""} disconnected from session ${attachment.serverId} (${code}: ${reason})` + `[Relay DO] v${version}:${attachment.role}${attachment.connectionId ? `(${attachment.connectionId})` : ""} disconnected from session ${attachment.serverId} (${code}: ${reason})` ); if (version === LEGACY_RELAY_VERSION) { return; } - if (attachment.role === "client" && attachment.clientId) { + if (attachment.role === "client" && attachment.connectionId) { const remainingClientSockets = this.state - .getWebSockets(`client:${attachment.clientId}`) + .getWebSockets(`client:${attachment.connectionId}`) .some((socket) => socket !== ws); if (remainingClientSockets) { return; } - this.pendingClientFrames.delete(attachment.clientId); + this.pendingFrames.delete(attachment.connectionId); // Last socket for this session closed: now clean up matching server-data socket. - for (const serverWs of this.state.getWebSockets(`server:${attachment.clientId}`)) { + for (const serverWs of this.state.getWebSockets(`server:${attachment.connectionId}`)) { try { serverWs.close(1001, "Client disconnected"); } catch { // ignore } } - this.notifyControls({ type: "client_disconnected", clientId: attachment.clientId }); + this.notifyControls({ type: "disconnected", connectionId: attachment.connectionId }); return; } - if (attachment.role === "server" && attachment.clientId) { + if (attachment.role === "server" && attachment.connectionId) { // Force the client to reconnect and re-handshake when the daemon side drops. - for (const clientWs of this.state.getWebSockets(`client:${attachment.clientId}`)) { + for (const clientWs of this.state.getWebSockets(`client:${attachment.connectionId}`)) { try { clientWs.close(1012, "Server disconnected"); } catch { diff --git a/packages/relay/src/e2e.test.ts b/packages/relay/src/e2e.test.ts index cd3292b9f..f7d514759 100644 --- a/packages/relay/src/e2e.test.ts +++ b/packages/relay/src/e2e.test.ts @@ -133,7 +133,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis it("full flow: daemon and client exchange encrypted messages through relay", { timeout: 90_000 }, async () => { const serverId = "test-session-" + Date.now(); - const clientId = "clt_test_" + Date.now() + "_" + Math.random().toString(36).slice(2); + const connectionId = "clt_test_" + Date.now() + "_" + Math.random().toString(36).slice(2); // === DAEMON SIDE === // Generate keypair (public key goes in QR) @@ -167,7 +167,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis const waitForClientSeen = new Promise((resolve, reject) => { const timeout = setTimeout( - () => reject(new Error("timed out waiting for client_connected")), + () => reject(new Error("timed out waiting for connected")), 5000 ); const onMessage = (raw: unknown) => { @@ -179,13 +179,13 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis ? (raw as any).toString() : ""; const msg = JSON.parse(text); - if (msg?.type === "client_connected" && msg.clientId === clientId) { + if (msg?.type === "connected" && msg.connectionId === connectionId) { clearTimeout(timeout); daemonControlWs.off("message", onMessage); resolve(); return; } - if (msg?.type === "sync" && Array.isArray(msg.clientIds) && msg.clientIds.includes(clientId)) { + if (msg?.type === "sync" && Array.isArray(msg.connectionIds) && msg.connectionIds.includes(connectionId)) { clearTimeout(timeout); daemonControlWs.off("message", onMessage); resolve(); @@ -197,9 +197,9 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis daemonControlWs.on("message", onMessage); }); - // Client connects to relay as "client" role (must include clientId) + // Client connects to relay as "client" role (must include connectionId) const clientWs = new WebSocket( - `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}&v=2` + `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&connectionId=${connectionId}&v=2` ); await new Promise((resolve, reject) => { @@ -210,7 +210,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis await waitForClientSeen; const daemonWs = new WebSocket( - `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}&v=2` + `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&connectionId=${connectionId}&v=2` ); await new Promise((resolve, reject) => { daemonWs.on("open", resolve); @@ -300,7 +300,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis it("relay only sees opaque bytes after handshake", { timeout: 90_000 }, async () => { const serverId = "opaque-test-" + Date.now(); - const clientId = "clt_opaque_" + Date.now() + "_" + Math.random().toString(36).slice(2); + const connectionId = "clt_opaque_" + Date.now() + "_" + Math.random().toString(36).slice(2); // Setup keys const daemonKeyPair = await generateKeyPair(); @@ -328,7 +328,7 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis const waitForClientSeen = new Promise((resolve, reject) => { const timeout = setTimeout( - () => reject(new Error("timed out waiting for client_connected")), + () => reject(new Error("timed out waiting for connected")), 5000 ); const onMessage = (raw: unknown) => { @@ -340,13 +340,13 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis ? (raw as any).toString() : ""; const msg = JSON.parse(text); - if (msg?.type === "client_connected" && msg.clientId === clientId) { + if (msg?.type === "connected" && msg.connectionId === connectionId) { clearTimeout(timeout); daemonControlWs.off("message", onMessage); resolve(); return; } - if (msg?.type === "sync" && Array.isArray(msg.clientIds) && msg.clientIds.includes(clientId)) { + if (msg?.type === "sync" && Array.isArray(msg.connectionIds) && msg.connectionIds.includes(connectionId)) { clearTimeout(timeout); daemonControlWs.off("message", onMessage); resolve(); @@ -359,13 +359,13 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis }); const clientWs = new WebSocket( - `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&clientId=${clientId}&v=2` + `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=client&connectionId=${connectionId}&v=2` ); await new Promise((r) => clientWs.on("open", r)); await waitForClientSeen; const daemonWs = new WebSocket( - `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&clientId=${clientId}&v=2` + `ws://127.0.0.1:${relayPort}/ws?serverId=${serverId}&role=server&connectionId=${connectionId}&v=2` ); await new Promise((r) => daemonWs.on("open", r)); diff --git a/packages/relay/src/live-relay.e2e.test.ts b/packages/relay/src/live-relay.e2e.test.ts index 2b5e073c1..c45d3a59f 100644 --- a/packages/relay/src/live-relay.e2e.test.ts +++ b/packages/relay/src/live-relay.e2e.test.ts @@ -36,14 +36,14 @@ describe("Live relay (relay.paseo.sh) E2E", () => { await withRetry( async () => { const serverId = `live-${Date.now()}-${Math.random().toString(16).slice(2)}`; - const clientId = `clt_live_${Date.now()}_${Math.random().toString(16).slice(2)}`; + const connectionId = `clt_live_${Date.now()}_${Math.random().toString(16).slice(2)}`; const serverControlUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent(serverId)}&role=server&v=2`; const serverDataUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent( serverId - )}&role=server&clientId=${encodeURIComponent(clientId)}&v=2`; + )}&role=server&connectionId=${encodeURIComponent(connectionId)}&v=2`; const clientUrl = `${RELAY_BASE_URL}/ws?serverId=${encodeURIComponent( serverId - )}&role=client&clientId=${encodeURIComponent(clientId)}&v=2`; + )}&role=client&connectionId=${encodeURIComponent(connectionId)}&v=2`; // === Key setup === const daemonKeyPair = await generateKeyPair(); @@ -87,13 +87,13 @@ describe("Live relay (relay.paseo.sh) E2E", () => { await new Promise((resolve, reject) => { const timeout = setTimeout( - () => reject(new Error("Timed out waiting for client_connected")), + () => reject(new Error("Timed out waiting for connected")), 10_000 ); daemonControlWs.on("message", (raw) => { try { const msg = JSON.parse(raw.toString()); - if (msg && msg.type === "client_connected" && msg.clientId === clientId) { + if (msg && msg.type === "connected" && msg.connectionId === connectionId) { clearTimeout(timeout); resolve(); } diff --git a/packages/relay/src/types.ts b/packages/relay/src/types.ts index 46e61525c..409435260 100644 --- a/packages/relay/src/types.ts +++ b/packages/relay/src/types.ts @@ -20,9 +20,9 @@ export interface RelaySessionAttachment { */ version?: "1" | "2"; /** - * Unique id for the client connection. Allows the daemon to create an - * independent socket + E2EE channel per connected client. + * Unique id for the connection. Allows the daemon to create an + * independent socket + E2EE channel per connected connection. */ - clientId?: string | null; + connectionId?: string | null; createdAt: number; } diff --git a/packages/server/src/client/daemon-client.test.ts b/packages/server/src/client/daemon-client.test.ts index e7b4b41ec..4d357564a 100644 --- a/packages/server/src/client/daemon-client.test.ts +++ b/packages/server/src/client/daemon-client.test.ts @@ -83,7 +83,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -171,7 +172,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -197,7 +199,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, connectTimeoutMs: 100, @@ -234,7 +237,8 @@ describe('DaemonClient', () => { let transportIndex = 0 const client = new DaemonClient({ - url: 'ws://relay.test/ws?role=client&serverId=srv_test&clientId=clsk_test&v=2', + url: 'ws://relay.test/ws?role=client&serverId=srv_test&v=2', + clientId: 'clsk_test', logger, reconnect: { enabled: true, @@ -267,42 +271,24 @@ describe('DaemonClient', () => { } }) - test('requires explicit relay session identity when URL clientId is missing', () => { + test('requires non-empty clientId', () => { expect(() => { new DaemonClient({ url: 'ws://relay.test/ws?role=client&serverId=srv_test&v=2', + clientId: '', reconnect: { enabled: false }, }) - }).toThrow('Relay client requires clientSessionKey or URL clientId') + }).toThrow('Daemon client requires a non-empty clientId') }) - test('rejects relay identity mismatch between URL clientId and clientSessionKey', () => { - expect(() => { - new DaemonClient({ - url: 'ws://relay.test/ws?role=client&serverId=srv_test&clientId=clsk_a&v=2', - clientSessionKey: 'clsk_b', - reconnect: { enabled: false }, - }) - }).toThrow('Relay clientId and clientSessionKey must match when both are provided') - }) - - test('requires explicit direct session identity when URL clientSessionKey is missing', () => { + test('requires non-empty clientId for direct connections', () => { expect(() => { new DaemonClient({ url: 'ws://127.0.0.1:6767/ws', + clientId: ' ', reconnect: { enabled: false }, }) - }).toThrow('Direct client requires clientSessionKey or URL clientSessionKey') - }) - - test('rejects direct identity mismatch between URL and config clientSessionKey', () => { - expect(() => { - new DaemonClient({ - url: 'ws://127.0.0.1:6767/ws?clientSessionKey=clsk_a', - clientSessionKey: 'clsk_b', - reconnect: { enabled: false }, - }) - }).toThrow('Direct URL clientSessionKey and config clientSessionKey must match when both are provided') + }).toThrow('Daemon client requires a non-empty clientId') }) test('logs configured runtime generation in connection transition events', async () => { @@ -310,7 +296,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, runtimeGeneration: 7, @@ -336,7 +323,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -399,7 +387,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -467,7 +456,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -526,7 +516,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -597,7 +588,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -659,7 +651,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -705,7 +698,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -786,7 +780,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -829,7 +824,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -881,7 +877,8 @@ describe('DaemonClient', () => { }) const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory, @@ -908,7 +905,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -962,7 +960,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1037,7 +1036,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1093,7 +1093,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1139,7 +1140,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1181,7 +1183,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1223,7 +1226,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1285,7 +1289,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1347,7 +1352,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1418,7 +1424,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1470,7 +1477,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1562,7 +1570,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1632,7 +1641,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, @@ -1668,7 +1678,8 @@ describe('DaemonClient', () => { const mock = createMockTransport() const client = new DaemonClient({ - url: 'ws://test?clientSessionKey=clsk_unit_test', + url: 'ws://test', + clientId: 'clsk_unit_test', logger, reconnect: { enabled: false }, transportFactory: () => mock.transport, diff --git a/packages/server/src/client/daemon-client.ts b/packages/server/src/client/daemon-client.ts index 2bd995354..efe9b6fe5 100644 --- a/packages/server/src/client/daemon-client.ts +++ b/packages/server/src/client/daemon-client.ts @@ -6,6 +6,7 @@ import { AgentResumedStatusPayloadSchema, RestartRequestedStatusPayloadSchema, SessionInboundMessageSchema, + type WSWelcomeMessage, WSOutboundMessageSchema, } from '../shared/messages.js' import type { @@ -144,7 +145,8 @@ export type DaemonEventHandler = (event: DaemonEvent) => void export type DaemonClientConfig = { url: string - clientSessionKey?: string + clientId: string + clientType?: 'mobile' | 'browser' | 'cli' | 'mcp' runtimeGeneration?: number | null authHeader?: string suppressSendErrors?: boolean @@ -312,7 +314,7 @@ function isWaiterTimeoutError(error: unknown): boolean { return error instanceof Error && error.message.startsWith('Timeout waiting for message') } -function normalizeClientSessionKey(value: unknown): string | null { +function normalizeClientId(value: unknown): string | null { if (typeof value !== 'string') { return null } @@ -390,8 +392,9 @@ export class DaemonClient { private terminalStreams: TerminalStreamManager private readonly logConnectionPath: 'direct' | 'relay' private readonly logServerId: string | null - private readonly logClientSessionKeyHash: string + private readonly logClientIdHash: string private readonly logGeneration: number | null + private lastWelcomeMessage: WSWelcomeMessage | null = null constructor(private config: DaemonClientConfig) { this.logger = config.logger ?? consoleLogger @@ -402,11 +405,16 @@ export class DaemonClient { } catch { parsedUrlForLog = null } - const parsedServerIdForLog = normalizeClientSessionKey( + const parsedServerIdForLog = normalizeClientId( parsedUrlForLog?.searchParams.get('serverId') ) this.logServerId = parsedServerIdForLog ?? parsedUrlForLog?.host ?? null - this.logClientSessionKeyHash = 'h_unresolved' + const resolvedClientId = normalizeClientId(this.config.clientId) + if (!resolvedClientId) { + throw new Error('Daemon client requires a non-empty clientId') + } + this.config.clientId = resolvedClientId + this.logClientIdHash = hashForLog(resolvedClientId) this.logGeneration = typeof this.config.runtimeGeneration === 'number' && Number.isFinite(this.config.runtimeGeneration) ? this.config.runtimeGeneration @@ -422,48 +430,6 @@ export class DaemonClient { }) }, }) - const configClientSessionKey = normalizeClientSessionKey(this.config.clientSessionKey) - let parsed: URL | null = null - try { - parsed = new URL(this.config.url) - } catch { - // ignore - invalid URL will be handled on connect - } - - // Relay requires a clientId so the daemon can create an independent - // socket + E2EE channel per connected client. - if (isRelayClientWebSocketUrl(this.config.url)) { - const urlClientId = normalizeClientSessionKey(parsed?.searchParams.get('clientId')) - if (urlClientId && configClientSessionKey && configClientSessionKey !== urlClientId) { - throw new Error('Relay clientId and clientSessionKey must match when both are provided') - } - const resolvedClientSessionKey = configClientSessionKey ?? urlClientId - if (!resolvedClientSessionKey) { - throw new Error('Relay client requires clientSessionKey or URL clientId') - } - this.config.clientSessionKey = resolvedClientSessionKey - this.logClientSessionKeyHash = hashForLog(resolvedClientSessionKey) - if (parsed && !urlClientId) { - parsed.searchParams.set('clientId', resolvedClientSessionKey) - this.config.url = parsed.toString() - } - return - } - - const urlClientSessionKey = normalizeClientSessionKey(parsed?.searchParams.get('clientSessionKey')) - if (urlClientSessionKey && configClientSessionKey && configClientSessionKey !== urlClientSessionKey) { - throw new Error('Direct URL clientSessionKey and config clientSessionKey must match when both are provided') - } - const resolvedClientSessionKey = configClientSessionKey ?? urlClientSessionKey - if (!resolvedClientSessionKey) { - throw new Error('Direct client requires clientSessionKey or URL clientSessionKey') - } - this.config.clientSessionKey = resolvedClientSessionKey - this.logClientSessionKeyHash = hashForLog(resolvedClientSessionKey) - if (parsed && !urlClientSessionKey) { - parsed.searchParams.set('clientSessionKey', resolvedClientSessionKey) - this.config.url = parsed.toString() - } } // ============================================================================ @@ -532,8 +498,10 @@ export class DaemonClient { logger: this.logger, }) } - const transport = transportFactory({ url: this.config.url, headers }) + const transportUrl = this.resolveTransportUrlForAttempt() + const transport = transportFactory({ url: transportUrl, headers }) this.transport = transport + this.lastWelcomeMessage = null this.updateConnectionState({ status: 'connecting', @@ -556,18 +524,12 @@ export class DaemonClient { this.transportCleanup = [ transport.onOpen(() => { - this.resetConnectTimeout() if (this.pendingGenericTransportErrorTimeout) { clearTimeout(this.pendingGenericTransportErrorTimeout) this.pendingGenericTransportErrorTimeout = null } this.lastErrorValue = null - this.reconnectAttempt = 0 - this.updateConnectionState({ status: 'connected' }, { event: 'TRANSPORT_OPEN' }) - this.resubscribeCheckoutDiffSubscriptions() - this.resubscribeTerminalDirectorySubscriptions() - this.flushPendingSendQueue() - this.resolveConnect() + this.sendHelloMessage() }), transport.onClose((event) => { this.resetConnectTimeout() @@ -674,6 +636,7 @@ export class DaemonClient { this.clearWaiters(new Error('Daemon client closed')) this.rejectPendingSendQueue(new Error('Daemon client closed')) this.terminalStreams.clearAll() + this.lastWelcomeMessage = null this.updateConnectionState( { status: 'disposed' }, { event: 'DISPOSE', reason: 'Client closed', reasonCode: 'disposed' } @@ -2599,6 +2562,45 @@ export class DaemonClient { return requestId ?? crypto.randomUUID() } + getLastWelcomeMessage(): WSWelcomeMessage | null { + return this.lastWelcomeMessage + } + + private resolveTransportUrlForAttempt(): string { + return this.config.url + } + + private sendHelloMessage(): void { + if (!this.transport) { + this.scheduleReconnect({ + reason: 'Transport unavailable before hello', + event: 'HELLO_TRANSPORT_MISSING', + reasonCode: 'transport_error', + }) + return + } + + try { + this.transport.send( + JSON.stringify({ + type: 'hello', + clientId: this.config.clientId, + clientType: this.config.clientType ?? 'cli', + protocolVersion: 1, + }) + ) + } catch (error) { + const message = + error instanceof Error ? error.message : 'Failed to send hello message' + this.lastErrorValue = message + this.scheduleReconnect({ + reason: message, + event: 'HELLO_SEND_FAILED', + reasonCode: 'transport_error', + }) + } + } + private disposeTransport(code = 1001, reason = 'Reconnecting'): void { this.cleanupTransport() if (this.transport) { @@ -2660,7 +2662,7 @@ export class DaemonClient { const parsed = WSOutboundMessageSchema.safeParse(parsedJson) if (!parsed.success) { - const msgType = (parsedJson as { message?: { type?: string } })?.message?.type ?? 'unknown' + const msgType = (parsedJson as { type?: string })?.type ?? 'unknown' this.logger.warn({ msgType, error: parsed.error.message }, 'Message validation failed') return } @@ -2669,6 +2671,18 @@ export class DaemonClient { return } + if (parsed.data.type === 'welcome') { + this.lastWelcomeMessage = parsed.data + this.resetConnectTimeout() + this.reconnectAttempt = 0 + this.updateConnectionState({ status: 'connected' }, { event: 'HELLO_WELCOME' }) + this.resubscribeCheckoutDiffSubscriptions() + this.resubscribeTerminalDirectorySubscriptions() + this.flushPendingSendQueue() + this.resolveConnect() + return + } + this.handleSessionMessage(parsed.data.message) } @@ -2704,7 +2718,7 @@ export class DaemonClient { this.logger.info( { serverId: this.logServerId, - clientSessionKeyHash: this.logClientSessionKeyHash, + clientIdHash: this.logClientIdHash, from: previous.status, to: next.status, event: metadata?.event ?? 'STATE_UPDATE', @@ -2745,6 +2759,7 @@ export class DaemonClient { this.clearWaiters(new Error(reason ?? 'Connection lost')) this.rejectPendingSendQueue(new Error(reason ?? 'Connection lost')) this.terminalStreams.clearAll() + this.lastWelcomeMessage = null if (wasDisposed) { this.rejectConnect(new Error(reason ?? 'Daemon client is disposed')) diff --git a/packages/server/src/server/agent/activity-curator.test.ts b/packages/server/src/server/agent/activity-curator.test.ts index 62d871a8e..416b0f81a 100644 --- a/packages/server/src/server/agent/activity-curator.test.ts +++ b/packages/server/src/server/agent/activity-curator.test.ts @@ -128,26 +128,51 @@ describe("curateAgentActivity", () => { callId: "task-1", name: "Task", status: "running", - input: { description: "Investigate" }, + detail: { + type: "sub_agent", + subAgentType: "Explore", + description: "Investigate repository", + log: "[Read] README.md", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + ], + }, }), toolCallItem({ callId: "task-1", name: "Task", status: "running", - metadata: { subAgentActivity: "Read" }, - }), - toolCallItem({ - callId: "task-1", - name: "Task", - status: "running", - metadata: { subAgentActivity: "Edit" }, + detail: { + type: "sub_agent", + subAgentType: "Explore", + description: "Investigate repository", + log: "[Read] README.md\n[Bash] ls", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + { + index: 2, + toolName: "Bash", + summary: "ls", + }, + ], + }, }), ]; const result = curateAgentActivity(timeline); const lines = result.split("\n"); - expect(lines.filter((line) => line.startsWith("[Task]"))).toEqual(["[Task] Edit"]); + expect(lines.filter((line) => line.startsWith("[Explore]"))).toEqual([ + "[Explore] Investigate repository", + ]); }); it("renders todo/error/compaction entries", () => { diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 6ef4c77fa..d5d01d4de 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -16,6 +16,63 @@ import type { AgentStreamEvent, } from "./agent-sdk-types.js"; +type Deferred = { + promise: Promise; + resolve: (value: T) => void; + reject: (reason?: unknown) => void; +}; + +function deferred(): Deferred { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +} + +class EventPushable implements AsyncIterable { + private queue: T[] = []; + private resolvers: Array<(value: IteratorResult) => void> = []; + private closed = false; + + push(value: T): void { + if (this.closed) { + return; + } + const resolver = this.resolvers.shift(); + if (resolver) { + resolver({ value, done: false }); + return; + } + this.queue.push(value); + } + + end(): void { + this.closed = true; + while (this.resolvers.length > 0) { + const resolver = this.resolvers.shift(); + resolver?.({ value: undefined, done: true }); + } + } + + [Symbol.asyncIterator](): AsyncIterator { + return { + next: () => { + if (this.queue.length > 0) { + const value = this.queue.shift()!; + return Promise.resolve({ value, done: false }); + } + if (this.closed) { + return Promise.resolve({ value: undefined, done: true }); + } + return new Promise((resolve) => this.resolvers.push(resolve)); + }, + }; + } +} + const TEST_CAPABILITIES = { supportsStreaming: false, supportsSessionPersistence: false, @@ -792,6 +849,207 @@ describe("AgentManager", () => { expect(refreshed?.runtimeInfo?.model).toBe("gpt-5.2-codex"); }); + test("waitForAgentEvent does not resolve idle until pendingRun is cleared", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-wait-coherence-")); + const storagePath = join(workdir, "agents"); + const storage = new AgentStorage(storagePath, logger); + const releaseTurnCompleted = deferred(); + const releaseStreamEnd = deferred(); + + class SlowTerminalSession extends TestAgentSession { + override async *stream(): AsyncGenerator { + yield { type: "turn_started", provider: this.provider }; + await releaseTurnCompleted.promise; + yield { type: "turn_completed", provider: this.provider }; + await releaseStreamEnd.promise; + } + } + + class SlowTerminalClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new SlowTerminalSession(config); + } + } + + const manager = new AgentManager({ + clients: { + codex: new SlowTerminalClient(), + }, + registry: storage, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000124", + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd: workdir, + }); + + const turnCompletedSeen = new Promise((resolve) => { + const unsubscribe = manager.subscribe( + (event) => { + if ( + event.type === "agent_stream" && + event.agentId === snapshot.id && + event.event.type === "turn_completed" + ) { + unsubscribe(); + resolve(); + } + }, + { agentId: snapshot.id, replayState: false } + ); + }); + + const stream = manager.streamAgent(snapshot.id, "hello"); + const consumePromise = (async () => { + for await (const _event of stream) { + // Drain events so manager lifecycle progresses naturally. + } + })(); + + await manager.waitForAgentRunStart(snapshot.id); + const waitPromise = manager.waitForAgentEvent(snapshot.id); + + releaseTurnCompleted.resolve(); + await turnCompletedSeen; + const earlyResolution = await Promise.race([ + waitPromise.then(() => "resolved"), + new Promise<"pending">((resolve) => setTimeout(() => resolve("pending"), 50)), + ]); + expect(earlyResolution).toBe("pending"); + + releaseStreamEnd.resolve(); + const waited = await waitPromise; + expect(waited.status).toBe("idle"); + + await consumePromise; + }); + + test("applies live autonomous events while no foreground run is active", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-live-events-")); + const storagePath = join(workdir, "agents"); + const storage = new AgentStorage(storagePath, logger); + const liveEvents = new EventPushable(); + + class LiveEventSession extends TestAgentSession { + async *streamLiveEvents(): AsyncGenerator { + for await (const event of liveEvents) { + yield event; + } + } + + override async close(): Promise { + liveEvents.end(); + } + } + + class LiveEventClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new LiveEventSession(config); + } + } + + const manager = new AgentManager({ + clients: { + codex: new LiveEventClient(), + }, + registry: storage, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000125", + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd: workdir, + }); + + const lifecycleUpdates: string[] = []; + let sawRunningState = false; + let resolveSettled!: () => void; + const settled = new Promise((resolve) => { + resolveSettled = resolve; + }); + manager.subscribe( + (event) => { + if (event.type === "agent_state" && event.agent.id === snapshot.id) { + lifecycleUpdates.push(event.agent.lifecycle); + if (event.agent.lifecycle === "running") { + sawRunningState = true; + } + if (sawRunningState && event.agent.lifecycle === "idle") { + resolveSettled(); + } + } + }, + { agentId: snapshot.id, replayState: false } + ); + + liveEvents.push({ type: "turn_started", provider: "codex" }); + liveEvents.push({ + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "AUTONOMOUS_PUMP_MESSAGE" }, + }); + liveEvents.push({ type: "turn_completed", provider: "codex" }); + await settled; + + const updated = manager.getAgent(snapshot.id); + expect(updated?.lifecycle).toBe("idle"); + expect(manager.getTimeline(snapshot.id)).toContainEqual({ + type: "assistant_message", + text: "AUTONOMOUS_PUMP_MESSAGE", + }); + expect(lifecycleUpdates).toContain("running"); + expect(lifecycleUpdates).toContain("idle"); + }); + + test("waitForAgentEvent waitForActive resolves for autonomous live-event run", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-live-wait-")); + const storagePath = join(workdir, "agents"); + const storage = new AgentStorage(storagePath, logger); + const liveEvents = new EventPushable(); + + class LiveEventSession extends TestAgentSession { + async *streamLiveEvents(): AsyncGenerator { + for await (const event of liveEvents) { + yield event; + } + } + + override async close(): Promise { + liveEvents.end(); + } + } + + class LiveEventClient extends TestAgentClient { + override async createSession(config: AgentSessionConfig): Promise { + return new LiveEventSession(config); + } + } + + const manager = new AgentManager({ + clients: { + codex: new LiveEventClient(), + }, + registry: storage, + logger, + idFactory: () => "00000000-0000-4000-8000-000000000126", + }); + + const snapshot = await manager.createAgent({ + provider: "codex", + cwd: workdir, + }); + + const waitPromise = manager.waitForAgentEvent(snapshot.id, { waitForActive: true }); + liveEvents.push({ type: "turn_started", provider: "codex" }); + liveEvents.push({ type: "turn_completed", provider: "codex" }); + + const result = await waitPromise; + expect(result.status).toBe("idle"); + }); + test("keeps updatedAt monotonic when user message and run start happen in the same millisecond", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); const storagePath = join(workdir, "agents"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index f230e3984..aea05ce2e 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -238,6 +238,10 @@ type SubscriptionRecord = { agentId: string | null; }; +type LiveEventStreamingSession = AgentSession & { + streamLiveEvents: () => AsyncGenerator; +}; + const DEFAULT_MAX_TIMELINE_ITEMS = 2000; const DEFAULT_TIMELINE_FETCH_LIMIT = 200; const BUSY_STATUSES: AgentLifecycleStatus[] = [ @@ -272,6 +276,16 @@ function validateAgentId(agentId: string, source: string): string { return result.data; } +function supportsLiveEventStream( + session: AgentSession +): session is LiveEventStreamingSession { + return ( + "streamLiveEvents" in session && + typeof (session as { streamLiveEvents?: unknown }).streamLiveEvents === + "function" + ); +} + export class AgentManager { private readonly clients = new Map(); private readonly agents = new Map(); @@ -281,6 +295,7 @@ export class AgentManager { private readonly registry?: AgentStorage; private readonly previousStatuses = new Map(); private readonly backgroundTasks = new Set>(); + private readonly liveEventPumps = new Map>(); private onAgentAttention?: AgentAttentionCallback; private logger: Logger; @@ -727,6 +742,7 @@ export class AgentManager { // Remove the existing agent entry before swapping sessions this.agents.delete(agentId); + this.liveEventPumps.delete(agentId); try { await existing.session.close(); } catch (error) { @@ -756,6 +772,7 @@ export class AgentManager { async closeAgent(agentId: string): Promise { const agent = this.requireAgent(agentId); this.agents.delete(agentId); + this.liveEventPumps.delete(agentId); // Clean up previousStatus to prevent memory leak this.previousStatuses.delete(agentId); const session = agent.session; @@ -1275,6 +1292,7 @@ export class AgentManager { let currentStatus: AgentLifecycleStatus = initialStatus; let hasStarted = initialBusy || hasPendingRun; + let terminalStatusOverride: AgentLifecycleStatus | null = null; // Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them let unsubscribe: (() => void) | null = null; @@ -1337,6 +1355,9 @@ export class AgentManager { return; } if (!waitForActive || hasStarted) { + if (terminalStatusOverride) { + currentStatus = terminalStatusOverride; + } finish(null); } return; @@ -1348,20 +1369,15 @@ export class AgentManager { return; } if (event.event.type === "turn_failed") { - currentStatus = "error"; hasStarted = true; - finish(null); + terminalStatusOverride = "error"; return; } if (event.event.type === "turn_completed") { - currentStatus = "idle"; hasStarted = true; - finish(null); } if (event.event.type === "turn_canceled") { - currentStatus = "idle"; hasStarted = true; - finish(null); } } }, @@ -1461,6 +1477,7 @@ export class AgentManager { managed.lifecycle = "idle"; await this.persistSnapshot(managed); this.emitState(managed); + this.startLiveEventPump(managed); return { ...managed }; } @@ -1619,9 +1636,16 @@ export class AgentManager { case "turn_completed": agent.lastUsage = event.usage; agent.lastError = undefined; + if (!agent.pendingRun && agent.lifecycle !== "idle") { + (agent as ActiveManagedAgent).lifecycle = "idle"; + this.emitState(agent); + } void this.refreshRuntimeInfo(agent); break; case "turn_failed": + if (!agent.pendingRun) { + agent.lifecycle = "error"; + } agent.lastError = event.error; for (const [requestId] of agent.pendingPermissions) { agent.pendingPermissions.delete(requestId); @@ -1637,6 +1661,9 @@ export class AgentManager { this.emitState(agent); break; case "turn_canceled": + if (!agent.pendingRun) { + (agent as ActiveManagedAgent).lifecycle = "idle"; + } agent.lastError = undefined; for (const [requestId] of agent.pendingPermissions) { agent.pendingPermissions.delete(requestId); @@ -1651,6 +1678,12 @@ export class AgentManager { } this.emitState(agent); break; + case "turn_started": + if (!agent.pendingRun) { + (agent as ActiveManagedAgent).lifecycle = "running"; + this.emitState(agent); + } + break; case "permission_requested": agent.pendingPermissions.set(event.request.id, event.request); this.emitState(agent); @@ -1888,4 +1921,42 @@ export class AgentManager { } return agent; } + + private startLiveEventPump(agent: ActiveManagedAgent): void { + if (!supportsLiveEventStream(agent.session)) { + return; + } + const liveSession = agent.session; + if (this.liveEventPumps.has(agent.id)) { + return; + } + const pump = (async () => { + try { + for await (const event of liveSession.streamLiveEvents()) { + const current = this.agents.get(agent.id); + if (!current) { + break; + } + // Foreground streamAgent owns live prompt events while running. + // Background live-event tailing is for autonomous activity while idle. + if (current.pendingRun) { + continue; + } + this.handleStreamEvent(current, event); + } + } catch (error) { + this.logger.warn( + { err: error, agentId: agent.id }, + "Live event pump failed" + ); + } + })(); + this.liveEventPumps.set(agent.id, pump); + pump.finally(() => { + const current = this.liveEventPumps.get(agent.id); + if (current === pump) { + this.liveEventPumps.delete(agent.id); + } + }); + } } diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index af075ecd7..eb6882a77 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -149,6 +149,17 @@ export type ToolCallDetail = }>; truncated?: boolean; } + | { + type: "sub_agent"; + subAgentType?: string; + description?: string; + log: string; + actions: Array<{ + index: number; + toolName: string; + summary?: string; + }>; + } | { type: "unknown"; input: unknown | null; diff --git a/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts b/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts index 757aeed1c..424c19aa5 100644 --- a/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts +++ b/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts @@ -1,4 +1,7 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { appendFileSync, mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; import { createTestLogger } from "../../../test-utils/test-logger.js"; import { ClaudeAgentClient } from "./claude-agent.js"; @@ -50,6 +53,23 @@ function buildUsage() { }; } +function createPromptUuidReader(prompt: AsyncIterable) { + const iterator = prompt[Symbol.asyncIterator](); + let cached: Promise | null = null; + return async () => { + if (!cached) { + cached = iterator.next().then((next) => { + if (next.done) { + return null; + } + const value = next.value as { uuid?: unknown } | undefined; + return typeof value?.uuid === "string" ? value.uuid : null; + }); + } + return cached; + }; +} + function buildFirstQueryMock( allowOldAssistant: Promise ): QueryMock { @@ -108,7 +128,8 @@ function buildFirstQueryMock( }; } -function buildSecondQueryMock(): QueryMock { +function buildSecondQueryMock(prompt: AsyncIterable): QueryMock { + const readPromptUuid = createPromptUuidReader(prompt); let step = 0; return { next: vi.fn(async () => { @@ -126,6 +147,21 @@ function buildSecondQueryMock(): QueryMock { }; } if (step === 1) { + step += 1; + const promptUuid = (await readPromptUuid()) ?? "missing-prompt-uuid"; + return { + done: false, + value: { + type: "user", + message: { role: "user", content: "second prompt" }, + parent_tool_use_id: null, + uuid: promptUuid, + session_id: "interrupt-regression-session", + isReplay: true, + }, + }; + } + if (step === 2) { step += 1; return { done: false, @@ -137,7 +173,7 @@ function buildSecondQueryMock(): QueryMock { }, }; } - if (step === 2) { + if (step === 3) { step += 1; return { done: false, @@ -188,22 +224,49 @@ function collectAssistantText(events: AgentStreamEvent[]): string { .join(""); } +function createTimedIteratorReader(params: { iterator: AsyncIterator }) { + const { iterator } = params; + let pendingNext: Promise> | null = null; + + return { + async nextWithTimeout(timeoutMs: number): Promise> { + if (!pendingNext) { + pendingNext = iterator.next(); + } + const timeout = new Promise((resolve) => { + setTimeout(() => resolve(null), timeoutMs); + }); + const outcome = await Promise.race([ + pendingNext.then((result) => ({ kind: "result" as const, result })), + timeout.then(() => ({ kind: "timeout" as const })), + ]); + if (outcome.kind === "timeout") { + throw new Error("Timed out waiting for live event"); + } + pendingNext = null; + return outcome.result; + }, + }; +} + describe("ClaudeAgentSession interrupt restart regression", () => { beforeEach(() => { const allowOldAssistant = deferred(); let queryCreateCount = 0; - sdkMocks.query.mockImplementation(() => { - queryCreateCount += 1; - if (queryCreateCount === 1) { - const mock = buildFirstQueryMock(allowOldAssistant.promise); - sdkMocks.firstQuery = mock; + sdkMocks.query.mockImplementation( + ({ prompt }: { prompt: AsyncIterable }) => { + queryCreateCount += 1; + if (queryCreateCount === 1) { + const mock = buildFirstQueryMock(allowOldAssistant.promise); + sdkMocks.firstQuery = mock; + return mock; + } + const mock = buildSecondQueryMock(prompt); + sdkMocks.secondQuery = mock; return mock; } - const mock = buildSecondQueryMock(); - sdkMocks.secondQuery = mock; - return mock; - }); + ); sdkMocks.releaseOldAssistant = () => allowOldAssistant.resolve(); }); @@ -239,4 +302,382 @@ describe("ClaudeAgentSession interrupt restart regression", () => { await firstTurn.return?.(); await session.close(); }); + + test("ignores stale task-notification assistant/result events queued before the current prompt", async () => { + const logger = createTestLogger(); + + sdkMocks.query.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + const readPromptUuid = createPromptUuidReader(prompt); + let step = 0; + return { + next: vi.fn(async () => { + if (step === 0) { + step += 1; + return { + done: false, + value: { + type: "system", + subtype: "init", + session_id: "task-notification-session", + permissionMode: "default", + model: "opus", + }, + }; + } + if (step === 1) { + step += 1; + return { + done: false, + value: { + type: "system", + subtype: "task_notification", + task_id: "task-123", + status: "completed", + output_file: "/tmp/task-123.txt", + summary: "Codex agent is done", + session_id: "task-notification-session", + uuid: "task-note-1", + }, + }; + } + if (step === 2) { + step += 1; + return { + done: false, + value: { + type: "assistant", + message: { + content: "STALE_TASK_NOTIFICATION_RESPONSE", + }, + }, + }; + } + if (step === 3) { + step += 1; + return { + done: false, + value: { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + }; + } + if (step === 4) { + step += 1; + const promptUuid = (await readPromptUuid()) ?? "missing-prompt-uuid"; + return { + done: false, + value: { + type: "user", + message: { role: "user", content: "current prompt" }, + parent_tool_use_id: null, + uuid: promptUuid, + session_id: "task-notification-session", + isReplay: true, + }, + }; + } + if (step === 5) { + step += 1; + return { + done: false, + value: { + type: "assistant", + message: { + content: "CURRENT_PROMPT_RESPONSE", + }, + }, + }; + } + if (step === 6) { + step += 1; + return { + done: false, + value: { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + }; + } + return { done: true, value: undefined }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => undefined), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + } satisfies QueryMock; + }); + + const client = new ClaudeAgentClient({ logger }); + const session = await client.createSession({ + provider: "claude", + cwd: process.cwd(), + }); + + const events = await collectUntilTerminal(session.stream("current prompt")); + const assistantText = collectAssistantText(events); + + expect(assistantText).toContain("CURRENT_PROMPT_RESPONSE"); + expect(assistantText).not.toContain("STALE_TASK_NOTIFICATION_RESPONSE"); + + await session.close(); + }); + + test("does not terminate the current prompt on a stale pre-prompt result event", async () => { + const logger = createTestLogger(); + + sdkMocks.query.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + const readPromptUuid = createPromptUuidReader(prompt); + let step = 0; + return { + next: vi.fn(async () => { + if (step === 0) { + step += 1; + return { + done: false, + value: { + type: "system", + subtype: "init", + session_id: "stale-result-session", + permissionMode: "default", + model: "opus", + }, + }; + } + if (step === 1) { + step += 1; + return { + done: false, + value: { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + }; + } + if (step === 2) { + step += 1; + const promptUuid = (await readPromptUuid()) ?? "missing-prompt-uuid"; + return { + done: false, + value: { + type: "user", + message: { role: "user", content: "current prompt" }, + parent_tool_use_id: null, + uuid: promptUuid, + session_id: "stale-result-session", + isReplay: true, + }, + }; + } + if (step === 3) { + step += 1; + return { + done: false, + value: { + type: "assistant", + message: { + content: "FRESH_AFTER_STALE_RESULT", + }, + }, + }; + } + if (step === 4) { + step += 1; + return { + done: false, + value: { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + }; + } + return { done: true, value: undefined }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => undefined), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + } satisfies QueryMock; + }); + + const client = new ClaudeAgentClient({ logger }); + const session = await client.createSession({ + provider: "claude", + cwd: process.cwd(), + }); + + const events = await collectUntilTerminal(session.stream("current prompt")); + const assistantText = collectAssistantText(events); + + expect(assistantText).toContain("FRESH_AFTER_STALE_RESULT"); + + await session.close(); + }); + + test("emits autonomous live events from history tail when Claude wakes itself", async () => { + const logger = createTestLogger(); + const projectDir = mkdtempSync(path.join(tmpdir(), "claude-live-tail-project-")); + const configDir = mkdtempSync(path.join(tmpdir(), "claude-live-tail-config-")); + const previousClaudeConfigDir = process.env.CLAUDE_CONFIG_DIR; + process.env.CLAUDE_CONFIG_DIR = configDir; + + sdkMocks.query.mockImplementation(() => { + let step = 0; + return { + next: vi.fn(async () => { + if (step === 0) { + step += 1; + return { + done: false, + value: { + type: "system", + subtype: "init", + session_id: "live-autonomous-session", + permissionMode: "default", + model: "opus", + }, + }; + } + if (step === 1) { + step += 1; + return { + done: false, + value: { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + }; + } + return { done: true, value: undefined }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => undefined), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + } satisfies QueryMock; + }); + + try { + const client = new ClaudeAgentClient({ logger }); + const session = await client.createSession({ + provider: "claude", + cwd: projectDir, + }); + + await collectUntilTerminal(session.stream("seed prompt")); + expect(session.describePersistence()?.sessionId).toBe( + "live-autonomous-session" + ); + + const sanitizedCwd = projectDir.replace(/[\\/\.]/g, "-").replace(/_/g, "-"); + const historyDir = path.join(configDir, "projects", sanitizedCwd); + mkdirSync(historyDir, { recursive: true }); + const historyPath = path.join(historyDir, "live-autonomous-session.jsonl"); + writeFileSync(historyPath, "", "utf8"); + + const liveIterator = ( + session as unknown as { + streamLiveEvents: () => AsyncGenerator; + } + ).streamLiveEvents(); + const timedReader = createTimedIteratorReader({ iterator: liveIterator }); + const appendAutonomousBurst = (suffix: string) => { + appendFileSync( + historyPath, + `${JSON.stringify({ + type: "system", + subtype: "task_notification", + task_id: `task-${suffix}`, + status: "completed", + output_file: "/tmp/out.txt", + summary: "Background task finished", + uuid: `task-note-uuid-${suffix}`, + session_id: "live-autonomous-session", + })}\n` + ); + appendFileSync( + historyPath, + `${JSON.stringify({ + type: "assistant", + message: { content: "AUTONOMOUS_WAKE_RESPONSE" }, + uuid: `assistant-uuid-${suffix}`, + session_id: "live-autonomous-session", + })}\n` + ); + appendFileSync( + historyPath, + `${JSON.stringify({ + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + uuid: `result-uuid-${suffix}`, + session_id: "live-autonomous-session", + })}\n` + ); + }; + + const liveEvents: AgentStreamEvent[] = []; + let completed = false; + for (let attempt = 0; attempt < 4 && !completed; attempt += 1) { + appendAutonomousBurst(String(attempt)); + while (liveEvents.length < 12) { + let next: IteratorResult; + try { + next = await timedReader.nextWithTimeout(2_000); + } catch { + break; + } + if (next.done) { + break; + } + liveEvents.push(next.value); + if (next.value.type === "turn_completed") { + completed = true; + break; + } + } + } + + expect(liveEvents.some((event) => event.type === "turn_started")).toBe(true); + expect( + liveEvents.some( + (event) => + event.type === "timeline" && + event.item.type === "assistant_message" && + event.item.text.includes("AUTONOMOUS_WAKE_RESPONSE") + ) + ).toBe(true); + expect(liveEvents.some((event) => event.type === "turn_completed")).toBe(true); + + await session.close(); + } finally { + if (previousClaudeConfigDir === undefined) { + delete process.env.CLAUDE_CONFIG_DIR; + } else { + process.env.CLAUDE_CONFIG_DIR = previousClaudeConfigDir; + } + rmSync(projectDir, { recursive: true, force: true }); + rmSync(configDir, { recursive: true, force: true }); + } + }); }); diff --git a/packages/server/src/server/agent/providers/claude-agent.sub-agent-sidechain.test.ts b/packages/server/src/server/agent/providers/claude-agent.sub-agent-sidechain.test.ts new file mode 100644 index 000000000..6f4eb8164 --- /dev/null +++ b/packages/server/src/server/agent/providers/claude-agent.sub-agent-sidechain.test.ts @@ -0,0 +1,351 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { createTestLogger } from "../../../test-utils/test-logger.js"; +import type { AgentStreamEvent } from "../agent-sdk-types.js"; +import type { AgentTimelineRow } from "../agent-manager.js"; +import { projectTimelineRows } from "../timeline-projection.js"; +import { ClaudeAgentClient } from "./claude-agent.js"; + +const sdkMocks = vi.hoisted(() => ({ + query: vi.fn(), +})); + +vi.mock("@anthropic-ai/claude-agent-sdk", () => ({ + query: sdkMocks.query, +})); + +type QueryMock = { + next: ReturnType; + interrupt: ReturnType; + return: ReturnType; + setPermissionMode: ReturnType; + setModel: ReturnType; + supportedModels: ReturnType; + supportedCommands: ReturnType; + rewindFiles: ReturnType; +}; + +function buildQueryMock(events: unknown[]): QueryMock { + let index = 0; + return { + next: vi.fn(async () => { + if (index >= events.length) { + return { done: true, value: undefined }; + } + const value = events[index]; + index += 1; + return { done: false, value }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => undefined), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + }; +} + +async function collectUntilTerminal( + stream: AsyncGenerator +): Promise { + const events: AgentStreamEvent[] = []; + for await (const event of stream) { + events.push(event); + if ( + event.type === "turn_completed" || + event.type === "turn_failed" || + event.type === "turn_canceled" + ) { + break; + } + } + return events; +} + +function buildTailScenarioEvents(actionCount: number): unknown[] { + const actionEvents = Array.from({ length: actionCount }, (_, index) => { + const actionNumber = index + 1; + return { + type: "stream_event", + parent_tool_use_id: "task-tail-1", + event: { + type: "content_block_start", + index: actionNumber, + content_block: { + type: "tool_use", + id: `sub-read-${actionNumber}`, + name: "Read", + input: { + file_path: `file-${actionNumber}.md`, + }, + }, + }, + }; + }); + + return [ + { + type: "system", + subtype: "init", + session_id: "sidechain-tail-session", + permissionMode: "default", + model: "opus", + }, + { + type: "stream_event", + parent_tool_use_id: null, + event: { + type: "content_block_start", + index: 0, + content_block: { + type: "tool_use", + id: "task-tail-1", + name: "Task", + input: { + subagent_type: "Explore", + description: "Tail latest sub-agent activity", + }, + }, + }, + }, + ...actionEvents, + { + type: "assistant", + parent_tool_use_id: null, + message: { + content: [ + { + type: "tool_result", + tool_use_id: "task-tail-1", + tool_name: "Task", + content: "done", + is_error: false, + }, + ], + }, + }, + { + type: "result", + subtype: "success", + usage: { + input_tokens: 1, + cache_read_input_tokens: 0, + output_tokens: 1, + }, + total_cost_usd: 0, + }, + ]; +} + +describe("ClaudeAgentSession sub-agent sidechain updates", () => { + const logger = createTestLogger(); + + beforeEach(() => { + const largeOldText = "VERY_LARGE_OLD_STRING".repeat(50); + sdkMocks.query.mockImplementation(() => + buildQueryMock([ + { + type: "system", + subtype: "init", + session_id: "sidechain-session", + permissionMode: "default", + model: "opus", + }, + { + type: "stream_event", + parent_tool_use_id: null, + event: { + type: "content_block_start", + index: 0, + content_block: { + type: "tool_use", + id: "task-call-1", + name: "Task", + input: { + subagent_type: "Explore", + description: "Inspect repository structure", + }, + }, + }, + }, + { + type: "stream_event", + parent_tool_use_id: "task-call-1", + event: { + type: "content_block_start", + index: 1, + content_block: { + type: "tool_use", + id: "sub-read-1", + name: "Read", + input: { + file_path: "README.md", + }, + }, + }, + }, + { + type: "stream_event", + parent_tool_use_id: "task-call-1", + event: { + type: "content_block_start", + index: 2, + content_block: { + type: "tool_use", + id: "sub-edit-1", + name: "Edit", + input: { + file_path: "src/index.ts", + old_string: largeOldText, + new_string: "replacement", + }, + }, + }, + }, + { + type: "tool_progress", + tool_use_id: "sub-edit-1", + tool_name: "Edit", + parent_tool_use_id: "task-call-1", + elapsed_time_seconds: 1, + }, + { + type: "assistant", + parent_tool_use_id: null, + message: { + content: [ + { + type: "tool_result", + tool_use_id: "task-call-1", + tool_name: "Task", + content: "done", + is_error: false, + }, + ], + }, + }, + { + type: "result", + subtype: "success", + usage: { + input_tokens: 1, + cache_read_input_tokens: 0, + output_tokens: 1, + }, + total_cost_usd: 0, + }, + ]) + ); + }); + + afterEach(() => { + sdkMocks.query.mockReset(); + }); + + test("accumulates lightweight sub_agent detail and preserves callId lifecycle collapse", async () => { + const session = await new ClaudeAgentClient({ logger }).createSession({ + provider: "claude", + cwd: process.cwd(), + }); + + const events = await collectUntilTerminal(session.stream("delegate work")); + await session.close(); + + const timelineToolCalls = events + .filter( + (event): event is Extract => + event.type === "timeline" && event.item.type === "tool_call" + ) + .map((event) => event.item) + .filter((item) => item.callId === "task-call-1"); + + expect(timelineToolCalls.length).toBeGreaterThanOrEqual(2); + + const subAgentUpdates = timelineToolCalls.filter( + (item) => item.detail.type === "sub_agent" + ); + expect(subAgentUpdates.length).toBeGreaterThanOrEqual(1); + + const latest = subAgentUpdates[subAgentUpdates.length - 1]; + expect(latest).toBeDefined(); + if (!latest || latest.detail.type !== "sub_agent") { + throw new Error("expected sub_agent detail"); + } + + expect(latest.detail.subAgentType).toBe("Explore"); + expect(latest.detail.description).toBe("Inspect repository structure"); + expect(latest.detail.actions).toEqual([ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + { + index: 2, + toolName: "Edit", + summary: "src/index.ts", + }, + ]); + expect(latest.detail.log).toContain("[Read] README.md"); + expect(latest.detail.log).toContain("[Edit] src/index.ts"); + expect(latest.detail.log).not.toContain("VERY_LARGE_OLD_STRING"); + + const rows: AgentTimelineRow[] = timelineToolCalls.map((item, index) => ({ + seq: index + 1, + timestamp: `2026-02-01T00:00:0${index}.000Z`, + item, + })); + const projected = projectTimelineRows(rows, "claude", "projected"); + const projectedTaskCalls = projected.filter( + (entry) => entry.item.type === "tool_call" && entry.item.callId === "task-call-1" + ); + + expect(projectedTaskCalls).toHaveLength(1); + }); + + test("tails sub-agent actions instead of dropping latest entries at cap", async () => { + sdkMocks.query.mockImplementation(() => buildQueryMock(buildTailScenarioEvents(205))); + + const session = await new ClaudeAgentClient({ logger }).createSession({ + provider: "claude", + cwd: process.cwd(), + }); + + const events = await collectUntilTerminal(session.stream("delegate work")); + await session.close(); + + const timelineToolCalls = events + .filter( + (event): event is Extract => + event.type === "timeline" && event.item.type === "tool_call" + ) + .map((event) => event.item) + .filter((item) => item.callId === "task-tail-1"); + const subAgentUpdates = timelineToolCalls.filter( + (item) => item.detail.type === "sub_agent" + ); + const latest = subAgentUpdates[subAgentUpdates.length - 1]; + expect(latest).toBeDefined(); + if (!latest || latest.detail.type !== "sub_agent") { + throw new Error("expected sub_agent detail"); + } + + expect(latest.detail.actions).toHaveLength(200); + expect(latest.detail.actions[0]).toEqual({ + index: 6, + toolName: "Read", + summary: "file-6.md", + }); + expect(latest.detail.actions[199]).toEqual({ + index: 205, + toolName: "Read", + summary: "file-205.md", + }); + + expect(latest.detail.log).not.toContain("[Read] file-1.md"); + expect(latest.detail.log).not.toContain("[Read] file-5.md"); + expect(latest.detail.log).toContain("[Read] file-6.md"); + expect(latest.detail.log).toContain("[Read] file-205.md"); + }); +}); diff --git a/packages/server/src/server/agent/providers/claude-agent.test.ts b/packages/server/src/server/agent/providers/claude-agent.test.ts index 51d4d343c..c53298cb2 100644 --- a/packages/server/src/server/agent/providers/claude-agent.test.ts +++ b/packages/server/src/server/agent/providers/claude-agent.test.ts @@ -1264,7 +1264,7 @@ async function startAgentMcpServer(): Promise { ); test( - "collapses sub-agent tool calls into Task metadata updates", + "collapses sub-agent tool calls into Task sub_agent detail updates", async () => { const cwd = tmpCwd(); const client = new ClaudeAgentClient({ logger }); @@ -1297,19 +1297,13 @@ async function startAgentMcpServer(): Promise { const taskCalls = toolCalls.filter((item) => item.name === "Task"); expect(taskCalls.length).toBeGreaterThanOrEqual(1); - // Sub-agent tool calls (Read, Bash, shell, etc.) should NOT appear as - // separate timeline items — they should only appear as Task metadata updates - const subAgentLeaks = toolCalls.filter( - (item) => item.name !== "Task" && item.metadata?.subAgentActivity === undefined - ); - // If there are non-Task tool calls, they must be from the main agent, not sub-agent - // We can't 100% guarantee Claude won't also use tools directly, but Task metadata + // We can't 100% guarantee Claude won't also use tools directly, but Task detail // updates should exist for the sub-agent activity - const taskWithMetadata = taskCalls.filter( - (item) => item.metadata?.subAgentActivity + const taskWithSubAgentDetail = taskCalls.filter( + (item) => item.detail.type === "sub_agent" ); if (taskCalls.length > 0) { - expect(taskWithMetadata.length).toBeGreaterThanOrEqual(1); + expect(taskWithSubAgentDetail.length).toBeGreaterThanOrEqual(1); } // Verify the curator produces clean output with collapsed Task entries diff --git a/packages/server/src/server/agent/providers/claude-agent.ts b/packages/server/src/server/agent/providers/claude-agent.ts index cfb0afea9..c5e192814 100644 --- a/packages/server/src/server/agent/providers/claude-agent.ts +++ b/packages/server/src/server/agent/providers/claude-agent.ts @@ -29,6 +29,7 @@ import { mapClaudeFailedToolCall, mapClaudeRunningToolCall, } from "./claude/tool-call-mapper.js"; +import { buildToolCallDisplayModel } from "../../../shared/tool-call-display.js"; import type { AgentCapabilityFlags, @@ -341,12 +342,43 @@ type ToolUseCacheEntry = { input?: AgentMetadata | null; }; +type SubAgentActionEntry = { + index: number; + toolName: string; + summary?: string; +}; + +type SubAgentActivityState = { + subAgentType?: string; + description?: string; + actions: SubAgentActionEntry[]; + actionKeys: string[]; + nextActionIndex: number; + actionIndexByKey: Map; +}; + +type SubAgentActionCandidate = { + key: string; + toolName: string; + input: unknown; +}; + const DEFAULT_PERMISSION_TIMEOUT_MS = 120_000; +const MAX_SUB_AGENT_LOG_ENTRIES = 200; +const MAX_SUB_AGENT_SUMMARY_CHARS = 160; function isMetadata(value: unknown): value is AgentMetadata { return typeof value === "object" && value !== null; } +function readTrimmedString(value: unknown): string | undefined { + if (typeof value !== "string") { + return undefined; + } + const trimmed = value.trim(); + return trimmed.length > 0 ? trimmed : undefined; +} + function isMcpServerConfig(value: unknown): value is McpServerConfig { if (!isMetadata(value)) { return false; @@ -684,10 +716,12 @@ class ClaudeAgentSession implements AgentSession { private cachedRuntimeInfo: AgentRuntimeInfo | null = null; private lastOptionsModel: string | null = null; private selectableModelIds: Set | null = null; - private activeSidechains = new Map(); + private activeSidechains = new Map(); private compacting = false; private queryRestartNeeded = false; private userMessageIds: string[] = []; + private sentUserMessageIds = new Set(); + private closed = false; constructor( config: ClaudeAgentConfig, @@ -887,6 +921,88 @@ class ClaudeAgentSession implements AgentSession { } } + async *streamLiveEvents(): AsyncGenerator { + let currentHistoryPath: string | null = null; + let lineCursor = 0; + let suppressLocalTurn = false; + let autonomousTurnActive = false; + + while (!this.closed) { + const sessionId = this.claudeSessionId; + const historyPath = sessionId ? this.resolveHistoryPath(sessionId) : null; + if (!historyPath) { + await this.waitForLiveHistoryPoll(); + continue; + } + + let lines: string[]; + try { + const content = fs.readFileSync(historyPath, "utf8"); + lines = content.split(/\r?\n/).filter((line) => line.trim().length > 0); + } catch { + await this.waitForLiveHistoryPoll(); + continue; + } + + if (currentHistoryPath !== historyPath) { + currentHistoryPath = historyPath; + lineCursor = lines.length; + suppressLocalTurn = false; + autonomousTurnActive = false; + await this.waitForLiveHistoryPoll(); + continue; + } + + if (lineCursor > lines.length) { + lineCursor = 0; + } + if (lineCursor === lines.length) { + await this.waitForLiveHistoryPoll(); + continue; + } + + const nextLines = lines.slice(lineCursor); + lineCursor = lines.length; + for (const rawLine of nextLines) { + let entry: any; + try { + entry = JSON.parse(rawLine); + } catch { + continue; + } + if (entry?.isSidechain) { + continue; + } + + if ( + entry?.type === "user" && + typeof entry.uuid === "string" && + this.sentUserMessageIds.has(entry.uuid) + ) { + this.sentUserMessageIds.delete(entry.uuid); + suppressLocalTurn = true; + continue; + } + + if (suppressLocalTurn) { + if (entry?.type === "result") { + suppressLocalTurn = false; + } + continue; + } + + const converted = this.convertLiveHistoryEntryToEvents({ + entry, + autonomousTurnActive, + }); + autonomousTurnActive = converted.autonomousTurnActive; + for (const event of converted.events) { + yield event; + } + } + } + } + async getAvailableModes(): Promise { return this.availableModes; } @@ -1018,6 +1134,7 @@ class ClaudeAgentSession implements AgentSession { } async close(): Promise { + this.closed = true; this.rejectAllPendingPermissions(new Error("Claude session closed")); this.input?.end(); await this.query?.interrupt?.(); @@ -1468,6 +1585,7 @@ class ClaudeAgentSession implements AgentSession { const messageId = randomUUID(); this.rememberUserMessageId(messageId); + this.sentUserMessageIds.add(messageId); return { type: "user", @@ -1504,6 +1622,11 @@ class ClaudeAgentSession implements AgentSession { } this.input.push(sdkMessage); + const expectedPromptUuid = + typeof sdkMessage.uuid === "string" && sdkMessage.uuid.length > 0 + ? sdkMessage.uuid + : null; + let promptActivated = expectedPromptUuid === null; while (true) { // Check if this turn has been superseded by a new one. @@ -1524,6 +1647,35 @@ class ClaudeAgentSession implements AgentSession { break; } + // Guard against stale queued events from a prior run (for example, + // task-notification side activity): do not treat assistant/result as part + // of this prompt until we observe a prompt-owned stream start. + if (!promptActivated) { + const isTopLevelAssistantStart = + value.type === "stream_event" && + (value as { event?: { type?: unknown }; parent_tool_use_id?: unknown }).event + ?.type === "message_start" && + ((value as { parent_tool_use_id?: unknown }).parent_tool_use_id ?? null) === + null; + if ( + value.type === "user" && + typeof value.uuid === "string" && + value.uuid === expectedPromptUuid + ) { + promptActivated = true; + continue; + } + if (isTopLevelAssistantStart) { + promptActivated = true; + } else if (value.type === "system" && value.subtype === "init") { + yield value; + continue; + } + if (!promptActivated) { + continue; + } + } + yield value; if (value.type === "result") { break; @@ -1617,63 +1769,268 @@ class ClaudeAgentSession implements AgentSession { message: SDKMessage, parentToolUseId: string ): AgentStreamEvent[] { - let toolName: string | undefined; + const state = + this.activeSidechains.get(parentToolUseId) ?? + ({ + actions: [], + actionKeys: [], + nextActionIndex: 1, + actionIndexByKey: new Map(), + } satisfies SubAgentActivityState); + this.activeSidechains.set(parentToolUseId, state); - if (message.type === "assistant") { - const content = message.message?.content; - if (Array.isArray(content)) { - for (const block of content) { - if (isClaudeContentChunk(block) && - (block.type === "tool_use" || block.type === "mcp_tool_use" || block.type === "server_tool_use") && - typeof block.name === "string" - ) { - toolName = block.name; - break; - } - } + const contextUpdated = this.updateSubAgentContextFromTaskInput( + state, + parentToolUseId + ); + const actionCandidates = this.extractSubAgentActionCandidates(message); + let actionUpdated = false; + for (const action of actionCandidates) { + if (this.appendSubAgentAction(state, action)) { + actionUpdated = true; } - } else if (message.type === "stream_event") { - const event = message.event; - if (event.type === "content_block_start") { - const cb = isClaudeContentChunk(event.content_block) ? event.content_block : null; - if (cb?.type === "tool_use" && typeof cb.name === "string") { - toolName = cb.name; - } - } - } else if (message.type === "tool_progress") { - toolName = message.tool_name; } - if (!toolName) { + if (!contextUpdated && !actionUpdated) { return []; } - const prev = this.activeSidechains.get(parentToolUseId); - if (prev === toolName) { - return []; - } - this.activeSidechains.set(parentToolUseId, toolName); - const toolCall = mapClaudeRunningToolCall({ name: "Task", callId: parentToolUseId, input: null, output: null, - metadata: { subAgentActivity: toolName }, }); if (!toolCall) { return []; } + const detail: Extract["detail"] = { + type: "sub_agent", + ...(state.subAgentType ? { subAgentType: state.subAgentType } : {}), + ...(state.description ? { description: state.description } : {}), + log: state.actions + .map((action) => + action.summary + ? `[${action.toolName}] ${action.summary}` + : `[${action.toolName}]` + ) + .join("\n"), + actions: state.actions.map((action) => ({ + index: action.index, + toolName: action.toolName, + ...(action.summary ? { summary: action.summary } : {}), + })), + }; + return [ { type: "timeline", - item: toolCall, + item: { + ...toolCall, + detail, + }, provider: "claude", }, ]; } + private updateSubAgentContextFromTaskInput( + state: SubAgentActivityState, + parentToolUseId: string + ): boolean { + const taskInput = this.toolUseCache.get(parentToolUseId)?.input; + const nextSubAgentType = this.normalizeSubAgentText(taskInput?.subagent_type); + const nextDescription = this.normalizeSubAgentText(taskInput?.description); + + let changed = false; + if (nextSubAgentType && nextSubAgentType !== state.subAgentType) { + state.subAgentType = nextSubAgentType; + changed = true; + } + if (nextDescription && nextDescription !== state.description) { + state.description = nextDescription; + changed = true; + } + return changed; + } + + private normalizeSubAgentText(value: unknown): string | undefined { + const normalized = readTrimmedString(value)?.replace(/\s+/g, " "); + if (!normalized) { + return undefined; + } + if (normalized.length <= MAX_SUB_AGENT_SUMMARY_CHARS) { + return normalized; + } + return `${normalized.slice(0, MAX_SUB_AGENT_SUMMARY_CHARS)}...`; + } + + private extractSubAgentActionCandidates( + message: SDKMessage + ): SubAgentActionCandidate[] { + if (message.type === "assistant") { + const content = message.message?.content; + if (!Array.isArray(content)) { + return []; + } + const actions: SubAgentActionCandidate[] = []; + for (const block of content) { + if ( + !isClaudeContentChunk(block) || + !( + block.type === "tool_use" || + block.type === "mcp_tool_use" || + block.type === "server_tool_use" + ) || + typeof block.name !== "string" + ) { + continue; + } + const key = + readTrimmedString(block.id) ?? + `assistant:${block.name}:${actions.length}`; + actions.push({ + key, + toolName: block.name, + input: block.input ?? null, + }); + } + return actions; + } + + if (message.type === "stream_event") { + const event = message.event; + if (event.type !== "content_block_start") { + return []; + } + const block = isClaudeContentChunk(event.content_block) + ? event.content_block + : null; + if ( + !block || + !( + block.type === "tool_use" || + block.type === "mcp_tool_use" || + block.type === "server_tool_use" + ) || + typeof block.name !== "string" + ) { + return []; + } + const key = + readTrimmedString(block.id) ?? + `stream:${block.name}:${typeof event.index === "number" ? event.index : 0}`; + return [ + { + key, + toolName: block.name, + input: block.input ?? null, + }, + ]; + } + + if (message.type === "tool_progress") { + const toolName = readTrimmedString(message.tool_name); + if (!toolName) { + return []; + } + const key = + readTrimmedString(message.tool_use_id) ?? `progress:${toolName}`; + return [{ key, toolName, input: null }]; + } + + return []; + } + + private appendSubAgentAction( + state: SubAgentActivityState, + candidate: SubAgentActionCandidate + ): boolean { + const normalizedToolName = readTrimmedString(candidate.toolName); + if (!normalizedToolName) { + return false; + } + + const summary = this.deriveSubAgentActionSummary( + normalizedToolName, + candidate.input + ); + const existingIndex = state.actionIndexByKey.get(candidate.key); + + if (existingIndex !== undefined) { + const existing = state.actions[existingIndex]; + if (!existing) { + return false; + } + const nextSummary = existing.summary ?? summary; + const unchanged = + existing.toolName === normalizedToolName && + existing.summary === nextSummary; + if (unchanged) { + return false; + } + state.actions[existingIndex] = { + ...existing, + toolName: normalizedToolName, + ...(nextSummary ? { summary: nextSummary } : {}), + }; + return true; + } + + const nextEntry: SubAgentActionEntry = { + index: state.nextActionIndex, + toolName: normalizedToolName, + ...(summary ? { summary } : {}), + }; + state.nextActionIndex += 1; + state.actions.push(nextEntry); + state.actionKeys.push(candidate.key); + this.trimSubAgentTail(state); + this.rebuildSubAgentActionIndex(state); + return true; + } + + private trimSubAgentTail(state: SubAgentActivityState): void { + while (state.actions.length > MAX_SUB_AGENT_LOG_ENTRIES) { + state.actions.shift(); + state.actionKeys.shift(); + } + } + + private rebuildSubAgentActionIndex(state: SubAgentActivityState): void { + state.actionIndexByKey.clear(); + for (let index = 0; index < state.actionKeys.length; index += 1) { + const key = state.actionKeys[index]; + if (key) { + state.actionIndexByKey.set(key, index); + } + } + } + + private deriveSubAgentActionSummary( + toolName: string, + input: unknown + ): string | undefined { + const runningToolCall = mapClaudeRunningToolCall({ + name: toolName, + callId: `sub-agent-summary-${toolName}`, + input, + output: null, + }); + if (!runningToolCall) { + return undefined; + } + const display = buildToolCallDisplayModel({ + name: runningToolCall.name, + status: runningToolCall.status, + error: runningToolCall.error, + detail: runningToolCall.detail, + metadata: runningToolCall.metadata, + }); + return this.normalizeSubAgentText(display.summary); + } + private translateMessageToEvents(message: SDKMessage, turnContext: TurnContext): AgentStreamEvent[] { const parentToolUseId = "parent_tool_use_id" in message ? (message as { parent_tool_use_id: string | null }).parent_tool_use_id @@ -1983,6 +2340,7 @@ class ClaudeAgentSession implements AgentSession { } } this.toolUseCache.clear(); + this.activeSidechains.clear(); } private pushToolCall( @@ -2023,6 +2381,60 @@ class ClaudeAgentSession implements AgentSession { } } + private waitForLiveHistoryPoll(): Promise { + return new Promise((resolve) => setTimeout(resolve, 250)); + } + + private convertLiveHistoryEntryToEvents(params: { + entry: any; + autonomousTurnActive: boolean; + }): { + events: AgentStreamEvent[]; + autonomousTurnActive: boolean; + } { + const { entry, autonomousTurnActive } = params; + const events: AgentStreamEvent[] = []; + + if (entry?.type === "result") { + if (!autonomousTurnActive) { + return { events, autonomousTurnActive: false }; + } + const usage = + entry && typeof entry === "object" && entry.usage + ? this.convertUsage(entry as SDKResultMessage) + : undefined; + if (entry?.subtype === "success") { + events.push({ type: "turn_completed", provider: "claude", usage }); + } else { + const errorMessage = + Array.isArray(entry?.errors) && entry.errors.length > 0 + ? entry.errors.join("\n") + : "Claude run failed"; + events.push({ + type: "turn_failed", + provider: "claude", + error: errorMessage, + }); + } + return { events, autonomousTurnActive: false }; + } + + let nextAutonomousActive = autonomousTurnActive; + const timelineItems = this.convertHistoryEntry(entry); + if (!nextAutonomousActive && timelineItems.length > 0) { + events.push({ type: "turn_started", provider: "claude" }); + nextAutonomousActive = true; + } + for (const item of timelineItems) { + events.push({ type: "timeline", item, provider: "claude" }); + } + + return { + events, + autonomousTurnActive: nextAutonomousActive, + }; + } + private loadPersistedHistory(sessionId: string) { try { const historyPath = this.resolveHistoryPath(sessionId); @@ -2207,6 +2619,7 @@ class ClaudeAgentSession implements AgentSession { if (typeof block.tool_use_id === "string") { this.toolUseCache.delete(block.tool_use_id); + this.activeSidechains.delete(block.tool_use_id); } } diff --git a/packages/server/src/server/daemon-client.e2e.test.ts b/packages/server/src/server/daemon-client.e2e.test.ts index b0a2ff0e0..898327d63 100644 --- a/packages/server/src/server/daemon-client.e2e.test.ts +++ b/packages/server/src/server/daemon-client.e2e.test.ts @@ -11,7 +11,6 @@ import { DaemonClient, } from "./test-utils/index.js"; import { getFullAccessConfig, getAskModeConfig } from "./daemon-e2e/agent-configs.js"; -import { parseServerInfoStatusPayload } from "./messages.js"; import { chunkPcm16, parsePcm16MonoWav, @@ -280,24 +279,17 @@ describe("daemon client E2E", () => { } }, 30000); - test("emits server_info on websocket connect", async () => { + test("receives welcome on websocket connect", async () => { const client = new DaemonClient({ url: `ws://127.0.0.1:${ctx.daemon.port}/ws`, + clientId: `cid-e2e-${randomUUID()}`, + clientType: "cli", }); - - const infoPromise = waitForSignal<{ serverId: string }>(5000, (resolve) => { - const unsubscribe = client.on("status", (message) => { - if (message.type !== "status") return; - const payload = parseServerInfoStatusPayload(message.payload); - if (!payload) return; - resolve({ serverId: payload.serverId }); - }); - return unsubscribe; - }); - await client.connect(); - const info = await infoPromise; - expect(info.serverId.length).toBeGreaterThan(0); + const welcome = client.getLastWelcomeMessage(); + expect(welcome).not.toBeNull(); + expect(welcome?.serverId.length).toBeGreaterThan(0); + expect(typeof welcome?.resumed).toBe("boolean"); await client.close(); }, 15000); @@ -315,38 +307,20 @@ describe("daemon client E2E", () => { const client = new DaemonClient({ url: `ws://127.0.0.1:${isolatedCtx.daemon.port}/ws`, + clientId: `cid-e2e-${randomUUID()}`, + clientType: "cli", }); try { - const infoPromise = waitForSignal<{ - dictationEnabled: boolean; - dictationReason: string; - voiceEnabled: boolean; - voiceReason: string; - }>(5000, (resolve) => { - const unsubscribe = client.on("status", (message) => { - if (message.type !== "status") return; - const payload = parseServerInfoStatusPayload(message.payload); - if (!payload) return; - const voice = payload.capabilities?.voice; - if (!voice) return; - resolve({ - dictationEnabled: voice.dictation.enabled, - dictationReason: voice.dictation.reason, - voiceEnabled: voice.voice.enabled, - voiceReason: voice.voice.reason, - }); - }); - return unsubscribe; - }); - await client.connect(); - const info = await infoPromise; + const welcome = client.getLastWelcomeMessage(); + const voice = welcome?.capabilities?.voice; + expect(voice).toBeTruthy(); - expect(info.dictationEnabled).toBe(false); - expect(info.dictationReason).toBe("Dictation is disabled in daemon config."); - expect(info.voiceEnabled).toBe(false); - expect(info.voiceReason).toBe("Realtime voice is disabled in daemon config."); + expect(voice?.dictation.enabled).toBe(false); + expect(voice?.dictation.reason).toBe("Dictation is disabled in daemon config."); + expect(voice?.voice.enabled).toBe(false); + expect(voice?.voice.reason).toBe("Realtime voice is disabled in daemon config."); } finally { await client.close().catch(() => undefined); await isolatedCtx.cleanup(); diff --git a/packages/server/src/server/daemon-e2e/checkout-debug.ts b/packages/server/src/server/daemon-e2e/checkout-debug.ts index b9eafaa1e..3b2b90b70 100644 --- a/packages/server/src/server/daemon-e2e/checkout-debug.ts +++ b/packages/server/src/server/daemon-e2e/checkout-debug.ts @@ -32,15 +32,15 @@ class LoggingWebSocket extends OriginalWebSocket { const PASEO_HOME = process.env.PASEO_HOME ?? `${os.homedir()}/.paseo`; const PASEO_LISTEN = process.env.PASEO_LISTEN ?? "127.0.0.1:6767"; const DAEMON_URL = `ws://${PASEO_LISTEN}/ws`; -const CLIENT_SESSION_KEY = "clsk_checkout_debug"; +const CLIENT_ID = "clsk_checkout_debug"; async function testMultiAgentSequence() { console.log("\n=== Testing multi-agent checkout sequence ==="); console.log(`Daemon URL: ${DAEMON_URL}`); const client = new DaemonClient({ - url: `${DAEMON_URL}?clientSessionKey=${CLIENT_SESSION_KEY}`, - clientSessionKey: CLIENT_SESSION_KEY, + url: DAEMON_URL, + clientId: CLIENT_ID, webSocketFactory: (url) => new LoggingWebSocket(url) as any, reconnect: { enabled: false }, }); diff --git a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts index 1ec3ec1ae..611b4abe4 100644 --- a/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/relay-transport.e2e.test.ts @@ -198,13 +198,12 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis const offerUrl = parseOfferUrlFromLogs(lines); const { serverId, daemonPublicKeyB64 } = decodeOfferFromFragmentUrl(offerUrl); - const clientId = `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`; + const stableClientId = `cid_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`; const ws = new WebSocket( buildRelayWebSocketUrl({ endpoint: `127.0.0.1:${relayPort}`, serverId, role: "client", - clientId, }) ); @@ -234,6 +233,8 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis ws.on("open", async () => { try { + let pingSent = false; + let channelRef: Awaited> | null = null; const channel = await createClientChannel( transport, daemonPublicKeyB64, @@ -242,8 +243,13 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis try { const payload = typeof data === "string" ? JSON.parse(data) : data; - // The daemon may send an initial `server_info` status message - // immediately upon connect; ignore everything until we see `pong`. + if (payload && typeof payload === "object" && (payload as any).type === "welcome") { + if (!pingSent && channelRef) { + pingSent = true; + void channelRef.send(JSON.stringify({ type: "ping" })); + } + return; + } if (payload && typeof payload === "object" && (payload as any).type === "pong") { clearTimeout(timeout); resolve(payload); @@ -260,7 +266,15 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis }, } ); - await channel.send(JSON.stringify({ type: "ping" })); + channelRef = channel; + await channel.send( + JSON.stringify({ + type: "hello", + clientId: stableClientId, + clientType: "cli", + protocolVersion: 1, + }) + ); } catch (err) { clearTimeout(timeout); reject(err); @@ -316,9 +330,9 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis endpoint: `127.0.0.1:${relayPort}`, serverId, role: "client", - clientId: `clt_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`, }) ); + const stableClientId = `cid_test_${Date.now().toString(36)}_${Math.random().toString(36).slice(2)}`; const received = await new Promise((resolve, reject) => { const timeout = setTimeout(() => { @@ -346,9 +360,18 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis ws.on("open", async () => { try { + let pingSent = false; + let channelRef: Awaited> | null = null; const channel = await createClientChannel(transport, daemonPublicKeyB64, { onmessage: (data) => { const payload = typeof data === "string" ? JSON.parse(data) : data; + if (payload && typeof payload === "object" && (payload as any).type === "welcome") { + if (!pingSent && channelRef) { + pingSent = true; + void channelRef.send(JSON.stringify({ type: "ping" })); + } + return; + } if (payload && typeof payload === "object" && (payload as any).type === "pong") { clearTimeout(timeout); resolve(payload); @@ -360,7 +383,15 @@ async function waitForRelayWebSocketReady(port: number, timeout = 60000): Promis reject(err); }, }); - await channel.send(JSON.stringify({ type: "ping" })); + channelRef = channel; + await channel.send( + JSON.stringify({ + type: "hello", + clientId: stableClientId, + clientType: "cli", + protocolVersion: 1, + }) + ); } catch (err) { clearTimeout(timeout); reject(err); diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts index d4614b8f2..345c311d6 100644 --- a/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts @@ -67,9 +67,19 @@ describe("daemon E2E (real claude) - send while running recovery", () => { const reconnected = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); try { await reconnected.connect(); - await reconnected.fetchAgents({ - subscribe: { subscriptionId: "reconnected" }, - }); + const applySnapshot = ( + snapshot: Parameters[0]["snapshot"] + ) => { + const next = applyAgentInputProcessingTransition({ + snapshot, + currentIsProcessing: isProcessing, + previousIsRunning, + latestUpdatedAt, + }); + isProcessing = next.isProcessing; + previousIsRunning = next.previousIsRunning; + latestUpdatedAt = next.latestUpdatedAt; + }; reconnected.on("agent_update", (message) => { if (message.type !== "agent_update" || message.payload.kind !== "upsert") { @@ -78,19 +88,24 @@ describe("daemon E2E (real claude) - send while running recovery", () => { if (message.payload.agent.id !== agent.id) { return; } - - const next = applyAgentInputProcessingTransition({ - snapshot: message.payload.agent, - currentIsProcessing: isProcessing, - previousIsRunning, - latestUpdatedAt, - }); - isProcessing = next.isProcessing; - previousIsRunning = next.previousIsRunning; - latestUpdatedAt = next.latestUpdatedAt; + applySnapshot(message.payload.agent); }); + const initial = await reconnected.fetchAgents({ + subscribe: { subscriptionId: "reconnected" }, + }); + const hydratedSnapshot = initial.entries.find( + (candidate) => candidate.agent.id === agent.id + )?.agent; + if (hydratedSnapshot) { + applySnapshot(hydratedSnapshot); + } + await secondary.waitForFinish(agent.id, 180_000); + const finalSnapshot = await secondary.fetchAgent(agent.id); + if (finalSnapshot) { + applySnapshot(finalSnapshot); + } // Sending while running should clear processing even if reconnect misses the // not-running -> running transition. diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.test.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.test.ts new file mode 100644 index 000000000..9f2378574 --- /dev/null +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, test } from "vitest"; + +import { applyAgentInputProcessingTransition } from "./send-while-running-stuck-test-utils.js"; + +function snapshot(status: "running" | "idle", updatedAtMs: number) { + return { + status, + updatedAt: new Date(updatedAtMs).toISOString(), + } as any; +} + +describe("applyAgentInputProcessingTransition", () => { + test("clears processing for stale non-running snapshot after reconnect", () => { + const result = applyAgentInputProcessingTransition({ + snapshot: snapshot("idle", 1_000), + currentIsProcessing: true, + previousIsRunning: true, + latestUpdatedAt: 2_000, + }); + + expect(result).toEqual({ + isProcessing: false, + previousIsRunning: false, + latestUpdatedAt: 2_000, + }); + }); + + test("keeps processing for stale running snapshot", () => { + const result = applyAgentInputProcessingTransition({ + snapshot: snapshot("running", 1_000), + currentIsProcessing: true, + previousIsRunning: true, + latestUpdatedAt: 2_000, + }); + + expect(result).toEqual({ + isProcessing: true, + previousIsRunning: true, + latestUpdatedAt: 2_000, + }); + }); + + test("clears processing when fresh update stops running", () => { + const result = applyAgentInputProcessingTransition({ + snapshot: snapshot("idle", 3_000), + currentIsProcessing: true, + previousIsRunning: true, + latestUpdatedAt: 2_000, + }); + + expect(result).toEqual({ + isProcessing: false, + previousIsRunning: false, + latestUpdatedAt: 3_000, + }); + }); +}); diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.ts index 0e148c6a4..96f4ea076 100644 --- a/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.ts +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck-test-utils.ts @@ -7,7 +7,18 @@ export function applyAgentInputProcessingTransition(input: { latestUpdatedAt: number; }): { isProcessing: boolean; previousIsRunning: boolean; latestUpdatedAt: number } { const updatedAt = new Date(input.snapshot.updatedAt).getTime(); + const isRunning = input.snapshot.status === "running"; if (updatedAt < input.latestUpdatedAt) { + // Reconnect flows can deliver an authoritative non-running snapshot + // whose timestamp predates this client's local "processing started" time. + // Clear processing to avoid getting stuck even when we miss an edge transition. + if (input.currentIsProcessing && !isRunning) { + return { + isProcessing: false, + previousIsRunning: false, + latestUpdatedAt: input.latestUpdatedAt, + }; + } return { isProcessing: input.currentIsProcessing, previousIsRunning: input.previousIsRunning, @@ -15,7 +26,6 @@ export function applyAgentInputProcessingTransition(input: { }; } - const isRunning = input.snapshot.status === "running"; const wasRunning = input.previousIsRunning; let isProcessing = input.currentIsProcessing; diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts index 6710343fd..659d879c0 100644 --- a/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts @@ -41,7 +41,7 @@ describe("daemon E2E (real codex) - send while running recovery", () => { ...getFullAccessConfig("codex"), }); - await primary.sendMessage(agent.id, "Run: sleep 30"); + await primary.sendMessage(agent.id, "Run: sleep 5"); await primary.waitForAgentUpsert( agent.id, (snapshot) => snapshot.status === "running", @@ -64,9 +64,19 @@ describe("daemon E2E (real codex) - send while running recovery", () => { const reconnected = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); try { await reconnected.connect(); - await reconnected.fetchAgents({ - subscribe: { subscriptionId: "reconnected" }, - }); + const applySnapshot = ( + snapshot: Parameters[0]["snapshot"] + ) => { + const next = applyAgentInputProcessingTransition({ + snapshot, + currentIsProcessing: isProcessing, + previousIsRunning, + latestUpdatedAt, + }); + isProcessing = next.isProcessing; + previousIsRunning = next.previousIsRunning; + latestUpdatedAt = next.latestUpdatedAt; + }; reconnected.on("agent_update", (message) => { if (message.type !== "agent_update" || message.payload.kind !== "upsert") { @@ -75,19 +85,24 @@ describe("daemon E2E (real codex) - send while running recovery", () => { if (message.payload.agent.id !== agent.id) { return; } - - const next = applyAgentInputProcessingTransition({ - snapshot: message.payload.agent, - currentIsProcessing: isProcessing, - previousIsRunning, - latestUpdatedAt, - }); - isProcessing = next.isProcessing; - previousIsRunning = next.previousIsRunning; - latestUpdatedAt = next.latestUpdatedAt; + applySnapshot(message.payload.agent); }); + const initial = await reconnected.fetchAgents({ + subscribe: { subscriptionId: "reconnected" }, + }); + const hydratedSnapshot = initial.entries.find( + (candidate) => candidate.agent.id === agent.id + )?.agent; + if (hydratedSnapshot) { + applySnapshot(hydratedSnapshot); + } + await secondary.waitForFinish(agent.id, 120_000); + const finalSnapshot = await secondary.fetchAgent(agent.id); + if (finalSnapshot) { + applySnapshot(finalSnapshot); + } // Sending while running should clear processing even if reconnect misses the // not-running -> running transition. diff --git a/packages/server/src/server/daemon-e2e/ui-action-stress.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/ui-action-stress.real.e2e.test.ts index c327de5e8..3c929334f 100644 --- a/packages/server/src/server/daemon-e2e/ui-action-stress.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/ui-action-stress.real.e2e.test.ts @@ -382,17 +382,7 @@ async function runUiScenario(params: { } if (action.type === "wait_for_finish") { - const startedAt = Date.now(); - let result = await client.waitForFinish(agentId, action.timeoutMs); - while ( - result.status === "idle" && - result.final?.status === "running" && - Date.now() - startedAt < action.timeoutMs - ) { - const elapsed = Date.now() - startedAt; - const remaining = Math.max(1_000, action.timeoutMs - elapsed); - result = await client.waitForFinish(agentId, remaining); - } + const result = await client.waitForFinish(agentId, action.timeoutMs); expect( result.status, diff --git a/packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts b/packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts index c256b63ef..eb0130b12 100644 --- a/packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/wait-for-idle.e2e.test.ts @@ -126,4 +126,64 @@ describe("waitForFinish edge cases", () => { rmSync(cwd1, { recursive: true, force: true }); rmSync(cwd2, { recursive: true, force: true }); }, 60000); + + test("waitForFinish keeps status and final snapshot coherent when a new run starts at idle edge", async () => { + const cwd = tmpCwd(); + + const agent = await ctx.client.createAgent({ + provider: "codex", + cwd, + title: "Idle Edge Coherence", + modeId: "bypassPermissions", + }); + + let sawRunning = false; + let spawnedSecondRun = false; + let secondRunDrain: Promise | null = null; + const unsubscribe = ctx.daemon.daemon.agentManager.subscribe( + (event) => { + if (event.type !== "agent_state" || event.agent.id !== agent.id) { + return; + } + if (event.agent.lifecycle === "running") { + sawRunning = true; + return; + } + if ( + event.agent.lifecycle !== "idle" || + !sawRunning || + spawnedSecondRun + ) { + return; + } + + spawnedSecondRun = true; + const stream = ctx.daemon.daemon.agentManager.streamAgent( + agent.id, + "Use your shell tool to run sleep 1 and then reply done." + ); + secondRunDrain = (async () => { + for await (const _event of stream) { + // Drain second run so manager can settle. + } + })(); + }, + { agentId: agent.id, replayState: false } + ); + + try { + await ctx.client.sendMessage(agent.id, "Reply with exactly: first"); + const state = await ctx.client.waitForFinish(agent.id, 30000); + + expect(spawnedSecondRun).toBe(true); + expect(state.status).toBe("idle"); + expect(state.final?.status).toBe("idle"); + expect(state.error).toBeNull(); + } finally { + unsubscribe(); + await secondRunDrain; + await ctx.client.deleteAgent(agent.id); + rmSync(cwd, { recursive: true, force: true }); + } + }, 60000); }); diff --git a/packages/server/src/server/relay-transport.test.ts b/packages/server/src/server/relay-transport.test.ts index 456475ced..925e7ef2d 100644 --- a/packages/server/src/server/relay-transport.test.ts +++ b/packages/server/src/server/relay-transport.test.ts @@ -203,7 +203,7 @@ describe("relay-transport control lifecycle", () => { const control = MockWebSocket.instances[0]; control.open(); control.message(JSON.stringify({ type: "pong", ts: Date.now() })); - control.message(JSON.stringify({ type: "client_connected", clientId: "clt_test" })); + control.message(JSON.stringify({ type: "connected", connectionId: "clt_test" })); const dataSocket = MockWebSocket.instances[1]; expect(dataSocket).toBeDefined(); diff --git a/packages/server/src/server/relay-transport.ts b/packages/server/src/server/relay-transport.ts index 834a47775..726d01676 100644 --- a/packages/server/src/server/relay-transport.ts +++ b/packages/server/src/server/relay-transport.ts @@ -32,9 +32,9 @@ type RelaySocketLike = { }; type ControlMessage = - | { type: "sync"; clientIds: string[] } - | { type: "client_connected"; clientId: string } - | { type: "client_disconnected"; clientId: string } + | { type: "sync"; connectionIds: string[] } + | { type: "connected"; connectionId: string } + | { type: "disconnected"; connectionId: string } | { type: "ping" } | { type: "pong" }; @@ -50,15 +50,15 @@ function tryParseControlMessage(raw: unknown): ControlMessage | null { if (!parsed || typeof parsed !== "object") return null; if (parsed.type === "ping") return { type: "ping" }; if (parsed.type === "pong") return { type: "pong" }; - if (parsed.type === "sync" && Array.isArray(parsed.clientIds)) { - const clientIds = parsed.clientIds.filter((id: unknown) => typeof id === "string" && id.trim().length > 0); - return { type: "sync", clientIds }; + if (parsed.type === "sync" && Array.isArray(parsed.connectionIds)) { + const connectionIds = parsed.connectionIds.filter((id: unknown) => typeof id === "string" && id.trim().length > 0); + return { type: "sync", connectionIds }; } - if (parsed.type === "client_connected" && typeof parsed.clientId === "string" && parsed.clientId.trim()) { - return { type: "client_connected", clientId: parsed.clientId.trim() }; + if (parsed.type === "connected" && typeof parsed.connectionId === "string" && parsed.connectionId.trim()) { + return { type: "connected", connectionId: parsed.connectionId.trim() }; } - if (parsed.type === "client_disconnected" && typeof parsed.clientId === "string" && parsed.clientId.trim()) { - return { type: "client_disconnected", clientId: parsed.clientId.trim() }; + if (parsed.type === "disconnected" && typeof parsed.connectionId === "string" && parsed.connectionId.trim()) { + return { type: "disconnected", connectionId: parsed.connectionId.trim() }; } return null; } catch { @@ -79,7 +79,7 @@ export function startRelayTransport({ let controlWs: WebSocket | null = null; let reconnectTimeout: ReturnType | null = null; let reconnectAttempt = 0; - const dataSockets = new Map(); // clientId -> ws + const dataSockets = new Map(); // connectionId -> ws let controlKeepaliveInterval: ReturnType | null = null; let controlReadyTimeout: ReturnType | null = null; let controlLastSeenAt = 0; @@ -256,24 +256,24 @@ export function startRelayTransport({ } if (msg.type === "pong") return; if (msg.type === "sync") { - for (const clientId of msg.clientIds) { - ensureClientDataSocket(clientId); + for (const connectionId of msg.connectionIds) { + ensureClientDataSocket(connectionId); } return; } - if (msg.type === "client_connected") { - ensureClientDataSocket(msg.clientId); + if (msg.type === "connected") { + ensureClientDataSocket(msg.connectionId); return; } - if (msg.type === "client_disconnected") { - const existing = dataSockets.get(msg.clientId); + if (msg.type === "disconnected") { + const existing = dataSockets.get(msg.connectionId); if (existing) { try { existing.close(1001, "Client disconnected"); } catch { // ignore } - dataSockets.delete(msg.clientId); + dataSockets.delete(msg.connectionId); } } }); @@ -291,25 +291,25 @@ export function startRelayTransport({ }, delayMs); }; - const ensureClientDataSocket = (clientId: string): void => { + const ensureClientDataSocket = (routingId: string): void => { if (stopped) return; - if (!clientId) return; - if (dataSockets.has(clientId)) return; + if (!routingId) return; + if (dataSockets.has(routingId)) return; const url = buildRelayWebSocketUrl({ endpoint: relayEndpoint, serverId, role: "server", - clientId, + connectionId: routingId, }); const socket = new WebSocket(url, { handshakeTimeout: 10_000, perMessageDeflate: false }); - dataSockets.set(clientId, socket); + dataSockets.set(routingId, socket); let attached = false; const openTimeout = setTimeout(() => { if (stopped) return; if (socket.readyState === WebSocket.OPEN) return; - relayLogger.warn({ url, clientId }, "relay_data_open_timeout_terminating"); + relayLogger.warn({ url, routingId }, "relay_data_open_timeout_terminating"); try { socket.terminate(); } catch { @@ -319,18 +319,17 @@ export function startRelayTransport({ socket.on("open", () => { clearTimeout(openTimeout); - relayLogger.info({ url, clientId }, "relay_data_connected"); + relayLogger.info({ url, routingId }, "relay_data_connected"); if (attached) return; attached = true; const externalMetadata: ExternalSocketMetadata = { transport: "relay", - externalSessionKey: `session:${clientId}`, }; if (daemonKeyPair) { void attachEncryptedSocket( socket, daemonKeyPair, - relayLogger.child({ clientId }), + relayLogger.child({ routingId }), attachSocket, externalMetadata ); @@ -342,16 +341,16 @@ export function startRelayTransport({ socket.on("close", (code, reason) => { clearTimeout(openTimeout); relayLogger.warn( - { code, reason: reason?.toString?.(), url, clientId }, + { code, reason: reason?.toString?.(), url, routingId }, "relay_data_disconnected" ); - if (dataSockets.get(clientId) === socket) { - dataSockets.delete(clientId); + if (dataSockets.get(routingId) === socket) { + dataSockets.delete(routingId); } }); socket.on("error", (err) => { - relayLogger.warn({ err, url, clientId }, "relay_data_error"); + relayLogger.warn({ err, url, routingId }, "relay_data_error"); }); }; diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index c87f48090..d6ec4dfc1 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -5421,16 +5421,29 @@ export class Session { }, effectiveTimeoutMs) try { - const result = await this.agentManager.waitForAgentEvent(agentId, { + let result = await this.agentManager.waitForAgentEvent(agentId, { signal: abortController.signal, }) - - const final = await this.getAgentPayloadById(agentId) + let final = await this.getAgentPayloadById(agentId) if (!final) { throw new Error(`Agent ${agentId} disappeared while waiting`) } - const status = result.permission ? 'permission' : result.status === 'error' ? 'error' : 'idle' + let status: 'permission' | 'error' | 'idle' = + result.permission ? 'permission' : result.status === 'error' ? 'error' : 'idle' + + // Keep wait_for_finish coherent: never return idle while the final snapshot is running. + // This can happen if a new run starts between wait resolution and final snapshot fetch. + while (status === 'idle' && final.status === 'running') { + result = await this.agentManager.waitForAgentEvent(agentId, { + signal: abortController.signal, + }) + final = await this.getAgentPayloadById(agentId) + if (!final) { + throw new Error(`Agent ${agentId} disappeared while waiting`) + } + status = result.permission ? 'permission' : result.status === 'error' ? 'error' : 'idle' + } this.emit({ type: 'wait_for_finish_response', diff --git a/packages/server/src/server/test-utils/daemon-client.ts b/packages/server/src/server/test-utils/daemon-client.ts index f635444cb..c37953b93 100644 --- a/packages/server/src/server/test-utils/daemon-client.ts +++ b/packages/server/src/server/test-utils/daemon-client.ts @@ -11,24 +11,24 @@ import { export type DaemonClientConfig = Omit< SharedDaemonClientConfig, - "webSocketFactory" | "transportFactory" + "webSocketFactory" | "transportFactory" | "clientId" >; export type CreateAgentOptions = CreateAgentRequestOptions; export { type SendMessageOptions, type DaemonEvent, type DaemonEventHandler }; let testClientCounter = 0; -function nextTestClientSessionKey(): string { +function nextTestClientId(): string { testClientCounter += 1; - return `clsk_test_client_${testClientCounter}`; + return `clid_test_client_${testClientCounter}`; } export class DaemonClient extends SharedDaemonClient { constructor(config: DaemonClientConfig) { - const clientSessionKey = config.clientSessionKey ?? nextTestClientSessionKey(); + const clientId = nextTestClientId(); super({ ...config, - clientSessionKey, + clientId, webSocketFactory: (url, options) => new WebSocket(url, { headers: options?.headers }) as unknown as WebSocketLike, }); diff --git a/packages/server/src/server/websocket-server.relay-reconnect.test.ts b/packages/server/src/server/websocket-server.relay-reconnect.test.ts index 852be1e83..786765cc0 100644 --- a/packages/server/src/server/websocket-server.relay-reconnect.test.ts +++ b/packages/server/src/server/websocket-server.relay-reconnect.test.ts @@ -74,7 +74,6 @@ vi.mock("./push/push-service.js", () => ({ import { VoiceAssistantWebSocketServer, - type ExternalSocketMetadata, } from "./websocket-server"; import { parseServerInfoStatusPayload } from "./messages.js"; import type { SpeechReadinessSnapshot } from "./speech/speech-runtime.js"; @@ -242,6 +241,65 @@ function createDownloadInProgressSpeechReadinessSnapshot(): SpeechReadinessSnaps }; } +function createHelloMessage(clientId: string) { + return { + type: "hello" as const, + clientId, + clientType: "cli" as const, + protocolVersion: 1, + }; +} + +function createDirectRequest() { + return { + headers: { + host: "localhost:6767", + origin: "http://localhost:6767", + "user-agent": "vitest", + }, + socket: { + remoteAddress: "127.0.0.1", + }, + url: "/ws", + }; +} + +async function attachRelayAndHello(params: { + server: VoiceAssistantWebSocketServer; + socket: MockSocket; + clientId: string; +}) { + await params.server.attachExternalSocket(params.socket, { transport: "relay" }); + params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId))); + await Promise.resolve(); + expect(params.socket.sent.length).toBeGreaterThan(0); + const welcome = JSON.parse(params.socket.sent[0] as string) as { + type?: unknown; + resumed?: unknown; + serverId?: unknown; + }; + expect(welcome.type).toBe("welcome"); + return welcome; +} + +async function attachDirectAndHello(params: { + server: VoiceAssistantWebSocketServer; + socket: MockSocket; + clientId: string; +}) { + await (params.server as any).attachSocket(params.socket, createDirectRequest()); + params.socket.emit("message", JSON.stringify(createHelloMessage(params.clientId))); + await Promise.resolve(); + expect(params.socket.sent.length).toBeGreaterThan(0); + const welcome = JSON.parse(params.socket.sent[0] as string) as { + type?: unknown; + resumed?: unknown; + serverId?: unknown; + }; + expect(welcome.type).toBe("welcome"); + return welcome; +} + describe("relay external socket reconnect behavior", () => { beforeEach(() => { sessionMock.instances.length = 0; @@ -254,13 +312,15 @@ describe("relay external socket reconnect behavior", () => { test("keeps the same session when relay reconnects within grace window", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-1", - }; + const clientId = "cid-relay-reconnect"; const socket1 = new MockSocket(); - await server.attachExternalSocket(socket1, metadata); + const firstWelcome = await attachRelayAndHello({ + server, + socket: socket1, + clientId, + }); + expect(firstWelcome.resumed).toBe(false); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; @@ -269,7 +329,12 @@ describe("relay external socket reconnect behavior", () => { expect(session.cleanup).not.toHaveBeenCalled(); const socket2 = new MockSocket(); - await server.attachExternalSocket(socket2, metadata); + const secondWelcome = await attachRelayAndHello({ + server, + socket: socket2, + clientId, + }); + expect(secondWelcome.resumed).toBe(true); expect(sessionMock.instances).toHaveLength(1); await vi.advanceTimersByTimeAsync(20_000); @@ -278,19 +343,8 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); - test("rejects direct socket attach when clientSessionKey is missing", async () => { + test("closes pending connection when hello timeout elapses", async () => { const server = createServer(); - const request = { - headers: { - host: "localhost:6767", - origin: "http://localhost:6767", - "user-agent": "vitest", - }, - socket: { - remoteAddress: "127.0.0.1", - }, - url: "/ws", - }; const socket = new MockSocket(); let closeCode: number | null = null; @@ -300,63 +354,105 @@ describe("relay external socket reconnect behavior", () => { closeReason = typeof reason === "string" ? reason : String(reason ?? ""); }); - await (server as any).attachSocket(socket, request); + await (server as any).attachSocket(socket, createDirectRequest()); + await vi.advanceTimersByTimeAsync(15_000); - expect(closeCode).toBe(1008); - expect(closeReason).toBe("Missing clientSessionKey"); + expect(closeCode).toBe(4001); + expect(closeReason).toBe("Hello timeout"); expect(sessionMock.instances).toHaveLength(0); await server.close(); }); - test("attaches multiple sockets to the same relay external session key", async () => { + test("marks hello as resumed when clientId already has a session", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-multi", - }; + const clientId = "cid-resume-flag"; - const socket1 = new MockSocket(); - await server.attachExternalSocket(socket1, metadata); - expect(sessionMock.instances).toHaveLength(1); - const session = sessionMock.instances[0]!; - - const socket2 = new MockSocket(); - await server.attachExternalSocket(socket2, metadata); - expect(sessionMock.instances).toHaveLength(1); - - const onMessage = session.args.onMessage as - | ((msg: { type: "status"; payload: { status: string } }) => void) - | undefined; - expect(onMessage).toBeTypeOf("function"); - - onMessage?.({ - type: "status", - payload: { status: "ok" }, + const firstSocket = new MockSocket(); + const firstWelcome = await attachRelayAndHello({ + server, + socket: firstSocket, + clientId, }); + expect(firstWelcome.resumed).toBe(false); - expect(socket1.sent.length).toBeGreaterThan(1); - expect(socket2.sent.length).toBeGreaterThan(1); + firstSocket.emit("close", 1006, ""); + await vi.advanceTimersByTimeAsync(1_000); + + const secondSocket = new MockSocket(); + const secondWelcome = await attachRelayAndHello({ + server, + socket: secondSocket, + clientId, + }); + expect(secondWelcome.resumed).toBe(true); await server.close(); }); - test("reuses direct session when clientSessionKey reconnects within grace window", async () => { + test("marks hello as not resumed for new clientIds", async () => { const server = createServer(); - const request = { - headers: { - host: "localhost:6767", - origin: "http://localhost:6767", - "user-agent": "vitest", - }, - socket: { - remoteAddress: "127.0.0.1", - }, - url: "/ws?clientSessionKey=clsk_direct_reconnect", - }; + + const firstSocket = new MockSocket(); + const firstWelcome = await attachRelayAndHello({ + server, + socket: firstSocket, + clientId: "cid-new-1", + }); + expect(firstWelcome.resumed).toBe(false); + + const secondSocket = new MockSocket(); + const secondWelcome = await attachRelayAndHello({ + server, + socket: secondSocket, + clientId: "cid-new-2", + }); + expect(secondWelcome.resumed).toBe(false); + expect(sessionMock.instances).toHaveLength(2); + + await server.close(); + }); + + test("rejects session messages before hello", async () => { + const server = createServer(); + const socket = new MockSocket(); + let closeCode: number | null = null; + let closeReason = ""; + socket.on("close", (code: unknown, reason: unknown) => { + closeCode = typeof code === "number" ? code : null; + closeReason = typeof reason === "string" ? reason : String(reason ?? ""); + }); + + await server.attachExternalSocket(socket, { transport: "relay" }); + socket.emit( + "message", + JSON.stringify({ + type: "session", + message: { + type: "ping", + }, + }) + ); + await Promise.resolve(); + + expect(closeCode).toBe(4002); + expect(["Invalid hello", "Session message before hello"]).toContain(closeReason); + expect(sessionMock.instances).toHaveLength(0); + + await server.close(); + }); + + test("reuses direct session when same clientId reconnects within grace window", async () => { + const server = createServer(); + const clientId = "cid-direct-reconnect"; const socket1 = new MockSocket(); - await (server as any).attachSocket(socket1, request); + const firstWelcome = await attachDirectAndHello({ + server, + socket: socket1, + clientId, + }); + expect(firstWelcome.resumed).toBe(false); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; @@ -365,7 +461,12 @@ describe("relay external socket reconnect behavior", () => { expect(session.cleanup).not.toHaveBeenCalled(); const socket2 = new MockSocket(); - await (server as any).attachSocket(socket2, request); + const secondWelcome = await attachDirectAndHello({ + server, + socket: socket2, + clientId, + }); + expect(secondWelcome.resumed).toBe(true); expect(sessionMock.instances).toHaveLength(1); await vi.advanceTimersByTimeAsync(20_000); @@ -374,32 +475,25 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); - test("reuses one session when switching from direct to relay with the same session key", async () => { + test("reuses one session when switching from direct to relay with the same clientId", async () => { const server = createServer(); - const clientSessionKey = "clsk_switch_path"; - const directRequest = { - headers: { - host: "localhost:6767", - origin: "http://localhost:6767", - "user-agent": "vitest", - }, - socket: { - remoteAddress: "127.0.0.1", - }, - url: `/ws?clientSessionKey=${clientSessionKey}`, - }; - const relayMetadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: `session:${clientSessionKey}`, - }; + const clientId = "cid-switch-path"; const directSocket = new MockSocket(); - await (server as any).attachSocket(directSocket, directRequest); + await attachDirectAndHello({ + server, + socket: directSocket, + clientId, + }); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; const relaySocket = new MockSocket(); - await server.attachExternalSocket(relaySocket, relayMetadata); + await attachRelayAndHello({ + server, + socket: relaySocket, + clientId, + }); expect(sessionMock.instances).toHaveLength(1); const onMessage = session.args.onMessage as @@ -411,8 +505,8 @@ describe("relay external socket reconnect behavior", () => { payload: { status: "ok" }, }); - expect(directSocket.sent.length).toBeGreaterThan(1); - expect(relaySocket.sent.length).toBeGreaterThan(1); + expect(directSocket.sent.length).toBeGreaterThan(0); + expect(relaySocket.sent.length).toBeGreaterThan(0); directSocket.emit("close", 1006, ""); await vi.advanceTimersByTimeAsync(1_000); @@ -427,13 +521,14 @@ describe("relay external socket reconnect behavior", () => { test("cleans up relay session when reconnect grace expires", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-2", - }; + const clientId = "cid-relay-grace-expire"; const socket1 = new MockSocket(); - await server.attachExternalSocket(socket1, metadata); + await attachRelayAndHello({ + server, + socket: socket1, + clientId, + }); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; @@ -444,53 +539,46 @@ describe("relay external socket reconnect behavior", () => { await server.close(); }); - test("includes voice capabilities in server_info when speech readiness exists", async () => { + test("includes voice capabilities in welcome when speech readiness exists", async () => { const speechReadiness = createReadySpeechReadinessSnapshot(); const server = createServer({ speechReadiness }); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-server-info-capabilities", - }; const socket = new MockSocket(); - await server.attachExternalSocket(socket, metadata); - - expect(socket.sent).toHaveLength(1); - expect(typeof socket.sent[0]).toBe("string"); - const envelope = JSON.parse(socket.sent[0] as string) as { - type?: unknown; - message?: { - type?: unknown; - payload?: unknown; + const welcome = await attachRelayAndHello({ + server, + socket, + clientId: "cid-server-info-capabilities", + }) as { + version?: unknown; + capabilities?: { + voice?: { + dictation?: { enabled?: unknown; reason?: unknown }; + voice?: { enabled?: unknown; reason?: unknown }; + }; }; }; - expect(envelope.type).toBe("session"); - expect(envelope.message?.type).toBe("status"); - - const payload = parseServerInfoStatusPayload(envelope.message?.payload); - expect(payload?.status).toBe("server_info"); - expect(payload?.version).toBe(TEST_DAEMON_VERSION); - expect(payload?.capabilities?.voice?.dictation.enabled).toBe( + expect(welcome.version).toBe(TEST_DAEMON_VERSION); + expect(welcome.capabilities?.voice?.dictation?.enabled).toBe( speechReadiness.dictation.enabled ); - expect(payload?.capabilities?.voice?.dictation.reason).toBe(""); - expect(payload?.capabilities?.voice?.voice.enabled).toBe( + expect(welcome.capabilities?.voice?.dictation?.reason).toBe(""); + expect(welcome.capabilities?.voice?.voice?.enabled).toBe( speechReadiness.realtimeVoice.enabled ); - expect(payload?.capabilities?.voice?.voice.reason).toBe(""); + expect(welcome.capabilities?.voice?.voice?.reason).toBe(""); await server.close(); }); test("broadcasts updated server_info when capabilities change", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-server-info-broadcast", - }; const socket = new MockSocket(); - await server.attachExternalSocket(socket, metadata); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-server-info-broadcast", + }); expect(socket.sent).toHaveLength(1); const speechReadiness = createReadySpeechReadinessSnapshot(); @@ -513,12 +601,12 @@ describe("relay external socket reconnect behavior", () => { test("includes temporary retry guidance while models are downloading", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-server-info-download-guidance", - }; const socket = new MockSocket(); - await server.attachExternalSocket(socket, metadata); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-server-info-download-guidance", + }); expect(socket.sent).toHaveLength(1); server.publishSpeechReadiness(createDownloadInProgressSpeechReadinessSnapshot()); @@ -542,13 +630,13 @@ describe("relay external socket reconnect behavior", () => { test("routes inbound binary mux frames to session.handleBinaryFrame", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-binary-inbound", - }; const socket = new MockSocket(); - await server.attachExternalSocket(socket, metadata); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-binary-inbound", + }); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; @@ -583,13 +671,13 @@ describe("relay external socket reconnect behavior", () => { test("sends outbound binary mux frames from session over websocket", async () => { const server = createServer(); - const metadata: ExternalSocketMetadata = { - transport: "relay", - externalSessionKey: "session:client-binary-outbound", - }; const socket = new MockSocket(); - await server.attachExternalSocket(socket, metadata); + await attachRelayAndHello({ + server, + socket, + clientId: "cid-binary-outbound", + }); expect(sessionMock.instances).toHaveLength(1); const session = sessionMock.instances[0]!; diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 8157d0bd8..d691390ab 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -10,6 +10,8 @@ import type { TerminalManager } from "../terminal/terminal-manager.js"; import type pino from "pino"; import { type ServerInfoStatusPayload, + type WSHelloMessage, + type WSWelcomeMessage, WSInboundMessageSchema, type ServerCapabilityState, type ServerCapabilities, @@ -51,7 +53,11 @@ import { export type AgentMcpTransportFactory = () => Promise; export type ExternalSocketMetadata = { transport: "relay"; - externalSessionKey: string; +}; + +type PendingConnection = { + connectionLogger: pino.Logger; + helloTimeout: ReturnType | null; }; type WebSocketServerConfig = { @@ -152,11 +158,15 @@ type SessionConnection = { clientId: string; connectionLogger: pino.Logger; sockets: Set; - externalSessionKey: string | null; externalDisconnectCleanupTimeout: ReturnType | null; }; const EXTERNAL_SESSION_DISCONNECT_GRACE_MS = 90_000; +const HELLO_TIMEOUT_MS = 15_000; +const WS_CLOSE_HELLO_TIMEOUT = 4001; +const WS_CLOSE_INVALID_HELLO = 4002; +const WS_CLOSE_INCOMPATIBLE_PROTOCOL = 4003; +const WS_PROTOCOL_VERSION = 1; export class MissingDaemonVersionError extends Error { constructor() { @@ -171,9 +181,9 @@ export class MissingDaemonVersionError extends Error { export class VoiceAssistantWebSocketServer { private readonly logger: pino.Logger; private readonly wss: WebSocketServer; + private readonly pendingConnections: Map = new Map(); private readonly sessions: Map = new Map(); private readonly externalSessionsByKey: Map = new Map(); - private clientIdCounter = 0; private readonly serverId: string; private readonly daemonVersion: string; private readonly agentManager: AgentManager; @@ -334,7 +344,7 @@ export class VoiceAssistantWebSocketServer { return; } this.serverCapabilities = next; - this.broadcastServerInfo(); + this.broadcastCapabilitiesUpdate(); } public async attachExternalSocket( @@ -350,6 +360,14 @@ export class VoiceAssistantWebSocketServer { ...this.externalSessionsByKey.values(), ]); + const pendingSockets = new Set(this.pendingConnections.keys()); + for (const pending of this.pendingConnections.values()) { + if (pending.helloTimeout) { + clearTimeout(pending.helloTimeout); + pending.helloTimeout = null; + } + } + const cleanupPromises: Promise[] = []; for (const connection of uniqueConnections) { if (connection.externalDisconnectCleanupTimeout) { @@ -372,7 +390,22 @@ export class VoiceAssistantWebSocketServer { ); } } + + for (const ws of pendingSockets) { + cleanupPromises.push( + new Promise((resolve) => { + if (ws.readyState === 3) { + resolve(); + return; + } + ws.once("close", () => resolve()); + ws.close(); + }) + ); + } + await Promise.all(cleanupPromises); + this.pendingConnections.clear(); this.sessions.clear(); this.externalSessionsByKey.clear(); this.wss.close(); @@ -416,61 +449,7 @@ export class VoiceAssistantWebSocketServer { metadata?: ExternalSocketMetadata ): Promise { const requestMetadata = extractSocketRequestMetadata(request); - const relayExternalSessionKey = - metadata?.transport === "relay" && metadata.externalSessionKey.trim().length > 0 - ? metadata.externalSessionKey - : null; - const directExternalSessionKey = - typeof requestMetadata.clientSessionKey === "string" && - requestMetadata.clientSessionKey.trim().length > 0 - ? `session:${requestMetadata.clientSessionKey.trim()}` - : null; - const externalSessionKey = relayExternalSessionKey ?? directExternalSessionKey; - - if (metadata?.transport !== "relay" && !directExternalSessionKey) { - this.logger.warn( - { - host: requestMetadata.host, - origin: requestMetadata.origin, - remoteAddress: requestMetadata.remoteAddress, - }, - "Rejected direct connection without clientSessionKey" - ); - try { - ws.close(1008, "Missing clientSessionKey"); - } catch { - // ignore close errors - } - return; - } - - if (externalSessionKey) { - const existing = this.externalSessionsByKey.get(externalSessionKey); - if (existing) { - if (existing.externalDisconnectCleanupTimeout) { - clearTimeout(existing.externalDisconnectCleanupTimeout); - existing.externalDisconnectCleanupTimeout = null; - } - - existing.sockets.add(ws); - this.sessions.set(ws, existing); - this.sendServerInfo(ws); - existing.connectionLogger.trace( - { - clientId: existing.clientId, - externalSessionKey, - totalSessions: this.sessions.size, - }, - "Client reconnected" - ); - this.bindSocketHandlers(ws, existing); - return; - } - } - - const clientId = `client-${++this.clientIdCounter}`; const connectionLoggerFields: Record = { - clientId, transport: metadata?.transport === "relay" ? "relay" : "direct", }; if (requestMetadata.host) { @@ -486,6 +465,47 @@ export class VoiceAssistantWebSocketServer { connectionLoggerFields.remoteAddress = requestMetadata.remoteAddress; } const connectionLogger = this.logger.child(connectionLoggerFields); + + const pending: PendingConnection = { + connectionLogger, + helloTimeout: null, + }; + const timeout = setTimeout(() => { + if (this.pendingConnections.get(ws) !== pending) { + return; + } + pending.helloTimeout = null; + this.pendingConnections.delete(ws); + pending.connectionLogger.warn( + { timeoutMs: HELLO_TIMEOUT_MS }, + "Closing connection due to missing hello" + ); + try { + ws.close(WS_CLOSE_HELLO_TIMEOUT, "Hello timeout"); + } catch { + // ignore close errors + } + }, HELLO_TIMEOUT_MS); + pending.helloTimeout = timeout; + (timeout as unknown as { unref?: () => void }).unref?.(); + + this.pendingConnections.set(ws, pending); + this.bindSocketHandlers(ws); + + pending.connectionLogger.trace( + { + totalPendingConnections: this.pendingConnections.size, + }, + "Client connected; awaiting hello" + ); + } + + private createSessionConnection(params: { + ws: WebSocketLike; + clientId: string; + connectionLogger: pino.Logger; + }): SessionConnection { + const { ws, clientId, connectionLogger } = params; let connection: SessionConnection | null = null; const session = new Session({ @@ -538,23 +558,109 @@ export class VoiceAssistantWebSocketServer { clientId, connectionLogger, sockets: new Set([ws]), - externalSessionKey, externalDisconnectCleanupTimeout: null, }; + return connection; + } - this.sessions.set(ws, connection); - if (externalSessionKey) { - this.externalSessionsByKey.set(externalSessionKey, connection); + private clearPendingConnection(ws: WebSocketLike): PendingConnection | null { + const pending = this.pendingConnections.get(ws); + if (!pending) { + return null; + } + if (pending.helloTimeout) { + clearTimeout(pending.helloTimeout); + pending.helloTimeout = null; + } + this.pendingConnections.delete(ws); + return pending; + } + + private buildWelcomeMessage(params: { resumed: boolean }): WSWelcomeMessage { + return { + type: "welcome", + serverId: this.serverId, + hostname: getHostname(), + version: this.daemonVersion, + resumed: params.resumed, + ...(this.serverCapabilities ? { capabilities: this.serverCapabilities } : {}), + }; + } + + private handleHello(params: { + ws: WebSocketLike; + message: WSHelloMessage; + pending: PendingConnection; + }): void { + const { ws, message, pending } = params; + + if (message.protocolVersion !== WS_PROTOCOL_VERSION) { + this.clearPendingConnection(ws); + pending.connectionLogger.warn( + { + receivedProtocolVersion: message.protocolVersion, + expectedProtocolVersion: WS_PROTOCOL_VERSION, + }, + "Rejected hello due to protocol version mismatch" + ); + try { + ws.close(WS_CLOSE_INCOMPATIBLE_PROTOCOL, "Incompatible protocol version"); + } catch { + // ignore close errors + } + return; } - this.sendServerInfo(ws); + const clientId = message.clientId.trim(); + if (clientId.length === 0) { + this.clearPendingConnection(ws); + pending.connectionLogger.warn("Rejected hello with empty clientId"); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Invalid hello"); + } catch { + // ignore close errors + } + return; + } - connectionLogger.trace( - { clientId, externalSessionKey, totalSessions: this.sessions.size }, - "Client connected" + this.clearPendingConnection(ws); + const existing = this.externalSessionsByKey.get(clientId); + if (existing) { + if (existing.externalDisconnectCleanupTimeout) { + clearTimeout(existing.externalDisconnectCleanupTimeout); + existing.externalDisconnectCleanupTimeout = null; + } + existing.sockets.add(ws); + this.sessions.set(ws, existing); + this.sendToClient(ws, this.buildWelcomeMessage({ resumed: true })); + existing.connectionLogger.trace( + { + clientId, + resumed: true, + totalSessions: this.sessions.size, + }, + "Client connected via hello" + ); + return; + } + + const connectionLogger = pending.connectionLogger.child({ clientId }); + const connection = this.createSessionConnection({ + ws, + clientId, + connectionLogger, + }); + this.sessions.set(ws, connection); + this.externalSessionsByKey.set(clientId, connection); + this.sendToClient(ws, this.buildWelcomeMessage({ resumed: false })); + connection.connectionLogger.trace( + { + clientId, + resumed: false, + totalSessions: this.sessions.size, + }, + "Client connected via hello" ); - - this.bindSocketHandlers(ws, connection); } private buildServerInfoStatusPayload(): ServerInfoStatusPayload { @@ -567,7 +673,7 @@ export class VoiceAssistantWebSocketServer { }; } - private broadcastServerInfo(): void { + private broadcastCapabilitiesUpdate(): void { this.broadcast( wrapSessionMessage({ type: "status", @@ -576,27 +682,13 @@ export class VoiceAssistantWebSocketServer { ); } - private sendServerInfo(ws: WebSocketLike): void { - // Advertise stable server identity immediately on connect (used for URL/shareable IDs). - this.sendToClient( - ws, - wrapSessionMessage({ - type: "status", - payload: this.buildServerInfoStatusPayload(), - }) - ); - } - - private bindSocketHandlers( - ws: WebSocketLike, - connection: SessionConnection - ): void { + private bindSocketHandlers(ws: WebSocketLike): void { ws.on("message", (data) => { void this.handleRawMessage(ws, data); }); ws.on("close", async (code: number, reason: unknown) => { - await this.detachSocket(ws, connection, { + await this.detachSocket(ws, { code: typeof code === "number" ? code : undefined, reason, }); @@ -604,8 +696,11 @@ export class VoiceAssistantWebSocketServer { ws.on("error", async (error) => { const err = error instanceof Error ? error : new Error(String(error)); - connection.connectionLogger.error({ err }, "Client error"); - await this.detachSocket(ws, connection, { error: err }); + const active = this.sessions.get(ws); + const pending = this.pendingConnections.get(ws); + const log = active?.connectionLogger ?? pending?.connectionLogger ?? this.logger; + log.error({ err }, "Client error"); + await this.detachSocket(ws, { error: err }); }); } @@ -623,19 +718,33 @@ export class VoiceAssistantWebSocketServer { private async detachSocket( ws: WebSocketLike, - connection: SessionConnection, details: { code?: number; reason?: unknown; error?: Error; } ): Promise { - const activeConnection = this.sessions.get(ws); - if (activeConnection !== connection) return; + const pending = this.clearPendingConnection(ws); + if (pending) { + pending.connectionLogger.trace( + { + code: details.code, + reason: stringifyCloseReason(details.reason), + }, + "Pending client disconnected" + ); + return; + } + + const connection = this.sessions.get(ws); + if (!connection) { + return; + } + this.sessions.delete(ws); connection.sockets.delete(ws); - if (connection.externalSessionKey && connection.sockets.size === 0) { + if (connection.sockets.size === 0) { if (connection.externalDisconnectCleanupTimeout) { clearTimeout(connection.externalDisconnectCleanupTimeout); } @@ -651,7 +760,6 @@ export class VoiceAssistantWebSocketServer { connection.connectionLogger.trace( { clientId: connection.clientId, - externalSessionKey: connection.externalSessionKey, code: details.code, reason: stringifyCloseReason(details.reason), reconnectGraceMs: EXTERNAL_SESSION_DISCONNECT_GRACE_MS, @@ -690,11 +798,9 @@ export class VoiceAssistantWebSocketServer { this.sessions.delete(socket); } connection.sockets.clear(); - if (connection.externalSessionKey) { - const existing = this.externalSessionsByKey.get(connection.externalSessionKey); - if (existing === connection) { - this.externalSessionsByKey.delete(connection.externalSessionKey); - } + const existing = this.externalSessionsByKey.get(connection.clientId); + if (existing === connection) { + this.externalSessionsByKey.delete(connection.clientId); } connection.connectionLogger.trace( @@ -708,15 +814,24 @@ export class VoiceAssistantWebSocketServer { ws: WebSocketLike, data: Buffer | ArrayBuffer | Buffer[] | string ): Promise { + const activeConnection = this.sessions.get(ws); + const pendingConnection = this.pendingConnections.get(ws); + const log = activeConnection?.connectionLogger ?? pendingConnection?.connectionLogger ?? this.logger; + try { - const activeConnection = this.sessions.get(ws); const buffer = bufferFromWsData(data); const asBytes = asUint8Array(buffer); if (asBytes) { const frame = decodeBinaryMuxFrame(asBytes); if (frame) { if (!activeConnection) { - this.logger.error("No session found for client"); + log.warn("Rejected binary frame before hello"); + this.clearPendingConnection(ws); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Session message before hello"); + } catch { + // ignore close errors + } return; } activeConnection.session.handleBinaryFrame(frame); @@ -726,6 +841,22 @@ export class VoiceAssistantWebSocketServer { const parsed = JSON.parse(buffer.toString()); const parsedMessage = WSInboundMessageSchema.safeParse(parsed); if (!parsedMessage.success) { + if (pendingConnection) { + pendingConnection.connectionLogger.warn( + { + error: parsedMessage.error.message, + }, + "Rejected pending message before hello" + ); + this.clearPendingConnection(ws); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Invalid hello"); + } catch { + // ignore close errors + } + return; + } + const requestInfo = extractRequestInfoFromUnknownWsInbound(parsed); const isUnknownSchema = requestInfo?.requestId != null && @@ -734,7 +865,6 @@ export class VoiceAssistantWebSocketServer { "type" in parsed && (parsed as { type?: unknown }).type === "session"; - const log = activeConnection?.connectionLogger ?? this.logger; log.warn( { clientId: activeConnection?.clientId, @@ -786,8 +916,43 @@ export class VoiceAssistantWebSocketServer { return; } + if (pendingConnection) { + if (message.type === "hello") { + this.handleHello({ + ws, + message, + pending: pendingConnection, + }); + return; + } + + pendingConnection.connectionLogger.warn( + { + messageType: message.type, + }, + "Rejected pending message before hello" + ); + this.clearPendingConnection(ws); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Session message before hello"); + } catch { + // ignore close errors + } + return; + } + if (!activeConnection) { - this.logger.error("No session found for client"); + this.logger.error("No connection found for websocket"); + return; + } + + if (message.type === "hello") { + activeConnection.connectionLogger.warn("Received hello on active connection"); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Unexpected hello"); + } catch { + // ignore close errors + } return; } @@ -816,7 +981,7 @@ export class VoiceAssistantWebSocketServer { ? `${rawPayload.slice(0, 2000)}... (truncated)` : rawPayload; - this.logger.error( + log.error( { err, rawPayload: trimmedRawPayload, @@ -825,6 +990,16 @@ export class VoiceAssistantWebSocketServer { "Failed to parse/handle message" ); + if (this.pendingConnections.has(ws)) { + this.clearPendingConnection(ws); + try { + ws.close(WS_CLOSE_INVALID_HELLO, "Invalid hello"); + } catch { + // ignore close errors + } + return; + } + const requestInfo = extractRequestInfoFromUnknownWsInbound(parsedPayload); if (requestInfo) { this.sendToClient( @@ -952,7 +1127,6 @@ type SocketRequestMetadata = { origin?: string; userAgent?: string; remoteAddress?: string; - clientSessionKey?: string; }; function extractSocketRequestMetadata(request: unknown): SocketRequestMetadata { @@ -983,30 +1157,12 @@ function extractSocketRequestMetadata(request: unknown): SocketRequestMetadata { typeof record.socket?.remoteAddress === "string" ? record.socket.remoteAddress : undefined; - const rawUrl = typeof record.url === "string" ? record.url : null; - const clientSessionKey = (() => { - if (!rawUrl) { - return undefined; - } - try { - const parsed = new URL(rawUrl, "http://localhost"); - const value = parsed.searchParams.get("clientSessionKey"); - if (!value) { - return undefined; - } - const trimmed = value.trim(); - return trimmed.length > 0 ? trimmed : undefined; - } catch { - return undefined; - } - })(); return { ...(host ? { host } : {}), ...(origin ? { origin } : {}), ...(userAgent ? { userAgent } : {}), ...(remoteAddress ? { remoteAddress } : {}), - ...(clientSessionKey ? { clientSessionKey } : {}), }; } diff --git a/packages/server/src/shared/daemon-endpoints.test.ts b/packages/server/src/shared/daemon-endpoints.test.ts index e28529afb..cb843db55 100644 --- a/packages/server/src/shared/daemon-endpoints.test.ts +++ b/packages/server/src/shared/daemon-endpoints.test.ts @@ -13,12 +13,24 @@ describe("relay websocket URL versioning", () => { endpoint: "relay.paseo.sh:443", serverId: "srv_test", role: "client", - clientSessionKey: "clsk_test", }) ); expect(url.searchParams.get("v")).toBe(CURRENT_RELAY_PROTOCOL_VERSION); - expect(url.searchParams.get("clientId")).toBe("clsk_test"); + expect(url.searchParams.has("connectionId")).toBe(false); + }); + + test("includes connectionId when provided (server data sockets)", () => { + const url = new URL( + buildRelayWebSocketUrl({ + endpoint: "relay.paseo.sh:443", + serverId: "srv_test", + role: "server", + connectionId: "conn_abc123", + }) + ); + + expect(url.searchParams.get("connectionId")).toBe("conn_abc123"); }); test("allows explicitly requesting v1 relay URLs", () => { diff --git a/packages/server/src/shared/daemon-endpoints.ts b/packages/server/src/shared/daemon-endpoints.ts index dc68387ff..7dbeb1936 100644 --- a/packages/server/src/shared/daemon-endpoints.ts +++ b/packages/server/src/shared/daemon-endpoints.ts @@ -83,28 +83,23 @@ function shouldUseSecureWebSocket(port: number): boolean { } export function buildDaemonWebSocketUrl( - endpoint: string, - params?: { clientSessionKey?: string } + endpoint: string ): string { const { host, port, isIpv6 } = parseHostPort(endpoint); const protocol = shouldUseSecureWebSocket(port) ? "wss" : "ws"; const hostPart = isIpv6 ? `[${host}]` : host; - const url = new URL(`${protocol}://${hostPart}:${port}/ws`); - if ( - typeof params?.clientSessionKey === "string" && - params.clientSessionKey.trim().length > 0 - ) { - url.searchParams.set("clientSessionKey", params.clientSessionKey.trim()); - } - return url.toString(); + return new URL(`${protocol}://${hostPart}:${port}/ws`).toString(); } export function buildRelayWebSocketUrl(params: { endpoint: string; serverId: string; role: RelayRole; - clientId?: string; - clientSessionKey?: string; + /** + * Per-connection routing identifier used by the daemon to open server data sockets. + * Clients should NOT provide this — the relay assigns a routing ID on connect. + */ + connectionId?: string; version?: RelayProtocolVersion | 1 | 2; }): string { const { host, port, isIpv6 } = parseHostPort(params.endpoint); @@ -114,16 +109,8 @@ export function buildRelayWebSocketUrl(params: { url.searchParams.set("serverId", params.serverId); url.searchParams.set("role", params.role); url.searchParams.set("v", normalizeRelayProtocolVersion(params.version)); - if ( - params.clientId && - params.clientSessionKey && - params.clientId !== params.clientSessionKey - ) { - throw new Error("clientId and clientSessionKey must match when both are provided"); - } - const resolvedClientId = params.clientId ?? params.clientSessionKey; - if (resolvedClientId) { - url.searchParams.set("clientId", resolvedClientId); + if (params.connectionId) { + url.searchParams.set("connectionId", params.connectionId); } return url.toString(); } diff --git a/packages/server/src/shared/messages.stream-parsing.test.ts b/packages/server/src/shared/messages.stream-parsing.test.ts index abcfa3c4c..9ef651270 100644 --- a/packages/server/src/shared/messages.stream-parsing.test.ts +++ b/packages/server/src/shared/messages.stream-parsing.test.ts @@ -78,6 +78,48 @@ describe('shared messages stream parsing', () => { } }) + it('parses representative sub_agent tool_call event', () => { + const parsed = AgentStreamMessageSchema.parse({ + type: 'agent_stream', + payload: { + agentId: 'agent_live', + timestamp: '2026-02-08T20:10:00.000Z', + event: { + type: 'timeline', + provider: 'claude', + item: { + type: 'tool_call', + callId: 'call_sub_agent_live', + name: 'Task', + status: 'running', + detail: { + type: 'sub_agent', + subAgentType: 'Explore', + description: 'Inspect repository structure', + log: '[Read] README.md', + actions: [ + { + index: 1, + toolName: 'Read', + summary: 'README.md', + }, + ], + }, + error: null, + }, + }, + }, + }) + + expect(parsed.payload.event.type).toBe('timeline') + if (parsed.payload.event.type === 'timeline') { + expect(parsed.payload.event.item.type).toBe('tool_call') + if (parsed.payload.event.item.type === 'tool_call') { + expect(parsed.payload.event.item.detail.type).toBe('sub_agent') + } + } + }) + it('rejects removed initialize_agent_request inbound payload', () => { const parsed = SessionInboundMessageSchema.safeParse({ type: 'initialize_agent_request', diff --git a/packages/server/src/shared/messages.tool-call-schema.test.ts b/packages/server/src/shared/messages.tool-call-schema.test.ts index 690a5cebe..8b553f383 100644 --- a/packages/server/src/shared/messages.tool-call-schema.test.ts +++ b/packages/server/src/shared/messages.tool-call-schema.test.ts @@ -109,4 +109,41 @@ describe("shared messages tool_call schema", () => { expect(missingDetail.success).toBe(false); expect(legacyStatus.success).toBe(false); }); + + it("parses canonical sub_agent detail payload", () => { + const parsed = AgentTimelineItemPayloadSchema.parse({ + type: "tool_call", + callId: "call_sub_agent_1", + name: "Task", + status: "running", + error: null, + detail: { + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository structure", + log: "[Read] README.md\n[Bash] ls", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + { + index: 2, + toolName: "Bash", + summary: "ls", + }, + ], + }, + }); + + expect(parsed.type).toBe("tool_call"); + if (parsed.type === "tool_call") { + expect(parsed.detail.type).toBe("sub_agent"); + if (parsed.detail.type === "sub_agent") { + expect(parsed.detail.subAgentType).toBe("Explore"); + expect(parsed.detail.actions).toHaveLength(2); + } + } + }); }); diff --git a/packages/server/src/shared/tool-call-display.test.ts b/packages/server/src/shared/tool-call-display.test.ts index 3b99051bf..c1f887d0b 100644 --- a/packages/server/src/shared/tool-call-display.test.ts +++ b/packages/server/src/shared/tool-call-display.test.ts @@ -38,24 +38,29 @@ describe("shared tool-call display mapping", () => { }); }); - it("keeps task metadata summary on unknown detail", () => { + it("uses sub-agent detail for task label and description", () => { const display = buildToolCallDisplayModel({ name: "task", status: "running", error: null, detail: { - type: "unknown", - input: null, - output: null, - }, - metadata: { - subAgentActivity: "Running tests", + type: "sub_agent", + subAgentType: "Explore", + description: "Inspect repository structure", + log: "[Read] README.md", + actions: [ + { + index: 1, + toolName: "Read", + summary: "README.md", + }, + ], }, }); expect(display).toEqual({ - displayName: "Task", - summary: "Running tests", + displayName: "Explore", + summary: "Inspect repository structure", }); }); diff --git a/packages/server/src/shared/tool-call-display.ts b/packages/server/src/shared/tool-call-display.ts index b6eb12233..63568bdf8 100644 --- a/packages/server/src/shared/tool-call-display.ts +++ b/packages/server/src/shared/tool-call-display.ts @@ -87,6 +87,10 @@ export function buildToolCallDisplayModel(input: ToolCallDisplayInput): ToolCall displayName = "Worktree Setup"; summary = input.detail.branchName; break; + case "sub_agent": + displayName = readString(input.detail.subAgentType) ?? "Task"; + summary = readString(input.detail.description); + break; case "unknown": break; }