@posthog/agent 2.3.104 → 2.3.116

package/dist/types.d.ts

@@ -88,7 +88,8 @@ type LogLevel = "debug" | "info" | "warn" | "error";
 type OnLogCallback = (level: LogLevel, scope: string, message: string, data?: unknown) => void;
 interface PostHogAPIConfig {
     apiUrl: string;
-    getApiKey: () => string;
+    getApiKey: () => string | Promise<string>;
+    refreshApiKey?: () => string | Promise<string>;
     projectId: number;
     userAgent?: string;
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/agent",
-  "version": "2.3.104",
+  "version": "2.3.116",
   "repository": "https://github.com/PostHog/code",
   "description": "TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog",
   "exports": {

package/src/agent.ts

@@ -45,17 +45,17 @@ export class Agent {
     }
   }
 
-  private _configureLlmGateway(overrideUrl?: string): {
+  private async _configureLlmGateway(overrideUrl?: string): Promise<{
     gatewayUrl: string;
     apiKey: string;
-  } | null {
+  } | null> {
     if (!this.posthogAPI) {
       return null;
     }
 
     try {
       const gatewayUrl = overrideUrl ?? this.posthogAPI.getLlmGatewayUrl();
-      const apiKey = this.posthogAPI.getApiKey();
+      const apiKey = await this.posthogAPI.getApiKey();
 
       process.env.OPENAI_BASE_URL = `${gatewayUrl}/v1`;
       process.env.OPENAI_API_KEY = apiKey;
@@ -74,7 +74,7 @@ export class Agent {
     taskRunId: string,
     options: TaskExecutionOptions = {},
   ): Promise<InProcessAcpConnection> {
-    const gatewayConfig = this._configureLlmGateway(options.gatewayUrl);
+    const gatewayConfig = await this._configureLlmGateway(options.gatewayUrl);
     this.logger.info("Configured LLM gateway", {
       adapter: options.adapter,
     });
@@ -175,7 +175,7 @@ export class Agent {
 
   async cleanup(): Promise<void> {
     if (this.sessionLogWriter && this.taskRunId) {
-      await this.sessionLogWriter.flush(this.taskRunId);
+      await this.sessionLogWriter.flush(this.taskRunId, { coalesce: true });
     }
     await this.acpConnection?.cleanup();
   }

package/src/posthog-api.test.ts

@@ -0,0 +1,48 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { PostHogAPIClient } from "./posthog-api";
+
+const mockFetch = vi.fn();
+
+vi.stubGlobal("fetch", mockFetch);
+
+describe("PostHogAPIClient", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("refreshes once when fetching task run logs gets an auth failure", async () => {
+    const getApiKey = vi.fn().mockResolvedValue("stale-token");
+    const refreshApiKey = vi.fn().mockResolvedValue("fresh-token");
+    const client = new PostHogAPIClient({
+      apiUrl: "https://app.posthog.com",
+      getApiKey,
+      refreshApiKey,
+      projectId: 1,
+    });
+
+    mockFetch
+      .mockResolvedValueOnce({
+        ok: false,
+        status: 401,
+        statusText: "Unauthorized",
+      })
+      .mockResolvedValueOnce({
+        ok: true,
+        text: vi
+          .fn()
+          .mockResolvedValue(
+            `${JSON.stringify({ type: "notification", notification: { method: "foo" } })}\n`,
+          ),
+      });
+
+    const logs = await client.fetchTaskRunLogs({
+      id: "run-1",
+      task: "task-1",
+    } as never);
+
+    expect(logs).toHaveLength(1);
+    expect(getApiKey).toHaveBeenCalledTimes(1);
+    expect(refreshApiKey).toHaveBeenCalledTimes(1);
+    expect(mockFetch).toHaveBeenCalledTimes(2);
+  });
+});

package/src/posthog-api.ts

@@ -47,27 +47,63 @@ export class PostHogAPIClient {
     return host;
   }
 
-  private get headers(): Record<string, string> {
-    return {
-      Authorization: `Bearer ${this.config.getApiKey()}`,
-      "Content-Type": "application/json",
-      "User-Agent": this.config.userAgent ?? DEFAULT_USER_AGENT,
-    };
+  private isAuthFailure(status: number): boolean {
+    return status === 401 || status === 403;
   }
 
-  private async apiRequest<T>(
+  private async resolveApiKey(forceRefresh = false): Promise<string> {
+    if (forceRefresh && this.config.refreshApiKey) {
+      return this.config.refreshApiKey();
+    }
+
+    return this.config.getApiKey();
+  }
+
+  private async buildHeaders(
+    options: RequestInit,
+    forceRefresh = false,
+  ): Promise<Headers> {
+    const headers = new Headers(options.headers);
+    headers.set(
+      "Authorization",
+      `Bearer ${await this.resolveApiKey(forceRefresh)}`,
+    );
+    headers.set("Content-Type", "application/json");
+    headers.set("User-Agent", this.config.userAgent ?? DEFAULT_USER_AGENT);
+    return headers;
+  }
+
+  private async performRequest(
     endpoint: string,
-    options: RequestInit = {},
-  ): Promise<T> {
+    options: RequestInit,
+    forceRefresh = false,
+  ): Promise<Response> {
     const url = `${this.baseUrl}${endpoint}`;
 
-    const response = await fetch(url, {
+    return fetch(url, {
       ...options,
-      headers: {
-        ...this.headers,
-        ...options.headers,
-      },
+      headers: await this.buildHeaders(options, forceRefresh),
     });
+  }
+
+  private async performRequestWithRetry(
+    endpoint: string,
+    options: RequestInit = {},
+  ): Promise<Response> {
+    let response = await this.performRequest(endpoint, options);
+
+    if (!response.ok && this.isAuthFailure(response.status)) {
+      response = await this.performRequest(endpoint, options, true);
+    }
+
+    return response;
+  }
+
+  private async apiRequest<T>(
+    endpoint: string,
+    options: RequestInit = {},
+  ): Promise<T> {
+    const response = await this.performRequestWithRetry(endpoint, options);
 
     if (!response.ok) {
       let errorMessage: string;
@@ -87,8 +123,8 @@ export class PostHogAPIClient {
     return this.config.projectId;
   }
 
-  getApiKey(): string {
-    return this.config.getApiKey();
+  async getApiKey(forceRefresh = false): Promise<string> {
+    return this.resolveApiKey(forceRefresh);
   }
 
   getLlmGatewayUrl(): string {
@@ -228,12 +264,10 @@ export class PostHogAPIClient {
    */
   async fetchTaskRunLogs(taskRun: TaskRun): Promise<StoredEntry[]> {
     const teamId = this.getTeamId();
+    const endpoint = `/api/projects/${teamId}/tasks/${taskRun.task}/runs/${taskRun.id}/logs`;
 
     try {
-      const response = await fetch(
-        `${this.baseUrl}/api/projects/${teamId}/tasks/${taskRun.task}/runs/${taskRun.id}/logs`,
-        { headers: this.headers },
-      );
+      const response = await this.performRequestWithRetry(endpoint);
 
       if (!response.ok) {
         if (response.status === 404) {

package/src/server/agent-server.ts

@@ -523,7 +523,9 @@ export class AgentServer {
         // against async log persistence to object storage.
         let assistantMessage: string | undefined;
         try {
-          await this.session.logWriter.flush(this.session.payload.run_id);
+          await this.session.logWriter.flush(this.session.payload.run_id, {
+            coalesce: true,
+          });
           assistantMessage = this.session.logWriter.getFullAgentResponse(
             this.session.payload.run_id,
           );
@@ -747,6 +749,9 @@ export class AgentServer {
     });
 
     this.logger.info("Session initialized successfully");
+    this.logger.info(
+      `Agent version: ${this.config.version ?? packageJson.version}`,
+    );
 
     // Signal in_progress so the UI can start polling for updates
     this.posthogAPI
@@ -1111,7 +1116,9 @@ Important:
   ): Promise<void> {
     if (this.session?.payload.run_id === payload.run_id) {
       try {
-        await this.session.logWriter.flush(payload.run_id);
+        await this.session.logWriter.flush(payload.run_id, {
+          coalesce: true,
+        });
       } catch (error) {
         this.logger.warn("Failed to flush session logs before completion", {
           taskId: payload.task_id,
@@ -1270,7 +1277,7 @@ Important:
     }
 
     try {
-      await this.session.logWriter.flush(payload.run_id);
+      await this.session.logWriter.flush(payload.run_id, { coalesce: true });
     } catch (error) {
       this.logger.warn("Failed to flush logs before Slack relay", {
         taskId: payload.task_id,
@@ -1473,7 +1480,9 @@ Important:
     }
 
     try {
-      await this.session.logWriter.flush(this.session.payload.run_id);
+      await this.session.logWriter.flush(this.session.payload.run_id, {
+        coalesce: true,
+      });
     } catch (error) {
       this.logger.error("Failed to flush session logs", error);
     }

package/src/session-log-writer.test.ts

@@ -159,6 +159,222 @@ describe("SessionLogWriter", () => {
     });
   });
 
+  describe("_doFlush does not prematurely coalesce", () => {
+    it("does not coalesce buffered chunks during a timed flush", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      // Buffer some chunks (no non-chunk event to trigger coalescing)
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "Hello " },
+        }),
+      );
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "world" },
+        }),
+      );
+
+      // Flush without any non-chunk event arriving — simulates
+      // the 500ms debounce timer firing mid-stream
+      await logWriter.flush(sessionId);
+
+      // No entries should have been sent — chunks are still buffered
+      expect(mockAppendLog).not.toHaveBeenCalled();
+
+      // Now a non-chunk event arrives, triggering natural coalescing
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("usage_update", { used: 100 }),
+      );
+
+      await logWriter.flush(sessionId);
+
+      expect(mockAppendLog).toHaveBeenCalledTimes(1);
+      const entries: StoredNotification[] = mockAppendLog.mock.calls[0][2];
+      expect(entries).toHaveLength(2); // coalesced agent_message + usage_update
+      const coalesced = entries[0].notification;
+      expect(coalesced.params?.update).toEqual({
+        sessionUpdate: "agent_message",
+        content: { type: "text", text: "Hello world" },
+      });
+    });
+  });
+
+  describe("flushAll coalesces on shutdown", () => {
+    it("coalesces remaining chunks before flushing", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "partial response" },
+        }),
+      );
+
+      await logWriter.flushAll();
+
+      expect(mockAppendLog).toHaveBeenCalledTimes(1);
+      const entries: StoredNotification[] = mockAppendLog.mock.calls[0][2];
+      expect(entries).toHaveLength(1);
+      const coalesced = entries[0].notification;
+      expect(coalesced.params?.update).toEqual({
+        sessionUpdate: "agent_message",
+        content: { type: "text", text: "partial response" },
+      });
+    });
+  });
+
+  describe("flush with coalesce option", () => {
+    it("drains chunk buffer when coalesce is true", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "complete text" },
+        }),
+      );
+
+      await logWriter.flush(sessionId, { coalesce: true });
+
+      expect(mockAppendLog).toHaveBeenCalledTimes(1);
+      const entries: StoredNotification[] = mockAppendLog.mock.calls[0][2];
+      const coalesced = entries[0].notification;
+      expect(coalesced.params?.update).toEqual({
+        sessionUpdate: "agent_message",
+        content: { type: "text", text: "complete text" },
+      });
+    });
+
+    it("does not coalesce when coalesce is false", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "buffered" },
+        }),
+      );
+
+      await logWriter.flush(sessionId, { coalesce: false });
+
+      expect(mockAppendLog).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("direct agent_message supersedes chunks", () => {
+    it("discards buffered chunks when a direct agent_message arrives", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      // Buffer partial chunks
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "partial " },
+        }),
+      );
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "text" },
+        }),
+      );
+
+      // Direct agent_message arrives — authoritative full text
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message", {
+          content: { type: "text", text: "complete full response" },
+        }),
+      );
+
+      await logWriter.flush(sessionId);
+
+      expect(mockAppendLog).toHaveBeenCalledTimes(1);
+      const entries: StoredNotification[] = mockAppendLog.mock.calls[0][2];
+      // Only the direct agent_message — no coalesced partial entry
+      expect(entries).toHaveLength(1);
+      const coalesced = entries[0].notification;
+      expect(coalesced.params?.update).toEqual({
+        sessionUpdate: "agent_message",
+        content: { type: "text", text: "complete full response" },
+      });
+      expect(logWriter.getLastAgentMessage(sessionId)).toBe(
+        "complete full response",
+      );
+    });
+
+    it("is additive with earlier coalesced text in multi-message turns", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      // First assistant message: chunks coalesced by a tool_call event
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "first message" },
+        }),
+      );
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("tool_call", { toolCallId: "tc1" }),
+      );
+      // "first message" is now coalesced into currentTurnMessages
+
+      // Second assistant message arrives as direct agent_message
+      // (e.g., after tool result, no active chunk buffer)
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message", {
+          content: { type: "text", text: "second message" },
+        }),
+      );
+
+      const response = logWriter.getFullAgentResponse(sessionId);
+      // Both messages are preserved — direct message is additive
+      expect(response).toBe("first message\n\nsecond message");
+    });
+
+    it("persisted log does not contain stale entries when chunks are superseded", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      // Chunks buffered, then direct agent_message supersedes before coalescing
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message_chunk", {
+          content: { type: "text", text: "partial" },
+        }),
+      );
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message", {
+          content: { type: "text", text: "complete" },
+        }),
+      );
+
+      await logWriter.flush(sessionId);
+
+      expect(mockAppendLog).toHaveBeenCalledTimes(1);
+      const entries: StoredNotification[] = mockAppendLog.mock.calls[0][2];
+      // Only the direct agent_message — no coalesced partial entry
+      expect(entries).toHaveLength(1);
+      const persisted = entries[0].notification;
+      expect(persisted.params?.update).toEqual({
+        sessionUpdate: "agent_message",
+        content: { type: "text", text: "complete" },
+      });
+    });
+  });
+
   describe("register", () => {
     it("does not re-register existing sessions", () => {
       const sessionId = "s1";

package/src/session-log-writer.ts

@@ -54,9 +54,12 @@ export class SessionLogWriter {
   }
 
   async flushAll(): Promise<void> {
-    const sessionIds = [...this.sessions.keys()];
+    // Coalesce any in-progress chunk buffers before the final flush
+    // During normal operation, chunks are coalesced when the next non-chunk
+    // event arrives, but on shutdown there may be no subsequent event
     const flushPromises: Promise<void>[] = [];
-    for (const sessionId of sessionIds) {
+    for (const [sessionId, session] of this.sessions) {
+      this.emitCoalescedMessage(sessionId, session);
       flushPromises.push(this.flush(sessionId));
     }
     await Promise.all(flushPromises);
@@ -123,8 +126,14 @@ export class SessionLogWriter {
         return;
       }
 
-      // Non-chunk event: flush any buffered chunks first
-      this.emitCoalescedMessage(sessionId, session);
+      // Non-chunk event: flush any buffered chunks first.
+      // If this is a direct agent_message AND there are buffered chunks,
+      // the direct message supersedes the partial chunks
+      if (this.isDirectAgentMessage(message) && session.chunkBuffer) {
+        session.chunkBuffer = undefined;
+      } else {
+        this.emitCoalescedMessage(sessionId, session);
+      }
 
       const nonChunkAgentText = this.extractAgentMessageText(message);
       if (nonChunkAgentText) {
@@ -155,7 +164,17 @@ export class SessionLogWriter {
     }
   }
 
-  async flush(sessionId: string): Promise<void> {
+  async flush(
+    sessionId: string,
+    { coalesce = false }: { coalesce?: boolean } = {},
+  ): Promise<void> {
+    if (coalesce) {
+      const session = this.sessions.get(sessionId);
+      if (session) {
+        this.emitCoalescedMessage(sessionId, session);
+      }
+    }
+
     // Serialize flushes per session
     const prev = this.flushQueues.get(sessionId) ?? Promise.resolve();
     const next = prev.catch(() => {}).then(() => this._doFlush(sessionId));
@@ -175,9 +194,6 @@ export class SessionLogWriter {
       return;
     }
 
-    // Emit any buffered chunks before flushing
-    this.emitCoalescedMessage(sessionId, session);
-
     const pending = this.pendingEntries.get(sessionId);
     if (!this.posthogAPI || !pending?.length) {
       return;
@@ -231,11 +247,21 @@ export class SessionLogWriter {
     }
   }
 
-  private isAgentMessageChunk(message: Record<string, unknown>): boolean {
-    if (message.method !== "session/update") return false;
+  private getSessionUpdateType(
+    message: Record<string, unknown>,
+  ): string | undefined {
+    if (message.method !== "session/update") return undefined;
     const params = message.params as Record<string, unknown> | undefined;
     const update = params?.update as Record<string, unknown> | undefined;
-    return update?.sessionUpdate === "agent_message_chunk";
+    return update?.sessionUpdate as string | undefined;
+  }
+
+  private isDirectAgentMessage(message: Record<string, unknown>): boolean {
+    return this.getSessionUpdateType(message) === "agent_message";
+  }
+
+  private isAgentMessageChunk(message: Record<string, unknown>): boolean {
+    return this.getSessionUpdateType(message) === "agent_message_chunk";
   }
 
   private extractChunkText(message: Record<string, unknown>): string {
@@ -290,6 +316,17 @@ export class SessionLogWriter {
   getFullAgentResponse(sessionId: string): string | undefined {
     const session = this.sessions.get(sessionId);
     if (!session || session.currentTurnMessages.length === 0) return undefined;
+
+    if (session.chunkBuffer) {
+      this.logger.warn(
+        "getFullAgentResponse called with non-empty chunk buffer",
+        {
+          sessionId,
+          bufferedLength: session.chunkBuffer.text.length,
+        },
+      );
+    }
+
     return session.currentTurnMessages.join("\n\n");
   }
 

package/src/types.ts

@@ -126,7 +126,8 @@ export type OnLogCallback = (
 
 export interface PostHogAPIConfig {
   apiUrl: string;
-  getApiKey: () => string;
+  getApiKey: () => string | Promise<string>;
+  refreshApiKey?: () => string | Promise<string>;
   projectId: number;
   userAgent?: string;
 }