pi-soly 1.9.3 → 1.11.0

package/mcp/__tests__/proxy-modes-auto-auth.upstream-test.ts

@@ -0,0 +1,253 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  authenticate: vi.fn(),
+  supportsOAuth: vi.fn(),
+  lazyConnect: vi.fn(),
+  updateServerMetadata: vi.fn(),
+  updateMetadataCache: vi.fn(),
+  getFailureAgeSeconds: vi.fn(),
+  updateStatusBar: vi.fn(),
+}));
+
+vi.mock("../mcp-auth-flow.ts", () => ({
+  authenticate: mocks.authenticate,
+  supportsOAuth: mocks.supportsOAuth,
+}));
+
+vi.mock("../init.ts", () => ({
+  lazyConnect: mocks.lazyConnect,
+  updateServerMetadata: mocks.updateServerMetadata,
+  updateMetadataCache: mocks.updateMetadataCache,
+  getFailureAgeSeconds: mocks.getFailureAgeSeconds,
+  updateStatusBar: mocks.updateStatusBar,
+}));
+
+describe("proxy auto auth", () => {
+  beforeEach(() => {
+    vi.resetModules();
+    mocks.authenticate.mockReset().mockResolvedValue("authenticated");
+    mocks.supportsOAuth.mockReset().mockReturnValue(true);
+    mocks.lazyConnect.mockReset().mockResolvedValue(false);
+    mocks.updateServerMetadata.mockReset();
+    mocks.updateMetadataCache.mockReset();
+    mocks.getFailureAgeSeconds.mockReset().mockReturnValue(null);
+    mocks.updateStatusBar.mockReset();
+  });
+
+  it("auto-authenticates and retries executeConnect once", async () => {
+    const { executeConnect } = await import("../proxy-modes.ts");
+
+    let current: any;
+    const connected = {
+      status: "connected",
+      tools: [{ name: "search", description: "Search" }],
+      resources: [],
+    };
+
+    const manager = {
+      connect: vi
+        .fn()
+        .mockImplementationOnce(async () => {
+          current = { status: "needs-auth" };
+          return current;
+        })
+        .mockImplementationOnce(async () => {
+          current = connected;
+          return current;
+        }),
+      close: vi.fn(async () => {
+        current = undefined;
+      }),
+      getConnection: vi.fn(() => current),
+    };
+
+    const state = {
+      config: {
+        settings: { autoAuth: true, toolPrefix: "server" },
+        mcpServers: {
+          demo: { url: "https://api.example.com/mcp", auth: "oauth" },
+        },
+      },
+      manager,
+      toolMetadata: new Map(),
+      failureTracker: new Map(),
+      ui: { setStatus: vi.fn() },
+    } as any;
+
+    const result = await executeConnect(state, "demo");
+
+    expect(mocks.authenticate).toHaveBeenCalledWith(
+      "demo",
+      "https://api.example.com/mcp",
+      state.config.mcpServers.demo,
+    );
+    expect(manager.close).toHaveBeenCalledWith("demo");
+    expect(manager.connect).toHaveBeenCalledTimes(2);
+    expect(result.content[0].text).toContain("demo (1 tools)");
+  });
+
+  it("fails fast for non-ui browser auth when autoAuth is enabled", async () => {
+    const { executeConnect } = await import("../proxy-modes.ts");
+
+    const manager = {
+      connect: vi.fn(async () => ({ status: "needs-auth" })),
+      close: vi.fn(async () => {}),
+      getConnection: vi.fn(() => ({ status: "needs-auth" })),
+    };
+
+    const state = {
+      config: {
+        settings: { autoAuth: true },
+        mcpServers: {
+          demo: { url: "https://api.example.com/mcp", auth: "oauth" },
+        },
+      },
+      manager,
+      toolMetadata: new Map(),
+      failureTracker: new Map(),
+      ui: undefined,
+    } as any;
+
+    const result = await executeConnect(state, "demo");
+
+    expect(mocks.authenticate).not.toHaveBeenCalled();
+    expect(result.content[0].text).toContain("auth-start");
+    expect(result.content[0].text).toContain("/mcp-auth demo");
+  });
+
+  it("uses custom authRequiredMessage for non-ui autoAuth failures", async () => {
+    const { executeConnect } = await import("../proxy-modes.ts");
+
+    const state = {
+      config: {
+        settings: {
+          autoAuth: true,
+          authRequiredMessage: "Reconnect ${server} from the host app.",
+        },
+        mcpServers: {
+          demo: { url: "https://api.example.com/mcp", auth: "oauth" },
+        },
+      },
+      manager: {
+        connect: vi.fn(async () => ({ status: "needs-auth" })),
+        close: vi.fn(async () => {}),
+        getConnection: vi.fn(() => ({ status: "needs-auth" })),
+      },
+      toolMetadata: new Map(),
+      failureTracker: new Map(),
+      ui: undefined,
+    } as any;
+
+    const result = await executeConnect(state, "demo");
+
+    expect(mocks.authenticate).not.toHaveBeenCalled();
+    expect(result.content[0].text).toBe("Reconnect demo from the host app.");
+  });
+
+  it("runs URL elicitations returned by proxy tool calls", async () => {
+    const { UrlElicitationRequiredError } = await import("@modelcontextprotocol/sdk/types.js");
+    const { executeCall } = await import("../proxy-modes.ts");
+    const error = new UrlElicitationRequiredError([{
+      mode: "url",
+      message: "Connect your account",
+      elicitationId: "connect-1",
+      url: "https://example.com/connect",
+    }]);
+    const connection = {
+      status: "connected",
+      client: { callTool: vi.fn().mockRejectedValue(error) },
+    };
+    const manager = {
+      getConnection: vi.fn(() => connection),
+      handleUrlElicitationRequired: vi.fn().mockResolvedValue("accept"),
+      touch: vi.fn(),
+      incrementInFlight: vi.fn(),
+      decrementInFlight: vi.fn(),
+    };
+    const state = {
+      config: { settings: {}, mcpServers: { demo: { command: "demo" } } },
+      manager,
+      toolMetadata: new Map([["demo", [{
+        name: "demo_search",
+        originalName: "search",
+        description: "Search",
+        inputSchema: { type: "object", properties: {} },
+      }]]]),
+      failureTracker: new Map(),
+      completedUiSessions: [],
+    } as any;
+
+    const result = await executeCall(state, "demo_search", {}, "demo");
+
+    expect(manager.handleUrlElicitationRequired).toHaveBeenCalledWith("demo", error);
+    expect(result.details).toMatchObject({ error: "url_elicitation_required", action: "accept" });
+  });
+
+  it("auto-authenticates and retries executeCall once", async () => {
+    const { executeCall } = await import("../proxy-modes.ts");
+
+    let current: any = { status: "needs-auth" };
+    const connected = {
+      status: "connected",
+      client: {
+        callTool: vi.fn(async () => ({
+          isError: false,
+          content: [{ type: "text", text: "ok" }],
+        })),
+      },
+      tools: [{ name: "search", description: "Search" }],
+      resources: [],
+    };
+
+    const manager = {
+      connect: vi.fn(async () => {
+        current = connected;
+        return connected;
+      }),
+      close: vi.fn(async () => {
+        current = undefined;
+      }),
+      getConnection: vi.fn(() => current),
+      touch: vi.fn(),
+      incrementInFlight: vi.fn(),
+      decrementInFlight: vi.fn(),
+    };
+
+    const state = {
+      config: {
+        settings: { autoAuth: true, toolPrefix: "server" },
+        mcpServers: {
+          demo: { url: "https://api.example.com/mcp", auth: "oauth" },
+        },
+      },
+      manager,
+      toolMetadata: new Map([
+        [
+          "demo",
+          [
+            {
+              name: "demo_search",
+              originalName: "search",
+              description: "Search",
+              inputSchema: { type: "object", properties: {} },
+            },
+          ],
+        ],
+      ]),
+      failureTracker: new Map(),
+      ui: { setStatus: vi.fn() },
+      completedUiSessions: [],
+    } as any;
+
+    const result = await executeCall(state, "demo_search", { q: "hello" }, "demo");
+
+    expect(mocks.authenticate).toHaveBeenCalledWith(
+      "demo",
+      "https://api.example.com/mcp",
+      state.config.mcpServers.demo,
+    );
+    expect(manager.connect).toHaveBeenCalledTimes(1);
+    expect(result.content[0].text).toContain("ok");
+  });
+});

