-
Notifications
You must be signed in to change notification settings - Fork 3k
feat(core): add memory recall delivery telemetry #7393
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
a4e6702
69c925e
2b83e6a
d3da151
f54d168
0c55dbe
1fd8a2f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -79,11 +79,17 @@ import { ToolNames } from '../tools/tool-names.js'; | |
| import { | ||
| NextSpeakerCheckEvent, | ||
| logNextSpeakerCheck, | ||
| logMemoryRecallDelivery, | ||
| startInteractionSpan, | ||
| endInteractionSpan, | ||
| getActiveInteractionSpan, | ||
| addUserPromptAttributes, | ||
| MemoryRecallDeliveryEvent, | ||
| } from '../telemetry/index.js'; | ||
| import type { | ||
| MemoryRecallDeliveryPoint, | ||
| MemoryRecallDiscardReason, | ||
| } from '../telemetry/types.js'; | ||
| import { uiTelemetryService } from '../telemetry/uiTelemetry.js'; | ||
|
|
||
| // Forked agent cache | ||
|
|
@@ -218,8 +224,13 @@ type MemoryPrefetchHandle = { | |
| promise: Promise<RelevantAutoMemoryPromptResult>; | ||
| /** Set by promise.finally(). null until the promise settles. */ | ||
| settledAt: number | null; | ||
| /** Set when the promise resolves, even if the consume point never runs. */ | ||
| result: RelevantAutoMemoryPromptResult | null; | ||
| /** True after memory has been injected — prevents double-inject. */ | ||
| consumed: boolean; | ||
| /** True after delivery/discard telemetry has recorded the terminal outcome. */ | ||
| terminalLogged: boolean; | ||
| firedAt: number; | ||
| controller: AbortController; | ||
| }; | ||
|
|
||
|
|
@@ -657,6 +668,7 @@ export class GeminiClient { | |
| */ | ||
| requestShutdown(): void { | ||
| this.shutdownRequested = true; | ||
| this.cancelPendingMemoryPrefetch('shutdown'); | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -669,12 +681,48 @@ export class GeminiClient { | |
| * hadn't run yet), the settled result is discarded — logged at debug so | ||
| * operators can diagnose missing-memory scenarios. | ||
| */ | ||
| private cancelPendingMemoryPrefetch(): void { | ||
| private logMemoryPrefetchDelivery( | ||
| handle: MemoryPrefetchHandle, | ||
| deliveryPoint: MemoryRecallDeliveryPoint, | ||
| result: RelevantAutoMemoryPromptResult, | ||
| discardReason?: MemoryRecallDiscardReason, | ||
| ): void { | ||
| if (handle.terminalLogged) return; | ||
| handle.terminalLogged = true; | ||
|
Comment on lines
+690
to
+691
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Suggestion] The Failure scenario: A regression that removes the guard would go undetected by the test suite until a rare timing-dependent double-log surfaces in production (e.g., a race between Suggested fix: Add a test that calls — qwen3.7-max via Qwen Code /review |
||
| logMemoryRecallDelivery( | ||
| this.config, | ||
| new MemoryRecallDeliveryEvent({ | ||
| phase: 'refined', | ||
| delivery_point: deliveryPoint, | ||
| discard_reason: discardReason, | ||
| strategy: result.strategy, | ||
| docs_selected: result.selectedDocs.length, | ||
| latency_ms: Date.now() - handle.firedAt, | ||
| }), | ||
| ); | ||
| } | ||
|
|
||
| private logMemoryPrefetchDiscard( | ||
| handle: MemoryPrefetchHandle, | ||
| discardReason: MemoryRecallDiscardReason, | ||
| ): void { | ||
| this.logMemoryPrefetchDelivery( | ||
| handle, | ||
| 'discarded', | ||
| handle.result ?? EMPTY_RELEVANT_AUTO_MEMORY_RESULT, | ||
| discardReason, | ||
| ); | ||
| } | ||
|
|
||
| private cancelPendingMemoryPrefetch( | ||
| discardReason: MemoryRecallDiscardReason, | ||
| ): void { | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| const handle = this.pendingMemoryPrefetch; | ||
| if (!handle) return; | ||
| if (handle.settledAt !== null && !handle.consumed) { | ||
| debugLogger.debug('Discarding settled but unconsumed memory prefetch.'); | ||
| } | ||
| this.logMemoryPrefetchDiscard(handle, discardReason); | ||
| handle.controller.abort(); | ||
| this.pendingMemoryPrefetch = undefined; | ||
| } | ||
|
|
@@ -687,7 +735,9 @@ export class GeminiClient { | |
| * Centralises the consume-and-mark dance so the UserQuery and ToolResult | ||
| * inject sites can't drift on the guard logic. | ||
| */ | ||
| private async tryConsumeMemoryPrefetch(): Promise<RelevantAutoMemoryPromptResult | null> { | ||
| private async tryConsumeMemoryPrefetch( | ||
| deliveryPoint: Exclude<MemoryRecallDeliveryPoint, 'discarded'>, | ||
| ): Promise<RelevantAutoMemoryPromptResult | null> { | ||
| const handle = this.pendingMemoryPrefetch; | ||
| if (!handle || handle.settledAt === null || handle.consumed) { | ||
| return null; | ||
|
|
@@ -699,6 +749,14 @@ export class GeminiClient { | |
| for (const doc of result.selectedDocs) { | ||
| this.surfacedRelevantAutoMemoryPaths.add(doc.filePath); | ||
| } | ||
| this.logMemoryPrefetchDelivery(handle, deliveryPoint, result); | ||
| } else { | ||
| this.logMemoryPrefetchDelivery( | ||
| handle, | ||
| 'discarded', | ||
| result, | ||
| 'no_relevant_results', | ||
| ); | ||
| } | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| return result; | ||
| } | ||
|
|
@@ -732,7 +790,7 @@ export class GeminiClient { | |
| this.config.getBaseLlmClient().clearPerModelGeneratorCache(); | ||
| // Abort any in-flight auto-memory recall so the stale controller | ||
| // does not leak into the next session. | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('reset'); | ||
| // Drop any deferred tools revealed this session so /clear really gives | ||
| // a clean slate. We don't clear inside startChat itself because that path | ||
| // is also taken by compression (which preserves the session), and | ||
|
|
@@ -2052,16 +2110,22 @@ export class GeminiClient { | |
| // A previous recall may still be pending (slow side-query, new user | ||
| // turn arrived before it settled). Abort it before installing the | ||
| // new handle so the orphan doesn't keep running indefinitely. | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('new_query'); | ||
| const controller = new AbortController(); | ||
| // Bridge the caller's signal into the prefetch controller so a user | ||
| // abort (Ctrl-C / Esc) on the parent turn also terminates the | ||
| // recall side-query. `{ once: true }` lets the listener clean itself | ||
| // up after firing; we still call removeEventListener on the promise's | ||
| // finally to cover the normal-completion case so a long-lived parent | ||
| // signal doesn't accumulate listeners across many turns. | ||
| const onParentAbort = () => controller.abort(); | ||
| let prefetchAbortReason: MemoryRecallDiscardReason | null = null; | ||
| const onParentAbort = () => { | ||
| prefetchAbortReason = 'abort'; | ||
| controller.abort(); | ||
| this.cancelPendingMemoryPrefetch('abort'); | ||
| }; | ||
| if (signal.aborted) { | ||
| prefetchAbortReason = 'abort'; | ||
| controller.abort(); | ||
| } else { | ||
|
Comment on lines
2127
to
2130
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Suggestion] When the parent signal is already aborted at handle creation, only Failure scenario: A caller enters with an already-aborted Suggested fix: Mirror the — qwen3.7-max via Qwen Code /review |
||
| signal.addEventListener('abort', onParentAbort, { once: true }); | ||
|
|
@@ -2097,14 +2161,23 @@ export class GeminiClient { | |
| const handle: MemoryPrefetchHandle = { | ||
| promise, | ||
| settledAt: null, | ||
| result: null, | ||
| consumed: false, | ||
| terminalLogged: false, | ||
| firedAt: Date.now(), | ||
| controller, | ||
| }; | ||
| void promise.then((result) => { | ||
| handle.result = result; | ||
| }); | ||
| void promise.finally(() => { | ||
| handle.settledAt = Date.now(); | ||
| signal.removeEventListener('abort', onParentAbort); | ||
| }); | ||
| this.pendingMemoryPrefetch = handle; | ||
| if (prefetchAbortReason) { | ||
| this.cancelPendingMemoryPrefetch(prefetchAbortReason); | ||
| } | ||
| } | ||
|
|
||
| // Track prompt count for commit attribution. Only the user typing a | ||
|
|
@@ -2189,7 +2262,7 @@ export class GeminiClient { | |
| this.config.getMaxSessionTurns() > 0 && | ||
| this.sessionTurnCount > this.config.getMaxSessionTurns() | ||
| ) { | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| yield { type: GeminiEventType.MaxSessionTurns }; | ||
| if (isTopLevelInteraction) | ||
| endInteractionSpan('error', { | ||
|
|
@@ -2202,7 +2275,7 @@ export class GeminiClient { | |
| // Ensure turns never exceeds MAX_TURNS to prevent infinite loops | ||
| const boundedTurns = Math.min(turns, MAX_TURNS); | ||
| if (!boundedTurns) { | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| if (isTopLevelInteraction) | ||
| endInteractionSpan('error', { errorMessage: 'max turns exhausted' }); | ||
| return new Turn(this.getChat(), prompt_id); | ||
|
|
@@ -2243,7 +2316,7 @@ export class GeminiClient { | |
| const lastPromptTokenCount = | ||
| uiTelemetryService.getLastPromptTokenCount(); | ||
| if (lastPromptTokenCount > sessionTokenLimit) { | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| yield { | ||
| type: GeminiEventType.SessionTokenLimitExceeded, | ||
| value: { | ||
|
|
@@ -2302,7 +2375,7 @@ export class GeminiClient { | |
| `Arena control signal received: ${controlSignal.type} - ${controlSignal.reason}`, | ||
| ); | ||
| await arenaAgentClient.reportCancelled(); | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('abort'); | ||
| if (isTopLevelInteraction) endInteractionSpan('cancelled'); | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| return new Turn(this.getChat(), prompt_id); | ||
| } | ||
|
|
@@ -2398,7 +2471,7 @@ export class GeminiClient { | |
| // after any await prior to this point — flatMapTextParts above is | ||
| // the natural drain.) If still not settled, skip — the ToolResult | ||
| // inject point will retry on the next turn. | ||
| const userQueryMemory = await this.tryConsumeMemoryPrefetch(); | ||
| const userQueryMemory = await this.tryConsumeMemoryPrefetch('initial'); | ||
| if (userQueryMemory?.prompt) { | ||
| // Unshift to the front of systemReminders: on a UserQuery turn | ||
| // requestToSend leads with user text, so positioning memory at | ||
|
|
@@ -2412,7 +2485,8 @@ export class GeminiClient { | |
| } | ||
|
|
||
| if (messageType === SendMessageType.ToolResult) { | ||
| const toolResultMemory = await this.tryConsumeMemoryPrefetch(); | ||
| const toolResultMemory = | ||
| await this.tryConsumeMemoryPrefetch('tool_result'); | ||
| if (toolResultMemory?.prompt) { | ||
| // Append (not prepend): on a ToolResult turn, requestToSend leads | ||
| // with functionResponse parts that must immediately follow the | ||
|
|
@@ -2479,8 +2553,17 @@ export class GeminiClient { | |
|
|
||
| const resultStream = turn.run(model, requestToSend, signal); | ||
| let didUpdateIdeContextState = false; | ||
| let hasToolCalls = false; | ||
| try { | ||
| for await (const event of resultStream) { | ||
| if (event.type === GeminiEventType.ToolCallRequest) { | ||
| hasToolCalls = true; | ||
| } else if ( | ||
| event.type === GeminiEventType.Retry || | ||
| event.type === GeminiEventType.ModelFallback | ||
| ) { | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| hasToolCalls = false; | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| } | ||
|
qwen-code-dev-bot marked this conversation as resolved.
|
||
| if (messageDisplay && event.type === GeminiEventType.Content) { | ||
| messageDisplay.addChunk(event.value); | ||
| } | ||
|
|
@@ -2514,7 +2597,7 @@ export class GeminiClient { | |
| this.lastApiCompletionTimestamp = Date.now(); | ||
| if (isTopLevelInteraction) | ||
| endInteractionSpan('error', { errorMessage: 'loop detected' }); | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| return turn; | ||
| } | ||
|
|
||
|
|
@@ -2545,7 +2628,7 @@ export class GeminiClient { | |
| endInteractionSpan('error', { errorMessage: 'loop detected' }); | ||
| // finally cleanup catches this, but cancel explicitly to match | ||
| // the cleanup pattern at other early-return sites. | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| return turn; | ||
| } | ||
| // Update arena status on Finished events — stats are derived | ||
|
|
@@ -2609,7 +2692,7 @@ export class GeminiClient { | |
| } | ||
| // finally cleanup catches this, but cancel explicitly to match | ||
| // the cleanup pattern at other early-return sites. | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| return turn; | ||
| } | ||
| } | ||
|
|
@@ -2999,9 +3082,14 @@ export class GeminiClient { | |
| if (isTopLevelInteraction) { | ||
| endInteractionSpan(signal?.aborted ? 'cancelled' : 'ok'); | ||
| } | ||
| // Reached the bottom of the try — this turn ended cleanly. Preserve | ||
| // any still-pending memory prefetch so the next ToolResult turn can | ||
| // consume it (the whole point of the fire-and-forget design). | ||
| // Reached the bottom of the try — this turn ended cleanly. If the | ||
| // model did not request tool calls, no future ToolResult will arrive | ||
| // to consume the prefetch, so close it out now. When tool calls ARE | ||
| // pending, preserve the handle so the next ToolResult turn can | ||
| // consume it (the fire-and-forget design). | ||
| if (!hasToolCalls) { | ||
| this.cancelPendingMemoryPrefetch('no_safe_delivery_point'); | ||
| } | ||
| normalCompletion = true; | ||
| return turn; | ||
| } finally { | ||
|
|
@@ -3017,7 +3105,9 @@ export class GeminiClient { | |
| // `return turn`. Catches uncaught exceptions and guards against | ||
| // future early-return sites that forget to call cancel. | ||
| if (!normalCompletion) { | ||
| this.cancelPendingMemoryPrefetch(); | ||
| this.cancelPendingMemoryPrefetch( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P1] A slow recall can still miss terminal delivery telemetry on a normal no-tool turn. This cleanup only runs when |
||
| signal?.aborted ? 'abort' : 'no_safe_delivery_point', | ||
| ); | ||
| } | ||
| if (isTopLevelInteraction) { | ||
| endInteractionSpan(signal?.aborted ? 'cancelled' : 'error', { | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[Suggestion]
requestShutdown()callscancelPendingMemoryPrefetch('shutdown')— a new discard reason — but no test covers this path. The existingrequestShutdowntests verify background memory task gating but none set up a pending prefetch handle, so the cancel path is never exercised in tests.Failure scenario: A future refactor of
requestShutdown()could remove or misplace thecancelPendingMemoryPrefetch('shutdown')call without any test detecting the regression. Operators would lose telemetry for this terminal outcome.Suggested fix: Add a test that sets up a pending prefetch, calls
requestShutdown(), and assertslogMemoryRecallDeliverywas called withdiscard_reason: 'shutdown'. Pattern after the existingresetdiscard test.— qwen3.7-max via Qwen Code /review