From a087f6b0bd458178c871b6dbcb5a59b8fd3231b2 Mon Sep 17 00:00:00 2001 From: Aadesh Kheria Date: Tue, 5 May 2026 13:16:47 -0700 Subject: [PATCH] Refactor AI query logging to encapsulate error handling within async tasks and improve serialization robustness --- .../mcp-review/update-correction/route.ts | 2 +- .../src/lib/ai/loggers/ai-query-logger.ts | 102 ++++++++++-------- 2 files changed, 58 insertions(+), 46 deletions(-) diff --git a/apps/backend/src/app/api/latest/internal/mcp-review/update-correction/route.ts b/apps/backend/src/app/api/latest/internal/mcp-review/update-correction/route.ts index 7b4023f1e..10bf2d17d 100644 --- a/apps/backend/src/app/api/latest/internal/mcp-review/update-correction/route.ts +++ b/apps/backend/src/app/api/latest/internal/mcp-review/update-correction/route.ts @@ -33,7 +33,7 @@ export const POST = createSmartRouteHandler({ const token = getEnvVariable("STACK_MCP_LOG_TOKEN"); const reviewer = user.display_name ?? user.primary_email ?? user.id; - + await callReducerStrict("upsert_qa_from_call", [ token, body.correlationId, diff --git a/apps/backend/src/lib/ai/loggers/ai-query-logger.ts b/apps/backend/src/lib/ai/loggers/ai-query-logger.ts index 89f9816d3..5cb08a051 100644 --- a/apps/backend/src/lib/ai/loggers/ai-query-logger.ts +++ b/apps/backend/src/lib/ai/loggers/ai-query-logger.ts @@ -108,20 +108,25 @@ export async function logAiQuery(entry: AiQueryLogEntry): Promise { } function serializeSteps(steps: ReadonlyArray>): string { - return JSON.stringify(steps.map((step, i) => ({ - step: i, - text: step.text || undefined, - toolCalls: step.toolCalls.map(tc => ({ - toolName: tc.toolName, - toolCallId: tc.toolCallId, - args: tc.input, - })), - toolResults: step.toolResults.map(tr => ({ - toolName: tr.toolName, - toolCallId: tr.toolCallId, - result: truncateLargeToolResult(tr.toolName, tr.output), - })), - }))); + try { + return JSON.stringify(steps.map((step, i) => ({ + step: i, + text: step.text || undefined, + toolCalls: step.toolCalls.map(tc => ({ + toolName: tc.toolName, + toolCallId: tc.toolCallId, + args: tc.input, + })), + toolResults: step.toolResults.map(tr => ({ + toolName: tr.toolName, + toolCallId: tr.toolCallId, + result: truncateLargeToolResult(tr.toolName, tr.output), + })), + }))); + } catch (e) { + captureError("ai-query-steps-serialize", e); + return JSON.stringify({ _serializationFailed: true, stepCount: steps.length }); + } } export function logAiQuerySuccess(args: { @@ -134,22 +139,27 @@ export function logAiQuerySuccess(args: { openrouterGenerationId: string | undefined, }): void { const { common, startedAt, steps, text, usage, providerMetadata, openrouterGenerationId } = args; - const rawCost = extractCostFromUsage(usage); - runAsynchronouslyAndWaitUntil(logAiQuery({ - ...common, - stepsJson: serializeSteps(steps), - finalText: text, - inputTokens: usage.inputTokens ?? undefined, - outputTokens: usage.outputTokens ?? undefined, - cachedInputTokens: extractCachedTokens(providerMetadata), - cacheCreationTokens: usage.inputTokenDetails.cacheWriteTokens ?? undefined, - costUsd: rawCost.costUsd ?? extractOpenRouterCost(providerMetadata), - cacheDiscountUsd: undefined, // backfilled by refineGenerationCost below - openrouterGenerationId, - stepCount: steps.length, - durationMs: BigInt(Math.round(performance.now() - startedAt)), - errorMessage: undefined, - })); + // Build the row inside the async task so any throw (serialization, + // metadata extraction, etc.) is contained by the async boundary instead + // of bubbling up into the user-facing success path. + runAsynchronouslyAndWaitUntil(async () => { + const rawCost = extractCostFromUsage(usage); + await logAiQuery({ + ...common, + stepsJson: serializeSteps(steps), + finalText: text, + inputTokens: usage.inputTokens ?? undefined, + outputTokens: usage.outputTokens ?? undefined, + cachedInputTokens: extractCachedTokens(providerMetadata), + cacheCreationTokens: usage.inputTokenDetails.cacheWriteTokens ?? undefined, + costUsd: rawCost.costUsd ?? extractOpenRouterCost(providerMetadata), + cacheDiscountUsd: undefined, // backfilled by refineGenerationCost below + openrouterGenerationId, + stepCount: steps.length, + durationMs: BigInt(Math.round(performance.now() - startedAt)), + errorMessage: undefined, + }); + }); if (openrouterGenerationId != null) { runAsynchronouslyAndWaitUntil(refineGenerationCost({ generationId: openrouterGenerationId, @@ -166,19 +176,21 @@ export function logAiQueryFailure(args: { }): void { const { common, startedAt, err, partialSteps } = args; captureError("ai-query-upstream", err); - runAsynchronouslyAndWaitUntil(logAiQuery({ - ...common, - stepsJson: partialSteps && partialSteps.length > 0 ? serializeSteps(partialSteps) : "[]", - finalText: "", - inputTokens: undefined, - outputTokens: undefined, - cachedInputTokens: undefined, - cacheCreationTokens: undefined, - costUsd: undefined, - cacheDiscountUsd: undefined, - openrouterGenerationId: undefined, - stepCount: partialSteps?.length ?? 0, - durationMs: BigInt(Math.round(performance.now() - startedAt)), - errorMessage: formatErrorForLog(err), - })); + runAsynchronouslyAndWaitUntil(async () => { + await logAiQuery({ + ...common, + stepsJson: partialSteps && partialSteps.length > 0 ? serializeSteps(partialSteps) : "[]", + finalText: "", + inputTokens: undefined, + outputTokens: undefined, + cachedInputTokens: undefined, + cacheCreationTokens: undefined, + costUsd: undefined, + cacheDiscountUsd: undefined, + openrouterGenerationId: undefined, + stepCount: partialSteps?.length ?? 0, + durationMs: BigInt(Math.round(performance.now() - startedAt)), + errorMessage: formatErrorForLog(err), + }); + }); }