package/mcp/__tests__/proxy-modes-discovery.upstream-test.ts

@@ -0,0 +1,84 @@
+import { describe, expect, it } from "vitest";
+import { executeCall, executeSearch } from "../proxy-modes.ts";
+import type { McpExtensionState } from "../state.ts";
+
+function createState(): McpExtensionState {
+  return {
+    config: {
+      mcpServers: {
+        demo: { command: "npx", args: ["demo"] },
+      },
+    },
+    toolMetadata: new Map([
+      [
+        "demo",
+        [
+          {
+            name: "demo_search",
+            originalName: "search",
+            description: "Search demo records",
+            inputSchema: { type: "object", properties: {} },
+          },
+        ],
+      ],
+    ]),
+    manager: {
+      getConnection: () => undefined,
+    },
+    failureTracker: new Map(),
+  } as unknown as McpExtensionState;
+}
+
+describe("proxy discovery", () => {
+  it("searches MCP tools only", () => {
+    const result = executeSearch(createState(), "read");
+
+    expect(result.content[0].text).toBe('No tools matching "read"');
+    expect(result.details).toMatchObject({ count: 0, matches: [] });
+  });
+
+  it("rejects regex queries longer than the safety cap", () => {
+    const result = executeSearch(createState(), "a".repeat(257), true);
+
+    expect(result.details).toMatchObject({ error: "query_too_long", maxLength: 256 });
+  });
+
+  it("reports malformed regex queries separately from unsafe patterns", () => {
+    const result = executeSearch(createState(), "[", true);
+
+    expect(result.details).toMatchObject({ error: "invalid_pattern" });
+  });
+
+  it("rejects catastrophic-backtracking regex queries", () => {
+    const result = executeSearch(createState(), "(a+)+$", true);
+
+    expect(result.details).toMatchObject({ error: "unsafe_pattern", safetyStatus: "vulnerable" });
+  });
+
+  it("accepts safe regex queries", () => {
+    const result = executeSearch(createState(), "^demo_[a-z]+$", true);
+
+    expect(result.details).toMatchObject({ count: 1, query: "^demo_[a-z]+$" });
+  });
+
+  it("keeps non-regex searches unaffected by the regex length cap", () => {
+    const result = executeSearch(createState(), "search terms ".repeat(40), false);
+
+    expect(result.details).not.toMatchObject({ error: "query_too_long" });
+  });
+
+  it("tells callers to invoke native Pi tools directly", async () => {
+    const result = await executeCall(
+      createState(),
+      "read",
+      undefined,
+      undefined,
+      () => [{ name: "read", description: "Read a file" } as any],
+    );
+
+    expect(result.content[0].text).toBe(
+      '"read" is a native Pi tool. Call read directly instead of using mcp({ tool: "read" }).',
+    );
+    expect(result.details).toMatchObject({ error: "native_tool", requestedTool: "read" });
+  });
+});

