pi-soly 1.9.3 → 1.11.0

package/mcp/__tests__/ui-session-messages.upstream-test.ts

@@ -0,0 +1,72 @@
+import { describe, it, expect } from "vitest";
+import { parseUiPromptHandoff, type UiSessionMessages } from "../types.ts";
+
+describe("UiSessionMessages", () => {
+  describe("type structure", () => {
+    it("can create empty session messages", () => {
+      const messages: UiSessionMessages = {
+        prompts: [],
+        notifications: [],
+        intents: [],
+      };
+
+      expect(messages.prompts).toHaveLength(0);
+      expect(messages.notifications).toHaveLength(0);
+      expect(messages.intents).toHaveLength(0);
+    });
+
+    it("can store prompts", () => {
+      const messages: UiSessionMessages = {
+        prompts: ["What is the weather?", "Tell me more"],
+        notifications: [],
+        intents: [],
+      };
+
+      expect(messages.prompts).toHaveLength(2);
+      expect(messages.prompts[0]).toBe("What is the weather?");
+    });
+
+    it("can store notifications", () => {
+      const messages: UiSessionMessages = {
+        prompts: [],
+        notifications: ["Task completed", "Error occurred"],
+        intents: [],
+      };
+
+      expect(messages.notifications).toHaveLength(2);
+    });
+
+    it("can store intents with params", () => {
+      const messages: UiSessionMessages = {
+        prompts: [],
+        notifications: [],
+        intents: [
+          { intent: "get_forecast", params: { days: 7, location: "NYC" } },
+          { intent: "refresh" },
+        ],
+      };
+
+      expect(messages.intents).toHaveLength(2);
+      expect(messages.intents[0].intent).toBe("get_forecast");
+      expect(messages.intents[0].params).toEqual({ days: 7, location: "NYC" });
+      expect(messages.intents[1].params).toBeUndefined();
+    });
+  });
+
+  describe("named handoff envelopes", () => {
+    it("parses canonical intent-newline-json payloads", () => {
+      expect(
+        parseUiPromptHandoff('visualization_annotations_submitted\n{"visualizationId":"flow","annotations":[]}')
+      ).toEqual({
+        intent: "visualization_annotations_submitted",
+        params: { visualizationId: "flow", annotations: [] },
+        raw: 'visualization_annotations_submitted\n{"visualizationId":"flow","annotations":[]}',
+      });
+    });
+
+    it("ignores free-form prompts", () => {
+      expect(parseUiPromptHandoff("Please analyze this chart")).toBeUndefined();
+      expect(parseUiPromptHandoff("visualization_annotations_submitted {}")).toBeUndefined();
+    });
+  });
+});

package/mcp/__tests__/ui-streaming.upstream-test.ts

