From c14d9ded370cd42a391a8519b626f5a87baa5f72 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Mon, 1 Jun 2026 18:03:22 +0200 Subject: [PATCH 1/3] wip --- .../main/services/agent/auth-adapter.test.ts | 12 ++ .../src/main/services/agent/auth-adapter.ts | 2 + .../main/services/mcp-apps/service.test.ts | 104 +++++++++++++++ .../src/main/services/mcp-apps/service.ts | 119 +++++++++++++++++- apps/code/src/main/trpc/routers/mcp-apps.ts | 29 +++++ .../mcp-apps/components/McpAppHost.tsx | 53 +++++++- .../session-update/McpToolBlock.tsx | 63 +++++++++- apps/code/src/shared/types/mcp-apps.test.ts | 48 +++++++ apps/code/src/shared/types/mcp-apps.ts | 59 +++++++++ 9 files changed, 477 insertions(+), 12 deletions(-) create mode 100644 apps/code/src/main/services/mcp-apps/service.test.ts create mode 100644 apps/code/src/shared/types/mcp-apps.test.ts diff --git a/apps/code/src/main/services/agent/auth-adapter.test.ts b/apps/code/src/main/services/agent/auth-adapter.test.ts index 4d3aaf1ff7..cd66974ede 100644 --- a/apps/code/src/main/services/agent/auth-adapter.test.ts +++ b/apps/code/src/main/services/agent/auth-adapter.test.ts @@ -105,6 +105,18 @@ describe("AgentAuthAdapter", () => { ); }); + it("identifies as the PostHog Code consumer so the MCP server emits UI-app metadata", async () => { + const { servers } = await adapter.buildMcpServers(baseCredentials); + + const posthogServer = servers.find((s) => s.name === "posthog"); + expect(posthogServer).toBeDefined(); + expect(posthogServer?.headers).toEqual( + expect.arrayContaining([ + { name: "x-posthog-mcp-consumer", value: "posthog-code" }, + ]), + ); + }); + it("routes authenticated installed MCP servers through the proxy URL", async () => { mockFetch.mockResolvedValue({ ok: true, diff --git a/apps/code/src/main/services/agent/auth-adapter.ts b/apps/code/src/main/services/agent/auth-adapter.ts index 1cfa711fe0..259e6264be 100644 --- a/apps/code/src/main/services/agent/auth-adapter.ts +++ b/apps/code/src/main/services/agent/auth-adapter.ts @@ -99,6 +99,8 @@ export class AgentAuthAdapter { value: String(credentials.projectId), }, { name: "x-posthog-mcp-version", value: "2" }, + // Identify as the PostHog Code UI-apps host. + { name: "x-posthog-mcp-consumer", value: "posthog-code" }, ], }); diff --git a/apps/code/src/main/services/mcp-apps/service.test.ts b/apps/code/src/main/services/mcp-apps/service.test.ts new file mode 100644 index 0000000000..e4139e5489 --- /dev/null +++ b/apps/code/src/main/services/mcp-apps/service.test.ts @@ -0,0 +1,104 @@ +import { + LEGACY_RESOURCE_URI_META_KEY, + McpAppsServiceEvent, + type McpAppsToolCallUiDiscoveredEvent, + POSTHOG_EXEC_TOOL_KEY, +} from "@shared/types/mcp-apps"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../../utils/logger", () => ({ + logger: { + scope: () => ({ + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }), + }, +})); + +import { McpAppsService } from "./service"; + +function makeService(): McpAppsService { + const urlLauncher = { launch: vi.fn() }; + return new McpAppsService(urlLauncher as never); +} + +describe("McpAppsService — per-call exec UI", () => { + let service: McpAppsService; + + beforeEach(() => { + service = makeService(); + }); + + it("registers a per-call association from a modern _meta.ui.resourceUri", () => { + const events: McpAppsToolCallUiDiscoveredEvent[] = []; + service.on(McpAppsServiceEvent.ToolCallUiDiscovered, (e) => events.push(e)); + + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-1", { + content: [{ type: "text", text: "{}" }], + _meta: { ui: { resourceUri: "ui://posthog/insight" } }, + }); + + expect(service.hasUiForToolCall("call-1")).toBe(true); + expect(events).toEqual([ + { + toolCallId: "call-1", + toolKey: POSTHOG_EXEC_TOOL_KEY, + resourceUri: "ui://posthog/insight", + }, + ]); + }); + + it("supports the legacy flat resourceUri key", () => { + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-legacy", { + _meta: { [LEGACY_RESOURCE_URI_META_KEY]: "ui://posthog/legacy" }, + }); + expect(service.hasUiForToolCall("call-legacy")).toBe(true); + }); + + it("does nothing for non-exec tools", () => { + service.notifyToolResult("mcp__posthog__query", "call-2", { + _meta: { ui: { resourceUri: "ui://posthog/insight" } }, + }); + expect(service.hasUiForToolCall("call-2")).toBe(false); + }); + + it("ignores errored exec results", () => { + service.notifyToolResult( + POSTHOG_EXEC_TOOL_KEY, + "call-3", + { _meta: { ui: { resourceUri: "ui://posthog/insight" } } }, + true, + ); + expect(service.hasUiForToolCall("call-3")).toBe(false); + }); + + it("ignores exec results without a UI resource", () => { + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-4", { + content: [{ type: "text", text: "no ui here" }], + }); + expect(service.hasUiForToolCall("call-4")).toBe(false); + }); + + it("does not emit twice for the same call", () => { + const events: McpAppsToolCallUiDiscoveredEvent[] = []; + service.on(McpAppsServiceEvent.ToolCallUiDiscovered, (e) => events.push(e)); + + const result = { _meta: { ui: { resourceUri: "ui://posthog/insight" } } }; + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-5", result); + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-5", result); + + expect(events).toHaveLength(1); + }); + + it("evicts the association when the call is cancelled", () => { + service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-6", { + _meta: { ui: { resourceUri: "ui://posthog/insight" } }, + }); + expect(service.hasUiForToolCall("call-6")).toBe(true); + + service.notifyToolCancelled(POSTHOG_EXEC_TOOL_KEY, "call-6"); + expect(service.hasUiForToolCall("call-6")).toBe(false); + }); +}); diff --git a/apps/code/src/main/services/mcp-apps/service.ts b/apps/code/src/main/services/mcp-apps/service.ts index 46c89bb266..e8ebfc02a0 100644 --- a/apps/code/src/main/services/mcp-apps/service.ts +++ b/apps/code/src/main/services/mcp-apps/service.ts @@ -3,9 +3,12 @@ import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/ import type { Tool } from "@modelcontextprotocol/sdk/types.js"; import type { IUrlLauncher } from "@posthog/platform/url-launcher"; import { + BUILTIN_POSTHOG_SERVER_NAME, + EXEC_TOOL_NAME, type McpAppsDiscoveryCompleteEvent, McpAppsServiceEvent, type McpAppsServiceEvents, + type McpAppsToolCallUiDiscoveredEvent, type McpAppsToolCancelledEvent, type McpAppsToolInputEvent, type McpAppsToolResultEvent, @@ -14,6 +17,8 @@ import { type McpToolUiAssociation, type McpToolUiMeta, type McpUiResource, + POSTHOG_EXEC_TOOL_KEY, + resolveResultResourceUri, } from "@shared/types/mcp-apps"; import { inject, injectable } from "inversify"; import { MAIN_TOKENS } from "../../di/tokens"; @@ -36,6 +41,13 @@ export class McpAppsService extends TypedEventEmitter { private connections = new Map(); private resourceCache = new Map(); private toolAssociations = new Map(); + /** + * Per-tool-call UI associations for PostHog's built-in `exec` tool, keyed by + * `toolCallId`. The `exec` tool can't be discovered via `listTools()` (its UI + * `resourceUri` rides on each call's response `_meta`, not its registration), + * so associations here are populated lazily in {@link notifyToolResult}. + */ + private execToolCallAssociations = new Map(); private toolDefinitions = new Map(); private serverConfigs = new Map(); private pendingConnections = new Map>(); @@ -106,6 +118,16 @@ export class McpAppsService extends TypedEventEmitter { ]); for (const tool of toolsList.tools) { + // The built-in PostHog `exec` tool carries no registration-time UI + // resource (it resolves per call), but we still cache its definition so + // the app bridge can surface it when a per-call UI is rendered. + if ( + serverName === BUILTIN_POSTHOG_SERVER_NAME && + tool.name === EXEC_TOOL_NAME + ) { + this.toolDefinitions.set(POSTHOG_EXEC_TOOL_KEY, tool); + } + const uiMeta = (tool as McpToolUiMeta)._meta?.ui; if (!uiMeta?.resourceUri) continue; @@ -319,6 +341,47 @@ export class McpAppsService extends TypedEventEmitter { return has; } + /** + * Whether a UI app was resolved for a specific tool call. Only the built-in + * PostHog `exec` tool populates this (via {@link notifyToolResult}); for + * registration-discovered tools use {@link hasUiForTool}. + */ + hasUiForToolCall(toolCallId: string): boolean { + const has = this.execToolCallAssociations.has(toolCallId); + log.debug("hasUiForToolCall", { toolCallId, result: has }); + return has; + } + + /** + * Fetch the UI resource resolved for a specific `exec` tool call. Reuses the + * same lazy-fetch + cache path as registration-discovered tools. + */ + async getUiResourceForToolCall( + toolCallId: string, + ): Promise { + const association = this.execToolCallAssociations.get(toolCallId); + if (!association) { + log.debug("getUiResourceForToolCall: no association found", { + toolCallId, + }); + return null; + } + + const cached = this.resourceCache.get(association.resourceUri); + if (cached) return cached; + + const pendingFetch = this.pendingFetches.get(association.resourceUri); + if (pendingFetch) return pendingFetch; + + const fetchPromise = this.fetchUiResource(association); + this.pendingFetches.set(association.resourceUri, fetchPromise); + try { + return await fetchPromise; + } finally { + this.pendingFetches.delete(association.resourceUri); + } + } + getToolDefinition(toolKey: string): Tool | null { return this.toolDefinitions.get(toolKey) ?? null; } @@ -383,6 +446,14 @@ export class McpAppsService extends TypedEventEmitter { isError?: boolean, ): void { log.info("notifyToolResult", { toolKey, toolCallId, isError }); + + // PostHog's built-in `exec` tool surfaces UI apps through the per-call + // response `_meta` rather than registration-time tool metadata, so resolve + // and register the association here before the result is forwarded. + if (toolKey === POSTHOG_EXEC_TOOL_KEY && !isError) { + this.registerExecToolCallUi(toolCallId, result); + } + this.emit(McpAppsServiceEvent.ToolResult, { toolKey, toolCallId, @@ -391,8 +462,43 @@ export class McpAppsService extends TypedEventEmitter { } satisfies McpAppsToolResultEvent); } + /** + * Parse a `resourceUri` out of an `exec` call's response `_meta` and, when + * present, record a per-call association and announce it so the renderer can + * mount the UI app for this specific call. + */ + private registerExecToolCallUi(toolCallId: string, result: unknown): void { + if (this.execToolCallAssociations.has(toolCallId)) return; + + const resourceUri = resolveResultResourceUri(result); + if (!resourceUri) { + log.debug("registerExecToolCallUi: no UI resourceUri on result", { + toolCallId, + }); + return; + } + + this.execToolCallAssociations.set(toolCallId, { + toolKey: POSTHOG_EXEC_TOOL_KEY, + serverName: BUILTIN_POSTHOG_SERVER_NAME, + toolName: EXEC_TOOL_NAME, + resourceUri, + }); + log.info("registerExecToolCallUi: resolved per-call UI", { + toolCallId, + resourceUri, + }); + + this.emit(McpAppsServiceEvent.ToolCallUiDiscovered, { + toolCallId, + toolKey: POSTHOG_EXEC_TOOL_KEY, + resourceUri, + } satisfies McpAppsToolCallUiDiscoveredEvent); + } + notifyToolCancelled(toolKey: string, toolCallId: string): void { log.info("notifyToolCancelled", { toolKey, toolCallId }); + this.execToolCallAssociations.delete(toolCallId); this.emit(McpAppsServiceEvent.ToolCancelled, { toolKey, toolCallId, @@ -415,6 +521,7 @@ export class McpAppsService extends TypedEventEmitter { this.resourceCache.clear(); this.resourceMetaCache.clear(); this.toolAssociations.clear(); + this.execToolCallAssociations.clear(); this.toolDefinitions.clear(); this.pendingConnections.clear(); this.pendingFetches.clear(); @@ -452,10 +559,19 @@ export class McpAppsService extends TypedEventEmitter { this.toolAssociations.delete(key); } } + for (const [callId, assoc] of this.execToolCallAssociations) { + if (assoc.serverName === serverName) { + urisToEvict.add(assoc.resourceUri); + this.execToolCallAssociations.delete(callId); + } + } // Only evict cached resources not referenced by remaining associations const stillReferenced = new Set( - [...this.toolAssociations.values()].map((a) => a.resourceUri), + [ + ...this.toolAssociations.values(), + ...this.execToolCallAssociations.values(), + ].map((a) => a.resourceUri), ); for (const uri of urisToEvict) { if (!stillReferenced.has(uri)) { @@ -472,6 +588,7 @@ export class McpAppsService extends TypedEventEmitter { this.resourceCache.clear(); this.resourceMetaCache.clear(); this.toolAssociations.clear(); + this.execToolCallAssociations.clear(); this.toolDefinitions.clear(); this.serverConfigs.clear(); this.pendingConnections.clear(); diff --git a/apps/code/src/main/trpc/routers/mcp-apps.ts b/apps/code/src/main/trpc/routers/mcp-apps.ts index 60d423d435..156179c52f 100644 --- a/apps/code/src/main/trpc/routers/mcp-apps.ts +++ b/apps/code/src/main/trpc/routers/mcp-apps.ts @@ -8,6 +8,7 @@ import { openLinkInput, proxyResourceReadInput, proxyToolCallInput, + toolCallIdInput, } from "@shared/types/mcp-apps"; import { container } from "../../di/container"; import { MAIN_TOKENS } from "../../di/tokens"; @@ -27,6 +28,19 @@ export const mcpAppsRouter = router({ .input(hasUiForToolInput) .query(({ input }) => getService().hasUiForTool(input.toolKey)), + // Per-call UI lookups for PostHog's built-in `exec` tool, which resolves its + // UI app from each call's response `_meta` rather than registration metadata. + hasUiForToolCall: publicProcedure + .input(toolCallIdInput) + .query(({ input }) => getService().hasUiForToolCall(input.toolCallId)), + + getUiResourceForToolCall: publicProcedure + .input(toolCallIdInput) + .output(mcpUiResourceSchema.nullable()) + .query(({ input }) => + getService().getUiResourceForToolCall(input.toolCallId), + ), + getToolDefinition: publicProcedure .input(getToolDefinitionInput) .query(({ input }) => getService().getToolDefinition(input.toolKey)), @@ -101,4 +115,19 @@ export const mcpAppsRouter = router({ yield event; } }), + + onToolCallUiDiscovered: publicProcedure + .input(toolCallIdInput) + .subscription(async function* (opts) { + const service = getService(); + const targetToolCallId = opts.input.toolCallId; + for await (const event of service.toIterable( + McpAppsServiceEvent.ToolCallUiDiscovered, + { signal: opts.signal }, + )) { + if (event.toolCallId === targetToolCallId) { + yield event; + } + } + }), }); diff --git a/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx b/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx index e909423cd3..568f4e27d6 100644 --- a/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx +++ b/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx @@ -8,11 +8,12 @@ import type { import { ArrowsIn, ArrowsOut, Plugs, X } from "@phosphor-icons/react"; import { Box, Flex, IconButton, Text } from "@radix-ui/themes"; import { useTRPC } from "@renderer/trpc/client"; +import { POSTHOG_EXEC_TOOL_KEY } from "@shared/types/mcp-apps"; import { useThemeStore } from "@stores/themeStore"; import { useMutation, useQuery } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; import { logger } from "@utils/logger"; -import { useEffect, useRef, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { createPortal } from "react-dom"; import { type Phase, useAppBridge } from "../hooks/useAppBridge"; import { toCallToolResult } from "../utils/mcp-app-host-utils"; @@ -40,11 +41,21 @@ export function McpAppHost({ const [iframeEl, setIframeEl] = useState(null); const isDarkMode = useThemeStore((s) => s.isDarkMode); + // PostHog's built-in `exec` tool resolves its UI app per call (from the + // response `_meta`) rather than from registration metadata, so it's keyed by + // toolCallId instead of the shared `mcp__posthog__exec` tool key. + const isExec = mcpToolName === POSTHOG_EXEC_TOOL_KEY; + const { data: uiResource, isLoading: resourceLoading } = useQuery( - trpcReact.mcpApps.getUiResource.queryOptions( - { toolKey: mcpToolName }, - { staleTime: Infinity }, - ), + isExec + ? trpcReact.mcpApps.getUiResourceForToolCall.queryOptions( + { toolCallId: toolCall.toolCallId }, + { staleTime: Infinity }, + ) + : trpcReact.mcpApps.getUiResource.queryOptions( + { toolKey: mcpToolName }, + { staleTime: Infinity }, + ), ); const { data: toolDefinition } = useQuery( @@ -96,12 +107,35 @@ export function McpAppHost({ openLink: openLinkMut.mutateAsync, }); + // For `exec`, the UI app is discovered *from* the result, so this host mounts + // only after the result already arrived (and the live subscription fired). + // Send each call's result exactly once, whether it comes from the prop replay + // below or a late subscription event. + const sentResultForCallRef = useRef(null); + const sendResultOnce = useCallback( + (raw: unknown) => { + if (sentResultForCallRef.current === toolCall.toolCallId) return; + sentResultForCallRef.current = toolCall.toolCallId; + const toolResult = toCallToolResult(raw); + log.info("Sending tool result to app", { mcpToolName, toolResult }); + sendWhenReady((bridge) => bridge.sendToolResult(toolResult)); + }, + [toolCall.toolCallId, sendWhenReady, mcpToolName], + ); + // Forward tool results from subscriptions useSubscription( trpcReact.mcpApps.onToolResult.subscriptionOptions( { toolKey: mcpToolName }, { onData: (event) => { + // `exec` shares one tool key across every call, so scope delivery to + // this call and dedupe against the prop replay. + if (isExec) { + if (event.toolCallId !== toolCall.toolCallId) return; + sendResultOnce(event.result); + return; + } const toolResult = toCallToolResult(event.result); log.info("Sending tool result to app", { mcpToolName, @@ -114,6 +148,15 @@ export function McpAppHost({ ), ); + // `exec` replay: the result is already on the tool call by the time this host + // mounts, so push it from the prop (the subscription event is long gone). + useEffect(() => { + if (!isExec) return; + if (toolCall.status !== "completed" && toolCall.status !== "failed") return; + if (toolCall.rawOutput == null) return; + sendResultOnce(toolCall.rawOutput); + }, [isExec, toolCall.status, toolCall.rawOutput, sendResultOnce]); + // Forward tool cancellations from subscriptions useSubscription( trpcReact.mcpApps.onToolCancelled.subscriptionOptions( diff --git a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx index 25d0e5e3d8..51168bde26 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx @@ -3,8 +3,10 @@ import { McpToolView } from "@features/mcp-apps/components/McpToolView"; import { parseMcpToolKey } from "@features/mcp-apps/utils/mcp-app-host-utils"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { useTRPC } from "@renderer/trpc/client"; +import { POSTHOG_EXEC_TOOL_KEY } from "@shared/types/mcp-apps"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; +import { useEffect } from "react"; import type { ToolViewProps } from "./toolCallUtils"; interface McpToolBlockProps extends ToolViewProps { @@ -12,29 +14,45 @@ interface McpToolBlockProps extends ToolViewProps { } export function McpToolBlock(props: McpToolBlockProps) { - const { mcpToolName } = props; + const { mcpToolName, toolCall } = props; const { serverName, toolName } = parseMcpToolKey(mcpToolName); + // PostHog's built-in `exec` tool surfaces UI apps through each call's response + // `_meta`, so it can't be discovered at registration time like other MCP + // tools. For it we gate on a per-call lookup instead of the per-tool one. + const isExec = mcpToolName === POSTHOG_EXEC_TOOL_KEY; + const toolCallId = toolCall.toolCallId; + const mcpAppsDisabled = useSettingsStore((s) => s.mcpAppsDisabledServers); const isDisabledForServer = mcpAppsDisabled.includes(serverName); + const enabled = !isDisabledForServer; const trpcReact = useTRPC(); const queryClient = useQueryClient(); - const { data: hasUi } = useQuery( + // Registration-discovered tools: a stable per-tool association. + const { data: hasUiByTool } = useQuery( trpcReact.mcpApps.hasUiForTool.queryOptions( { toolKey: mcpToolName }, - { - staleTime: Infinity, - enabled: !isDisabledForServer, - }, + { staleTime: Infinity, enabled: enabled && !isExec }, ), ); + // `exec` tool: a per-call association resolved once the result arrives. + const { data: hasUiByCall } = useQuery( + trpcReact.mcpApps.hasUiForToolCall.queryOptions( + { toolCallId }, + { staleTime: Infinity, enabled: enabled && isExec }, + ), + ); + + const hasUi = isExec ? hasUiByCall : hasUiByTool; + // When MCP Apps discovery completes (possibly after this component mounted), // invalidate the hasUiForTool query so we pick up newly-discovered UIs. useSubscription( trpcReact.mcpApps.onDiscoveryComplete.subscriptionOptions(undefined, { + enabled: enabled && !isExec, onData: (_event) => { void queryClient.invalidateQueries( trpcReact.mcpApps.hasUiForTool.pathFilter(), @@ -46,6 +64,39 @@ export function McpToolBlock(props: McpToolBlockProps) { }), ); + // `exec`: a UI app is announced for this specific call once its response + // `_meta` is parsed in the main process. Refresh the per-call lookups. + useSubscription( + trpcReact.mcpApps.onToolCallUiDiscovered.subscriptionOptions( + { toolCallId }, + { + enabled: enabled && isExec, + onData: (_event) => { + void queryClient.invalidateQueries( + trpcReact.mcpApps.hasUiForToolCall.pathFilter(), + ); + void queryClient.invalidateQueries( + trpcReact.mcpApps.getUiResourceForToolCall.pathFilter(), + ); + }, + }, + ), + ); + + // Fallback for the race where this call completes before the subscription + // above is established: once the call settles, re-check the per-call lookup. + const status = toolCall.status; + useEffect(() => { + if (!enabled || !isExec) return; + if (status !== "completed" && status !== "failed") return; + void queryClient.invalidateQueries( + trpcReact.mcpApps.hasUiForToolCall.pathFilter(), + ); + void queryClient.invalidateQueries( + trpcReact.mcpApps.getUiResourceForToolCall.pathFilter(), + ); + }, [enabled, isExec, status, queryClient, trpcReact]); + return ( <> diff --git a/apps/code/src/shared/types/mcp-apps.test.ts b/apps/code/src/shared/types/mcp-apps.test.ts new file mode 100644 index 0000000000..a01e5a56ae --- /dev/null +++ b/apps/code/src/shared/types/mcp-apps.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it } from "vitest"; +import { + LEGACY_RESOURCE_URI_META_KEY, + POSTHOG_EXEC_TOOL_KEY, + resolveResultResourceUri, +} from "./mcp-apps"; + +describe("resolveResultResourceUri", () => { + it("reads the modern nested _meta.ui.resourceUri", () => { + expect( + resolveResultResourceUri({ _meta: { ui: { resourceUri: "ui://x" } } }), + ).toBe("ui://x"); + }); + + it("falls back to the legacy flat key", () => { + expect( + resolveResultResourceUri({ + _meta: { [LEGACY_RESOURCE_URI_META_KEY]: "ui://y" }, + }), + ).toBe("ui://y"); + }); + + it("prefers the modern key over the legacy one", () => { + expect( + resolveResultResourceUri({ + _meta: { + ui: { resourceUri: "ui://modern" }, + [LEGACY_RESOURCE_URI_META_KEY]: "ui://legacy", + }, + }), + ).toBe("ui://modern"); + }); + + it("returns undefined when there is no UI resource", () => { + expect(resolveResultResourceUri({ content: [] })).toBeUndefined(); + expect(resolveResultResourceUri({ _meta: {} })).toBeUndefined(); + expect( + resolveResultResourceUri({ _meta: { ui: { resourceUri: "" } } }), + ).toBeUndefined(); + expect(resolveResultResourceUri("a string result")).toBeUndefined(); + expect(resolveResultResourceUri(null)).toBeUndefined(); + expect(resolveResultResourceUri(undefined)).toBeUndefined(); + }); + + it("pins the built-in exec tool key", () => { + expect(POSTHOG_EXEC_TOOL_KEY).toBe("mcp__posthog__exec"); + }); +}); diff --git a/apps/code/src/shared/types/mcp-apps.ts b/apps/code/src/shared/types/mcp-apps.ts index 532bdd0d29..5201b5d998 100644 --- a/apps/code/src/shared/types/mcp-apps.ts +++ b/apps/code/src/shared/types/mcp-apps.ts @@ -68,6 +68,47 @@ export interface McpResourceUiMeta { }; } +// --- Built-in PostHog MCP server (single `exec` tool) --- +// +// The built-in PostHog MCP server is registered under this name (see +// `AgentAuthAdapter.buildMcpServers`). Unlike the MCP Apps spec — which binds a +// tool to its UI app upfront via the tool's registration `_meta.ui.resourceUri` +// — PostHog surfaces every UI app through a single generic `exec` tool and rides +// the `resourceUri` on each *tool-call response* `_meta` instead. Discovery via +// `listTools()` therefore never sees it, so for this server + tool the host has +// to resolve the UI app per call from the result metadata. +export const BUILTIN_POSTHOG_SERVER_NAME = "posthog"; +export const EXEC_TOOL_NAME = "exec"; +export const POSTHOG_EXEC_TOOL_KEY = `mcp__${BUILTIN_POSTHOG_SERVER_NAME}__${EXEC_TOOL_NAME}`; + +/** + * Legacy flat `_meta` key for a UI resource URI on a tool-call response. Mirrors + * `RESOURCE_URI_META_KEY` from `@modelcontextprotocol/ext-apps`. The modern form + * is the nested `_meta.ui.resourceUri`; servers may emit either (PostHog emits + * both). Hardcoded rather than imported so this module stays free of the + * ext-apps server entrypoint. + */ +export const LEGACY_RESOURCE_URI_META_KEY = "ui/resourceUri"; + +/** + * Resolve a UI resource URI from a tool-call response, preferring the modern + * nested `_meta.ui.resourceUri` and falling back to the legacy flat key — + * matching the host-side resolution recommended by `@modelcontextprotocol/ext-apps`. + * Returns `undefined` when the result carries no UI resource. + */ +export function resolveResultResourceUri(result: unknown): string | undefined { + if (result == null || typeof result !== "object") return undefined; + const meta = (result as { _meta?: Record })._meta; + if (meta == null || typeof meta !== "object") return undefined; + const ui = (meta as { ui?: { resourceUri?: unknown } }).ui; + const modern = ui?.resourceUri; + if (typeof modern === "string" && modern.length > 0) return modern; + const legacy = (meta as Record)[ + LEGACY_RESOURCE_URI_META_KEY + ]; + return typeof legacy === "string" && legacy.length > 0 ? legacy : undefined; +} + /** Tool-to-UI associations */ export const mcpToolUiAssociationSchema = z.object({ toolKey: z.string(), @@ -112,6 +153,11 @@ export const mcpAppsSubscriptionInput = z.object({ toolKey: z.string(), }); +/** Identifies a single tool *call* — used by the per-call exec UI path. */ +export const toolCallIdInput = z.object({ + toolCallId: z.string(), +}); + // --- Service event types --- export interface McpAppsToolInputEvent { @@ -136,11 +182,23 @@ export interface McpAppsDiscoveryCompleteEvent { toolKeys: string[]; } +/** + * Emitted when a UI app is resolved for a single tool *call* (the exec path). + * Unlike `DiscoveryComplete`, which fires once after registration-time + * discovery, this fires per call as soon as the response `_meta` is parsed. + */ +export interface McpAppsToolCallUiDiscoveredEvent { + toolCallId: string; + toolKey: string; + resourceUri: string; +} + export const McpAppsServiceEvent = { ToolInput: "tool-input", ToolResult: "tool-result", ToolCancelled: "tool-cancelled", DiscoveryComplete: "discovery-complete", + ToolCallUiDiscovered: "tool-call-ui-discovered", } as const; export interface McpAppsServiceEvents { @@ -148,6 +206,7 @@ export interface McpAppsServiceEvents { [McpAppsServiceEvent.ToolResult]: McpAppsToolResultEvent; [McpAppsServiceEvent.ToolCancelled]: McpAppsToolCancelledEvent; [McpAppsServiceEvent.DiscoveryComplete]: McpAppsDiscoveryCompleteEvent; + [McpAppsServiceEvent.ToolCallUiDiscovered]: McpAppsToolCallUiDiscoveredEvent; } // --- MCP server connection config --- From dc34b7c117a25ad980637b2f706f819cb2587890 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Tue, 2 Jun 2026 10:50:50 +0200 Subject: [PATCH 2/3] wip --- .../main/services/mcp-apps/service.test.ts | 88 +---- .../src/main/services/mcp-apps/service.ts | 323 ++++++++---------- .../src/main/services/mcp-proxy/service.ts | 16 +- apps/code/src/main/trpc/routers/mcp-apps.ts | 32 +- .../mcp-apps/components/McpAppHost.tsx | 56 ++- .../session-update/McpToolBlock.tsx | 99 +++--- apps/code/src/shared/types/mcp-apps.ts | 25 +- 7 files changed, 286 insertions(+), 353 deletions(-) diff --git a/apps/code/src/main/services/mcp-apps/service.test.ts b/apps/code/src/main/services/mcp-apps/service.test.ts index e4139e5489..d9de961e3e 100644 --- a/apps/code/src/main/services/mcp-apps/service.test.ts +++ b/apps/code/src/main/services/mcp-apps/service.test.ts @@ -1,9 +1,3 @@ -import { - LEGACY_RESOURCE_URI_META_KEY, - McpAppsServiceEvent, - type McpAppsToolCallUiDiscoveredEvent, - POSTHOG_EXEC_TOOL_KEY, -} from "@shared/types/mcp-apps"; import { beforeEach, describe, expect, it, vi } from "vitest"; vi.mock("../../utils/logger", () => ({ @@ -24,81 +18,27 @@ function makeService(): McpAppsService { return new McpAppsService(urlLauncher as never); } -describe("McpAppsService — per-call exec UI", () => { +describe("McpAppsService.getUiResourceByUri", () => { let service: McpAppsService; beforeEach(() => { service = makeService(); }); - it("registers a per-call association from a modern _meta.ui.resourceUri", () => { - const events: McpAppsToolCallUiDiscoveredEvent[] = []; - service.on(McpAppsServiceEvent.ToolCallUiDiscovered, (e) => events.push(e)); - - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-1", { - content: [{ type: "text", text: "{}" }], - _meta: { ui: { resourceUri: "ui://posthog/insight" } }, - }); - - expect(service.hasUiForToolCall("call-1")).toBe(true); - expect(events).toEqual([ - { - toolCallId: "call-1", - toolKey: POSTHOG_EXEC_TOOL_KEY, - resourceUri: "ui://posthog/insight", - }, - ]); - }); - - it("supports the legacy flat resourceUri key", () => { - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-legacy", { - _meta: { [LEGACY_RESOURCE_URI_META_KEY]: "ui://posthog/legacy" }, - }); - expect(service.hasUiForToolCall("call-legacy")).toBe(true); - }); - - it("does nothing for non-exec tools", () => { - service.notifyToolResult("mcp__posthog__query", "call-2", { - _meta: { ui: { resourceUri: "ui://posthog/insight" } }, - }); - expect(service.hasUiForToolCall("call-2")).toBe(false); - }); - - it("ignores errored exec results", () => { - service.notifyToolResult( - POSTHOG_EXEC_TOOL_KEY, - "call-3", - { _meta: { ui: { resourceUri: "ui://posthog/insight" } } }, - true, - ); - expect(service.hasUiForToolCall("call-3")).toBe(false); + it("rejects non-ui:// URIs without attempting a fetch", async () => { + await expect( + service.getUiResourceByUri("posthog", "https://evil.example/app.html"), + ).resolves.toBeNull(); + await expect( + service.getUiResourceByUri("posthog", "file:///etc/passwd"), + ).resolves.toBeNull(); }); - it("ignores exec results without a UI resource", () => { - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-4", { - content: [{ type: "text", text: "no ui here" }], - }); - expect(service.hasUiForToolCall("call-4")).toBe(false); - }); - - it("does not emit twice for the same call", () => { - const events: McpAppsToolCallUiDiscoveredEvent[] = []; - service.on(McpAppsServiceEvent.ToolCallUiDiscovered, (e) => events.push(e)); - - const result = { _meta: { ui: { resourceUri: "ui://posthog/insight" } } }; - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-5", result); - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-5", result); - - expect(events).toHaveLength(1); - }); - - it("evicts the association when the call is cancelled", () => { - service.notifyToolResult(POSTHOG_EXEC_TOOL_KEY, "call-6", { - _meta: { ui: { resourceUri: "ui://posthog/insight" } }, - }); - expect(service.hasUiForToolCall("call-6")).toBe(true); - - service.notifyToolCancelled(POSTHOG_EXEC_TOOL_KEY, "call-6"); - expect(service.hasUiForToolCall("call-6")).toBe(false); + it("returns null when the server has no connection config", async () => { + // ui:// passes the guard, but with no configured server the lazy connection + // fails and the fetch resolves to null rather than throwing. + await expect( + service.getUiResourceByUri("posthog", "ui://posthog/survey-list.html"), + ).resolves.toBeNull(); }); }); diff --git a/apps/code/src/main/services/mcp-apps/service.ts b/apps/code/src/main/services/mcp-apps/service.ts index e8ebfc02a0..aa153b6a87 100644 --- a/apps/code/src/main/services/mcp-apps/service.ts +++ b/apps/code/src/main/services/mcp-apps/service.ts @@ -5,10 +5,10 @@ import type { IUrlLauncher } from "@posthog/platform/url-launcher"; import { BUILTIN_POSTHOG_SERVER_NAME, EXEC_TOOL_NAME, + LEGACY_RESOURCE_URI_META_KEY, type McpAppsDiscoveryCompleteEvent, McpAppsServiceEvent, type McpAppsServiceEvents, - type McpAppsToolCallUiDiscoveredEvent, type McpAppsToolCancelledEvent, type McpAppsToolInputEvent, type McpAppsToolResultEvent, @@ -27,6 +27,29 @@ import { TypedEventEmitter } from "../../utils/typed-event-emitter"; const log = logger.scope("mcp-apps-service"); +/** + * Safe diagnostic snapshot of a tool-call result — surfaces whether/where a UI + * `resourceUri` is hiding in `_meta` without dumping large payloads (HTML, rows). + */ +function summarizeResult(result: unknown): Record { + if (result == null || typeof result !== "object") { + return { resultType: typeof result }; + } + const obj = result as Record; + const meta = obj._meta; + const hasMeta = meta != null && typeof meta === "object"; + const metaObj = hasMeta ? (meta as Record) : undefined; + return { + resultType: "object", + resultKeys: Object.keys(obj), + hasMeta, + metaKeys: metaObj ? Object.keys(metaObj) : undefined, + metaUi: metaObj?.ui, + legacyResourceUri: metaObj?.[LEGACY_RESOURCE_URI_META_KEY], + resolvedResourceUri: resolveResultResourceUri(result), + }; +} + const UI_MIME_TYPE = "text/html;profile=mcp-app"; const MAX_HTML_SIZE = 5 * 1024 * 1024; // 5MB @@ -41,13 +64,6 @@ export class McpAppsService extends TypedEventEmitter { private connections = new Map(); private resourceCache = new Map(); private toolAssociations = new Map(); - /** - * Per-tool-call UI associations for PostHog's built-in `exec` tool, keyed by - * `toolCallId`. The `exec` tool can't be discovered via `listTools()` (its UI - * `resourceUri` rides on each call's response `_meta`, not its registration), - * so associations here are populated lazily in {@link notifyToolResult}. - */ - private execToolCallAssociations = new Map(); private toolDefinitions = new Map(); private serverConfigs = new Map(); private pendingConnections = new Map>(); @@ -117,6 +133,15 @@ export class McpAppsService extends TypedEventEmitter { }), ]); + log.info("discoverServerUiTools: listed tools", { + serverName, + toolNames: toolsList.tools.map((t) => t.name), + hasExecTool: + serverName === BUILTIN_POSTHOG_SERVER_NAME && + toolsList.tools.some((t) => t.name === EXEC_TOOL_NAME), + resourceUris: resourcesList?.resources.map((r) => r.uri), + }); + for (const tool of toolsList.tools) { // The built-in PostHog `exec` tool carries no registration-time UI // resource (it resolves per call), but we still cache its definition so @@ -206,7 +231,7 @@ export class McpAppsService extends TypedEventEmitter { }); const client = new Client( - { name: "Twig", version: "1.0.0" }, + { name: "posthog-code", version: "1.0.0" }, { capabilities: { extensions: { @@ -229,9 +254,7 @@ export class McpAppsService extends TypedEventEmitter { } /** - * Get the UI resource for a tool. Fetches lazily on first access: - * creates an MCP connection if needed, then reads the resource HTML. - * Deduplicates concurrent fetches for the same resource URI. + * Fetch the UI resource for a registration-discovered tool, by its tool key. */ async getUiResourceForTool(toolKey: string): Promise { const association = this.toolAssociations.get(toolKey); @@ -239,100 +262,137 @@ export class McpAppsService extends TypedEventEmitter { log.debug("getUiResourceForTool: no association found", { toolKey }); return null; } + return this.fetchUiResourceByUri( + association.serverName, + association.resourceUri, + ); + } - // Return cached resource immediately - const cached = this.resourceCache.get(association.resourceUri); + /** + * Fetch a UI resource directly by its `ui://` URI. Used by the built-in + * PostHog `exec` path, where the resource URI is resolved per call from the + * tool result's `_meta` (in the renderer) rather than from a registered + * tool->UI association. Because the renderer derives it from the persisted + * conversation, exec UI apps survive app restarts — unlike the old in-memory + * per-call association map. + */ + async getUiResourceByUri( + serverName: string, + resourceUri: string, + ): Promise { + if (!resourceUri.startsWith("ui://")) { + log.warn("getUiResourceByUri: rejecting non-ui:// URI", { + serverName, + resourceUri, + }); + return null; + } + return this.fetchUiResourceByUri(serverName, resourceUri); + } + + /** + * Lazily fetch + cache a UI resource's HTML, deduplicating concurrent fetches + * for the same URI. Shared by the registration and per-call exec paths. + */ + private async fetchUiResourceByUri( + serverName: string, + resourceUri: string, + ): Promise { + const cached = this.resourceCache.get(resourceUri); if (cached) { - log.debug("getUiResourceForTool: cache hit", { toolKey }); + log.debug("fetchUiResourceByUri: cache hit", { serverName, resourceUri }); return cached; } - // Deduplicate concurrent fetches for the same resource URI - const pendingFetch = this.pendingFetches.get(association.resourceUri); + const pendingFetch = this.pendingFetches.get(resourceUri); if (pendingFetch) { - log.debug("getUiResourceForTool: joining pending fetch", { - toolKey, - uri: association.resourceUri, + log.debug("fetchUiResourceByUri: joining pending fetch", { + serverName, + resourceUri, }); return pendingFetch; } - // Start the fetch for this resource URI - log.debug("getUiResourceForTool: starting lazy fetch", { - toolKey, - serverName: association.serverName, - uri: association.resourceUri, + log.debug("fetchUiResourceByUri: starting lazy fetch", { + serverName, + resourceUri, }); - const fetchPromise = this.fetchUiResource(association); - this.pendingFetches.set(association.resourceUri, fetchPromise); - + const fetchPromise = this.doFetchUiResource(serverName, resourceUri); + this.pendingFetches.set(resourceUri, fetchPromise); try { return await fetchPromise; } finally { - this.pendingFetches.delete(association.resourceUri); + this.pendingFetches.delete(resourceUri); } } - private async fetchUiResource( - association: McpToolUiAssociation, + private async doFetchUiResource( + serverName: string, + resourceUri: string, ): Promise { + let resourceResult: Awaited>; try { - const conn = await this.getOrCreateConnection(association.serverName); - const resourceResult = await conn.client.readResource({ - uri: association.resourceUri, + const conn = await this.getOrCreateConnection(serverName); + resourceResult = await conn.client.readResource({ uri: resourceUri }); + } catch (err) { + // Connection/read failures are transient — most notably "No server config + // for: posthog" during the boot race, before a session populates configs. + // Rethrow so the caller's query surfaces an error and retries, instead of + // caching a permanent `null` for this (shared) resource URI and poisoning + // every later call that reuses it. + log.warn("Failed to fetch UI resource (transient — will retry)", { + serverName, + uri: resourceUri, + error: err instanceof Error ? err.message : String(err), }); + throw err; + } - const textContent = resourceResult.contents.find( - (c) => "text" in c && c.mimeType === UI_MIME_TYPE, - ); - if (!textContent || !("text" in textContent)) { - log.warn("UI resource had no matching text content", { - serverName: association.serverName, - uri: association.resourceUri, - contentsCount: resourceResult.contents.length, - }); - return null; - } - - if (textContent.text.length > MAX_HTML_SIZE) { - log.warn("UI resource HTML exceeds size limit", { - uri: association.resourceUri, - size: textContent.text.length, - limit: MAX_HTML_SIZE, - }); - return null; - } - - // Use metadata cached during discovery - const resourceMeta = this.resourceMetaCache.get(association.resourceUri); - - const resource: McpUiResource = { - uri: association.resourceUri, - name: resourceMeta?.name, - mimeType: UI_MIME_TYPE, - csp: resourceMeta?._meta?.ui?.csp, - permissions: resourceMeta?._meta?.ui?.permissions, - html: textContent.text, - serverName: association.serverName, - }; - - this.resourceCache.set(association.resourceUri, resource); - log.info("Lazily fetched and cached UI resource", { - serverName: association.serverName, - uri: association.resourceUri, - htmlLength: textContent.text.length, - hasCsp: !!resource.csp, + const textContent = resourceResult.contents.find( + (c) => "text" in c && c.mimeType === UI_MIME_TYPE, + ); + if (!textContent || !("text" in textContent)) { + // The server answered but has no usable UI content — a definitive "no UI" + // for this URI, so caching `null` is correct. + log.warn("UI resource had no matching text content", { + serverName, + uri: resourceUri, + contentsCount: resourceResult.contents.length, }); + return null; + } - return resource; - } catch (err) { - log.warn("Failed to lazily fetch UI resource", { - serverName: association.serverName, - uri: association.resourceUri, - error: err instanceof Error ? err.message : String(err), + if (textContent.text.length > MAX_HTML_SIZE) { + log.warn("UI resource HTML exceeds size limit", { + uri: resourceUri, + size: textContent.text.length, + limit: MAX_HTML_SIZE, }); return null; } + + // Use metadata cached during discovery (CSP/permissions), if available. + const resourceMeta = this.resourceMetaCache.get(resourceUri); + + const resource: McpUiResource = { + uri: resourceUri, + name: resourceMeta?.name, + mimeType: UI_MIME_TYPE, + csp: resourceMeta?._meta?.ui?.csp, + permissions: resourceMeta?._meta?.ui?.permissions, + html: textContent.text, + serverName, + }; + + this.resourceCache.set(resourceUri, resource); + log.info("Lazily fetched and cached UI resource", { + serverName, + uri: resourceUri, + htmlLength: textContent.text.length, + hasCsp: !!resource.csp, + }); + + return resource; } hasUiForTool(toolKey: string): boolean { @@ -341,47 +401,6 @@ export class McpAppsService extends TypedEventEmitter { return has; } - /** - * Whether a UI app was resolved for a specific tool call. Only the built-in - * PostHog `exec` tool populates this (via {@link notifyToolResult}); for - * registration-discovered tools use {@link hasUiForTool}. - */ - hasUiForToolCall(toolCallId: string): boolean { - const has = this.execToolCallAssociations.has(toolCallId); - log.debug("hasUiForToolCall", { toolCallId, result: has }); - return has; - } - - /** - * Fetch the UI resource resolved for a specific `exec` tool call. Reuses the - * same lazy-fetch + cache path as registration-discovered tools. - */ - async getUiResourceForToolCall( - toolCallId: string, - ): Promise { - const association = this.execToolCallAssociations.get(toolCallId); - if (!association) { - log.debug("getUiResourceForToolCall: no association found", { - toolCallId, - }); - return null; - } - - const cached = this.resourceCache.get(association.resourceUri); - if (cached) return cached; - - const pendingFetch = this.pendingFetches.get(association.resourceUri); - if (pendingFetch) return pendingFetch; - - const fetchPromise = this.fetchUiResource(association); - this.pendingFetches.set(association.resourceUri, fetchPromise); - try { - return await fetchPromise; - } finally { - this.pendingFetches.delete(association.resourceUri); - } - } - getToolDefinition(toolKey: string): Tool | null { return this.toolDefinitions.get(toolKey) ?? null; } @@ -445,14 +464,12 @@ export class McpAppsService extends TypedEventEmitter { result: unknown, isError?: boolean, ): void { - log.info("notifyToolResult", { toolKey, toolCallId, isError }); - - // PostHog's built-in `exec` tool surfaces UI apps through the per-call - // response `_meta` rather than registration-time tool metadata, so resolve - // and register the association here before the result is forwarded. - if (toolKey === POSTHOG_EXEC_TOOL_KEY && !isError) { - this.registerExecToolCallUi(toolCallId, result); - } + log.info("notifyToolResult", { + toolKey, + toolCallId, + isError, + ...summarizeResult(result), + }); this.emit(McpAppsServiceEvent.ToolResult, { toolKey, @@ -462,43 +479,8 @@ export class McpAppsService extends TypedEventEmitter { } satisfies McpAppsToolResultEvent); } - /** - * Parse a `resourceUri` out of an `exec` call's response `_meta` and, when - * present, record a per-call association and announce it so the renderer can - * mount the UI app for this specific call. - */ - private registerExecToolCallUi(toolCallId: string, result: unknown): void { - if (this.execToolCallAssociations.has(toolCallId)) return; - - const resourceUri = resolveResultResourceUri(result); - if (!resourceUri) { - log.debug("registerExecToolCallUi: no UI resourceUri on result", { - toolCallId, - }); - return; - } - - this.execToolCallAssociations.set(toolCallId, { - toolKey: POSTHOG_EXEC_TOOL_KEY, - serverName: BUILTIN_POSTHOG_SERVER_NAME, - toolName: EXEC_TOOL_NAME, - resourceUri, - }); - log.info("registerExecToolCallUi: resolved per-call UI", { - toolCallId, - resourceUri, - }); - - this.emit(McpAppsServiceEvent.ToolCallUiDiscovered, { - toolCallId, - toolKey: POSTHOG_EXEC_TOOL_KEY, - resourceUri, - } satisfies McpAppsToolCallUiDiscoveredEvent); - } - notifyToolCancelled(toolKey: string, toolCallId: string): void { log.info("notifyToolCancelled", { toolKey, toolCallId }); - this.execToolCallAssociations.delete(toolCallId); this.emit(McpAppsServiceEvent.ToolCancelled, { toolKey, toolCallId, @@ -521,7 +503,6 @@ export class McpAppsService extends TypedEventEmitter { this.resourceCache.clear(); this.resourceMetaCache.clear(); this.toolAssociations.clear(); - this.execToolCallAssociations.clear(); this.toolDefinitions.clear(); this.pendingConnections.clear(); this.pendingFetches.clear(); @@ -559,19 +540,10 @@ export class McpAppsService extends TypedEventEmitter { this.toolAssociations.delete(key); } } - for (const [callId, assoc] of this.execToolCallAssociations) { - if (assoc.serverName === serverName) { - urisToEvict.add(assoc.resourceUri); - this.execToolCallAssociations.delete(callId); - } - } // Only evict cached resources not referenced by remaining associations const stillReferenced = new Set( - [ - ...this.toolAssociations.values(), - ...this.execToolCallAssociations.values(), - ].map((a) => a.resourceUri), + [...this.toolAssociations.values()].map((a) => a.resourceUri), ); for (const uri of urisToEvict) { if (!stillReferenced.has(uri)) { @@ -588,7 +560,6 @@ export class McpAppsService extends TypedEventEmitter { this.resourceCache.clear(); this.resourceMetaCache.clear(); this.toolAssociations.clear(); - this.execToolCallAssociations.clear(); this.toolDefinitions.clear(); this.serverConfigs.clear(); this.pendingConnections.clear(); diff --git a/apps/code/src/main/services/mcp-proxy/service.ts b/apps/code/src/main/services/mcp-proxy/service.ts index 1cf267355e..7d851fdd22 100644 --- a/apps/code/src/main/services/mcp-proxy/service.ts +++ b/apps/code/src/main/services/mcp-proxy/service.ts @@ -166,6 +166,12 @@ export class McpProxyService { options: RequestInit, res: http.ServerResponse, ): Promise { + log.debug("MCP proxy forwarding request", { + id, + url, + method: options.method, + requestBody: truncateRequestBody(options.body), + }); try { let response = await this.authService.authenticatedFetch( fetch, @@ -187,6 +193,8 @@ export class McpProxyService { log.warn("MCP auth failure — refreshing token and retrying", { id, url, + method: options.method, + requestBody: truncateRequestBody(options.body), status: response.status, }); await this.authService.refreshAccessToken(); @@ -228,7 +236,13 @@ export class McpProxyService { this.writeStreamingResponse(response, res); } catch (err) { - log.error("MCP proxy forward error", { id, url, err }); + log.error("MCP proxy forward error", { + id, + url, + method: options.method, + requestBody: truncateRequestBody(options.body), + err, + }); if (!res.headersSent) { res.writeHead(502); } diff --git a/apps/code/src/main/trpc/routers/mcp-apps.ts b/apps/code/src/main/trpc/routers/mcp-apps.ts index 156179c52f..ad92479e92 100644 --- a/apps/code/src/main/trpc/routers/mcp-apps.ts +++ b/apps/code/src/main/trpc/routers/mcp-apps.ts @@ -1,5 +1,6 @@ import { getToolDefinitionInput, + getUiResourceByUriInput, getUiResourceInput, hasUiForToolInput, McpAppsServiceEvent, @@ -8,7 +9,6 @@ import { openLinkInput, proxyResourceReadInput, proxyToolCallInput, - toolCallIdInput, } from "@shared/types/mcp-apps"; import { container } from "../../di/container"; import { MAIN_TOKENS } from "../../di/tokens"; @@ -28,17 +28,14 @@ export const mcpAppsRouter = router({ .input(hasUiForToolInput) .query(({ input }) => getService().hasUiForTool(input.toolKey)), - // Per-call UI lookups for PostHog's built-in `exec` tool, which resolves its - // UI app from each call's response `_meta` rather than registration metadata. - hasUiForToolCall: publicProcedure - .input(toolCallIdInput) - .query(({ input }) => getService().hasUiForToolCall(input.toolCallId)), - - getUiResourceForToolCall: publicProcedure - .input(toolCallIdInput) + // Fetch a UI resource by URI. The built-in PostHog `exec` tool resolves its + // UI app per call from the result's `_meta` (in the renderer) rather than + // registration metadata, so the host fetches the resource by URI directly. + getUiResourceByUri: publicProcedure + .input(getUiResourceByUriInput) .output(mcpUiResourceSchema.nullable()) .query(({ input }) => - getService().getUiResourceForToolCall(input.toolCallId), + getService().getUiResourceByUri(input.serverName, input.resourceUri), ), getToolDefinition: publicProcedure @@ -115,19 +112,4 @@ export const mcpAppsRouter = router({ yield event; } }), - - onToolCallUiDiscovered: publicProcedure - .input(toolCallIdInput) - .subscription(async function* (opts) { - const service = getService(); - const targetToolCallId = opts.input.toolCallId; - for await (const event of service.toIterable( - McpAppsServiceEvent.ToolCallUiDiscovered, - { signal: opts.signal }, - )) { - if (event.toolCallId === targetToolCallId) { - yield event; - } - } - }), }); diff --git a/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx b/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx index 568f4e27d6..1b9851068a 100644 --- a/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx +++ b/apps/code/src/renderer/features/mcp-apps/components/McpAppHost.tsx @@ -8,7 +8,10 @@ import type { import { ArrowsIn, ArrowsOut, Plugs, X } from "@phosphor-icons/react"; import { Box, Flex, IconButton, Text } from "@radix-ui/themes"; import { useTRPC } from "@renderer/trpc/client"; -import { POSTHOG_EXEC_TOOL_KEY } from "@shared/types/mcp-apps"; +import { + POSTHOG_EXEC_TOOL_KEY, + resolveResultResourceUri, +} from "@shared/types/mcp-apps"; import { useThemeStore } from "@stores/themeStore"; import { useMutation, useQuery } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; @@ -41,16 +44,21 @@ export function McpAppHost({ const [iframeEl, setIframeEl] = useState(null); const isDarkMode = useThemeStore((s) => s.isDarkMode); - // PostHog's built-in `exec` tool resolves its UI app per call (from the - // response `_meta`) rather than from registration metadata, so it's keyed by - // toolCallId instead of the shared `mcp__posthog__exec` tool key. + // PostHog's built-in `exec` tool resolves its UI app per call from the + // response `_meta` (carried on `toolCall.rawOutput`, which is persisted in the + // conversation) rather than from registration metadata. We fetch the resource + // by its `ui://` URI so it works for both live and rehydrated (post-restart) + // calls; registration-discovered tools fetch by their tool key. const isExec = mcpToolName === POSTHOG_EXEC_TOOL_KEY; + const execResourceUri = isExec + ? resolveResultResourceUri(toolCall.rawOutput) + : undefined; const { data: uiResource, isLoading: resourceLoading } = useQuery( isExec - ? trpcReact.mcpApps.getUiResourceForToolCall.queryOptions( - { toolCallId: toolCall.toolCallId }, - { staleTime: Infinity }, + ? trpcReact.mcpApps.getUiResourceByUri.queryOptions( + { serverName, resourceUri: execResourceUri ?? "" }, + { staleTime: Infinity, enabled: !!execResourceUri }, ) : trpcReact.mcpApps.getUiResource.queryOptions( { toolKey: mcpToolName }, @@ -66,13 +74,24 @@ export function McpAppHost({ ); useEffect(() => { - log.debug("McpAppHost render", { + log.info("McpAppHost render", { mcpToolName, + isExec, + toolCallId: toolCall.toolCallId, + status: toolCall.status, resourceLoading, hasResource: !!uiResource, resourceUri: uiResource?.uri, }); - }, [mcpToolName, resourceLoading, uiResource, uiResource?.uri]); + }, [ + mcpToolName, + isExec, + toolCall.toolCallId, + toolCall.status, + resourceLoading, + uiResource, + uiResource?.uri, + ]); const proxyToolCallMut = useMutation( trpcReact.mcpApps.proxyToolCall.mutationOptions(), @@ -153,9 +172,24 @@ export function McpAppHost({ useEffect(() => { if (!isExec) return; if (toolCall.status !== "completed" && toolCall.status !== "failed") return; - if (toolCall.rawOutput == null) return; + if (toolCall.rawOutput == null) { + log.info("exec replay skipped: no rawOutput on toolCall", { + toolCallId: toolCall.toolCallId, + status: toolCall.status, + }); + return; + } + log.info("exec replay: sending result from toolCall prop", { + toolCallId: toolCall.toolCallId, + }); sendResultOnce(toolCall.rawOutput); - }, [isExec, toolCall.status, toolCall.rawOutput, sendResultOnce]); + }, [ + isExec, + toolCall.status, + toolCall.rawOutput, + toolCall.toolCallId, + sendResultOnce, + ]); // Forward tool cancellations from subscriptions useSubscription( diff --git a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx index 51168bde26..b52af368df 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx @@ -3,12 +3,18 @@ import { McpToolView } from "@features/mcp-apps/components/McpToolView"; import { parseMcpToolKey } from "@features/mcp-apps/utils/mcp-app-host-utils"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { useTRPC } from "@renderer/trpc/client"; -import { POSTHOG_EXEC_TOOL_KEY } from "@shared/types/mcp-apps"; +import { + POSTHOG_EXEC_TOOL_KEY, + resolveResultResourceUri, +} from "@shared/types/mcp-apps"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; +import { logger } from "@utils/logger"; import { useEffect } from "react"; import type { ToolViewProps } from "./toolCallUtils"; +const log = logger.scope("mcp-tool-block"); + interface McpToolBlockProps extends ToolViewProps { mcpToolName: string; } @@ -18,42 +24,46 @@ export function McpToolBlock(props: McpToolBlockProps) { const { serverName, toolName } = parseMcpToolKey(mcpToolName); // PostHog's built-in `exec` tool surfaces UI apps through each call's response - // `_meta`, so it can't be discovered at registration time like other MCP - // tools. For it we gate on a per-call lookup instead of the per-tool one. + // `_meta` rather than registration-time tool metadata. The resolved `ui://` + // URI lives on `toolCall.rawOutput`, which is persisted in the conversation — + // so deriving it from the prop makes exec UI apps survive app restarts. const isExec = mcpToolName === POSTHOG_EXEC_TOOL_KEY; - const toolCallId = toolCall.toolCallId; + const execResourceUri = isExec + ? resolveResultResourceUri(toolCall.rawOutput) + : undefined; const mcpAppsDisabled = useSettingsStore((s) => s.mcpAppsDisabledServers); const isDisabledForServer = mcpAppsDisabled.includes(serverName); - const enabled = !isDisabledForServer; const trpcReact = useTRPC(); const queryClient = useQueryClient(); - // Registration-discovered tools: a stable per-tool association. + // Registration-discovered tools: a stable per-tool association rebuilt at boot. const { data: hasUiByTool } = useQuery( trpcReact.mcpApps.hasUiForTool.queryOptions( { toolKey: mcpToolName }, - { staleTime: Infinity, enabled: enabled && !isExec }, - ), - ); - - // `exec` tool: a per-call association resolved once the result arrives. - const { data: hasUiByCall } = useQuery( - trpcReact.mcpApps.hasUiForToolCall.queryOptions( - { toolCallId }, - { staleTime: Infinity, enabled: enabled && isExec }, + { staleTime: Infinity, enabled: !isDisabledForServer && !isExec }, ), ); - const hasUi = isExec ? hasUiByCall : hasUiByTool; + const hasUi = isExec ? !!execResourceUri : hasUiByTool; - // When MCP Apps discovery completes (possibly after this component mounted), - // invalidate the hasUiForTool query so we pick up newly-discovered UIs. + // Discovery completing signals that MCP server configs are now populated and + // a connection can be opened. Two reasons to react: + // - registration path: a UI may have been newly discovered → refresh the gate. + // - exec path: when viewing a past conversation on app boot, the resource + // fetch can run before any session populates the server config and fail + // with "No server config". Re-fetch once discovery lands so it succeeds. useSubscription( trpcReact.mcpApps.onDiscoveryComplete.subscriptionOptions(undefined, { - enabled: enabled && !isExec, + enabled: !isDisabledForServer, onData: (_event) => { + if (isExec) { + void queryClient.invalidateQueries( + trpcReact.mcpApps.getUiResourceByUri.pathFilter(), + ); + return; + } void queryClient.invalidateQueries( trpcReact.mcpApps.hasUiForTool.pathFilter(), ); @@ -64,38 +74,27 @@ export function McpToolBlock(props: McpToolBlockProps) { }), ); - // `exec`: a UI app is announced for this specific call once its response - // `_meta` is parsed in the main process. Refresh the per-call lookups. - useSubscription( - trpcReact.mcpApps.onToolCallUiDiscovered.subscriptionOptions( - { toolCallId }, - { - enabled: enabled && isExec, - onData: (_event) => { - void queryClient.invalidateQueries( - trpcReact.mcpApps.hasUiForToolCall.pathFilter(), - ); - void queryClient.invalidateQueries( - trpcReact.mcpApps.getUiResourceForToolCall.pathFilter(), - ); - }, - }, - ), - ); - - // Fallback for the race where this call completes before the subscription - // above is established: once the call settles, re-check the per-call lookup. - const status = toolCall.status; useEffect(() => { - if (!enabled || !isExec) return; - if (status !== "completed" && status !== "failed") return; - void queryClient.invalidateQueries( - trpcReact.mcpApps.hasUiForToolCall.pathFilter(), - ); - void queryClient.invalidateQueries( - trpcReact.mcpApps.getUiResourceForToolCall.pathFilter(), - ); - }, [enabled, isExec, status, queryClient, trpcReact]); + log.info("render state", { + mcpToolName, + isExec, + toolCallId: toolCall.toolCallId, + status: toolCall.status, + isDisabledForServer, + hasUiByTool, + execResourceUri, + willRenderApp: !!hasUi && !isDisabledForServer, + }); + }, [ + mcpToolName, + isExec, + toolCall.toolCallId, + toolCall.status, + isDisabledForServer, + hasUiByTool, + execResourceUri, + hasUi, + ]); return ( <> diff --git a/apps/code/src/shared/types/mcp-apps.ts b/apps/code/src/shared/types/mcp-apps.ts index 5201b5d998..2a3db018cb 100644 --- a/apps/code/src/shared/types/mcp-apps.ts +++ b/apps/code/src/shared/types/mcp-apps.ts @@ -153,9 +153,15 @@ export const mcpAppsSubscriptionInput = z.object({ toolKey: z.string(), }); -/** Identifies a single tool *call* — used by the per-call exec UI path. */ -export const toolCallIdInput = z.object({ - toolCallId: z.string(), +/** + * Fetch a UI resource directly by URI. Used by the built-in PostHog `exec` + * path, where the renderer resolves the `ui://` URI from the tool result's + * `_meta` (see {@link resolveResultResourceUri}) rather than a registered + * tool->UI association. + */ +export const getUiResourceByUriInput = z.object({ + serverName: z.string(), + resourceUri: z.string(), }); // --- Service event types --- @@ -182,23 +188,11 @@ export interface McpAppsDiscoveryCompleteEvent { toolKeys: string[]; } -/** - * Emitted when a UI app is resolved for a single tool *call* (the exec path). - * Unlike `DiscoveryComplete`, which fires once after registration-time - * discovery, this fires per call as soon as the response `_meta` is parsed. - */ -export interface McpAppsToolCallUiDiscoveredEvent { - toolCallId: string; - toolKey: string; - resourceUri: string; -} - export const McpAppsServiceEvent = { ToolInput: "tool-input", ToolResult: "tool-result", ToolCancelled: "tool-cancelled", DiscoveryComplete: "discovery-complete", - ToolCallUiDiscovered: "tool-call-ui-discovered", } as const; export interface McpAppsServiceEvents { @@ -206,7 +200,6 @@ export interface McpAppsServiceEvents { [McpAppsServiceEvent.ToolResult]: McpAppsToolResultEvent; [McpAppsServiceEvent.ToolCancelled]: McpAppsToolCancelledEvent; [McpAppsServiceEvent.DiscoveryComplete]: McpAppsDiscoveryCompleteEvent; - [McpAppsServiceEvent.ToolCallUiDiscovered]: McpAppsToolCallUiDiscoveredEvent; } // --- MCP server connection config --- From e71c2b7eb73243b6ca389df475b05e9068580691 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Thu, 4 Jun 2026 11:32:03 +0200 Subject: [PATCH 3/3] fix(mcp-apps): address review comments - Rename MCP host identity from "Twig" to "posthog-code" in useAppBridge, matching the client identity already used in the main service. - Remove the debug "render state" useEffect from McpToolBlock. - Fix the getUiResourceByUri test to expect a rejection: a missing server config rethrows (transient boot-race), it does not resolve to null. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../main/services/mcp-apps/service.test.ts | 7 ++--- .../features/mcp-apps/hooks/useAppBridge.ts | 2 +- .../session-update/McpToolBlock.tsx | 26 ------------------- 3 files changed, 5 insertions(+), 30 deletions(-) diff --git a/apps/code/src/main/services/mcp-apps/service.test.ts b/apps/code/src/main/services/mcp-apps/service.test.ts index d9de961e3e..fed7d3b2a5 100644 --- a/apps/code/src/main/services/mcp-apps/service.test.ts +++ b/apps/code/src/main/services/mcp-apps/service.test.ts @@ -34,11 +34,12 @@ describe("McpAppsService.getUiResourceByUri", () => { ).resolves.toBeNull(); }); - it("returns null when the server has no connection config", async () => { + it("rejects when the server has no connection config", async () => { // ui:// passes the guard, but with no configured server the lazy connection - // fails and the fetch resolves to null rather than throwing. + // fails. The fetch rethrows rather than caching a permanent null, so the + // caller's query can surface the error and retry once boot populates configs. await expect( service.getUiResourceByUri("posthog", "ui://posthog/survey-list.html"), - ).resolves.toBeNull(); + ).rejects.toThrow("No server config for: posthog"); }); }); diff --git a/apps/code/src/renderer/features/mcp-apps/hooks/useAppBridge.ts b/apps/code/src/renderer/features/mcp-apps/hooks/useAppBridge.ts index c70fa57e7c..1791918d90 100644 --- a/apps/code/src/renderer/features/mcp-apps/hooks/useAppBridge.ts +++ b/apps/code/src/renderer/features/mcp-apps/hooks/useAppBridge.ts @@ -62,7 +62,7 @@ interface UseAppBridgeReturn { sendWhenReady: (fn: (bridge: AppBridge) => void) => void; } -const HOST_INFO = { name: "Twig", version: "1.0.0" }; +const HOST_INFO = { name: "posthog-code", version: "1.0.0" }; const HOST_CAPABILITIES: McpUiHostCapabilities = { openLinks: {}, diff --git a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx index b52af368df..d618c6e272 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/McpToolBlock.tsx @@ -9,12 +9,8 @@ import { } from "@shared/types/mcp-apps"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; -import { logger } from "@utils/logger"; -import { useEffect } from "react"; import type { ToolViewProps } from "./toolCallUtils"; -const log = logger.scope("mcp-tool-block"); - interface McpToolBlockProps extends ToolViewProps { mcpToolName: string; } @@ -74,28 +70,6 @@ export function McpToolBlock(props: McpToolBlockProps) { }), ); - useEffect(() => { - log.info("render state", { - mcpToolName, - isExec, - toolCallId: toolCall.toolCallId, - status: toolCall.status, - isDisabledForServer, - hasUiByTool, - execResourceUri, - willRenderApp: !!hasUi && !isDisabledForServer, - }); - }, [ - mcpToolName, - isExec, - toolCall.toolCallId, - toolCall.status, - isDisabledForServer, - hasUiByTool, - execResourceUri, - hasUi, - ]); - return ( <>