diff --git a/ts/packages/agents/code/src/codeActionHandler.ts b/ts/packages/agents/code/src/codeActionHandler.ts index 0fdbbb9bf..e83d39e37 100644 --- a/ts/packages/agents/code/src/codeActionHandler.ts +++ b/ts/packages/agents/code/src/codeActionHandler.ts @@ -24,7 +24,7 @@ import { } from "@typeagent/agent-sdk/helpers/action"; import { evaluateCodeReadiness, - resolveCodePort, + resolveCodePortOverride, setupCode, whichExists, } from "./readiness.js"; @@ -32,13 +32,23 @@ import { const debug = registerDebug("typeagent:code"); // Shared WebSocket server that bridges this code agent to the Coda VS Code -// extension (ts/packages/coda) on port 8082. Created on first enable, closed -// when the last session disables the code agent. Storing it per-session caused -// "No websocket connection" errors when an action ran on a session different -// from the one that originally created the server (e.g. after schema enable on -// a different conversation), and also masked EADDRINUSE failures from a second -// bind attempt on port 8082. +// extension (ts/packages/coda). Created on first session-enable, closed when +// the last session disables. Storing it per-session caused "No websocket +// connection" errors when an action ran on a session different from the one +// that originally created the server (e.g. after schema enable on a different +// conversation), and also masked EADDRINUSE failures from a second bind +// attempt on the configured port. +// +// Port allocation: by default the OS picks a free ephemeral port (port=0). +// Each session that uses the shared server registers it under its own +// `sessionContextId`, so the PortRegistrar's `closeSessionContext` backstop +// auto-releases per-session entries and `lookup("code")` keeps returning the +// shared port as long as ≥1 session has it enabled. `CODE_WEBSOCKET_PORT` +// remains an explicit override (useful for pinning the port when debugging +// or when an external client expects a known address). let sharedWebSocketServer: CodeAgentWebSocketServer | undefined; +let sharedStartingPromise: Promise | undefined; +let sharedClosingPromise: Promise | undefined; let sharedWebSocketRefCount = 0; const sharedPendingCalls: Map< number, @@ -66,7 +76,7 @@ export function instantiate(): AppAgent { actionContext, ctx.choiceManager, () => ctx.webSocketServer?.isConnected() === true, - resolveCodePort(process.env), + getKnownCodePort(), ); }, handleChoice: async (choiceId, response, context) => { @@ -90,6 +100,11 @@ type CodeActionContext = { // Manages yes/no choice callbacks (currently only the setup-flow card). // Hooked up via the AppAgent.handleChoice in instantiate() above. choiceManager: ChoiceManager; + // Handle returned by sessionContext.registerPort, kept so we can release + // exactly this session's registration on disable. The + // closeSessionContext backstop will also release it if the disable path + // is skipped. + portRegistration?: { release: () => void }; }; async function initializeCodeContext(): Promise { @@ -126,10 +141,101 @@ async function checkCodeReadiness( return evaluateCodeReadiness({ clientConnected, vsCodeCliInstalled, - port: resolveCodePort(process.env), + port: getKnownCodePort(), }); } +// Returns the port the code agent's WS server is/will be reachable on, +// for display in readiness/setup messaging. Two phases: +// - After bind: `getSharedCodePort()` returns the actual bound port +// (OS-assigned by default, or `CODE_WEBSOCKET_PORT` if set; either way, +// this is the authoritative answer). +// - Before bind: no live port exists, so we fall back to the static +// prediction from `CODE_WEBSOCKET_PORT` if set, else `undefined` +// (the UI shows "port unknown until bind"). +function getKnownCodePort(): number | undefined { + return getSharedCodePort() ?? resolveCodePortOverride(process.env); +} + +// Bind hint for the shared server. Returns the explicit override if +// CODE_WEBSOCKET_PORT is set (useful for pinning the port when debugging); +// otherwise 0 so the OS picks a free port and the registrar/discovery +// channel publishes it. +// +// Note: we only validate the *shape* of the env var here (numeric, >= 0). +// If the caller asks for a specific port and the OS can't bind it +// (EADDRINUSE), `CodeAgentWebSocketServer.start()` rejects with that error +// and the schema-enable fails loudly — we deliberately do NOT silently +// fall back to an OS-assigned port, since the user explicitly asked for +// a specific one. +function getCodeBindPort(): number { + const raw = process.env["CODE_WEBSOCKET_PORT"]; + if (raw === undefined) return 0; + const n = parseInt(raw, 10); + if (!Number.isFinite(n) || n < 0) { + debug( + `Ignoring malformed CODE_WEBSOCKET_PORT=${raw}; using OS-assigned port instead`, + ); + return 0; + } + return n; +} + +// Wire the shared server's onMessage handler. Module-scoped because the +// server itself is module-scoped — all sessions route their pending-call +// completions through the same handler. +function attachSharedOnMessage(server: CodeAgentWebSocketServer): void { + server.onMessage = (message: string) => { + try { + const data = JSON.parse(message) as WebSocketMessageV2; + + if (data.id !== undefined && data.result !== undefined) { + const pendingCall = sharedPendingCalls.get(Number(data.id)); + + if (pendingCall) { + sharedPendingCalls.delete(Number(data.id)); + const { resolve, context } = pendingCall; + if (context?.actionIO) { + context.actionIO.setDisplay(data.result); + } + resolve(); + } + } + } catch (error) { + debug("Error parsing WebSocket message:", error); + } + }; +} + +// Start (or attach to an in-flight start of) the shared WebSocket server. +// Concurrent enables from different sessions can race; serialize via +// sharedStartingPromise so only one bind attempt is in flight. +async function ensureSharedServer(): Promise { + // If a previous teardown is still releasing the port (matters under + // CODE_WEBSOCKET_PORT override), await it before binding again. + if (sharedClosingPromise !== undefined) { + await sharedClosingPromise; + } + if (sharedWebSocketServer !== undefined) { + return sharedWebSocketServer; + } + if (sharedStartingPromise !== undefined) { + return sharedStartingPromise; + } + sharedStartingPromise = (async () => { + try { + const server = + await CodeAgentWebSocketServer.start(getCodeBindPort()); + attachSharedOnMessage(server); + sharedWebSocketServer = server; + return server; + } finally { + sharedStartingPromise = undefined; + } + })(); + return sharedStartingPromise; +} + async function updateCodeContext( enable: boolean, context: SessionContext, @@ -140,43 +246,33 @@ async function updateCodeContext( if (agentContext.enabled.has(schemaName)) { return; } - if (agentContext.enabled.size === 0) { - sharedWebSocketRefCount++; - } + const isFirstSchemaForSession = agentContext.enabled.size === 0; agentContext.enabled.add(schemaName); - - if (!sharedWebSocketServer) { - // TODO: stop hardcoding the port. The dispatcher should hand each - // agent a free port at initialize time so multiple TypeAgent - // installs / sessions on one host can't collide on 8082. - const port = parseInt(process.env["CODE_WEBSOCKET_PORT"] || "8082"); - sharedWebSocketServer = new CodeAgentWebSocketServer(port); - - sharedWebSocketServer.onMessage = (message: string) => { - try { - const data = JSON.parse(message) as WebSocketMessageV2; - - if (data.id !== undefined && data.result !== undefined) { - const pendingCall = sharedPendingCalls.get( - Number(data.id), - ); - - if (pendingCall) { - sharedPendingCalls.delete(Number(data.id)); - const { resolve, context } = pendingCall; - if (context?.actionIO) { - context.actionIO.setDisplay(data.result); - } - resolve(); - } - } - } catch (error) { - debug("Error parsing WebSocket message:", error); - } - }; + try { + const server = await ensureSharedServer(); + agentContext.webSocketServer = server; + agentContext.pendingCall = sharedPendingCalls; + if (isFirstSchemaForSession) { + // Per-session registration: the registrar allows multiple + // entries for `(code, default)` across sessions and lookup + // returns the most recent, so each active session + // independently keeps the shared port discoverable. The + // backstop in closeSessionContext releases ours if disable + // is skipped. + agentContext.portRegistration = context.registerPort( + "default", + server.port, + ); + sharedWebSocketRefCount++; + } + } catch (e) { + // Roll back the per-session schema bookkeeping so a subsequent + // retry sees a clean slate. Don't touch shared module state — + // the bind itself failed, so we never incremented the refcount + // or registered. + agentContext.enabled.delete(schemaName); + throw e; } - agentContext.webSocketServer = sharedWebSocketServer; - agentContext.pendingCall = sharedPendingCalls; } else { if (!agentContext.enabled.has(schemaName)) { return; @@ -184,16 +280,35 @@ async function updateCodeContext( agentContext.enabled.delete(schemaName); if (agentContext.enabled.size === 0) { agentContext.webSocketServer = undefined; + // Release this session's registration before potentially closing + // the server. Release is idempotent and a no-op if already + // released by the backstop. + agentContext.portRegistration?.release(); + delete agentContext.portRegistration; + sharedWebSocketRefCount = Math.max(0, sharedWebSocketRefCount - 1); if (sharedWebSocketRefCount === 0 && sharedWebSocketServer) { - sharedWebSocketServer.close(); + const server = sharedWebSocketServer; sharedWebSocketServer = undefined; sharedPendingCalls.clear(); + // Track the in-flight close so a rapid re-enable awaits + // port release under a fixed-port override. + sharedClosingPromise = server.close().finally(() => { + sharedClosingPromise = undefined; + }); + await sharedClosingPromise; } } } } +// Exposed for readiness/setup messaging — undefined when the shared server +// isn't bound yet, otherwise the actual bound port. Lets readiness messages +// always reflect the real listener. +export function getSharedCodePort(): number | undefined { + return sharedWebSocketServer?.port; +} + function getVSCodeStoragePath(): string { const platform = os.platform(); if (platform === "darwin") diff --git a/ts/packages/agents/code/src/codeAgentWebSocketServer.ts b/ts/packages/agents/code/src/codeAgentWebSocketServer.ts index 4cbbe44e5..221bda195 100644 --- a/ts/packages/agents/code/src/codeAgentWebSocketServer.ts +++ b/ts/packages/agents/code/src/codeAgentWebSocketServer.ts @@ -2,23 +2,81 @@ // Licensed under the MIT License. import { WebSocketServer, WebSocket } from "ws"; +import { AddressInfo } from "net"; import registerDebug from "debug"; const debug = registerDebug("typeagent:code:websocket"); export class CodeAgentWebSocketServer { - private server: WebSocketServer; private clients: Map = new Map(); private clientIdCounter = 0; public onMessage?: (message: string) => void; - constructor(port: number = 8082) { - this.server = new WebSocketServer({ port }); + /** + * @param server the underlying ws server, already bound and listening. + * @param port the actually bound port (OS-assigned when the caller + * passed 0). + * + * Construction is private — use {@link CodeAgentWebSocketServer.start} + * so callers always get a server that is guaranteed to be bound + * before they read {@link port} or pass it to the registrar. + */ + private constructor( + private readonly server: WebSocketServer, + public readonly port: number, + ) { this.setupHandlers(); debug(`CodeAgentWebSocketServer listening on port ${port}`); + } - this.server.on("error", (error) => { - debug("Server error:", error); + /** + * Bind a new server on `port`. Resolves only after the + * `listening` event so callers can synchronously read + * {@link port}; rejects on the first `error` event so bind + * failures (EADDRINUSE under fixed-port overrides) surface + * loudly instead of being swallowed by an attached error + * handler. + * + * Pass `0` to let the OS pick a free ephemeral port; the + * actual port is then available via {@link port}. + */ + public static start(port: number = 0): Promise { + return new Promise((resolve, reject) => { + const server = new WebSocketServer({ port }); + let settled = false; + const onError = (error: Error) => { + if (settled) { + debug("Server error after listening:", error); + return; + } + settled = true; + server.removeListener("listening", onListening); + debug("Server bind error:", error); + reject(error); + }; + const onListening = () => { + if (settled) return; + settled = true; + server.removeListener("error", onError); + const address = server.address() as AddressInfo | null; + if (!address || typeof address === "string") { + server.close(); + reject( + new Error( + "ws server.address() did not return an AddressInfo", + ), + ); + return; + } + // Re-attach a permanent error handler so post-listen errors + // are logged rather than crashing the process. + server.on("error", (error) => { + debug("Server error:", error); + }); + resolve(new CodeAgentWebSocketServer(server, address.port)); + }; + server.once("error", onError); + server.once("listening", onListening); }); } @@ -48,10 +106,6 @@ export class CodeAgentWebSocketServer { this.clients.delete(clientId); }); }); - - this.server.on("error", (error) => { - debug("Server error:", error); - }); } public broadcast(message: string): number { @@ -117,7 +171,14 @@ export class CodeAgentWebSocketServer { return states; } - public close(): void { + /** + * Close all client connections and the underlying server. + * Resolves when the server has fully released its port — important + * for a rapid disable→enable cycle under a fixed-port override + * (`CODE_WEBSOCKET_PORT`), where a synchronous return would race + * the new bind into EADDRINUSE. + */ + public close(): Promise { debug("Closing CodeAgentWebSocketServer"); for (const [, client] of this.clients.entries()) { if (client.readyState === WebSocket.OPEN) { @@ -125,6 +186,8 @@ export class CodeAgentWebSocketServer { } } this.clients.clear(); - this.server.close(); + return new Promise((resolve) => { + this.server.close(() => resolve()); + }); } } diff --git a/ts/packages/agents/code/src/readiness.ts b/ts/packages/agents/code/src/readiness.ts index 620ea0d41..2395c634a 100644 --- a/ts/packages/agents/code/src/readiness.ts +++ b/ts/packages/agents/code/src/readiness.ts @@ -9,6 +9,7 @@ // server — not that we can dial out to a port. import { spawn } from "child_process"; +import registerDebug from "debug"; import { ActionContext, ActionResult, @@ -21,6 +22,8 @@ import { createYesNoChoiceResult, } from "@typeagent/agent-sdk/helpers/action"; +const debug = registerDebug("typeagent:code"); + // HH:MM timestamp prefix for status updates — same convention used by // screencapture / desktop / calendar so progress reads consistently. function ts(): string { @@ -41,8 +44,12 @@ export type CodeReadinessProbe = { // user has set everything up correctly and just closed VS Code is // misleading. vsCodeCliInstalled: boolean; - // For messaging only. Pulled from CODE_WEBSOCKET_PORT (default 8082). - port: number; + // For messaging only. Undefined until the shared server is bound; the + // user-facing message then omits the port number. Pulled from the + // actual bound port (via getSharedCodePort()) so readiness stays in + // sync with the registered listener, including when the OS picked a + // free ephemeral port (port=0 default). + port: number | undefined; }; // Pure decision function — exported for unit tests so we don't have to @@ -77,9 +84,10 @@ export function evaluateCodeReadiness( // This is a transient runtime state, not a config problem; `setup` // resolves it by launching VS Code and waiting for the extension to // connect. + const portSuffix = probe.port !== undefined ? ` on port ${probe.port}` : ""; return { state: "setup-required", - message: `VS Code isn't currently running (or the Coda extension isn't connected on port ${probe.port}).`, + message: `VS Code isn't currently running (or the Coda extension isn't connected${portSuffix}).`, details: "Run `@config agent setup code` to launch VS Code — the Coda extension will auto-connect — or open VS Code yourself and your code commands will start working immediately.", }; @@ -102,16 +110,30 @@ export async function whichExists(tool: string): Promise { }); } -// Resolves the configured port the same way updateAgentContext does. -// Centralized here so readiness messaging stays in sync with the actual -// listener configuration. -export function resolveCodePort(env: NodeJS.ProcessEnv): number { +// Resolves the explicit port override (CODE_WEBSOCKET_PORT) for readiness +// messaging when the shared server isn't bound yet. Returns undefined to +// signal "no static port" — the caller (readiness probe) should query +// `getSharedCodePort()` first and fall through here only when the server +// hasn't been started yet. +// +// Validation matches getCodeBindPort() in codeActionHandler.ts. +export function resolveCodePortOverride( + env: NodeJS.ProcessEnv, +): number | undefined { const raw = env.CODE_WEBSOCKET_PORT; - if (raw !== undefined) { - const n = parseInt(raw); - if (Number.isFinite(n) && n > 0) return n; + if (raw === undefined) return undefined; + const n = parseInt(raw, 10); + if (Number.isFinite(n) && n > 0) { + debug( + `CODE_WEBSOCKET_PORT override active: using port ${n} (set in environment)`, + ); + return n; } - return 8082; + debug( + `CODE_WEBSOCKET_PORT override ignored: invalid value %o; falling back to OS-assigned port`, + raw, + ); + return undefined; } // ============================================================================ @@ -169,7 +191,7 @@ export async function setupCode( actionContext: ActionContext, choiceManager: ChoiceManager, isConnected: () => boolean, - port: number, + port: number | undefined, ): Promise { if (isConnected()) { // Defense-in-depth: dispatcher only invokes setup() when readiness @@ -181,9 +203,10 @@ export async function setupCode( "VS Code is already connected.", ); } + const portSuffix = port !== undefined ? ` on port ${port}` : ""; return createYesNoChoiceResult( choiceManager, - `Launch VS Code? The Coda extension will auto-connect to the code agent's WebSocket server on port ${port}. I'll wait up to 30 seconds for the connection — you can keep working in the meantime.`, + `Launch VS Code? The Coda extension will auto-connect to the code agent's WebSocket server${portSuffix}. I'll wait up to 30 seconds for the connection — you can keep working in the meantime.`, async (confirmed, liveActionContext) => { if (!confirmed) { return createActionResultFromTextDisplay( @@ -204,7 +227,7 @@ export async function setupCode( export async function runLaunchAndWait( actionContext: ActionContext, isConnected: () => boolean, - port: number, + port: number | undefined, options?: { timeoutMs?: number; launchImpl?: () => Promise; @@ -229,10 +252,11 @@ export async function runLaunchAndWait( ); } + const waitOn = port !== undefined ? ` on port ${port}` : ""; actionContext.actionIO.appendDisplay( { type: "text", - content: `[${ts()}] Waiting for the Coda extension to connect on port ${port} (up to ${Math.round(timeoutMs / 1000)}s)…`, + content: `[${ts()}] Waiting for the Coda extension to connect${waitOn} (up to ${Math.round(timeoutMs / 1000)}s)…`, kind: "status", }, "block", diff --git a/ts/packages/agents/code/test/codeReadiness.spec.ts b/ts/packages/agents/code/test/codeReadiness.spec.ts index 32e0e7559..27e27d673 100644 --- a/ts/packages/agents/code/test/codeReadiness.spec.ts +++ b/ts/packages/agents/code/test/codeReadiness.spec.ts @@ -5,14 +5,14 @@ * Tests for the code agent's readiness/setup wiring. * * Pure functions only — `evaluateCodeReadiness` (decision) and - * `resolveCodePort` (env parse) — plus `runLaunchAndWait` exercised with - * an injected `launchImpl` so we don't actually spawn `code --new-window` - * during tests. + * `resolveCodePortOverride` (env parse) — plus `runLaunchAndWait` + * exercised with an injected `launchImpl` so we don't actually spawn + * `code --new-window` during tests. */ import { evaluateCodeReadiness, - resolveCodePort, + resolveCodePortOverride, runLaunchAndWait, } from "../src/readiness.js"; import type { ActionContext } from "@typeagent/agent-sdk"; @@ -75,43 +75,58 @@ describe("evaluateCodeReadiness", () => { }); expect(r.message).toMatch(/port 9999/); }); + + test("port omitted from message when undefined (server not yet bound)", () => { + // Initial readiness probe can fire before updateAgentContext has + // bound the shared server. Messaging must still be coherent — no + // dangling "port undefined" string. + const r = evaluateCodeReadiness({ + clientConnected: false, + vsCodeCliInstalled: true, + port: undefined, + }); + expect(r.message).not.toMatch(/port/i); + expect(r.message).not.toMatch(/undefined/); + }); }); -describe("resolveCodePort", () => { - test("defaults to 8082 when CODE_WEBSOCKET_PORT is unset", () => { - expect(resolveCodePort({} as NodeJS.ProcessEnv)).toBe(8082); +describe("resolveCodePortOverride", () => { + test("returns undefined when CODE_WEBSOCKET_PORT is unset (dynamic port)", () => { + expect( + resolveCodePortOverride({} as NodeJS.ProcessEnv), + ).toBeUndefined(); }); test("parses a valid integer override", () => { expect( - resolveCodePort({ + resolveCodePortOverride({ CODE_WEBSOCKET_PORT: "9100", } as NodeJS.ProcessEnv), ).toBe(9100); }); - test("falls back to default when CODE_WEBSOCKET_PORT is non-numeric", () => { - // Defensive: parseInt on garbage returns NaN. Without the - // Number.isFinite guard we'd return NaN as the port and the - // listener would explode. + test("returns undefined when CODE_WEBSOCKET_PORT is non-numeric", () => { + // Defensive: parseInt on garbage returns NaN. We must not propagate + // NaN as a port — return undefined so the agent falls back to + // OS-assigned binding. expect( - resolveCodePort({ + resolveCodePortOverride({ CODE_WEBSOCKET_PORT: "abc", } as NodeJS.ProcessEnv), - ).toBe(8082); + ).toBeUndefined(); }); - test("falls back to default for zero / negative values", () => { + test("returns undefined for zero / negative values", () => { expect( - resolveCodePort({ + resolveCodePortOverride({ CODE_WEBSOCKET_PORT: "0", } as NodeJS.ProcessEnv), - ).toBe(8082); + ).toBeUndefined(); expect( - resolveCodePort({ + resolveCodePortOverride({ CODE_WEBSOCKET_PORT: "-1", } as NodeJS.ProcessEnv), - ).toBe(8082); + ).toBeUndefined(); }); }); diff --git a/ts/packages/agents/code/test/codeUpdateContext.spec.ts b/ts/packages/agents/code/test/codeUpdateContext.spec.ts new file mode 100644 index 000000000..56d0e69f7 --- /dev/null +++ b/ts/packages/agents/code/test/codeUpdateContext.spec.ts @@ -0,0 +1,219 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +/** + * Integration-style tests for the code agent's shared-server lifecycle + * and per-session port registration. These tests actually bind a + * WebSocket server on a random free port (port=0) per the design — they + * exercise the real `CodeAgentWebSocketServer.start` + `close` paths + * plus the per-session `registerPort` accounting end-to-end. + * + * Coverage matrix (per the PR 2 plan's test list): + * - enable A: shared server bound, A's registerPort called. + * - enable A + B: both sessions register, lookup returns same port. + * - A disables: A's registration released; server still up; B's + * registration intact. + * - both disable: server closed, no live registrations. + * - non-owner disable: shared server stays up under the other session. + * - skipping disable (backstop simulation): manually release A's + * registration; B still works. + * + * The shared server is module-scoped state, so tests run serially via + * Jest's default sequential mode for a single suite. + */ + +import { instantiate, getSharedCodePort } from "../src/codeActionHandler.js"; +import type { SessionContext } from "@typeagent/agent-sdk"; + +// Minimal stub of SessionContext — only the surface the code agent uses +// during updateAgentContext. registerPort records every call so we can +// assert per-session accounting; the returned handle's release is real +// (no-op default; tests override when they want to inspect release). +type RegisterCall = { role: string; port: number }; +type StubContext = { + sessionContext: SessionContext; + registerCalls: RegisterCall[]; + releaseSpy: () => void; + releaseCount: number; +}; + +function makeStubContext(agentContext: any): StubContext { + const registerCalls: RegisterCall[] = []; + let releaseCount = 0; + const releaseSpy = () => { + releaseCount++; + }; + const sessionContext = { + agentContext, + registerPort(role: string, port: number) { + registerCalls.push({ role, port }); + return { release: releaseSpy }; + }, + // The rest of SessionContext isn't touched by updateCodeContext. + } as unknown as SessionContext; + return { + sessionContext, + registerCalls, + get releaseCount() { + return releaseCount; + }, + releaseSpy, + } as any; +} + +describe("code agent shared server + per-session registration", () => { + const agent = instantiate(); + // Track all sessions created during a test so afterEach can defensively + // tear down anything a test left enabled (e.g., on assertion failure + // before the test's own try/finally cleanup runs). Without this, a + // surviving shared server contaminates later tests. + const liveSessions = new Set(); + + async function newSession(): Promise { + const agentContext = await agent.initializeAgentContext!(); + const ctx = makeStubContext(agentContext); + liveSessions.add(ctx); + return ctx; + } + + afterEach(async () => { + // Disable any session still enabled from the test body, then drop + // it from the live set. updateAgentContext(false, ...) is idempotent + // and refcount-safe, so it's fine to call on sessions whose + // try/finally already ran. + for (const ctx of liveSessions) { + try { + await agent.updateAgentContext!( + false, + ctx.sessionContext, + "code", + ); + } catch { + // Best-effort cleanup; swallow so other sessions still tear down. + } + } + liveSessions.clear(); + // Sanity: the shared server should now be down. If a test left it up + // for a reason we don't model, surface it loudly so we don't silently + // leak across tests. + if (getSharedCodePort() !== undefined) { + throw new Error( + "shared WebSocket server still bound after afterEach cleanup", + ); + } + }); + + test("enable on a fresh session binds the shared server and registers under (code, default)", async () => { + const s = await newSession(); + await agent.updateAgentContext!(true, s.sessionContext, "code"); + try { + // Bind happened — getSharedCodePort returns the bound port. + const port = getSharedCodePort(); + expect(port).toBeDefined(); + expect(port).toBeGreaterThan(0); + // Exactly one register call for the session's first schema. + expect(s.registerCalls).toEqual([{ role: "default", port }]); + // Web socket server is wired into the per-session agentContext. + expect(s.sessionContext.agentContext.webSocketServer).toBeDefined(); + } finally { + await agent.updateAgentContext!(false, s.sessionContext, "code"); + // After last session disables, shared server is torn down. + expect(getSharedCodePort()).toBeUndefined(); + } + }); + + test("two sessions enabling produce two registrations on the SAME port", async () => { + const a = await newSession(); + const b = await newSession(); + await agent.updateAgentContext!(true, a.sessionContext, "code"); + const portAfterA = getSharedCodePort(); + await agent.updateAgentContext!(true, b.sessionContext, "code"); + const portAfterB = getSharedCodePort(); + try { + expect(portAfterA).toBeDefined(); + expect(portAfterB).toBe(portAfterA); + expect(a.registerCalls).toEqual([ + { role: "default", port: portAfterA }, + ]); + expect(b.registerCalls).toEqual([ + { role: "default", port: portAfterA }, + ]); + } finally { + await agent.updateAgentContext!(false, a.sessionContext, "code"); + await agent.updateAgentContext!(false, b.sessionContext, "code"); + } + }); + + test("disabling one session releases only its own registration; server stays up under the other", async () => { + const a = await newSession(); + const b = await newSession(); + await agent.updateAgentContext!(true, a.sessionContext, "code"); + await agent.updateAgentContext!(true, b.sessionContext, "code"); + const port = getSharedCodePort(); + expect(port).toBeDefined(); + + // A disables — its release runs; B is untouched. + await agent.updateAgentContext!(false, a.sessionContext, "code"); + expect((a as any).releaseCount).toBe(1); + expect((b as any).releaseCount).toBe(0); + // Shared server still bound to the same port (B keeps it alive). + expect(getSharedCodePort()).toBe(port); + + // B disables — server closes. + await agent.updateAgentContext!(false, b.sessionContext, "code"); + expect((b as any).releaseCount).toBe(1); + expect(getSharedCodePort()).toBeUndefined(); + }); + + test("multiple schemas on one session only register once and release once", async () => { + const s = await newSession(); + // updateCodeContext is called per schema; with refcount + // bookkeeping based on `enabled.size === 0`, the registration is + // only created on the FIRST schema and released on the LAST. + await agent.updateAgentContext!(true, s.sessionContext, "code"); + await agent.updateAgentContext!(true, s.sessionContext, "code-editor"); + try { + expect(s.registerCalls).toHaveLength(1); + } finally { + // Disable one schema — server stays up because the session + // still has another schema enabled. + await agent.updateAgentContext!( + false, + s.sessionContext, + "code-editor", + ); + expect((s as any).releaseCount).toBe(0); + expect(getSharedCodePort()).toBeDefined(); + // Disable the last schema — release fires, server closes. + await agent.updateAgentContext!(false, s.sessionContext, "code"); + expect((s as any).releaseCount).toBe(1); + expect(getSharedCodePort()).toBeUndefined(); + } + }); + + test("backstop simulation: releasing A's handle externally still leaves B's registration intact", async () => { + // Simulates closeSessionContext's finally backstop running for A + // (which calls portRegistrar.releaseAllForSession bypassing the + // agent's updateAgentContext(false, ...) path) while B remains + // active. The agent must tolerate this — its own state assumes + // the registration may already be released. + const a = await newSession(); + const b = await newSession(); + await agent.updateAgentContext!(true, a.sessionContext, "code"); + await agent.updateAgentContext!(true, b.sessionContext, "code"); + const port = getSharedCodePort(); + + // Pretend the backstop released A's handle directly. + a.sessionContext.agentContext.portRegistration?.release(); + + // Now drive A's disable explicitly. The release call inside + // updateCodeContext should be a no-op (idempotent stub), + // bookkeeping must still complete cleanly. + await agent.updateAgentContext!(false, a.sessionContext, "code"); + // Server still up under B. + expect(getSharedCodePort()).toBe(port); + + await agent.updateAgentContext!(false, b.sessionContext, "code"); + expect(getSharedCodePort()).toBeUndefined(); + }); +}); diff --git a/ts/packages/coda/package.json b/ts/packages/coda/package.json index 5f75cf078..a460aa05c 100644 --- a/ts/packages/coda/package.json +++ b/ts/packages/coda/package.json @@ -50,6 +50,8 @@ "onStartupFinished" ], "dependencies": { + "@typeagent/agent-server-client": "workspace:*", + "@typeagent/agent-server-protocol": "workspace:*", "body-parser": "^1.20.3", "chalk": "^5.4.1", "cors": "^2.8.5", diff --git a/ts/packages/coda/src/webSocket.ts b/ts/packages/coda/src/webSocket.ts index a97ef1db4..ba8e488ed 100644 --- a/ts/packages/coda/src/webSocket.ts +++ b/ts/packages/coda/src/webSocket.ts @@ -2,6 +2,7 @@ // Licensed under the MIT License. import WebSocket from "ws"; +import { discoverPort } from "@typeagent/agent-server-client/discovery"; export type WebSocketMessage = { source: string; @@ -11,31 +12,80 @@ export type WebSocketMessage = { body: any; }; +// One-shot resolution of the code agent's WebSocket endpoint via the +// agent-server's discovery channel. Returns undefined if the endpoint +// can't be resolved (discovery unreachable, or `code` agent not yet +// enabled in any session) — the caller's reconnect loop will retry. +// Honors: +// - `CODE_WEBSOCKET_HOST` env: explicit override; if set, skips +// discovery and dials the given URL directly. +// - `AGENT_SERVER_URL` env: where to reach the agent-server's +// discovery WS (defaults to ws://localhost:8999, the +// AGENT_SERVER_DEFAULT_URL constant). +async function resolveCodeEndpoint(): Promise { + const explicit = process.env["CODE_WEBSOCKET_HOST"]; + if (explicit) return explicit; + + const result = await discoverPort("code", undefined, { + url: process.env["AGENT_SERVER_URL"], + }); + if (result.kind === "found") { + return `ws://localhost:${result.port}`; + } + if (result.kind === "not-registered") { + // Agent server is running but the `code` agent isn't enabled + // (no one has activated the `code` schema in any session yet). + // Transient — the reconnect loop will re-discover periodically. + console.log( + "code agent not yet registered with agent-server; will retry.", + ); + return undefined; + } + // Discovery unreachable (agentServer not running). Reconnect loop + // will retry. + console.log( + `Discovery channel unreachable (${result.error.message}); will retry.`, + ); + return undefined; +} + export async function createWebSocket( channel: string, role: string, clientId?: string, -) { - return new Promise((resolve, reject) => { - let endpoint = - process.env["CODE_WEBSOCKET_HOST"] ?? "ws://localhost:8082"; - endpoint += `?channel=${channel}&role=${role}`; +): Promise { + let endpoint: string; + try { + const base = await resolveCodeEndpoint(); + if (!base) return undefined; + endpoint = `${base}?channel=${channel}&role=${role}`; if (clientId) { - endpoint += `clientId=${clientId}`; + endpoint += `&clientId=${clientId}`; } - - const webSocket = new WebSocket(endpoint); - - webSocket.onopen = (event: object) => { + } catch (error) { + console.error("Error resolving code agent endpoint:", error); + return undefined; + } + return new Promise((resolve) => { + let webSocket: WebSocket; + try { + webSocket = new WebSocket(endpoint); + } catch (error) { + // e.g. invalid URL from CODE_WEBSOCKET_HOST override. + console.error("Error constructing WebSocket:", error); + resolve(undefined); + return; + } + webSocket.onopen = () => { console.log("websocket open"); resolve(webSocket); }; - webSocket.onmessage = (event: object) => {}; - webSocket.onclose = (event: object) => { + webSocket.onmessage = () => {}; + webSocket.onclose = () => { console.log("websocket connection closed"); resolve(undefined); }; - webSocket.onerror = (event: object) => { + webSocket.onerror = () => { console.error("websocket error"); resolve(undefined); }; diff --git a/ts/packages/coda/src/wsConnect.ts b/ts/packages/coda/src/wsConnect.ts index 2dc4e3497..51602a919 100644 --- a/ts/packages/coda/src/wsConnect.ts +++ b/ts/packages/coda/src/wsConnect.ts @@ -112,13 +112,27 @@ function arrayBufferToJson(arrayBuffer: any) { } } +// Single in-flight reconnect interval. Pre-discovery code path spawned +// a new setInterval each time reconnect was called (from initializeWS +// and from every onclose); on rapid disconnects that leaked one +// interval per disconnect, all racing to call ensureWebsocketConnected. +// Guard with a module-scoped handle so only one timer is alive at a time. +let reconnectIntervalId: ReturnType | undefined; + function reconnectWebSocket() { - const connectionCheckIntervalId = setInterval(async () => { + if (reconnectIntervalId !== undefined) { + return; + } + reconnectIntervalId = setInterval(async () => { if (webSocket && webSocket.readyState === WebSocket.OPEN) { console.log("Clearing reconnect retry interval"); - clearInterval(connectionCheckIntervalId); + clearInterval(reconnectIntervalId); + reconnectIntervalId = undefined; } else { console.log("Retrying connection"); + // Re-runs the discovery handshake — important because the + // code agent's port can change across agent-server + // restarts (now that it binds to an OS-assigned port). await ensureWebsocketConnected(); } }, 5 * 1000); diff --git a/ts/packages/commandExecutor/README.md b/ts/packages/commandExecutor/README.md index d6c176896..c8088118c 100644 --- a/ts/packages/commandExecutor/README.md +++ b/ts/packages/commandExecutor/README.md @@ -211,7 +211,7 @@ Command Executor MCP Server TypeAgent Dispatcher (WebSocket) ↓ ├─ TypeAgent Agents (Music, Lists, Calendar, etc.) - └─ Coda VSCode Extension (via WebSocket on port 8082) + └─ Coda VSCode Extension (via WebSocket; port discovered through agent-server) └─ VSCode APIs (theme, editor, files, terminal, etc.) ``` diff --git a/ts/packages/commandExecutor/VSCODE_CAPABILITIES.md b/ts/packages/commandExecutor/VSCODE_CAPABILITIES.md index ecb97378c..d50a6acfc 100644 --- a/ts/packages/commandExecutor/VSCODE_CAPABILITIES.md +++ b/ts/packages/commandExecutor/VSCODE_CAPABILITIES.md @@ -10,7 +10,7 @@ The Command Executor MCP server can control VSCode through the Coda extension. B ``` User → Claude Code → execute_command MCP tool → → TypeAgent Dispatcher → - → Coda Extension (WebSocket on port 8082) → + → Coda Extension (WebSocket; port discovered through agent-server) → → VSCode APIs ``` @@ -295,7 +295,7 @@ To use these VSCode capabilities: 2. **Coda Extension** must be installed and activated in VSCode: - Published as `aisystems.copilot-coda` - - Auto-connects to dispatcher on port 8082 + - Discovers the code agent's WebSocket port via the agent-server (default `ws://localhost:8999`) 3. **Command Executor MCP Server** configured in `.mcp.json`: ```json diff --git a/ts/packages/shell/demo/DEMO_SETUP.md b/ts/packages/shell/demo/DEMO_SETUP.md index 21ccc4268..6e5dc56f5 100644 --- a/ts/packages/shell/demo/DEMO_SETUP.md +++ b/ts/packages/shell/demo/DEMO_SETUP.md @@ -31,8 +31,9 @@ Reinstall vscode-shell after any changes to `ts/packages/vscode-shell` (includin ## 2. Processes that must be running before you start 1. **Agent server** — `pnpm --filter agent-server start` (listens on `ws://localhost:8999`). -2. **Code agent's WebSocket** on port `8082` — started automatically when the - `code` schema is enabled in your session, and consumed by the coda extension. +2. **Code agent's WebSocket** — started automatically when the `code` schema + is enabled in your session (binds an OS-assigned port). The coda extension + discovers it through the agent-server's discovery channel. 3. **VS Code** open, with the **TypeAgent activity-bar icon visible** (vscode-shell installed). Do **not** click the icon until Part 1 hands off. 4. **gh CLI authenticated** — `gh auth status` should be green. Set the default @@ -94,7 +95,7 @@ listener installed by the webview. | `show check runs for that PR` runs `gh run view ` and 404s | Reasoning maps "for that PR" to a workflow-run id, and there's no conversation memory of the prior PR | Use the explicit phrasing `show check runs for PR N [in OWNER/REPO]` (now backed by `gh pr checks`) | | `create scratch.ts` returns "Unknown action name: code.code-general" | Reasoning hallucinates a sub-schema name as an action name | Rephrase as `create a typescript file scratch.ts with a hello world function` (drop "called" and "new"); long-term tracked in plan.md | | `splitEditor` returns "Did not handle the action" | Installed coda is older than `ts/packages/coda` source | Rebuild & reinstall coda | -| `splitEditor` returns "No websocket connection" | Code agent WS server isn't up in this session | `@config schema code on` (then verify port 8082 is listening) | +| `splitEditor` returns "No websocket connection" | Code agent WS server isn't up in this session | `@config schema code on` (the agent-server's discovery channel will then publish the dynamic port to the coda extension) | | `create scratch.ts` returns "Integration ... not found" | Onboarding agent grabbed the request | `@config schema onboarding off` | | `change my color theme to ...` toggles the title-bar instead of the theme | Desktop agent grabbed the request because the prompt didn't say "vscode" | Use the exact wording in the demo file | | "remind me which PR we were just looking at" hallucinates PR #123 / `example/repo` | Reasoning has no conversation memory | `@config request reasoning copilot` | diff --git a/ts/pnpm-lock.yaml b/ts/pnpm-lock.yaml index 7f0819091..920f10e8f 100644 --- a/ts/pnpm-lock.yaml +++ b/ts/pnpm-lock.yaml @@ -3548,6 +3548,12 @@ importers: packages/coda: dependencies: + '@typeagent/agent-server-client': + specifier: workspace:* + version: link:../agentServer/client + '@typeagent/agent-server-protocol': + specifier: workspace:* + version: link:../agentServer/protocol body-parser: specifier: ^1.20.3 version: 1.20.3