Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions apps/code/electron-builder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
],

Expand Down
2 changes: 2 additions & 0 deletions apps/code/electron.vite.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import {
copyCodexAcpBinaries,
copyDrizzleMigrations,
copyEnricherGrammars,
copyPiRpcHost,
copyPosthogPlugin,
fixFilenameCircularRef,
getBuildDate,
Expand Down Expand Up @@ -93,6 +94,7 @@ export default defineConfig(({ mode }) => {
autoServicesPlugin(path.join(__dirname, "src/main/services")),
fixFilenameCircularRef(),
copyClaudeExecutable(),
copyPiRpcHost(),
copyPosthogPlugin(isDev),
copyDrizzleMigrations(),
copyCodexAcpBinaries(),
Expand Down
1 change: 1 addition & 0 deletions apps/code/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
5 changes: 5 additions & 0 deletions apps/code/src/main/di/bindings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down
35 changes: 30 additions & 5 deletions apps/code/src/main/di/container.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -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<RootLogger>(ROOT_LOGGER).toConstantValue(logger);
Expand Down Expand Up @@ -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<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
cancelSessionsByTaskId: async (taskId: string) => {
await Promise.all([
ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
]);
},
}));
container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
stopWatching: async (worktreePath: string) => {
Expand All @@ -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<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
cancelSessionsByTaskId: async (taskId: string) => {
await Promise.all([
ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
]);
},
}));
container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({
stopWatching: async (worktreePath: string) => {
Expand Down Expand Up @@ -685,7 +705,12 @@ container.load(workspaceModule);
container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => {
const agent = ctx.get<AgentService>(AGENT_SERVICE);
return {
cancelSessionsByTaskId: (taskId) => agent.cancelSessionsByTaskId(taskId),
cancelSessionsByTaskId: async (taskId) => {
await Promise.all([
agent.cancelSessionsByTaskId(taskId),
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
]);
},
onAgentFileActivity: (handler) =>
agent.on(AgentServiceEvent.AgentFileActivity, handler),
};
Expand Down
Original file line number Diff line number Diff line change
@@ -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",
},
});
});
});
Original file line number Diff line number Diff line change
@@ -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<string>;

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<PiRpcClient> {
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,
Comment thread
veria-ai[bot] marked this conversation as resolved.
},
});
}

private getProxyUrl(region: CloudRegion): Promise<string> {
if (this.proxyRegion !== region || !this.proxyUrlPromise) {
this.proxyRegion = region;
const gatewayUrl = getLlmGatewayUrl(getCloudUrlFromRegion(region));
this.proxyUrlPromise = this.authProxy.start(gatewayUrl);
}

return this.proxyUrlPromise;
}
}
2 changes: 2 additions & 0 deletions apps/code/src/main/trpc/router.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -97,6 +98,7 @@ export const trpcRouter = router({
onboardingImport: onboardingImportRouter,
logs: logsRouter,
os: osRouter,
piSession: piSessionRouter,
processTracking: processTrackingRouter,
provisioning: provisioningRouter,
sleep: sleepRouter,
Expand Down
3 changes: 3 additions & 0 deletions apps/code/src/renderer/di/bindings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions apps/code/src/renderer/di/container.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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";

Expand Down Expand Up @@ -290,6 +293,7 @@ container

// Bind services
container.bind<ITaskCreationHost>(TASK_CREATION_HOST).to(TrpcTaskCreationHost);
container.bind<PiRunner>(PI_RUNNER).to(TrpcPiRunner);
container.bind(TASK_CREATION_EFFECTS).toConstantValue(taskCreationEffects);
container.bind<TaskService>(RENDERER_TASK_SERVICE).to(TaskService);
container.bind<TaskService>(TASK_SERVICE).toService(RENDERER_TASK_SERVICE);
Expand Down
28 changes: 28 additions & 0 deletions apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts
Original file line number Diff line number Diff line change
@@ -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<HostTrpcClient>(HOST_TRPC_CLIENT);
}

export class TrpcPiRunner implements PiRunner {
async create(input: PiRunInput): Promise<void> {
await hostClient().piSession.start.mutate(input);
}

resume(input: PiResumeInput): Promise<void> {
return hostClient().piSession.resume.mutate(input);
}

stop(taskId: string): Promise<void> {
return hostClient().piSession.stop.mutate({ taskId });
}
}
8 changes: 6 additions & 2 deletions apps/code/src/renderer/trpc/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<TrpcRouter>({
links: [ipcInstrumentationLink<TrpcRouter>(), ipcLink()],
links: [
ipcInstrumentationLink<TrpcRouter>(),
ipcLink({ transformer: superjson }),
],
});

export const hostTrpcClient = createTRPCClient<HostRouter>({
links: [ipcLink()],
links: [ipcLink({ transformer: superjson })],
});

const context = createTRPCContext<TrpcRouter>();
Expand Down
23 changes: 23 additions & 0 deletions apps/code/vite-main-plugins.mts
Original file line number Diff line number Diff line change
Expand Up @@ -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"));
Comment thread
jonathanlab marked this conversation as resolved.
},
};
}

export function copyClaudeExecutable(): Plugin {
return {
name: "copy-claude-executable",
Expand Down
1 change: 1 addition & 0 deletions apps/web/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading
Loading