package/mcp/__tests__/proxy-modes-manual-auth.upstream-test.ts

@@ -0,0 +1,91 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  completeAuthFromInput: vi.fn(),
+  startAuth: vi.fn(),
+  supportsOAuth: vi.fn(),
+  lazyConnect: vi.fn(),
+  updateServerMetadata: vi.fn(),
+  updateMetadataCache: vi.fn(),
+  getFailureAgeSeconds: vi.fn(),
+  updateStatusBar: vi.fn(),
+}));
+
+vi.mock("../mcp-auth-flow.ts", () => ({
+  authenticate: vi.fn(),
+  completeAuthFromInput: mocks.completeAuthFromInput,
+  startAuth: mocks.startAuth,
+  supportsOAuth: mocks.supportsOAuth,
+}));
+
+vi.mock("../init.ts", () => ({
+  lazyConnect: mocks.lazyConnect,
+  updateServerMetadata: mocks.updateServerMetadata,
+  updateMetadataCache: mocks.updateMetadataCache,
+  getFailureAgeSeconds: mocks.getFailureAgeSeconds,
+  updateStatusBar: mocks.updateStatusBar,
+}));
+
+function createState(overrides: Record<string, unknown> = {}) {
+  return {
+    config: {
+      settings: {},
+      mcpServers: {
+        demo: { url: "https://api.example.com/mcp", auth: "oauth" },
+        bearer: { url: "https://api.example.com/mcp", auth: "bearer" },
+      },
+    },
+    manager: { close: vi.fn(async () => {}) },
+    toolMetadata: new Map(),
+    failureTracker: new Map([["demo", Date.now()]]),
+    ...overrides,
+  } as any;
+}
+
+describe("manual OAuth proxy actions", () => {
+  beforeEach(() => {
+    vi.resetModules();
+    mocks.completeAuthFromInput.mockReset().mockResolvedValue("authenticated");
+    mocks.startAuth.mockReset().mockResolvedValue({
+      authorizationUrl: "https://auth.example.com/authorize?redirect_uri=http%3A%2F%2Flocalhost%3A19876%2Fcallback",
+    });
+    mocks.supportsOAuth.mockReset().mockImplementation((definition) => definition.auth === "oauth");
+    mocks.updateStatusBar.mockReset();
+  });
+
+  it("returns copyable instructions and authorization URL", async () => {
+    const { executeAuthStart } = await import("../proxy-modes.ts");
+    const state = createState();
+
+    const result = await executeAuthStart(state, "demo");
+
+    expect(mocks.startAuth).toHaveBeenCalledWith("demo", "https://api.example.com/mcp", state.config.mcpServers.demo);
+    expect(result.content[0].text).toContain("Open this URL in your local browser");
+    expect(result.content[0].text).toContain("https://auth.example.com/authorize");
+    expect(result.content[0].text).toContain("auth-complete");
+    expect(result.details).toMatchObject({ mode: "auth-start", server: "demo" });
+  });
+
+  it("rejects auth-start for non-OAuth servers", async () => {
+    const { executeAuthStart } = await import("../proxy-modes.ts");
+
+    const result = await executeAuthStart(createState(), "bearer");
+
+    expect(mocks.startAuth).not.toHaveBeenCalled();
+    expect(result.content[0].text).toContain("not configured for OAuth");
+    expect(result.details).toMatchObject({ error: "oauth_not_supported" });
+  });
+
+  it("completes auth from a copied redirect URL and resets connection state", async () => {
+    const { executeAuthComplete } = await import("../proxy-modes.ts");
+    const state = createState();
+
+    const result = await executeAuthComplete(state, "demo", "http://localhost:19876/callback?code=abc&state=state");
+
+    expect(mocks.completeAuthFromInput).toHaveBeenCalledWith("demo", "http://localhost:19876/callback?code=abc&state=state");
+    expect(state.manager.close).toHaveBeenCalledWith("demo");
+    expect(state.failureTracker.has("demo")).toBe(false);
+    expect(mocks.updateStatusBar).toHaveBeenCalledWith(state);
+    expect(result.content[0].text).toContain("OAuth authentication successful");
+  });
+});

