@posthog/agent 2.1.115 → 2.1.120

package/src/server/question-relay.test.ts

@@ -0,0 +1,343 @@
+import { type SetupServerApi, setupServer } from "msw/node";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { createTestRepo, type TestRepo } from "../test/fixtures/api.js";
+import { createPostHogHandlers } from "../test/mocks/msw-handlers.js";
+import { AgentServer } from "./agent-server.js";
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+const TEST_PAYLOAD = {
+  run_id: "test-run-id",
+  task_id: "test-task-id",
+  team_id: 1,
+  user_id: 1,
+  distinct_id: "test-distinct-id",
+  mode: "interactive" as const,
+};
+
+const QUESTION_META = {
+  twigToolKind: "question",
+  questions: [
+    {
+      question: "Which license should I use?",
+      options: [
+        { label: "MIT", description: "Permissive license" },
+        { label: "Apache 2.0", description: "Patent grant included" },
+        { label: "GPL v3", description: "Copyleft license" },
+      ],
+    },
+  ],
+};
+
+describe("Question relay", () => {
+  let repo: TestRepo;
+  let server: any;
+  let mswServer: SetupServerApi;
+  const port = 3098;
+
+  beforeEach(async () => {
+    repo = await createTestRepo("question-relay");
+    mswServer = setupServer(
+      ...createPostHogHandlers({ baseUrl: "http://localhost:8000" }),
+    );
+    mswServer.listen({ onUnhandledRequest: "bypass" });
+
+    server = new AgentServer({
+      port,
+      jwtPublicKey: "unused-in-unit-tests",
+      repositoryPath: repo.path,
+      apiUrl: "http://localhost:8000",
+      apiKey: "test-api-key",
+      projectId: 1,
+      mode: "interactive",
+      taskId: "test-task-id",
+      runId: "test-run-id",
+    });
+  });
+
+  afterEach(async () => {
+    mswServer.close();
+    await repo.cleanup();
+  });
+
+  describe("isQuestionMeta", () => {
+    it.each([
+      ["null", null],
+      ["undefined", undefined],
+      ["number", 42],
+      ["string", "not a question"],
+      ["object without question field", { options: [] }],
+      ["object with non-string question", { question: 123 }],
+      ["object with non-array options", { question: "Q?", options: "bad" }],
+      [
+        "object with invalid option items",
+        { question: "Q?", options: [{ notLabel: "x" }] },
+      ],
+    ])("rejects %s", (_label, value) => {
+      expect(server.isQuestionMeta(value)).toBe(false);
+    });
+
+    it.each([
+      [
+        "question with options",
+        {
+          question: "Pick one",
+          options: [{ label: "A", description: "desc" }, { label: "B" }],
+        },
+      ],
+      ["question without options", { question: "What do you think?" }],
+      ["question with empty options", { question: "Confirm?", options: [] }],
+    ])("accepts %s", (_label, value) => {
+      expect(server.isQuestionMeta(value)).toBe(true);
+    });
+  });
+
+  describe("getFirstQuestionMeta", () => {
+    it.each([
+      ["null meta", null],
+      ["undefined meta", undefined],
+      ["meta without questions", { other: "field" }],
+      ["meta with empty questions array", { questions: [] }],
+      ["meta with non-array questions", { questions: "not-array" }],
+    ])("returns null for %s", (_label, meta) => {
+      expect(server.getFirstQuestionMeta(meta)).toBeNull();
+    });
+
+    it("returns first question from valid meta", () => {
+      const result = server.getFirstQuestionMeta(QUESTION_META);
+      expect(result).toEqual(QUESTION_META.questions[0]);
+    });
+  });
+
+  describe("relaySlackQuestion", () => {
+    it("relays formatted question with options via posthogAPI", () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.relaySlackQuestion(TEST_PAYLOAD, QUESTION_META);
+
+      expect(relaySpy).toHaveBeenCalledOnce();
+      const [taskId, runId, message] = relaySpy.mock.calls[0];
+      expect(taskId).toBe("test-task-id");
+      expect(runId).toBe("test-run-id");
+      expect(message).toContain("*Which license should I use?*");
+      expect(message).toContain("1. *MIT*");
+      expect(message).toContain("Permissive license");
+      expect(message).toContain("2. *Apache 2.0*");
+      expect(message).toContain("3. *GPL v3*");
+      expect(message).toContain("Reply in this thread");
+    });
+
+    it("sets questionRelayedToSlack flag", () => {
+      vi.spyOn(server.posthogAPI, "relayMessage").mockResolvedValue(undefined);
+
+      server.relaySlackQuestion(TEST_PAYLOAD, QUESTION_META);
+      expect(server.questionRelayedToSlack).toBe(true);
+    });
+
+    it("does not relay when meta has no valid question", () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.relaySlackQuestion(TEST_PAYLOAD, { twigToolKind: "question" });
+      expect(server.questionRelayedToSlack).toBe(false);
+      expect(relaySpy).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("createCloudClient requestPermission", () => {
+    const ALLOW_OPTIONS = [
+      { kind: "allow_once", optionId: "allow", name: "Allow" },
+    ];
+
+    describe("with TWIG_INTERACTION_ORIGIN=slack", () => {
+      beforeEach(() => {
+        process.env.TWIG_INTERACTION_ORIGIN = "slack";
+      });
+
+      afterEach(() => {
+        delete process.env.TWIG_INTERACTION_ORIGIN;
+      });
+
+      it("returns cancelled with relay message for question tool", async () => {
+        vi.spyOn(server.posthogAPI, "relayMessage").mockResolvedValue(
+          undefined,
+        );
+        const client = server.createCloudClient(TEST_PAYLOAD);
+
+        const result = await client.requestPermission({
+          options: ALLOW_OPTIONS,
+          toolCall: { _meta: QUESTION_META },
+        });
+
+        expect(result.outcome.outcome).toBe("cancelled");
+        expect(result._meta?.message).toContain("relayed to the Slack thread");
+        expect(result._meta?.message).toContain("Do NOT re-ask the question");
+      });
+
+      it("auto-approves non-question tools", async () => {
+        const client = server.createCloudClient(TEST_PAYLOAD);
+
+        const result = await client.requestPermission({
+          options: ALLOW_OPTIONS,
+          toolCall: { _meta: { twigToolKind: "bash" } },
+        });
+
+        expect(result.outcome.outcome).toBe("selected");
+      });
+
+      it("auto-approves tools without meta", async () => {
+        const client = server.createCloudClient(TEST_PAYLOAD);
+
+        const result = await client.requestPermission({
+          options: ALLOW_OPTIONS,
+          toolCall: { _meta: null },
+        });
+
+        expect(result.outcome.outcome).toBe("selected");
+      });
+    });
+
+    describe("without TWIG_INTERACTION_ORIGIN", () => {
+      beforeEach(() => {
+        delete process.env.TWIG_INTERACTION_ORIGIN;
+      });
+
+      it("auto-approves question tools (no Slack relay)", async () => {
+        const client = server.createCloudClient(TEST_PAYLOAD);
+
+        const result = await client.requestPermission({
+          options: ALLOW_OPTIONS,
+          toolCall: { _meta: QUESTION_META },
+        });
+
+        expect(result.outcome.outcome).toBe("selected");
+      });
+    });
+  });
+
+  describe("relayAgentResponse duplicate suppression", () => {
+    it("skips relay when questionRelayedToSlack is set", async () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.session = {
+        payload: TEST_PAYLOAD,
+        logWriter: {
+          flush: vi.fn().mockResolvedValue(undefined),
+          getLastAgentMessage: vi.fn().mockReturnValue("agent response"),
+          isRegistered: vi.fn().mockReturnValue(true),
+        },
+      };
+
+      server.questionRelayedToSlack = true;
+      await server.relayAgentResponse(TEST_PAYLOAD);
+
+      expect(server.questionRelayedToSlack).toBe(false);
+      expect(relaySpy).not.toHaveBeenCalled();
+    });
+
+    it("relays normally when questionRelayedToSlack is not set", async () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.session = {
+        payload: TEST_PAYLOAD,
+        logWriter: {
+          flush: vi.fn().mockResolvedValue(undefined),
+          getLastAgentMessage: vi.fn().mockReturnValue("agent response"),
+          isRegistered: vi.fn().mockReturnValue(true),
+        },
+      };
+
+      server.questionRelayedToSlack = false;
+      await server.relayAgentResponse(TEST_PAYLOAD);
+
+      expect(relaySpy).toHaveBeenCalledWith(
+        "test-task-id",
+        "test-run-id",
+        "agent response",
+      );
+    });
+
+    it("does not relay when no agent message is available", async () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.session = {
+        payload: TEST_PAYLOAD,
+        logWriter: {
+          flush: vi.fn().mockResolvedValue(undefined),
+          getLastAgentMessage: vi.fn().mockReturnValue(null),
+          isRegistered: vi.fn().mockReturnValue(true),
+        },
+      };
+
+      server.questionRelayedToSlack = false;
+      await server.relayAgentResponse(TEST_PAYLOAD);
+
+      expect(relaySpy).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("sendInitialTaskMessage prompt source", () => {
+    it("uses run state initial_prompt_override when present", async () => {
+      vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({
+        id: "test-task-id",
+        title: "t",
+        description: "original task description",
+      } as any);
+      vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({
+        id: "test-run-id",
+        task: "test-task-id",
+        state: { initial_prompt_override: "override instruction" },
+      } as any);
+
+      const promptSpy = vi.fn().mockResolvedValue({ stopReason: "max_tokens" });
+      server.session = {
+        payload: TEST_PAYLOAD,
+        acpSessionId: "acp-session",
+        clientConnection: { prompt: promptSpy },
+      };
+
+      await server.sendInitialTaskMessage(TEST_PAYLOAD);
+
+      expect(promptSpy).toHaveBeenCalledWith({
+        sessionId: "acp-session",
+        prompt: [{ type: "text", text: "override instruction" }],
+      });
+    });
+
+    it("falls back to task description when override is missing", async () => {
+      vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({
+        id: "test-task-id",
+        title: "t",
+        description: "original task description",
+      } as any);
+      vi.spyOn(server.posthogAPI, "getTaskRun").mockResolvedValue({
+        id: "test-run-id",
+        task: "test-task-id",
+        state: {},
+      } as any);
+
+      const promptSpy = vi.fn().mockResolvedValue({ stopReason: "max_tokens" });
+      server.session = {
+        payload: TEST_PAYLOAD,
+        acpSessionId: "acp-session",
+        clientConnection: { prompt: promptSpy },
+      };
+
+      await server.sendInitialTaskMessage(TEST_PAYLOAD);
+
+      expect(promptSpy).toHaveBeenCalledWith({
+        sessionId: "acp-session",
+        prompt: [{ type: "text", text: "original task description" }],
+      });
+    });
+  });
+});

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

@@ -137,6 +137,25 @@ describe("SessionLogWriter", () => {
         sessionUpdate: "agent_message",
         content: { type: "text", text: "Hello world" },
       });
+      expect(logWriter.getLastAgentMessage(sessionId)).toBe("Hello world");
+    });
+
+    it("tracks direct agent_message updates", async () => {
+      const sessionId = "s1";
+      logWriter.register(sessionId, { taskId: "t1", runId: sessionId });
+
+      logWriter.appendRawLine(
+        sessionId,
+        makeSessionUpdate("agent_message", {
+          content: { type: "text", text: "Pick MIT or Apache" },
+        }),
+      );
+
+      await logWriter.flush(sessionId);
+
+      expect(logWriter.getLastAgentMessage(sessionId)).toBe(
+        "Pick MIT or Apache",
+      );
     });
   });
 

