From e7e8db7950ac1164a3d02a2ead73d673c83784cb Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 16:57:56 -0700 Subject: [PATCH 01/11] Add client-count column to @config ports + notifyClientCountChanged SDK Upgrade the existing `@config ports` command to surface PortRegistrar allocations directly (instead of the legacy `getLocalHostPort` shim) and add a Clients column that mirrors the emoji-rich formatting used by `@config agent`. Rows are grouped by (agent, role, port) and `:system` rows are tagged `[system]`. Empty registrar renders `No ports registered`. To drive the new column, add `SessionContext.notifyClientCountChanged(role, count)` to the agent SDK following the existing `notifyReadinessChanged` pattern. The implementation delegates to a new `PortRegistrar.setClientCount` / `getClientCount` pair scoped by `(agentName, role, sessionContextId)`. Writes with no live allocation are silently dropped (defends against late notifications after `releaseAllForSession`); `release` and `releaseAllForSession` both drop matching count entries. `getClientCount` returns `undefined` for `never reported` so the render can distinguish it from an explicit `0`. Wire the SDK method through: - agentRpc: caller + callee for `notifyClientCountChanged`. - browser agent: `AgentWebSocketServer` gains an `onClientCountChanged` callback that fires AFTER every sessionMap mutation (fixes an off-by-one had we reused the existing `onClientDisconnected` hook which fires pre-delete). - code agent: `CodeAgentWebSocketServer` exposes `onClientCountChanged`; `codeActionHandler` maintains a module-level `sharedActiveSessions` set so the single shared server fans the same global count out to every active session. Tests: - portRegistrar.spec: setClientCount round-trip, drop-on-no-allocation, accepts 0, rejects negative + non-integer, release + releaseAllForSession drop counts, per-session isolation. - sessionContext.spec: notifyClientCountChanged delegates with agent name + sessionContextId; swallows setClientCount errors. - codeUpdateContext.spec: stub SessionContext gets the new method. Build green for agent-dispatcher, browser-typeagent, code-agent. Tests: dispatcher 681/681, browser 255/255, code 18/18. Prettier + repo-policy-check clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ts/packages/agentRpc/src/client.ts | 8 + ts/packages/agentRpc/src/server.ts | 10 ++ ts/packages/agentRpc/src/types.ts | 5 + ts/packages/agentSdk/src/agentInterface.ts | 23 +++ ts/packages/agentServer/protocol/README.md | 2 + .../src/agent/agentWebSocketServer.mts | 28 ++++ .../src/agent/browserActionHandler.mts | 5 + .../browser/src/agent/websiteMemory.mts | 1 + .../agents/code/src/codeActionHandler.ts | 24 +++ .../code/src/codeAgentWebSocketServer.ts | 15 +- .../code/test/codeUpdateContext.spec.ts | 3 + .../dispatcher/src/context/portRegistrar.ts | 76 +++++++++ .../system/handlers/configCommandHandlers.ts | 156 ++++++++++++++++-- .../dispatcher/src/execute/sessionContext.ts | 24 +++ .../dispatcher/test/portRegistrar.spec.ts | 66 ++++++++ .../dispatcher/test/sessionContext.spec.ts | 74 +++++++++ 16 files changed, 504 insertions(+), 16 deletions(-) diff --git a/ts/packages/agentRpc/src/client.ts b/ts/packages/agentRpc/src/client.ts index be63a7d46..92f486b21 100644 --- a/ts/packages/agentRpc/src/client.ts +++ b/ts/packages/agentRpc/src/client.ts @@ -385,6 +385,14 @@ export async function createAgentRpcClient( const context = contextMap.get(param.contextId); return context.notifyReadinessChanged(); }, + notifyClientCountChanged: async (param: { + contextId: number; + role: string; + count: number; + }) => { + const context = contextMap.get(param.contextId); + return context.notifyClientCountChanged(param.role, param.count); + }, storageRead: async (param: { contextId: number; session: boolean; diff --git a/ts/packages/agentRpc/src/server.ts b/ts/packages/agentRpc/src/server.ts index ae011f572..76ed6b029 100644 --- a/ts/packages/agentRpc/src/server.ts +++ b/ts/packages/agentRpc/src/server.ts @@ -641,6 +641,16 @@ export function createAgentRpcServer( contextId, }); }, + notifyClientCountChanged: async ( + role: string, + count: number, + ): Promise => { + return rpc.invoke("notifyClientCountChanged", { + contextId, + role, + count, + }); + }, }; } diff --git a/ts/packages/agentRpc/src/types.ts b/ts/packages/agentRpc/src/types.ts index 9bac5a87b..9ca3d2dd3 100644 --- a/ts/packages/agentRpc/src/types.ts +++ b/ts/packages/agentRpc/src/types.ts @@ -160,6 +160,11 @@ export type AgentContextInvokeFunctions = { indexes: (param: { contextId: number; type: string }) => Promise; reloadAgentSchema: (param: { contextId: number }) => Promise; notifyReadinessChanged: (param: { contextId: number }) => Promise; + notifyClientCountChanged: (param: { + contextId: number; + role: string; + count: number; + }) => Promise; popupQuestion: (param: { contextId: number; message: string; diff --git a/ts/packages/agentSdk/src/agentInterface.ts b/ts/packages/agentSdk/src/agentInterface.ts index 652e6043b..246f3687c 100644 --- a/ts/packages/agentSdk/src/agentInterface.ts +++ b/ts/packages/agentSdk/src/agentInterface.ts @@ -323,6 +323,29 @@ export interface SessionContext { // surface in the trigger path (the next refresh will retry). notifyReadinessChanged(): Promise; + /** + * Notify the dispatcher that the number of clients currently + * connected to one of this agent's registered listeners has + * changed. Surfaced by the `@system ports` diagnostic command. + * + * `role` must match a role this agent previously passed to + * {@link registerPort}; writes for which no live registration + * exists are dropped silently (defends against late notifications + * arriving after the agent's session context has closed). + * + * `count` is the current total of connected clients for the + * `(agent, role, this session)` tuple, **after** the connect / + * disconnect that triggered this call has been applied to the + * agent's internal tracking. Pass `0` when the last client goes + * away so the column doesn't show a stale positive count. + * + * Best-effort: errors are swallowed (and debug-logged) so an + * event-driven trigger (e.g. WebSocket onClose) never throws into + * the event-emitter path. Agents that don't open ports never need + * to call this. + */ + notifyClientCountChanged(role: string, count: number): Promise; + /** * Register a port this agent has just bound (typically with * `bind(0)` so the OS picks a free ephemeral port). The dispatcher diff --git a/ts/packages/agentServer/protocol/README.md b/ts/packages/agentServer/protocol/README.md index 154e8015c..60db95d98 100644 --- a/ts/packages/agentServer/protocol/README.md +++ b/ts/packages/agentServer/protocol/README.md @@ -100,6 +100,8 @@ createDiscoveryHandlers((agentName, role) => Passing a lookup callback (rather than the `IPortRegistrar` itself) keeps this package free of an `agent-dispatcher` dependency. +Live port allocations (including any client counts agents have published via `SessionContext.notifyClientCountChanged`) are surfaced in-process by the dispatcher's `@config ports` command. + --- ## Trademarks diff --git a/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts b/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts index 39c159e45..71e2c5e7d 100644 --- a/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts +++ b/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts @@ -40,6 +40,16 @@ interface SessionHandlers { getPreferredClientType?: () => "extension" | "electron" | undefined; onClientConnected?: (client: BrowserClient) => void; onClientDisconnected?: (client: BrowserClient) => void; + /** + * Fired after the {@link clients} map mutation completes for any + * connect / disconnect affecting this session, with the post- + * mutation total of tracked clients for the session. Used by the + * agent to push counts up through `SessionContext.notifyClientCountChanged` + * so `@system ports` can surface them. Off-by-one safe: the new + * count is computed AFTER the connect / disconnect is reflected in + * the map. + */ + onClientCountChanged?: (count: number) => void; onWebAgentMessage?: (client: BrowserClient, data: any) => void; activeClientId: string | null; } @@ -190,6 +200,12 @@ export class AgentWebSocketServer { handlers.onClientConnected(client); } } + + // Push the initial count up now that the session knows + // about its pre-connected clients. + if (handlers.onClientCountChanged) { + handlers.onClientCountChanged(preConnected.size); + } } debug(`Session registered: ${sessionId}`); @@ -324,6 +340,11 @@ export class AgentWebSocketServer { session.onClientConnected(client); } + // Off-by-one safe: fired AFTER the sessionMap mutation above. + if (session?.onClientCountChanged) { + session.onClientCountChanged(sessionMap.size); + } + ws.on("message", (message: string) => { client.lastActivity = new Date(); @@ -368,11 +389,18 @@ export class AgentWebSocketServer { } const sm = this.clients.get(client.sessionId); + let postCount = 0; if (sm) { sm.delete(clientId); + postCount = sm.size; if (sm.size === 0) this.clients.delete(client.sessionId); } + // Off-by-one safe: fired AFTER the sessionMap delete. + if (s?.onClientCountChanged) { + s.onClientCountChanged(postCount); + } + if (s && s.activeClientId === clientId) { this.selectNewActiveClientForSession(client.sessionId); } diff --git a/ts/packages/agents/browser/src/agent/browserActionHandler.mts b/ts/packages/agents/browser/src/agent/browserActionHandler.mts index 5518d8056..1a412517e 100644 --- a/ts/packages/agents/browser/src/agent/browserActionHandler.mts +++ b/ts/packages/agents/browser/src/agent/browserActionHandler.mts @@ -678,6 +678,11 @@ async function updateBrowserContext( // timeout from a stale "ready" cache. void context.notifyReadinessChanged(); }, + onClientCountChanged: (count: number) => { + // Surface to `@system ports`. Best-effort; the SDK + // method swallows errors internally. + void context.notifyClientCountChanged("default", count); + }, onWebAgentMessage: async (client: BrowserClient, data: any) => { if ( data.method === "webAgent/message" && diff --git a/ts/packages/agents/browser/src/agent/websiteMemory.mts b/ts/packages/agents/browser/src/agent/websiteMemory.mts index eafd925c5..f3995c4e9 100644 --- a/ts/packages/agents/browser/src/agent/websiteMemory.mts +++ b/ts/packages/agents/browser/src/agent/websiteMemory.mts @@ -170,6 +170,7 @@ export async function resolveURLWithHistory( indexes: async () => [], reloadAgentSchema: async () => {}, notifyReadinessChanged: async () => {}, + notifyClientCountChanged: async () => {}, }; // Use searchWebMemories with URL resolution optimized parameters diff --git a/ts/packages/agents/code/src/codeActionHandler.ts b/ts/packages/agents/code/src/codeActionHandler.ts index e83d39e37..7ad176b2c 100644 --- a/ts/packages/agents/code/src/codeActionHandler.ts +++ b/ts/packages/agents/code/src/codeActionHandler.ts @@ -50,6 +50,13 @@ let sharedWebSocketServer: CodeAgentWebSocketServer | undefined; let sharedStartingPromise: Promise | undefined; let sharedClosingPromise: Promise | undefined; let sharedWebSocketRefCount = 0; +// Sessions currently sharing the bound port. The shared WS server has +// no per-session client tracking (code's WS protocol carries no session +// id), so a single global count is fanned out to every active session +// — each session's registrar entry surfaces the same global number via +// `@system ports`. Sessions are added on first-schema-enable and +// removed on last-schema-disable. +const sharedActiveSessions = new Set>(); const sharedPendingCalls: Map< number, { @@ -205,6 +212,14 @@ function attachSharedOnMessage(server: CodeAgentWebSocketServer): void { debug("Error parsing WebSocket message:", error); } }; + // Fan out client-count updates to every active session so each + // session's `@system ports` row reflects the (global) count. The + // SDK method swallows errors internally. + server.onClientCountChanged = (count: number) => { + for (const sc of sharedActiveSessions) { + void sc.notifyClientCountChanged("default", count); + } + }; } // Start (or attach to an in-flight start of) the shared WebSocket server. @@ -264,6 +279,14 @@ async function updateCodeContext( server.port, ); sharedWebSocketRefCount++; + sharedActiveSessions.add(context); + // Publish the current (global) count to this session's + // newly-registered row so `@system ports` doesn't show + // "?" until the next connect/disconnect event. + void context.notifyClientCountChanged( + "default", + server.getConnectedCount(), + ); } } catch (e) { // Roll back the per-session schema bookkeeping so a subsequent @@ -285,6 +308,7 @@ async function updateCodeContext( // released by the backstop. agentContext.portRegistration?.release(); delete agentContext.portRegistration; + sharedActiveSessions.delete(context); sharedWebSocketRefCount = Math.max(0, sharedWebSocketRefCount - 1); if (sharedWebSocketRefCount === 0 && sharedWebSocketServer) { diff --git a/ts/packages/agents/code/src/codeAgentWebSocketServer.ts b/ts/packages/agents/code/src/codeAgentWebSocketServer.ts index c31660ce0..4ebc21dc1 100644 --- a/ts/packages/agents/code/src/codeAgentWebSocketServer.ts +++ b/ts/packages/agents/code/src/codeAgentWebSocketServer.ts @@ -12,6 +12,15 @@ export class CodeAgentWebSocketServer { private clients: Map = new Map(); private clientIdCounter = 0; public onMessage?: (message: string) => void; + /** + * Fired after the {@link clients} map mutation completes for any + * connect / disconnect, with the post-mutation total. Used by the + * code agent to push counts up via `SessionContext.notifyClientCountChanged` + * so `@system ports` can surface them. NOTE: code's WS protocol + * has no session identity, so this count is global across every + * session sharing the bound port — not per-session. + */ + public onClientCountChanged?: (count: number) => void; /** * @param server the underlying ws server, already bound and listening. @@ -104,6 +113,7 @@ export class CodeAgentWebSocketServer { const clientId = `client-${++this.clientIdCounter}-${Date.now()}`; debug("New client connected"); this.clients.set(clientId, ws); + this.onClientCountChanged?.(this.clients.size); // Store client ID on the WebSocket for reference (ws as any).clientId = clientId; @@ -118,11 +128,14 @@ export class CodeAgentWebSocketServer { ws.on("close", () => { debug("Client disconnected"); this.clients.delete(clientId); + this.onClientCountChanged?.(this.clients.size); }); ws.on("error", (error) => { debug("Client error:", error); - this.clients.delete(clientId); + if (this.clients.delete(clientId)) { + this.onClientCountChanged?.(this.clients.size); + } }); }); } diff --git a/ts/packages/agents/code/test/codeUpdateContext.spec.ts b/ts/packages/agents/code/test/codeUpdateContext.spec.ts index 56d0e69f7..2f4d181f5 100644 --- a/ts/packages/agents/code/test/codeUpdateContext.spec.ts +++ b/ts/packages/agents/code/test/codeUpdateContext.spec.ts @@ -49,6 +49,9 @@ function makeStubContext(agentContext: any): StubContext { registerCalls.push({ role, port }); return { release: releaseSpy }; }, + notifyClientCountChanged(_role: string, _count: number) { + // no-op stub; tested elsewhere via registrar unit tests + }, // The rest of SessionContext isn't touched by updateCodeContext. } as unknown as SessionContext; return { diff --git a/ts/packages/dispatcher/dispatcher/src/context/portRegistrar.ts b/ts/packages/dispatcher/dispatcher/src/context/portRegistrar.ts index 50c5abf6f..3309ff9fc 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/portRegistrar.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/portRegistrar.ts @@ -64,6 +64,32 @@ export interface IPortRegistrar { lookup(agentName: string, role?: string): number | undefined; hasActiveAllocations(): boolean; list(): readonly Allocation[]; + /** + * Update the cached client count for the live allocation matching + * `(agentName, role, sessionContextId)`. Writes for which no live + * allocation exists are silently dropped (defends against late + * notifications arriving after a session closed its registrations). + * Updates are best-effort and informational only — the count is + * surfaced via `@system ports` and used nowhere else. + */ + setClientCount( + agentName: string, + role: string, + sessionContextId: string, + count: number, + ): void; + /** + * Read the most recently reported client count for + * `(agentName, role, sessionContextId)`. `undefined` means + * "no count ever reported" — distinct from `0`, which means an + * agent has explicitly reported that no clients are currently + * connected. + */ + getClientCount( + agentName: string, + role: string, + sessionContextId: string, + ): number | undefined; } /** @@ -101,6 +127,16 @@ export class PortRegistrar implements IPortRegistrar { */ private readonly tripleIndex = new Map(); + /** + * Cached client counts keyed by the same `(agentName, role, + * sessionContextId)` triple as {@link tripleIndex}. Populated by + * agents calling `SessionContext.notifyClientCountChanged`; + * consumed by the `@system ports` diagnostic. Cleared whenever + * the matching allocation is released so a stale count can't + * outlive its registration. + */ + private readonly clientCounts = new Map(); + /** * Record a port that an agent has just bound. Idempotent on the * `(agentName, role, sessionContextId)` triple: a second call with @@ -211,6 +247,7 @@ export class PortRegistrar implements IPortRegistrar { } this.allocations.delete(regId); this.tripleIndex.delete(keyOf(allocation)); + this.clientCounts.delete(keyOf(allocation)); debug( `release ${allocation.agentName}/${allocation.role} session=${allocation.sessionContextId} port=${allocation.port} regId=${regId}`, ); @@ -300,6 +337,45 @@ export class PortRegistrar implements IPortRegistrar { public list(): readonly Allocation[] { return Array.from(this.allocations.values()).map((a) => ({ ...a })); } + + public setClientCount( + agentName: string, + role: string, + sessionContextId: string, + count: number, + ): void { + if (!Number.isInteger(count) || count < 0) { + debugWarn( + `setClientCount ignored: invalid count ${count} for ${agentName}/${role} session=${sessionContextId}`, + ); + return; + } + const tripleKey = keyOf({ agentName, role, sessionContextId }); + if (!this.tripleIndex.has(tripleKey)) { + // Late notification after release, or agent published a + // count before registering its port. Either way, no live + // allocation to attach the count to — drop it. If the agent + // re-registers later it will publish a fresh count. + debugWarn( + `setClientCount ignored: no live allocation for ${agentName}/${role} session=${sessionContextId}`, + ); + return; + } + this.clientCounts.set(tripleKey, count); + debug( + `setClientCount ${agentName}/${role} session=${sessionContextId} count=${count}`, + ); + } + + public getClientCount( + agentName: string, + role: string, + sessionContextId: string, + ): number | undefined { + return this.clientCounts.get( + keyOf({ agentName, role, sessionContextId }), + ); + } } /** diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts index 588a579ee..3b4914c1f 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts @@ -46,6 +46,7 @@ import { displayWarn, } from "@typeagent/agent-sdk/helpers/display"; import { alwaysEnabledAgents } from "../../appAgentManager.js"; +import { SYSTEM_SESSION_CONTEXT_ID } from "../../portRegistrar.js"; import { getCacheFactory } from "../../../utils/cacheFactory.js"; import { resolveCommand } from "../../../command/command.js"; import { toggleActivityContext } from "../../../execute/activityContext.js"; @@ -992,30 +993,155 @@ class ConfigTranslationNumberOfInitialActionsCommandHandler } class ConfigPortsCommandHandler implements CommandHandler { - public readonly description = "Lists the ports assigned to agents."; + public readonly description = + "Lists ports registered by agents and the number of clients connected to each."; public readonly parameters = {}; - public async run(context: ActionContext) { - const ports: string[][] = [["", "Agent", "Port"]]; - const cmdContext = context.sessionContext.agentContext as any; - - cmdContext.agents.getAppAgentNames().forEach((name: string) => { - const port = cmdContext.agents.getLocalHostPort(name); - - if (port !== undefined) { - ports.push([ - cmdContext.agents.getAppAgentEmoji(name), - name, - port.toString(), - ]); + public async run(context: ActionContext) { + const cmdContext = context.sessionContext.agentContext; + const registrar = cmdContext.portRegistrar; + const agents = cmdContext.agents; + + // Group registrar entries by (agentName, role, port). Code's + // shared WS server has no per-session client identity, so its + // count is global and the same number is reported across every + // session that registered the shared port. Grouping collapses + // those duplicates into one row with the (shared) count. + type Row = { + agentName: string; + role: string; + port: number; + isSystem: boolean; + // undefined === "agent never reported a count" (distinct from 0) + clientCount: number | undefined; + }; + const grouped = new Map(); + for (const alloc of registrar.list()) { + const key = `${alloc.agentName}\u0000${alloc.role}\u0000${alloc.port}`; + const isSystem = + alloc.sessionContextId === SYSTEM_SESSION_CONTEXT_ID; + const cc = registrar.getClientCount( + alloc.agentName, + alloc.role, + alloc.sessionContextId, + ); + const existing = grouped.get(key); + if (existing === undefined) { + grouped.set(key, { + agentName: alloc.agentName, + role: alloc.role, + port: alloc.port, + isSystem, + clientCount: cc, + }); + } else { + // System tag is sticky: if any underlying allocation is + // system-owned, treat the group as system. + if (isSystem) existing.isSystem = true; + // Sum across sessions when ≥1 has reported a count. + // Leave undefined only if NO underlying allocation has + // ever reported a count. + if (cc !== undefined) { + existing.clientCount = (existing.clientCount ?? 0) + cc; + } } + } + + const rows = Array.from(grouped.values()).sort((a, b) => { + const an = a.agentName.localeCompare(b.agentName); + if (an !== 0) return an; + const ar = a.role.localeCompare(b.role); + if (ar !== 0) return ar; + return a.port - b.port; }); - displayResult(ports, context); + if (rows.length === 0) { + displayWarn("No ports registered", context); + return; + } + + // Plain-text (CLI / console) — fixed-width via chalk for alignment. + const text: string[][] = [["", "Agent", "Role", "Port", "Clients"]]; + for (const r of rows) { + const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; + const name = r.isSystem + ? `${r.agentName} ${chalk.gray("[system]")}` + : r.agentName; + const clients = + r.clientCount === undefined + ? chalk.gray("?") + : r.clientCount.toString(); + text.push([emoji, name, r.role, r.port.toString(), clients]); + } + + // Rich HTML for the shell. + const html = buildPortsHtml(rows, agents); + + context.actionIO.appendDisplay({ + type: "text", + content: text, + alternates: [{ type: "html", content: html }], + }); } } +function escapeHtml(s: string): string { + return s + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} + +function buildPortsHtml( + rows: ReadonlyArray<{ + agentName: string; + role: string; + port: number; + isSystem: boolean; + clientCount: number | undefined; + }>, + agents: { getAppAgentEmoji(name: string): string }, +): string { + const thStyle = `text-align:left;padding:4px 8px;font-weight:600;font-size:13px;color:#64748b;border-bottom:2px solid #e2e8f0`; + const headerCols = [ + `Agent`, + `Role`, + `Port`, + `Clients`, + ]; + + const tdBase = `padding:3px 12px;border-bottom:1px solid #f1f5f9;white-space:nowrap`; + const rowsHtml = rows + .map((r) => { + const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; + const systemTag = r.isSystem + ? ` [system]` + : ""; + const clientCell = + r.clientCount === undefined + ? `?` + : escapeHtml(r.clientCount.toString()); + return ( + `` + + `${emoji ? emoji + " " : ""}${escapeHtml(r.agentName)}${systemTag}` + + `${escapeHtml(r.role)}` + + `${r.port}` + + `${clientCell}` + + `` + ); + }) + .join(""); + + return ( + `` + + `${headerCols.join("")}` + + `${rowsHtml}
` + ); +} + class FixedSchemaCommandHandler implements CommandHandler { public readonly description = "Set a fixed schema disable switching"; public readonly parameters = { diff --git a/ts/packages/dispatcher/dispatcher/src/execute/sessionContext.ts b/ts/packages/dispatcher/dispatcher/src/execute/sessionContext.ts index d942c43ef..f57855442 100644 --- a/ts/packages/dispatcher/dispatcher/src/execute/sessionContext.ts +++ b/ts/packages/dispatcher/dispatcher/src/execute/sessionContext.ts @@ -21,6 +21,9 @@ import { validateGrammarPatternsImpl } from "../validation/grammarValidationServ import registerDebug from "debug"; const debug = registerDebug("typeagent:dispatcher:sessionContext"); +const debugClientCountWarn = registerDebug( + "typeagent:dispatcher:clientCount:warn", +); import { randomUUID } from "node:crypto"; export function createSessionContext( @@ -243,6 +246,27 @@ export function createSessionContext( // ignore } }, + async notifyClientCountChanged( + role: string, + count: number, + ): Promise { + // Best-effort: counts are informational (surfaced only via + // `@system ports`) so any failure is swallowed; logged + // under typeagent:dispatcher:clientCount:warn so it isn't + // entirely invisible. + try { + context.portRegistrar.setClientCount( + name, + role, + sessionContextId, + count, + ); + } catch (e) { + debugClientCountWarn( + `notifyClientCountChanged failed for ${name}/${role}: ${e instanceof Error ? e.message : String(e)}`, + ); + } + }, async validateGrammarPatterns( request: GrammarValidationRequest, ): Promise { diff --git a/ts/packages/dispatcher/dispatcher/test/portRegistrar.spec.ts b/ts/packages/dispatcher/dispatcher/test/portRegistrar.spec.ts index f6c04568c..10e304db8 100644 --- a/ts/packages/dispatcher/dispatcher/test/portRegistrar.spec.ts +++ b/ts/packages/dispatcher/dispatcher/test/portRegistrar.spec.ts @@ -216,4 +216,70 @@ describe("PortRegistrar", () => { expect(r.hasActiveAllocations()).toBe(false); }); }); + + describe("client counts", () => { + test("getClientCount returns undefined when never set", () => { + const r = new PortRegistrar(); + r.register("browser", "ws", 1, SID_A); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + }); + + test("setClientCount round-trips on a live allocation", () => { + const r = new PortRegistrar(); + r.register("browser", "ws", 1, SID_A); + r.setClientCount("browser", "ws", SID_A, 3); + expect(r.getClientCount("browser", "ws", SID_A)).toBe(3); + }); + + test("setClientCount drops writes with no live allocation", () => { + const r = new PortRegistrar(); + r.setClientCount("browser", "ws", SID_A, 5); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + }); + + test("setClientCount accepts 0 and distinguishes it from undefined", () => { + const r = new PortRegistrar(); + r.register("browser", "ws", 1, SID_A); + r.setClientCount("browser", "ws", SID_A, 0); + expect(r.getClientCount("browser", "ws", SID_A)).toBe(0); + }); + + test("setClientCount rejects negative and non-integer values", () => { + const r = new PortRegistrar(); + r.register("browser", "ws", 1, SID_A); + r.setClientCount("browser", "ws", SID_A, -1); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + r.setClientCount("browser", "ws", SID_A, 1.5); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + }); + + test("release drops the matching count entry", () => { + const r = new PortRegistrar(); + const id = r.register("browser", "ws", 1, SID_A); + r.setClientCount("browser", "ws", SID_A, 4); + r.release(id); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + }); + + test("releaseAllForSession drops counts for that session only", () => { + const r = new PortRegistrar(); + r.register("browser", "ws", 1, SID_A); + r.register("browser", "ws", 2, SID_B); + r.setClientCount("browser", "ws", SID_A, 7); + r.setClientCount("browser", "ws", SID_B, 9); + r.releaseAllForSession(SID_A); + expect(r.getClientCount("browser", "ws", SID_A)).toBeUndefined(); + expect(r.getClientCount("browser", "ws", SID_B)).toBe(9); + }); + + test("counts are per-session under the same (agent, role)", () => { + const r = new PortRegistrar(); + r.register("code", "default", 8082, SID_A); + r.register("code", "default", 8082, SID_B); + r.setClientCount("code", "default", SID_A, 1); + r.setClientCount("code", "default", SID_B, 2); + expect(r.getClientCount("code", "default", SID_A)).toBe(1); + expect(r.getClientCount("code", "default", SID_B)).toBe(2); + }); + }); }); diff --git a/ts/packages/dispatcher/dispatcher/test/sessionContext.spec.ts b/ts/packages/dispatcher/dispatcher/test/sessionContext.spec.ts index b747e3fd5..98b94c852 100644 --- a/ts/packages/dispatcher/dispatcher/test/sessionContext.spec.ts +++ b/ts/packages/dispatcher/dispatcher/test/sessionContext.spec.ts @@ -249,6 +249,80 @@ describe("notifyReadinessChanged", () => { }); }); +describe("notifyClientCountChanged", () => { + function makeContextWithRegistrar() { + const setCalls: Array<{ + agentName: string; + role: string; + sessionContextId: string; + count: number; + }> = []; + const ctx = { + session: { + getSessionDirPath: () => undefined, + getConfig: () => ({}), + }, + storageProvider: { getStorage: () => ({}) as any }, + persistDir: undefined, + instanceDir: undefined, + commandLock: async (fn: any) => fn(), + agents: { + getTransientState: () => undefined, + getSharedLocalHostPort: async () => undefined, + setLocalHostPort: () => {}, + }, + clientIO: { + notify: () => {}, + question: async () => 0, + }, + translatorCache: { clear: () => {} }, + lastActionSchemaName: undefined, + conversationManager: undefined, + portRegistrar: { + setClientCount: ( + agentName: string, + role: string, + sessionContextId: string, + count: number, + ) => { + setCalls.push({ + agentName, + role, + sessionContextId, + count, + }); + }, + }, + } as any; + return { ctx, setCalls }; + } + + test("delegates to portRegistrar.setClientCount with the agent's own name and sessionContextId", async () => { + const { ctx, setCalls } = makeContextWithRegistrar(); + const sc = createSessionContext("myAgent", {}, ctx, false, "ncc-1"); + await sc.notifyClientCountChanged!("default", 3); + expect(setCalls).toEqual([ + { + agentName: "myAgent", + role: "default", + sessionContextId: "ncc-1", + count: 3, + }, + ]); + }); + + test("swallows errors from setClientCount so event-driven callers do not throw", async () => { + const { ctx } = makeContextWithRegistrar(); + ctx.portRegistrar.setClientCount = () => { + throw new Error("boom"); + }; + const sc = createSessionContext("myAgent", {}, ctx, false, "ncc-2"); + await expect( + sc.notifyClientCountChanged!("default", 1), + ).resolves.toBeUndefined(); + }); +}); + describe("initializeCommandHandlerContext option validation", () => { test("instanceDir without storageProvider throws", async () => { await expect( From 8afe576cf874f0bdf4939c286fe69d4311849904 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:08:17 -0700 Subject: [PATCH 02/11] Move ports command from @config ports to @system ports The previous commit upgraded the existing @config ports stub. Per discussion, the right home for this is the top-level @system namespace alongside the other diagnostic commands (@system session, @system history, @system env, etc.) since it's a system-wide diagnostic, not a configuration toggle. - Extract the handler into its own file system/handlers/portsCommandHandler.ts exporting PortsCommandHandler (the prior commit's class). - Register it under systemHandlers in systemAgent.ts. - Remove the now-empty 'ports' entry from getConfigCommandHandlers. - Update agentServer/protocol/README.md to mention @system ports (not @config ports). Build clean for agent-dispatcher; tests still 681/681. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ts/packages/agentServer/protocol/README.md | 2 +- .../system/handlers/configCommandHandlers.ts | 153 ----------------- .../system/handlers/portsCommandHandler.ts | 159 ++++++++++++++++++ .../src/context/system/systemAgent.ts | 2 + 4 files changed, 162 insertions(+), 154 deletions(-) create mode 100644 ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts diff --git a/ts/packages/agentServer/protocol/README.md b/ts/packages/agentServer/protocol/README.md index 60db95d98..dbe5dbfa6 100644 --- a/ts/packages/agentServer/protocol/README.md +++ b/ts/packages/agentServer/protocol/README.md @@ -100,7 +100,7 @@ createDiscoveryHandlers((agentName, role) => Passing a lookup callback (rather than the `IPortRegistrar` itself) keeps this package free of an `agent-dispatcher` dependency. -Live port allocations (including any client counts agents have published via `SessionContext.notifyClientCountChanged`) are surfaced in-process by the dispatcher's `@config ports` command. +Live port allocations (including any client counts agents have published via `SessionContext.notifyClientCountChanged`) are surfaced in-process by the dispatcher's `@system ports` command. --- diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts index 3b4914c1f..1c7ee9e24 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts @@ -46,7 +46,6 @@ import { displayWarn, } from "@typeagent/agent-sdk/helpers/display"; import { alwaysEnabledAgents } from "../../appAgentManager.js"; -import { SYSTEM_SESSION_CONTEXT_ID } from "../../portRegistrar.js"; import { getCacheFactory } from "../../../utils/cacheFactory.js"; import { resolveCommand } from "../../../command/command.js"; import { toggleActivityContext } from "../../../execute/activityContext.js"; @@ -992,156 +991,6 @@ class ConfigTranslationNumberOfInitialActionsCommandHandler } } -class ConfigPortsCommandHandler implements CommandHandler { - public readonly description = - "Lists ports registered by agents and the number of clients connected to each."; - - public readonly parameters = {}; - - public async run(context: ActionContext) { - const cmdContext = context.sessionContext.agentContext; - const registrar = cmdContext.portRegistrar; - const agents = cmdContext.agents; - - // Group registrar entries by (agentName, role, port). Code's - // shared WS server has no per-session client identity, so its - // count is global and the same number is reported across every - // session that registered the shared port. Grouping collapses - // those duplicates into one row with the (shared) count. - type Row = { - agentName: string; - role: string; - port: number; - isSystem: boolean; - // undefined === "agent never reported a count" (distinct from 0) - clientCount: number | undefined; - }; - const grouped = new Map(); - for (const alloc of registrar.list()) { - const key = `${alloc.agentName}\u0000${alloc.role}\u0000${alloc.port}`; - const isSystem = - alloc.sessionContextId === SYSTEM_SESSION_CONTEXT_ID; - const cc = registrar.getClientCount( - alloc.agentName, - alloc.role, - alloc.sessionContextId, - ); - const existing = grouped.get(key); - if (existing === undefined) { - grouped.set(key, { - agentName: alloc.agentName, - role: alloc.role, - port: alloc.port, - isSystem, - clientCount: cc, - }); - } else { - // System tag is sticky: if any underlying allocation is - // system-owned, treat the group as system. - if (isSystem) existing.isSystem = true; - // Sum across sessions when ≥1 has reported a count. - // Leave undefined only if NO underlying allocation has - // ever reported a count. - if (cc !== undefined) { - existing.clientCount = (existing.clientCount ?? 0) + cc; - } - } - } - - const rows = Array.from(grouped.values()).sort((a, b) => { - const an = a.agentName.localeCompare(b.agentName); - if (an !== 0) return an; - const ar = a.role.localeCompare(b.role); - if (ar !== 0) return ar; - return a.port - b.port; - }); - - if (rows.length === 0) { - displayWarn("No ports registered", context); - return; - } - - // Plain-text (CLI / console) — fixed-width via chalk for alignment. - const text: string[][] = [["", "Agent", "Role", "Port", "Clients"]]; - for (const r of rows) { - const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; - const name = r.isSystem - ? `${r.agentName} ${chalk.gray("[system]")}` - : r.agentName; - const clients = - r.clientCount === undefined - ? chalk.gray("?") - : r.clientCount.toString(); - text.push([emoji, name, r.role, r.port.toString(), clients]); - } - - // Rich HTML for the shell. - const html = buildPortsHtml(rows, agents); - - context.actionIO.appendDisplay({ - type: "text", - content: text, - alternates: [{ type: "html", content: html }], - }); - } -} - -function escapeHtml(s: string): string { - return s - .replace(/&/g, "&") - .replace(//g, ">") - .replace(/"/g, """) - .replace(/'/g, "'"); -} - -function buildPortsHtml( - rows: ReadonlyArray<{ - agentName: string; - role: string; - port: number; - isSystem: boolean; - clientCount: number | undefined; - }>, - agents: { getAppAgentEmoji(name: string): string }, -): string { - const thStyle = `text-align:left;padding:4px 8px;font-weight:600;font-size:13px;color:#64748b;border-bottom:2px solid #e2e8f0`; - const headerCols = [ - `Agent`, - `Role`, - `Port`, - `Clients`, - ]; - - const tdBase = `padding:3px 12px;border-bottom:1px solid #f1f5f9;white-space:nowrap`; - const rowsHtml = rows - .map((r) => { - const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; - const systemTag = r.isSystem - ? ` [system]` - : ""; - const clientCell = - r.clientCount === undefined - ? `?` - : escapeHtml(r.clientCount.toString()); - return ( - `` + - `${emoji ? emoji + " " : ""}${escapeHtml(r.agentName)}${systemTag}` + - `${escapeHtml(r.role)}` + - `${r.port}` + - `${clientCell}` + - `` - ); - }) - .join(""); - - return ( - `` + - `${headerCols.join("")}` + - `${rowsHtml}
` - ); -} - class FixedSchemaCommandHandler implements CommandHandler { public readonly description = "Set a fixed schema disable switching"; public readonly parameters = { @@ -2209,8 +2058,6 @@ export function getConfigCommandHandlers(): CommandHandlerTable { ), }, }, - - ports: new ConfigPortsCommandHandler(), }, }; } diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts new file mode 100644 index 000000000..f2690f481 --- /dev/null +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -0,0 +1,159 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { ActionContext } from "@typeagent/agent-sdk"; +import { CommandHandler } from "@typeagent/agent-sdk/helpers/command"; +import { displayWarn } from "@typeagent/agent-sdk/helpers/display"; +import chalk from "chalk"; +import { CommandHandlerContext } from "../../commandHandlerContext.js"; +import { SYSTEM_SESSION_CONTEXT_ID } from "../../portRegistrar.js"; + +export class PortsCommandHandler implements CommandHandler { + public readonly description = + "Lists ports registered by agents and the number of clients connected to each."; + + public readonly parameters = {}; + + public async run(context: ActionContext) { + const cmdContext = context.sessionContext.agentContext; + const registrar = cmdContext.portRegistrar; + const agents = cmdContext.agents; + + // Group registrar entries by (agentName, role, port). Code's + // shared WS server has no per-session client identity, so its + // count is global and the same number is reported across every + // session that registered the shared port. Grouping collapses + // those duplicates into one row with the (shared) count. + type Row = { + agentName: string; + role: string; + port: number; + isSystem: boolean; + // undefined === "agent never reported a count" (distinct from 0) + clientCount: number | undefined; + }; + const grouped = new Map(); + for (const alloc of registrar.list()) { + const key = `${alloc.agentName}\u0000${alloc.role}\u0000${alloc.port}`; + const isSystem = + alloc.sessionContextId === SYSTEM_SESSION_CONTEXT_ID; + const cc = registrar.getClientCount( + alloc.agentName, + alloc.role, + alloc.sessionContextId, + ); + const existing = grouped.get(key); + if (existing === undefined) { + grouped.set(key, { + agentName: alloc.agentName, + role: alloc.role, + port: alloc.port, + isSystem, + clientCount: cc, + }); + } else { + // System tag is sticky: if any underlying allocation is + // system-owned, treat the group as system. + if (isSystem) existing.isSystem = true; + // Sum across sessions when ≥1 has reported a count. + // Leave undefined only if NO underlying allocation has + // ever reported a count. + if (cc !== undefined) { + existing.clientCount = (existing.clientCount ?? 0) + cc; + } + } + } + + const rows = Array.from(grouped.values()).sort((a, b) => { + const an = a.agentName.localeCompare(b.agentName); + if (an !== 0) return an; + const ar = a.role.localeCompare(b.role); + if (ar !== 0) return ar; + return a.port - b.port; + }); + + if (rows.length === 0) { + displayWarn("No ports registered", context); + return; + } + + // Plain-text (CLI / console) — fixed-width via chalk for alignment. + const text: string[][] = [["", "Agent", "Role", "Port", "Clients"]]; + for (const r of rows) { + const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; + const name = r.isSystem + ? `${r.agentName} ${chalk.gray("[system]")}` + : r.agentName; + const clients = + r.clientCount === undefined + ? chalk.gray("?") + : r.clientCount.toString(); + text.push([emoji, name, r.role, r.port.toString(), clients]); + } + + // Rich HTML for the shell. + const html = buildPortsHtml(rows, agents); + + context.actionIO.appendDisplay({ + type: "text", + content: text, + alternates: [{ type: "html", content: html }], + }); + } +} + +function escapeHtml(s: string): string { + return s + .replace(/&/g, "&") + .replace(//g, ">") + .replace(/"/g, """) + .replace(/'/g, "'"); +} + +function buildPortsHtml( + rows: ReadonlyArray<{ + agentName: string; + role: string; + port: number; + isSystem: boolean; + clientCount: number | undefined; + }>, + agents: { getAppAgentEmoji(name: string): string }, +): string { + const thStyle = `text-align:left;padding:4px 8px;font-weight:600;font-size:13px;color:#64748b;border-bottom:2px solid #e2e8f0`; + const headerCols = [ + `Agent`, + `Role`, + `Port`, + `Clients`, + ]; + + const tdBase = `padding:3px 12px;border-bottom:1px solid #f1f5f9;white-space:nowrap`; + const rowsHtml = rows + .map((r) => { + const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; + const systemTag = r.isSystem + ? ` [system]` + : ""; + const clientCell = + r.clientCount === undefined + ? `?` + : escapeHtml(r.clientCount.toString()); + return ( + `` + + `${emoji ? emoji + " " : ""}${escapeHtml(r.agentName)}${systemTag}` + + `${escapeHtml(r.role)}` + + `${r.port}` + + `${clientCell}` + + `` + ); + }) + .join(""); + + return ( + `` + + `${headerCols.join("")}` + + `${rowsHtml}
` + ); +} diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/systemAgent.ts b/ts/packages/dispatcher/dispatcher/src/context/system/systemAgent.ts index f71e4eab6..0dd669aa1 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/systemAgent.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/systemAgent.ts @@ -59,6 +59,7 @@ import { OpenCommandHandler } from "./handlers/openCommandHandler.js"; import { getIndexCommandHandlers } from "./handlers/indexCommandHandler.js"; import { getMemoryCommandHandlers } from "../memory.js"; import { getSettingsCommandHandlers } from "./handlers/settingsCommandHandlers.js"; +import { PortsCommandHandler } from "./handlers/portsCommandHandler.js"; export const systemHandlers: CommandHandlerTable = { description: "Type Agent System Commands", @@ -106,6 +107,7 @@ export const systemHandlers: CommandHandlerTable = { open: new OpenCommandHandler(), index: getIndexCommandHandlers(), settings: getSettingsCommandHandlers(), + ports: new PortsCommandHandler(), }, }; From 8d0bc4986faeea3128d415d660ae25dba544ddbb Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:14:42 -0700 Subject: [PATCH 03/11] @system ports: tolerate non-app-agent entries in the registrar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The registrar contains an entry for the agent-server's own listening port registered under the well-known 'agent-server' name (see server.ts:464). That name is not a real app-agent, so calling agents.getAppAgentEmoji('agent-server') throws 'Unknown app agent: agent-server' and the whole @system ports command bombs out — only visible when the dispatcher is running inside the agent-server (detached mode), since standalone shell mode has no such entry. Wrap the emoji lookup in a try/catch fallback to an empty string so synthetic / non-app-agent entries render with no emoji instead of killing the command. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../system/handlers/portsCommandHandler.ts | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index f2690f481..4e3983e4c 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -19,6 +19,19 @@ export class PortsCommandHandler implements CommandHandler { const registrar = cmdContext.portRegistrar; const agents = cmdContext.agents; + // Best-effort emoji lookup: the registrar may contain entries for + // pseudo-agents that aren't real app-agents (notably the + // agent-server's own listening port registered under the + // well-known "agent-server" name), in which case + // `getAppAgentEmoji` throws "Unknown app agent: ...". + const safeEmoji = (name: string): string => { + try { + return agents.getAppAgentEmoji(name) ?? ""; + } catch { + return ""; + } + }; + // Group registrar entries by (agentName, role, port). Code's // shared WS server has no per-session client identity, so its // count is global and the same number is reported across every @@ -80,7 +93,7 @@ export class PortsCommandHandler implements CommandHandler { // Plain-text (CLI / console) — fixed-width via chalk for alignment. const text: string[][] = [["", "Agent", "Role", "Port", "Clients"]]; for (const r of rows) { - const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; + const emoji = safeEmoji(r.agentName); const name = r.isSystem ? `${r.agentName} ${chalk.gray("[system]")}` : r.agentName; @@ -92,7 +105,7 @@ export class PortsCommandHandler implements CommandHandler { } // Rich HTML for the shell. - const html = buildPortsHtml(rows, agents); + const html = buildPortsHtml(rows, { getAppAgentEmoji: safeEmoji }); context.actionIO.appendDisplay({ type: "text", From 04add85300169d85b998df7ac2c06650c00d9a4f Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:16:58 -0700 Subject: [PATCH 04/11] =?UTF-8?q?@system=20ports:=20show=20=F0=9F=A4=96=20?= =?UTF-8?q?for=20the=20agent-server's=20own=20row?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/context/system/handlers/portsCommandHandler.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index 4e3983e4c..21d2f54a0 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -6,7 +6,10 @@ import { CommandHandler } from "@typeagent/agent-sdk/helpers/command"; import { displayWarn } from "@typeagent/agent-sdk/helpers/display"; import chalk from "chalk"; import { CommandHandlerContext } from "../../commandHandlerContext.js"; -import { SYSTEM_SESSION_CONTEXT_ID } from "../../portRegistrar.js"; +import { + AGENT_SERVER_REGISTRAR_NAME, + SYSTEM_SESSION_CONTEXT_ID, +} from "../../portRegistrar.js"; export class PortsCommandHandler implements CommandHandler { public readonly description = @@ -23,8 +26,11 @@ export class PortsCommandHandler implements CommandHandler { // pseudo-agents that aren't real app-agents (notably the // agent-server's own listening port registered under the // well-known "agent-server" name), in which case - // `getAppAgentEmoji` throws "Unknown app agent: ...". + // `getAppAgentEmoji` throws "Unknown app agent: ...". Render + // the agent-server's own row with 🤖 so it's not visually + // empty; everything else falls back to no emoji. const safeEmoji = (name: string): string => { + if (name === AGENT_SERVER_REGISTRAR_NAME) return "🤖"; try { return agents.getAppAgentEmoji(name) ?? ""; } catch { From 7d52828a37c4f3baa9b0982e98da1f6d44727bc0 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:21:32 -0700 Subject: [PATCH 05/11] @system ports: take max instead of sum when collapsing same-port groups Multiple conversation sessions registered to the SAME physical port (shared servers like browser and code) all observe the same global connected-client count and each fans it out independently. Summing them double-counts: 2 sessions x 2 clients reported = 4. Taking the max within a (agent, role, port) group gives the true count for shared servers and remains correct for per-session servers (where each group has exactly one contributor). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../system/handlers/portsCommandHandler.ts | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index 21d2f54a0..000f6f5a7 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -74,11 +74,18 @@ export class PortsCommandHandler implements CommandHandler { // System tag is sticky: if any underlying allocation is // system-owned, treat the group as system. if (isSystem) existing.isSystem = true; - // Sum across sessions when ≥1 has reported a count. - // Leave undefined only if NO underlying allocation has - // ever reported a count. + // Within a (agent, role, port) group, every contributing + // session is registered to the SAME physical server, so + // each session's reported count is the same global + // number — not an independent value to add. Take the max + // (treating "never reported" as 0 only when comparing). + // For per-session servers each group has a single + // contributor and this collapses to that count anyway. if (cc !== undefined) { - existing.clientCount = (existing.clientCount ?? 0) + cc; + existing.clientCount = Math.max( + existing.clientCount ?? 0, + cc, + ); } } } From 320126f1ec9059b94b9111bf2424f68fc9f89316 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:25:43 -0700 Subject: [PATCH 06/11] @system ports: show 'N/A' instead of '?' for agents that don't publish counts Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/context/system/handlers/portsCommandHandler.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index 000f6f5a7..b43f21eee 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -112,7 +112,7 @@ export class PortsCommandHandler implements CommandHandler { : r.agentName; const clients = r.clientCount === undefined - ? chalk.gray("?") + ? chalk.gray("N/A") : r.clientCount.toString(); text.push([emoji, name, r.role, r.port.toString(), clients]); } @@ -164,7 +164,7 @@ function buildPortsHtml( : ""; const clientCell = r.clientCount === undefined - ? `?` + ? `N/A` : escapeHtml(r.clientCount.toString()); return ( `` + From 004871065eb50fa57f3b676175c53034c8e8f3d1 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:28:56 -0700 Subject: [PATCH 07/11] @system ports: drop redundant [system] tag from agent-server row Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../system/handlers/portsCommandHandler.ts | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index b43f21eee..f287536f6 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -107,7 +107,12 @@ export class PortsCommandHandler implements CommandHandler { const text: string[][] = [["", "Agent", "Role", "Port", "Clients"]]; for (const r of rows) { const emoji = safeEmoji(r.agentName); - const name = r.isSystem + // The agent-server row is always system-scoped and the name + // already makes that obvious, so the [system] tag would just + // be visual noise. + const showSystemTag = + r.isSystem && r.agentName !== AGENT_SERVER_REGISTRAR_NAME; + const name = showSystemTag ? `${r.agentName} ${chalk.gray("[system]")}` : r.agentName; const clients = @@ -118,7 +123,10 @@ export class PortsCommandHandler implements CommandHandler { } // Rich HTML for the shell. - const html = buildPortsHtml(rows, { getAppAgentEmoji: safeEmoji }); + const html = buildPortsHtml(rows, { + getAppAgentEmoji: safeEmoji, + shouldShowSystemTag: (name) => name !== AGENT_SERVER_REGISTRAR_NAME, + }); context.actionIO.appendDisplay({ type: "text", @@ -145,7 +153,10 @@ function buildPortsHtml( isSystem: boolean; clientCount: number | undefined; }>, - agents: { getAppAgentEmoji(name: string): string }, + agents: { + getAppAgentEmoji(name: string): string; + shouldShowSystemTag(name: string): boolean; + }, ): string { const thStyle = `text-align:left;padding:4px 8px;font-weight:600;font-size:13px;color:#64748b;border-bottom:2px solid #e2e8f0`; const headerCols = [ @@ -159,9 +170,10 @@ function buildPortsHtml( const rowsHtml = rows .map((r) => { const emoji = agents.getAppAgentEmoji(r.agentName) ?? ""; - const systemTag = r.isSystem - ? ` [system]` - : ""; + const systemTag = + r.isSystem && agents.shouldShowSystemTag(r.agentName) + ? ` [system]` + : ""; const clientCell = r.clientCount === undefined ? `N/A` From 0ec323d9dfd6ada4f4154997f0db62e3604ff635 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 17:34:21 -0700 Subject: [PATCH 08/11] docs: document the new @ports diagnostic command Add an entry for @ports to docs/architecture/dispatcher.md's system-agent command list and to dispatcher/README.md under a new Diagnostics section. Both call out the (agent, role, port) grouping, the agent-server's own listen-port row, and the N/A semantic for agents that don't publish client counts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ts/docs/architecture/dispatcher.md | 6 ++++++ ts/packages/dispatcher/dispatcher/README.md | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/ts/docs/architecture/dispatcher.md b/ts/docs/architecture/dispatcher.md index e2bc6b17a..1cad07b82 100644 --- a/ts/docs/architecture/dispatcher.md +++ b/ts/docs/architecture/dispatcher.md @@ -507,6 +507,12 @@ Handles `@`-prefixed system commands: - `@explain` — Explanation of cached translations - `@feedback` — Inspect and export user-feedback entries recorded by the chat UI (`list`, `top`, `filter`, `export`, `count`) +- `@ports` — List all registered TCP ports (per `(agent, role, port)` + group) with the agent-server's own listen port and the current number + of clients connected to each agent's WS server. Agents that don't + publish a count via `SessionContext.notifyClientCountChanged` render + `N/A` in the Clients column. Currently the browser and code agents + publish counts; others are diagnostics-only. Each command has a `CommandDescriptor` that defines expected parameters, subcommands, and help text. diff --git a/ts/packages/dispatcher/dispatcher/README.md b/ts/packages/dispatcher/dispatcher/README.md index 8a199827f..5256a6ed0 100644 --- a/ts/packages/dispatcher/dispatcher/README.md +++ b/ts/packages/dispatcher/dispatcher/README.md @@ -192,6 +192,12 @@ By default agents runs out of proc in their own process. This is to ensure that | `@config explanation on\|off` | Toggle LLM explanation (Turn off to stop updating construction store) | | `@config log db on\|off` | Toggle sending logging information to a remote database (default: on) | +### Diagnostics + +| Command | Description | +| -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `@ports` | List every TCP port registered with the dispatcher's `PortRegistrar`, grouped by `(agent, role, port)`. Includes the agent-server's own listen port and the connected-client count for any agent that publishes one via `SessionContext.notifyClientCountChanged` (currently the browser and code agents). Rows from agents that don't publish a count render `N/A`. | + ### User feedback When the user rates an agent message via the chat UI's thumbs-up/down buttons or moves a bubble to the trash, the dispatcher persists each event to the per-session `displayLog.json` (as `user-feedback` and `user-message-hidden` entries) and emits a `userFeedback` telemetry event through `Logger.logEvent`. The `@feedback` command group lets you inspect and export those entries. From 637bea033b31f4e48153a869144266b4917a3fdb Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Mon, 18 May 2026 18:25:05 -0700 Subject: [PATCH 09/11] fix: address PR #2358 Copilot review comments - portsCommandHandler: sum per-session client counts (treating undefined as 0) instead of Math.max. Browser agent reports genuinely per-session counts; max undercounts when multiple sessions each have clients. - codeActionHandler: attribute the (global) shared-server count to a single primary session (first in insertion order); other sessions report 0. Prevents double-counting now that the handler sums. Transfers ownership when the primary disables. - agentWebSocketServer (browser): coerce eq.headers.origin array values to a single string before passing to the allowlist (Node typings allow string | string[] | undefined). - browserIpc: use bracketed hostname when reconstructing the inline-browser WS URL so IPv6 literals like ::1 round-trip as ws://[::1]:port/... instead of the invalid ws://::1:port/... - originAllowlist (browser + code): also accept unbracketed ::1 for robustness against URL parser/serializer differences across runtimes (precedent: examples/workflow/engine/builtinTasks.ts). Doc-nit suggestions (rename @system ports -> @ports) were skipped: ports is registered inside systemHandlers so @system ports is the correct invocation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/agent/agentWebSocketServer.mts | 9 +++- .../browser/src/agent/originAllowlist.mts | 9 +++- .../agents/code/src/codeActionHandler.ts | 41 +++++++++++++++---- .../agents/code/src/originAllowlist.ts | 9 +++- .../system/handlers/portsCommandHandler.ts | 22 +++++----- ts/packages/shell/src/main/browserIpc.ts | 6 ++- 6 files changed, 71 insertions(+), 25 deletions(-) diff --git a/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts b/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts index 71e2c5e7d..a99a83c64 100644 --- a/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts +++ b/ts/packages/agents/browser/src/agent/agentWebSocketServer.mts @@ -96,7 +96,14 @@ export class AgentWebSocketServer { const server = new WebSocketServer({ port, verifyClient: (info, cb) => { - const origin = info.origin || info.req.headers.origin; + // `info.req.headers.origin` is `string | string[] | + // undefined`; coerce arrays to the first element so + // `isAllowedAgentOrigin` (which calls `startsWith` / + // `new URL`) only ever sees a string. + const rawOrigin = info.origin || info.req.headers.origin; + const origin = Array.isArray(rawOrigin) + ? rawOrigin[0] + : rawOrigin; if (isAllowedAgentOrigin(origin)) { cb(true); } else { diff --git a/ts/packages/agents/browser/src/agent/originAllowlist.mts b/ts/packages/agents/browser/src/agent/originAllowlist.mts index 17df4c7d0..d19473487 100644 --- a/ts/packages/agents/browser/src/agent/originAllowlist.mts +++ b/ts/packages/agents/browser/src/agent/originAllowlist.mts @@ -37,11 +37,16 @@ export function isAllowedAgentOrigin(origin: string | undefined): boolean { } // Node's URL parser preserves IPv6 brackets in `hostname` // (e.g. `new URL("http://[::1]:8080").hostname === "[::1]"`), - // so match the bracketed form. + // so match the bracketed form. Also accept the unbracketed + // `::1` for robustness against URL parser/serializer + // differences across runtimes (other SSRF guards in the repo, + // e.g. examples/workflow/engine/src/builtinTasks.ts, accept + // both). return ( u.hostname === "localhost" || u.hostname === "127.0.0.1" || - u.hostname === "[::1]" + u.hostname === "[::1]" || + u.hostname === "::1" ); } catch { return false; diff --git a/ts/packages/agents/code/src/codeActionHandler.ts b/ts/packages/agents/code/src/codeActionHandler.ts index 7ad176b2c..b035438ed 100644 --- a/ts/packages/agents/code/src/codeActionHandler.ts +++ b/ts/packages/agents/code/src/codeActionHandler.ts @@ -212,12 +212,19 @@ function attachSharedOnMessage(server: CodeAgentWebSocketServer): void { debug("Error parsing WebSocket message:", error); } }; - // Fan out client-count updates to every active session so each - // session's `@system ports` row reflects the (global) count. The - // SDK method swallows errors internally. + // Fan out client-count updates to active sessions. To prevent the + // `@system ports` summing from double-counting (each session is + // registered to the SAME physical server), attribute the global + // count to a single "primary" session (the first one in insertion + // order) and report 0 from the rest. The SDK method swallows + // errors internally. server.onClientCountChanged = (count: number) => { + const primary = sharedActiveSessions.values().next().value; for (const sc of sharedActiveSessions) { - void sc.notifyClientCountChanged("default", count); + void sc.notifyClientCountChanged( + "default", + sc === primary ? count : 0, + ); } }; } @@ -280,12 +287,17 @@ async function updateCodeContext( ); sharedWebSocketRefCount++; sharedActiveSessions.add(context); - // Publish the current (global) count to this session's - // newly-registered row so `@system ports` doesn't show - // "?" until the next connect/disconnect event. + // Publish the current (global) count to the primary + // session (first in insertion order) and 0 to others + // so `@system ports` summing doesn't double-count. If + // this session is now becoming the primary (i.e. it's + // the first to enable), it gets the real count; + // otherwise it reports 0 and any future + // onClientCountChanged fanout will keep it at 0. + const primary = sharedActiveSessions.values().next().value; void context.notifyClientCountChanged( "default", - server.getConnectedCount(), + context === primary ? server.getConnectedCount() : 0, ); } } catch (e) { @@ -308,6 +320,8 @@ async function updateCodeContext( // released by the backstop. agentContext.portRegistration?.release(); delete agentContext.portRegistration; + const wasPrimary = + sharedActiveSessions.values().next().value === context; sharedActiveSessions.delete(context); sharedWebSocketRefCount = Math.max(0, sharedWebSocketRefCount - 1); @@ -321,6 +335,17 @@ async function updateCodeContext( sharedClosingPromise = undefined; }); await sharedClosingPromise; + } else if (wasPrimary && sharedWebSocketServer) { + // Primary session went away — transfer the (global) count + // to the new primary so `@system ports` keeps reporting + // the real number instead of 0. + const newPrimary = sharedActiveSessions.values().next().value; + if (newPrimary) { + void newPrimary.notifyClientCountChanged( + "default", + sharedWebSocketServer.getConnectedCount(), + ); + } } } } diff --git a/ts/packages/agents/code/src/originAllowlist.ts b/ts/packages/agents/code/src/originAllowlist.ts index 6deb08f62..f7a3e0dc7 100644 --- a/ts/packages/agents/code/src/originAllowlist.ts +++ b/ts/packages/agents/code/src/originAllowlist.ts @@ -45,11 +45,16 @@ export function isAllowedAgentOrigin(origin: string | undefined): boolean { } // Node's URL parser preserves IPv6 brackets in `hostname` // (e.g. `new URL("http://[::1]:8080").hostname === "[::1]"`), - // so match the bracketed form. + // so match the bracketed form. Also accept the unbracketed + // `::1` for robustness against URL parser/serializer + // differences across runtimes (other SSRF guards in the repo, + // e.g. examples/workflow/engine/src/builtinTasks.ts, accept + // both). return ( u.hostname === "localhost" || u.hostname === "127.0.0.1" || - u.hostname === "[::1]" + u.hostname === "[::1]" || + u.hostname === "::1" ); } catch { return false; diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index f287536f6..2fc541703 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -74,18 +74,18 @@ export class PortsCommandHandler implements CommandHandler { // System tag is sticky: if any underlying allocation is // system-owned, treat the group as system. if (isSystem) existing.isSystem = true; - // Within a (agent, role, port) group, every contributing - // session is registered to the SAME physical server, so - // each session's reported count is the same global - // number — not an independent value to add. Take the max - // (treating "never reported" as 0 only when comparing). - // For per-session servers each group has a single - // contributor and this collapses to that count anyway. + // Within a (agent, role, port) group, sum contributing + // counts. Some agents (e.g. browser) report genuine + // per-session counts that must be added to get the + // total. Agents that share a single physical server + // and would otherwise double-count (e.g. code) are + // expected to publish the global count from a single + // "owner" session and 0 from the rest — see the + // primary-session pattern in codeActionHandler. A + // group whose contributors all report `undefined` + // stays `undefined` (rendered as "N/A"). if (cc !== undefined) { - existing.clientCount = Math.max( - existing.clientCount ?? 0, - cc, - ); + existing.clientCount = (existing.clientCount ?? 0) + cc; } } } diff --git a/ts/packages/shell/src/main/browserIpc.ts b/ts/packages/shell/src/main/browserIpc.ts index 9321a0758..0f176de34 100644 --- a/ts/packages/shell/src/main/browserIpc.ts +++ b/ts/packages/shell/src/main/browserIpc.ts @@ -287,7 +287,11 @@ async function createInlineBrowserWebSocket(): Promise { let endpoint: string; try { const u = new URL(agentServerUrl); - endpoint = `${u.protocol}//${u.hostname}:${result.port}/?channel=browser&role=client&clientId=inlineBrowser&sessionId=${sessionId}`; + // Use `host` (which preserves IPv6 brackets, e.g. `[::1]`) and + // strip the existing port — `u.host` includes it. `u.hostname` + // would round-trip IPv6 literals without brackets, producing an + // invalid URL like `ws://::1:8999/...`. + endpoint = `${u.protocol}//${u.hostname.includes(":") && !u.hostname.startsWith("[") ? `[${u.hostname}]` : u.hostname}:${result.port}/?channel=browser&role=client&clientId=inlineBrowser&sessionId=${sessionId}`; } catch (e) { debugBrowserIPCError("Invalid agent-server URL: %s", e); return undefined; From ab1cbf75011bf5b0e3e25f2b750e85bcb7d54404 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Tue, 19 May 2026 13:37:49 -0700 Subject: [PATCH 10/11] fix: address second round of PR #2358 review comments - portsCommandHandler: escape emoji in HTML output (XSS hardening; agent-provided metadata was interpolated raw) - dispatcher README: rename @ports -> @system ports (correct invocation) - browserIpc: clarify the IPv6 bracket comment to match the actual hostname-based implementation Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ts/packages/dispatcher/dispatcher/README.md | 6 +++--- .../src/context/system/handlers/portsCommandHandler.ts | 2 +- ts/packages/shell/src/main/browserIpc.ts | 10 ++++++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/README.md b/ts/packages/dispatcher/dispatcher/README.md index 5256a6ed0..d68c37d93 100644 --- a/ts/packages/dispatcher/dispatcher/README.md +++ b/ts/packages/dispatcher/dispatcher/README.md @@ -194,9 +194,9 @@ By default agents runs out of proc in their own process. This is to ensure that ### Diagnostics -| Command | Description | -| -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `@ports` | List every TCP port registered with the dispatcher's `PortRegistrar`, grouped by `(agent, role, port)`. Includes the agent-server's own listen port and the connected-client count for any agent that publishes one via `SessionContext.notifyClientCountChanged` (currently the browser and code agents). Rows from agents that don't publish a count render `N/A`. | +| Command | Description | +| --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `@system ports` | List every TCP port registered with the dispatcher's `PortRegistrar`, grouped by `(agent, role, port)`. Includes the agent-server's own listen port and the connected-client count for any agent that publishes one via `SessionContext.notifyClientCountChanged` (currently the browser and code agents). Rows from agents that don't publish a count render `N/A`. | ### User feedback diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts index 2fc541703..c76d13861 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/portsCommandHandler.ts @@ -180,7 +180,7 @@ function buildPortsHtml( : escapeHtml(r.clientCount.toString()); return ( `` + - `${emoji ? emoji + " " : ""}${escapeHtml(r.agentName)}${systemTag}` + + `${emoji ? escapeHtml(emoji) + " " : ""}${escapeHtml(r.agentName)}${systemTag}` + `${escapeHtml(r.role)}` + `${r.port}` + `${clientCell}` + diff --git a/ts/packages/shell/src/main/browserIpc.ts b/ts/packages/shell/src/main/browserIpc.ts index 0f176de34..413abb512 100644 --- a/ts/packages/shell/src/main/browserIpc.ts +++ b/ts/packages/shell/src/main/browserIpc.ts @@ -287,10 +287,12 @@ async function createInlineBrowserWebSocket(): Promise { let endpoint: string; try { const u = new URL(agentServerUrl); - // Use `host` (which preserves IPv6 brackets, e.g. `[::1]`) and - // strip the existing port — `u.host` includes it. `u.hostname` - // would round-trip IPv6 literals without brackets, producing an - // invalid URL like `ws://::1:8999/...`. + // `URL.hostname` strips IPv6 brackets in some Node versions + // (e.g. returns `::1` instead of `[::1]`), which would produce + // an invalid URL like `ws://::1:8999/...`. Re-bracket bare IPv6 + // literals before composing the endpoint. We use `hostname` + // (not `host`) because `host` includes the existing port and we + // need to substitute the discovered port from the registrar. endpoint = `${u.protocol}//${u.hostname.includes(":") && !u.hostname.startsWith("[") ? `[${u.hostname}]` : u.hostname}:${result.port}/?channel=browser&role=client&clientId=inlineBrowser&sessionId=${sessionId}`; } catch (e) { debugBrowserIPCError("Invalid agent-server URL: %s", e); From 5f32ec2158c26802ae152a355889cfc89d288923 Mon Sep 17 00:00:00 2001 From: Tal Zaccai Date: Thu, 21 May 2026 12:02:22 -0700 Subject: [PATCH 11/11] Factor agent originAllowlist into shared websocket-utils helper Address @robgruen review feedback on PR #2358: the browser and code agent originAllowlist files are near-clones differing only in which URL-scheme prefixes (chrome-extension, vscode-webview, etc.) they accept on top of the shared loopback + no-Origin baseline. Extract that baseline into websocket-utils as a new subpath export 'websocket-utils/originAllowlist' exposing createAgentOriginAllowlist({ extensionSchemes }). The per-agent files become one-line factory calls that name only their extra accepted schemes. The subpath export (rather than re-exporting from the package index) keeps Jest tests for these agents from transitively loading @typeagent/config, which uses import.meta and breaks under Jest's CommonJS runtime. Behavior preserved: code agent's prior defense-in-depth allowance for the unbracketed '::1' hostname is now part of the shared baseline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../browser/src/agent/originAllowlist.mts | 51 +++--------- .../agents/code/src/originAllowlist.ts | 70 +++++----------- ts/packages/utils/webSocketUtils/package.json | 3 +- .../webSocketUtils/src/originAllowlist.ts | 81 +++++++++++++++++++ 4 files changed, 111 insertions(+), 94 deletions(-) create mode 100644 ts/packages/utils/webSocketUtils/src/originAllowlist.ts diff --git a/ts/packages/agents/browser/src/agent/originAllowlist.mts b/ts/packages/agents/browser/src/agent/originAllowlist.mts index d19473487..c4a9de80a 100644 --- a/ts/packages/agents/browser/src/agent/originAllowlist.mts +++ b/ts/packages/agents/browser/src/agent/originAllowlist.mts @@ -1,54 +1,23 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +import { createAgentOriginAllowlist } from "websocket-utils/originAllowlist"; + /** * Origin allowlist for the browser agent's WebSocket server. * * Allowed: * - `chrome-extension://...` and `moz-extension://...` (the typeagent * Chrome / Edge browser extensions). - * - `http(s)://localhost(:port)`, `http(s)://127.0.0.1(:port)`, and - * `http(s)://[::1](:port)` (the Electron shell's inline browser, - * plus loopback dev clients on either IPv4 or IPv6). - * - **No Origin header** — Node `ws` clients (and any non-browser caller - * that hits the bridge over loopback) don't send Origin. The bridge - * binds to localhost, so this is loopback-restricted at the OS level. + * - The shared loopback + no-Origin baseline documented on + * {@link createAgentOriginAllowlist} (Electron shell's inline + * browser, Node `ws` clients). * * Anything else is rejected with HTTP 403 before the `connection` event - * fires. Every per-agent listener that binds to an ephemeral port via the - * PortRegistrar must gate Origin so those ports can't be dialed by + * fires. Every per-agent listener that binds to an ephemeral port via + * the PortRegistrar must gate Origin so those ports can't be dialed by * arbitrary web pages on the same host. */ -export function isAllowedAgentOrigin(origin: string | undefined): boolean { - if (origin === undefined || origin === "" || origin === "null") { - // No Origin header: legitimate for Node `ws` clients. - return true; - } - if ( - origin.startsWith("chrome-extension://") || - origin.startsWith("moz-extension://") - ) { - return true; - } - try { - const u = new URL(origin); - if (u.protocol !== "http:" && u.protocol !== "https:") { - return false; - } - // Node's URL parser preserves IPv6 brackets in `hostname` - // (e.g. `new URL("http://[::1]:8080").hostname === "[::1]"`), - // so match the bracketed form. Also accept the unbracketed - // `::1` for robustness against URL parser/serializer - // differences across runtimes (other SSRF guards in the repo, - // e.g. examples/workflow/engine/src/builtinTasks.ts, accept - // both). - return ( - u.hostname === "localhost" || - u.hostname === "127.0.0.1" || - u.hostname === "[::1]" || - u.hostname === "::1" - ); - } catch { - return false; - } -} +export const isAllowedAgentOrigin = createAgentOriginAllowlist({ + extensionSchemes: ["chrome-extension://", "moz-extension://"], +}); diff --git a/ts/packages/agents/code/src/originAllowlist.ts b/ts/packages/agents/code/src/originAllowlist.ts index f7a3e0dc7..fd26a853f 100644 --- a/ts/packages/agents/code/src/originAllowlist.ts +++ b/ts/packages/agents/code/src/originAllowlist.ts @@ -1,62 +1,28 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +import { createAgentOriginAllowlist } from "websocket-utils/originAllowlist"; + /** * Origin allowlist for the code agent's WebSocket server. * * Allowed: * - `vscode-webview://...`, `vscode-file://...`, and - * `vscode-resource://...` (the VS Code extension host's - * sandboxed surfaces). - * - `http(s)://localhost(:port)`, `http(s)://127.0.0.1(:port)`, and - * `http(s)://[::1](:port)` (loopback dev clients on either IPv4 or - * IPv6). - * - **No Origin header** — Node `ws` clients (and the VS Code - * extension's own `ws` client) don't send Origin. The server - * binds to localhost, so this is loopback-restricted at the OS - * level. - * - * Anything else is rejected with HTTP 403 before the `connection` - * event fires. Every per-agent listener that binds to an ephemeral port - * via the PortRegistrar must gate Origin so those ports can't be dialed - * by arbitrary web pages on the same host. + * `vscode-resource://...` (the VS Code extension host's sandboxed + * surfaces). + * - The shared loopback + no-Origin baseline documented on + * {@link createAgentOriginAllowlist} (the VS Code extension's own + * `ws` client, manual loopback debugging). * - * Kept in sync with `agents/browser/src/agent/originAllowlist.mts`; - * duplicated rather than shared because the policies differ in which - * extension scheme prefixes are accepted (Chrome/Firefox vs. VS Code). + * Anything else is rejected with HTTP 403 before the `connection` event + * fires. Every per-agent listener that binds to an ephemeral port via + * the PortRegistrar must gate Origin so those ports can't be dialed by + * arbitrary web pages on the same host. */ -export function isAllowedAgentOrigin(origin: string | undefined): boolean { - if (origin === undefined || origin === "" || origin === "null") { - // No Origin header: legitimate for Node `ws` clients (the - // VS Code extension uses one). - return true; - } - if ( - origin.startsWith("vscode-webview://") || - origin.startsWith("vscode-file://") || - origin.startsWith("vscode-resource://") - ) { - return true; - } - try { - const u = new URL(origin); - if (u.protocol !== "http:" && u.protocol !== "https:") { - return false; - } - // Node's URL parser preserves IPv6 brackets in `hostname` - // (e.g. `new URL("http://[::1]:8080").hostname === "[::1]"`), - // so match the bracketed form. Also accept the unbracketed - // `::1` for robustness against URL parser/serializer - // differences across runtimes (other SSRF guards in the repo, - // e.g. examples/workflow/engine/src/builtinTasks.ts, accept - // both). - return ( - u.hostname === "localhost" || - u.hostname === "127.0.0.1" || - u.hostname === "[::1]" || - u.hostname === "::1" - ); - } catch { - return false; - } -} +export const isAllowedAgentOrigin = createAgentOriginAllowlist({ + extensionSchemes: [ + "vscode-webview://", + "vscode-file://", + "vscode-resource://", + ], +}); diff --git a/ts/packages/utils/webSocketUtils/package.json b/ts/packages/utils/webSocketUtils/package.json index dd6e40b93..078723e48 100644 --- a/ts/packages/utils/webSocketUtils/package.json +++ b/ts/packages/utils/webSocketUtils/package.json @@ -13,7 +13,8 @@ "author": "Microsoft", "type": "module", "exports": { - ".": "./dist/index.js" + ".": "./dist/index.js", + "./originAllowlist": "./dist/originAllowlist.js" }, "files": [ "dist", diff --git a/ts/packages/utils/webSocketUtils/src/originAllowlist.ts b/ts/packages/utils/webSocketUtils/src/originAllowlist.ts new file mode 100644 index 000000000..fa94e6c19 --- /dev/null +++ b/ts/packages/utils/webSocketUtils/src/originAllowlist.ts @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +/** + * Build the Origin gate used by per-agent WebSocket bridges that bind to + * an ephemeral, loopback-only port via the dispatcher's PortRegistrar. + * + * Every predicate returned from this factory shares the same baseline: + * + * - `http(s)://localhost(:port)`, `http(s)://127.0.0.1(:port)`, and + * `http(s)://[::1](:port)` (also the unbracketed `::1` form) are + * allowed (loopback web clients on either IPv4 or IPv6; Node's URL + * parser preserves IPv6 brackets in `hostname`, so we match the + * bracketed form, while the unbracketed form is accepted defensively + * against URL parser/serializer differences across runtimes). + * - A missing/empty/`"null"` Origin is allowed — Node `ws` clients and + * other non-browser callers don't send Origin, and the listener binds + * to loopback so this remains OS-level restricted. + * - All other `http(s)` origins are rejected; non-`http(s)` schemes are + * rejected unless explicitly named in {@link extensionSchemes}. + * + * Agents layer in their own additional client surfaces (browser + * extensions, VS Code webviews, etc.) by passing scheme prefixes via + * {@link OriginAllowlistOptions.extensionSchemes}. Each entry is + * matched as a literal prefix against the Origin string (e.g. + * `"chrome-extension://"` accepts `chrome-extension://abc123`). + * + * Returned predicates are pure and reusable; build once per agent and + * call from the WS server's `verifyClient` so denied clients are + * rejected with HTTP 403 before the `connection` event fires. + */ +export type OriginAllowlistOptions = { + /** + * Additional URL-scheme prefixes (e.g. `"chrome-extension://"`, + * `"vscode-webview://"`) whose origins should be accepted. Matched + * as a literal prefix on the raw Origin string, so callers must + * include the trailing `://`. + */ + extensionSchemes?: readonly string[]; +}; + +/** + * Returns a predicate that decides whether an incoming WebSocket + * upgrade's `Origin` header should be accepted. See + * {@link OriginAllowlistOptions} for the shared policy. + */ +export function createAgentOriginAllowlist( + options: OriginAllowlistOptions = {}, +): (origin: string | undefined) => boolean { + const schemes = options.extensionSchemes ?? []; + return (origin: string | undefined): boolean => { + if (origin === undefined || origin === "" || origin === "null") { + // No Origin header: legitimate for Node `ws` and other + // non-browser clients. + return true; + } + for (const scheme of schemes) { + if (origin.startsWith(scheme)) { + return true; + } + } + try { + const u = new URL(origin); + if (u.protocol !== "http:" && u.protocol !== "https:") { + return false; + } + // Also accept the unbracketed `::1` for robustness against + // URL parser/serializer differences across runtimes (other + // SSRF guards in the repo, e.g. + // examples/workflow/engine/src/builtinTasks.ts, accept both). + return ( + u.hostname === "localhost" || + u.hostname === "127.0.0.1" || + u.hostname === "[::1]" || + u.hostname === "::1" + ); + } catch { + return false; + } + }; +}