feat(agent): echo initiating message_id on relay_message (#3434)

VojtechBartos · web-flow · commit c14064effc71 · 2026-07-20T12:10:45.000+02:00
diff --git a/packages/agent/src/posthog-api.test.ts b/packages/agent/src/posthog-api.test.ts
@@ -78,6 +78,55 @@ describe("PostHogAPIClient", () => {
     );
   });
 
+  it.each([
+    [
+      "includes message_id and text_parts when provided",
+      ["part one", "final answer"],
+      "msg-1",
+      {
+        text: "final answer",
+        text_parts: ["part one", "final answer"],
+        message_id: "msg-1",
+      },
+    ],
+    [
+      "omits optional fields when unknown",
+      undefined,
+      undefined,
+      { text: "final answer" },
+    ],
+  ])(
+    "relay_message body %s",
+    async (_label, textParts, messageId, expectedBody) => {
+      const client = new PostHogAPIClient({
+        apiUrl: "https://app.posthog.com",
+        getApiKey: vi.fn().mockResolvedValue("token"),
+        projectId: 7,
+      });
+
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: vi.fn().mockResolvedValue({ status: "ok" }),
+      });
+
+      await client.relayMessage(
+        "task-1",
+        "run-1",
+        "final answer",
+        textParts,
+        messageId,
+      );
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        "https://app.posthog.com/api/projects/7/tasks/task-1/runs/run-1/relay_message/",
+        expect.objectContaining({
+          method: "POST",
+          body: JSON.stringify(expectedBody),
+        }),
+      );
+    },
+  );
+
   it("returns only the artifacts created by the current upload request", async () => {
     const client = new PostHogAPIClient({
       apiUrl: "https://app.posthog.com",
diff --git a/packages/agent/src/posthog-api.ts b/packages/agent/src/posthog-api.ts
@@ -230,15 +230,23 @@ export class PostHogAPIClient {
     runId: string,
     text: string,
     textParts?: string[],
+    messageId?: string,
   ): Promise<void> {
     const teamId = this.getTeamId();
     // Send `text_parts` alongside the joined `text` so backends that understand
     // the new schema can pick just the post-last-tool-use answer, while older
     // backends still get the flat `text` field they already handle.
-    const body: { text: string; text_parts?: string[] } = { text };
+    // `message_id` correlates the relay with the user message that initiated
+    // the turn; it is omitted when no message id is known (e.g. boot prompt).
+    const body: { text: string; text_parts?: string[]; message_id?: string } = {
+      text,
+    };
     if (textParts && textParts.length > 0) {
       body.text_parts = textParts;
     }
+    if (messageId) {
+      body.message_id = messageId;
+    }
     await this.apiRequest<{ status: string }>(
       `/api/projects/${teamId}/tasks/${taskId}/runs/${runId}/relay_message/`,
       {
diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts
@@ -2306,6 +2306,111 @@ describe("AgentServer HTTP Mode", () => {
       expect(body.result?.stopReason).toBe("end_turn");
       expect(prompt).toHaveBeenCalledTimes(2);
     }, 20000);
+
+    // Shared plumbing for the relay-echo tests: install a controllable
+    // prompt, stub the log writer so relayAgentResponse has an answer to
+    // relay, and spy on the relay_message client call.
+    const setupRelayEchoServer = async (
+      prompt: () => Promise<{ stopReason: string }>,
+    ) => {
+      const s = createServer();
+      await s.start();
+      const serverInternals = s as unknown as {
+        session: {
+          clientConnection: { prompt: typeof prompt };
+          logWriter: {
+            getFullAgentResponse: (runId: string) => string | undefined;
+            getAgentResponseParts: (runId: string) => string[];
+          };
+        };
+        posthogAPI: PostHogAPIClient;
+      };
+      serverInternals.session.clientConnection.prompt = prompt;
+      vi.spyOn(
+        serverInternals.session.logWriter,
+        "getFullAgentResponse",
+      ).mockReturnValue("final answer");
+      vi.spyOn(
+        serverInternals.session.logWriter,
+        "getAgentResponseParts",
+      ).mockReturnValue(["final answer"]);
+      const relaySpy = vi
+        .spyOn(serverInternals.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      const token = createToken();
+      const send = (messageId?: string) =>
+        fetch(`http://localhost:${port}/command`, {
+          method: "POST",
+          headers: {
+            Authorization: `Bearer ${token}`,
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify({
+            jsonrpc: "2.0",
+            id: messageId ?? "no-id",
+            method: "user_message",
+            params: {
+              content: "do the thing",
+              ...(messageId ? { messageId } : {}),
+            },
+          }),
+        });
+
+      return { relaySpy, send };
+    };
+
+    it("echoes each turn's own initiating messageId on relay_message", async () => {
+      const pendingTurns: Array<(result: { stopReason: string }) => void> = [];
+      const prompt = vi.fn(
+        () =>
+          new Promise<{ stopReason: string }>((resolve) => {
+            pendingTurns.push(resolve);
+          }),
+      );
+      const { relaySpy, send } = await setupRelayEchoServer(prompt);
+
+      // The second message lands while the first turn is still in flight;
+      // each relay carries its own sender's id, not the first turn's.
+      const first = send("m-first");
+      await vi.waitFor(() => expect(prompt).toHaveBeenCalledTimes(1));
+      const second = send("m-second");
+      await vi.waitFor(() => expect(prompt).toHaveBeenCalledTimes(2));
+
+      pendingTurns[0]({ stopReason: "end_turn" });
+      await first;
+      await vi.waitFor(() => expect(relaySpy).toHaveBeenCalledTimes(1));
+      expect(relaySpy.mock.calls[0][4]).toBe("m-first");
+
+      pendingTurns[1]({ stopReason: "end_turn" });
+      await second;
+      await vi.waitFor(() => expect(relaySpy).toHaveBeenCalledTimes(2));
+      expect(relaySpy.mock.calls[1][4]).toBe("m-second");
+
+      // A message without an id relays without correlation (backward
+      // compatible with backends that don't know message_id).
+      relaySpy.mockClear();
+      const anonymous = send(undefined);
+      await vi.waitFor(() => expect(prompt).toHaveBeenCalledTimes(3));
+      pendingTurns[2]({ stopReason: "end_turn" });
+      await anonymous;
+      await vi.waitFor(() => expect(relaySpy).toHaveBeenCalledTimes(1));
+      expect(relaySpy.mock.calls[0][4]).toBeUndefined();
+    }, 20000);
+
+    it("does not leak a failed turn's messageId into the next turn", async () => {
+      const prompt = vi
+        .fn(async () => ({ stopReason: "end_turn" }))
+        .mockRejectedValueOnce(new Error("sdk connection lost"));
+      const { relaySpy, send } = await setupRelayEchoServer(prompt);
+
+      await send("m-fail");
+      expect(relaySpy).not.toHaveBeenCalled();
+
+      await send("m-next");
+      await vi.waitFor(() => expect(relaySpy).toHaveBeenCalledTimes(1));
+      expect(relaySpy.mock.calls[0][4]).toBe("m-next");
+    }, 20000);
   });
 
   describe("404 handling", () => {
diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts
@@ -1091,9 +1091,11 @@ export class AgentServer {
 
         if (result.stopReason === "end_turn") {
           // Relay the response to Slack. For follow-ups this is the primary
-          // delivery path — the HTTP caller only handles reactions.
-          this.relayAgentResponse(this.session.payload).catch((err) =>
-            this.logger.debug("Failed to relay follow-up response", err),
+          // delivery path — the HTTP caller only handles reactions. Echo the
+          // initiating message's id so the backend can attribute the answer.
+          this.relayAgentResponse(this.session.payload, messageId).catch(
+            (err) =>
+              this.logger.debug("Failed to relay follow-up response", err),
           );
         }
 
@@ -3818,7 +3820,10 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
     };
   }
 
-  private async relayAgentResponse(payload: JwtPayload): Promise<void> {
+  private async relayAgentResponse(
+    payload: JwtPayload,
+    messageId?: string,
+  ): Promise<void> {
     if (!this.session) {
       return;
     }
@@ -3862,6 +3867,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
         payload.run_id,
         message,
         messageParts,
+        messageId,
       );
     } catch (error) {
       this.logger.debug("Failed to relay initial agent response to Slack", {
diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts
@@ -23,7 +23,10 @@ interface TestableAgentServer {
   };
   questionRelayedToSlack: boolean;
   session: unknown;
-  relayAgentResponse: (payload: Record<string, unknown>) => Promise<void>;
+  relayAgentResponse: (
+    payload: Record<string, unknown>,
+    messageId?: string,
+  ) => Promise<void>;
   sendInitialTaskMessage: (payload: Record<string, unknown>) => Promise<void>;
 }
 
@@ -555,6 +558,34 @@ describe("Question relay", () => {
         "test-run-id",
         "agent response",
         ["first part", "agent response"],
+        undefined,
+      );
+    });
+
+    it("passes the initiating message id through to relayMessage", async () => {
+      const relaySpy = vi
+        .spyOn(server.posthogAPI, "relayMessage")
+        .mockResolvedValue(undefined);
+
+      server.session = {
+        payload: TEST_PAYLOAD,
+        logWriter: {
+          flush: vi.fn().mockResolvedValue(undefined),
+          getFullAgentResponse: vi.fn().mockReturnValue("agent response"),
+          getAgentResponseParts: vi.fn().mockReturnValue(["agent response"]),
+          isRegistered: vi.fn().mockReturnValue(true),
+        },
+      };
+
+      server.questionRelayedToSlack = false;
+      await server.relayAgentResponse(TEST_PAYLOAD, "msg-123");
+
+      expect(relaySpy).toHaveBeenCalledWith(
+        "test-task-id",
+        "test-run-id",
+        "agent response",
+        ["agent response"],
+        "msg-123",
       );
     });
 
