From f11ab9f5971e20bf5be0fc2d933f6f60ef3132b4 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Tue, 14 Jul 2026 19:40:02 +0200 Subject: [PATCH] feat(agent): echo initiating message_id on relay_message Pass the id of the user_message command that initiated a turn from the handler closure to relayAgentResponse and include it as an optional message_id field in the relay_message POST body. Each turn's relay carries its own initiating message's id; the field is omitted when no id is known (e.g. the boot prompt), so older backends are unaffected. --- packages/agent/src/posthog-api.test.ts | 49 ++++++++ packages/agent/src/posthog-api.ts | 10 +- .../agent/src/server/agent-server.test.ts | 105 ++++++++++++++++++ packages/agent/src/server/agent-server.ts | 14 ++- .../agent/src/server/question-relay.test.ts | 33 +++++- 5 files changed, 205 insertions(+), 6 deletions(-) diff --git a/packages/agent/src/posthog-api.test.ts b/packages/agent/src/posthog-api.test.ts index 16f2b37061..651f539b3c 100644 --- 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 index 5ad0d1c1e5..9c983f84f5 100644 --- a/packages/agent/src/posthog-api.ts +++ b/packages/agent/src/posthog-api.ts @@ -238,15 +238,23 @@ export class PostHogAPIClient { runId: string, text: string, textParts?: string[], + messageId?: string, ): Promise { 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 index 89f5c2c15c..f15c2feb9a 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -1828,6 +1828,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 index 1c14af21ae..176b5c17a1 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -989,9 +989,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), ); } @@ -3635,7 +3637,10 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} }; } - private async relayAgentResponse(payload: JwtPayload): Promise { + private async relayAgentResponse( + payload: JwtPayload, + messageId?: string, + ): Promise { if (!this.session) { return; } @@ -3679,6 +3684,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 index 3f8648f158..558bd1b284 100644 --- 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) => Promise; + relayAgentResponse: ( + payload: Record, + messageId?: string, + ) => Promise; sendInitialTaskMessage: (payload: Record) => Promise; } @@ -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", ); });