Refactor session state machine and zod boundaries

proxy-server/src/launchd/agent.ts

@@ -4,12 +4,12 @@ import { dirname, join, resolve } from "node:path";
 import { homedir } from "node:os";
 import { Command } from "commander";
 import { build as buildPlist, parse as parsePlist } from "plist";
+import { z } from "zod";
 import type { Logger } from "copilot-sdk-proxy";
 import type { ProviderMode } from "copilot-sdk-proxy";
 import { isProviderName } from "../cli-validators.js";
 import { patchSettings, restoreSettings } from "../settings-patcher/index.js";
 import { defaultExec, type ExecFn } from "../utils/child-process.js";
-import { isRecord } from "../utils/type-guards.js";
 
 export const AGENT_LABEL = "com.xcode-copilot-server";
 
@@ -105,6 +105,10 @@ interface ParsedPlistArgs {
   autoPatch: boolean;
 }
 
+const PlistArgsSchema = z.object({
+  ProgramArguments: z.array(z.unknown()),
+});
+
 export function parsePlistArgs(plistContent: string): ParsedPlistArgs {
   let raw: unknown;
   try {
@@ -114,19 +118,14 @@ export function parsePlistArgs(plistContent: string): ParsedPlistArgs {
     return { proxy: null, autoPatch: false };
   }
 
-  if (!isRecord(raw)) {
-    return { proxy: null, autoPatch: false };
-  }
-
-  const parsed = raw;
-  const args = parsed["ProgramArguments"];
-  if (!Array.isArray(args)) {
+  const result = PlistArgsSchema.safeParse(raw);
+  if (!result.success) {
     return { proxy: null, autoPatch: false };
   }
 
-  const flagArgs = args
-    .slice(2)
-    .filter((a): a is string => typeof a === "string");
+  const flagArgs = result.data.ProgramArguments.slice(2).filter(
+    (a): a is string => typeof a === "string",
+  );
 
   const cmd = new Command()
     .exitOverride()

proxy-server/src/providers/shared/streaming-core.ts

@@ -7,7 +7,6 @@ import type {
 } from "copilot-sdk-proxy";
 import { createCommonEventHandler } from "copilot-sdk-proxy";
 import type { ToolBridgeState } from "../../tool-bridge/state.js";
-import { isRecord } from "../../utils/type-guards.js";
 import { BRIDGE_TOOL_PREFIX } from "../../tool-bridge/bridge-constants.js";
 
 // Xcode sends tool names without the bridge prefix.
@@ -17,6 +16,14 @@ function stripBridgePrefix(name: string): string {
     : name;
 }
 
+// Subset of the SDK's tool request we consume. arguments is a JSON object
+// (the SDK types it as Record<string, unknown>), not an arbitrary unknown.
+interface ToolRequest {
+  toolCallId: string;
+  name: string;
+  arguments?: Record<string, unknown>;
+}
+
 export interface StrippedToolRequest {
   toolCallId: string;
   name: string;
@@ -123,9 +130,7 @@ class StreamingHandler {
     });
   }
 
-  private stripAndNormalize(
-    requests: { toolCallId: string; name: string; arguments?: unknown }[],
-  ): StrippedToolRequest[] {
+  private stripAndNormalize(requests: ToolRequest[]): StrippedToolRequest[] {
     const filtered = this.hasBridge
       ? requests.filter((tr) => tr.name.startsWith(BRIDGE_TOOL_PREFIX))
       : requests;
@@ -134,20 +139,18 @@ class StreamingHandler {
       const resolved = this.state.toolCache.resolveToolName(
         stripBridgePrefix(tr.name),
       );
-      const args: Record<string, unknown> = isRecord(tr.arguments)
-        ? tr.arguments
-        : {};
       return {
         toolCallId: tr.toolCallId,
         name: resolved,
-        arguments: this.state.toolCache.normalizeArgs(resolved, args),
+        arguments: this.state.toolCache.normalizeArgs(
+          resolved,
+          tr.arguments ?? {},
+        ),
       };
     });
   }
 
-  private onMessage(data: {
-    toolRequests?: { toolCallId: string; name: string; arguments?: unknown }[];
-  }): void {
+  private onMessage(data: { toolRequests?: ToolRequest[] }): void {
     if (!data.toolRequests || data.toolRequests.length === 0) {
       const r = this.getReply();
       if (r) this.common.flushDeltas();

proxy-server/src/settings-patcher/codex.ts

@@ -18,23 +18,14 @@ import type {
 } from "./types.js";
 import { extractLocalhostPort } from "./url-utils.js";
 import { defaultExec, type ExecFn } from "../utils/child-process.js";
-import { isRecord } from "../utils/type-guards.js";
+import { z } from "zod";
 
-interface EnvBackup {
-  OPENAI_BASE_URL: string | null;
-  OPENAI_API_KEY: string | null;
-}
+const EnvBackupSchema = z.object({
+  OPENAI_BASE_URL: z.string().nullable(),
+  OPENAI_API_KEY: z.string().nullable(),
+});
 
-function isEnvBackup(value: unknown): value is EnvBackup {
-  if (!isRecord(value)) return false;
-  return (
-    "OPENAI_BASE_URL" in value &&
-    (typeof value.OPENAI_BASE_URL === "string" ||
-      value.OPENAI_BASE_URL === null) &&
-    "OPENAI_API_KEY" in value &&
-    (typeof value.OPENAI_API_KEY === "string" || value.OPENAI_API_KEY === null)
-  );
-}
+type EnvBackup = z.infer<typeof EnvBackupSchema>;
 
 function defaultCodexBackupPath(): string {
   return join(
@@ -136,11 +127,12 @@ async function readEnvBackup(
     logger.warn("Backup file contains invalid JSON, unsetting env vars");
     return null;
   }
-  if (!isEnvBackup(parsed)) {
+  const result = EnvBackupSchema.safeParse(parsed);
+  if (!result.success) {
     logger.warn("Backup file has unexpected shape, unsetting env vars");
     return null;
   }
-  return parsed;
+  return result.data;
 }
 
 export async function restoreCodexSettings(

proxy-server/src/tool-bridge/routes.ts

@@ -3,7 +3,7 @@ import { z } from "zod";
 import type { ToolStateProvider } from "../conversation-manager.js";
 import type { Logger } from "copilot-sdk-proxy";
 import { BRIDGE_SERVER_NAME } from "./bridge-constants.js";
-import { isRecord, errorMessage } from "../utils/type-guards.js";
+import { errorMessage } from "../utils/type-guards.js";
 import {
   JSONRPC_PARSE_ERROR,
   JSONRPC_INVALID_PARAMS,
@@ -18,6 +18,13 @@ const JsonRpcRequestSchema = z.object({
   params: z.record(z.string(), z.unknown()).optional(),
 });
 
+// A non-object or missing arguments field is coerced to {} rather than
+// rejected, matching what tools/call callers send in practice.
+const ToolCallParamsSchema = z.object({
+  name: z.string().min(1),
+  arguments: z.record(z.string(), z.unknown()).default({}).catch({}),
+});
+
 function jsonRpcResult(id: number | string, result: unknown) {
   return { jsonrpc: "2.0" as const, id, result };
 }
@@ -74,14 +81,15 @@ async function handleToolCall(
     return jsonRpcError(id, JSONRPC_INTERNAL_ERROR, "Conversation not found");
   }
 
-  const rawName = params?.["name"];
-  const name = typeof rawName === "string" ? rawName : undefined;
-  const rawArgs = params?.["arguments"];
-  const args: Record<string, unknown> = isRecord(rawArgs) ? rawArgs : {};
-
-  if (!name) {
-    return jsonRpcError(id, JSONRPC_INVALID_PARAMS, "Missing tool name");
+  const parsed = ToolCallParamsSchema.safeParse(params);
+  if (!parsed.success) {
+    return jsonRpcError(
+      id,
+      JSONRPC_INVALID_PARAMS,
+      "Missing or invalid tool name",
+    );
   }
+  const { name, arguments: args } = parsed.data;
 
   const resolved = state.toolCache.resolveToolName(name);
   if (resolved !== name) {

proxy-server/src/tool-bridge/session-lifecycle.ts

@@ -1,41 +1,80 @@
 import type { ToolRouter } from "./tool-router.js";
+import { assertNever } from "../utils/type-guards.js";
+
+// A session is idle, actively streaming, or finished with an error. Holding one
+// value means "active" and "errored" can never be set at the same time.
+export type SessionStatus = "idle" | "active" | "errored";
 
 export class SessionLifecycle {
   private readonly toolRouter: ToolRouter;
-  private _sessionActive = false;
-  private _hadError = false;
+  private _status: SessionStatus = "idle";
 
   constructor(toolRouter: ToolRouter) {
     this.toolRouter = toolRouter;
   }
 
+  get status(): SessionStatus {
+    return this._status;
+  }
+
   get sessionActive(): boolean {
-    return this._sessionActive;
+    switch (this._status) {
+      case "active":
+        return true;
+      case "idle":
+      case "errored":
+        return false;
+      default:
+        return assertNever(this._status);
+    }
   }
 
   get hadError(): boolean {
-    return this._hadError;
+    switch (this._status) {
+      case "errored":
+        return true;
+      case "idle":
+      case "active":
+        return false;
+      default:
+        return assertNever(this._status);
+    }
   }
 
   markSessionActive(): void {
     // Clear leftover entries from abandoned tool cycles so they don't
     // sit at the front of the FIFO queue and bind to wrong call IDs.
     this.toolRouter.rejectAll("New session cycle");
-    this._sessionActive = true;
+    this._status = "active";
   }
 
   markSessionErrored(): void {
-    this._hadError = true;
+    this._status = "errored";
   }
 
   markSessionInactive(): void {
     // Don't reject expected entries here. The session goes idle before the
     // bridge's tools/call requests arrive, so they are still needed.
-    this._sessionActive = false;
+    this.deactivate();
   }
 
   cleanup(): void {
-    this._sessionActive = false;
     this.toolRouter.rejectAll("Session cleanup");
+    this.deactivate();
+  }
+
+  // Keep an errored status so callers can still see the failure after the
+  // stream winds down. Only an active session falls back to idle.
+  private deactivate(): void {
+    switch (this._status) {
+      case "active":
+        this._status = "idle";
+        return;
+      case "idle":
+      case "errored":
+        return;
+      default:
+        assertNever(this._status);
+    }
   }
 }

proxy-server/src/utils/type-guards.ts

@@ -1,15 +1,13 @@
-export function isRecord(value: unknown): value is Record<string, unknown> {
-  return typeof value === "object" && value !== null && !Array.isArray(value);
+// Forces callers to handle every case of a union, so if a new variant is added,
+// the unhandled value stops being `never` and this call fails to compile.
+export function assertNever(value: never): never {
+  throw new Error(`Unhandled case: ${String(value)}`);
 }
 
 export function errorMessage(err: unknown): string {
   return err instanceof Error ? err.message : String(err);
 }
 
 export function isErrnoException(err: unknown): err is NodeJS.ErrnoException {
-  return (
-    err instanceof Error &&
-    "code" in err &&
-    typeof (err as { code: unknown }).code === "string"
-  );
+  return err instanceof Error && "code" in err && typeof err.code === "string";
 }

proxy-server/test/tool-bridge/mcp-routes.test.ts

@@ -218,7 +218,37 @@ describe("POST /mcp/:convId — tools/call", () => {
     const body = res.json();
     expect(body.id).toBe(6);
     expect(body.error.code).toBe(JSONRPC_INVALID_PARAMS);
-    expect(body.error.message).toContain("Missing tool name");
+    expect(body.error.message).toContain("tool name");
+  });
+
+  it("coerces a non-object arguments field to an empty object", async () => {
+    const conv = manager.create();
+    conv.state.toolCache.cacheTools([
+      {
+        name: "Read",
+        description: "Read a file",
+        input_schema: { type: "object", properties: {} },
+      },
+    ]);
+
+    vi.spyOn(conv.state.toolRouter, "registerMCPRequest").mockImplementation(
+      (_name, resolve) => {
+        resolve("ok");
+      },
+    );
+
+    const res = await app.inject({
+      method: "POST",
+      url: `/mcp/${conv.id}`,
+      payload: jsonRpc("tools/call", 8, {
+        name: "Read",
+        arguments: "not an object",
+      }),
+    });
+    expect(res.statusCode).toBe(200);
+    const body = res.json();
+    expect(body.id).toBe(8);
+    expect(body.result.content).toEqual([{ type: "text", text: "ok" }]);
   });
 
   it("returns error for unknown conversation", async () => {