Skip to content

Commit 38a0c40

Browse files
jonathanlabjonathanlab
committed
feat: hook up pi runtime to ChatThread
1 parent 9485f63 commit 38a0c40

44 files changed

Lines changed: 2849 additions & 315 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎apps/code/src/main/di/container.ts‎

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,17 @@ import {
318318
WORKTREE_REPOSITORY as MAIN_WORKTREE_REPOSITORY,
319319
} from "./tokens";
320320

321+
async function cancelTaskSessions(
322+
agentService: AgentService,
323+
piSessionService: PiSessionService,
324+
taskId: string,
325+
): Promise<void> {
326+
await Promise.all([
327+
agentService.cancelSessionsByTaskId(taskId),
328+
piSessionService.stop(taskId),
329+
]);
330+
}
331+
321332
export const container = new TypedContainer<MainBindings>({
322333
defaultScope: "Singleton",
323334
});
@@ -395,12 +406,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => {
395406
});
396407
container.load(archiveModule);
397408
container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({
398-
cancelSessionsByTaskId: async (taskId: string) => {
399-
await Promise.all([
400-
ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
401-
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
402-
]);
403-
},
409+
cancelSessionsByTaskId: (taskId: string) =>
410+
cancelTaskSessions(
411+
ctx.get<AgentService>(AGENT_SERVICE),
412+
ctx.get<PiSessionService>(PI_SESSION_SERVICE),
413+
taskId,
414+
),
404415
}));
405416
container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
406417
stopWatching: async (worktreePath: string) => {
@@ -411,12 +422,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({
411422
}));
412423
container.load(suspensionModule);
413424
container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({
414-
cancelSessionsByTaskId: async (taskId: string) => {
415-
await Promise.all([
416-
ctx.get<AgentService>(AGENT_SERVICE).cancelSessionsByTaskId(taskId),
417-
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
418-
]);
419-
},
425+
cancelSessionsByTaskId: (taskId: string) =>
426+
cancelTaskSessions(
427+
ctx.get<AgentService>(AGENT_SERVICE),
428+
ctx.get<PiSessionService>(PI_SESSION_SERVICE),
429+
taskId,
430+
),
420431
}));
421432
container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({
422433
stopWatching: async (worktreePath: string) => {
@@ -687,12 +698,12 @@ container.load(workspaceModule);
687698
container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => {
688699
const agent = ctx.get<AgentService>(AGENT_SERVICE);
689700
return {
690-
cancelSessionsByTaskId: async (taskId) => {
691-
await Promise.all([
692-
agent.cancelSessionsByTaskId(taskId),
693-
ctx.get<PiSessionService>(PI_SESSION_SERVICE).stop(taskId),
694-
]);
695-
},
701+
cancelSessionsByTaskId: (taskId) =>
702+
cancelTaskSessions(
703+
agent,
704+
ctx.get<PiSessionService>(PI_SESSION_SERVICE),
705+
taskId,
706+
),
696707
onAgentFileActivity: (handler) =>
697708
agent.on(AgentServiceEvent.AgentFileActivity, handler),
698709
};

‎apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,26 +3,27 @@ import type {
33
PiRunInput,
44
PiRunner,
55
} from "@posthog/core/pi-runtime/piRunner";
6-
import { resolveService } from "@posthog/di/container";
76
import {
87
HOST_TRPC_CLIENT,
98
type HostTrpcClient,
109
} from "@posthog/host-router/client";
10+
import { inject, injectable } from "inversify";
1111

12-
function hostClient(): HostTrpcClient {
13-
return resolveService<HostTrpcClient>(HOST_TRPC_CLIENT);
14-
}
15-
12+
@injectable()
1613
export class TrpcPiRunner implements PiRunner {
14+
constructor(
15+
@inject(HOST_TRPC_CLIENT) private readonly hostClient: HostTrpcClient,
16+
) {}
17+
1718
async create(input: PiRunInput): Promise<void> {
18-
await hostClient().piSession.start.mutate(input);
19+
await this.hostClient.piSession.start.mutate(input);
1920
}
2021

2122
resume(input: PiResumeInput): Promise<void> {
22-
return hostClient().piSession.resume.mutate(input);
23+
return this.hostClient.piSession.resume.mutate(input);
2324
}
2425

2526
stop(taskId: string): Promise<void> {
26-
return hostClient().piSession.stop.mutate({ taskId });
27+
return this.hostClient.piSession.stop.mutate({ taskId });
2728
}
2829
}

‎packages/core/package.json‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
"clean": "node ../../scripts/rimraf.mjs .turbo"
1717
},
1818
"dependencies": {
19+
"@earendil-works/pi-ai": "catalog:",
20+
"@earendil-works/pi-coding-agent": "catalog:",
1921
"@json-render/core": "^0.19.0",
2022
"@modelcontextprotocol/ext-apps": "^1.1.2",
2123
"@modelcontextprotocol/sdk": "^1.12.1",

‎packages/core/src/pi-runtime/conversation/toolKind.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import type { AgentToolKind } from "@posthog/shared";
2+
3+
export type PiToolName =
4+
| "read"
5+
| "bash"
6+
| "edit"
7+
| "write"
8+
| "grep"
9+
| "find"
10+
| "ls";
11+
12+
export const TOOL_KIND_BY_NAME: Record<PiToolName, AgentToolKind> = {
13+
read: "read",
14+
edit: "edit",
15+
write: "edit",
16+
bash: "execute",
17+
grep: "search",
18+
find: "search",
19+
ls: "read",
20+
};

‎packages/core/src/pi-runtime/conversation/toolTranslator.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
2+
import type {
3+
AgentToolCallContent,
4+
AgentToolCallLocation,
5+
} from "@posthog/shared";
6+
7+
export interface PiToolTranslatorInput {
8+
toolCallId: string;
9+
arguments: unknown;
10+
resultContent?: (TextContent | ImageContent)[];
11+
details?: unknown;
12+
isError?: boolean;
13+
}
14+
15+
export interface PiToolTranslatorOutput {
16+
locations?: AgentToolCallLocation[];
17+
content?: AgentToolCallContent[];
18+
}
19+
20+
export type PiToolTranslator = (
21+
input: PiToolTranslatorInput,
22+
) => PiToolTranslatorOutput;

‎packages/core/src/pi-runtime/conversation/tools/bashTranslator.test.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import { describe, expect, it } from "vitest";
2+
import { bashTranslator } from "./bashTranslator";
3+
4+
describe("bashTranslator", () => {
5+
it("surfaces command output as text content on success", () => {
6+
const result = bashTranslator({
7+
toolCallId: "tool-1",
8+
arguments: { command: "echo hi" },
9+
resultContent: [{ type: "text", text: "hi\n" }],
10+
details: { truncation: undefined, fullOutputPath: undefined },
11+
isError: false,
12+
});
13+
14+
expect(result).toEqual({
15+
content: [
16+
{
17+
type: "content",
18+
content: { type: "text", text: "hi\n" },
19+
},
20+
],
21+
});
22+
});
23+
24+
it("surfaces stderr-style output for a failed command", () => {
25+
const result = bashTranslator({
26+
toolCallId: "tool-2",
27+
arguments: { command: "false" },
28+
resultContent: [
29+
{ type: "text", text: "command failed with exit code 1" },
30+
],
31+
isError: true,
32+
});
33+
34+
expect(result).toEqual({
35+
content: [
36+
{
37+
type: "content",
38+
content: { type: "text", text: "command failed with exit code 1" },
39+
},
40+
],
41+
});
42+
});
43+
44+
it("returns no content when the result has no text blocks", () => {
45+
const result = bashTranslator({
46+
toolCallId: "tool-3",
47+
arguments: { command: "echo hi" },
48+
resultContent: [],
49+
});
50+
51+
expect(result).toEqual({});
52+
});
53+
54+
it("joins multiple text blocks and ignores image blocks", () => {
55+
const result = bashTranslator({
56+
toolCallId: "tool-4",
57+
arguments: { command: "cat file.txt" },
58+
resultContent: [
59+
{ type: "text", text: "line one\n" },
60+
{ type: "image", data: "abc", mimeType: "image/png" },
61+
{ type: "text", text: "line two\n" },
62+
],
63+
});
64+
65+
expect(result).toEqual({
66+
content: [
67+
{
68+
type: "content",
69+
content: { type: "text", text: "line one\nline two\n" },
70+
},
71+
],
72+
});
73+
});
74+
75+
it("returns no content when resultContent is missing", () => {
76+
const result = bashTranslator({
77+
toolCallId: "tool-5",
78+
arguments: { command: "echo hi" },
79+
});
80+
81+
expect(result).toEqual({});
82+
});
83+
});

‎packages/core/src/pi-runtime/conversation/tools/bashTranslator.ts‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import type { AgentToolCallContent } from "@posthog/shared";
2+
import type { PiToolTranslator } from "../toolTranslator";
3+
4+
export const bashTranslator: PiToolTranslator = ({ resultContent }) => {
5+
const outputText = resultContent
6+
?.filter((block) => block.type === "text")
7+
.map((block) => block.text)
8+
.join("");
9+
10+
if (!outputText) {
11+
return {};
12+
}
13+
14+
const content: AgentToolCallContent[] = [
15+
{
16+
type: "content",
17+
content: { type: "text", text: outputText },
18+
},
19+
];
20+
21+
return { content };
22+
};

‎packages/core/src/pi-runtime/conversation/tools/editTranslator.test.ts‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import { describe, expect, it } from "vitest";
2+
import { editTranslator } from "./editTranslator";
3+
4+
describe("editTranslator", () => {
5+
it("produces a diff from the edit arguments and a location for the path", () => {
6+
const output = editTranslator({
7+
toolCallId: "1",
8+
arguments: {
9+
path: "src/foo.ts",
10+
edits: [{ oldText: "const a = 1;", newText: "const a = 2;" }],
11+
},
12+
details: {
13+
diff: "- const a = 1;\n+ const a = 2;",
14+
patch: "@@ -1 +1 @@\n-const a = 1;\n+const a = 2;",
15+
firstChangedLine: 3,
16+
},
17+
});
18+
19+
expect(output.locations).toEqual([{ path: "src/foo.ts", line: 3 }]);
20+
expect(output.content).toEqual([
21+
{
22+
type: "diff",
23+
path: "src/foo.ts",
24+
oldText: "const a = 1;",
25+
newText: "const a = 2;",
26+
},
27+
]);
28+
});
29+
30+
it("joins multiple edits into a single diff", () => {
31+
const output = editTranslator({
32+
toolCallId: "1",
33+
arguments: {
34+
path: "src/foo.ts",
35+
edits: [
36+
{ oldText: "const a = 1;", newText: "const a = 2;" },
37+
{ oldText: "const b = 1;", newText: "const b = 2;" },
38+
],
39+
},
40+
details: {
41+
diff: "irrelevant",
42+
patch: "irrelevant",
43+
},
44+
});
45+
46+
expect(output.content).toEqual([
47+
{
48+
type: "diff",
49+
path: "src/foo.ts",
50+
oldText: "const a = 1;\nconst b = 1;",
51+
newText: "const a = 2;\nconst b = 2;",
52+
},
53+
]);
54+
});
55+
56+
it("falls back to the details diff string when edits are missing", () => {
57+
const output = editTranslator({
58+
toolCallId: "1",
59+
arguments: { path: "src/foo.ts" },
60+
details: {
61+
diff: "- const a = 1;\n+ const a = 2;",
62+
patch: "@@ -1 +1 @@",
63+
},
64+
});
65+
66+
expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]);
67+
expect(output.content).toEqual([
68+
{
69+
type: "diff",
70+
path: "src/foo.ts",
71+
newText: "- const a = 1;\n+ const a = 2;",
72+
},
73+
]);
74+
});
75+
76+
it("falls back to result text content on error with no details", () => {
77+
const output = editTranslator({
78+
toolCallId: "1",
79+
arguments: { path: "src/foo.ts" },
80+
resultContent: [{ type: "text", text: "permission denied" }],
81+
isError: true,
82+
});
83+
84+
expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]);
85+
expect(output.content).toEqual([
86+
{
87+
type: "content",
88+
content: { type: "text", text: "permission denied" },
89+
},
90+
]);
91+
});
92+
93+
it("returns no locations or content when arguments are missing entirely", () => {
94+
const output = editTranslator({
95+
toolCallId: "1",
96+
arguments: undefined,
97+
});
98+
99+
expect(output.locations).toBeUndefined();
100+
expect(output.content).toBeUndefined();
101+
});
102+
});

0 commit comments

Comments
 (0)