package/src/session-log-writer.ts

@@ -22,6 +22,7 @@ interface ChunkBuffer {
 interface SessionState {
   context: SessionContext;
   chunkBuffer?: ChunkBuffer;
+  lastAgentMessage?: string;
 }
 
 export class SessionLogWriter {
@@ -50,11 +51,15 @@ export class SessionLogWriter {
 
   async flushAll(): Promise<void> {
     const sessionIds = [...this.sessions.keys()];
-    const pendingCounts = sessionIds.map((id) => ({
-      id,
-      pending: this.pendingEntries.get(id)?.length ?? 0,
-      messages: this.messageCounts.get(id) ?? 0,
-    }));
+    const pendingCounts = sessionIds.map((id) => {
+      const session = this.sessions.get(id);
+      return {
+        taskId: session?.context.taskId,
+        runId: session?.context.runId,
+        pending: this.pendingEntries.get(id)?.length ?? 0,
+        messages: this.messageCounts.get(id) ?? 0,
+      };
+    });
     this.logger.info("flushAll called", {
       sessions: sessionIds.length,
       pending: pendingCounts,
@@ -73,8 +78,8 @@ export class SessionLogWriter {
     }
 
     this.logger.info("Session registered", {
-      sessionId,
       taskId: context.taskId,
+      runId: context.runId,
     });
     this.sessions.set(sessionId, { context });
 
@@ -113,7 +118,11 @@ export class SessionLogWriter {
     const count = (this.messageCounts.get(sessionId) ?? 0) + 1;
     this.messageCounts.set(sessionId, count);
     if (count % 10 === 1) {
-      this.logger.info("Messages received", { count, sessionId });
+      this.logger.info("Messages received", {
+        count,
+        taskId: session.context.taskId,
+        runId: session.context.runId,
+      });
     }
 
     try {
@@ -137,6 +146,11 @@ export class SessionLogWriter {
       // Non-chunk event: flush any buffered chunks first
       this.emitCoalescedMessage(sessionId, session);
 
+      const nonChunkAgentText = this.extractAgentMessageText(message);
+      if (nonChunkAgentText) {
+        session.lastAgentMessage = nonChunkAgentText;
+      }
+
       const entry: StoredNotification = {
         type: "notification",
         timestamp,
@@ -153,7 +167,8 @@ export class SessionLogWriter {
       }
     } catch {
       this.logger.warn("Failed to parse raw line for persistence", {
-        sessionId,
+        taskId: session.context.taskId,
+        runId: session.context.runId,
         lineLength: line.length,
       });
     }
@@ -172,7 +187,8 @@ export class SessionLogWriter {
     const pending = this.pendingEntries.get(sessionId);
     if (!this.posthogAPI || !pending?.length) {
       this.logger.info("flush: nothing to persist", {
-        sessionId,
+        taskId: session.context.taskId,
+        runId: session.context.runId,
         hasPosthogAPI: !!this.posthogAPI,
         pendingCount: pending?.length ?? 0,
       });
@@ -196,7 +212,8 @@ export class SessionLogWriter {
       );
       this.retryCounts.set(sessionId, 0);
       this.logger.info("Flushed session logs", {
-        sessionId,
+        taskId: session.context.taskId,
+        runId: session.context.runId,
         entryCount: pending.length,
       });
     } catch (error) {
@@ -206,7 +223,11 @@ export class SessionLogWriter {
       if (retryCount >= SessionLogWriter.MAX_FLUSH_RETRIES) {
         this.logger.error(
           `Dropping ${pending.length} session log entries after ${retryCount} failed flush attempts`,
-          { sessionId, error },
+          {
+            taskId: session.context.taskId,
+            runId: session.context.runId,
+            error,
+          },
         );
         this.retryCounts.set(sessionId, 0);
       } else {
@@ -245,6 +266,7 @@ export class SessionLogWriter {
 
     const { text, firstTimestamp } = session.chunkBuffer;
     session.chunkBuffer = undefined;
+    session.lastAgentMessage = text;
 
     const entry: StoredNotification = {
       type: "notification",
@@ -271,6 +293,39 @@ export class SessionLogWriter {
     }
   }
 
+  getLastAgentMessage(sessionId: string): string | undefined {
+    return this.sessions.get(sessionId)?.lastAgentMessage;
+  }
+
+  private extractAgentMessageText(
+    message: Record<string, unknown>,
+  ): string | null {
+    if (message.method !== "session/update") {
+      return null;
+    }
+
+    const params = message.params as Record<string, unknown> | undefined;
+    const update = params?.update as Record<string, unknown> | undefined;
+    if (update?.sessionUpdate !== "agent_message") {
+      return null;
+    }
+
+    const content = update.content as
+      | { type?: string; text?: string }
+      | undefined;
+    if (content?.type === "text" && typeof content.text === "string") {
+      const trimmed = content.text.trim();
+      return trimmed.length > 0 ? trimmed : null;
+    }
+
+    if (typeof update.message === "string") {
+      const trimmed = update.message.trim();
+      return trimmed.length > 0 ? trimmed : null;
+    }
+
+    return null;
+  }
+
   private scheduleFlush(sessionId: string): void {
     const existing = this.flushTimeouts.get(sessionId);
     if (existing) clearTimeout(existing);
@@ -316,7 +371,12 @@ export class SessionLogWriter {
     try {
       fs.appendFileSync(logPath, `${JSON.stringify(entry)}\n`);
     } catch (error) {
-      this.logger.warn("Failed to write to local cache", { logPath, error });
+      this.logger.warn("Failed to write to local cache", {
+        taskId: session.context.taskId,
+        runId: session.context.runId,
+        logPath,
+        error,
+      });
     }
   }
 }

package/src/test/mocks/msw-handlers.ts

@@ -5,6 +5,7 @@ type AnyHttpResponse = Response | ReturnType<typeof HttpResponse.json>;
 export interface PostHogHandlersOptions {
   baseUrl?: string;
   onAppendLog?: (entries: unknown[]) => void;
+  getTask?: () => unknown;
   getTaskRun?: () => unknown;
   appendLogResponse?: () => AnyHttpResponse;
 }
@@ -13,6 +14,7 @@ export function createPostHogHandlers(options: PostHogHandlersOptions = {}) {
   const {
     baseUrl = "http://localhost:8000",
     onAppendLog,
+    getTask,
     getTaskRun,
     appendLogResponse,
   } = options;
@@ -33,13 +35,35 @@ export function createPostHogHandlers(options: PostHogHandlersOptions = {}) {
       },
     ),
 
+    // GET /tasks/:taskId - Fetch task details
+    http.get(`${baseUrl}/api/projects/:projectId/tasks/:taskId/`, () => {
+      const task = getTask?.() ?? {
+        id: "test-task-id",
+        title: "Test task",
+        description: null,
+        origin_product: "user_created",
+        repository: "test/repo",
+        created_at: new Date().toISOString(),
+        updated_at: new Date().toISOString(),
+      };
+      return HttpResponse.json(task);
+    }),
+
     // GET /runs/:runId - Fetch task run details
     http.get(
-      `${baseUrl}/api/projects/:projectId/tasks/:taskId/runs/:runId`,
+      `${baseUrl}/api/projects/:projectId/tasks/:taskId/runs/:runId/`,
       () => {
         const taskRun = getTaskRun?.() ?? { log_url: "" };
         return HttpResponse.json(taskRun);
       },
     ),
+
+    // PATCH /runs/:runId - Update task run
+    http.patch(
+      `${baseUrl}/api/projects/:projectId/tasks/:taskId/runs/:runId/`,
+      () => {
+        return HttpResponse.json({});
+      },
+    ),
   ];
 }