@posthog/agent 2.0.2 → 2.0.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/agent",
-  "version": "2.0.2",
+  "version": "2.0.3",
   "repository": "https://github.com/PostHog/twig",
   "description": "TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog",
   "exports": {
@@ -46,7 +46,7 @@
     }
   },
   "bin": {
-    "agent-server": "./dist/server/bin.js"
+    "agent-server": "./dist/server/bin.cjs"
   },
   "type": "module",
   "keywords": [
@@ -72,8 +72,8 @@
     "tsx": "^4.20.6",
     "typescript": "^5.5.0",
     "vitest": "^2.1.8",
-    "@twig/git": "1.0.0",
-    "@posthog/shared": "1.0.0"
+    "@posthog/shared": "1.0.0",
+    "@twig/git": "1.0.0"
   },
   "dependencies": {
     "@opentelemetry/api-logs": "^0.208.0",

package/src/sagas/resume-saga.test.ts

@@ -5,6 +5,7 @@ import type { PostHogAPIClient } from "../posthog-api.js";
 import { ResumeSaga } from "./resume-saga.js";
 import {
   createAgentChunk,
+  createAgentMessage,
   createArchiveBuffer,
   createMockApiClient,
   createMockLogger,
@@ -141,6 +142,149 @@ describe("ResumeSaga", () => {
       expect(content).toEqual({ type: "text", text: "Hello world!" });
     });
 
+    it("rebuilds from coalesced agent_message events", async () => {
+      (mockApiClient.getTaskRun as ReturnType<typeof vi.fn>).mockResolvedValue(
+        createTaskRun(),
+      );
+      (
+        mockApiClient.fetchTaskRunLogs as ReturnType<typeof vi.fn>
+      ).mockResolvedValue([
+        createUserMessage("Hello"),
+        createAgentMessage("Hi there! Let me help."),
+        createUserMessage("Thanks"),
+        createAgentMessage("Sure thing."),
+      ]);
+
+      const saga = new ResumeSaga(mockLogger);
+      const result = await saga.run({
+        taskId: "task-1",
+        runId: "run-1",
+        repositoryPath: repo.path,
+        apiClient: mockApiClient,
+      });
+
+      expect(result.success).toBe(true);
+      if (!result.success) return;
+
+      expect(result.data.conversation).toHaveLength(4);
+      expect(result.data.conversation[0].role).toBe("user");
+      expect(result.data.conversation[1].role).toBe("assistant");
+      expect(result.data.conversation[1].content[0]).toEqual({
+        type: "text",
+        text: "Hi there! Let me help.",
+      });
+      expect(result.data.conversation[3].content[0]).toEqual({
+        type: "text",
+        text: "Sure thing.",
+      });
+    });
+
+    it("merges multiple agent_message events within a single turn", async () => {
+      (mockApiClient.getTaskRun as ReturnType<typeof vi.fn>).mockResolvedValue(
+        createTaskRun(),
+      );
+      (
+        mockApiClient.fetchTaskRunLogs as ReturnType<typeof vi.fn>
+      ).mockResolvedValue([
+        createUserMessage("Fix the bug"),
+        createAgentMessage("I'll look into this."),
+        createToolCall("call-1", "ReadFile", { path: "/bug.ts" }),
+        createToolResult("call-1", "buggy code"),
+        createAgentMessage(" Here's the fix."),
+        createToolCall("call-2", "Edit", { path: "/bug.ts" }),
+        createToolResult("call-2", "done"),
+        createAgentMessage(" All fixed now."),
+      ]);
+
+      const saga = new ResumeSaga(mockLogger);
+      const result = await saga.run({
+        taskId: "task-1",
+        runId: "run-1",
+        repositoryPath: repo.path,
+        apiClient: mockApiClient,
+      });
+
+      expect(result.success).toBe(true);
+      if (!result.success) return;
+
+      expect(result.data.conversation).toHaveLength(2);
+      const assistantTurn = result.data.conversation[1];
+      expect(assistantTurn.role).toBe("assistant");
+      // All text segments should merge into one block
+      expect(assistantTurn.content).toHaveLength(1);
+      expect(assistantTurn.content[0]).toEqual({
+        type: "text",
+        text: "I'll look into this. Here's the fix. All fixed now.",
+      });
+      expect(assistantTurn.toolCalls).toHaveLength(2);
+    });
+
+    it("produces same result from chunks and coalesced messages", async () => {
+      const chunkEntries = [
+        createUserMessage("Hello"),
+        createAgentChunk("First "),
+        createAgentChunk("part."),
+        createToolCall("call-1", "Read", { path: "/a" }),
+        createToolResult("call-1", "content"),
+        createAgentChunk("Second "),
+        createAgentChunk("part."),
+      ];
+
+      const coalescedEntries = [
+        createUserMessage("Hello"),
+        createAgentMessage("First part."),
+        createToolCall("call-1", "Read", { path: "/a" }),
+        createToolResult("call-1", "content"),
+        createAgentMessage("Second part."),
+      ];
+
+      // Run with chunks (old format)
+      (mockApiClient.getTaskRun as ReturnType<typeof vi.fn>).mockResolvedValue(
+        createTaskRun(),
+      );
+      (
+        mockApiClient.fetchTaskRunLogs as ReturnType<typeof vi.fn>
+      ).mockResolvedValue(chunkEntries);
+
+      const saga1 = new ResumeSaga(mockLogger);
+      const result1 = await saga1.run({
+        taskId: "task-1",
+        runId: "run-1",
+        repositoryPath: repo.path,
+        apiClient: mockApiClient,
+      });
+
+      // Run with coalesced (new format)
+      (
+        mockApiClient.fetchTaskRunLogs as ReturnType<typeof vi.fn>
+      ).mockResolvedValue(coalescedEntries);
+
+      const saga2 = new ResumeSaga(mockLogger);
+      const result2 = await saga2.run({
+        taskId: "task-1",
+        runId: "run-1",
+        repositoryPath: repo.path,
+        apiClient: mockApiClient,
+      });
+
+      expect(result1.success).toBe(true);
+      expect(result2.success).toBe(true);
+      if (!result1.success || !result2.success) return;
+
+      // Conversation structure should be identical
+      expect(result1.data.conversation.length).toBe(
+        result2.data.conversation.length,
+      );
+      for (let i = 0; i < result1.data.conversation.length; i++) {
+        expect(result1.data.conversation[i].role).toBe(
+          result2.data.conversation[i].role,
+        );
+        expect(result1.data.conversation[i].content).toEqual(
+          result2.data.conversation[i].content,
+        );
+      }
+    });
+
     it("tracks tool calls with results", async () => {
       (mockApiClient.getTaskRun as ReturnType<typeof vi.fn>).mockResolvedValue(
         createTaskRun(),

package/src/sagas/resume-saga.ts

@@ -244,7 +244,30 @@ export class ResumeSaga extends Saga<ResumeInput, ResumeOutput> {
             break;
           }
 
+          case "agent_message": {
+            const content = update.content as ContentBlock | undefined;
+            if (content) {
+              if (
+                content.type === "text" &&
+                currentAssistantContent.length > 0 &&
+                currentAssistantContent[currentAssistantContent.length - 1]
+                  .type === "text"
+              ) {
+                const lastBlock = currentAssistantContent[
+                  currentAssistantContent.length - 1
+                ] as { type: "text"; text: string };
+                lastBlock.text += (
+                  content as { type: "text"; text: string }
+                ).text;
+              } else {
+                currentAssistantContent.push(content);
+              }
+            }
+            break;
+          }
+
           case "agent_message_chunk": {
+            // Backward compatibility with older logs that have individual chunks
             const content = update.content as ContentBlock | undefined;
             if (content) {
               if (

package/src/sagas/test-fixtures.ts

@@ -205,6 +205,15 @@ export function createAgentChunk(text: string): StoredNotification {
   });
 }
 
+export function createAgentMessage(text: string): StoredNotification {
+  return createNotification("session/update", {
+    update: {
+      sessionUpdate: "agent_message",
+      content: { type: "text", text },
+    },
+  });
+}
+
 export function createToolCall(
   toolCallId: string,
   toolName: string,

package/src/session-log-writer.ts

@@ -16,9 +16,15 @@ export interface SessionLogWriterOptions {
   logger?: Logger;
 }
 
+interface ChunkBuffer {
+  text: string;
+  firstTimestamp: string;
+}
+
 interface SessionState {
   context: SessionContext;
   otelWriter?: OtelLogWriter;
+  chunkBuffer?: ChunkBuffer;
 }
 
 export class SessionLogWriter {
@@ -75,9 +81,28 @@ export class SessionLogWriter {
 
     try {
       const message = JSON.parse(line);
+      const timestamp = new Date().toISOString();
+
+      // Check if this is an agent_message_chunk event
+      if (this.isAgentMessageChunk(message)) {
+        const text = this.extractChunkText(message);
+        if (text) {
+          if (!session.chunkBuffer) {
+            session.chunkBuffer = { text, firstTimestamp: timestamp };
+          } else {
+            session.chunkBuffer.text += text;
+          }
+        }
+        // Don't emit chunk events
+        return;
+      }
+
+      // Non-chunk event: flush any buffered chunks first
+      this.emitCoalescedMessage(sessionId, session);
+
       const entry: StoredNotification = {
         type: "notification",
-        timestamp: new Date().toISOString(),
+        timestamp,
         notification: message,
       };
 
@@ -103,6 +128,9 @@ export class SessionLogWriter {
     const session = this.sessions.get(sessionId);
     if (!session) return;
 
+    // Emit any buffered chunks before flushing
+    this.emitCoalescedMessage(sessionId, session);
+
     if (session.otelWriter) {
       await session.otelWriter.flush();
     }
@@ -128,6 +156,58 @@ export class SessionLogWriter {
     }
   }
 
+  private isAgentMessageChunk(message: Record<string, unknown>): boolean {
+    if (message.method !== "session/update") return false;
+    const params = message.params as Record<string, unknown> | undefined;
+    const update = params?.update as Record<string, unknown> | undefined;
+    return update?.sessionUpdate === "agent_message_chunk";
+  }
+
+  private extractChunkText(message: Record<string, unknown>): string {
+    const params = message.params as Record<string, unknown> | undefined;
+    const update = params?.update as Record<string, unknown> | undefined;
+    const content = update?.content as
+      | { type: string; text?: string }
+      | undefined;
+    if (content?.type === "text" && content.text) {
+      return content.text;
+    }
+    return "";
+  }
+
+  private emitCoalescedMessage(sessionId: string, session: SessionState): void {
+    if (!session.chunkBuffer) return;
+
+    const { text, firstTimestamp } = session.chunkBuffer;
+    session.chunkBuffer = undefined;
+
+    const entry: StoredNotification = {
+      type: "notification",
+      timestamp: firstTimestamp,
+      notification: {
+        jsonrpc: "2.0",
+        method: "session/update",
+        params: {
+          update: {
+            sessionUpdate: "agent_message",
+            content: { type: "text", text },
+          },
+        },
+      },
+    };
+
+    if (session.otelWriter) {
+      session.otelWriter.emit({ notification: entry });
+    }
+
+    if (this.posthogAPI) {
+      const pending = this.pendingEntries.get(sessionId) ?? [];
+      pending.push(entry);
+      this.pendingEntries.set(sessionId, pending);
+      this.scheduleFlush(sessionId);
+    }
+  }
+
   private scheduleFlush(sessionId: string): void {
     const existing = this.flushTimeouts.get(sessionId);
     if (existing) clearTimeout(existing);

package/dist/server/bin.d.ts

@@ -1 +0,0 @@
-#!/usr/bin/env node