diff --git a/apps/code/src/main/di/container.ts b/apps/code/src/main/di/container.ts index c121b765d8..15056cb1ee 100644 --- a/apps/code/src/main/di/container.ts +++ b/apps/code/src/main/di/container.ts @@ -318,6 +318,17 @@ import { WORKTREE_REPOSITORY as MAIN_WORKTREE_REPOSITORY, } from "./tokens"; +async function cancelTaskSessions( + agentService: AgentService, + piSessionService: PiSessionService, + taskId: string, +): Promise { + await Promise.all([ + agentService.cancelSessionsByTaskId(taskId), + piSessionService.stop(taskId), + ]); +} + export const container = new TypedContainer({ defaultScope: "Singleton", }); @@ -395,12 +406,12 @@ container.bind(MCP_PROXY_AUTH).toDynamicValue((ctx) => { }); container.load(archiveModule); container.bind(ARCHIVE_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: async (taskId: string) => { - await Promise.all([ - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId: string) => + cancelTaskSessions( + ctx.get(AGENT_SERVICE), + ctx.get(PI_SESSION_SERVICE), + taskId, + ), })); container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -411,12 +422,12 @@ container.bind(ARCHIVE_FILE_WATCHER).toDynamicValue((ctx) => ({ })); container.load(suspensionModule); container.bind(SUSPENSION_SESSION_CANCELLER).toDynamicValue((ctx) => ({ - cancelSessionsByTaskId: async (taskId: string) => { - await Promise.all([ - ctx.get(AGENT_SERVICE).cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId: string) => + cancelTaskSessions( + ctx.get(AGENT_SERVICE), + ctx.get(PI_SESSION_SERVICE), + taskId, + ), })); container.bind(SUSPENSION_FILE_WATCHER).toDynamicValue((ctx) => ({ stopWatching: async (worktreePath: string) => { @@ -687,12 +698,12 @@ container.load(workspaceModule); container.bind(WORKSPACE_AGENT).toDynamicValue((ctx): WorkspaceAgent => { const agent = ctx.get(AGENT_SERVICE); return { - cancelSessionsByTaskId: async (taskId) => { - await Promise.all([ - agent.cancelSessionsByTaskId(taskId), - ctx.get(PI_SESSION_SERVICE).stop(taskId), - ]); - }, + cancelSessionsByTaskId: (taskId) => + cancelTaskSessions( + agent, + ctx.get(PI_SESSION_SERVICE), + taskId, + ), onAgentFileActivity: (handler) => agent.on(AgentServiceEvent.AgentFileActivity, handler), }; diff --git a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts index 3f4311f7fb..438f0b485a 100644 --- a/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts +++ b/apps/code/src/renderer/platform-adapters/trpc-pi-runner.ts @@ -3,26 +3,27 @@ import type { 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"; +import { inject, injectable } from "inversify"; -function hostClient(): HostTrpcClient { - return resolveService(HOST_TRPC_CLIENT); -} - +@injectable() export class TrpcPiRunner implements PiRunner { + constructor( + @inject(HOST_TRPC_CLIENT) private readonly hostClient: HostTrpcClient, + ) {} + async create(input: PiRunInput): Promise { - await hostClient().piSession.start.mutate(input); + await this.hostClient.piSession.start.mutate(input); } resume(input: PiResumeInput): Promise { - return hostClient().piSession.resume.mutate(input); + return this.hostClient.piSession.resume.mutate(input); } stop(taskId: string): Promise { - return hostClient().piSession.stop.mutate({ taskId }); + return this.hostClient.piSession.stop.mutate({ taskId }); } } diff --git a/packages/core/package.json b/packages/core/package.json index 32d069cbae..aa8bc9ff19 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -16,6 +16,8 @@ "clean": "node ../../scripts/rimraf.mjs .turbo" }, "dependencies": { + "@earendil-works/pi-ai": "catalog:", + "@earendil-works/pi-coding-agent": "catalog:", "@json-render/core": "^0.19.0", "@modelcontextprotocol/ext-apps": "^1.1.2", "@modelcontextprotocol/sdk": "^1.12.1", diff --git a/packages/core/src/pi-runtime/conversation/toolKind.ts b/packages/core/src/pi-runtime/conversation/toolKind.ts new file mode 100644 index 0000000000..d39c7b1e26 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/toolKind.ts @@ -0,0 +1,20 @@ +import type { AgentToolKind } from "@posthog/shared"; + +export type PiToolName = + | "read" + | "bash" + | "edit" + | "write" + | "grep" + | "find" + | "ls"; + +export const TOOL_KIND_BY_NAME: Record = { + read: "read", + edit: "edit", + write: "edit", + bash: "execute", + grep: "search", + find: "search", + ls: "read", +}; diff --git a/packages/core/src/pi-runtime/conversation/toolTranslator.ts b/packages/core/src/pi-runtime/conversation/toolTranslator.ts new file mode 100644 index 0000000000..be9017193f --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/toolTranslator.ts @@ -0,0 +1,22 @@ +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; + +export interface PiToolTranslatorInput { + toolCallId: string; + arguments: unknown; + resultContent?: (TextContent | ImageContent)[]; + details?: unknown; + isError?: boolean; +} + +export interface PiToolTranslatorOutput { + locations?: AgentToolCallLocation[]; + content?: AgentToolCallContent[]; +} + +export type PiToolTranslator = ( + input: PiToolTranslatorInput, +) => PiToolTranslatorOutput; diff --git a/packages/core/src/pi-runtime/conversation/tools/bashTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/bashTranslator.test.ts new file mode 100644 index 0000000000..8b58be6381 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/bashTranslator.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, it } from "vitest"; +import { bashTranslator } from "./bashTranslator"; + +describe("bashTranslator", () => { + it("surfaces command output as text content on success", () => { + const result = bashTranslator({ + toolCallId: "tool-1", + arguments: { command: "echo hi" }, + resultContent: [{ type: "text", text: "hi\n" }], + details: { truncation: undefined, fullOutputPath: undefined }, + isError: false, + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "hi\n" }, + }, + ], + }); + }); + + it("surfaces stderr-style output for a failed command", () => { + const result = bashTranslator({ + toolCallId: "tool-2", + arguments: { command: "false" }, + resultContent: [ + { type: "text", text: "command failed with exit code 1" }, + ], + isError: true, + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "command failed with exit code 1" }, + }, + ], + }); + }); + + it("returns no content when the result has no text blocks", () => { + const result = bashTranslator({ + toolCallId: "tool-3", + arguments: { command: "echo hi" }, + resultContent: [], + }); + + expect(result).toEqual({}); + }); + + it("joins multiple text blocks and ignores image blocks", () => { + const result = bashTranslator({ + toolCallId: "tool-4", + arguments: { command: "cat file.txt" }, + resultContent: [ + { type: "text", text: "line one\n" }, + { type: "image", data: "abc", mimeType: "image/png" }, + { type: "text", text: "line two\n" }, + ], + }); + + expect(result).toEqual({ + content: [ + { + type: "content", + content: { type: "text", text: "line one\nline two\n" }, + }, + ], + }); + }); + + it("returns no content when resultContent is missing", () => { + const result = bashTranslator({ + toolCallId: "tool-5", + arguments: { command: "echo hi" }, + }); + + expect(result).toEqual({}); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/bashTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/bashTranslator.ts new file mode 100644 index 0000000000..694dc0de16 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/bashTranslator.ts @@ -0,0 +1,22 @@ +import type { AgentToolCallContent } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const bashTranslator: PiToolTranslator = ({ resultContent }) => { + const outputText = resultContent + ?.filter((block) => block.type === "text") + .map((block) => block.text) + .join(""); + + if (!outputText) { + return {}; + } + + const content: AgentToolCallContent[] = [ + { + type: "content", + content: { type: "text", text: outputText }, + }, + ]; + + return { content }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/editTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/editTranslator.test.ts new file mode 100644 index 0000000000..75f0e4a38f --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/editTranslator.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from "vitest"; +import { editTranslator } from "./editTranslator"; + +describe("editTranslator", () => { + it("produces a diff from the edit arguments and a location for the path", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { + path: "src/foo.ts", + edits: [{ oldText: "const a = 1;", newText: "const a = 2;" }], + }, + details: { + diff: "- const a = 1;\n+ const a = 2;", + patch: "@@ -1 +1 @@\n-const a = 1;\n+const a = 2;", + firstChangedLine: 3, + }, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: 3 }]); + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: "const a = 1;", + newText: "const a = 2;", + }, + ]); + }); + + it("joins multiple edits into a single diff", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { + path: "src/foo.ts", + edits: [ + { oldText: "const a = 1;", newText: "const a = 2;" }, + { oldText: "const b = 1;", newText: "const b = 2;" }, + ], + }, + details: { + diff: "irrelevant", + patch: "irrelevant", + }, + }); + + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: "const a = 1;\nconst b = 1;", + newText: "const a = 2;\nconst b = 2;", + }, + ]); + }); + + it("falls back to the details diff string when edits are missing", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { path: "src/foo.ts" }, + details: { + diff: "- const a = 1;\n+ const a = 2;", + patch: "@@ -1 +1 @@", + }, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]); + expect(output.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + newText: "- const a = 1;\n+ const a = 2;", + }, + ]); + }); + + it("falls back to result text content on error with no details", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: { path: "src/foo.ts" }, + resultContent: [{ type: "text", text: "permission denied" }], + isError: true, + }); + + expect(output.locations).toEqual([{ path: "src/foo.ts", line: undefined }]); + expect(output.content).toEqual([ + { + type: "content", + content: { type: "text", text: "permission denied" }, + }, + ]); + }); + + it("returns no locations or content when arguments are missing entirely", () => { + const output = editTranslator({ + toolCallId: "1", + arguments: undefined, + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/editTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/editTranslator.ts new file mode 100644 index 0000000000..1b717624d0 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/editTranslator.ts @@ -0,0 +1,55 @@ +import type { + EditToolDetails, + EditToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const editTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as EditToolInput | undefined; + const details = rawDetails as EditToolDetails | undefined; + + const locations: AgentToolCallLocation[] | undefined = args?.path + ? [{ path: args.path, line: details?.firstChangedLine }] + : undefined; + + if (args?.path && args.edits && args.edits.length > 0) { + const diff: AgentToolCallContent = { + type: "diff", + path: args.path, + oldText: args.edits.map((edit) => edit.oldText).join("\n"), + newText: args.edits.map((edit) => edit.newText).join("\n"), + }; + + return { locations, content: [diff] }; + } + + if (args?.path && details?.diff) { + const diff: AgentToolCallContent = { + type: "diff", + path: args.path, + newText: details.diff, + }; + + return { locations, content: [diff] }; + } + + const textContent = resultContent + ?.filter((block) => block.type === "text") + .map((block) => ({ + type: "content" as const, + content: { type: "text" as const, text: block.text }, + })); + + return { + locations, + content: textContent && textContent.length > 0 ? textContent : undefined, + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/findTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/findTranslator.test.ts new file mode 100644 index 0000000000..f862ae05c2 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/findTranslator.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it } from "vitest"; +import { findTranslator } from "./findTranslator"; + +describe("findTranslator", () => { + it("emits a location from the search path and text content from results", () => { + const output = findTranslator({ + toolCallId: "call-1", + arguments: { pattern: "*.ts", path: "packages/core/src" }, + resultContent: [ + { + type: "text", + text: "packages/core/src/foo.ts\npackages/core/src/bar.ts", + }, + ], + }); + + expect(output.locations).toEqual([{ path: "packages/core/src" }]); + expect(output.content).toEqual([ + { + type: "content", + content: { + type: "text", + text: "packages/core/src/foo.ts\npackages/core/src/bar.ts", + }, + }, + ]); + }); + + it("omits locations and content when path and resultContent are absent", () => { + const output = findTranslator({ + toolCallId: "call-2", + arguments: { pattern: "*.ts" }, + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); + + it("returns no content when result blocks are all images", () => { + const output = findTranslator({ + toolCallId: "call-3", + arguments: { pattern: "*.png", path: "assets" }, + resultContent: [ + { type: "image", data: "base64data", mimeType: "image/png" }, + ], + }); + + expect(output.locations).toEqual([{ path: "assets" }]); + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/findTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/findTranslator.ts new file mode 100644 index 0000000000..cfa526b69e --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/findTranslator.ts @@ -0,0 +1,26 @@ +import type { FindToolInput } from "@earendil-works/pi-coding-agent"; +import type { AgentToolCallLocation } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const findTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, +}) => { + const args = rawArguments as FindToolInput | undefined; + + const locations: AgentToolCallLocation[] | undefined = args?.path + ? [{ path: args.path }] + : undefined; + + const content = resultContent + ?.filter((block) => block.type === "text") + .map((block) => ({ + type: "content" as const, + content: { type: "text" as const, text: block.text }, + })); + + return { + locations, + content: content && content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/grepTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/grepTranslator.test.ts new file mode 100644 index 0000000000..820f9bcae7 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/grepTranslator.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; +import { grepTranslator } from "./grepTranslator"; + +describe("grepTranslator", () => { + it("returns locations and text content on success", () => { + const result = grepTranslator({ + toolCallId: "call-1", + arguments: { pattern: "foo", path: "src" }, + resultContent: [ + { type: "text", text: "src/a.ts:1:foo\nsrc/b.ts:2:foo bar" }, + ], + details: undefined, + isError: false, + }); + + expect(result.locations).toEqual([{ path: "src" }]); + expect(result.content).toEqual([ + { + type: "content", + content: { type: "text", text: "src/a.ts:1:foo\nsrc/b.ts:2:foo bar" }, + }, + ]); + }); + + it("appends truncation notes and omits locations when path is missing", () => { + const result = grepTranslator({ + toolCallId: "call-2", + arguments: { pattern: "foo" }, + resultContent: [{ type: "text", text: "match" }], + details: { matchLimitReached: 100, linesTruncated: true }, + isError: false, + }); + + expect(result.locations).toBeUndefined(); + expect(result.content).toEqual([ + { type: "content", content: { type: "text", text: "match" } }, + { + type: "content", + content: { + type: "text", + text: "Match limit reached at 100 matches. Some lines were truncated.", + }, + }, + ]); + }); + + it("returns undefined content and locations when there is nothing to report", () => { + const result = grepTranslator({ + toolCallId: "call-3", + arguments: { pattern: "foo" }, + resultContent: [], + details: undefined, + isError: true, + }); + + expect(result.locations).toBeUndefined(); + expect(result.content).toBeUndefined(); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/grepTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/grepTranslator.ts new file mode 100644 index 0000000000..0686d7bfb4 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/grepTranslator.ts @@ -0,0 +1,55 @@ +import type { + GrepToolDetails, + GrepToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContentBlock, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const grepTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as GrepToolInput; + const details = rawDetails as GrepToolDetails | undefined; + + const locations: AgentToolCallLocation[] = []; + if (args?.path) { + locations.push({ path: args.path }); + } + + const content: AgentToolCallContentBlock[] = []; + const resultText = (resultContent ?? []) + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("\n"); + + if (resultText) { + content.push({ + type: "content", + content: { type: "text", text: resultText }, + }); + } + + const notes: string[] = []; + if (details?.matchLimitReached !== undefined) { + notes.push(`Match limit reached at ${details.matchLimitReached} matches.`); + } + if (details?.linesTruncated) { + notes.push("Some lines were truncated."); + } + if (notes.length > 0) { + content.push({ + type: "content", + content: { type: "text", text: notes.join(" ") }, + }); + } + + return { + locations: locations.length > 0 ? locations : undefined, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/lsTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/lsTranslator.test.ts new file mode 100644 index 0000000000..bfe71d3a21 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/lsTranslator.test.ts @@ -0,0 +1,57 @@ +import { describe, expect, it } from "vitest"; +import { lsTranslator } from "./lsTranslator"; + +describe("lsTranslator", () => { + it("returns a location for the listed path and text content from the result", () => { + const output = lsTranslator({ + toolCallId: "call-1", + arguments: { path: "src/features" }, + resultContent: [{ type: "text", text: "sessions/\ntools/" }], + }); + + expect(output.locations).toEqual([{ path: "src/features" }]); + expect(output.content).toEqual([ + { type: "content", content: { type: "text", text: "sessions/\ntools/" } }, + ]); + }); + + it("omits locations when no path is given and content when result is empty", () => { + const output = lsTranslator({ + toolCallId: "call-2", + arguments: {}, + resultContent: [], + }); + + expect(output.locations).toBeUndefined(); + expect(output.content).toBeUndefined(); + }); + + it("still surfaces content on an error result", () => { + const output = lsTranslator({ + toolCallId: "call-3", + arguments: { path: "missing-dir" }, + resultContent: [{ type: "text", text: "ENOENT: no such directory" }], + isError: true, + }); + + expect(output.locations).toEqual([{ path: "missing-dir" }]); + expect(output.content).toEqual([ + { + type: "content", + content: { type: "text", text: "ENOENT: no such directory" }, + }, + ]); + }); + + it("ignores image content blocks", () => { + const output = lsTranslator({ + toolCallId: "call-4", + arguments: { path: "src" }, + resultContent: [ + { type: "image", data: "base64data", mimeType: "image/png" }, + ], + }); + + expect(output.content).toBeUndefined(); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/lsTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/lsTranslator.ts new file mode 100644 index 0000000000..d50893580e --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/lsTranslator.ts @@ -0,0 +1,33 @@ +import type { LsToolInput } from "@earendil-works/pi-coding-agent"; +import type { + AgentToolCallContent, + AgentToolCallLocation, +} from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const lsTranslator: PiToolTranslator = ({ + arguments: args, + resultContent, +}) => { + const input = args as LsToolInput | undefined; + + const locations: AgentToolCallLocation[] = []; + if (input?.path) { + locations.push({ path: input.path }); + } + + const content: AgentToolCallContent[] = []; + for (const block of resultContent ?? []) { + if (block.type === "text") { + content.push({ + type: "content", + content: { type: "text", text: block.text }, + }); + } + } + + return { + locations: locations.length > 0 ? locations : undefined, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/readTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/readTranslator.test.ts new file mode 100644 index 0000000000..37e7398540 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/readTranslator.test.ts @@ -0,0 +1,94 @@ +import type { + ReadToolDetails, + ReadToolInput, +} from "@earendil-works/pi-coding-agent"; +import { describe, expect, it } from "vitest"; +import { readTranslator } from "./readTranslator"; + +describe("readTranslator", () => { + it("returns the read path as a location and the file text as content", () => { + const args: ReadToolInput = { path: "src/index.ts" }; + + const output = readTranslator({ + toolCallId: "1", + arguments: args, + resultContent: [{ type: "text", text: "export const x = 1;" }], + isError: false, + }); + + expect(output).toEqual({ + locations: [{ path: "src/index.ts" }], + content: [ + { + type: "content", + content: { type: "text", text: "export const x = 1;" }, + }, + ], + }); + }); + + it("appends truncation info when the read result was truncated", () => { + const args: ReadToolInput = { path: "src/big.ts", offset: 0, limit: 100 }; + const details: ReadToolDetails = { + truncation: { + content: "line1\nline2", + truncated: true, + truncatedBy: "lines", + totalLines: 5000, + totalBytes: 100000, + outputLines: 100, + outputBytes: 2000, + lastLinePartial: false, + firstLineExceedsLimit: false, + maxLines: 2000, + maxBytes: 50000, + }, + }; + + const output = readTranslator({ + toolCallId: "2", + arguments: args, + resultContent: [{ type: "text", text: "line1\nline2" }], + details, + isError: false, + }); + + expect(output.locations).toEqual([{ path: "src/big.ts" }]); + expect(output.content).toHaveLength(1); + const [content] = output.content ?? []; + expect(content?.type).toBe("content"); + if (content?.type === "content" && content.content.type === "text") { + expect(content.content.text).toContain("line1\nline2"); + expect(content.content.text).toContain( + "truncated: showing 100 of 5000 lines", + ); + } + }); + + it("omits content when the result has no text block, e.g. an error result", () => { + const args: ReadToolInput = { path: "src/missing.ts" }; + + const output = readTranslator({ + toolCallId: "3", + arguments: args, + resultContent: undefined, + isError: true, + }); + + expect(output).toEqual({ + locations: [{ path: "src/missing.ts" }], + content: undefined, + }); + }); + + it("omits locations when arguments are missing", () => { + const output = readTranslator({ + toolCallId: "4", + arguments: undefined, + resultContent: [{ type: "text", text: "content" }], + isError: false, + }); + + expect(output.locations).toBeUndefined(); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/readTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/readTranslator.ts new file mode 100644 index 0000000000..4086394c94 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/readTranslator.ts @@ -0,0 +1,35 @@ +import type { + ReadToolDetails, + ReadToolInput, +} from "@earendil-works/pi-coding-agent"; +import type { AgentToolCallContent } from "@posthog/shared"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const readTranslator: PiToolTranslator = ({ + arguments: rawArguments, + resultContent, + details: rawDetails, +}) => { + const args = rawArguments as ReadToolInput | undefined; + const details = rawDetails as ReadToolDetails | undefined; + + const locations = args?.path ? [{ path: args.path }] : undefined; + + const textBlock = resultContent?.find((block) => block.type === "text"); + const content: AgentToolCallContent[] = []; + + if (textBlock && textBlock.type === "text") { + let text = textBlock.text; + + if (details?.truncation?.truncated) { + text = `${text}\n[truncated: showing ${details.truncation.outputLines} of ${details.truncation.totalLines} lines]`; + } + + content.push({ type: "content", content: { type: "text", text } }); + } + + return { + locations, + content: content.length > 0 ? content : undefined, + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/tools/writeTranslator.test.ts b/packages/core/src/pi-runtime/conversation/tools/writeTranslator.test.ts new file mode 100644 index 0000000000..caffa7403e --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/writeTranslator.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from "vitest"; +import { writeTranslator } from "./writeTranslator"; + +describe("writeTranslator", () => { + it("produces a diff content block and location for the written file", () => { + const result = writeTranslator({ + toolCallId: "call-1", + arguments: { path: "src/foo.ts", content: "export const foo = 1;\n" }, + }); + + expect(result.locations).toEqual([{ path: "src/foo.ts" }]); + expect(result.content).toEqual([ + { + type: "diff", + path: "src/foo.ts", + oldText: null, + newText: "export const foo = 1;\n", + }, + ]); + }); + + it("still produces a diff when the result is an error", () => { + const result = writeTranslator({ + toolCallId: "call-2", + arguments: { path: "src/bar.ts", content: "" }, + isError: true, + resultContent: [{ type: "text", text: "permission denied" }], + }); + + expect(result.locations).toEqual([{ path: "src/bar.ts" }]); + expect(result.content).toEqual([ + { + type: "diff", + path: "src/bar.ts", + oldText: null, + newText: "", + }, + ]); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/tools/writeTranslator.ts b/packages/core/src/pi-runtime/conversation/tools/writeTranslator.ts new file mode 100644 index 0000000000..7e2e4b4bf8 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/tools/writeTranslator.ts @@ -0,0 +1,18 @@ +import type { WriteToolInput } from "@earendil-works/pi-coding-agent"; +import type { PiToolTranslator } from "../toolTranslator"; + +export const writeTranslator: PiToolTranslator = ({ arguments: args }) => { + const input = args as WriteToolInput; + + return { + locations: [{ path: input.path }], + content: [ + { + type: "diff", + path: input.path, + oldText: null, + newText: input.content, + }, + ], + }; +}; diff --git a/packages/core/src/pi-runtime/conversation/translatePiConversation.test.ts b/packages/core/src/pi-runtime/conversation/translatePiConversation.test.ts new file mode 100644 index 0000000000..cfbbcbd7cf --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/translatePiConversation.test.ts @@ -0,0 +1,205 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { createPiConversationTranslator } from "./translatePiConversation"; + +function assistant( + content: AssistantMessage["content"], + stopReason: AssistantMessage["stopReason"] = "stop", + timestamp = 10, +): AssistantMessage { + return { + role: "assistant", + content, + api: "anthropic-messages" as AssistantMessage["api"], + provider: "anthropic" as AssistantMessage["provider"], + model: "test-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason, + timestamp, + }; +} + +describe("createPiConversationTranslator", () => { + it("keeps complete assistant content when translating history", () => { + const translator = createPiConversationTranslator(); + + expect( + translator.translateHistoryMessage( + assistant([{ type: "text", text: "complete" }]), + ), + ).toContainEqual({ + type: "assistant_message_chunk", + timestamp: 10, + content: { type: "text", text: "complete" }, + }); + }); + + it("uses message_update deltas without repeating cumulative text at message_end", () => { + const translator = createPiConversationTranslator(); + const message = assistant([{ type: "text", text: "complete" }]); + + const streamed = translator.translateEvent({ + type: "message_update", + message, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "complete", + partial: message, + }, + }); + const ended = translator.translateEvent({ type: "message_end", message }); + + expect(streamed).toEqual([ + { + type: "assistant_message_chunk", + timestamp: 10, + content: { type: "text", text: "complete" }, + }, + ]); + expect(ended).toEqual([]); + }); + + it("does not repeat streamed content when assistant timestamps collide", () => { + const translator = createPiConversationTranslator(); + const first = assistant([{ type: "text", text: "first" }]); + const second = assistant([{ type: "text", text: "second" }]); + + translator.translateEvent({ + type: "message_update", + message: first, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "first", + partial: first, + }, + }); + translator.translateEvent({ type: "message_end", message: first }); + translator.translateEvent({ + type: "message_update", + message: second, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "second", + partial: second, + }, + }); + + expect( + translator.translateEvent({ type: "message_end", message: second }), + ).toEqual([]); + }); + + it("completes a turn using the latest runtime timestamp", () => { + const translator = createPiConversationTranslator(); + const laterMessage = assistant( + [{ type: "text", text: "later" }], + "stop", + 20, + ); + const earlierMessage = assistant( + [{ type: "text", text: "earlier" }], + "stop", + 10, + ); + + translator.translateEvent({ type: "message_end", message: laterMessage }); + translator.translateEvent({ type: "message_end", message: earlierMessage }); + + expect(translator.translateEvent({ type: "agent_settled" })).toEqual([ + { type: "turn_completed", timestamp: 20 }, + ]); + }); + + it("translates direct bash history into the generic execute tool UI", () => { + const translator = createPiConversationTranslator(); + + expect( + translator.translateHistoryMessage({ + role: "bashExecution", + command: "pwd", + output: "/tmp/project", + exitCode: 0, + cancelled: false, + truncated: false, + timestamp: 20, + }), + ).toEqual([ + { + type: "tool_call_started", + timestamp: 20, + toolCall: { + id: "pi-bash-20", + title: "pwd", + kind: "execute", + status: "in_progress", + rawInput: { command: "pwd" }, + }, + }, + { + type: "tool_call_updated", + timestamp: 20, + toolCall: { + id: "pi-bash-20", + status: "completed", + rawOutput: "/tmp/project", + content: [ + { + type: "content", + content: { type: "text", text: "/tmp/project" }, + }, + ], + }, + }, + ]); + }); + + it("preserves tool calls when filtering streamed assistant content", () => { + const translator = createPiConversationTranslator(); + const message = assistant([ + { type: "text", text: "running" }, + { + type: "toolCall", + id: "tool-1", + name: "bash", + arguments: { command: "pwd" }, + }, + ]); + + translator.translateEvent({ + type: "message_update", + message, + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "running", + partial: message, + }, + }); + + expect(translator.translateEvent({ type: "message_end", message })).toEqual( + [ + { + type: "tool_call_started", + timestamp: 10, + toolCall: { + id: "tool-1", + title: "bash", + kind: "execute", + status: "pending", + rawInput: { command: "pwd" }, + }, + }, + ], + ); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/translatePiConversation.ts b/packages/core/src/pi-runtime/conversation/translatePiConversation.ts new file mode 100644 index 0000000000..ca56c90b40 --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/translatePiConversation.ts @@ -0,0 +1,214 @@ +import type { AssistantMessage, Message } from "@earendil-works/pi-ai"; +import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { createPiMessageTranslator } from "./translatePiMessage"; + +type AgentMessage = Extract< + AgentSessionEvent, + { type: "message_end" } +>["message"]; + +function isMessage(message: AgentMessage): message is Message { + return ( + message.role === "user" || + message.role === "assistant" || + message.role === "toolResult" + ); +} + +function customMessageEvents(message: AgentMessage): AgentConversationEvent[] { + if (message.role === "bashExecution") { + const id = `pi-bash-${message.timestamp}`; + const failed = message.cancelled || (message.exitCode ?? 0) !== 0; + + return [ + { + type: "tool_call_started", + timestamp: message.timestamp, + toolCall: { + id, + title: message.command, + kind: "execute", + status: "in_progress", + rawInput: { command: message.command }, + }, + }, + { + type: "tool_call_updated", + timestamp: message.timestamp, + toolCall: { + id, + status: failed ? "failed" : "completed", + rawOutput: message.output, + content: message.output + ? [ + { + type: "content", + content: { type: "text", text: message.output }, + }, + ] + : [], + }, + }, + ]; + } + + let text: string | undefined; + + if ( + message.role === "branchSummary" || + message.role === "compactionSummary" + ) { + text = message.summary; + } else if (message.role === "custom" && message.display) { + text = + typeof message.content === "string" + ? message.content + : message.content + .flatMap((content) => + content.type === "text" ? [content.text] : [], + ) + .join("\n"); + } + + if (!text) { + return []; + } + + return [ + { + type: "assistant_message_chunk", + timestamp: message.timestamp, + content: { type: "text", text }, + }, + ]; +} + +function isAssistantMessage( + message: AgentMessage, +): message is AssistantMessage { + return message.role === "assistant"; +} + +export interface PiConversationTranslator { + translateHistoryMessage(message: AgentMessage): AgentConversationEvent[]; + translateEvent(event: AgentSessionEvent): AgentConversationEvent[]; +} + +export function createPiConversationTranslator(): PiConversationTranslator { + const messageTranslator = createPiMessageTranslator(); + const streamedAssistantTimestamps = new Set(); + let historyTurnActive = false; + let latestRuntimeTimestamp = 0; + + function translateHistoryMessage( + message: AgentMessage, + ): AgentConversationEvent[] { + const events: AgentConversationEvent[] = []; + + if (message.role === "user" && historyTurnActive) { + events.push({ + type: "turn_completed", + timestamp: message.timestamp, + }); + historyTurnActive = false; + } + + if (isMessage(message)) { + events.push(...messageTranslator.translate(message)); + } else { + events.push(...customMessageEvents(message)); + } + + if (message.role === "user") { + historyTurnActive = true; + } + + if ( + isAssistantMessage(message) && + message.stopReason !== "toolUse" && + historyTurnActive + ) { + events.push({ + type: "turn_completed", + timestamp: message.timestamp, + stopReason: message.stopReason, + }); + historyTurnActive = false; + } + + return events; + } + + function translateEvent(event: AgentSessionEvent): AgentConversationEvent[] { + if (event.type === "message_update") { + const update = event.assistantMessageEvent; + latestRuntimeTimestamp = Math.max( + latestRuntimeTimestamp, + event.message.timestamp, + ); + + if (update.type === "text_delta" && update.delta) { + streamedAssistantTimestamps.add(event.message.timestamp); + return [ + { + type: "assistant_message_chunk", + timestamp: event.message.timestamp, + content: { type: "text", text: update.delta }, + }, + ]; + } + + if (update.type === "thinking_delta" && update.delta) { + streamedAssistantTimestamps.add(event.message.timestamp); + return [ + { + type: "assistant_thought_chunk", + timestamp: event.message.timestamp, + content: { type: "text", text: update.delta }, + }, + ]; + } + + return []; + } + + if (event.type === "message_end") { + latestRuntimeTimestamp = Math.max( + latestRuntimeTimestamp, + event.message.timestamp, + ); + + if (!isMessage(event.message)) { + return customMessageEvents(event.message); + } + + const events = messageTranslator.translate(event.message); + if ( + event.message.role !== "assistant" || + !streamedAssistantTimestamps.has(event.message.timestamp) + ) { + return events; + } + + return events.filter( + (translated) => + translated.type !== "assistant_message_chunk" && + translated.type !== "assistant_thought_chunk", + ); + } + + if (event.type === "agent_settled") { + streamedAssistantTimestamps.clear(); + + const timestamp = latestRuntimeTimestamp; + latestRuntimeTimestamp = 0; + + return timestamp > 0 ? [{ type: "turn_completed", timestamp }] : []; + } + + return []; + } + + return { translateHistoryMessage, translateEvent }; +} diff --git a/packages/core/src/pi-runtime/conversation/translatePiMessage.test.ts b/packages/core/src/pi-runtime/conversation/translatePiMessage.test.ts new file mode 100644 index 0000000000..fe9dc67edf --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/translatePiMessage.test.ts @@ -0,0 +1,97 @@ +import type { AssistantMessage, UserMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { createPiMessageTranslator } from "./translatePiMessage"; + +function makeAssistant(content: AssistantMessage["content"]): AssistantMessage { + return { + role: "assistant", + content, + api: "anthropic-messages" as AssistantMessage["api"], + provider: "anthropic" as AssistantMessage["provider"], + model: "test-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 0, + }; +} + +describe("createPiMessageTranslator", () => { + it("translates a string user message into a user message chunk", () => { + const translator = createPiMessageTranslator(); + + const message: UserMessage = { + role: "user", + content: "hello there", + timestamp: 0, + }; + + expect(translator.translate(message)).toEqual([ + { + type: "user_message", + id: "pi-user-0-1", + timestamp: 0, + content: [{ type: "text", text: "hello there" }], + }, + ]); + }); + + it("translates user message content blocks into user message chunks", () => { + const translator = createPiMessageTranslator(); + + const message: UserMessage = { + role: "user", + content: [ + { type: "text", text: "first" }, + { type: "image", data: "abc", mimeType: "image/png" }, + ], + timestamp: 0, + }; + + expect(translator.translate(message)).toEqual([ + { + type: "user_message", + id: "pi-user-0-1", + timestamp: 0, + content: [ + { type: "text", text: "first" }, + { type: "image", data: "abc", mimeType: "image/png" }, + ], + }, + ]); + }); + + it("translates a plain text assistant message into an agent message chunk", () => { + const translator = createPiMessageTranslator(); + + const message = makeAssistant([{ type: "text", text: "working on it" }]); + + expect(translator.translate(message)).toEqual([ + { + type: "assistant_message_chunk", + timestamp: 0, + content: { type: "text", text: "working on it" }, + }, + ]); + }); + + it("translates assistant thinking into an agent thought chunk", () => { + const translator = createPiMessageTranslator(); + + const message = makeAssistant([{ type: "thinking", thinking: "hmm" }]); + + expect(translator.translate(message)).toEqual([ + { + type: "assistant_thought_chunk", + timestamp: 0, + content: { type: "text", text: "hmm" }, + }, + ]); + }); +}); diff --git a/packages/core/src/pi-runtime/conversation/translatePiMessage.ts b/packages/core/src/pi-runtime/conversation/translatePiMessage.ts new file mode 100644 index 0000000000..4ffada68ee --- /dev/null +++ b/packages/core/src/pi-runtime/conversation/translatePiMessage.ts @@ -0,0 +1,208 @@ +import type { + AssistantMessage, + Message, + ToolResultMessage, + UserMessage, +} from "@earendil-works/pi-ai"; +import type { + AgentContent, + AgentConversationEvent, + AgentToolCallStatus, +} from "@posthog/shared"; +import { type PiToolName, TOOL_KIND_BY_NAME } from "./toolKind"; +import { bashTranslator } from "./tools/bashTranslator"; +import { editTranslator } from "./tools/editTranslator"; +import { findTranslator } from "./tools/findTranslator"; +import { grepTranslator } from "./tools/grepTranslator"; +import { lsTranslator } from "./tools/lsTranslator"; +import { readTranslator } from "./tools/readTranslator"; +import { writeTranslator } from "./tools/writeTranslator"; +import type { PiToolTranslator } from "./toolTranslator"; + +const TRANSLATOR_BY_NAME: Record = { + read: readTranslator, + bash: bashTranslator, + edit: editTranslator, + write: writeTranslator, + grep: grepTranslator, + find: findTranslator, + ls: lsTranslator, +}; + +interface PendingToolCall { + name: string; + arguments: unknown; +} + +function isPiToolName(name: string): name is PiToolName { + return name in TOOL_KIND_BY_NAME; +} + +function toContent(block: { + type: string; + text?: string; + data?: string; + mimeType?: string; +}): AgentContent | undefined { + if (block.type === "text" && typeof block.text === "string") { + return { type: "text", text: block.text }; + } + + if ( + block.type === "image" && + typeof block.data === "string" && + typeof block.mimeType === "string" + ) { + return { type: "image", data: block.data, mimeType: block.mimeType }; + } + + return undefined; +} + +export interface PiMessageTranslator { + translate(message: Message): AgentConversationEvent[]; +} + +export function createPiMessageTranslator(): PiMessageTranslator { + const pendingToolCalls = new Map(); + let userMessageId = 0; + + function translateUser(message: UserMessage): AgentConversationEvent[] { + const content = + typeof message.content === "string" + ? [{ type: "text" as const, text: message.content }] + : message.content.flatMap((block) => { + const translated = toContent(block); + return translated ? [translated] : []; + }); + + if (content.length === 0) { + return []; + } + + userMessageId += 1; + + return [ + { + type: "user_message", + id: `pi-user-${message.timestamp}-${userMessageId}`, + timestamp: message.timestamp, + content, + }, + ]; + } + + function translateAssistant( + message: AssistantMessage, + ): AgentConversationEvent[] { + const events: AgentConversationEvent[] = []; + + for (const block of message.content) { + if (block.type === "text") { + events.push({ + type: "assistant_message_chunk", + timestamp: message.timestamp, + content: { type: "text", text: block.text }, + }); + continue; + } + + if (block.type === "thinking") { + events.push({ + type: "assistant_thought_chunk", + timestamp: message.timestamp, + content: { type: "text", text: block.thinking }, + }); + continue; + } + + if (block.type === "toolCall") { + pendingToolCalls.set(block.id, { + name: block.name, + arguments: block.arguments, + }); + + const kind = isPiToolName(block.name) + ? TOOL_KIND_BY_NAME[block.name] + : null; + + events.push({ + type: "tool_call_started", + timestamp: message.timestamp, + toolCall: { + id: block.id, + title: block.name, + kind, + status: "pending", + rawInput: block.arguments, + }, + }); + } + } + + return events; + } + + function translateToolResult( + message: ToolResultMessage, + ): AgentConversationEvent[] { + const pending = pendingToolCalls.get(message.toolCallId); + pendingToolCalls.delete(message.toolCallId); + + const status: AgentToolCallStatus = message.isError + ? "failed" + : "completed"; + const toolCall: Extract< + AgentConversationEvent, + { type: "tool_call_updated" } + >["toolCall"] = { + id: message.toolCallId, + status, + rawOutput: message.content, + }; + + const translator = isPiToolName(message.toolName) + ? TRANSLATOR_BY_NAME[message.toolName] + : undefined; + + if (translator) { + const output = translator({ + toolCallId: message.toolCallId, + arguments: pending?.arguments, + resultContent: message.content, + details: message.details, + isError: message.isError, + }); + + if (output.content) { + toolCall.content = output.content; + } + + if (output.locations) { + toolCall.locations = output.locations; + } + } + + return [ + { + type: "tool_call_updated", + timestamp: message.timestamp, + toolCall, + }, + ]; + } + + return { + translate(message: Message): AgentConversationEvent[] { + if (message.role === "user") { + return translateUser(message); + } + + if (message.role === "assistant") { + return translateAssistant(message); + } + + return translateToolResult(message); + }, + }; +} diff --git a/packages/host-router/src/routers/pi-session.router.ts b/packages/host-router/src/routers/pi-session.router.ts index a45b2c86c9..0df3c37a96 100644 --- a/packages/host-router/src/routers/pi-session.router.ts +++ b/packages/host-router/src/routers/pi-session.router.ts @@ -2,11 +2,39 @@ 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 { + piSessionAvailableModelsOutput, + piSessionBashInput, + piSessionBashOutput, + piSessionCancelledOutput, + piSessionCommandsOutput, + piSessionCompactInput, + piSessionCycleModelOutput, + piSessionEnabledInput, piSessionEntriesInput, + piSessionEntryInput, + piSessionExportInput, + piSessionExportOutput, + piSessionForkMessagesOutput, + piSessionForkOutput, + piSessionLastAssistantTextOutput, + piSessionMessageInput, + piSessionModelInput, + piSessionModelOutput, + piSessionNameInput, + piSessionNewInput, + piSessionPathInput, + piSessionPromptAndWaitInput, piSessionPromptInput, + piSessionQueueModeInput, piSessionRuntimeOutput, piSessionStartOutput, + piSessionStatusOutput, + piSessionStderrOutput, + piSessionThinkingCycleOutput, + piSessionThinkingLevelInput, + piSessionTimeoutInput, piSessionTranscriptInput, + piSessionUnknownOutput, resumePiSessionInput, startPiSessionInput, } from "@posthog/workspace-server/services/pi-session/schemas"; @@ -27,7 +55,31 @@ export const piSessionRouter = router({ prompt: publicProcedure .input(piSessionPromptInput) .mutation(({ ctx, input }) => - getService(ctx.container).prompt(input.taskId, input.prompt), + getService(ctx.container).prompt( + input.taskId, + input.prompt, + input.images, + ), + ), + + steer: publicProcedure + .input(piSessionMessageInput) + .mutation(({ ctx, input }) => + getService(ctx.container).steer( + input.taskId, + input.message, + input.images, + ), + ), + + followUp: publicProcedure + .input(piSessionMessageInput) + .mutation(({ ctx, input }) => + getService(ctx.container).followUp( + input.taskId, + input.message, + input.images, + ), ), abort: publicProcedure @@ -36,17 +88,152 @@ export const piSessionRouter = router({ getService(ctx.container).abort(input.taskId), ), - stop: publicProcedure + newSession: publicProcedure + .input(piSessionNewInput) + .output(piSessionCancelledOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).newSession(input.taskId, input.parentSession), + ), + + setModel: publicProcedure + .input(piSessionModelInput) + .output(piSessionModelOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).setModel( + input.taskId, + input.provider, + input.modelId, + ), + ), + + cycleModel: publicProcedure .input(piSessionTranscriptInput) - .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + .output(piSessionCycleModelOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).cycleModel(input.taskId), + ), - runtime: publicProcedure + availableModels: publicProcedure .input(piSessionTranscriptInput) - .output(piSessionRuntimeOutput) - .query(({ ctx, input }) => getService(ctx.container).runtime(input.taskId)), + .output(piSessionAvailableModelsOutput) + .query(({ ctx, input }) => + getService(ctx.container).availableModels(input.taskId), + ), + + setThinkingLevel: publicProcedure + .input(piSessionThinkingLevelInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setThinkingLevel(input.taskId, input.level), + ), + + cycleThinkingLevel: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionThinkingCycleOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).cycleThinkingLevel(input.taskId), + ), + + setSteeringMode: publicProcedure + .input(piSessionQueueModeInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setSteeringMode(input.taskId, input.mode), + ), + + setFollowUpMode: publicProcedure + .input(piSessionQueueModeInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setFollowUpMode(input.taskId, input.mode), + ), + + compact: publicProcedure + .input(piSessionCompactInput) + .output(piSessionUnknownOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).compact(input.taskId, input.customInstructions), + ), + + setAutoCompaction: publicProcedure + .input(piSessionEnabledInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setAutoCompaction(input.taskId, input.enabled), + ), + + setAutoRetry: publicProcedure + .input(piSessionEnabledInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setAutoRetry(input.taskId, input.enabled), + ), + + abortRetry: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).abortRetry(input.taskId), + ), + + bash: publicProcedure + .input(piSessionBashInput) + .output(piSessionBashOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).bash(input.taskId, input.command), + ), + + abortBash: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => + getService(ctx.container).abortBash(input.taskId), + ), + + sessionStats: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionUnknownOutput) + .query(({ ctx, input }) => + getService(ctx.container).sessionStats(input.taskId), + ), + + exportHtml: publicProcedure + .input(piSessionExportInput) + .output(piSessionExportOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).exportHtml(input.taskId, input.outputPath), + ), + + switchSession: publicProcedure + .input(piSessionPathInput) + .output(piSessionCancelledOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).switchSession(input.taskId, input.sessionPath), + ), + + fork: publicProcedure + .input(piSessionEntryInput) + .output(piSessionForkOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).fork(input.taskId, input.entryId), + ), + + clone: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionCancelledOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).clone(input.taskId), + ), + + forkMessages: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionForkMessagesOutput) + .query(({ ctx, input }) => + getService(ctx.container).forkMessages(input.taskId), + ), + + setSessionName: publicProcedure + .input(piSessionNameInput) + .mutation(({ ctx, input }) => + getService(ctx.container).setSessionName(input.taskId, input.name), + ), status: publicProcedure .input(piSessionTranscriptInput) + .output(piSessionStatusOutput) .query(({ ctx, input }) => getService(ctx.container).status(input.taskId)), entries: publicProcedure @@ -55,6 +242,71 @@ export const piSessionRouter = router({ getService(ctx.container).entries(input.taskId, input.since), ), + tree: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionUnknownOutput) + .query(({ ctx, input }) => getService(ctx.container).tree(input.taskId)), + + lastAssistantText: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionLastAssistantTextOutput) + .query(({ ctx, input }) => + getService(ctx.container).lastAssistantText(input.taskId), + ), + + messages: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionUnknownOutput) + .query(({ ctx, input }) => + getService(ctx.container).messages(input.taskId), + ), + + commands: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionCommandsOutput) + .query(({ ctx, input }) => + getService(ctx.container).commands(input.taskId), + ), + + waitForIdle: publicProcedure + .input(piSessionTimeoutInput) + .mutation(({ ctx, input }) => + getService(ctx.container).waitForIdle(input.taskId, input.timeout), + ), + + collectEvents: publicProcedure + .input(piSessionTimeoutInput) + .output(piSessionUnknownOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).collectEvents(input.taskId, input.timeout), + ), + + promptAndWait: publicProcedure + .input(piSessionPromptAndWaitInput) + .output(piSessionUnknownOutput) + .mutation(({ ctx, input }) => + getService(ctx.container).promptAndWait( + input.taskId, + input.prompt, + input.images, + input.timeout, + ), + ), + + stderr: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionStderrOutput) + .query(({ ctx, input }) => getService(ctx.container).stderr(input.taskId)), + + stop: publicProcedure + .input(piSessionTranscriptInput) + .mutation(({ ctx, input }) => getService(ctx.container).stop(input.taskId)), + + runtime: publicProcedure + .input(piSessionTranscriptInput) + .output(piSessionRuntimeOutput) + .query(({ ctx, input }) => getService(ctx.container).runtime(input.taskId)), + onEvent: publicProcedure .input(piSessionTranscriptInput) .subscription(async function* (opts) { diff --git a/packages/shared/src/agent-conversation.ts b/packages/shared/src/agent-conversation.ts new file mode 100644 index 0000000000..6a09230ed0 --- /dev/null +++ b/packages/shared/src/agent-conversation.ts @@ -0,0 +1,141 @@ +export type AgentToolKind = + | "read" + | "edit" + | "delete" + | "move" + | "search" + | "execute" + | "think" + | "fetch" + | "switch_mode" + | "question" + | "other"; + +export type AgentToolCallStatus = + | "pending" + | "in_progress" + | "completed" + | "failed"; + +export interface AgentTextContent { + type: "text"; + text: string; +} + +export interface AgentImageContent { + type: "image"; + data: string; + mimeType: string; +} + +export interface AgentAudioContent { + type: "audio"; + data: string; + mimeType: string; +} + +export interface AgentResourceLinkContent { + type: "resource_link"; + uri: string; + name: string; + description?: string | null; + mimeType?: string | null; + size?: number | null; + title?: string | null; +} + +export interface AgentTextResource { + uri: string; + mimeType?: string | null; + text: string; +} + +export interface AgentBlobResource { + uri: string; + mimeType?: string | null; + blob: string; +} + +export interface AgentEmbeddedResourceContent { + type: "resource"; + resource: AgentTextResource | AgentBlobResource; +} + +export type AgentContent = + | AgentTextContent + | AgentImageContent + | AgentAudioContent + | AgentResourceLinkContent + | AgentEmbeddedResourceContent; + +export interface AgentToolCallContentBlock { + type: "content"; + content: AgentContent; +} + +export interface AgentToolCallDiff { + type: "diff"; + path: string; + oldText?: string | null; + newText: string; +} + +export interface AgentToolCallTerminal { + type: "terminal"; + terminalId: string; +} + +export type AgentToolCallContent = + | AgentToolCallContentBlock + | AgentToolCallDiff + | AgentToolCallTerminal; + +export interface AgentToolCallLocation { + path: string; + line?: number | null; +} + +export interface AgentToolCall { + id: string; + title: string; + kind?: AgentToolKind | null; + status?: AgentToolCallStatus | null; + content?: AgentToolCallContent[]; + locations?: AgentToolCallLocation[]; + rawInput?: unknown; + rawOutput?: unknown; + parentId?: string; +} + +export type AgentConversationEvent = + | { + type: "user_message"; + id: string; + timestamp: number; + content: AgentContent[]; + } + | { + type: "assistant_message_chunk"; + timestamp: number; + content: AgentContent; + } + | { + type: "assistant_thought_chunk"; + timestamp: number; + content: AgentContent; + } + | { + type: "tool_call_started"; + timestamp: number; + toolCall: AgentToolCall; + } + | { + type: "tool_call_updated"; + timestamp: number; + toolCall: Pick & Partial>; + } + | { + type: "turn_completed"; + timestamp: number; + stopReason?: string; + }; diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 33d2f564ed..cf6ce7845a 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,4 +1,23 @@ export * from "./adapter"; +export type { + AgentAudioContent, + AgentBlobResource, + AgentContent, + AgentConversationEvent, + AgentEmbeddedResourceContent, + AgentImageContent, + AgentResourceLinkContent, + AgentTextContent, + AgentTextResource, + AgentToolCall, + AgentToolCallContent, + AgentToolCallContentBlock, + AgentToolCallDiff, + AgentToolCallLocation, + AgentToolCallStatus, + AgentToolCallTerminal, + AgentToolKind, +} from "./agent-conversation"; export * from "./agent-runtime"; export * from "./analytics-events"; export { type ArchivedTask, archivedTaskSchema } from "./archive-domain"; diff --git a/packages/ui/src/features/pi-sessions/PiSessionView.tsx b/packages/ui/src/features/pi-sessions/PiSessionView.tsx index f97cdaa3c4..36370aac60 100644 --- a/packages/ui/src/features/pi-sessions/PiSessionView.tsx +++ b/packages/ui/src/features/pi-sessions/PiSessionView.tsx @@ -1,84 +1,55 @@ +import { createPiConversationTranslator } from "@posthog/core/pi-runtime/conversation/translatePiConversation"; import { useHostTRPC, useHostTRPCClient } from "@posthog/host-router/react"; import { - Button, Empty, EmptyDescription, EmptyHeader, EmptyTitle, } from "@posthog/quill"; +import type { AgentConversationEvent } from "@posthog/shared"; +import { PromptInput } from "@posthog/ui/features/message-editor/components/PromptInput"; +import { useDraftStore } from "@posthog/ui/features/message-editor/draftStore"; +import { ChatThread } from "@posthog/ui/features/sessions/components/chat-thread/ChatThread"; +import { CHAT_CONTENT_MAX_WIDTH } from "@posthog/ui/features/sessions/constants"; +import { useWorkspace } from "@posthog/ui/features/workspace/useWorkspace"; +import { Box, Flex } from "@radix-ui/themes"; import { useQuery } from "@tanstack/react-query"; import { useSubscription } from "@trpc/tanstack-react-query"; -import { useEffect, useMemo, useState } from "react"; -import { - applyPiEvent, - emptyLiveFeed, - type PiEntries, - PiEntriesSyncer, - type PiEvent, - type PiLiveFeed, - type PiMessage, -} from "./piSessionFeed"; +import { useEffect, useMemo, useRef, useState } from "react"; +import { type PiEntries, PiEntriesSyncer, type PiEvent } from "./piSessionFeed"; import { useEnsurePiSession } from "./useEnsurePiSession"; interface PiSessionViewProps { taskId: string; } -type PiMessageWithContent = Extract; - -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 draftActions = useDraftStore((state) => state.actions); + const workspace = useWorkspace(taskId); + const repoPath = workspace?.worktreePath ?? workspace?.folderPath; - const [prompt, setPrompt] = useState(""); - const [liveFeed, setLiveFeed] = useState(emptyLiveFeed); - const [syncedEntries, setSyncedEntries] = useState( - undefined, - ); + const [liveEvents, setLiveEvents] = useState([]); + const liveEventsRef = useRef([]); + const [syncedEntries, setSyncedEntries] = useState(); + const [isPromptPending, setIsPromptPending] = useState(); + const [isBashRunning, setIsBashRunning] = useState(false); const { data: fetchedEntries } = useQuery({ ...trpc.piSession.entries.queryOptions({ taskId }), enabled: sessionReady, }); - const { error: statusError } = useQuery({ + const { data: status, error: statusError } = useQuery({ ...trpc.piSession.status.queryOptions({ taskId }), enabled: sessionReady, }); + const { data: commands } = useQuery({ + ...trpc.piSession.commands.queryOptions({ taskId }), + enabled: sessionReady, + }); const syncer = useMemo( () => @@ -93,7 +64,53 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { syncer.seed(fetchedEntries); }, [syncer, fetchedEntries]); + useEffect(() => { + draftActions.setContext(taskId, { + taskId, + repoPath, + disabled: !sessionReady, + isLoading: !!isPromptPending || isBashRunning, + }); + }, [ + draftActions, + isBashRunning, + isPromptPending, + repoPath, + sessionReady, + taskId, + ]); + + useEffect(() => { + if (!commands) { + return; + } + + draftActions.setCommands( + taskId, + commands.map((command) => ({ + name: command.name, + description: command.description ?? "", + })), + ); + }, [commands, draftActions, taskId]); + const history = syncedEntries ?? fetchedEntries; + const conversationEvents = useMemo(() => { + const translator = createPiConversationTranslator(); + const translated: AgentConversationEvent[] = []; + + for (const entry of history?.entries ?? []) { + if (entry.type === "message") { + translated.push(...translator.translateHistoryMessage(entry.message)); + } + } + + for (const event of liveEvents) { + translated.push(...translator.translateEvent(event)); + } + + return translated; + }, [history, liveEvents]); useSubscription( trpc.piSession.onEvent.subscriptionOptions( @@ -101,23 +118,66 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { { enabled: sessionReady, onData: (event: PiEvent) => { - setLiveFeed((feed) => applyPiEvent(feed, event)); + if (event.type === "agent_start") { + setIsPromptPending(true); + } if (event.type === "agent_settled") { - void syncer.sync().then(() => setLiveFeed(emptyLiveFeed)); + const settledEventCount = liveEventsRef.current.length; + + void syncer + .sync() + .catch(() => undefined) + .finally(() => { + const remainingEvents = + liveEventsRef.current.slice(settledEventCount); + + liveEventsRef.current = remainingEvents; + setLiveEvents(remainingEvents); + setIsPromptPending(false); + }); + return; } + + const nextEvents = [...liveEventsRef.current, event]; + liveEventsRef.current = nextEvents; + setLiveEvents(nextEvents); }, }, ), ); - const send = async () => { - const text = prompt.trim(); - if (!text) { + const sendPrompt = (text: string) => { + const prompt = text.trim(); + if (!prompt) { return; } - await client.piSession.prompt.mutate({ taskId, prompt: text }); - setPrompt(""); + + setIsPromptPending(true); + void client.piSession.prompt + .mutate({ taskId, prompt }) + .catch(() => setIsPromptPending(false)); + }; + + const runBashCommand = (command: string) => { + setIsBashRunning(true); + void client.piSession.bash + .mutate({ taskId, command }) + .then(() => syncer.sync()) + .finally(() => setIsBashRunning(false)); + }; + + const cancelPrompt = () => { + if (isBashRunning) { + void client.piSession.abortBash + .mutate({ taskId }) + .finally(() => setIsBashRunning(false)); + return; + } + + void client.piSession.abort + .mutate({ taskId }) + .finally(() => setIsPromptPending(false)); }; const sessionError = ensureError ?? statusError; @@ -132,7 +192,7 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { ); } - if (!sessionReady) { + if (!sessionReady || !history || !status) { return ( @@ -142,55 +202,36 @@ export function PiSessionView({ taskId }: PiSessionViewProps) { ); } - return ( -
-
- {history?.entries.map((entry) => { - if (entry.type !== "message") { - return null; - } + const pending = (isPromptPending ?? status.isStreaming) || isBashRunning; - return ( -
- -
- ); - })} - {liveFeed.liveMessages.map((message) => ( -
- -
- ))} - {liveFeed.streamingMessage ? ( -
- -
- ) : null} -
-
-