diff --git a/apps/code/electron-builder.ts b/apps/code/electron-builder.ts index d124f10bba..d0fd293eb3 100644 --- a/apps/code/electron-builder.ts +++ b/apps/code/electron-builder.ts @@ -49,6 +49,8 @@ const config: Configuration = { ".vite/build/plugins/posthog/**", ".vite/build/codex-acp/**", ".vite/build/grammars/**", + ".vite/build/rpc-host.js", + ".vite/build/rpc-host.js.map", ...asarUnpackGlobs, ], diff --git a/apps/code/electron.vite.config.ts b/apps/code/electron.vite.config.ts index 2e8172dd17..16bce4a5c2 100644 --- a/apps/code/electron.vite.config.ts +++ b/apps/code/electron.vite.config.ts @@ -22,6 +22,7 @@ import { copyCodexAcpBinaries, copyDrizzleMigrations, copyEnricherGrammars, + copyPiRpcHost, copyPosthogPlugin, fixFilenameCircularRef, getBuildDate, @@ -93,6 +94,7 @@ export default defineConfig(({ mode }) => { autoServicesPlugin(path.join(__dirname, "src/main/services")), fixFilenameCircularRef(), copyClaudeExecutable(), + copyPiRpcHost(), copyPosthogPlugin(isDev), copyDrizzleMigrations(), copyCodexAcpBinaries(), diff --git a/apps/code/package.json b/apps/code/package.json index a0a50e5b15..37961dcd6e 100644 --- a/apps/code/package.json +++ b/apps/code/package.json @@ -152,6 +152,7 @@ "reflect-metadata": "^0.2.2", "semver": "^7.8.1", "shadcn": "^4.1.2", + "superjson": "catalog:", "smol-toml": "^1.6.1", "tailwindcss-scroll-mask": "^0.0.3", "tw-animate-css": "^1.4.0", diff --git a/apps/code/src/main/di/bindings.ts b/apps/code/src/main/di/bindings.ts index 9d59c1d7b9..7bd793ee9a 100644 --- a/apps/code/src/main/di/bindings.ts +++ b/apps/code/src/main/di/bindings.ts @@ -186,6 +186,10 @@ import type { MCP_RELAY_SERVICE, McpRelayService, } from "@posthog/workspace-server/services/mcp-relay/identifiers"; +import type { + PI_RPC_CLIENT_FACTORY, + PiRpcClientFactory, +} from "@posthog/workspace-server/services/pi-session/identifiers"; import type { PosthogPluginService } from "@posthog/workspace-server/services/posthog-plugin/posthog-plugin"; import type { ProcessTrackingService } from "@posthog/workspace-server/services/process-tracking/process-tracking"; import type { @@ -350,6 +354,7 @@ export interface MainBindings { [AGENT_REPO_FILES]: unknown; [AGENT_AUTH]: unknown; [AGENT_LOGGER]: RootLogger; + [PI_RPC_CLIENT_FACTORY]: PiRpcClientFactory; // Logger [ROOT_LOGGER]: RootLogger; diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index 63878c6604..00922bc392 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -187,6 +187,12 @@ import { OAUTH_CALLBACK_SERVER } from "@posthog/workspace-server/services/oauth- import { oauthCallbackModule } from "@posthog/workspace-server/services/oauth-callback/oauth-callback.module"; import { onboardingImportModule } from "@posthog/workspace-server/services/onboarding-import/onboarding-import.module"; import { osModule } from "@posthog/workspace-server/services/os/os.module"; +import { + PI_RPC_CLIENT_FACTORY, + PI_SESSION_SERVICE, +} from "@posthog/workspace-server/services/pi-session/identifiers"; +import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; +import { piSessionModule } from "@posthog/workspace-server/services/pi-session/pi-session.module"; import { POSTHOG_PLUGIN_SERVICE } from "@posthog/workspace-server/services/posthog-plugin/identifiers"; import { posthogPluginModule } from "@posthog/workspace-server/services/posthog-plugin/posthog-plugin.module"; import { PROCESS_TRACKING_SERVICE } from "@posthog/workspace-server/services/process-tracking/identifiers"; @@ -223,6 +229,7 @@ import { workspaceModule } from "@posthog/workspace-server/services/workspace/wo import { workspaceMetadataModule } from "@posthog/workspace-server/services/workspace-metadata/workspace-metadata.module"; import ExternalAppsStoreImpl from "electron-store"; import type { FileWatcherBridge } from "../index"; +import { DesktopPiRpcClientFactory } from "../platform-adapters/desktop-pi-rpc-client-factory"; import { ElectronAppLifecycle } from "../platform-adapters/electron-app-lifecycle"; import { ElectronAppMeta } from "../platform-adapters/electron-app-meta"; import { ElectronAppMetrics } from "../platform-adapters/electron-app-metrics"; @@ -359,10 +366,15 @@ container .bind(MAIN_DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY) .toService(DEFAULT_ADDITIONAL_DIRECTORY_REPOSITORY); container.load(agentModule); +container.load(piSessionModule); container.bind(AGENT_SLEEP_COORDINATOR).toService(MAIN_SLEEP_SERVICE); container.bind(AGENT_MCP_APPS).toService(MCP_APPS_SERVICE); container.bind(AGENT_REPO_FILES).toService(MAIN_FS_SERVICE); container.bind(AGENT_AUTH).toService(MAIN_AUTH_SERVICE); +container + .bind(PI_RPC_CLIENT_FACTORY) + .to(DesktopPiRpcClientFactory) + .inSingletonScope(); container.bind(AGENT_LOGGER).toConstantValue(logger); container.load(osModule); container.bind(ROOT_LOGGER).toConstantValue(logger); @@ -394,8 +406,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => { }); container.load(archiveModule); container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: (taskId: string) => - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId: string) => { + await Promise.all([ + ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, })); container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -406,8 +422,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ })); container.load(suspensionModule); container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: (taskId: string) => - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId: string) => { + await Promise.all([ + ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, })); container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -685,7 +705,12 @@ container.load(workspaceModule); container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => { const agent = ctx.get(AGENT_SERVICE); return { - cancelSessionsByTaskId: (taskId) => agent.cancelSessionsByTaskId(taskId), + cancelSessionsByTaskId: async (taskId) => { + await Promise.all([ + agent.cancelSessionsByTaskId(taskId), + ctx.get(PI_SESSION_SERVICE).stop(taskId), + ]); + }, onAgentFileActivity: (handler) => agent.on(AgentServiceEvent.AgentFileActivity, handler), }; diff --git a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts new file mode 100644 index 0000000000..323237e3ef --- /dev/null +++ b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.test.ts @@ -0,0 +1,43 @@ +import type { PiRpcClient } from "@posthog/agent/pi/rpc-client"; +import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; +import { getCloudUrlFromRegion } from "@posthog/shared"; +import type { AgentAuth } from "@posthog/workspace-server/services/agent/ports"; +import type { AuthProxyService } from "@posthog/workspace-server/services/auth-proxy/auth-proxy"; +import { describe, expect, it, vi } from "vitest"; +import { DesktopPiRpcClientFactory } from "./desktop-pi-rpc-client-factory"; + +const createPiRpcClient = vi.hoisted(() => vi.fn()); + +vi.mock("@posthog/agent/pi/rpc-client", () => ({ createPiRpcClient })); + +describe("DesktopPiRpcClientFactory", () => { + it("routes Pi through the shared host auth proxy", async () => { + const auth = { + getOAuthCredentials: vi.fn(async () => ({ + access: "access-token", + refresh: "refresh-token", + expires: 1, + region: "eu" as const, + })), + } as unknown as AgentAuth; + const authProxy = { + start: vi.fn(async () => "http://127.0.0.1:1234"), + } as unknown as AuthProxyService; + const client = {} as PiRpcClient; + createPiRpcClient.mockReturnValue(client); + const factory = new DesktopPiRpcClientFactory(auth, authProxy); + + await expect(factory.create({ cwd: "/workspace" })).resolves.toBe(client); + expect(authProxy.start).toHaveBeenCalledWith( + getLlmGatewayUrl(getCloudUrlFromRegion("eu")), + ); + expect(createPiRpcClient).toHaveBeenCalledWith({ + cwd: "/workspace", + providerOptions: { + region: "eu", + baseUrl: "http://127.0.0.1:1234", + apiKey: "posthog-code-auth-proxy", + }, + }); + }); +}); diff --git a/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts new file mode 100644 index 0000000000..37d36f54dc --- /dev/null +++ b/apps/code/src/main/platform-adapters/desktop-pi-rpc-client-factory.ts @@ -0,0 +1,58 @@ +import { + createPiRpcClient, + type PiRpcClient, +} from "@posthog/agent/pi/rpc-client"; +import { getLlmGatewayUrl } from "@posthog/agent/posthog-api"; +import { type CloudRegion, getCloudUrlFromRegion } from "@posthog/shared"; +import { AGENT_AUTH } from "@posthog/workspace-server/services/agent/identifiers"; +import type { AgentAuth } from "@posthog/workspace-server/services/agent/ports"; +import type { AuthProxyService } from "@posthog/workspace-server/services/auth-proxy/auth-proxy"; +import { AUTH_PROXY_SERVICE } from "@posthog/workspace-server/services/auth-proxy/identifiers"; +import type { PiRpcClientFactory } from "@posthog/workspace-server/services/pi-session/identifiers"; +import { inject, injectable } from "inversify"; + +const PROXY_API_KEY = "posthog-code-auth-proxy"; + +@injectable() +export class DesktopPiRpcClientFactory implements PiRpcClientFactory { + private proxyRegion?: CloudRegion; + private proxyUrlPromise?: Promise; + + constructor( + @inject(AGENT_AUTH) private readonly auth: AgentAuth, + @inject(AUTH_PROXY_SERVICE) + private readonly authProxy: AuthProxyService, + ) {} + + async create(input: { + cwd: string; + model?: string; + sessionFile?: string; + }): Promise { + const credentials = await this.auth.getOAuthCredentials(); + if (!credentials) { + throw new Error("Pi requires PostHog authentication"); + } + + const baseUrl = await this.getProxyUrl(credentials.region); + + return createPiRpcClient({ + ...input, + providerOptions: { + region: credentials.region, + baseUrl, + apiKey: PROXY_API_KEY, + }, + }); + } + + private getProxyUrl(region: CloudRegion): Promise { + if (this.proxyRegion !== region || !this.proxyUrlPromise) { + this.proxyRegion = region; + const gatewayUrl = getLlmGatewayUrl(getCloudUrlFromRegion(region)); + this.proxyUrlPromise = this.authProxy.start(gatewayUrl); + } + + return this.proxyUrlPromise; + } +} diff --git a/apps/code/src/main/trpc/router.ts b/apps/code/src/main/trpc/router.ts index 68ef63321c..f2209456c6 100644 --- a/apps/code/src/main/trpc/router.ts +++ b/apps/code/src/main/trpc/router.ts @@ -37,6 +37,7 @@ import { notificationRouter } from "@posthog/host-router/routers/notification.ro import { oauthRouter } from "@posthog/host-router/routers/oauth.router"; import { onboardingImportRouter } from "@posthog/host-router/routers/onboarding-import.router"; import { osRouter } from "@posthog/host-router/routers/os.router"; +import { piSessionRouter } from "@posthog/host-router/routers/pi-session.router"; import { processTrackingRouter } from "@posthog/host-router/routers/process-tracking.router"; import { provisioningRouter } from "@posthog/host-router/routers/provisioning.router"; import { secureStoreRouter } from "@posthog/host-router/routers/secure-store.router"; @@ -97,6 +98,7 @@ export const trpcRouter = router({ onboardingImport: onboardingImportRouter, logs: logsRouter, os: osRouter, + piSession: piSessionRouter, processTracking: processTrackingRouter, provisioning: provisioningRouter, sleep: sleepRouter, diff --git a/apps/code/src/renderer/di/bindings.ts b/apps/code/src/renderer/di/bindings.ts index 89926b3b7f..3d25557e3f 100644 --- a/apps/code/src/renderer/di/bindings.ts +++ b/apps/code/src/renderer/di/bindings.ts @@ -73,6 +73,8 @@ import { GITHUB_CONNECT_CLIENT, type GithubConnectClient, } from "@posthog/core/onboarding/identifiers"; +import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; +import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { type BundleLocalSkill, CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, @@ -288,6 +290,7 @@ export interface RendererBindings { [SHELL_PROCESS_READER]: ShellProcessReader; [ANALYTICS_TRACKER]: AnalyticsTracker; [TASK_CREATION_HOST]: ITaskCreationHost; + [PI_RUNNER]: PiRunner; [TASK_CREATION_EFFECTS]: TaskCreationEffects; [RENDERER_TASK_SERVICE]: TaskService; [TASK_SERVICE]: TaskService; diff --git a/apps/code/src/renderer/di/container.ts b/apps/code/src/renderer/di/container.ts index dd18faaa23..9006191ad1 100644 --- a/apps/code/src/renderer/di/container.ts +++ b/apps/code/src/renderer/di/container.ts @@ -32,6 +32,8 @@ import type { LlmGatewayService } from "@posthog/core/llm-gateway/llm-gateway"; import type { LlmMessage } from "@posthog/core/llm-gateway/schemas"; import { LOCAL_MCP_WORKSPACE_CLIENT } from "@posthog/core/local-mcp/identifiers"; import type { LocalMcpWorkspaceClient } from "@posthog/core/local-mcp/localMcpImport"; +import { PI_RUNNER } from "@posthog/core/pi-runtime/identifiers"; +import type { PiRunner } from "@posthog/core/pi-runtime/piRunner"; import { CLOUD_ARTIFACT_BUNDLE_LOCAL_SKILL, CLOUD_ARTIFACT_READ_FILE_AS_BASE64, @@ -151,6 +153,7 @@ import { trpcClient } from "@renderer/trpc"; import { hostTrpcClient } from "@renderer/trpc/client"; import type { TRPCClient } from "@trpc/client"; import { hostLog, logger } from "@utils/logger"; +import { TrpcPiRunner } from "../platform-adapters/trpc-pi-runner"; import type { RendererBindings } from "./bindings"; import { TASK_SERVICE as RENDERER_TASK_SERVICE, TRPC_CLIENT } from "./tokens"; @@ -290,6 +293,7 @@ container // Bind services container.bind(TASK_CREATION_HOST).to(TrpcTaskCreationHost); +container.bind(PI_RUNNER).to(TrpcPiRunner); container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects); container.bind(RENDERER_TASK_SERVICE).to(TaskService); container.bind(TASK_SERVICE).toService(RENDERER_TASK_SERVICE); diff --git a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts new file mode 100644 index 0000000000..3f4311f7fb --- /dev/null +++ b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts @@ -0,0 +1,28 @@ +import type { + PiResumeInput, + PiRunInput, + PiRunner, +} from "@posthog/core/pi-runtime/piRunner"; +import { resolveService } from "@posthog/di/container"; +import { + HOST_TRPC_CLIENT, + type HostTrpcClient, +} from "@posthog/host-router/client"; + +function hostClient(): HostTrpcClient { + return resolveService(HOST_TRPC_CLIENT); +} + +export class TrpcPiRunner implements PiRunner { + async create(input: PiRunInput): Promise { + await hostClient().piSession.start.mutate(input); + } + + resume(input: PiResumeInput): Promise { + return hostClient().piSession.resume.mutate(input); + } + + stop(taskId: string): Promise { + return hostClient().piSession.stop.mutate({ taskId }); + } +} diff --git a/apps/code/src/renderer/trpc/client.ts b/apps/code/src/renderer/trpc/client.ts index 34151a047e..8cc86cdc64 100644 --- a/apps/code/src/renderer/trpc/client.ts +++ b/apps/code/src/renderer/trpc/client.ts @@ -7,14 +7,18 @@ import { createTRPCOptionsProxy, } from "@trpc/tanstack-react-query"; import { queryClient } from "@utils/queryClient"; +import superjson from "superjson"; import type { TrpcRouter } from "../../main/trpc/router"; export const trpcClient = createTRPCClient({ - links: [ipcInstrumentationLink(), ipcLink()], + links: [ + ipcInstrumentationLink(), + ipcLink({ transformer: superjson }), + ], }); export const hostTrpcClient = createTRPCClient({ - links: [ipcLink()], + links: [ipcLink({ transformer: superjson })], }); const context = createTRPCContext(); diff --git a/apps/code/vite-main-plugins.mts b/apps/code/vite-main-plugins.mts index f83bea74b7..8d3d23e295 100644 --- a/apps/code/vite-main-plugins.mts +++ b/apps/code/vite-main-plugins.mts @@ -163,6 +163,29 @@ function copyClaudeSupportAssets(sourcePath: string, destDir: string): void { } } +export function copyPiRpcHost(): Plugin { + return { + name: "copy-pi-rpc-host", + writeBundle() { + const candidates = [ + join(__dirname, "node_modules/@posthog/agent/dist/pi/rpc-host.js"), + join( + __dirname, + "../../node_modules/@posthog/agent/dist/pi/rpc-host.js", + ), + join(__dirname, "../../packages/agent/dist/pi/rpc-host.js"), + ]; + const source = candidates.find((candidate) => existsSync(candidate)); + if (!source) { + throw new Error( + `[copy-pi-rpc-host] Unable to find Pi RPC host, required at runtime by createPiRpcClient. Build @posthog/agent first. Checked:\n ${candidates.join("\n ")}`, + ); + } + copyFileSync(source, join(__dirname, ".vite/build/rpc-host.js")); + }, + }; +} + export function copyClaudeExecutable(): Plugin { return { name: "copy-claude-executable", diff --git a/apps/web/package.json b/apps/web/package.json index d06a50e056..d30e7da97f 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -20,6 +20,7 @@ "@posthog/workspace-client": "workspace:*", "@tanstack/react-query": "^5.100.14", "@trpc/client": "^11.17.0", + "superjson": "catalog:", "@trpc/tanstack-react-query": "^11.17.0", "inversify": "^7.10.6", "react": "19.2.6", diff --git a/apps/web/src/web-trpc.ts b/apps/web/src/web-trpc.ts index c5ffce58ac..b4e7d36a83 100644 --- a/apps/web/src/web-trpc.ts +++ b/apps/web/src/web-trpc.ts @@ -5,6 +5,7 @@ import { httpSubscriptionLink, splitLink, } from "@trpc/client"; +import superjson from "superjson"; // The ENTIRE electron->web transport difference. The renderer builds the same // client with `links: [ipcLink()]`; web swaps in HTTP: httpBatchLink for @@ -16,8 +17,8 @@ export const hostTrpcClient = createTRPCClient({ links: [ splitLink({ condition: (op) => op.type === "subscription", - true: httpSubscriptionLink({ url: API_URL }), - false: httpBatchLink({ url: API_URL }), + true: httpSubscriptionLink({ url: API_URL, transformer: superjson }), + false: httpBatchLink({ url: API_URL, transformer: superjson }), }), ], }); diff --git a/packages/agent/package.json b/packages/agent/package.json index 67352582de..a2b81f1ee6 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -28,6 +28,10 @@ "types": "./dist/posthog-products.d.ts", "import": "./dist/posthog-products.js" }, + "./pi/rpc-client": { + "types": "./dist/pi/rpc-client.d.ts", + "import": "./dist/pi/rpc-client.js" + }, "./pr-url-detector": { "types": "./dist/pr-url-detector.d.ts", "import": "./dist/pr-url-detector.js" diff --git a/packages/agent/src/pi/rpc-client.test.ts b/packages/agent/src/pi/rpc-client.test.ts index 449b9b3d25..4c8239d698 100644 --- a/packages/agent/src/pi/rpc-client.test.ts +++ b/packages/agent/src/pi/rpc-client.test.ts @@ -7,7 +7,11 @@ describe("createPiRpcClient", () => { const client = createPiRpcClient({ cwd: "/workspace", model: "claude-opus-4-8", - providerOptions: { apiKey: "token", region: "us" }, + providerOptions: { + region: "us", + baseUrl: "http://127.0.0.1:1234", + apiKey: "proxy-key", + }, }); expect(client).toBeInstanceOf(RpcClient); diff --git a/packages/agent/src/pi/rpc-client.ts b/packages/agent/src/pi/rpc-client.ts index c93e946f89..85896bf6f4 100644 --- a/packages/agent/src/pi/rpc-client.ts +++ b/packages/agent/src/pi/rpc-client.ts @@ -11,6 +11,10 @@ import { safePiEnvironment } from "./rpc-environment"; export type PiRpcClient = RpcClient; +type RpcClientProcessAccess = { + process?: ChildProcess; +}; + interface RpcClientInternals { process?: ChildProcess; stopReadingStdout?: () => void; @@ -47,7 +51,7 @@ function attachJsonlReader( class SecurePiRpcClient extends RpcClient { constructor( private readonly secureOptions: RpcClientOptions, - private readonly providerOptions?: PosthogProviderOptions, + private readonly providerOptions: PosthogProviderOptions, ) { super(secureOptions); } @@ -133,17 +137,24 @@ class SecurePiRpcClient extends RpcClient { } } +export function getPiRpcClientProcess( + client: PiRpcClient, +): ChildProcess | null { + return (client as unknown as RpcClientProcessAccess).process ?? null; +} + export type PiRpcClientOptions = Pick & { - providerOptions?: PosthogProviderOptions; + sessionFile?: string; + providerOptions: PosthogProviderOptions & { apiKey: string }; }; -export function createPiRpcClient( - options: PiRpcClientOptions = {}, -): PiRpcClient { - const { providerOptions, ...rpcOptions } = options; +export function createPiRpcClient(options: PiRpcClientOptions): PiRpcClient { + const { sessionFile, providerOptions, ...rpcOptions } = options; + const args = sessionFile ? ["--session-file", sessionFile] : []; return new SecurePiRpcClient( { ...rpcOptions, + args, cliPath: fileURLToPath(new URL("./rpc-host.js", import.meta.url)), provider: "posthog", }, diff --git a/packages/agent/src/pi/rpc-host.ts b/packages/agent/src/pi/rpc-host.ts index 951739bbe5..0ce15c5e68 100644 --- a/packages/agent/src/pi/rpc-host.ts +++ b/packages/agent/src/pi/rpc-host.ts @@ -1,4 +1,5 @@ import { readFileSync } from "node:fs"; +import { SessionManager } from "@earendil-works/pi-coding-agent"; import { createHarnessRuntime, runRpcMode } from "@posthog/harness"; import type { PosthogProviderOptions } from "@posthog/harness/extensions/posthog-provider/provider"; import { sanitizePiHostEnvironment } from "./rpc-environment"; @@ -13,10 +14,21 @@ function argumentValue(name: string): string | undefined { } const bootstrap = JSON.parse(readFileSync(3, "utf8")) as PiRpcBootstrap; +const providerOptions = bootstrap.providerOptions; +if (!providerOptions?.apiKey) { + throw new Error("Pi RPC host requires PostHog provider credentials"); +} sanitizePiHostEnvironment(); + +const cwd = process.cwd(); +const sessionFile = argumentValue("--session-file"); +const sessionManager = sessionFile + ? SessionManager.open(sessionFile, undefined, cwd) + : undefined; const runtime = await createHarnessRuntime({ - cwd: process.cwd(), - ...bootstrap.providerOptions, + cwd, + sessionManager, + ...providerOptions, }); const requestedModel = argumentValue("--model")?.replace(/^posthog\//, ""); diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index b1e053fa84..93a1e2ce88 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -87,6 +87,7 @@ const sharedOptions = { "@posthog/shared", "@posthog/git", "@posthog/enricher", + "@posthog/harness", "fflate", ], external: [ @@ -113,7 +114,6 @@ export default defineConfig([ "src/posthog-products.ts", "src/pr-url-detector.ts", "src/pi/rpc-client.ts", - "src/pi/rpc-host.ts", "src/resume.ts", "src/types.ts", "src/adapters/claude/questions/utils.ts", @@ -169,4 +169,16 @@ export default defineConfig([ clean: false, ...sharedOptions, }, + { + entry: { "pi/rpc-host": "src/pi/rpc-host.ts" }, + format: ["esm"], + dts: false, + clean: false, + banner: { + js: 'import { createRequire as __createRequire } from "node:module"; const require = __createRequire(import.meta.url);', + }, + ...sharedOptions, + noExternal: [/^(?!node:)/], + external: [...builtinModules, ...builtinModules.map((m) => `node:${m}`)], + }, ]); diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index a9c5e72318..5478769308 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -2426,6 +2426,7 @@ export class PostHogAPIClient { | "repository" | "json_schema" | "origin_product" + | "runtime" | "signal_report" > > & { diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index 9f9452ec04..ec0512ed4e 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -155,6 +155,22 @@ export class AuthService extends TypedEventEmitter { apiHost: getCloudUrlFromRegion(session.cloudRegion), }; } + async getOAuthCredentials(): Promise<{ + access: string; + refresh: string; + expires: number; + region: CloudRegion; + } | null> { + if (this.tokenOverride) return null; + await this.initialize(); + const session = await this.ensureValidSession(); + return { + access: session.accessToken, + refresh: session.refreshToken, + expires: session.accessTokenExpiresAt, + region: session.cloudRegion, + }; + } async refreshAccessToken(): Promise { const override = this.tokenOverride; if (override) { diff --git a/packages/core/src/inbox/reportTaskCreation.ts b/packages/core/src/inbox/reportTaskCreation.ts index ee85056652..67f5d9cb95 100644 --- a/packages/core/src/inbox/reportTaskCreation.ts +++ b/packages/core/src/inbox/reportTaskCreation.ts @@ -9,7 +9,7 @@ export interface PreviewConfigChoice { /** Minimal shape of a preview-config option we scan for the default model. */ export interface PreviewConfigOption { id?: string; - category?: string; + category?: string | null; type?: string; currentValue?: string | boolean | null; options?: PreviewConfigChoice[]; diff --git a/packages/core/src/pi-runtime/identifiers.ts b/packages/core/src/pi-runtime/identifiers.ts new file mode 100644 index 0000000000..42e0c14ab3 --- /dev/null +++ b/packages/core/src/pi-runtime/identifiers.ts @@ -0,0 +1 @@ +export const PI_RUNNER = Symbol.for("posthog.pi.runner"); diff --git a/packages/core/src/pi-runtime/piRunner.ts b/packages/core/src/pi-runtime/piRunner.ts new file mode 100644 index 0000000000..b174fb95fd --- /dev/null +++ b/packages/core/src/pi-runtime/piRunner.ts @@ -0,0 +1,17 @@ +export interface PiRunInput { + taskId: string; + cwd: string; + prompt: string; + model?: string; +} + +export interface PiResumeInput { + taskId: string; + cwd: string; +} + +export interface PiRunner { + create(input: PiRunInput): Promise; + resume(input: PiResumeInput): Promise; + stop(taskId: string): Promise; +} diff --git a/packages/core/src/task-detail/piTaskCreator.ts b/packages/core/src/task-detail/piTaskCreator.ts new file mode 100644 index 0000000000..b453b6f5aa --- /dev/null +++ b/packages/core/src/task-detail/piTaskCreator.ts @@ -0,0 +1,197 @@ +import { + Saga, + type SagaLogger, + type TaskCreationInput, + type TaskCreationOutput, + type Workspace, +} from "@posthog/shared"; +import type { Task } from "@posthog/shared/domain-types"; +import type { PiRunner } from "../pi-runtime/piRunner"; +import type { TaskCreationApiClient } from "./taskCreationApiClient"; +import type { ITaskCreationHost } from "./taskCreationHost"; +import { resolveTaskRepository } from "./taskRepository"; + +export interface PiTaskCreatorDeps { + posthogClient: TaskCreationApiClient; + host: ITaskCreationHost; + piRunner: PiRunner; + onTaskReady?: (output: TaskCreationOutput) => void; +} + +export class PiTaskCreator extends Saga { + readonly sagaName = "PiTaskCreator"; + + constructor( + private readonly deps: PiTaskCreatorDeps, + logger?: SagaLogger, + ) { + super(logger); + } + + protected async execute( + input: TaskCreationInput, + ): Promise { + if (input.workspaceMode === "cloud") { + throw new Error("Pi tasks are only supported in local workspaces"); + } + + const task = await this.createTask(input); + const repoPath = input.repoPath; + let workspace: Workspace | null = null; + + if (repoPath) { + workspace = await this.createWorkspace(task, repoPath, input); + } else if (input.allowNoRepo) { + workspace = await this.createScratchWorkspace(task); + } + + if (!workspace) { + throw new Error("Pi tasks require a workspace or scratch directory"); + } + + const cwd = workspace.worktreePath ?? workspace.folderPath; + const additionalDirectories = (input.additionalDirectories ?? []).filter( + (path) => path && path !== input.repoPath, + ); + if (additionalDirectories.length > 0) { + await this.step({ + name: "additional_directories", + execute: async () => { + await Promise.all( + additionalDirectories.map((path) => + this.deps.host.addAdditionalDirectory({ taskId: task.id, path }), + ), + ); + return { taskId: task.id, paths: additionalDirectories }; + }, + rollback: async ({ taskId, paths }) => { + await Promise.all( + paths.map((path) => + this.deps.host.removeAdditionalDirectory({ taskId, path }), + ), + ); + }, + }); + } + + await this.step({ + name: "pi_session", + execute: async () => { + await this.deps.piRunner.create({ + taskId: task.id, + cwd, + prompt: input.content ?? "", + model: input.model, + }); + return { taskId: task.id }; + }, + rollback: async ({ taskId }) => this.deps.piRunner.stop(taskId), + }); + + this.deps.onTaskReady?.({ task, workspace }); + return { task, workspace }; + } + + private async createTask(input: TaskCreationInput): Promise { + const repository = await resolveTaskRepository( + input, + this.deps.host, + this.log, + ); + + return this.step({ + name: "task_creation", + execute: async () => + (await this.deps.posthogClient.createTask({ + description: input.content ?? "", + repository: repository ?? undefined, + origin_product: input.signalReportId + ? "signal_report" + : "user_created", + signal_report: input.signalReportId ?? undefined, + channel: input.channelId ?? undefined, + runtime: "pi", + })) as unknown as Task, + rollback: async (task) => this.deps.posthogClient.deleteTask(task.id), + }); + } + + private async createWorkspace( + task: Task, + repoPath: string, + input: TaskCreationInput, + ): Promise { + const folder = await this.deps.host.getFolders().then(async (folders) => { + const existing = folders.find((candidate) => candidate.path === repoPath); + return existing ?? this.deps.host.addFolder({ folderPath: repoPath }); + }); + const workspaceInfo = await this.step({ + name: "workspace_creation", + execute: () => + this.deps.host.createWorkspace({ + taskId: task.id, + mainRepoPath: repoPath, + folderId: folder.id, + folderPath: repoPath, + mode: input.workspaceMode ?? "local", + branch: input.branch ?? undefined, + allowRemoteBranchCheckout: input.allowRemoteBranchCheckout, + reuseExistingWorktree: input.reuseExistingWorktree, + }), + rollback: () => + this.deps.host.deleteWorkspace({ + taskId: task.id, + mainRepoPath: repoPath, + }), + }); + + const workspaceMode = input.workspaceMode ?? "local"; + const worktree = workspaceInfo.worktree; + if (workspaceMode === "worktree" && !worktree) { + throw new Error("Pi worktree creation did not return a worktree"); + } + if (worktree) { + return { + taskId: task.id, + folderId: folder.id, + folderPath: repoPath, + mode: workspaceMode, + worktreePath: worktree.worktreePath, + worktreeName: worktree.worktreeName, + branchName: worktree.branchName, + baseBranch: worktree.baseBranch, + linkedBranch: workspaceInfo.linkedBranch, + createdAt: worktree.createdAt, + }; + } + + return { + taskId: task.id, + folderId: folder.id, + folderPath: repoPath, + mode: "local", + worktreePath: null, + worktreeName: null, + branchName: workspaceInfo.branchName, + baseBranch: input.branch ?? null, + linkedBranch: workspaceInfo.linkedBranch, + createdAt: new Date().toISOString(), + }; + } + + private async createScratchWorkspace(task: Task): Promise { + const folderPath = await this.deps.host.ensureScratchDir(task.id); + return { + taskId: task.id, + folderId: "", + folderPath, + mode: "local", + worktreePath: null, + worktreeName: null, + branchName: null, + baseBranch: null, + linkedBranch: null, + createdAt: new Date().toISOString(), + }; + } +} diff --git a/packages/core/src/task-detail/taskCreationHost.ts b/packages/core/src/task-detail/taskCreationHost.ts index a419ccf60a..92aefe9d1c 100644 --- a/packages/core/src/task-detail/taskCreationHost.ts +++ b/packages/core/src/task-detail/taskCreationHost.ts @@ -1,6 +1,6 @@ import type { ContentBlock } from "@agentclientprotocol/sdk"; import type { CloudSkillBundleRef } from "@posthog/core/sessions/cloudArtifactIdentifiers"; -import type { Workspace, WorkspaceMode } from "@posthog/shared"; +import type { Workspace, WorkspaceInfo, WorkspaceMode } from "@posthog/shared"; import type { TaskCreationApiClient } from "./taskCreationApiClient"; export interface CloudPromptTransport { @@ -21,16 +21,7 @@ export interface CreateWorkspaceArgs { reuseExistingWorktree?: boolean; } -export interface CreatedWorkspaceInfo { - worktree?: { - worktreePath?: string | null; - worktreeName?: string | null; - branchName?: string | null; - baseBranch?: string | null; - createdAt?: string | null; - } | null; - linkedBranch?: string | null; -} +export type CreatedWorkspaceInfo = WorkspaceInfo; export interface TaskFolderInfo { id: string; diff --git a/packages/core/src/task-detail/taskCreationSaga.test.ts b/packages/core/src/task-detail/taskCreationSaga.test.ts index 6a06da8206..5558ab6580 100644 --- a/packages/core/src/task-detail/taskCreationSaga.test.ts +++ b/packages/core/src/task-detail/taskCreationSaga.test.ts @@ -10,6 +10,8 @@ const mockHost = vi.hoisted(() => ({ getAuthenticatedClient: vi.fn(), getTaskDirectory: vi.fn(), ensureScratchDir: vi.fn(), + startPiSession: vi.fn(), + stopPiSession: vi.fn(), getWorkspace: vi.fn(), createWorkspace: vi.fn(), deleteWorkspace: vi.fn(), @@ -35,6 +37,7 @@ const mockHost = vi.hoisted(() => ({ linkTaskBranch: vi.fn(), })); +import { PiTaskCreator } from "./piTaskCreator"; import { TaskCreationSaga } from "./taskCreationSaga"; import { buildWorktreeAdoptionInput } from "./taskInput"; @@ -339,6 +342,42 @@ describe("TaskCreationSaga", () => { ); }); + it("starts a Pi session without creating an ACP session", async () => { + const createdTask = createTask({ repository: undefined }); + const createTaskRequest = vi.fn().mockResolvedValue(createdTask); + const saga = new PiTaskCreator({ + posthogClient: { + createTask: createTaskRequest, + deleteTask: vi.fn(), + } as never, + host, + piRunner: { + create: mockHost.startPiSession, + stop: mockHost.stopPiSession, + } as never, + }); + + const result = await saga.run({ + content: "Draft a launch email", + workspaceMode: "local", + runtime: "pi", + model: "claude-sonnet", + allowNoRepo: true, + }); + + expect(result.success).toBe(true); + expect(createTaskRequest).toHaveBeenCalledWith( + expect.objectContaining({ runtime: "pi" }), + ); + expect(mockHost.startPiSession).toHaveBeenCalledWith({ + taskId: "task-123", + cwd: "/tmp/scratch/task-123", + prompt: "Draft a launch email", + model: "claude-sonnet", + }); + expect(sessionService.connectToTask).not.toHaveBeenCalled(); + }); + it("uploads initial cloud attachments before starting the run", async () => { const createdTask = createTask(); const startedTask = createTask({ latest_run: createRun() }); diff --git a/packages/core/src/task-detail/taskCreationSaga.ts b/packages/core/src/task-detail/taskCreationSaga.ts index d74a341d35..e213fda451 100644 --- a/packages/core/src/task-detail/taskCreationSaga.ts +++ b/packages/core/src/task-detail/taskCreationSaga.ts @@ -24,6 +24,7 @@ import type { ImportedClaudeCliSession, ITaskCreationHost, } from "./taskCreationHost"; +import { resolveTaskRepository } from "./taskRepository"; export interface TaskCreationDeps { posthogClient: TaskCreationApiClient; @@ -739,27 +740,9 @@ export class TaskCreationSaga extends Saga< input: TaskCreationInput, warmPayload: WarmActivationPayload | null, ): Promise { - let repository = input.repository; - - const repoPathForDetection = input.repoPath; - if (!repository && repoPathForDetection) { - // Detection only fills the org/repo metadata on the task; a transient - // failure (e.g. the workspace-server transport dropping mid-call) must - // not abort task creation, so degrade to an untagged task instead. - const detected = await this.readOnlyStep("repo_detection", () => - this.deps.host - .detectRepo({ directoryPath: repoPathForDetection }) - .catch((error) => { - this.log.warn("Repo detection failed; creating task without one", { - error, - }); - return null; - }), - ); - if (detected) { - repository = `${detected.organization}/${detected.repository}`; - } - } + const repository = await this.readOnlyStep("repo_detection", () => + resolveTaskRepository(input, this.deps.host, this.log), + ); return this.step({ name: "task_creation", @@ -807,6 +790,7 @@ export class TaskCreationSaga extends Saga< : undefined, signal_report: input.signalReportId ?? undefined, channel: input.channelId ?? undefined, + runtime: "acp", pending_user_message: warmPayload?.pendingUserMessage, pending_user_artifact_ids: warmPayload?.pendingUserArtifactIds, // If creation activates a pre-warmed run, this is the only request diff --git a/packages/core/src/task-detail/taskInput.test.ts b/packages/core/src/task-detail/taskInput.test.ts index aa6d3d6dc1..e30a5a4c39 100644 --- a/packages/core/src/task-detail/taskInput.test.ts +++ b/packages/core/src/task-detail/taskInput.test.ts @@ -21,6 +21,23 @@ describe("prepareTaskInput", () => { }, ); + it("defaults task creation to the ACP runtime", () => { + const input = prepareTaskInput("do the thing", [], { + workspaceMode: "local", + }); + + expect(input.runtime).toBe("acp"); + }); + + it("preserves the selected Pi runtime", () => { + const input = prepareTaskInput("do the thing", [], { + workspaceMode: "local", + runtime: "pi", + }); + + expect(input.runtime).toBe("pi"); + }); + it("drops customInstructions for cloud when none is set", () => { const input = prepareTaskInput("do the thing", [], { workspaceMode: "cloud", diff --git a/packages/core/src/task-detail/taskInput.ts b/packages/core/src/task-detail/taskInput.ts index 63902749ca..67f2edb432 100644 --- a/packages/core/src/task-detail/taskInput.ts +++ b/packages/core/src/task-detail/taskInput.ts @@ -1,6 +1,7 @@ import { buildCloudTaskDescription } from "@posthog/core/editor/cloud-prompt"; import type { Adapter, + AgentRuntime, CloudMcpServerImport, CloudMcpServerRelayDesignation, TaskCreationInput, @@ -19,6 +20,7 @@ export interface PrepareTaskInputOptions { reuseExistingWorktree?: boolean; executionMode?: ExecutionMode; adapter?: Adapter; + runtime?: AgentRuntime; model?: string; reasoningLevel?: string; environmentId?: string | null; @@ -59,6 +61,7 @@ export function prepareTaskInput( reuseExistingWorktree: options.reuseExistingWorktree, executionMode: options.executionMode, adapter: options.adapter, + runtime: options.runtime ?? "acp", model: options.model, reasoningLevel: options.reasoningLevel, environmentId: options.environmentId ?? undefined, diff --git a/packages/core/src/task-detail/taskRepository.ts b/packages/core/src/task-detail/taskRepository.ts new file mode 100644 index 0000000000..ccd22a9bcf --- /dev/null +++ b/packages/core/src/task-detail/taskRepository.ts @@ -0,0 +1,26 @@ +import type { SagaLogger, TaskCreationInput } from "@posthog/shared"; +import type { ITaskCreationHost } from "./taskCreationHost"; + +export async function resolveTaskRepository( + input: TaskCreationInput, + host: Pick, + logger: Pick, +): Promise { + if (input.repository) { + return input.repository; + } + if (!input.repoPath) { + return undefined; + } + + try { + const detected = await host.detectRepo({ directoryPath: input.repoPath }); + if (!detected) { + return undefined; + } + return `${detected.organization}/${detected.repository}`; + } catch (error) { + logger.warn("Repo detection failed; creating task without one", { error }); + return undefined; + } +} diff --git a/packages/core/src/task-detail/taskService.test.ts b/packages/core/src/task-detail/taskService.test.ts index 049e775170..7883f3f4b5 100644 --- a/packages/core/src/task-detail/taskService.test.ts +++ b/packages/core/src/task-detail/taskService.test.ts @@ -1,6 +1,7 @@ import type { SessionService } from "@posthog/core/sessions/sessionService"; import type { RootLogger } from "@posthog/di/logger"; import { describe, expect, it, vi } from "vitest"; +import type { PiRunner } from "../pi-runtime/piRunner"; import type { TaskCreationEffects } from "./taskCreationEffects"; import type { ITaskCreationHost } from "./taskCreationHost"; import { buildWorktreeAdoptionInput } from "./taskInput"; @@ -45,7 +46,12 @@ function makeService(): TaskService { onWorkspaceCreated: vi.fn(), onCreateSuccess: vi.fn(), } as unknown as TaskCreationEffects; - return new TaskService(host, sessionService, effects, rootLogger); + const piRunner = { + create: vi.fn(), + resume: vi.fn(), + stop: vi.fn(), + } as unknown as PiRunner; + return new TaskService(host, sessionService, effects, piRunner, rootLogger); } describe("TaskService.createTask validation", () => { diff --git a/packages/core/src/task-detail/taskService.ts b/packages/core/src/task-detail/taskService.ts index 78e6e0c25f..73ea84990f 100644 --- a/packages/core/src/task-detail/taskService.ts +++ b/packages/core/src/task-detail/taskService.ts @@ -11,7 +11,10 @@ import type { } from "@posthog/shared"; import type { Task } from "@posthog/shared/domain-types"; import { inject, injectable } from "inversify"; +import { PI_RUNNER } from "../pi-runtime/identifiers"; +import type { PiRunner } from "../pi-runtime/piRunner"; import { TASK_CREATION_EFFECTS, TASK_CREATION_HOST } from "./identifiers"; +import { PiTaskCreator } from "./piTaskCreator"; import type { TaskCreationEffects } from "./taskCreationEffects"; import type { ITaskCreationHost } from "./taskCreationHost"; import { TaskCreationSaga } from "./taskCreationSaga"; @@ -42,6 +45,8 @@ export class TaskService { private readonly sessionService: SessionService, @inject(TASK_CREATION_EFFECTS) private readonly effects: TaskCreationEffects, + @inject(PI_RUNNER) + private readonly piRunner: PiRunner, @inject(ROOT_LOGGER) rootLogger: RootLogger, ) { @@ -103,30 +108,35 @@ export class TaskService { } } - const saga = new TaskCreationSaga( - { - posthogClient, - host: this.host, - sessionService: this.sessionService, - track: (event, props) => this.host.track(event, props), - onTaskReady: onTaskReady - ? (output) => { - this.effects.onWorkspaceCreated(output); - this.effects.onCreateSuccess(output, input); - onTaskReady(output); - } - : undefined, - }, - this.log, - ); - - const result = await saga.run(input); + let result: CreateTaskResult; + if (input.runtime === "pi") { + const creator = new PiTaskCreator( + { + posthogClient, + host: this.host, + piRunner: this.piRunner, + onTaskReady, + }, + this.log, + ); + result = await creator.run(input); + } else { + const creator = new TaskCreationSaga( + { + posthogClient, + host: this.host, + sessionService: this.sessionService, + track: (event, props) => this.host.track(event, props), + onTaskReady, + }, + this.log, + ); + result = await creator.run(input); + } if (result.success) { this.effects.onWorkspaceCreated(result.data); - if (!onTaskReady) { - this.effects.onCreateSuccess(result.data, input); - } + this.effects.onCreateSuccess(result.data, input); } return result; @@ -147,31 +157,63 @@ export class TaskService { }; } + let task: Task; + try { + task = await posthogClient.getTask(taskId); + if (taskRunId) { + this.log.info("Fetching specific task run", { taskId, taskRunId }); + task.latest_run = await posthogClient.getTaskRun(taskId, taskRunId); + } + } catch (error) { + return { + success: false, + error: error instanceof Error ? error.message : "Failed to fetch task", + failedStep: "fetch_task", + }; + } + + const runtime = task.runtime === "pi" ? "pi" : "acp"; const existingWorkspace = await this.host.getWorkspace(taskId); if (existingWorkspace) { this.log.info("Workspace already exists, fetching task only", { taskId }); try { - const task = await posthogClient.getTask(taskId); - - if (taskRunId) { - this.log.info("Fetching specific task run", { taskId, taskRunId }); - const run = await posthogClient.getTaskRun(taskId, taskRunId); - task.latest_run = run; + if (runtime === "pi") { + await this.piRunner.resume({ + taskId, + cwd: existingWorkspace.worktreePath ?? existingWorkspace.folderPath, + }); } return { success: true, - data: { - task: task as unknown as Task, - workspace: existingWorkspace, - }, + data: { task, workspace: existingWorkspace }, + }; + } catch (error) { + return { + success: false, + error: + error instanceof Error ? error.message : "Failed to resume task", + failedStep: runtime === "pi" ? "pi_session" : "fetch_task", + }; + } + } + + if (runtime === "pi") { + try { + const cwd = await this.host.ensureScratchDir(taskId); + await this.piRunner.resume({ taskId, cwd }); + return { + success: true, + data: { task, workspace: null }, }; } catch (error) { return { success: false, error: - error instanceof Error ? error.message : "Failed to fetch task", - failedStep: "fetch_task", + error instanceof Error + ? error.message + : "Failed to resume Pi session", + failedStep: "pi_session", }; } } diff --git a/packages/harness/src/extensions/posthog-provider/gateway-auth.test.ts b/packages/harness/src/extensions/posthog-provider/gateway-auth.test.ts index 9264f9879c..f4afbdb219 100644 --- a/packages/harness/src/extensions/posthog-provider/gateway-auth.test.ts +++ b/packages/harness/src/extensions/posthog-provider/gateway-auth.test.ts @@ -23,6 +23,21 @@ describe("resolveGatewayAuth", () => { }); }); + it("uses an explicit gateway override", async () => { + const auth = await resolveGatewayAuth( + { + region: "us", + apiKey: "proxy-key", + baseUrl: "http://127.0.0.1:1234", + }, + fakeCtx(undefined), + ); + expect(auth).toEqual({ + baseUrl: "http://127.0.0.1:1234", + apiKey: "proxy-key", + }); + }); + it("falls back to the posthog provider's resolved api key", async () => { const auth = await resolveGatewayAuth( { region: "eu" }, diff --git a/packages/harness/src/extensions/posthog-provider/gateway-auth.ts b/packages/harness/src/extensions/posthog-provider/gateway-auth.ts index 3e37b16ef7..c431953e02 100644 --- a/packages/harness/src/extensions/posthog-provider/gateway-auth.ts +++ b/packages/harness/src/extensions/posthog-provider/gateway-auth.ts @@ -25,7 +25,7 @@ export async function resolveGatewayAuth( ctx: ExtensionContext, ): Promise { const region = resolveRegion(options.region); - const baseUrl = getLlmGatewayUrl(region); + const baseUrl = options.baseUrl ?? getLlmGatewayUrl(region); const apiKey = options.apiKey ?? diff --git a/packages/harness/src/extensions/posthog-provider/models.test.ts b/packages/harness/src/extensions/posthog-provider/models.test.ts index 048e088e8d..6b9067a6bd 100644 --- a/packages/harness/src/extensions/posthog-provider/models.test.ts +++ b/packages/harness/src/extensions/posthog-provider/models.test.ts @@ -21,6 +21,18 @@ describe("gatewayBaseUrlForApi", () => { getLlmGatewayUrl("dev"), ); }); + + it("routes models through an explicit gateway override", () => { + expect( + gatewayBaseUrlForApi("anthropic-messages", "us", "http://proxy/"), + ).toBe("http://proxy"); + expect(gatewayBaseUrlForApi("openai-responses", "us", "http://proxy")).toBe( + "http://proxy/v1", + ); + expect( + gatewayBaseUrlForApi("openai-responses", "us", "http://proxy/v1"), + ).toBe("http://proxy/v1"); + }); }); describe("resolveModelConfigs", () => { @@ -139,13 +151,37 @@ describe("resolveModelConfigs", () => { expect(opus?.input).toEqual(["text", "image"]); expect(opus?.contextWindow).toBe(500000); expect(opus?.compat).toEqual({ forceAdaptiveThinking: true }); + expect(opus?.thinkingLevelMap).toMatchObject({ + xhigh: "xhigh", + }); const gpt = configs.find((model) => model.id === "gpt-5.5"); expect(gpt?.api).toBe("openai-responses"); + expect(gpt?.thinkingLevelMap).toMatchObject({ + off: "none", + xhigh: "xhigh", + }); expect(gpt?.input).toEqual(["text"]); expect(gpt?.baseUrl).toBe(`${getLlmGatewayUrl("dev")}/v1`); }); + it("fetches models through an authenticated gateway override", async () => { + const fetchSpy = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ data: [] }), + }); + global.fetch = fetchSpy as unknown as typeof fetch; + + await resolveModelConfigs("us", "http://127.0.0.1:1234", "proxy-key"); + + expect(fetchSpy).toHaveBeenCalledWith( + "http://127.0.0.1:1234/v1/models", + expect.objectContaining({ + headers: { Authorization: "Bearer proxy-key" }, + }), + ); + }); + it("falls back to the model id as the display name and default context window", async () => { global.fetch = vi.fn().mockResolvedValue({ ok: true, diff --git a/packages/harness/src/extensions/posthog-provider/models.ts b/packages/harness/src/extensions/posthog-provider/models.ts index 3c96c3bfd5..493c56b0ed 100644 --- a/packages/harness/src/extensions/posthog-provider/models.ts +++ b/packages/harness/src/extensions/posthog-provider/models.ts @@ -1,3 +1,4 @@ +import { getBuiltinModels } from "@earendil-works/pi-ai/providers/all"; import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent"; import type { CloudRegion } from "@posthog/shared"; import { getLlmGatewayUrl } from "./gateway"; @@ -21,6 +22,19 @@ type ModelFamily = "anthropic" | "openai" | "cloudflare"; const ZERO_COST = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }; +function findBuiltinModel(family: ModelFamily, id: string) { + if (family === "cloudflare") { + return undefined; + } + + const builtins = + family === "openai" + ? getBuiltinModels("openai") + : getBuiltinModels("anthropic"); + + return builtins.find((model) => model.id === id); +} + function detectFamily(model: GatewayModel): ModelFamily { if (model.owned_by === "openai" || model.id.startsWith("gpt-")) { return "openai"; @@ -37,10 +51,17 @@ function detectFamily(model: GatewayModel): ModelFamily { * `/v1` surface; every other API this provider uses is served off the * product root. */ -export function gatewayBaseUrlForApi(api: string, region: CloudRegion): string { - return api === "openai-responses" - ? `${getLlmGatewayUrl(region)}/v1` - : getLlmGatewayUrl(region); +export function gatewayBaseUrlForApi( + api: string, + region: CloudRegion, + baseUrl = getLlmGatewayUrl(region), +): string { + const normalizedBaseUrl = baseUrl.replace(/\/$/, ""); + if (api !== "openai-responses" || normalizedBaseUrl.endsWith("/v1")) { + return normalizedBaseUrl; + } + + return `${normalizedBaseUrl}/v1`; } function toModelConfig( @@ -54,13 +75,19 @@ function toModelConfig( ? ["text", "image"] : ["text"]; + const builtin = findBuiltinModel(family, model.id); + const thinkingLevelMap = builtin?.thinkingLevelMap + ? { thinkingLevelMap: builtin.thinkingLevelMap } + : {}; + if (family === "openai") { return { id: model.id, name, api: "openai-responses", baseUrl: gatewayBaseUrlForApi("openai-responses", region), - reasoning: true, + reasoning: builtin?.reasoning ?? true, + ...thinkingLevelMap, input, cost: ZERO_COST, contextWindow, @@ -86,7 +113,8 @@ function toModelConfig( id: model.id, name, api: "anthropic-messages", - reasoning: true, + reasoning: builtin?.reasoning ?? true, + ...thinkingLevelMap, input, cost: ZERO_COST, contextWindow, @@ -184,13 +212,14 @@ export function fallbackModelConfigs( async function fetchGatewayModels( region: CloudRegion, + baseUrl = getLlmGatewayUrl(region), apiKey?: string, ): Promise { if (process.env.PI_OFFLINE || process.env.HARNESS_STATIC_MODELS) { return []; } try { - const response = await fetch(`${getLlmGatewayUrl(region)}/v1/models`, { + const response = await fetch(`${baseUrl.replace(/\/$/, "")}/v1/models`, { headers: apiKey ? { Authorization: `Bearer ${apiKey}` } : undefined, signal: AbortSignal.timeout(MODELS_FETCH_TIMEOUT_MS), }); @@ -206,9 +235,10 @@ async function fetchGatewayModels( export async function resolveModelConfigs( region: CloudRegion, + baseUrl?: string, apiKey?: string, ): Promise { - const live = await fetchGatewayModels(region, apiKey); + const live = await fetchGatewayModels(region, baseUrl, apiKey); if (live.length === 0) { return fallbackModelConfigs(region); } diff --git a/packages/harness/src/extensions/posthog-provider/provider.test.ts b/packages/harness/src/extensions/posthog-provider/provider.test.ts index ab5cc804a4..cdad9d7a14 100644 --- a/packages/harness/src/extensions/posthog-provider/provider.test.ts +++ b/packages/harness/src/extensions/posthog-provider/provider.test.ts @@ -51,6 +51,26 @@ describe("buildPosthogProvider", () => { expect(config.apiKey).toBe("pha_static"); }); + it("routes every provider model through an explicit gateway override", () => { + const config = buildPosthogProvider(models, { + region: "us", + apiKey: "proxy-key", + baseUrl: "http://127.0.0.1:1234", + }); + + expect(config.baseUrl).toBe("http://127.0.0.1:1234"); + expect( + (config.models ?? []) + .filter((model) => model.api === "anthropic-messages") + .every((model) => model.baseUrl === "http://127.0.0.1:1234"), + ).toBe(true); + expect( + (config.models ?? []) + .filter((model) => model.api === "openai-responses") + .every((model) => model.baseUrl === "http://127.0.0.1:1234/v1"), + ).toBe(true); + }); + it("omits apiKey entirely when none is provided", () => { const config = buildPosthogProvider(models); expect(config.apiKey).toBeUndefined(); @@ -134,6 +154,25 @@ describe("oauth.modifyModels", () => { expect(result?.[0]?.baseUrl).toBe(`${getLlmGatewayUrl("eu")}/v1`); }); + it("keeps an explicit gateway override when OAuth credentials have another region", () => { + const config = buildPosthogProvider(models, { + region: "us", + baseUrl: "http://127.0.0.1:1234", + }); + const runtimeModels = [ + makeModel({ id: "claude-opus-4-8", api: "anthropic-messages" }), + ]; + + const result = config.oauth?.modifyModels?.(runtimeModels, { + access: "a", + refresh: "r", + expires: 0, + region: "eu", + }); + + expect(result?.[0]?.baseUrl).toBe("http://127.0.0.1:1234"); + }); + it("leaves other providers' models untouched", () => { const config = buildPosthogProvider(models, { region: "us" }); const otherModel = makeModel({ id: "gpt-4o", provider: "openai" }); diff --git a/packages/harness/src/extensions/posthog-provider/provider.ts b/packages/harness/src/extensions/posthog-provider/provider.ts index 1620511118..07980c777b 100644 --- a/packages/harness/src/extensions/posthog-provider/provider.ts +++ b/packages/harness/src/extensions/posthog-provider/provider.ts @@ -1,5 +1,6 @@ import type { Api, Model, OAuthCredentials } from "@earendil-works/pi-ai"; import type { + AuthStorage, ProviderConfig, ProviderModelConfig, } from "@earendil-works/pi-coding-agent"; @@ -17,6 +18,29 @@ export const POSTHOG_PROVIDER_NAME = "posthog"; export interface PosthogProviderOptions { region?: CloudRegion; apiKey?: string; + baseUrl?: string; +} + +export type PosthogOAuthCredentials = Pick< + OAuthCredentials, + "access" | "refresh" | "expires" +> & { + region: CloudRegion; +}; + +export function parsePosthogOAuthCredentials( + serialized: string | undefined, +): PosthogOAuthCredentials | null { + return serialized + ? (JSON.parse(serialized) as PosthogOAuthCredentials) + : null; +} + +export function setPosthogOAuthCredentials( + storage: AuthStorage, + credentials: PosthogOAuthCredentials, +): void { + storage.set(POSTHOG_PROVIDER_NAME, { type: "oauth", ...credentials }); } /** @@ -30,12 +54,16 @@ function remapModelsToCredentialRegion( models: Model[], credentials: OAuthCredentials, fallbackRegion: CloudRegion, + baseUrl?: string, ): Model[] { const region = (credentials.region as CloudRegion | undefined) ?? fallbackRegion; return models.map((model) => model.provider === POSTHOG_PROVIDER_NAME - ? { ...model, baseUrl: gatewayBaseUrlForApi(model.api, region) } + ? { + ...model, + baseUrl: gatewayBaseUrlForApi(model.api, region, baseUrl), + } : model, ); } @@ -46,18 +74,32 @@ export function buildPosthogProvider( ): ProviderConfig { const region = resolveRegion(options.region); const explicitRegion = resolveExplicitRegion(options.region); + const baseUrl = options.baseUrl ?? getLlmGatewayUrl(region); + const routedModels = models.map((model) => ({ + ...model, + baseUrl: gatewayBaseUrlForApi( + model.api ?? "anthropic-messages", + region, + baseUrl, + ), + })); const config: ProviderConfig = { name: "PostHog", - baseUrl: getLlmGatewayUrl(region), + baseUrl, api: "anthropic-messages", - models, + models: routedModels, oauth: { name: "PostHog", login: (callbacks) => loginPosthog(callbacks, explicitRegion), refreshToken: (credentials) => refreshPosthog(region, credentials), getApiKey: (credentials) => String(credentials.access), modifyModels: (models, credentials) => - remapModelsToCredentialRegion(models, credentials, region), + remapModelsToCredentialRegion( + models, + credentials, + region, + options.baseUrl, + ), }, }; if (options.apiKey) { @@ -70,6 +112,10 @@ export async function resolvePosthogProvider( options: PosthogProviderOptions = {}, ): Promise { const region = resolveRegion(options.region); - const models = await resolveModelConfigs(region, options.apiKey); + const models = await resolveModelConfigs( + region, + options.baseUrl, + options.apiKey, + ); return buildPosthogProvider(models, options); } diff --git a/packages/harness/src/index.ts b/packages/harness/src/index.ts index 1e9865e3af..e731d7d3b9 100644 --- a/packages/harness/src/index.ts +++ b/packages/harness/src/index.ts @@ -1,5 +1,10 @@ import type { AgentSessionRuntime } from "@earendil-works/pi-coding-agent"; +export { + type PosthogOAuthCredentials, + parsePosthogOAuthCredentials, + setPosthogOAuthCredentials, +} from "./extensions/posthog-provider/provider"; export { createHarnessRuntime, type HarnessRuntimeOptions, diff --git a/packages/harness/src/runtime.test.ts b/packages/harness/src/runtime.test.ts index 8cee49f2e3..1411a67363 100644 --- a/packages/harness/src/runtime.test.ts +++ b/packages/harness/src/runtime.test.ts @@ -1,4 +1,5 @@ -import { mkdtemp, rm } from "node:fs/promises"; +import { existsSync } from "node:fs"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -57,4 +58,140 @@ describe("createHarnessRuntime", () => { await runtime.dispose(); } }); + + it("keeps desktop-provided OAuth credentials in memory without touching auth.json", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + posthogOAuthCredentials: { + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + region: "us", + }, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + expect(runtime.services.authStorage.get("posthog")).toMatchObject({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + }); + expect(existsSync(join(agentDir, "auth.json"))).toBe(false); + } finally { + await runtime.dispose(); + } + }); + + it("seeds the in-memory store from auth.json without writing back to it", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + const authPath = join(agentDir, "auth.json"); + const storedCredentials = { + anthropic: { type: "api_key", key: "anthropic-key" }, + posthog: { + type: "oauth", + access: "stale-access", + refresh: "stale-refresh", + expires: 0, + }, + }; + await writeFile(authPath, JSON.stringify(storedCredentials)); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + posthogOAuthCredentials: { + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + region: "us", + }, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + expect(runtime.services.authStorage.get("anthropic")).toMatchObject({ + type: "api_key", + key: "anthropic-key", + }); + expect(runtime.services.authStorage.get("posthog")).toMatchObject({ + access: "access-token", + refresh: "refresh-token", + }); + expect(JSON.parse(await readFile(authPath, "utf8"))).toEqual( + storedCredentials, + ); + } finally { + await runtime.dispose(); + } + }); + + it("uses a static provider key ahead of stored OAuth credentials", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + await writeFile( + join(agentDir, "auth.json"), + JSON.stringify({ + posthog: { + type: "oauth", + access: "stale-access", + refresh: "stale-refresh", + expires: 0, + }, + }), + ); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + apiKey: "proxy-key", + baseUrl: "http://127.0.0.1:1234", + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + await expect( + runtime.services.modelRegistry.getApiKeyForProvider("posthog"), + ).resolves.toBe("proxy-key"); + } finally { + await runtime.dispose(); + } + }); + + it("uses file-backed auth storage when no desktop credentials are provided", async () => { + vi.stubEnv("PI_OFFLINE", "1"); + const pi = await import("@earendil-works/pi-coding-agent"); + const cwd = await temporaryDirectory(); + const agentDir = await temporaryDirectory(); + + const runtime = await createHarnessRuntime({ + agentDir, + cwd, + sessionManager: pi.SessionManager.inMemory(cwd), + }); + + try { + runtime.services.authStorage.set("posthog", { + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }); + + expect(existsSync(join(agentDir, "auth.json"))).toBe(true); + } finally { + await runtime.dispose(); + } + }); }); diff --git a/packages/harness/src/runtime.ts b/packages/harness/src/runtime.ts index 9997ce5275..91b691c23d 100644 --- a/packages/harness/src/runtime.ts +++ b/packages/harness/src/runtime.ts @@ -1,17 +1,36 @@ +import { readFileSync } from "node:fs"; import { join } from "node:path"; import type { AgentSessionRuntime, + AuthStorage, CreateAgentSessionFromServicesOptions, CreateAgentSessionRuntimeFactory, CreateAgentSessionServicesOptions, } from "@earendil-works/pi-coding-agent"; import { installHogBrandEnv } from "./extensions/hog-branding/brand-env"; +import { + POSTHOG_PROVIDER_NAME, + type PosthogOAuthCredentials, + setPosthogOAuthCredentials, +} from "./extensions/posthog-provider/provider"; import type { HarnessExtensionOptions } from "./extensions/registry"; type PiRuntimeTarget = Parameters[0]; +type AuthStorageSnapshot = Parameters[0]; -export type HarnessRuntimeOptions = HarnessExtensionOptions & - Partial< +function loadAuthStorageSnapshot( + authPath: string, +): AuthStorageSnapshot | undefined { + try { + return JSON.parse(readFileSync(authPath, "utf8")) as AuthStorageSnapshot; + } catch { + return undefined; + } +} + +export type HarnessRuntimeOptions = HarnessExtensionOptions & { + posthogOAuthCredentials?: PosthogOAuthCredentials; +} & Partial< Pick< PiRuntimeTarget, "cwd" | "agentDir" | "sessionManager" | "sessionStartEvent" @@ -35,6 +54,7 @@ export type HarnessRuntimeOptions = HarnessExtensionOptions & export async function createHarnessRuntime( options: HarnessRuntimeOptions = {}, ): Promise { + const { posthogOAuthCredentials, ...runtimeOptions } = options; // Pi reads its application branding when the SDK is first evaluated. Keep // every runtime import below dynamic so this always happens first. installHogBrandEnv(); @@ -45,8 +65,8 @@ export async function createHarnessRuntime( import("./extensions/posthog-provider/models"), ]); - const cwd = options.cwd ?? process.cwd(); - const agentDir = options.agentDir ?? pi.getAgentDir(); + const cwd = runtimeOptions.cwd ?? process.cwd(); + const agentDir = runtimeOptions.agentDir ?? pi.getAgentDir(); const createRuntime: CreateAgentSessionRuntimeFactory = async ({ cwd: runtimeCwd, @@ -54,12 +74,21 @@ export async function createHarnessRuntime( sessionManager, sessionStartEvent, }) => { + const authPath = join(runtimeAgentDir, "auth.json"); const authStorage = - options.authStorage ?? - pi.AuthStorage.create(join(runtimeAgentDir, "auth.json")); + runtimeOptions.authStorage ?? + (posthogOAuthCredentials + ? pi.AuthStorage.inMemory(loadAuthStorageSnapshot(authPath)) + : pi.AuthStorage.create(authPath)); + if (posthogOAuthCredentials) { + setPosthogOAuthCredentials(authStorage, posthogOAuthCredentials); + } + if (options.apiKey) { + authStorage.setRuntimeApiKey(POSTHOG_PROVIDER_NAME, options.apiKey); + } const services = await pi.createAgentSessionServices({ - ...options, + ...runtimeOptions, cwd: runtimeCwd, agentDir: runtimeAgentDir, authStorage, @@ -69,9 +98,9 @@ export async function createHarnessRuntime( projectTrusted: false, }), resourceLoaderOptions: { - ...options.resourceLoaderOptions, + ...runtimeOptions.resourceLoaderOptions, extensionFactories: [ - ...(options.resourceLoaderOptions?.extensionFactories ?? []), + ...(runtimeOptions.resourceLoaderOptions?.extensionFactories ?? []), ...harnessExtensions(options), ], }, @@ -86,11 +115,11 @@ export async function createHarnessRuntime( .find((model) => model.provider === "posthog"); const created = await pi.createAgentSessionFromServices({ - ...options, + ...runtimeOptions, services, sessionManager, sessionStartEvent, - model: options.model ?? preferredModel ?? fallbackModel, + model: runtimeOptions.model ?? preferredModel ?? fallbackModel, }); return { @@ -109,12 +138,12 @@ export async function createHarnessRuntime( }; const sessionManager = - options.sessionManager ?? pi.SessionManager.create(cwd); + runtimeOptions.sessionManager ?? pi.SessionManager.create(cwd); return pi.createAgentSessionRuntime(createRuntime, { cwd: sessionManager.getCwd(), agentDir, sessionManager, - sessionStartEvent: options.sessionStartEvent, + sessionStartEvent: runtimeOptions.sessionStartEvent, }); } diff --git a/packages/host-router/src/router.ts b/packages/host-router/src/router.ts index 020ac1b210..a9582f4172 100644 --- a/packages/host-router/src/router.ts +++ b/packages/host-router/src/router.ts @@ -36,6 +36,7 @@ import { notificationRouter } from "./routers/notification.router"; import { oauthRouter } from "./routers/oauth.router"; import { onboardingImportRouter } from "./routers/onboarding-import.router"; import { osRouter } from "./routers/os.router"; +import { piSessionRouter } from "./routers/pi-session.router"; import { processTrackingRouter } from "./routers/process-tracking.router"; import { provisioningRouter } from "./routers/provisioning.router"; import { secureStoreRouter } from "./routers/secure-store.router"; @@ -88,6 +89,7 @@ export const hostRouter = router({ oauth: oauthRouter, onboardingImport: onboardingImportRouter, os: osRouter, + piSession: piSessionRouter, processTracking: processTrackingRouter, provisioning: provisioningRouter, secureStore: secureStoreRouter, diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts new file mode 100644 index 0000000000..67eba086d1 --- /dev/null +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -0,0 +1,63 @@ +import { publicProcedure, router } from "@posthog/host-trpc/trpc"; +import { PI_SESSION_SERVICE } from "@posthog/workspace-server/services/pi-session/identifiers"; +import type { PiSessionService } from "@posthog/workspace-server/services/pi-session/pi-session"; +import { + piSessionEntriesInput, + piSessionPromptInput, + piSessionStartOutput, + piSessionTranscriptInput, + resumePiSessionInput, + startPiSessionInput, +} from "@posthog/workspace-server/services/pi-session/schemas"; + +const getService = (container: { get(token: symbol): T }) => + container.get(PI_SESSION_SERVICE); + +export const piSessionRouter = router({ + start: publicProcedure + .input(startPiSessionInput) + .output(piSessionStartOutput) + .mutation(({ ctx, input }) => getService(ctx.container).start(input)), + + resume: publicProcedure + .input(resumePiSessionInput) + .mutation(({ ctx, input }) => getService(ctx.container).resume(input)), + + prompt: publicProcedure + .input(piSessionPromptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).prompt(input.taskId, input.prompt), + ), + + abort: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).abort(input.taskId), + ), + + stop: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + + status: publicProcedure + .input(piSessionTranscriptInput) + .query(({ ctx, input }) => getService(ctx.container).status(input.taskId)), + + entries: publicProcedure + .input(piSessionEntriesInput) + .query(({ ctx, input }) => + getService(ctx.container).entries(input.taskId, input.since), + ), + + onEvent: publicProcedure + .input(piSessionTranscriptInput) + .subscription(async function* (opts) { + const service = getService(opts.ctx.container); + const iterable = service.toIterable("event", { signal: opts.signal }); + for await (const payload of iterable) { + if (payload.taskId === opts.input.taskId) { + yield payload.event; + } + } + }), +}); diff --git a/packages/host-trpc/package.json b/packages/host-trpc/package.json index d567d05c36..e27a5b490a 100644 --- a/packages/host-trpc/package.json +++ b/packages/host-trpc/package.json @@ -15,7 +15,8 @@ "clean": "node ../../scripts/rimraf.mjs .turbo" }, "dependencies": { - "@trpc/server": "catalog:" + "@trpc/server": "catalog:", + "superjson": "catalog:" }, "peerDependencies": { "inversify": "catalog:" diff --git a/packages/host-trpc/src/trpc.ts b/packages/host-trpc/src/trpc.ts index d209a88580..2150e00df9 100644 --- a/packages/host-trpc/src/trpc.ts +++ b/packages/host-trpc/src/trpc.ts @@ -1,8 +1,10 @@ import { initTRPC } from "@trpc/server"; +import superjson from "superjson"; import type { HostContext } from "./context"; const t = initTRPC.context().create({ isServer: true, + transformer: superjson, }); export const router = t.router; diff --git a/packages/shared/src/agent-runtime.ts b/packages/shared/src/agent-runtime.ts new file mode 100644 index 0000000000..a85a4a2ab1 --- /dev/null +++ b/packages/shared/src/agent-runtime.ts @@ -0,0 +1,3 @@ +export const AGENT_RUNTIMES = ["acp", "pi"] as const; + +export type AgentRuntime = (typeof AGENT_RUNTIMES)[number]; diff --git a/packages/shared/src/domain-types.ts b/packages/shared/src/domain-types.ts index 98f9d90ff9..273bf9a87b 100644 --- a/packages/shared/src/domain-types.ts +++ b/packages/shared/src/domain-types.ts @@ -1,5 +1,6 @@ import { z } from "zod"; import type { Adapter } from "./adapter"; +import type { AgentRuntime } from "./agent-runtime"; import type { DismissalReasonOptionValue } from "./dismissal-reasons"; import type { StoredLogEntry } from "./session-events"; @@ -60,6 +61,7 @@ export interface Task { json_schema?: Record | null; signal_report?: string | null; internal?: boolean; + runtime?: AgentRuntime; /** Backend channel (tasks product Channel UUID) this task is owned by. */ channel?: string | null; latest_run?: TaskRun; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index cffe7a2b5d..3897ec0cd7 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,4 +1,5 @@ export * from "./adapter"; +export * from "./agent-runtime"; export * from "./analytics-events"; export { type ArchivedTask, archivedTaskSchema } from "./archive-domain"; export { withTimeout } from "./async"; diff --git a/packages/shared/src/task-creation-domain.ts b/packages/shared/src/task-creation-domain.ts index 8b073027a9..f4660ac8f2 100644 --- a/packages/shared/src/task-creation-domain.ts +++ b/packages/shared/src/task-creation-domain.ts @@ -1,4 +1,5 @@ import type { Adapter } from "./adapter"; +import type { AgentRuntime } from "./agent-runtime"; import type { CloudRunSource, PrAuthorshipMode } from "./cloud"; import type { Task } from "./domain-types"; import type { ExecutionMode } from "./exec-types"; @@ -35,6 +36,7 @@ export interface TaskCreationInput { githubUserIntegrationId?: string; executionMode?: ExecutionMode; adapter?: Adapter; + runtime?: AgentRuntime; model?: string; reasoningLevel?: string; environmentId?: string; diff --git a/packages/shared/src/task.ts b/packages/shared/src/task.ts index 585d44bb1c..58f9e7975f 100644 --- a/packages/shared/src/task.ts +++ b/packages/shared/src/task.ts @@ -1,4 +1,5 @@ // PostHog Task model (matches PostHog Code's OpenAPI schema) +import type { AgentRuntime } from "./agent-runtime"; import type { UploadableSkillSource } from "./skills"; export interface Task { @@ -21,6 +22,7 @@ export interface Task { repository: string; // Format: "organization/repository" (e.g., "posthog/posthog-js") json_schema?: Record | null; // JSON schema for task output validation internal?: boolean; + runtime?: AgentRuntime; created_at: string; updated_at: string; created_by?: { diff --git a/packages/ui/package.json b/packages/ui/package.json index a93ec5c3a5..ef0cfd7514 100644 --- a/packages/ui/package.json +++ b/packages/ui/package.json @@ -18,6 +18,7 @@ }, "dependencies": { "@agentclientprotocol/sdk": "0.22.1", + "@earendil-works/pi-coding-agent": "catalog:", "@base-ui/react": "^1.3.0", "@codemirror/lang-angular": "^0.1.4", "@codemirror/lang-cpp": "^6.0.3", diff --git a/packages/ui/src/features/canvas/components/ChannelIntro.tsx b/packages/ui/src/features/canvas/components/ChannelIntro.tsx index c39d3d4e1e..75cd9583ba 100644 --- a/packages/ui/src/features/canvas/components/ChannelIntro.tsx +++ b/packages/ui/src/features/canvas/components/ChannelIntro.tsx @@ -63,7 +63,7 @@ export function ChannelIntro({
{contextMdState === "created" && ( - + @@ -76,7 +76,7 @@ export function ChannelIntro({ )} {contextMdState === "building" && ( - + diff --git a/packages/ui/src/features/canvas/components/CreateChannelModal.tsx b/packages/ui/src/features/canvas/components/CreateChannelModal.tsx index 9c80bb03e3..a325e7a933 100644 --- a/packages/ui/src/features/canvas/components/CreateChannelModal.tsx +++ b/packages/ui/src/features/canvas/components/CreateChannelModal.tsx @@ -271,7 +271,7 @@ export function CreateChannelModal({ Create a channel - + Name ; + +function PiMessageView({ message }: { message: PiMessage }) { + if (message.role === "bashExecution") { + return <>{message.output}; + } + if ( + message.role === "branchSummary" || + message.role === "compactionSummary" + ) { + return <>{message.summary}; + } + if ("content" in message) { + return <>{messageContentText(message)}; + } + return null; +} + +function messageContentText(message: PiMessageWithContent): string { + if (typeof message.content === "string") { + return message.content; + } + + return message.content + .flatMap((part) => (part.type === "text" ? [part.text] : [])) + .join("\n"); +} + +function messageBubbleClass(role: string): string { + if (role === "user") { + return "mb-3 ml-auto max-w-[80%] rounded-lg bg-accent-3 p-3 text-sm"; + } + return "mb-3 max-w-[80%] whitespace-pre-wrap rounded-lg bg-gray-3 p-3 text-sm"; +} + +export function PiSessionView({ taskId }: PiSessionViewProps) { + const trpc = useHostTRPC(); + const client = useHostTRPCClient(); + const { error: ensureError, isSuccess: sessionReady } = + useEnsurePiSession(taskId); + + const [prompt, setPrompt] = useState(""); + const [liveFeed, setLiveFeed] = useState(emptyLiveFeed); + const [syncedEntries, setSyncedEntries] = useState( + undefined, + ); + + const { data: fetchedEntries } = useQuery({ + ...trpc.piSession.entries.queryOptions({ taskId }), + enabled: sessionReady, + }); + const { error: statusError } = useQuery({ + ...trpc.piSession.status.queryOptions({ taskId }), + enabled: sessionReady, + }); + + const syncer = useMemo( + () => + new PiEntriesSyncer( + (since) => client.piSession.entries.query({ taskId, since }), + setSyncedEntries, + ), + [client, taskId], + ); + + useEffect(() => { + syncer.seed(fetchedEntries); + }, [syncer, fetchedEntries]); + + const history = syncedEntries ?? fetchedEntries; + + useSubscription( + trpc.piSession.onEvent.subscriptionOptions( + { taskId }, + { + enabled: sessionReady, + onData: (event: PiEvent) => { + setLiveFeed((feed) => applyPiEvent(feed, event)); + + if (event.type === "agent_settled") { + void syncer.sync().then(() => setLiveFeed(emptyLiveFeed)); + } + }, + }, + ), + ); + + const send = async () => { + const text = prompt.trim(); + if (!text) { + return; + } + await client.piSession.prompt.mutate({ taskId, prompt: text }); + setPrompt(""); + }; + + const sessionError = ensureError ?? statusError; + if (sessionError) { + return ( + + + Pi session failed to start + {sessionError.message} + + + ); + } + + if (!sessionReady) { + return ( + + + Starting Pi session… + + + ); + } + + return ( +
+
+ {history?.entries.map((entry) => { + if (entry.type !== "message") { + return null; + } + + return ( +
+ +
+ ); + })} + {liveFeed.liveMessages.map((message) => ( +
+ +
+ ))} + {liveFeed.streamingMessage ? ( +
+ +
+ ) : null} +
+
+