From 6f78db36ba66d57f80ec607d1d182c00e24c4993 Mon Sep 17 00:00:00 2001 From: Jovan Sakovic <49978945+sakce@users.noreply.github.com> Date: Thu, 28 May 2026 15:17:48 +0100 Subject: [PATCH 1/2] feat(agent): add self-modification via posthog-code-internal MCP server Adds a localhost HTTP MCP server, started with the main process, that exposes the agent's own internals as MCP tools generated from tRPC procedures via a YAML allowlist. Initial tools (custom instructions read/write, MCP server list/install) are hand-picked, but adding more is just flipping enabled: true in mcp-tools.yaml. Generated-By: PostHog Code Task-Id: 934a7fa1-fbb6-4f1d-89d3-f6abc69b7e23 --- apps/code/package.json | 1 + apps/code/scripts/electron-stub.mjs | 78 ++++ .../scripts/scaffold-mcp-tools-loader.mjs | 55 +++ .../scripts/scaffold-mcp-tools-preload.mjs | 10 + apps/code/scripts/scaffold-mcp-tools.ts | 183 ++++++++ apps/code/src/main/di/container.ts | 10 + apps/code/src/main/di/tokens.ts | 5 + apps/code/src/main/index.ts | 6 + .../main/services/agent/auth-adapter.test.ts | 7 + .../src/main/services/agent/auth-adapter.ts | 16 + .../src/main/services/agent/service.test.ts | 4 + apps/code/src/main/services/agent/service.ts | 56 +++ .../services/custom-instructions/schemas.ts | 45 ++ .../services/custom-instructions/service.ts | 64 +++ .../services/mcp-installations/schemas.ts | 79 ++++ .../mcp-installations/service.test.ts | 105 +++++ .../services/mcp-installations/service.ts | 189 +++++++++ .../posthog-code-internal-mcp/mcp-tools.yaml | 58 +++ .../posthog-code-internal-mcp/service.ts | 167 ++++++++ .../tool-registry.test.ts | 342 +++++++++++++++ .../tool-registry.ts | 391 ++++++++++++++++++ .../posthog-code-internal-mcp/yaml-schema.ts | 45 ++ apps/code/src/main/trpc/router.ts | 4 + .../main/trpc/routers/custom-instructions.ts | 39 ++ .../main/trpc/routers/mcp-installations.ts | 34 ++ apps/code/src/renderer/App.tsx | 6 + .../features/settings/stores/settingsStore.ts | 24 ++ packages/agent/src/index.ts | 6 +- 28 files changed, 2028 insertions(+), 1 deletion(-) create mode 100644 apps/code/scripts/electron-stub.mjs create mode 100644 apps/code/scripts/scaffold-mcp-tools-loader.mjs create mode 100644 apps/code/scripts/scaffold-mcp-tools-preload.mjs create mode 100644 apps/code/scripts/scaffold-mcp-tools.ts create mode 100644 apps/code/src/main/services/custom-instructions/schemas.ts create mode 100644 apps/code/src/main/services/custom-instructions/service.ts create mode 100644 apps/code/src/main/services/mcp-installations/schemas.ts create mode 100644 apps/code/src/main/services/mcp-installations/service.test.ts create mode 100644 apps/code/src/main/services/mcp-installations/service.ts create mode 100644 apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml create mode 100644 apps/code/src/main/services/posthog-code-internal-mcp/service.ts create mode 100644 apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.test.ts create mode 100644 apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.ts create mode 100644 apps/code/src/main/services/posthog-code-internal-mcp/yaml-schema.ts create mode 100644 apps/code/src/main/trpc/routers/custom-instructions.ts create mode 100644 apps/code/src/main/trpc/routers/mcp-installations.ts diff --git a/apps/code/package.json b/apps/code/package.json index 3c1f3d09bb..fafe2eb4a7 100644 --- a/apps/code/package.json +++ b/apps/code/package.json @@ -22,6 +22,7 @@ "build-icons": "bash scripts/generate-icns.sh", "typecheck": "tsc -p tsconfig.node.json --noEmit && tsc -p tsconfig.web.json --noEmit", "generate-client": "tsx scripts/update-openapi-client.ts", + "scaffold-mcp-tools": "tsx --import ./scripts/scaffold-mcp-tools-preload.mjs scripts/scaffold-mcp-tools.ts", "test": "vitest run", "test:e2e": "playwright test --config=tests/e2e/playwright.config.ts", "test:e2e:headed": "playwright test --config=tests/e2e/playwright.config.ts --headed", diff --git a/apps/code/scripts/electron-stub.mjs b/apps/code/scripts/electron-stub.mjs new file mode 100644 index 0000000000..32c3afa962 --- /dev/null +++ b/apps/code/scripts/electron-stub.mjs @@ -0,0 +1,78 @@ +// Stub module returned in place of `electron` when the scaffold script runs +// outside of Electron. The router walk only needs the procedures' input +// schemas — none of the methods on `app`, `BrowserWindow`, etc. are actually +// called at import time, so a no-op object suffices. +// +// Used by scaffold-mcp-tools-preload.mjs via a Node loader hook. +const noop = () => {}; +const emptyObj = new Proxy( + {}, + { + get() { + return noop; + }, + }, +); + +const stub = new Proxy( + { + app: emptyObj, + BrowserWindow: () => emptyObj, + ipcMain: emptyObj, + ipcRenderer: emptyObj, + Menu: emptyObj, + MenuItem: emptyObj, + dialog: emptyObj, + shell: emptyObj, + nativeImage: emptyObj, + clipboard: emptyObj, + safeStorage: emptyObj, + powerMonitor: emptyObj, + powerSaveBlocker: emptyObj, + autoUpdater: emptyObj, + crashReporter: emptyObj, + Notification: () => emptyObj, + Tray: () => emptyObj, + nativeTheme: emptyObj, + session: emptyObj, + screen: emptyObj, + protocol: emptyObj, + webContents: emptyObj, + systemPreferences: emptyObj, + contextBridge: emptyObj, + }, + { + get(target, prop) { + if (prop in target) { + return target[prop]; + } + return noop; + }, + }, +); + +export default stub; +export const app = stub.app; +export const BrowserWindow = stub.BrowserWindow; +export const ipcMain = stub.ipcMain; +export const ipcRenderer = stub.ipcRenderer; +export const Menu = stub.Menu; +export const MenuItem = stub.MenuItem; +export const dialog = stub.dialog; +export const shell = stub.shell; +export const nativeImage = stub.nativeImage; +export const clipboard = stub.clipboard; +export const safeStorage = stub.safeStorage; +export const powerMonitor = stub.powerMonitor; +export const powerSaveBlocker = stub.powerSaveBlocker; +export const autoUpdater = stub.autoUpdater; +export const crashReporter = stub.crashReporter; +export const Notification = stub.Notification; +export const Tray = stub.Tray; +export const nativeTheme = stub.nativeTheme; +export const session = stub.session; +export const screen = stub.screen; +export const protocol = stub.protocol; +export const webContents = stub.webContents; +export const systemPreferences = stub.systemPreferences; +export const contextBridge = stub.contextBridge; diff --git a/apps/code/scripts/scaffold-mcp-tools-loader.mjs b/apps/code/scripts/scaffold-mcp-tools-loader.mjs new file mode 100644 index 0000000000..b642a1c32c --- /dev/null +++ b/apps/code/scripts/scaffold-mcp-tools-loader.mjs @@ -0,0 +1,55 @@ +import * as fs from "node:fs"; +import * as path from "node:path"; +import { pathToFileURL } from "node:url"; + +const STUB_URL = pathToFileURL(`${import.meta.dirname}/electron-stub.mjs`).href; +const SRC_DIR = path.resolve(import.meta.dirname, "..", "src"); + +const PATH_ALIASES = { + "@main/": `${path.join(SRC_DIR, "main")}/`, + "@renderer/": `${path.join(SRC_DIR, "renderer")}/`, + "@shared/": `${path.join(SRC_DIR, "shared")}/`, + "@features/": `${path.join(SRC_DIR, "renderer", "features")}/`, + "@components/": `${path.join(SRC_DIR, "renderer", "components")}/`, + "@stores/": `${path.join(SRC_DIR, "renderer", "stores")}/`, + "@hooks/": `${path.join(SRC_DIR, "renderer", "hooks")}/`, + "@utils/": `${path.join(SRC_DIR, "renderer", "utils")}/`, + "@test/": `${path.join(SRC_DIR, "shared", "test")}/`, +}; + +// Some import targets exist as BOTH `foo.ts` AND `foo/` (sibling file + +// directory). Node ESM's default resolution picks the directory and looks for +// `index.json` — wrong. `bundler` moduleResolution (which tsconfig sets) and +// Vite prefer the `.ts` sibling. Replicate that by checking for the `.ts` +// file first and short-circuiting if it exists. +function preferFileSibling(absPath) { + for (const ext of [".ts", ".tsx", ".mjs", ".js"]) { + const candidate = `${absPath}${ext}`; + if (fs.existsSync(candidate)) { + return candidate; + } + } + return null; +} + +export function resolve(specifier, context, nextResolve) { + if (specifier === "electron") { + return { url: STUB_URL, format: "module", shortCircuit: true }; + } + for (const [prefix, target] of Object.entries(PATH_ALIASES)) { + if (specifier.startsWith(prefix)) { + const rel = specifier.slice(prefix.length); + const abs = path.join(target, rel); + const fileSibling = preferFileSibling(abs); + if (fileSibling) { + return { + url: pathToFileURL(fileSibling).href, + format: fileSibling.endsWith(".json") ? "json" : "module", + shortCircuit: true, + }; + } + return nextResolve(abs, context); + } + } + return nextResolve(specifier, context); +} diff --git a/apps/code/scripts/scaffold-mcp-tools-preload.mjs b/apps/code/scripts/scaffold-mcp-tools-preload.mjs new file mode 100644 index 0000000000..c5fb460899 --- /dev/null +++ b/apps/code/scripts/scaffold-mcp-tools-preload.mjs @@ -0,0 +1,10 @@ +import { register } from "node:module"; +import { pathToFileURL } from "node:url"; + +// Redirects `electron` imports to a local stub so the scaffold script can +// import the tRPC router (which transitively touches Electron-bound modules) +// without an Electron runtime present. +register( + "./scaffold-mcp-tools-loader.mjs", + pathToFileURL(`${import.meta.dirname}/`), +); diff --git a/apps/code/scripts/scaffold-mcp-tools.ts b/apps/code/scripts/scaffold-mcp-tools.ts new file mode 100644 index 0000000000..44fefb4dd0 --- /dev/null +++ b/apps/code/scripts/scaffold-mcp-tools.ts @@ -0,0 +1,183 @@ +#!/usr/bin/env tsx +/** + * Sync `apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml` + * with the live tRPC router. + * + * - Walks the router via `_def.procedures` and emits an `enabled: false` stub + * for every procedure that isn't already in the YAML. + * - Leaves existing entries untouched — your hand-authored config (title, + * description, annotations, param_overrides) is preserved. + * - Does NOT remove entries whose procedure has disappeared. It prints them + * as warnings; you decide whether to delete. Boot will hard-fail until you + * do, which is the forcing function. + * + * Usage: + * pnpm --filter code scaffold-mcp-tools + * pnpm --filter code scaffold-mcp-tools --check # exit 1 if out of date + */ + +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; + +// Imports below transitively load main-process services that read env vars +// at module-load time (see apps/code/src/main/utils/env.ts). Set defaults +// here so the script works outside an Electron context. Done before any +// dynamic import below. +if (!process.env.POSTHOG_CODE_DATA_DIR) { + process.env.POSTHOG_CODE_DATA_DIR = path.join( + os.tmpdir(), + "posthog-code-scaffold-mcp-tools", + ); +} +if (!process.env.POSTHOG_CODE_IS_DEV) process.env.POSTHOG_CODE_IS_DEV = "true"; +if (!process.env.POSTHOG_CODE_VERSION) { + process.env.POSTHOG_CODE_VERSION = "0.0.0-scaffold"; +} + +const YAML_PATH = path.resolve( + __dirname, + "..", + "src", + "main", + "services", + "posthog-code-internal-mcp", + "mcp-tools.yaml", +); + +const YAML_HEADER = `# Bridge from tRPC procedures to MCP tools exposed to the running agent. +# +# Re-run \`pnpm --filter code scaffold-mcp-tools\` after adding or removing +# tRPC procedures. New entries are scaffolded as enabled: false; the boot-time +# registry hard-fails if an entry references a procedure that no longer +# exists, so stale entries must be deleted by hand. +# +# Default-deny: every enabled tool is callable by the agent. Review carefully +# before flipping enabled: true on anything beyond the curated defaults. +`; + +interface Procedure { + path: string; + type: "query" | "mutation" | "subscription"; +} + +async function main(): Promise { + const check = process.argv.includes("--check"); + + // Dynamic import so the env defaults above are in place before the router's + // transitive deps load. + const { trpcRouter } = await import("../src/main/trpc/router"); + const { McpToolsYamlSchema } = await import( + "../src/main/services/posthog-code-internal-mcp/yaml-schema" + ); + const { parse: parseYaml, stringify: stringifyYaml } = await import("yaml"); + + const record = trpcRouter._def.procedures as Record< + string, + { _def: { type: "query" | "mutation" | "subscription" } } + >; + const procedures: Procedure[] = Object.entries(record) + .map(([p, proc]) => ({ path: p, type: proc._def.type })) + .filter((p) => p.type !== "subscription"); + + let existing: { tools: Record } = { + tools: {}, + }; + if (fs.existsSync(YAML_PATH)) { + const raw = fs.readFileSync(YAML_PATH, "utf-8"); + const parsed = parseYaml(raw); + const result = McpToolsYamlSchema.safeParse(parsed); + if (!result.success) { + console.error("Invalid existing mcp-tools.yaml:"); + for (const issue of result.error.issues) { + console.error(` ${issue.path.join(".")}: ${issue.message}`); + } + process.exit(1); + } + existing = result.data as { tools: Record }; + } + + const proceduresByPath = new Map(procedures.map((p) => [p.path, p])); + const existingByOperation = new Map(); + for (const [name, config] of Object.entries(existing.tools)) { + existingByOperation.set(config.operation, [name, config]); + } + + const mergedTools: Record = {}; + let added = 0; + let unchanged = 0; + const stale: string[] = []; + + for (const proc of procedures) { + const existingEntry = existingByOperation.get(proc.path); + if (existingEntry) { + const [name, config] = existingEntry; + mergedTools[name] = config; + unchanged++; + } else { + mergedTools[proc.path] = { + operation: proc.path, + enabled: false, + }; + added++; + } + } + + for (const [name, config] of Object.entries(existing.tools)) { + if (!proceduresByPath.has(config.operation)) { + mergedTools[name] = config; + stale.push(`${name} → ${config.operation}`); + } + } + + const sortedTools = Object.fromEntries( + Object.entries(mergedTools).sort(([a], [b]) => a.localeCompare(b)), + ); + + const nextContent = + YAML_HEADER + stringifyYaml({ tools: sortedTools }, { lineWidth: 120 }); + const currentContent = fs.existsSync(YAML_PATH) + ? fs.readFileSync(YAML_PATH, "utf-8") + : ""; + + const isUpToDate = currentContent === nextContent; + + if (check) { + if (!isUpToDate) { + console.error( + "mcp-tools.yaml is out of date with the tRPC router. Run `pnpm --filter code scaffold-mcp-tools` and commit the result.", + ); + console.error( + ` unchanged=${unchanged} added=${added} stale=${stale.length}`, + ); + process.exit(1); + } + console.log( + `mcp-tools.yaml is up to date (${procedures.length} procedures).`, + ); + return; + } + + if (!isUpToDate) { + fs.writeFileSync(YAML_PATH, nextContent); + } + + console.log( + `mcp-tools.yaml: ${procedures.length} procedures total — ${unchanged} unchanged, ${added} added.`, + ); + if (stale.length > 0) { + console.warn( + `\n⚠ ${stale.length} stale tool(s) in YAML reference procedures that no longer exist:`, + ); + for (const s of stale) console.warn(` - ${s}`); + console.warn( + " These were left in place. Delete them or boot will hard-fail.", + ); + process.exitCode = 2; + } +} + +void main().catch((err) => { + console.error(err); + process.exit(1); +}); diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index b2e2379419..66d50cdc5a 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -34,6 +34,7 @@ import { AuthProxyService } from "../services/auth-proxy/service"; import { CloudTaskService } from "../services/cloud-task/service"; import { ConnectivityService } from "../services/connectivity/service"; import { ContextMenuService } from "../services/context-menu/service"; +import { CustomInstructionsService } from "../services/custom-instructions/service"; import { DeepLinkService } from "../services/deep-link/service"; import { EnrichmentService } from "../services/enrichment/service"; import { EnvironmentService } from "../services/environment/service"; @@ -52,10 +53,12 @@ import { LlmGatewayService } from "../services/llm-gateway/service"; import { LocalLogsService } from "../services/local-logs/service"; import { McpAppsService } from "../services/mcp-apps/service"; import { McpCallbackService } from "../services/mcp-callback/service"; +import { McpInstallationsService } from "../services/mcp-installations/service"; import { McpProxyService } from "../services/mcp-proxy/service"; import { NewTaskLinkService } from "../services/new-task-link/service"; import { NotificationService } from "../services/notification/service"; import { OAuthService } from "../services/oauth/service"; +import { PostHogCodeInternalMcpService } from "../services/posthog-code-internal-mcp/service"; import { PosthogPluginService } from "../services/posthog-plugin/service"; import { ProcessTrackingService } from "../services/process-tracking/service"; import { ProvisioningService } from "../services/provisioning/service"; @@ -109,7 +112,14 @@ container.bind(MAIN_TOKENS.AgentAuthAdapter).to(AgentAuthAdapter); container.bind(MAIN_TOKENS.AgentService).to(AgentService); container.bind(MAIN_TOKENS.AuthService).to(AuthService); container.bind(MAIN_TOKENS.AuthProxyService).to(AuthProxyService); +container + .bind(MAIN_TOKENS.CustomInstructionsService) + .to(CustomInstructionsService); +container.bind(MAIN_TOKENS.McpInstallationsService).to(McpInstallationsService); container.bind(MAIN_TOKENS.McpProxyService).to(McpProxyService); +container + .bind(MAIN_TOKENS.PostHogCodeInternalMcpService) + .to(PostHogCodeInternalMcpService); container.bind(MAIN_TOKENS.ArchiveService).to(ArchiveService); container.bind(MAIN_TOKENS.SuspensionService).to(SuspensionService); container.bind(MAIN_TOKENS.AppLifecycleService).to(AppLifecycleService); diff --git a/apps/code/src/main/di/tokens.ts b/apps/code/src/main/di/tokens.ts index 69ea894b37..a6c4e4172d 100644 --- a/apps/code/src/main/di/tokens.ts +++ b/apps/code/src/main/di/tokens.ts @@ -43,7 +43,12 @@ export const MAIN_TOKENS = Object.freeze({ AgentService: Symbol.for("Main.AgentService"), AuthService: Symbol.for("Main.AuthService"), AuthProxyService: Symbol.for("Main.AuthProxyService"), + CustomInstructionsService: Symbol.for("Main.CustomInstructionsService"), + McpInstallationsService: Symbol.for("Main.McpInstallationsService"), McpProxyService: Symbol.for("Main.McpProxyService"), + PostHogCodeInternalMcpService: Symbol.for( + "Main.PostHogCodeInternalMcpService", + ), ArchiveService: Symbol.for("Main.ArchiveService"), SuspensionService: Symbol.for("Main.SuspensionService"), AppLifecycleService: Symbol.for("Main.AppLifecycleService"), diff --git a/apps/code/src/main/index.ts b/apps/code/src/main/index.ts index 6a005d365e..806623bb56 100644 --- a/apps/code/src/main/index.ts +++ b/apps/code/src/main/index.ts @@ -24,6 +24,7 @@ import { initializePostHog, trackAppEvent, } from "./services/posthog-analytics"; +import type { PostHogCodeInternalMcpService } from "./services/posthog-code-internal-mcp/service"; import type { PosthogPluginService } from "./services/posthog-plugin/service"; import type { SlackIntegrationService } from "./services/slack-integration/service"; import type { SuspensionService } from "./services/suspension/service"; @@ -159,6 +160,11 @@ async function initializeServices(): Promise { await authService.initialize(); + const internalMcp = container.get( + MAIN_TOKENS.PostHogCodeInternalMcpService, + ); + await internalMcp.start(); + // Initialize workspace branch watcher for live branch rename detection const workspaceService = container.get( MAIN_TOKENS.WorkspaceService, 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..dbd4c37cc5 100644 --- a/apps/code/src/main/services/agent/auth-adapter.test.ts +++ b/apps/code/src/main/services/agent/auth-adapter.test.ts @@ -58,6 +58,12 @@ function createDependencies() { (id: string) => `http://127.0.0.1:9998/${encodeURIComponent(id)}`, ), }, + internalMcp: { + getUrl: vi.fn().mockReturnValue("http://127.0.0.1:9997/mcp"), + getAuthHeader: vi + .fn() + .mockReturnValue({ name: "authorization", value: "Bearer test" }), + }, }; } @@ -77,6 +83,7 @@ describe("AgentAuthAdapter", () => { deps.authService as never, deps.authProxy as never, deps.mcpProxy as never, + deps.internalMcp as never, ); }); diff --git a/apps/code/src/main/services/agent/auth-adapter.ts b/apps/code/src/main/services/agent/auth-adapter.ts index 1cfa711fe0..afc1d618e2 100644 --- a/apps/code/src/main/services/agent/auth-adapter.ts +++ b/apps/code/src/main/services/agent/auth-adapter.ts @@ -11,6 +11,7 @@ import { logger } from "../../utils/logger"; import type { AuthService } from "../auth/service"; import type { AuthProxyService } from "../auth-proxy/service"; import type { McpProxyService } from "../mcp-proxy/service"; +import type { PostHogCodeInternalMcpService } from "../posthog-code-internal-mcp/service"; import type { Credentials } from "./schemas"; const log = logger.scope("agent-auth-adapter"); @@ -63,6 +64,8 @@ export class AgentAuthAdapter { private readonly authProxy: AuthProxyService, @inject(MAIN_TOKENS.McpProxyService) private readonly mcpProxy: McpProxyService, + @inject(MAIN_TOKENS.PostHogCodeInternalMcpService) + private readonly internalMcp: PostHogCodeInternalMcpService, ) {} createPosthogConfig(credentials: Credentials): AgentPosthogConfig { @@ -102,6 +105,19 @@ export class AgentAuthAdapter { ], }); + try { + servers.push({ + name: "posthog-code-internal", + type: "http", + url: this.internalMcp.getUrl(), + headers: [this.internalMcp.getAuthHeader()], + }); + } catch (err) { + // Service should always be running by the time the agent starts a task, + // but don't take down the whole MCP config if it isn't. + log.warn("posthog-code-internal MCP not available", { error: err }); + } + const installations = await this.fetchMcpInstallations(credentials); for (const installation of installations) { diff --git a/apps/code/src/main/services/agent/service.test.ts b/apps/code/src/main/services/agent/service.test.ts index 8507cc6075..2419ba71bd 100644 --- a/apps/code/src/main/services/agent/service.test.ts +++ b/apps/code/src/main/services/agent/service.test.ts @@ -200,6 +200,9 @@ function createMockDependencies() { addAdditionalDirectory: vi.fn(), removeAdditionalDirectory: vi.fn(), }, + internalMcp: { + on: vi.fn(), + }, }; } @@ -232,6 +235,7 @@ describe("AgentService", () => { deps.storagePaths as never, deps.defaultAdditionalDirectoryRepository as never, deps.workspaceRepository as never, + deps.internalMcp as never, ); }); diff --git a/apps/code/src/main/services/agent/service.ts b/apps/code/src/main/services/agent/service.ts index 1596f9ff5b..bac22cd95f 100644 --- a/apps/code/src/main/services/agent/service.ts +++ b/apps/code/src/main/services/agent/service.ts @@ -15,6 +15,7 @@ import { import { isMcpToolReadOnly, isNotification, + POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "@posthog/agent"; import type { McpToolApprovals } from "@posthog/agent/adapters/claude/mcp/tool-metadata"; @@ -53,6 +54,8 @@ import { logger } from "../../utils/logger"; import { TypedEventEmitter } from "../../utils/typed-event-emitter"; import type { FsService } from "../fs/service"; import type { McpAppsService } from "../mcp-apps/service"; +import { McpInstallationsServiceEvent } from "../mcp-installations/schemas"; +import type { McpInstallationsService } from "../mcp-installations/service"; import type { PosthogPluginService } from "../posthog-plugin/service"; import type { ProcessTrackingService } from "../process-tracking/service"; import { loadSessionEnvOverrides } from "../session-env/loader"; @@ -249,6 +252,8 @@ interface ManagedSession { mcpToolApprovals: McpToolApprovals; /** Maps tool keys to their installation for backend approval updates */ toolInstallations: McpToolInstallations; + /** Set when an MCP server is installed mid-turn; refresh runs after the turn ends. */ + pendingMcpRefresh: boolean; } /** Get the agent session ID from a managed session, throwing if not set. */ @@ -323,6 +328,8 @@ export class AgentService extends TypedEventEmitter { private readonly defaultAdditionalDirectoryRepository: IDefaultAdditionalDirectoryRepository, @inject(MAIN_TOKENS.WorkspaceRepository) private readonly workspaceRepository: IWorkspaceRepository, + @inject(MAIN_TOKENS.McpInstallationsService) + mcpInstallations: McpInstallationsService, ) { super(); this.processTracking = processTracking; @@ -333,6 +340,9 @@ export class AgentService extends TypedEventEmitter { this.mcpAppsService = mcpAppsService; powerManager.onResume(() => this.checkIdleDeadlines()); + mcpInstallations.on(McpInstallationsServiceEvent.Installed, () => { + void this.refreshAllSessionMcpServers(); + }); } private getClaudeCliPath(): string { @@ -414,6 +424,46 @@ export class AgentService extends TypedEventEmitter { this.recordActivity(taskRunId); } + private async refreshSessionMcpServers( + session: ManagedSession, + ): Promise { + try { + const { servers } = await this.agentAuthAdapter.buildMcpServers( + session.config.credentials, + ); + await session.clientSideConnection.extMethod( + POSTHOG_METHODS.REFRESH_SESSION, + { mcpServers: servers }, + ); + log.info("Refreshed MCP servers for session", { + taskRunId: session.taskRunId, + serverCount: servers.length, + }); + } catch (err) { + log.warn("Failed to refresh MCP servers for session", { + taskRunId: session.taskRunId, + err, + }); + } + } + + private async refreshAllSessionMcpServers(): Promise { + const refreshable: ManagedSession[] = []; + for (const session of this.sessions.values()) { + if (session.promptPending) { + // ACP refresh contract requires no prompt in flight; defer until the + // turn completes (see prompt() finally block). + session.pendingMcpRefresh = true; + log.info("Deferring MCP refresh until current turn ends", { + taskRunId: session.taskRunId, + }); + continue; + } + refreshable.push(session); + } + await Promise.all(refreshable.map((s) => this.refreshSessionMcpServers(s))); + } + /** * Check if any sessions are currently active (i.e. have a prompt pending). */ @@ -847,6 +897,7 @@ When creating pull requests, add the following footer at the end of the PR descr inFlightMcpToolCalls: new Map(), mcpToolApprovals: toolApprovals, toolInstallations, + pendingMcpRefresh: false, }; this.sessions.set(taskRunId, session); @@ -935,6 +986,11 @@ When creating pull requests, add the following footer at the end of the PR descr this.recordActivity(sessionId); this.sleepService.release(sessionId); + if (session.pendingMcpRefresh) { + session.pendingMcpRefresh = false; + void this.refreshSessionMcpServers(session); + } + if (!this.hasActiveSessions()) { this.emit(AgentServiceEvent.SessionsIdle, undefined); } diff --git a/apps/code/src/main/services/custom-instructions/schemas.ts b/apps/code/src/main/services/custom-instructions/schemas.ts new file mode 100644 index 0000000000..56be5b2f46 --- /dev/null +++ b/apps/code/src/main/services/custom-instructions/schemas.ts @@ -0,0 +1,45 @@ +import { z } from "zod"; + +export const readCustomInstructionsOutput = z.object({ + customInstructions: z.string(), +}); + +export type ReadCustomInstructionsOutput = z.infer< + typeof readCustomInstructionsOutput +>; + +export const writeCustomInstructionsInput = z.object({ + instructions: z + .string() + .describe( + "Full replacement text for the user's custom instructions. Pass empty string to clear.", + ), +}); + +export type WriteCustomInstructionsInput = z.infer< + typeof writeCustomInstructionsInput +>; + +export const writeCustomInstructionsOutput = z.object({ + ok: z.literal(true), +}); + +export type WriteCustomInstructionsOutput = z.infer< + typeof writeCustomInstructionsOutput +>; + +export const customInstructionsChanged = z.object({ + customInstructions: z.string(), +}); + +export type CustomInstructionsChanged = z.infer< + typeof customInstructionsChanged +>; + +export const CustomInstructionsServiceEvent = { + Changed: "changed", +} as const; + +export interface CustomInstructionsServiceEvents { + [CustomInstructionsServiceEvent.Changed]: CustomInstructionsChanged; +} diff --git a/apps/code/src/main/services/custom-instructions/service.ts b/apps/code/src/main/services/custom-instructions/service.ts new file mode 100644 index 0000000000..c0e1243f09 --- /dev/null +++ b/apps/code/src/main/services/custom-instructions/service.ts @@ -0,0 +1,64 @@ +import { injectable } from "inversify"; +import { decrypt, encrypt } from "../../utils/encryption"; +import { logger } from "../../utils/logger"; +import { rendererStore } from "../../utils/store"; +import { TypedEventEmitter } from "../../utils/typed-event-emitter"; +import { + CustomInstructionsServiceEvent, + type CustomInstructionsServiceEvents, +} from "./schemas"; + +const log = logger.scope("custom-instructions"); + +const SETTINGS_STORE_KEY = "settings-storage"; + +/** + * Owns reading/writing the user's custom instructions (extra guidance the + * user has appended to every agent prompt). Persisted inside the encrypted + * `settings-storage` bucket alongside other settings the renderer manages, + * so a write here is visible to the settings store on next reload — and + * immediately via the {@link CustomInstructionsServiceEvent.Changed} event. + */ +@injectable() +export class CustomInstructionsService extends TypedEventEmitter { + read(): string { + if (!rendererStore.has(SETTINGS_STORE_KEY)) return ""; + const encrypted = rendererStore.get(SETTINGS_STORE_KEY) as string; + const raw = decrypt(encrypted); + if (!raw) return ""; + try { + const parsed = JSON.parse(raw) as { + state?: { customInstructions?: string }; + }; + return parsed.state?.customInstructions ?? ""; + } catch (err) { + log.warn("Failed to parse settings-storage", { err }); + return ""; + } + } + + write(value: string): void { + let parsed: { state?: Record; version?: number } = { + state: {}, + version: 0, + }; + if (rendererStore.has(SETTINGS_STORE_KEY)) { + const encrypted = rendererStore.get(SETTINGS_STORE_KEY) as string; + const raw = decrypt(encrypted); + if (raw) { + try { + parsed = JSON.parse(raw); + } catch (err) { + log.warn("Settings store corrupted, overwriting with new state", { + err, + }); + } + } + } + parsed.state = { ...(parsed.state ?? {}), customInstructions: value }; + rendererStore.set(SETTINGS_STORE_KEY, encrypt(JSON.stringify(parsed))); + this.emit(CustomInstructionsServiceEvent.Changed, { + customInstructions: value, + }); + } +} diff --git a/apps/code/src/main/services/mcp-installations/schemas.ts b/apps/code/src/main/services/mcp-installations/schemas.ts new file mode 100644 index 0000000000..404813f3d7 --- /dev/null +++ b/apps/code/src/main/services/mcp-installations/schemas.ts @@ -0,0 +1,79 @@ +import { z } from "zod"; + +export const mcpServerInstallation = z.object({ + id: z.string(), + name: z.string(), + url: z.string(), + auth_type: z.string(), + is_enabled: z.boolean(), + pending_oauth: z.boolean(), + needs_reauth: z.boolean(), +}); + +export type McpServerInstallation = z.infer; + +export const listMcpInstallationsOutput = z.object({ + servers: z.array(mcpServerInstallation), +}); + +export type ListMcpInstallationsOutput = z.infer< + typeof listMcpInstallationsOutput +>; + +export const installCustomInput = z.object({ + name: z + .string() + .min(1) + .describe("Display name for the MCP server installation."), + url: z.string().url().describe("HTTPS URL of the MCP server endpoint."), + auth_type: z + .enum(["api_key", "oauth"]) + .default("api_key") + .describe( + "Authentication scheme. Use 'api_key' with `api_key` for static bearer tokens; use 'oauth' to start an OAuth handshake (response will include redirectUrl).", + ), + api_key: z + .string() + .optional() + .describe( + "Static bearer token. Required when auth_type='api_key' and the target server requires authentication.", + ), + description: z.string().optional(), +}); + +export type InstallCustomInput = z.infer; + +export const installCustomOk = z.object({ + kind: z.literal("ok"), + id: z.string(), + message: z.string(), +}); + +export const installCustomOAuth = z.object({ + kind: z.literal("oauth"), + id: z.string().optional(), + redirectUrl: z.string(), + message: z.string(), +}); + +export const installCustomError = z.object({ + kind: z.literal("error"), + status: z.number().optional(), + message: z.string(), +}); + +export const installCustomOutput = z.discriminatedUnion("kind", [ + installCustomOk, + installCustomOAuth, + installCustomError, +]); + +export type InstallCustomOutput = z.infer; + +export const McpInstallationsServiceEvent = { + Installed: "installed", +} as const; + +export interface McpInstallationsServiceEvents { + [McpInstallationsServiceEvent.Installed]: Record; +} diff --git a/apps/code/src/main/services/mcp-installations/service.test.ts b/apps/code/src/main/services/mcp-installations/service.test.ts new file mode 100644 index 0000000000..d7a6512e51 --- /dev/null +++ b/apps/code/src/main/services/mcp-installations/service.test.ts @@ -0,0 +1,105 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +// electron-store mkdir's userDataDir at import time, which fails in CI where +// the default mocked path (/mock/userData) isn't writable. The tests below +// don't exercise the store paths, so a no-op mock is safe. +vi.mock("../../utils/store", () => ({ + rendererStore: { + has: () => false, + get: () => undefined, + set: () => {}, + }, +})); + +import { McpInstallationsServiceEvent } from "./schemas"; +import { McpInstallationsService } from "./service"; + +interface FakeAuthService { + getValidAccessToken: () => Promise<{ apiHost: string; token: string }>; + getState: () => { projectId: number }; + authenticatedFetch: ( + fetchImpl: typeof fetch, + url: string, + init?: RequestInit, + ) => Promise; +} + +const createFakeAuth = ( + fetchImpl: (url: string) => Promise, +): FakeAuthService => ({ + getValidAccessToken: async () => ({ + apiHost: "https://example.com", + token: "t", + }), + getState: () => ({ projectId: 1 }), + authenticatedFetch: async (_f, url) => fetchImpl(String(url)), +}); + +const okJson = (body: unknown): Response => + new Response(JSON.stringify(body), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + +describe("McpInstallationsService.pollForOauthCompletion", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("emits Installed when pending_oauth flips to false", async () => { + const responses = [ + okJson({ + results: [ + { id: "abc", name: "linear", pending_oauth: true, is_enabled: true }, + ], + }), + okJson({ + results: [ + { id: "abc", name: "linear", pending_oauth: false, is_enabled: true }, + ], + }), + ]; + const auth = createFakeAuth(async () => { + const next = responses.shift(); + if (!next) throw new Error("no more responses"); + return next; + }); + const service = new McpInstallationsService(auth as never); + const handler = vi.fn(); + service.on(McpInstallationsServiceEvent.Installed, handler); + + const poll = ( + service as unknown as { + pollForOauthCompletion: (id: string, name: string) => Promise; + } + ).pollForOauthCompletion("abc", "linear"); + + await vi.advanceTimersByTimeAsync(3500); + await vi.advanceTimersByTimeAsync(3500); + await poll; + + expect(handler).toHaveBeenCalledOnce(); + }); + + it("stops polling when installation disappears", async () => { + const auth = createFakeAuth(async () => okJson({ results: [] })); + const service = new McpInstallationsService(auth as never); + const handler = vi.fn(); + service.on(McpInstallationsServiceEvent.Installed, handler); + + const poll = ( + service as unknown as { + pollForOauthCompletion: (id: string, name: string) => Promise; + } + ).pollForOauthCompletion("abc", "linear"); + + await vi.advanceTimersByTimeAsync(3500); + await poll; + + expect(handler).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/code/src/main/services/mcp-installations/service.ts b/apps/code/src/main/services/mcp-installations/service.ts new file mode 100644 index 0000000000..0f285bfb9b --- /dev/null +++ b/apps/code/src/main/services/mcp-installations/service.ts @@ -0,0 +1,189 @@ +import { inject, injectable } from "inversify"; +import { MAIN_TOKENS } from "../../di/tokens"; +import { logger } from "../../utils/logger"; +import { TypedEventEmitter } from "../../utils/typed-event-emitter"; +import type { AuthService } from "../auth/service"; +import { + type InstallCustomInput, + type InstallCustomOutput, + type ListMcpInstallationsOutput, + McpInstallationsServiceEvent, + type McpInstallationsServiceEvents, +} from "./schemas"; + +const log = logger.scope("mcp-installations"); + +const OAUTH_POLL_INTERVAL_MS = 3000; +const OAUTH_POLL_TIMEOUT_MS = 10 * 60 * 1000; + +/** + * Owns the PostHog REST surface for MCP server installations on the active + * project. Reads/writes against `/api/environments/{projectId}/mcp_server_installations/`. + * + * Emits {@link McpInstallationsServiceEvent.Installed} once an installation + * becomes usable — either immediately for api_key installs, or after the + * background OAuth poll observes `pending_oauth: false`. + */ +@injectable() +export class McpInstallationsService extends TypedEventEmitter { + constructor( + @inject(MAIN_TOKENS.AuthService) + private readonly authService: AuthService, + ) { + super(); + } + + async list(): Promise { + const { baseUrl, projectId } = await this.requireProject(); + const url = `${baseUrl}/api/environments/${projectId}/mcp_server_installations/`; + const response = await this.authService.authenticatedFetch(fetch, url, { + headers: { "Content-Type": "application/json" }, + }); + if (!response.ok) { + const errText = await response.text().catch(() => ""); + throw new Error( + `Failed to list MCP servers (${response.status}): ${errText.slice(0, 500)}`, + ); + } + const data = (await response.json()) as { + results?: Array<{ + id: string; + name?: string; + display_name?: string; + url?: string; + auth_type?: string; + is_enabled?: boolean; + pending_oauth?: boolean; + needs_reauth?: boolean; + }>; + }; + const servers = (data.results ?? []).map((i) => ({ + id: i.id, + name: i.name ?? i.display_name ?? "(unnamed)", + url: i.url ?? "", + auth_type: i.auth_type ?? "unknown", + is_enabled: i.is_enabled !== false, + pending_oauth: !!i.pending_oauth, + needs_reauth: !!i.needs_reauth, + })); + return { servers }; + } + + async installCustom(input: InstallCustomInput): Promise { + const { baseUrl, projectId } = await this.requireProject(); + const url = `${baseUrl}/api/environments/${projectId}/mcp_server_installations/install_custom/`; + const response = await this.authService.authenticatedFetch(fetch, url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + name: input.name, + url: input.url, + auth_type: input.auth_type, + api_key: input.api_key, + description: input.description, + install_source: "posthog-code", + }), + }); + if (!response.ok) { + const errText = await response.text().catch(() => ""); + return { + kind: "error", + status: response.status, + message: `Failed to install MCP server (${response.status}): ${errText.slice(0, 500)}`, + }; + } + const data = (await response.json()) as Record; + const id = typeof data.id === "string" ? data.id : undefined; + if (typeof data.redirect_url === "string") { + void this.pollForOauthCompletion(id, input.name); + return { + kind: "oauth", + id, + redirectUrl: data.redirect_url, + message: `OAuth flow required. The user must visit: ${data.redirect_url} to finish installing "${input.name}". Once authorized, the session will refresh automatically.`, + }; + } + this.emit(McpInstallationsServiceEvent.Installed, {}); + return { + kind: "ok", + id: id ?? "unknown", + message: `Installed MCP server "${input.name}" (id=${id ?? "unknown"}). Refreshing session to make it available immediately.`, + }; + } + + private async requireProject(): Promise<{ + baseUrl: string; + projectId: number; + }> { + const { apiHost } = await this.authService.getValidAccessToken(); + const projectId = this.authService.getState().projectId; + if (!projectId) { + throw new Error( + "No project selected. Sign in and pick a project before listing MCP servers.", + ); + } + return { + baseUrl: apiHost.replace(/\/+$/, ""), + projectId, + }; + } + + private async pollForOauthCompletion( + installationId: string | undefined, + name: string, + ): Promise { + const { apiHost } = await this.authService.getValidAccessToken(); + const projectId = this.authService.getState().projectId; + if (!projectId) return; + const baseUrl = apiHost.replace(/\/+$/, ""); + const url = `${baseUrl}/api/environments/${projectId}/mcp_server_installations/`; + + log.info("Polling for OAuth completion", { installationId, name }); + + const start = Date.now(); + while (Date.now() - start < OAUTH_POLL_TIMEOUT_MS) { + await new Promise((resolve) => + setTimeout(resolve, OAUTH_POLL_INTERVAL_MS), + ); + + try { + const response = await this.authService.authenticatedFetch(fetch, url, { + headers: { "Content-Type": "application/json" }, + }); + if (!response.ok) continue; + const data = (await response.json()) as { + results?: Array<{ + id: string; + name?: string; + display_name?: string; + pending_oauth?: boolean; + is_enabled?: boolean; + }>; + }; + const inst = (data.results ?? []).find((i) => + installationId + ? i.id === installationId + : i.name === name || i.display_name === name, + ); + if (!inst) { + log.info("OAuth installation no longer in list, stopping poll", { + installationId, + name, + }); + return; + } + if (!inst.pending_oauth && inst.is_enabled !== false) { + log.info("OAuth install completed, triggering session refresh", { + installationId: inst.id, + name, + }); + this.emit(McpInstallationsServiceEvent.Installed, {}); + return; + } + } catch (err) { + log.warn("OAuth poll error", { err }); + } + } + log.info("OAuth poll timed out", { installationId, name }); + } +} diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml b/apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml new file mode 100644 index 0000000000..c8e908bbe1 --- /dev/null +++ b/apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml @@ -0,0 +1,58 @@ +# Bridge from tRPC procedures to MCP tools exposed to the running agent. +# +# Re-run `pnpm --filter code scaffold-mcp-tools` after adding or removing +# tRPC procedures. New entries are scaffolded as enabled: false; the boot-time +# registry hard-fails if an entry references a procedure that no longer +# exists, so stale entries must be deleted by hand. +# +# Default-deny: every enabled tool is callable by the agent. Review carefully +# before flipping enabled: true on anything beyond the curated defaults. +tools: + customInstructions.read: + operation: customInstructions.read + enabled: true + title: Read custom instructions + description: > + Return the user's custom instructions — extra guidance the user has appended to every agent prompt. + Returns an empty string if no custom instructions are configured. + annotations: + readOnly: true + destructive: false + idempotent: true + + customInstructions.write: + operation: customInstructions.write + enabled: true + title: Write custom instructions + description: > + Replace the user's custom instructions. Pass the full new text — this overwrites the existing value. + Pass an empty string to clear. + annotations: + readOnly: false + destructive: true + idempotent: false + + mcpInstallations.list: + operation: mcpInstallations.list + enabled: true + title: List installed MCP servers + description: > + List the MCP server installations on the current PostHog project. Returns id, name, url, auth_type, + and status flags for each installation. + annotations: + readOnly: true + destructive: false + idempotent: true + + mcpInstallations.installCustom: + operation: mcpInstallations.installCustom + enabled: true + title: Install a custom MCP server + description: > + Install a new MCP server on the current PostHog project. Use auth_type="api_key" with `api_key` + for static bearer tokens; use auth_type="oauth" to start an OAuth handshake — the response will + include a redirectUrl the user must visit. + annotations: + readOnly: false + destructive: false + idempotent: false diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/service.ts b/apps/code/src/main/services/posthog-code-internal-mcp/service.ts new file mode 100644 index 0000000000..da76ad3911 --- /dev/null +++ b/apps/code/src/main/services/posthog-code-internal-mcp/service.ts @@ -0,0 +1,167 @@ +import { randomBytes } from "node:crypto"; +import http from "node:http"; +import * as path from "node:path"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import { injectable, preDestroy } from "inversify"; +import { trpcRouter } from "../../trpc/router"; +import { logger } from "../../utils/logger"; +import { McpToolRegistry } from "./tool-registry"; + +const log = logger.scope("posthog-code-internal-mcp"); + +const SERVER_NAME = "posthog-code-internal"; +const SERVER_VERSION = "1.0.0"; + +const YAML_PATH = path.join(__dirname, "mcp-tools.yaml"); + +/** + * Local-only HTTP MCP server that exposes a curated subset of the app's + * tRPC router to the running agent. The bridge from tRPC procedure to MCP + * tool is driven entirely by `mcp-tools.yaml` — see {@link McpToolRegistry}. + * + * Mirrors `McpProxyService`: listens on 127.0.0.1, generates a random bearer + * token at boot, and dies with the app via `@preDestroy`. Configuration is + * loaded once at start; YAML or procedure changes require an app restart. + */ +@injectable() +export class PostHogCodeInternalMcpService { + private server: http.Server | null = null; + private port: number | null = null; + private bearerToken: string | null = null; + private startPromise: Promise | null = null; + private registry: McpToolRegistry | null = null; + + async start(): Promise { + if (this.server && this.port) return; + if (this.startPromise) return this.startPromise; + this.startPromise = this.doStart().catch((err) => { + this.startPromise = null; + throw err; + }); + return this.startPromise; + } + + @preDestroy() + async stop(): Promise { + if (!this.server) return; + const server = this.server; + await new Promise((resolve) => { + server.close(() => { + log.info("PostHog Code internal MCP stopped"); + resolve(); + }); + }); + this.server = null; + this.port = null; + this.bearerToken = null; + this.startPromise = null; + this.registry = null; + } + + getUrl(): string { + if (!this.port) { + throw new Error("posthog-code-internal MCP server not started"); + } + return `http://127.0.0.1:${this.port}/mcp`; + } + + getAuthHeader(): { name: string; value: string } { + if (!this.bearerToken) { + throw new Error("posthog-code-internal MCP server not started"); + } + return { name: "authorization", value: `Bearer ${this.bearerToken}` }; + } + + private async doStart(): Promise { + // Build the registry before we listen — any YAML/router mismatch should + // fail the whole service start with a readable error in main logs. + this.registry = McpToolRegistry.build({ + router: trpcRouter, + yamlPath: YAML_PATH, + }); + + this.bearerToken = randomBytes(32).toString("hex"); + + const server = http.createServer((req, res) => { + void this.handleRequest(req, res); + }); + this.server = server; + + await new Promise((resolve, reject) => { + server.listen(0, "127.0.0.1", () => { + const addr = server.address(); + if (typeof addr === "object" && addr) { + this.port = addr.port; + log.info("PostHog Code internal MCP started", { port: this.port }); + resolve(); + } else { + reject(new Error("Failed to get internal MCP address")); + } + }); + server.on("error", (err) => { + log.error("Internal MCP server error", err); + reject(err); + }); + }); + } + + private async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + ): Promise { + const auth = req.headers.authorization; + if (!auth || auth !== `Bearer ${this.bearerToken}`) { + res.writeHead(401).end("Unauthorized"); + return; + } + + let mcpServer: McpServer | null = null; + let transport: StreamableHTTPServerTransport | null = null; + try { + // Stateless per-request: each HTTP request gets a fresh server + + // transport. Avoids cross-request session state and matches the SDK's + // documented stateless pattern. + mcpServer = this.buildServer(); + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + }); + const owned = { mcpServer, transport }; + res.on("close", () => { + try { + owned.transport.close(); + } catch {} + try { + owned.mcpServer.close(); + } catch {} + }); + await mcpServer.connect(transport); + await transport.handleRequest(req, res); + } catch (err) { + log.error("Internal MCP request error", err); + try { + transport?.close(); + } catch {} + try { + mcpServer?.close(); + } catch {} + if (!res.headersSent) { + res.writeHead(500).end("Internal error"); + } else { + res.end(); + } + } + } + + private buildServer(): McpServer { + if (!this.registry) { + throw new Error("posthog-code-internal MCP registry not initialized"); + } + const server = new McpServer( + { name: SERVER_NAME, version: SERVER_VERSION }, + { capabilities: { tools: {} } }, + ); + this.registry.registerAll(server, trpcRouter); + return server; + } +} diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.test.ts b/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.test.ts new file mode 100644 index 0000000000..f165fd2163 --- /dev/null +++ b/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.test.ts @@ -0,0 +1,342 @@ +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { initTRPC, TRPCError } from "@trpc/server"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { z } from "zod"; +import { McpToolRegistry } from "./tool-registry"; + +const t = initTRPC.create(); +const procedure = t.procedure; + +function makeRouter() { + return t.router({ + ping: procedure + .input(z.object({ msg: z.string() })) + .query(({ input }) => ({ pong: input.msg })), + settings: t.router({ + read: procedure.query(() => ({ value: "hello" })), + write: procedure + .input(z.object({ value: z.string() })) + .mutation(({ input }) => ({ ok: true as const, value: input.value })), + writeNumber: procedure + .input(z.object({ count: z.number().describe("How many") })) + .mutation(({ input }) => ({ count: input.count })), + }), + nonObjectInput: procedure + .input(z.string()) + .query(({ input }) => ({ input })), + sub: procedure.subscription(async function* () { + yield { tick: 1 }; + }), + failing: procedure.query(() => { + throw new TRPCError({ code: "FORBIDDEN", message: "not allowed" }); + }), + }); +} + +let tmpDir: string; +let yamlPath: string; + +beforeEach(() => { + tmpDir = mkdtempSync(join(tmpdir(), "mcp-tool-registry-")); + yamlPath = join(tmpDir, "mcp-tools.yaml"); +}); + +afterEach(() => { + rmSync(tmpDir, { recursive: true, force: true }); +}); + +function writeYaml(content: string) { + writeFileSync(yamlPath, content); +} + +function makeFakeMcpServer() { + const registered: Array<{ + name: string; + config: { + description?: string; + inputSchema?: unknown; + annotations?: unknown; + title?: string; + }; + handler: (args: Record) => Promise; + }> = []; + return { + server: { + registerTool: vi.fn((name, config, handler) => { + registered.push({ name, config, handler }); + }), + } as never, + registered, + }; +} + +describe("McpToolRegistry.build", () => { + it("hard-fails when YAML references a missing procedure", () => { + writeYaml(`tools: + bogus: + operation: not.a.real.path + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/references operation "not.a.real.path" which does not exist/); + }); + + it("hard-fails when an enabled tool targets a subscription", () => { + writeYaml(`tools: + sub: + operation: sub + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/subscription procedure/); + }); + + it("hard-fails when an enabled tool's input is not a ZodObject", () => { + writeYaml(`tools: + nonObjectInput: + operation: nonObjectInput + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/not a z.object/); + }); + + it("hard-fails when YAML lacks annotations on enabled tool", () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/enabled tools must declare annotations/); + }); + + it("hard-fails when param_overrides reference an unknown field", () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } + param_overrides: + missing: + description: x +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/overrides param "missing"/); + }); + + it("hard-fails when exclude_params reference an unknown field", () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } + exclude_params: ["missing"] +`); + expect(() => + McpToolRegistry.build({ router: makeRouter(), yamlPath }), + ).toThrow(/excludes param "missing"/); + }); + + it("skips disabled tools even if their procedure is broken", () => { + writeYaml(`tools: + sub: + operation: sub + enabled: false + ping: + operation: ping + enabled: false +`); + const reg = McpToolRegistry.build({ router: makeRouter(), yamlPath }); + expect(reg.getEntries()).toHaveLength(0); + }); +}); + +describe("McpToolRegistry.registerAll", () => { + it("registers an enabled query with description and annotations", async () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + title: Ping + description: Echoes the message back. + annotations: + readOnly: true + destructive: false + idempotent: true +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + + reg.registerAll(server, router); + + expect(registered).toHaveLength(1); + expect(registered[0].name).toBe("ping"); + expect(registered[0].config.title).toBe("Ping"); + expect(registered[0].config.description).toContain("Echoes"); + expect(registered[0].config.annotations).toEqual({ + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + }); + + const result = (await registered[0].handler({ msg: "hi" })) as { + content: Array<{ text: string }>; + }; + expect(result.content[0].text).toContain('"pong"'); + expect(result.content[0].text).toContain('"hi"'); + }); + + it("invokes the underlying procedure via createCaller", async () => { + writeYaml(`tools: + settings.write: + operation: settings.write + enabled: true + annotations: { readOnly: false, destructive: true, idempotent: false } +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + const result = (await registered[0].handler({ + value: "foo", + })) as { content: Array<{ text: string }> }; + expect(result.content[0].text).toContain('"ok": true'); + expect(result.content[0].text).toContain('"value": "foo"'); + }); + + it("formats TRPCError as MCP isError", async () => { + writeYaml(`tools: + failing: + operation: failing + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + const result = (await registered[0].handler({})) as { + isError: boolean; + content: Array<{ text: string }>; + }; + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain("FORBIDDEN"); + expect(result.content[0].text).toContain("not allowed"); + }); + + it("formats Zod validation failures as MCP isError", async () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + // msg is required; pass nothing and let tRPC's Zod parse reject it. + const result = (await registered[0].handler({})) as { + isError: boolean; + content: Array<{ text: string }>; + }; + expect(result.isError).toBe(true); + expect(result.content[0].text.toLowerCase()).toMatch( + /msg|required|invalid/, + ); + }); + + it("applies param_overrides descriptions to inputSchema shape", async () => { + writeYaml(`tools: + settings.writeNumber: + operation: settings.writeNumber + enabled: true + annotations: { readOnly: false, destructive: false, idempotent: false } + param_overrides: + count: + description: Overridden description for count +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + const shape = ( + registered[0].config.inputSchema as Record + ).count; + expect(shape.description).toBe("Overridden description for count"); + }); + + it("rename_params remaps the MCP-facing key but calls procedure with original", async () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } + rename_params: + msg: message +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + const shape = registered[0].config.inputSchema as Record; + expect(Object.keys(shape)).toEqual(["message"]); + + const result = (await registered[0].handler({ + message: "renamed", + })) as { content: Array<{ text: string }> }; + expect(result.content[0].text).toContain("renamed"); + }); + + it("exclude_params omits keys from the exposed inputSchema", async () => { + writeYaml(`tools: + ping: + operation: ping + enabled: true + annotations: { readOnly: true, destructive: false, idempotent: true } + exclude_params: ["msg"] +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + const shape = registered[0].config.inputSchema as Record; + expect(Object.keys(shape)).toEqual([]); + }); + + it("loads description from a sibling description_file", async () => { + const descPath = join(tmpDir, "ping-desc.txt"); + writeFileSync(descPath, " Description from file \n"); + writeYaml(`tools: + ping: + operation: ping + enabled: true + description_file: ping-desc.txt + annotations: { readOnly: true, destructive: false, idempotent: true } +`); + const router = makeRouter(); + const reg = McpToolRegistry.build({ router, yamlPath }); + const { server, registered } = makeFakeMcpServer(); + reg.registerAll(server, router); + + expect(registered[0].config.description).toBe("Description from file"); + }); +}); diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.ts b/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.ts new file mode 100644 index 0000000000..7f0f929b2b --- /dev/null +++ b/apps/code/src/main/services/posthog-code-internal-mcp/tool-registry.ts @@ -0,0 +1,391 @@ +import * as fs from "node:fs"; +import * as path from "node:path"; +import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; +import type { AnyProcedure, AnyRouter } from "@trpc/server"; +import { TRPCError } from "@trpc/server"; +import { parse as parseYaml } from "yaml"; +import { z } from "zod"; +import { logger } from "../../utils/logger"; +import { McpToolsYamlSchema, type ToolConfig } from "./yaml-schema"; + +const log = logger.scope("mcp-tool-registry"); + +interface ProcedureEntry { + path: string; + type: "query" | "mutation" | "subscription"; + inputSchema: z.ZodTypeAny | undefined; +} + +interface RegisteredEntry { + toolName: string; + config: Required> & + ToolConfig; + procedure: ProcedureEntry; + /** Final Zod object after applying exclude_params + param_overrides. */ + inputSchema: z.ZodObject | null; + /** Inverse map of rename_params: MCP-side key → tRPC-side key. */ + renameInverse: Record; +} + +/** + * Generic bridge from tRPC procedures to MCP tools, driven by `mcp-tools.yaml`. + * + * Design constraints: + * - Default-deny: every procedure must be explicitly enabled in YAML. + * - Hard-fail at boot on any mismatch (missing procedure, malformed config, + * non-ZodObject input on an enabled tool, etc.) — silent drift is worse + * than a startup crash that names exactly what to fix. + * - Context-free: procedures are invoked via `router.createCaller({})`, so + * they must resolve their dependencies via the main-process DI container + * rather than tRPC context. True today for every router under + * `apps/code/src/main/trpc/routers/` — if that ever changes, this registry + * must thread a real ctx. + */ +export class McpToolRegistry { + constructor(private readonly entries: RegisteredEntry[]) {} + + static build({ + router, + yamlPath, + }: { + router: AnyRouter; + yamlPath: string; + }): McpToolRegistry { + if (!fs.existsSync(yamlPath)) { + throw new Error(`mcp-tools.yaml not found at ${yamlPath}`); + } + const raw = fs.readFileSync(yamlPath, "utf-8"); + const parsed = parseYaml(raw); + const validated = McpToolsYamlSchema.safeParse(parsed); + if (!validated.success) { + const issues = validated.error.issues + .map((i) => ` ${i.path.join(".")}: ${i.message}`) + .join("\n"); + throw new Error(`Invalid mcp-tools.yaml:\n${issues}`); + } + + const procedures = collectProcedures(router); + const proceduresByPath = new Map(procedures.map((p) => [p.path, p])); + const yamlDir = path.dirname(yamlPath); + + const entries: RegisteredEntry[] = []; + for (const [toolName, config] of Object.entries(validated.data.tools)) { + const procedure = proceduresByPath.get(config.operation); + if (!procedure) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" references operation "${config.operation}" which does not exist on the tRPC router. ` + + `Run \`pnpm --filter code scaffold-mcp-tools\` to sync.`, + ); + } + if (!config.enabled) continue; + + if (procedure.type === "subscription") { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" targets subscription procedure "${procedure.path}". ` + + `Subscriptions cannot be exposed as MCP tools — set enabled: false.`, + ); + } + + const { inputSchema, renameInverse } = composeInputSchema({ + toolName, + procedure, + config, + }); + + const resolvedDescription = resolveDescription(config, yamlDir); + // refinements above guarantee annotations is present when enabled, but + // narrow defensively rather than asserting. + const annotations = config.annotations; + if (!annotations) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" is enabled but missing annotations.`, + ); + } + + entries.push({ + toolName, + config: { + ...config, + annotations, + description: resolvedDescription, + }, + procedure, + inputSchema, + renameInverse, + }); + } + + log.info("MCP tool registry built", { + enabled: entries.length, + total: Object.keys(validated.data.tools).length, + proceduresInRouter: procedures.length, + }); + + return new McpToolRegistry(entries); + } + + /** Register every enabled tool with the given MCP server instance. */ + registerAll(server: McpServer, router: AnyRouter): void { + for (const entry of this.entries) { + this.registerOne(server, router, entry); + } + } + + /** Useful for tests. */ + getEntries(): readonly RegisteredEntry[] { + return this.entries; + } + + private registerOne( + server: McpServer, + router: AnyRouter, + entry: RegisteredEntry, + ): void { + const handler = async ( + args: Record, + ): Promise => { + try { + const renamed = renameArgs(args, entry.renameInverse); + const caller = ( + router as unknown as { + createCaller: (ctx: unknown) => Record; + } + ).createCaller({}); + const procedureFn = resolveCallerPath(caller, entry.procedure.path); + const result = await procedureFn(renamed); + return { + content: [{ type: "text", text: safeStringify(result) }], + }; + } catch (err) { + return formatError(err); + } + }; + + const annotations = mapAnnotations(entry.config.annotations); + + if (entry.inputSchema) { + server.registerTool( + entry.toolName, + { + title: entry.config.title, + description: entry.config.description ?? entry.toolName, + inputSchema: entry.inputSchema.shape, + annotations, + }, + handler as never, + ); + } else { + server.registerTool( + entry.toolName, + { + title: entry.config.title, + description: entry.config.description ?? entry.toolName, + annotations, + }, + handler as never, + ); + } + } +} + +function collectProcedures(router: AnyRouter): ProcedureEntry[] { + const record = router._def.procedures as Record; + const out: ProcedureEntry[] = []; + for (const [path, procedure] of Object.entries(record)) { + const inputs = (procedure._def.inputs ?? []) as unknown[]; + const inputSchema = + inputs.length > 0 ? (inputs[0] as z.ZodTypeAny) : undefined; + out.push({ + path, + type: procedure._def.type, + inputSchema, + }); + } + return out; +} + +function composeInputSchema({ + toolName, + procedure, + config, +}: { + toolName: string; + procedure: ProcedureEntry; + config: ToolConfig; +}): { + inputSchema: z.ZodObject | null; + renameInverse: Record; +} { + const renameInverse: Record = {}; + if (config.rename_params) { + for (const [orig, alias] of Object.entries(config.rename_params)) { + renameInverse[alias] = orig; + } + } + + if (!procedure.inputSchema) { + if ( + config.param_overrides || + config.exclude_params || + config.rename_params + ) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" specifies param overrides/excludes/renames but procedure "${procedure.path}" has no input.`, + ); + } + return { inputSchema: null, renameInverse }; + } + + const baseSchema = procedure.inputSchema; + if (!(baseSchema instanceof z.ZodObject)) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" targets procedure "${procedure.path}" whose input is not a z.object — cannot expand into MCP params.`, + ); + } + let schema: z.ZodObject = + baseSchema as z.ZodObject; + + if (config.exclude_params) { + for (const key of config.exclude_params) { + if (!(key in schema.shape)) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" excludes param "${key}" not present on procedure "${procedure.path}".`, + ); + } + } + const omitShape: Record = {}; + for (const key of config.exclude_params) omitShape[key] = true; + schema = schema.omit(omitShape as never) as z.ZodObject; + } + + if (config.param_overrides) { + const nextShape: Record = { + ...(schema.shape as Record), + }; + for (const [key, override] of Object.entries(config.param_overrides)) { + if (!(key in nextShape)) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" overrides param "${key}" not present on procedure "${procedure.path}".`, + ); + } + if (override.description !== undefined) { + nextShape[key] = nextShape[key].describe(override.description); + } + } + schema = z.object(nextShape); + } + + if (config.rename_params) { + const nextShape: Record = {}; + for (const [origKey, field] of Object.entries(schema.shape)) { + const alias = config.rename_params[origKey] ?? origKey; + if (alias in nextShape) { + throw new Error( + `mcp-tools.yaml: tool "${toolName}" rename collision — "${alias}" appears twice in the renamed shape.`, + ); + } + nextShape[alias] = field as z.ZodTypeAny; + } + schema = z.object(nextShape); + } + + return { inputSchema: schema, renameInverse }; +} + +function resolveDescription( + config: ToolConfig, + yamlDir: string, +): string | undefined { + if (config.description) return config.description; + if (config.description_file) { + const filePath = path.resolve(yamlDir, config.description_file); + if (!fs.existsSync(filePath)) { + throw new Error( + `mcp-tools.yaml: description_file "${config.description_file}" not found (resolved to ${filePath}).`, + ); + } + return fs.readFileSync(filePath, "utf-8").trim(); + } + return undefined; +} + +function renameArgs( + args: Record | undefined, + inverse: Record, +): Record | undefined { + if (!args) return args; + if (Object.keys(inverse).length === 0) return args; + const out: Record = {}; + for (const [key, value] of Object.entries(args)) { + out[inverse[key] ?? key] = value; + } + return out; +} + +function resolveCallerPath( + caller: Record, + procedurePath: string, +): (input: unknown) => Promise { + const parts = procedurePath.split("."); + let current: unknown = caller; + for (const part of parts) { + if (current === undefined || current === null) { + throw new Error(`Cannot resolve "${procedurePath}" on caller`); + } + current = (current as Record)[part]; + } + if (typeof current !== "function") { + throw new Error(`"${procedurePath}" did not resolve to a procedure`); + } + return current as (input: unknown) => Promise; +} + +function safeStringify(value: unknown): string { + if (value === undefined) return ""; + if (typeof value === "string") return value; + try { + return JSON.stringify(value, null, 2); + } catch { + return String(value); + } +} + +function formatError(err: unknown): CallToolResult { + if (err instanceof TRPCError) { + return { + isError: true, + content: [{ type: "text", text: `[${err.code}] ${err.message}` }], + }; + } + if (err instanceof z.ZodError) { + return { + isError: true, + content: [ + { + type: "text", + text: `Invalid input:\n${err.issues + .map((i) => ` ${i.path.join(".")}: ${i.message}`) + .join("\n")}`, + }, + ], + }; + } + const message = err instanceof Error ? err.message : String(err); + return { + isError: true, + content: [{ type: "text", text: message }], + }; +} + +function mapAnnotations(annotations: NonNullable): { + readOnlyHint: boolean; + destructiveHint: boolean; + idempotentHint: boolean; +} { + return { + readOnlyHint: annotations.readOnly, + destructiveHint: annotations.destructive, + idempotentHint: annotations.idempotent, + }; +} diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/yaml-schema.ts b/apps/code/src/main/services/posthog-code-internal-mcp/yaml-schema.ts new file mode 100644 index 0000000000..6028dcfd45 --- /dev/null +++ b/apps/code/src/main/services/posthog-code-internal-mcp/yaml-schema.ts @@ -0,0 +1,45 @@ +import { z } from "zod"; + +const ParamOverride = z + .object({ + description: z.string().optional(), + }) + .strict(); + +const ToolAnnotations = z + .object({ + readOnly: z.boolean(), + destructive: z.boolean(), + idempotent: z.boolean(), + }) + .strict(); + +const ToolConfig = z + .object({ + operation: z.string().min(1), + enabled: z.boolean(), + title: z.string().optional(), + description: z.string().optional(), + description_file: z.string().optional(), + annotations: ToolAnnotations.optional(), + param_overrides: z.record(z.string(), ParamOverride).optional(), + exclude_params: z.array(z.string()).optional(), + rename_params: z.record(z.string(), z.string()).optional(), + }) + .strict() + .refine((d) => !(d.description && d.description_file), { + message: "description and description_file are mutually exclusive", + }) + .refine((d) => !d.enabled || d.annotations, { + message: "enabled tools must declare annotations", + }); + +export type ToolConfig = z.infer; + +export const McpToolsYamlSchema = z + .object({ + tools: z.record(z.string(), ToolConfig), + }) + .strict(); + +export type McpToolsYaml = z.infer; diff --git a/apps/code/src/main/trpc/router.ts b/apps/code/src/main/trpc/router.ts index f0f8dd9eb5..9beb165116 100644 --- a/apps/code/src/main/trpc/router.ts +++ b/apps/code/src/main/trpc/router.ts @@ -6,6 +6,7 @@ import { authRouter } from "./routers/auth"; import { cloudTaskRouter } from "./routers/cloud-task"; import { connectivityRouter } from "./routers/connectivity"; import { contextMenuRouter } from "./routers/context-menu"; +import { customInstructionsRouter } from "./routers/custom-instructions"; import { deepLinkRouter } from "./routers/deep-link"; import { encryptionRouter } from "./routers/encryption"; import { enrichmentRouter } from "./routers/enrichment"; @@ -23,6 +24,7 @@ import { llmGatewayRouter } from "./routers/llm-gateway"; import { logsRouter } from "./routers/logs"; import { mcpAppsRouter } from "./routers/mcp-apps"; import { mcpCallbackRouter } from "./routers/mcp-callback"; +import { mcpInstallationsRouter } from "./routers/mcp-installations"; import { notificationRouter } from "./routers/notification"; import { oauthRouter } from "./routers/oauth"; import { osRouter } from "./routers/os"; @@ -49,6 +51,7 @@ export const trpcRouter = router({ cloudTask: cloudTaskRouter, connectivity: connectivityRouter, contextMenu: contextMenuRouter, + customInstructions: customInstructionsRouter, enrichment: enrichmentRouter, environment: environmentRouter, @@ -65,6 +68,7 @@ export const trpcRouter = router({ llmGateway: llmGatewayRouter, mcpApps: mcpAppsRouter, mcpCallback: mcpCallbackRouter, + mcpInstallations: mcpInstallationsRouter, notification: notificationRouter, oauth: oauthRouter, logs: logsRouter, diff --git a/apps/code/src/main/trpc/routers/custom-instructions.ts b/apps/code/src/main/trpc/routers/custom-instructions.ts new file mode 100644 index 0000000000..310bc08808 --- /dev/null +++ b/apps/code/src/main/trpc/routers/custom-instructions.ts @@ -0,0 +1,39 @@ +import { container } from "../../di/container"; +import { MAIN_TOKENS } from "../../di/tokens"; +import { + CustomInstructionsServiceEvent, + readCustomInstructionsOutput, + writeCustomInstructionsInput, + writeCustomInstructionsOutput, +} from "../../services/custom-instructions/schemas"; +import type { CustomInstructionsService } from "../../services/custom-instructions/service"; +import { publicProcedure, router } from "../trpc"; + +const getService = () => + container.get( + MAIN_TOKENS.CustomInstructionsService, + ); + +export const customInstructionsRouter = router({ + read: publicProcedure + .output(readCustomInstructionsOutput) + .query(() => ({ customInstructions: getService().read() })), + + write: publicProcedure + .input(writeCustomInstructionsInput) + .output(writeCustomInstructionsOutput) + .mutation(({ input }) => { + getService().write(input.instructions); + return { ok: true as const }; + }), + + onChanged: publicProcedure.subscription(async function* (opts) { + const service = getService(); + for await (const data of service.toIterable( + CustomInstructionsServiceEvent.Changed, + { signal: opts.signal }, + )) { + yield data; + } + }), +}); diff --git a/apps/code/src/main/trpc/routers/mcp-installations.ts b/apps/code/src/main/trpc/routers/mcp-installations.ts new file mode 100644 index 0000000000..1901e74f5a --- /dev/null +++ b/apps/code/src/main/trpc/routers/mcp-installations.ts @@ -0,0 +1,34 @@ +import { container } from "../../di/container"; +import { MAIN_TOKENS } from "../../di/tokens"; +import { + installCustomInput, + installCustomOutput, + listMcpInstallationsOutput, + McpInstallationsServiceEvent, +} from "../../services/mcp-installations/schemas"; +import type { McpInstallationsService } from "../../services/mcp-installations/service"; +import { publicProcedure, router } from "../trpc"; + +const getService = () => + container.get(MAIN_TOKENS.McpInstallationsService); + +export const mcpInstallationsRouter = router({ + list: publicProcedure + .output(listMcpInstallationsOutput) + .query(() => getService().list()), + + installCustom: publicProcedure + .input(installCustomInput) + .output(installCustomOutput) + .mutation(({ input }) => getService().installCustom(input)), + + onInstalled: publicProcedure.subscription(async function* (opts) { + const service = getService(); + for await (const data of service.toIterable( + McpInstallationsServiceEvent.Installed, + { signal: opts.signal }, + )) { + yield data; + } + }), +}); diff --git a/apps/code/src/renderer/App.tsx b/apps/code/src/renderer/App.tsx index a5748db25b..30f50a5b8b 100644 --- a/apps/code/src/renderer/App.tsx +++ b/apps/code/src/renderer/App.tsx @@ -16,6 +16,7 @@ import { registerBillingSubscriptions } from "@features/billing/subscriptions"; import { AddDirectoryDialog } from "@features/folder-picker/components/AddDirectoryDialog"; import { OnboardingFlow } from "@features/onboarding/components/OnboardingFlow"; import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; +import { initializeSettingsStore } from "@features/settings/stores/settingsStore"; import { Flex, Spinner, Text } from "@radix-ui/themes"; import { initializeConnectivityToast } from "@renderer/features/connectivity/connectivityToast"; import { initializeConnectivityStore } from "@renderer/stores/connectivityStore"; @@ -80,6 +81,11 @@ function App() { return initializeUpdateStore(); }, []); + // Sync settings store when the internal MCP server writes custom instructions + useEffect(() => { + return initializeSettingsStore(); + }, []); + // Dev-only inbox demo command for local QA from the renderer console. useEffect(() => { if (import.meta.env.PROD) { diff --git a/apps/code/src/renderer/features/settings/stores/settingsStore.ts b/apps/code/src/renderer/features/settings/stores/settingsStore.ts index 69d626fcc1..09e88a0d4e 100644 --- a/apps/code/src/renderer/features/settings/stores/settingsStore.ts +++ b/apps/code/src/renderer/features/settings/stores/settingsStore.ts @@ -1,9 +1,13 @@ import type { WorkspaceMode } from "@main/services/workspace/schemas"; +import { trpcClient } from "@renderer/trpc/client"; import type { ExecutionMode } from "@shared/types"; import { electronStorage } from "@utils/electronStorage"; +import { logger } from "@utils/logger"; import { create } from "zustand"; import { persist } from "zustand/middleware"; +const log = logger.scope("settings-store"); + // ---------- Types ---------- export type DefaultRunMode = "local" | "cloud" | "last_used"; @@ -327,3 +331,23 @@ export const useSettingsStore = create()( }, ), ); + +/** + * Subscribe to custom-instructions writes coming from the agent's internal + * MCP server, so the in-memory store stays in sync after the persisted bucket + * is rewritten in the main process. + */ +export function initializeSettingsStore(): () => void { + const subscription = trpcClient.customInstructions.onChanged.subscribe( + undefined, + { + onData: ({ customInstructions }) => { + useSettingsStore.setState({ customInstructions }); + }, + onError: (error) => { + log.error("Custom instructions subscription error", { error }); + }, + }, + ); + return () => subscription.unsubscribe(); +} diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 05dc8bae59..bf5b02bc94 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,4 +1,8 @@ -export { isNotification, POSTHOG_NOTIFICATIONS } from "./acp-extensions"; +export { + isNotification, + POSTHOG_METHODS, + POSTHOG_NOTIFICATIONS, +} from "./acp-extensions"; export { getMcpToolMetadata, isMcpToolReadOnly, From 4f11893e64c18d99aff7500026e4b6e40bd05d1e Mon Sep 17 00:00:00 2001 From: Jovan Sakovic <49978945+sakce@users.noreply.github.com> Date: Thu, 28 May 2026 16:36:16 +0100 Subject: [PATCH 2/2] address review feedback: timing-safe auth, abortable oauth poll, AST-based scaffolder Greptile review on PR #2415: - Timing-safe bearer-token check via crypto.timingSafeEqual (constant-time buffer compare) instead of plain string equality. - pollForOauthCompletion now wires through an AbortController owned by McpInstallationsService and fired in @preDestroy, so in-flight polls stop when the app quits instead of running for up to 10 minutes. - Renamed the mock variable internalMcp -> mcpInstallations in agent/service.test.ts so it matches the actual injected service name. Also: rewrote scripts/scaffold-mcp-tools.ts to use the TypeScript compiler API for static AST parsing of router.ts + sub-routers, instead of dynamic import via a tsx loader. The previous approach hit ESM/CJS interop issues (node-machine-id, then named-export detection on .ts files loaded as CJS) that needed brittle per-module stubs. AST parsing has no runtime imports and runs in any Node version. The mcp-tools.yaml now lists all 256 live procedures (4 curated entries unchanged, 252 enabled: false stubs). Generated-By: PostHog Code Task-Id: 934a7fa1-fbb6-4f1d-89d3-f6abc69b7e23 --- apps/code/package.json | 2 +- apps/code/scripts/electron-stub.mjs | 78 -- .../scripts/scaffold-mcp-tools-loader.mjs | 55 -- .../scripts/scaffold-mcp-tools-preload.mjs | 10 - apps/code/scripts/scaffold-mcp-tools.ts | 215 +++-- .../src/main/services/agent/service.test.ts | 4 +- .../services/mcp-installations/service.ts | 45 +- .../posthog-code-internal-mcp/mcp-tools.yaml | 795 +++++++++++++++++- .../posthog-code-internal-mcp/service.ts | 20 +- 9 files changed, 998 insertions(+), 226 deletions(-) delete mode 100644 apps/code/scripts/electron-stub.mjs delete mode 100644 apps/code/scripts/scaffold-mcp-tools-loader.mjs delete mode 100644 apps/code/scripts/scaffold-mcp-tools-preload.mjs diff --git a/apps/code/package.json b/apps/code/package.json index fafe2eb4a7..29181f8bae 100644 --- a/apps/code/package.json +++ b/apps/code/package.json @@ -22,7 +22,7 @@ "build-icons": "bash scripts/generate-icns.sh", "typecheck": "tsc -p tsconfig.node.json --noEmit && tsc -p tsconfig.web.json --noEmit", "generate-client": "tsx scripts/update-openapi-client.ts", - "scaffold-mcp-tools": "tsx --import ./scripts/scaffold-mcp-tools-preload.mjs scripts/scaffold-mcp-tools.ts", + "scaffold-mcp-tools": "tsx scripts/scaffold-mcp-tools.ts", "test": "vitest run", "test:e2e": "playwright test --config=tests/e2e/playwright.config.ts", "test:e2e:headed": "playwright test --config=tests/e2e/playwright.config.ts --headed", diff --git a/apps/code/scripts/electron-stub.mjs b/apps/code/scripts/electron-stub.mjs deleted file mode 100644 index 32c3afa962..0000000000 --- a/apps/code/scripts/electron-stub.mjs +++ /dev/null @@ -1,78 +0,0 @@ -// Stub module returned in place of `electron` when the scaffold script runs -// outside of Electron. The router walk only needs the procedures' input -// schemas — none of the methods on `app`, `BrowserWindow`, etc. are actually -// called at import time, so a no-op object suffices. -// -// Used by scaffold-mcp-tools-preload.mjs via a Node loader hook. -const noop = () => {}; -const emptyObj = new Proxy( - {}, - { - get() { - return noop; - }, - }, -); - -const stub = new Proxy( - { - app: emptyObj, - BrowserWindow: () => emptyObj, - ipcMain: emptyObj, - ipcRenderer: emptyObj, - Menu: emptyObj, - MenuItem: emptyObj, - dialog: emptyObj, - shell: emptyObj, - nativeImage: emptyObj, - clipboard: emptyObj, - safeStorage: emptyObj, - powerMonitor: emptyObj, - powerSaveBlocker: emptyObj, - autoUpdater: emptyObj, - crashReporter: emptyObj, - Notification: () => emptyObj, - Tray: () => emptyObj, - nativeTheme: emptyObj, - session: emptyObj, - screen: emptyObj, - protocol: emptyObj, - webContents: emptyObj, - systemPreferences: emptyObj, - contextBridge: emptyObj, - }, - { - get(target, prop) { - if (prop in target) { - return target[prop]; - } - return noop; - }, - }, -); - -export default stub; -export const app = stub.app; -export const BrowserWindow = stub.BrowserWindow; -export const ipcMain = stub.ipcMain; -export const ipcRenderer = stub.ipcRenderer; -export const Menu = stub.Menu; -export const MenuItem = stub.MenuItem; -export const dialog = stub.dialog; -export const shell = stub.shell; -export const nativeImage = stub.nativeImage; -export const clipboard = stub.clipboard; -export const safeStorage = stub.safeStorage; -export const powerMonitor = stub.powerMonitor; -export const powerSaveBlocker = stub.powerSaveBlocker; -export const autoUpdater = stub.autoUpdater; -export const crashReporter = stub.crashReporter; -export const Notification = stub.Notification; -export const Tray = stub.Tray; -export const nativeTheme = stub.nativeTheme; -export const session = stub.session; -export const screen = stub.screen; -export const protocol = stub.protocol; -export const webContents = stub.webContents; -export const systemPreferences = stub.systemPreferences; -export const contextBridge = stub.contextBridge; diff --git a/apps/code/scripts/scaffold-mcp-tools-loader.mjs b/apps/code/scripts/scaffold-mcp-tools-loader.mjs deleted file mode 100644 index b642a1c32c..0000000000 --- a/apps/code/scripts/scaffold-mcp-tools-loader.mjs +++ /dev/null @@ -1,55 +0,0 @@ -import * as fs from "node:fs"; -import * as path from "node:path"; -import { pathToFileURL } from "node:url"; - -const STUB_URL = pathToFileURL(`${import.meta.dirname}/electron-stub.mjs`).href; -const SRC_DIR = path.resolve(import.meta.dirname, "..", "src"); - -const PATH_ALIASES = { - "@main/": `${path.join(SRC_DIR, "main")}/`, - "@renderer/": `${path.join(SRC_DIR, "renderer")}/`, - "@shared/": `${path.join(SRC_DIR, "shared")}/`, - "@features/": `${path.join(SRC_DIR, "renderer", "features")}/`, - "@components/": `${path.join(SRC_DIR, "renderer", "components")}/`, - "@stores/": `${path.join(SRC_DIR, "renderer", "stores")}/`, - "@hooks/": `${path.join(SRC_DIR, "renderer", "hooks")}/`, - "@utils/": `${path.join(SRC_DIR, "renderer", "utils")}/`, - "@test/": `${path.join(SRC_DIR, "shared", "test")}/`, -}; - -// Some import targets exist as BOTH `foo.ts` AND `foo/` (sibling file + -// directory). Node ESM's default resolution picks the directory and looks for -// `index.json` — wrong. `bundler` moduleResolution (which tsconfig sets) and -// Vite prefer the `.ts` sibling. Replicate that by checking for the `.ts` -// file first and short-circuiting if it exists. -function preferFileSibling(absPath) { - for (const ext of [".ts", ".tsx", ".mjs", ".js"]) { - const candidate = `${absPath}${ext}`; - if (fs.existsSync(candidate)) { - return candidate; - } - } - return null; -} - -export function resolve(specifier, context, nextResolve) { - if (specifier === "electron") { - return { url: STUB_URL, format: "module", shortCircuit: true }; - } - for (const [prefix, target] of Object.entries(PATH_ALIASES)) { - if (specifier.startsWith(prefix)) { - const rel = specifier.slice(prefix.length); - const abs = path.join(target, rel); - const fileSibling = preferFileSibling(abs); - if (fileSibling) { - return { - url: pathToFileURL(fileSibling).href, - format: fileSibling.endsWith(".json") ? "json" : "module", - shortCircuit: true, - }; - } - return nextResolve(abs, context); - } - } - return nextResolve(specifier, context); -} diff --git a/apps/code/scripts/scaffold-mcp-tools-preload.mjs b/apps/code/scripts/scaffold-mcp-tools-preload.mjs deleted file mode 100644 index c5fb460899..0000000000 --- a/apps/code/scripts/scaffold-mcp-tools-preload.mjs +++ /dev/null @@ -1,10 +0,0 @@ -import { register } from "node:module"; -import { pathToFileURL } from "node:url"; - -// Redirects `electron` imports to a local stub so the scaffold script can -// import the tRPC router (which transitively touches Electron-bound modules) -// without an Electron runtime present. -register( - "./scaffold-mcp-tools-loader.mjs", - pathToFileURL(`${import.meta.dirname}/`), -); diff --git a/apps/code/scripts/scaffold-mcp-tools.ts b/apps/code/scripts/scaffold-mcp-tools.ts index 44fefb4dd0..218d2f0fa9 100644 --- a/apps/code/scripts/scaffold-mcp-tools.ts +++ b/apps/code/scripts/scaffold-mcp-tools.ts @@ -3,13 +3,11 @@ * Sync `apps/code/src/main/services/posthog-code-internal-mcp/mcp-tools.yaml` * with the live tRPC router. * - * - Walks the router via `_def.procedures` and emits an `enabled: false` stub - * for every procedure that isn't already in the YAML. - * - Leaves existing entries untouched — your hand-authored config (title, - * description, annotations, param_overrides) is preserved. - * - Does NOT remove entries whose procedure has disappeared. It prints them - * as warnings; you decide whether to delete. Boot will hard-fail until you - * do, which is the forcing function. + * Uses the TypeScript compiler API to STATICALLY parse the router and + * sub-router source files — no runtime import, no Electron, no Node ESM/CJS + * interop. Walks `router({ ... })` object literals, detects whether each + * procedure chain ends in `.query`, `.mutation`, or `.subscription`, and + * emits `enabled: false` stubs for procedures missing from the YAML. * * Usage: * pnpm --filter code scaffold-mcp-tools @@ -17,32 +15,17 @@ */ import * as fs from "node:fs"; -import * as os from "node:os"; import * as path from "node:path"; +import * as ts from "typescript"; +import { parse as parseYaml, stringify as stringifyYaml } from "yaml"; +import { McpToolsYamlSchema } from "../src/main/services/posthog-code-internal-mcp/yaml-schema"; -// Imports below transitively load main-process services that read env vars -// at module-load time (see apps/code/src/main/utils/env.ts). Set defaults -// here so the script works outside an Electron context. Done before any -// dynamic import below. -if (!process.env.POSTHOG_CODE_DATA_DIR) { - process.env.POSTHOG_CODE_DATA_DIR = path.join( - os.tmpdir(), - "posthog-code-scaffold-mcp-tools", - ); -} -if (!process.env.POSTHOG_CODE_IS_DEV) process.env.POSTHOG_CODE_IS_DEV = "true"; -if (!process.env.POSTHOG_CODE_VERSION) { - process.env.POSTHOG_CODE_VERSION = "0.0.0-scaffold"; -} - -const YAML_PATH = path.resolve( - __dirname, - "..", - "src", - "main", - "services", - "posthog-code-internal-mcp", - "mcp-tools.yaml", +const APP_ROOT = path.resolve(__dirname, ".."); +const ROUTER_FILE = path.join(APP_ROOT, "src/main/trpc/router.ts"); +const ROUTERS_DIR = path.join(APP_ROOT, "src/main/trpc/routers"); +const YAML_PATH = path.join( + APP_ROOT, + "src/main/services/posthog-code-internal-mcp/mcp-tools.yaml", ); const YAML_HEADER = `# Bridge from tRPC procedures to MCP tools exposed to the running agent. @@ -56,29 +39,160 @@ const YAML_HEADER = `# Bridge from tRPC procedures to MCP tools exposed to the r # before flipping enabled: true on anything beyond the curated defaults. `; +type ProcedureType = "query" | "mutation" | "subscription"; + interface Procedure { path: string; - type: "query" | "mutation" | "subscription"; + type: ProcedureType; } -async function main(): Promise { - const check = process.argv.includes("--check"); +function parseSource(filePath: string): ts.SourceFile { + const text = fs.readFileSync(filePath, "utf-8"); + return ts.createSourceFile(filePath, text, ts.ScriptTarget.Latest, true); +} - // Dynamic import so the env defaults above are in place before the router's - // transitive deps load. - const { trpcRouter } = await import("../src/main/trpc/router"); - const { McpToolsYamlSchema } = await import( - "../src/main/services/posthog-code-internal-mcp/yaml-schema" - ); - const { parse: parseYaml, stringify: stringifyYaml } = await import("yaml"); +/** + * Find the object literal passed to a `router({...})` call inside the + * given source file. We look for the FIRST top-level `router(...)` call + * expression; that's the canonical pattern in this codebase (one router + * declaration per file). + */ +function findRouterObjectLiteral( + source: ts.SourceFile, +): ts.ObjectLiteralExpression | undefined { + let result: ts.ObjectLiteralExpression | undefined; + + function visit(node: ts.Node): void { + if (result) return; + if ( + ts.isCallExpression(node) && + ts.isIdentifier(node.expression) && + node.expression.text === "router" && + node.arguments.length === 1 && + ts.isObjectLiteralExpression(node.arguments[0]) + ) { + result = node.arguments[0]; + return; + } + ts.forEachChild(node, visit); + } + + visit(source); + return result; +} + +/** + * Build a map of namespace → router-file-basename by parsing the root + * router.ts. The pattern is: + * + * import { fooRouter } from "./routers/foo"; + * export const trpcRouter = router({ foo: fooRouter, ... }); + * + * We use the property assignments to map namespace → identifier, then the + * import declarations to map identifier → file path. + */ +function discoverSubRouters(): Map { + const source = parseSource(ROUTER_FILE); + + const importMap = new Map(); + for (const stmt of source.statements) { + if (!ts.isImportDeclaration(stmt)) continue; + if (!ts.isStringLiteral(stmt.moduleSpecifier)) continue; + const moduleSpec = stmt.moduleSpecifier.text; + if (!moduleSpec.startsWith("./routers/")) continue; + const basename = moduleSpec.slice("./routers/".length).replace(/\.js$/, ""); + const fileAbs = path.join(ROUTERS_DIR, `${basename}.ts`); + if (!stmt.importClause?.namedBindings) continue; + if (!ts.isNamedImports(stmt.importClause.namedBindings)) continue; + for (const spec of stmt.importClause.namedBindings.elements) { + importMap.set(spec.name.text, fileAbs); + } + } - const record = trpcRouter._def.procedures as Record< - string, - { _def: { type: "query" | "mutation" | "subscription" } } - >; - const procedures: Procedure[] = Object.entries(record) - .map(([p, proc]) => ({ path: p, type: proc._def.type })) - .filter((p) => p.type !== "subscription"); + const literal = findRouterObjectLiteral(source); + if (!literal) { + throw new Error(`Could not find router({...}) call in ${ROUTER_FILE}`); + } + + const namespaces = new Map(); + for (const prop of literal.properties) { + if (!ts.isPropertyAssignment(prop)) continue; + const key = propertyName(prop.name); + if (!key) continue; + if (!ts.isIdentifier(prop.initializer)) continue; + const filePath = importMap.get(prop.initializer.text); + if (!filePath) { + throw new Error( + `Router namespace '${key}' maps to identifier '${prop.initializer.text}' which has no matching import in router.ts`, + ); + } + namespaces.set(key, filePath); + } + + return namespaces; +} + +function propertyName(name: ts.PropertyName): string | undefined { + if (ts.isIdentifier(name) || ts.isStringLiteral(name)) return name.text; + return undefined; +} + +/** + * Walk the procedure chain's outermost call to determine its type. The chain + * looks like: + * publicProcedure.input(X).output(Y).query(handler) + * publicProcedure.subscription(handler) + * The OUTERMOST call's property name tells us the type. Anything that + * doesn't end in query/mutation/subscription is skipped (e.g. helpers). + */ +function classifyProcedure(expr: ts.Expression): ProcedureType | undefined { + if (!ts.isCallExpression(expr)) return undefined; + const callee = expr.expression; + if (!ts.isPropertyAccessExpression(callee)) return undefined; + const method = callee.name.text; + if ( + method === "query" || + method === "mutation" || + method === "subscription" + ) { + return method; + } + return undefined; +} + +function parseRouterProcedures( + namespace: string, + filePath: string, +): Procedure[] { + const source = parseSource(filePath); + const literal = findRouterObjectLiteral(source); + if (!literal) { + throw new Error(`Could not find router({...}) call in ${filePath}`); + } + const procedures: Procedure[] = []; + for (const prop of literal.properties) { + if (!ts.isPropertyAssignment(prop)) continue; + const key = propertyName(prop.name); + if (!key) continue; + const type = classifyProcedure(prop.initializer); + if (!type) continue; + procedures.push({ path: `${namespace}.${key}`, type }); + } + return procedures; +} + +function discoverProcedures(): Procedure[] { + const namespaces = discoverSubRouters(); + const all: Procedure[] = []; + for (const [namespace, filePath] of namespaces) { + all.push(...parseRouterProcedures(namespace, filePath)); + } + return all.filter((p) => p.type !== "subscription"); +} + +function main(): void { + const check = process.argv.includes("--check"); + const procedures = discoverProcedures(); let existing: { tools: Record } = { tools: {}, @@ -177,7 +291,4 @@ async function main(): Promise { } } -void main().catch((err) => { - console.error(err); - process.exit(1); -}); +main(); diff --git a/apps/code/src/main/services/agent/service.test.ts b/apps/code/src/main/services/agent/service.test.ts index 2419ba71bd..842c610dbd 100644 --- a/apps/code/src/main/services/agent/service.test.ts +++ b/apps/code/src/main/services/agent/service.test.ts @@ -200,7 +200,7 @@ function createMockDependencies() { addAdditionalDirectory: vi.fn(), removeAdditionalDirectory: vi.fn(), }, - internalMcp: { + mcpInstallations: { on: vi.fn(), }, }; @@ -235,7 +235,7 @@ describe("AgentService", () => { deps.storagePaths as never, deps.defaultAdditionalDirectoryRepository as never, deps.workspaceRepository as never, - deps.internalMcp as never, + deps.mcpInstallations as never, ); }); diff --git a/apps/code/src/main/services/mcp-installations/service.ts b/apps/code/src/main/services/mcp-installations/service.ts index 0f285bfb9b..2146019881 100644 --- a/apps/code/src/main/services/mcp-installations/service.ts +++ b/apps/code/src/main/services/mcp-installations/service.ts @@ -1,4 +1,4 @@ -import { inject, injectable } from "inversify"; +import { inject, injectable, preDestroy } from "inversify"; import { MAIN_TOKENS } from "../../di/tokens"; import { logger } from "../../utils/logger"; import { TypedEventEmitter } from "../../utils/typed-event-emitter"; @@ -16,6 +16,24 @@ const log = logger.scope("mcp-installations"); const OAUTH_POLL_INTERVAL_MS = 3000; const OAUTH_POLL_TIMEOUT_MS = 10 * 60 * 1000; +function wait(ms: number, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal.aborted) { + reject(new Error("aborted")); + return; + } + const timer = setTimeout(() => { + signal.removeEventListener("abort", onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(new Error("aborted")); + }; + signal.addEventListener("abort", onAbort, { once: true }); + }); +} + /** * Owns the PostHog REST surface for MCP server installations on the active * project. Reads/writes against `/api/environments/{projectId}/mcp_server_installations/`. @@ -26,6 +44,10 @@ const OAUTH_POLL_TIMEOUT_MS = 10 * 60 * 1000; */ @injectable() export class McpInstallationsService extends TypedEventEmitter { + // Aborts any in-flight OAuth polls when the app shuts down so we don't + // keep authenticated HTTP requests running for up to 10 minutes after quit. + private readonly shutdown = new AbortController(); + constructor( @inject(MAIN_TOKENS.AuthService) private readonly authService: AuthService, @@ -33,6 +55,11 @@ export class McpInstallationsService extends TypedEventEmitter { const { baseUrl, projectId } = await this.requireProject(); const url = `${baseUrl}/api/environments/${projectId}/mcp_server_installations/`; @@ -132,6 +159,8 @@ export class McpInstallationsService extends TypedEventEmitter { + const signal = this.shutdown.signal; + if (signal.aborted) return; const { apiHost } = await this.authService.getValidAccessToken(); const projectId = this.authService.getState().projectId; if (!projectId) return; @@ -142,13 +171,17 @@ export class McpInstallationsService extends TypedEventEmitter - setTimeout(resolve, OAUTH_POLL_INTERVAL_MS), - ); + try { + await wait(OAUTH_POLL_INTERVAL_MS, signal); + } catch { + log.info("OAuth poll aborted (shutdown)", { installationId, name }); + return; + } try { const response = await this.authService.authenticatedFetch(fetch, url, { headers: { "Content-Type": "application/json" }, + signal, }); if (!response.ok) continue; const data = (await response.json()) as { @@ -181,6 +214,10 @@ export class McpInstallationsService extends TypedEventEmitter - Return the user's custom instructions — extra guidance the user has appended to every agent prompt. - Returns an empty string if no custom instructions are configured. + Return the user's custom instructions — extra guidance the user has appended to every agent prompt. Returns an + empty string if no custom instructions are configured. annotations: readOnly: true destructive: false idempotent: true - customInstructions.write: operation: customInstructions.write enabled: true title: Write custom instructions description: > - Replace the user's custom instructions. Pass the full new text — this overwrites the existing value. - Pass an empty string to clear. + Replace the user's custom instructions. Pass the full new text — this overwrites the existing value. Pass an empty + string to clear. annotations: readOnly: false destructive: true idempotent: false - - mcpInstallations.list: - operation: mcpInstallations.list - enabled: true - title: List installed MCP servers - description: > - List the MCP server installations on the current PostHog project. Returns id, name, url, auth_type, - and status flags for each installation. - annotations: - readOnly: true - destructive: false - idempotent: true - + deepLink.getPendingDeepLink: + operation: deepLink.getPendingDeepLink + enabled: false + deepLink.getPendingNewTaskLink: + operation: deepLink.getPendingNewTaskLink + enabled: false + deepLink.getPendingReportLink: + operation: deepLink.getPendingReportLink + enabled: false + encryption.decrypt: + operation: encryption.decrypt + enabled: false + encryption.encrypt: + operation: encryption.encrypt + enabled: false + enrichment.detectPosthogInstallState: + operation: enrichment.detectPosthogInstallState + enabled: false + enrichment.enrichFile: + operation: enrichment.enrichFile + enabled: false + enrichment.findStaleFlagSuggestions: + operation: enrichment.findStaleFlagSuggestions + enabled: false + environment.create: + operation: environment.create + enabled: false + environment.delete: + operation: environment.delete + enabled: false + environment.get: + operation: environment.get + enabled: false + environment.list: + operation: environment.list + enabled: false + environment.update: + operation: environment.update + enabled: false + externalApps.copyPath: + operation: externalApps.copyPath + enabled: false + externalApps.getDetectedApps: + operation: externalApps.getDetectedApps + enabled: false + externalApps.getLastUsed: + operation: externalApps.getLastUsed + enabled: false + externalApps.openInApp: + operation: externalApps.openInApp + enabled: false + externalApps.setLastUsed: + operation: externalApps.setLastUsed + enabled: false + fileWatcher.listDirectory: + operation: fileWatcher.listDirectory + enabled: false + fileWatcher.start: + operation: fileWatcher.start + enabled: false + fileWatcher.stop: + operation: fileWatcher.stop + enabled: false + focus.checkout: + operation: focus.checkout + enabled: false + focus.cleanWorkingTree: + operation: focus.cleanWorkingTree + enabled: false + focus.deleteSession: + operation: focus.deleteSession + enabled: false + focus.detachWorktree: + operation: focus.detachWorktree + enabled: false + focus.findWorktreeByBranch: + operation: focus.findWorktreeByBranch + enabled: false + focus.getCommitSha: + operation: focus.getCommitSha + enabled: false + focus.getSession: + operation: focus.getSession + enabled: false + focus.isDirty: + operation: focus.isDirty + enabled: false + focus.isFocusActive: + operation: focus.isFocusActive + enabled: false + focus.reattachWorktree: + operation: focus.reattachWorktree + enabled: false + focus.saveSession: + operation: focus.saveSession + enabled: false + focus.startSync: + operation: focus.startSync + enabled: false + focus.startWatchingMainRepo: + operation: focus.startWatchingMainRepo + enabled: false + focus.stash: + operation: focus.stash + enabled: false + focus.stashApply: + operation: focus.stashApply + enabled: false + focus.stashPop: + operation: focus.stashPop + enabled: false + focus.stopSync: + operation: focus.stopSync + enabled: false + focus.stopWatchingMainRepo: + operation: focus.stopWatchingMainRepo + enabled: false + focus.toAbsoluteWorktreePath: + operation: focus.toAbsoluteWorktreePath + enabled: false + focus.toRelativeWorktreePath: + operation: focus.toRelativeWorktreePath + enabled: false + focus.validateFocusOperation: + operation: focus.validateFocusOperation + enabled: false + focus.worktreeExistsAtPath: + operation: focus.worktreeExistsAtPath + enabled: false + folders.addFolder: + operation: folders.addFolder + enabled: false + folders.clearAllData: + operation: folders.clearAllData + enabled: false + folders.getFolders: + operation: folders.getFolders + enabled: false + folders.getMostRecentlyAccessedRepository: + operation: folders.getMostRecentlyAccessedRepository + enabled: false + folders.getRepositoryByRemoteUrl: + operation: folders.getRepositoryByRemoteUrl + enabled: false + folders.removeFolder: + operation: folders.removeFolder + enabled: false + folders.updateFolderAccessed: + operation: folders.updateFolderAccessed + enabled: false + fs.listRepoFiles: + operation: fs.listRepoFiles + enabled: false + fs.readAbsoluteFile: + operation: fs.readAbsoluteFile + enabled: false + fs.readFileAsBase64: + operation: fs.readFileAsBase64 + enabled: false + fs.readRepoFile: + operation: fs.readRepoFile + enabled: false + fs.readRepoFileBounded: + operation: fs.readRepoFileBounded + enabled: false + fs.readRepoFiles: + operation: fs.readRepoFiles + enabled: false + fs.readRepoFilesBounded: + operation: fs.readRepoFilesBounded + enabled: false + fs.writeRepoFile: + operation: fs.writeRepoFile + enabled: false + git.checkoutBranch: + operation: git.checkoutBranch + enabled: false + git.cloneRepository: + operation: git.cloneRepository + enabled: false + git.commit: + operation: git.commit + enabled: false + git.createBranch: + operation: git.createBranch + enabled: false + git.createPr: + operation: git.createPr + enabled: false + git.detectRepo: + operation: git.detectRepo + enabled: false + git.discardFileChanges: + operation: git.discardFileChanges + enabled: false + git.generateCommitMessage: + operation: git.generateCommitMessage + enabled: false + git.generatePrTitleAndBody: + operation: git.generatePrTitleAndBody + enabled: false + git.getAllBranches: + operation: git.getAllBranches + enabled: false + git.getBranchChangedFiles: + operation: git.getBranchChangedFiles + enabled: false + git.getChangedFilesHead: + operation: git.getChangedFilesHead + enabled: false + git.getCommitConventions: + operation: git.getCommitConventions + enabled: false + git.getCurrentBranch: + operation: git.getCurrentBranch + enabled: false + git.getDiffCached: + operation: git.getDiffCached + enabled: false + git.getDiffHead: + operation: git.getDiffHead + enabled: false + git.getDiffStats: + operation: git.getDiffStats + enabled: false + git.getDiffUnstaged: + operation: git.getDiffUnstaged + enabled: false + git.getFileAtHead: + operation: git.getFileAtHead + enabled: false + git.getGhAuthToken: + operation: git.getGhAuthToken + enabled: false + git.getGhStatus: + operation: git.getGhStatus + enabled: false + git.getGitBusyState: + operation: git.getGitBusyState + enabled: false + git.getGithubIssue: + operation: git.getGithubIssue + enabled: false + git.getGithubPullRequest: + operation: git.getGithubPullRequest + enabled: false + git.getGitRepoInfo: + operation: git.getGitRepoInfo + enabled: false + git.getGitStatus: + operation: git.getGitStatus + enabled: false + git.getGitSyncStatus: + operation: git.getGitSyncStatus + enabled: false + git.getLatestCommit: + operation: git.getLatestCommit + enabled: false + git.getLocalBranchChangedFiles: + operation: git.getLocalBranchChangedFiles + enabled: false + git.getPrChangedFiles: + operation: git.getPrChangedFiles + enabled: false + git.getPrDetailsByUrl: + operation: git.getPrDetailsByUrl + enabled: false + git.getPrReviewComments: + operation: git.getPrReviewComments + enabled: false + git.getPrStatus: + operation: git.getPrStatus + enabled: false + git.getPrTemplate: + operation: git.getPrTemplate + enabled: false + git.getPrUrlForBranch: + operation: git.getPrUrlForBranch + enabled: false + git.openPr: + operation: git.openPr + enabled: false + git.publish: + operation: git.publish + enabled: false + git.pull: + operation: git.pull + enabled: false + git.push: + operation: git.push + enabled: false + git.replyToPrComment: + operation: git.replyToPrComment + enabled: false + git.resolveReviewThread: + operation: git.resolveReviewThread + enabled: false + git.searchGithubRefs: + operation: git.searchGithubRefs + enabled: false + git.stageFiles: + operation: git.stageFiles + enabled: false + git.sync: + operation: git.sync + enabled: false + git.unstageFiles: + operation: git.unstageFiles + enabled: false + git.updatePrByUrl: + operation: git.updatePrByUrl + enabled: false + git.validateRepo: + operation: git.validateRepo + enabled: false + githubIntegration.consumePendingCallback: + operation: githubIntegration.consumePendingCallback + enabled: false + githubIntegration.startFlow: + operation: githubIntegration.startFlow + enabled: false + handoff.execute: + operation: handoff.execute + enabled: false + handoff.executeToCloud: + operation: handoff.executeToCloud + enabled: false + handoff.preflight: + operation: handoff.preflight + enabled: false + handoff.preflightToCloud: + operation: handoff.preflightToCloud + enabled: false + linearIntegration.startFlow: + operation: linearIntegration.startFlow + enabled: false + llmGateway.invalidatePlanCache: + operation: llmGateway.invalidatePlanCache + enabled: false + llmGateway.prompt: + operation: llmGateway.prompt + enabled: false + logs.fetchS3Logs: + operation: logs.fetchS3Logs + enabled: false + logs.readLocalLogs: + operation: logs.readLocalLogs + enabled: false + logs.writeLocalLogs: + operation: logs.writeLocalLogs + enabled: false + mcpApps.getToolDefinition: + operation: mcpApps.getToolDefinition + enabled: false + mcpApps.getUiResource: + operation: mcpApps.getUiResource + enabled: false + mcpApps.hasUiForTool: + operation: mcpApps.hasUiForTool + enabled: false + mcpApps.openLink: + operation: mcpApps.openLink + enabled: false + mcpApps.proxyResourceRead: + operation: mcpApps.proxyResourceRead + enabled: false + mcpApps.proxyToolCall: + operation: mcpApps.proxyToolCall + enabled: false + mcpCallback.getCallbackUrl: + operation: mcpCallback.getCallbackUrl + enabled: false + mcpCallback.openAndWaitForCallback: + operation: mcpCallback.openAndWaitForCallback + enabled: false mcpInstallations.installCustom: operation: mcpInstallations.installCustom enabled: true title: Install a custom MCP server description: > - Install a new MCP server on the current PostHog project. Use auth_type="api_key" with `api_key` - for static bearer tokens; use auth_type="oauth" to start an OAuth handshake — the response will - include a redirectUrl the user must visit. + Install a new MCP server on the current PostHog project. Use auth_type="api_key" with `api_key` for static bearer + tokens; use auth_type="oauth" to start an OAuth handshake — the response will include a redirectUrl the user must + visit. annotations: readOnly: false destructive: false idempotent: false + mcpInstallations.list: + operation: mcpInstallations.list + enabled: true + title: List installed MCP servers + description: > + List the MCP server installations on the current PostHog project. Returns id, name, url, auth_type, and status + flags for each installation. + annotations: + readOnly: true + destructive: false + idempotent: true + notification.bounceDock: + operation: notification.bounceDock + enabled: false + notification.send: + operation: notification.send + enabled: false + notification.showDockBadge: + operation: notification.showDockBadge + enabled: false + oauth.cancelFlow: + operation: oauth.cancelFlow + enabled: false + os.checkWriteAccess: + operation: os.checkWriteAccess + enabled: false + os.downscaleImageFile: + operation: os.downscaleImageFile + enabled: false + os.getAppVersion: + operation: os.getAppVersion + enabled: false + os.getClaudePermissions: + operation: os.getClaudePermissions + enabled: false + os.getWorktreeLocation: + operation: os.getWorktreeLocation + enabled: false + os.openExternal: + operation: os.openExternal + enabled: false + os.readFileAsDataUrl: + operation: os.readFileAsDataUrl + enabled: false + os.saveClipboardFile: + operation: os.saveClipboardFile + enabled: false + os.saveClipboardImage: + operation: os.saveClipboardImage + enabled: false + os.saveClipboardText: + operation: os.saveClipboardText + enabled: false + os.searchDirectories: + operation: os.searchDirectories + enabled: false + os.selectAttachments: + operation: os.selectAttachments + enabled: false + os.selectDirectory: + operation: os.selectDirectory + enabled: false + os.selectFiles: + operation: os.selectFiles + enabled: false + os.showMessageBox: + operation: os.showMessageBox + enabled: false + processTracking.getSnapshot: + operation: processTracking.getSnapshot + enabled: false + processTracking.kill: + operation: processTracking.kill + enabled: false + processTracking.killAll: + operation: processTracking.killAll + enabled: false + processTracking.killByCategory: + operation: processTracking.killByCategory + enabled: false + processTracking.killByTaskId: + operation: processTracking.killByTaskId + enabled: false + processTracking.list: + operation: processTracking.list + enabled: false + processTracking.listByTaskId: + operation: processTracking.listByTaskId + enabled: false + secureStore.clear: + operation: secureStore.clear + enabled: false + secureStore.getItem: + operation: secureStore.getItem + enabled: false + secureStore.removeItem: + operation: secureStore.removeItem + enabled: false + secureStore.setItem: + operation: secureStore.setItem + enabled: false + shell.check: + operation: shell.check + enabled: false + shell.create: + operation: shell.create + enabled: false + shell.createCommand: + operation: shell.createCommand + enabled: false + shell.destroy: + operation: shell.destroy + enabled: false + shell.execute: + operation: shell.execute + enabled: false + shell.getProcess: + operation: shell.getProcess + enabled: false + shell.resize: + operation: shell.resize + enabled: false + shell.write: + operation: shell.write + enabled: false + skills.list: + operation: skills.list + enabled: false + slackIntegration.consumePendingCallback: + operation: slackIntegration.consumePendingCallback + enabled: false + slackIntegration.startFlow: + operation: slackIntegration.startFlow + enabled: false + sleep.getEnabled: + operation: sleep.getEnabled + enabled: false + sleep.setEnabled: + operation: sleep.setEnabled + enabled: false + suspension.list: + operation: suspension.list + enabled: false + suspension.restore: + operation: suspension.restore + enabled: false + suspension.settings: + operation: suspension.settings + enabled: false + suspension.suspend: + operation: suspension.suspend + enabled: false + suspension.suspendedTaskIds: + operation: suspension.suspendedTaskIds + enabled: false + suspension.updateSettings: + operation: suspension.updateSettings + enabled: false + updates.check: + operation: updates.check + enabled: false + updates.install: + operation: updates.install + enabled: false + updates.isEnabled: + operation: updates.isEnabled + enabled: false + usageMonitor.getLatest: + operation: usageMonitor.getLatest + enabled: false + usageMonitor.refresh: + operation: usageMonitor.refresh + enabled: false + workspace.create: + operation: workspace.create + enabled: false + workspace.delete: + operation: workspace.delete + enabled: false + workspace.deleteWorktree: + operation: workspace.deleteWorktree + enabled: false + workspace.getAll: + operation: workspace.getAll + enabled: false + workspace.getAllTaskTimestamps: + operation: workspace.getAllTaskTimestamps + enabled: false + workspace.getInfo: + operation: workspace.getInfo + enabled: false + workspace.getLocalTasks: + operation: workspace.getLocalTasks + enabled: false + workspace.getPinnedTaskIds: + operation: workspace.getPinnedTaskIds + enabled: false + workspace.getTaskPrStatus: + operation: workspace.getTaskPrStatus + enabled: false + workspace.getTaskTimestamps: + operation: workspace.getTaskTimestamps + enabled: false + workspace.getWorktreeFileUsage: + operation: workspace.getWorktreeFileUsage + enabled: false + workspace.getWorktreeSize: + operation: workspace.getWorktreeSize + enabled: false + workspace.getWorktreeTasks: + operation: workspace.getWorktreeTasks + enabled: false + workspace.linkBranch: + operation: workspace.linkBranch + enabled: false + workspace.listGitWorktrees: + operation: workspace.listGitWorktrees + enabled: false + workspace.markActivity: + operation: workspace.markActivity + enabled: false + workspace.markViewed: + operation: workspace.markViewed + enabled: false + workspace.reconcileCloudWorkspaces: + operation: workspace.reconcileCloudWorkspaces + enabled: false + workspace.togglePin: + operation: workspace.togglePin + enabled: false + workspace.unlinkBranch: + operation: workspace.unlinkBranch + enabled: false + workspace.verify: + operation: workspace.verify + enabled: false diff --git a/apps/code/src/main/services/posthog-code-internal-mcp/service.ts b/apps/code/src/main/services/posthog-code-internal-mcp/service.ts index da76ad3911..00c2d5eb8d 100644 --- a/apps/code/src/main/services/posthog-code-internal-mcp/service.ts +++ b/apps/code/src/main/services/posthog-code-internal-mcp/service.ts @@ -1,4 +1,4 @@ -import { randomBytes } from "node:crypto"; +import { randomBytes, timingSafeEqual } from "node:crypto"; import http from "node:http"; import * as path from "node:path"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; @@ -110,8 +110,7 @@ export class PostHogCodeInternalMcpService { req: http.IncomingMessage, res: http.ServerResponse, ): Promise { - const auth = req.headers.authorization; - if (!auth || auth !== `Bearer ${this.bearerToken}`) { + if (!this.isAuthorized(req)) { res.writeHead(401).end("Unauthorized"); return; } @@ -153,6 +152,21 @@ export class PostHogCodeInternalMcpService { } } + /** + * Constant-time bearer-token check. The token is per-boot random + the + * server only binds to 127.0.0.1, but credential checks should still avoid + * leaking length/equality timing. + */ + private isAuthorized(req: http.IncomingMessage): boolean { + if (!this.bearerToken) return false; + const header = req.headers.authorization; + if (!header) return false; + const expected = Buffer.from(`Bearer ${this.bearerToken}`, "utf8"); + const provided = Buffer.from(header, "utf8"); + if (provided.length !== expected.length) return false; + return timingSafeEqual(provided, expected); + } + private buildServer(): McpServer { if (!this.registry) { throw new Error("posthog-code-internal MCP registry not initialized");