@@ -0,0 +1,543 @@
+import { describe, it, expect, vi, afterEach } from "vitest";
+import http from "node:http";
+import { McpServerManager } from "../server-manager.ts";
+import { startUiServer, type UiServerOptions, type UiServerHandle } from "../ui-server.ts";
+import type { ConsentManager } from "../consent-manager.ts";
+import {
+  UI_STREAM_HOST_CONTEXT_KEY,
+  UI_STREAM_STRUCTURED_CONTENT_KEY,
+  SERVER_STREAM_RESULT_PATCH_METHOD,
+  getVisualizationStreamEnvelope,
+  getUiStreamHostContext,
+  serverStreamResultPatchNotificationSchema,
+  type UiResourceContent,
+  type VisualizationStreamEnvelope,
+} from "../types.ts";
+
+// Helper to connect to SSE and collect events
+function connectSSE(
+  url: string,
+  onEvent: (name: string, data: unknown, eventId?: string) => void,
+  headers: Record<string, string> = {}
+): Promise<{ close: () => void }> {
+  return new Promise((resolve, reject) => {
+    const parsed = new URL(url);
+    const req = http.request(
+      {
+        hostname: parsed.hostname,
+        port: parsed.port,
+        path: parsed.pathname + parsed.search,
+        method: "GET",
+        headers: { Accept: "text/event-stream", ...headers },
+      },
+      (res) => {
+        if (res.statusCode !== 200) {
+          reject(new Error(`SSE connection failed: ${res.statusCode}`));
+          return;
+        }
+        let buffer = "";
+        let eventName = "message";
+        let eventId: string | undefined;
+        res.on("data", (chunk) => {
+          buffer += chunk.toString();
+          const lines = buffer.split("\n");
+          buffer = lines.pop() ?? "";
+          for (const line of lines) {
+            if (line.startsWith("id: ")) {
+              eventId = line.slice(4);
+            } else if (line.startsWith("event: ")) {
+              eventName = line.slice(7);
+            } else if (line.startsWith("data: ")) {
+              try {
+                onEvent(eventName, JSON.parse(line.slice(6)), eventId);
+              } catch {
+                onEvent(eventName, line.slice(6), eventId);
+              }
+              eventName = "message";
+              eventId = undefined;
+            }
+          }
+        });
+        resolve({ close: () => req.destroy() });
+      }
+    );
+    req.on("error", reject);
+    req.end();
+  });
+}
+
+function createMockResource(): UiResourceContent {
+  return {
+    uri: "ui://test/widget",
+    html: "<h1>Test</h1>",
+    mimeType: "text/html",
+    meta: { permissions: {} },
+  };
+}
+
+function createServerOptions(overrides: Partial<UiServerOptions> = {}): UiServerOptions {
+  return {
+    serverName: "test-server",
+    toolName: "test_tool",
+    toolArgs: { key: "value" },
+    resource: createMockResource(),
+    manager: {} as McpServerManager,
+    consentManager: {} as ConsentManager,
+    ...overrides,
+  };
+}
+
+describe("UI Streaming", () => {
+  describe("getUiStreamHostContext", () => {
+    it("returns undefined for missing host context", () => {
+      expect(getUiStreamHostContext(undefined)).toBeUndefined();
+      expect(getUiStreamHostContext({})).toBeUndefined();
+    });
+
+    it("returns undefined for invalid stream context", () => {
+      expect(getUiStreamHostContext({ [UI_STREAM_HOST_CONTEXT_KEY]: "invalid" })).toBeUndefined();
+      expect(getUiStreamHostContext({ [UI_STREAM_HOST_CONTEXT_KEY]: { mode: "invalid" } })).toBeUndefined();
+    });
+
+    it("parses valid eager stream context", () => {
+      const context = getUiStreamHostContext({
+        [UI_STREAM_HOST_CONTEXT_KEY]: {
+          mode: "eager",
+          streamId: "abc-123",
+          intermediateResultPatches: false,
+          partialInput: false,
+        },
+      });
+      expect(context).toEqual({
+        mode: "eager",
+        streamId: "abc-123",
+        intermediateResultPatches: false,
+        partialInput: false,
+      });
+    });
+
+    it("parses valid stream-first context", () => {
+      const context = getUiStreamHostContext({
+        [UI_STREAM_HOST_CONTEXT_KEY]: {
+          mode: "stream-first",
+          streamId: "xyz-789",
+          intermediateResultPatches: true,
+          partialInput: false,
+        },
+      });
+      expect(context).toEqual({
+        mode: "stream-first",
+        streamId: "xyz-789",
+        intermediateResultPatches: true,
+        partialInput: false,
+      });
+    });
+  });
+
+  describe("getVisualizationStreamEnvelope", () => {
+    it("returns undefined for non-object input", () => {
+      expect(getVisualizationStreamEnvelope(undefined)).toBeUndefined();
+      expect(getVisualizationStreamEnvelope(null)).toBeUndefined();
+      expect(getVisualizationStreamEnvelope("string")).toBeUndefined();
+      expect(getVisualizationStreamEnvelope([])).toBeUndefined();
+    });
+
+    it("returns undefined when envelope key is missing", () => {
+      expect(getVisualizationStreamEnvelope({})).toBeUndefined();
+      expect(getVisualizationStreamEnvelope({ other: "data" })).toBeUndefined();
+    });
+
+    it("returns undefined for invalid envelope", () => {
+      expect(getVisualizationStreamEnvelope({
+        [UI_STREAM_STRUCTURED_CONTENT_KEY]: { streamId: "abc" }, // missing required fields
+      })).toBeUndefined();
+    });
+
+    it("parses valid envelope", () => {
+      const envelope = getVisualizationStreamEnvelope({
+        [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+          streamId: "stream-1",
+          sequence: 0,
+          frameType: "patch",
+          phase: "shell",
+          status: "ok",
+        },
+      });
+      expect(envelope).toEqual({
+        streamId: "stream-1",
+        sequence: 0,
+        frameType: "patch",
+        phase: "shell",
+        status: "ok",
+      });
+    });
+
+    it("parses envelope with optional fields", () => {
+      const envelope = getVisualizationStreamEnvelope({
+        [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+          streamId: "stream-2",
+          sequence: 5,
+          frameType: "checkpoint",
+          phase: "detail",
+          status: "ok",
+          message: "Checkpoint ready",
+          spec: { kind: "mermaid" },
+          checkpoint: { kind: "mermaid", code: "graph LR" },
+        },
+      });
+      expect(envelope?.message).toBe("Checkpoint ready");
+      expect(envelope?.spec).toEqual({ kind: "mermaid" });
+      expect(envelope?.checkpoint).toEqual({ kind: "mermaid", code: "graph LR" });
+    });
+  });
+
+  describe("McpServerManager stream listeners", () => {
+    function attachNotificationHandler(manager: McpServerManager, serverName = "test-server") {
+      const client = { setNotificationHandler: vi.fn() };
+      (manager as unknown as {
+        attachAdapterNotificationHandlers: (serverName: string, client: { setNotificationHandler: typeof client.setNotificationHandler }) => void;
+      }).attachAdapterNotificationHandlers(serverName, client);
+      expect(client.setNotificationHandler).toHaveBeenCalledOnce();
+      return client.setNotificationHandler.mock.calls[0][1] as (notification: {
+        method: string;
+        params: {
+          streamToken: string;
+          result: { content?: unknown[]; structuredContent?: Record<string, unknown> };
+        };
+      }) => void;
+    }
+
+    it("routes notifications to the matching listener", () => {
+      const manager = new McpServerManager();
+      const listener = vi.fn();
+      const handleNotification = attachNotificationHandler(manager, "server-a");
+
+      manager.registerUiStreamListener("token-123", listener);
+      const notification = {
+        method: SERVER_STREAM_RESULT_PATCH_METHOD,
+        params: {
+          streamToken: "token-123",
+          result: { content: [{ type: "text", text: "patch" }] },
+        },
+      };
+
+      handleNotification(notification);
+
+      expect(listener).toHaveBeenCalledTimes(1);
+      expect(listener).toHaveBeenCalledWith("server-a", notification.params);
+    });
+
+    it("does not call removed listeners", () => {
+      const manager = new McpServerManager();
+      const listener = vi.fn();
+      const handleNotification = attachNotificationHandler(manager);
+
+      manager.registerUiStreamListener("token-456", listener);
+      manager.removeUiStreamListener("token-456");
+
+      handleNotification({
+        method: SERVER_STREAM_RESULT_PATCH_METHOD,
+        params: {
+          streamToken: "token-456",
+          result: { content: [{ type: "text", text: "patch" }] },
+        },
+      });
+
+      expect(listener).not.toHaveBeenCalled();
+    });
+
+    it("keeps multiple listeners isolated by stream token", () => {
+      const manager = new McpServerManager();
+      const listener1 = vi.fn();
+      const listener2 = vi.fn();
+      const handleNotification = attachNotificationHandler(manager, "server-b");
+
+      manager.registerUiStreamListener("token-1", listener1);
+      manager.registerUiStreamListener("token-2", listener2);
+
+      handleNotification({
+        method: SERVER_STREAM_RESULT_PATCH_METHOD,
+        params: {
+          streamToken: "token-2",
+          result: { content: [{ type: "text", text: "patch-2" }] },
+        },
+      });
+
+      expect(listener1).not.toHaveBeenCalled();
+      expect(listener2).toHaveBeenCalledTimes(1);
+      expect(listener2).toHaveBeenCalledWith("server-b", {
+        streamToken: "token-2",
+        result: { content: [{ type: "text", text: "patch-2" }] },
+      });
+    });
+  });
+
+  describe("UiServer SSE streaming", () => {
+    let handle: UiServerHandle | null = null;
+
+    afterEach(() => {
+      if (handle) {
+        handle.close("test-cleanup");
+        handle = null;
+      }
+    });
+
+    it("sends result-patch events with stream envelope", async () => {
+      handle = await startUiServer(createServerOptions({
+        hostContext: {
+          [UI_STREAM_HOST_CONTEXT_KEY]: {
+            mode: "stream-first",
+            streamId: "test-stream",
+            intermediateResultPatches: true,
+            partialInput: false,
+          },
+        },
+      }));
+
+      const events: Array<{ name: string; data: unknown; id?: string }> = [];
+      const sse = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        (name, data, id) => events.push({ name, data, id })
+      );
+
+      // Send a result patch
+      handle.sendResultPatch({
+        content: [{ type: "text", text: "Streaming..." }],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "test-stream",
+            sequence: 1,
+            frameType: "patch",
+            phase: "shell",
+            status: "ok",
+            message: "Building shell",
+          },
+        },
+      });
+
+      // Wait for event propagation
+      await new Promise((r) => setTimeout(r, 50));
+      sse.close();
+
+      const patchEvents = events.filter((e) => e.name === "result-patch");
+      expect(patchEvents).toHaveLength(1);
+
+      const envelope = getVisualizationStreamEnvelope(
+        (patchEvents[0].data as { structuredContent?: unknown })?.structuredContent
+      );
+      expect(envelope?.frameType).toBe("patch");
+      expect(envelope?.phase).toBe("shell");
+    });
+
+    it("assigns sequential event IDs", async () => {
+      handle = await startUiServer(createServerOptions());
+
+      const eventIds: string[] = [];
+      const sse = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        (_name, _data, id) => { if (id) eventIds.push(id); }
+      );
+
+      handle.sendResultPatch({ content: [] });
+      handle.sendResultPatch({ content: [] });
+      handle.sendResultPatch({ content: [] });
+
+      await new Promise((r) => setTimeout(r, 50));
+      sse.close();
+
+      const ids = eventIds.map(Number).filter((n) => !Number.isNaN(n));
+      expect(ids).toHaveLength(3);
+      for (let i = 1; i < ids.length; i++) {
+        expect(ids[i]).toBeGreaterThan(ids[i - 1]);
+      }
+    });
+
+    it("replays events from Last-Event-ID", async () => {
+      handle = await startUiServer(createServerOptions());
+
+      // First connection to send some events
+      const firstEvents: Array<{ name: string; id?: string }> = [];
+      const sse1 = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        (name, _data, id) => firstEvents.push({ name, id })
+      );
+
+      handle.sendResultPatch({ content: [{ type: "text", text: "patch-1" }] });
+      handle.sendResultPatch({ content: [{ type: "text", text: "patch-2" }] });
+      handle.sendResultPatch({ content: [{ type: "text", text: "patch-3" }] });
+
+      await new Promise((r) => setTimeout(r, 50));
+      sse1.close();
+
+      // Get the ID of the first patch event
+      const patchEvents = firstEvents.filter((e) => e.name === "result-patch");
+      expect(patchEvents.length).toBe(3);
+      const firstPatchId = patchEvents[0].id;
+      expect(firstPatchId).toBeDefined();
+
+      // Second connection with Last-Event-ID should replay from that point
+      const replayedEvents: Array<{ name: string; id?: string }> = [];
+      const sse2 = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        (name, _data, id) => replayedEvents.push({ name, id }),
+        { "Last-Event-ID": firstPatchId! }
+      );
+
+      await new Promise((r) => setTimeout(r, 50));
+      sse2.close();
+
+      // Should replay events AFTER the provided ID
+      const replayedPatches = replayedEvents.filter((e) => e.name === "result-patch");
+      expect(replayedPatches.length).toBe(2); // patch-2 and patch-3
+    });
+
+    it("replays from latest checkpoint for fresh connections", async () => {
+      handle = await startUiServer(createServerOptions());
+
+      // Send patches and a checkpoint
+      const sse1 = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        () => {}
+      );
+
+      handle.sendResultPatch({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "s1", sequence: 0, frameType: "patch", phase: "shell", status: "ok",
+          },
+        },
+      });
+      handle.sendResultPatch({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "s1", sequence: 1, frameType: "checkpoint", phase: "detail", status: "ok",
+          },
+        },
+      });
+      handle.sendResultPatch({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "s1", sequence: 2, frameType: "patch", phase: "detail", status: "ok",
+          },
+        },
+      });
+
+      await new Promise((r) => setTimeout(r, 50));
+      sse1.close();
+
+      // Fresh connection (no Last-Event-ID) should start from checkpoint
+      const freshEvents: Array<{ data: unknown }> = [];
+      const sse2 = await connectSSE(
+        `http://localhost:${handle.port}/events?session=${handle.sessionToken}`,
+        (_name, data) => freshEvents.push({ data })
+      );
+
+      await new Promise((r) => setTimeout(r, 50));
+      sse2.close();
+
+      // Should have replayed checkpoint and subsequent patches
+      const envelopes = freshEvents
+        .map((e) => getVisualizationStreamEnvelope((e.data as { structuredContent?: unknown })?.structuredContent))
+        .filter(Boolean) as VisualizationStreamEnvelope[];
+
+      expect(envelopes).toHaveLength(2);
+      expect(envelopes.map((envelope) => envelope.frameType)).toEqual(["checkpoint", "patch"]);
+    });
+
+    it("tracks stream summary", async () => {
+      handle = await startUiServer(createServerOptions({
+        hostContext: {
+          [UI_STREAM_HOST_CONTEXT_KEY]: {
+            mode: "stream-first",
+            streamId: "summary-test",
+            intermediateResultPatches: true,
+            partialInput: false,
+          },
+        },
+      }));
+
+      handle.sendResultPatch({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "summary-test", sequence: 0, frameType: "patch", phase: "shell", status: "ok",
+          },
+        },
+      });
+      handle.sendResultPatch({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "summary-test", sequence: 1, frameType: "patch", phase: "structure", status: "ok",
+          },
+        },
+      });
+      handle.sendToolResult({
+        content: [],
+        structuredContent: {
+          [UI_STREAM_STRUCTURED_CONTENT_KEY]: {
+            streamId: "summary-test", sequence: 2, frameType: "final", phase: "settled", status: "ok",
+            message: "Complete",
+          },
+        },
+      });
+
+      const summary = handle.getStreamSummary();
+      expect(summary).toBeDefined();
+      expect(summary?.streamId).toBe("summary-test");
+      expect(summary?.mode).toBe("stream-first");
+      expect(summary?.frames).toBe(3);
+      expect(summary?.phases).toContain("shell");
+      expect(summary?.phases).toContain("structure");
+      expect(summary?.phases).toContain("settled");
+      expect(summary?.finalStatus).toBe("ok");
+      expect(summary?.lastMessage).toBe("Complete");
+    });
+  });
+
+  describe("Server stream result patch notification schema", () => {
+    it("validates correct notification format", () => {
+      const notification = {
+        method: SERVER_STREAM_RESULT_PATCH_METHOD,
+        params: {
+          streamToken: "token-123",
+          result: {
+            content: [{ type: "text", text: "test" }],
+            structuredContent: { data: "value" },
+          },
+        },
+      };
+
+      const parsed = serverStreamResultPatchNotificationSchema.safeParse(notification);
+      expect(parsed.success).toBe(true);
+    });
+
+    it("rejects notification with wrong method", () => {
+      const notification = {
+        method: "wrong/method",
+        params: {
+          streamToken: "token-123",
+          result: { content: [] },
+        },
+      };
+
+      const parsed = serverStreamResultPatchNotificationSchema.safeParse(notification);
+      expect(parsed.success).toBe(false);
+    });
+
+    it("rejects notification with missing streamToken", () => {
+      const notification = {
+        method: SERVER_STREAM_RESULT_PATCH_METHOD,
+        params: {
+          result: { content: [] },
+        },
+      };
+
+      const parsed = serverStreamResultPatchNotificationSchema.safeParse(notification);
+      expect(parsed.success).toBe(false);
+    });
+  });
+});

package/mcp/agent-dir.ts

@@ -0,0 +1,20 @@
+import { homedir } from "node:os";
+import { join, resolve } from "node:path";
+
+export function getAgentDir(): string {
+  const configured = process.env.PI_CODING_AGENT_DIR?.trim();
+  if (!configured) {
+    return join(homedir(), ".pi", "agent");
+  }
+  if (configured === "~") {
+    return homedir();
+  }
+  if (configured.startsWith("~/")) {
+    return resolve(homedir(), configured.slice(2));
+  }
+  return resolve(configured);
+}
+
+export function getAgentPath(...segments: string[]): string {
+  return join(getAgentDir(), ...segments);
+}