diff --git a/apps/backend/src/lib/ai/ai-query-handlers.test.ts b/apps/backend/src/lib/ai/ai-query-handlers.test.ts index a7a5df424..4a4d866fe 100644 --- a/apps/backend/src/lib/ai/ai-query-handlers.test.ts +++ b/apps/backend/src/lib/ai/ai-query-handlers.test.ts @@ -49,10 +49,12 @@ vi.mock("ai", async (importOriginal) => { }); import { logAiQuery } from "@/lib/ai/loggers/ai-query-logger"; +import { logIfMcpToolCall } from "@/lib/ai/loggers/mcp-call-logger"; import { handleStreamMode } from "./ai-query-handlers"; import type { CommonLogFields } from "./types"; const mockedLogAiQuery = vi.mocked(logAiQuery); +const mockedLogIfMcpToolCall = vi.mocked(logIfMcpToolCall); const baseCommon: CommonLogFields = { correlationId: "test-correlation-id", @@ -89,6 +91,7 @@ function makeCtx(): Parameters[0] { beforeEach(() => { hoisted.capturedCallbacks.current = undefined; mockedLogAiQuery.mockClear(); + mockedLogIfMcpToolCall.mockClear(); }); describe("handleStreamMode terminal callback logging", () => { @@ -180,4 +183,39 @@ describe("handleStreamMode terminal callback logging", () => { expect(mockedLogAiQuery).toHaveBeenCalledTimes(1); expect(mockedLogAiQuery).toHaveBeenCalledWith(expect.objectContaining({ type: "failure" })); }); + + // --- MCP call logging on stream failure --------------------------------- + // A failed MCP-backed question must still land in mcp_call_log (with its + // error) so it appears in the MCP Review workflow instead of vanishing. + + const mcpCtx = () => ({ + ...makeCtx(), + mcpCallMetadata: { toolName: "ask_hexclave", reason: "test", userPrompt: "prompt" }, + conversationIdForLog: "conversation-1", + }); + + it("logs the MCP call with the error when a stream with mcpCallMetadata errors", () => { + handleStreamMode(mcpCtx()); + hoisted.capturedCallbacks.current!.onError!({ error: new Error("upstream model error") }); + expect(mockedLogIfMcpToolCall).toHaveBeenCalledTimes(1); + expect(mockedLogIfMcpToolCall).toHaveBeenCalledWith(expect.objectContaining({ + errorMessage: expect.stringContaining("upstream model error"), + })); + }); + + it("logs the MCP call with the error when a stream with mcpCallMetadata aborts", () => { + handleStreamMode(mcpCtx()); + hoisted.capturedCallbacks.current!.onAbort!(); + expect(mockedLogIfMcpToolCall).toHaveBeenCalledTimes(1); + expect(mockedLogIfMcpToolCall).toHaveBeenCalledWith(expect.objectContaining({ + errorMessage: expect.stringContaining("Stream aborted"), + })); + }); + + it("logs the MCP call at most once on rapid-fire abort+error", () => { + handleStreamMode(mcpCtx()); + hoisted.capturedCallbacks.current!.onAbort!(); + hoisted.capturedCallbacks.current!.onError!({ error: new DOMException("aborted", "AbortError") }); + expect(mockedLogIfMcpToolCall).toHaveBeenCalledTimes(1); + }); }); diff --git a/apps/backend/src/lib/ai/ai-query-handlers.ts b/apps/backend/src/lib/ai/ai-query-handlers.ts index d7fd10d31..83ef13ec8 100644 --- a/apps/backend/src/lib/ai/ai-query-handlers.ts +++ b/apps/backend/src/lib/ai/ai-query-handlers.ts @@ -48,6 +48,23 @@ export function handleStreamMode(ctx: ModeContext & { // SpacetimeDB and trip the UNIQUE(correlation_id) constraint on the second // insert. Guard with a single boolean — the first terminal callback wins. let logged = false; + const logStreamFailure = (err: unknown) => { + clearTimeout(timeoutId); + if (logged) return; + logged = true; + logAiQuery({ type: "failure", common, startedAt, err, partialSteps: completedSteps }); + logIfMcpToolCall({ + mcpCallMetadata, + conversationIdForLog, + correlationId, + messages, + steps: completedSteps, + text: "", + startedAt, + modelId: String(model.modelId), + errorMessage: err instanceof Error ? err.stack ?? `${err.name}: ${err.message}` : String(err), + }); + }; const result = streamText({ model, messages: messagesWithCachedSystemPrompt, @@ -80,24 +97,8 @@ export function handleStreamMode(ctx: ModeContext & { modelId: String(model.modelId), }); }, - onError: ({ error }) => { - clearTimeout(timeoutId); - if (logged) return; - logged = true; - logAiQuery({ type: "failure", common, startedAt, err: error, partialSteps: completedSteps }); - }, - onAbort: () => { - clearTimeout(timeoutId); - if (logged) return; - logged = true; - logAiQuery({ - type: "failure", - common, - startedAt, - err: new Error("Stream aborted (client disconnect or timeout)"), - partialSteps: completedSteps, - }); - }, + onError: ({ error }) => logStreamFailure(error), + onAbort: () => logStreamFailure(new Error("Stream aborted (client disconnect or timeout)")), }); return { statusCode: 200,