package/mcp/__tests__/proxy-modes-ui-messages.upstream-test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from "vitest";
+import { executeUiMessages } from "../proxy-modes.ts";
+import type { McpExtensionState } from "../state.ts";
+
+function createState(prompts: string[]): McpExtensionState {
+  return {
+    completedUiSessions: [
+      {
+        serverName: "interactive-visualizer",
+        toolName: "show_visualization",
+        completedAt: new Date("2026-03-12T16:00:00Z"),
+        reason: "done",
+        messages: {
+          prompts,
+          notifications: [],
+          intents: [],
+        },
+      },
+    ],
+  } as McpExtensionState;
+}
+
+describe("executeUiMessages", () => {
+  it("normalizes canonical handoff prompts into structured intents", () => {
+    const state = createState([
+      'visualization_annotations_submitted\n{"visualizationId":"flow","annotations":[{"id":"a1","kind":"pin","text":"Check this"}]}',
+    ]);
+
+    const result = executeUiMessages(state);
+    expect(result.content[0]).toMatchObject({
+      type: "text",
+      text: expect.stringContaining("visualization_annotations_submitted"),
+    });
+    expect(result.content[0]).toMatchObject({
+      text: expect.not.stringContaining("### Prompts:\n- visualization_annotations_submitted"),
+    });
+    expect(result.details).toMatchObject({
+      intents: [
+        {
+          intent: "visualization_annotations_submitted",
+          params: {
+            visualizationId: "flow",
+            annotations: [{ id: "a1", kind: "pin", text: "Check this" }],
+          },
+        },
+      ],
+      handoffs: [
+        {
+          intent: "visualization_annotations_submitted",
+          params: {
+            visualizationId: "flow",
+            annotations: [{ id: "a1", kind: "pin", text: "Check this" }],
+          },
+        },
+      ],
+      cleared: true,
+    });
+    expect(state.completedUiSessions).toEqual([]);
+  });
+
+  it("preserves ordinary prompts as prompts", () => {
+    const state = createState(["Please analyze this flow"]);
+    const result = executeUiMessages(state);
+    expect(result.content[0]).toMatchObject({
+      text: expect.stringContaining("### Prompts:\n- Please analyze this flow"),
+    });
+    expect(result.details).toMatchObject({
+      prompts: ["Please analyze this flow"],
+      intents: [],
+    });
+  });
+});