From e71212cdc6b98ac9b920922f1f16b8cd834a0442 Mon Sep 17 00:00:00 2001 From: agent_coder Date: Sat, 11 Jul 2026 10:43:35 +0300 Subject: [PATCH] =?UTF-8?q?perf(ai-chat):=20=D0=BA=D1=8D=D1=88=20compactTo?= =?UTF-8?q?olOutput=20=D0=BF=D0=BE=20identity=20=D1=88=D0=B0=D0=B3=D0=B0?= =?UTF-8?q?=20(#490)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compactToolOutput делает JSON.stringify каждого output на КАЖДОМ flush. Т.к. onStepFinish на шаге N перестраивает всю assistant-строку по всем N накопленным шагам, а каждый output — 50–200 KB, это O(N²) stringify за ход. Мемоизация по identity шага: finished-шаг в capturedSteps неизменен и держит стабильную ссылку между flush'ами, поэтому его parts (и дорогой stringify output) строятся ровно раз за ход. buildStepParts вынесен в чистую функцию; assistantParts принимает опциональный StepPartsCache (WeakMap), flushAssistant пробрасывает его, stream() заводит один WeakMap на ход и передаёт во все flush'и. Промах кэша (или его отсутствие в тестах/легаси-вызовах) просто пересобирает — байтового расхождения нет. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/core/ai-chat/ai-chat.service.spec.ts | 49 ++++++ .../src/core/ai-chat/ai-chat.service.ts | 165 +++++++++++------- 2 files changed, 153 insertions(+), 61 deletions(-) diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index cd32331a..bee3139f 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -13,6 +13,7 @@ import { compactToolOutput, assistantParts, serializeSteps, + type StepPartsCache, rowToUiMessage, prepareAgentStep, stepBudgetWarning, @@ -114,6 +115,54 @@ describe('compactToolOutput', () => { describe('assistantParts', () => { type AnyPart = Record; + // #490 memoization: assistantParts builds each step's parts once and caches + // them by the step OBJECT's identity, so a mid-stream flush does not + // re-stringify every prior step's (large) output. Observable property: with a + // shared cache, the second call over the SAME step object returns the cached + // (identical) part array even if the step's underlying output was swapped — + // proving the work was memoized, not redone. + it('memoizes a step by identity (shared cache => one build per step)', () => { + const cache: StepPartsCache = new WeakMap(); + const step = { + text: 'x', + toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }], + toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }], + }; + const first = assistantParts([step], '', cache) as AnyPart[]; + expect((first.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe( + 1, + ); + // Swap the output for a NEW value; a re-build would pick it up, a cache hit + // keeps the first result. + step.toolResults[0] = { + toolCallId: 'c1', + toolName: 'getPage', + output: { v: 2 }, + }; + const second = assistantParts([step], '', cache) as AnyPart[]; + expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe( + 1, + ); + // Same cached part objects are reused. + expect(second.find((p) => p.type === 'tool-getPage')).toBe( + first.find((p) => p.type === 'tool-getPage'), + ); + }); + + it('without a cache, each call rebuilds (no stale memo)', () => { + const step = { + text: 'x', + toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: {} }], + toolResults: [{ toolCallId: 'c1', toolName: 'getPage', output: { v: 1 } }], + }; + const first = assistantParts([step], '') as AnyPart[]; + step.toolResults[0].output = { v: 2 }; + const second = assistantParts([step], '') as AnyPart[]; + expect((second.find((p) => p.type === 'tool-getPage')!.output as any).v).toBe( + 2, + ); + }); + it('emits output-available for a tool-call WITH a paired result', () => { const steps = [ { diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 3c2bd6ec..4f8cb2dc 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1347,6 +1347,11 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { const capturedSteps: StepLike[] = []; let inProgressText = ''; + // Per-turn step->parts memo (#490): shared across every flushAssistant call + // this turn so each finished step's (large) output is JSON-stringified ONCE, + // not re-stringified on every subsequent onStepFinish flush (was O(N²)). + const partsCache: StepPartsCache = new WeakMap(); + // Token-degeneration guard (#444). When the final-step lockdown is OFF, a // runaway repetition loop (the 255KB "loadTools." incident) is aborted via // this internal controller, unioned with the run/socket signal below. The @@ -1416,7 +1421,10 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { await this.aiChatMessageRepo.update( assistantId, workspace.id, - flushAssistant(capturedSteps, '', 'streaming', { pageChanged }), + flushAssistant(capturedSteps, '', 'streaming', { + pageChanged, + partsCache, + }), { onlyIfStreaming: true }, ); } catch (err) { @@ -1649,6 +1657,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { // closure scope here). Omitted/0 = no limit. maxContextTokens: resolved?.chatContextWindow, pageChanged, + partsCache, }), ); // #184/#487: the RUN is finalized ALWAYS (never gated on the message). @@ -1714,6 +1723,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { flushAssistant(capturedSteps, inProgressText, 'error', { error: errorText, pageChanged, + partsCache, }), ); // #184: settle the RUN as failed, carrying the provider/transport cause. @@ -1737,6 +1747,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { flushAssistant(capturedSteps, truncated, 'error', { error: OUTPUT_DEGENERATION_ERROR, pageChanged, + partsCache, }), ); if (runId) @@ -1771,6 +1782,7 @@ export class AiChatService implements OnModuleInit, OnModuleDestroy { await finalizeAssistant( flushAssistant(capturedSteps, inProgressText, 'aborted', { pageChanged, + partsCache, }), ); // #184: settle the RUN as aborted (an explicit user stop reached the @@ -2368,71 +2380,97 @@ function normalizeToolError(error: unknown): string { */ // Exported only so the unit tests can import these pure helpers; exporting // them does not change runtime behavior. +/** + * Per-turn memo for {@link assistantParts}: a step's rebuilt parts keyed by the + * step OBJECT's identity (#490). A finished step in `capturedSteps` keeps a stable + * reference across every mid-stream flush, and `compactToolOutput` inside it does a + * `JSON.stringify` of the whole (often 50–200 KB) output — so without a memo each + * `onStepFinish` re-stringifies EVERY prior step's output (O(N²) stringify over a + * turn). Keyed by step identity => one stringify per step per turn. WeakMap so a + * turn's steps are GC'd with the turn. + */ +export type StepPartsCache = WeakMap>>; + +/** Build the parts for ONE step (text + a part per tool call). Pure. */ +function buildStepParts(step: StepLike): Array> { + const parts: Array> = []; + if (step.text) { + parts.push({ type: 'text', text: step.text }); + } + // Index this step's results by tool call id to pair calls with outputs. + const resultsById = new Map(); + for (const r of step.toolResults ?? []) { + if (r.toolCallId) resultsById.set(r.toolCallId, r.output); + } + // Index this step's THROWN tool failures (ai@6 `tool-error` content parts) + // by tool call id, so a call that failed replays with its real error text. + const errorsById = new Map(); + for (const part of step.content ?? []) { + if (part.type === 'tool-error' && part.toolCallId) { + errorsById.set(part.toolCallId, part.error); + } + } + for (const call of step.toolCalls ?? []) { + if (!call.toolName || !call.toolCallId) continue; + const hasResult = resultsById.has(call.toolCallId); + if (hasResult) { + // output-available: the tool returned; the next turn replays its result. + parts.push({ + type: `tool-${call.toolName}`, + toolCallId: call.toolCallId, + state: 'output-available', + input: call.input, + output: compactToolOutput(resultsById.get(call.toolCallId)), + }); + } else if (errorsById.has(call.toolCallId)) { + // The tool THREW: replay the REAL error so the model on the next turn + // knows WHY the call failed (and does not blindly repeat it). An + // output-error round-trips through convertToModelMessages as a balanced + // tool-call + tool-result, keeping the rebuilt history valid. + parts.push({ + type: `tool-${call.toolName}`, + toolCallId: call.toolCallId, + state: 'output-error', + input: call.input, + errorText: normalizeToolError(errorsById.get(call.toolCallId)), + }); + } else { + // No paired result AND no tool-error (e.g. aborted mid-step). Persisting + // a bare tool-call (input-available) would replay as an unpaired call and + // throw MissingToolResultsError on the next turn (convertToModelMessages + // emits no tool-result for it). Emit a SYNTHETIC paired result instead: + // an output-error round-trips through convertToModelMessages as a + // balanced tool-call + tool-result, keeping the rebuilt history valid. + parts.push({ + type: `tool-${call.toolName}`, + toolCallId: call.toolCallId, + state: 'output-error', + input: call.input, + errorText: TOOL_CALL_INCOMPLETE_TEXT, + }); + } + } + return parts; +} + export function assistantParts( steps: ReadonlyArray | undefined, fallbackText: string, + cache?: StepPartsCache, ): UIMessage['parts'] { const parts: Array> = []; - let sawText = false; for (const step of steps ?? []) { - if (step.text) { - parts.push({ type: 'text', text: step.text }); - sawText = true; - } - // Index this step's results by tool call id to pair calls with outputs. - const resultsById = new Map(); - for (const r of step.toolResults ?? []) { - if (r.toolCallId) resultsById.set(r.toolCallId, r.output); - } - // Index this step's THROWN tool failures (ai@6 `tool-error` content parts) - // by tool call id, so a call that failed replays with its real error text. - const errorsById = new Map(); - for (const part of step.content ?? []) { - if (part.type === 'tool-error' && part.toolCallId) { - errorsById.set(part.toolCallId, part.error); - } - } - for (const call of step.toolCalls ?? []) { - if (!call.toolName || !call.toolCallId) continue; - const hasResult = resultsById.has(call.toolCallId); - if (hasResult) { - // output-available: the tool returned; the next turn replays its result. - parts.push({ - type: `tool-${call.toolName}`, - toolCallId: call.toolCallId, - state: 'output-available', - input: call.input, - output: compactToolOutput(resultsById.get(call.toolCallId)), - }); - } else if (errorsById.has(call.toolCallId)) { - // The tool THREW: replay the REAL error so the model on the next turn - // knows WHY the call failed (and does not blindly repeat it). An - // output-error round-trips through convertToModelMessages as a balanced - // tool-call + tool-result, keeping the rebuilt history valid. - parts.push({ - type: `tool-${call.toolName}`, - toolCallId: call.toolCallId, - state: 'output-error', - input: call.input, - errorText: normalizeToolError(errorsById.get(call.toolCallId)), - }); - } else { - // No paired result AND no tool-error (e.g. aborted mid-step). Persisting - // a bare tool-call (input-available) would replay as an unpaired call and - // throw MissingToolResultsError on the next turn (convertToModelMessages - // emits no tool-result for it). Emit a SYNTHETIC paired result instead: - // an output-error round-trips through convertToModelMessages as a - // balanced tool-call + tool-result, keeping the rebuilt history valid. - parts.push({ - type: `tool-${call.toolName}`, - toolCallId: call.toolCallId, - state: 'output-error', - input: call.input, - errorText: TOOL_CALL_INCOMPLETE_TEXT, - }); - } + // Memoize per step object (#490): a finished step is immutable and keeps its + // reference across flushes, so its parts (and the costly output stringify) are + // built exactly once per turn. A cache miss (or no cache) just rebuilds. + let stepParts = cache?.get(step as object); + if (!stepParts) { + stepParts = buildStepParts(step); + cache?.set(step as object, stepParts); } + parts.push(...stepParts); } + const sawText = parts.some((p) => p.type === 'text'); if (!sawText && fallbackText) { // No per-step text (e.g. a single final block): append the final text after // any tool parts so the natural call -> result -> answer order is preserved. @@ -2595,6 +2633,9 @@ export function flushAssistant( maxContextTokens?: number; error?: string; pageChanged?: { title: string; diff: string } | null; + // Per-turn step->parts memo (#490): pass the SAME cache on every flush of a + // turn so each finished step's output is stringified once, not once per flush. + partsCache?: StepPartsCache; }, ): AssistantFlush { const finished = capturedSteps ?? []; @@ -2604,9 +2645,11 @@ export function flushAssistant( // in-progress step's text (the partial answer cut off by an error/abort, or // simply not yet flushed mid-stream) as the last text part so the persisted // parts match what streamed to the client. - const parts = assistantParts(finished, '') as unknown as Array< - Record - >; + const parts = assistantParts( + finished, + '', + extra?.partsCache, + ) as unknown as Array>; if (trailing) parts.push({ type: 'text', text: trailing }); const metadata: Record = {