Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
733 changes: 731 additions & 2 deletions packages/core/src/core/client.test.ts

Large diffs are not rendered by default.

126 changes: 108 additions & 18 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
};

Expand Down Expand Up @@ -657,6 +668,7 @@ export class GeminiClient {
*/
requestShutdown(): void {
this.shutdownRequested = true;
this.cancelPendingMemoryPrefetch('shutdown');
}
Comment on lines +671 to 672

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] requestShutdown() calls cancelPendingMemoryPrefetch('shutdown') — a new discard reason — but no test covers this path. The existing requestShutdown tests 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 the cancelPendingMemoryPrefetch('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 asserts logMemoryRecallDelivery was called with discard_reason: 'shutdown'. Pattern after the existing reset discard test.

— qwen3.7-max via Qwen Code /review


/**
Expand All @@ -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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] The terminalLogged idempotency guard is not directly tested. No test constructs a scenario where a handle's terminalLogged is set true before a second call to logMemoryPrefetchDelivery.

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 cancelPendingMemoryPrefetch and tryConsumeMemoryPrefetch).

Suggested fix: Add a test that calls logMemoryPrefetchDelivery (or triggers both delivery and discard) twice on the same handle and asserts logMemoryRecallDelivery is called exactly once.

— 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 {
Comment thread
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;
}
Expand All @@ -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;
Expand All @@ -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',
);
}
Comment thread
qwen-code-dev-bot marked this conversation as resolved.
return result;
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] When the parent signal is already aborted at handle creation, only controller.abort() fires — cancelPendingMemoryPrefetch('abort') is not called. The onParentAbort handler in the else branch includes both calls, but the pre-aborted path only mirrors the first.

Failure scenario: A caller enters with an already-aborted AbortSignal. The recall controller is aborted and the catch handler returns EMPTY_RELEVANT_AUTO_MEMORY_RESULT, but the delivery event won't carry discard_reason: 'abort' — operators cannot distinguish pre-abort discard from other empty-result discards.

Suggested fix: Mirror the onParentAbort handler's this.cancelPendingMemoryPrefetch('abort') call in the pre-aborted branch. Note the handle may not yet be installed on this.pendingMemoryPrefetch at this point, so this may require restructuring.

— qwen3.7-max via Qwen Code /review

signal.addEventListener('abort', onParentAbort, { once: true });
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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', {
Expand All @@ -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);
Expand Down Expand Up @@ -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: {
Expand Down Expand Up @@ -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');
Comment thread
qwen-code-dev-bot marked this conversation as resolved.
return new Turn(this.getChat(), prompt_id);
}
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
) {
Comment thread
qwen-code-dev-bot marked this conversation as resolved.
hasToolCalls = false;
Comment thread
qwen-code-dev-bot marked this conversation as resolved.
}
Comment thread
qwen-code-dev-bot marked this conversation as resolved.
if (messageDisplay && event.type === GeminiEventType.Content) {
messageDisplay.addChunk(event.value);
}
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
}
}
Expand Down Expand Up @@ -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 {
Expand All @@ -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(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The 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 normalCompletion is false; when recall is still pending at the initial consume point and the model finishes without tool calls or another continuation, the bottom-of-try path sets normalCompletion = true and preserves the handle for a future ToolResult. But no ToolResult will be scheduled for that no-tool turn, so the prefetch has no terminal delivery/discard event until a later unrelated new query/reset logs a misleading reason, or never if the session ends. That violates the PR's one-terminal-outcome goal; the no-future-delivery path should close the pending prefetch as no_safe_delivery_point (or equivalent) before returning.

signal?.aborted ? 'abort' : 'no_safe_delivery_point',
);
}
if (isTopLevelInteraction) {
endInteractionSpan(signal?.aborted ? 'cancelled' : 'error', {
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/telemetry/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ export const EVENT_PERFORMANCE_REGRESSION = 'qwen-code.performance.regression';
export const EVENT_MEMORY_EXTRACT = 'qwen-code.memory.extract';
export const EVENT_MEMORY_DREAM = 'qwen-code.memory.dream';
export const EVENT_MEMORY_RECALL = 'qwen-code.memory.recall';
export const EVENT_MEMORY_RECALL_DELIVERY = 'qwen-code.memory.recall.delivery';

// Session Tracing Span Names
export const SPAN_INTERACTION = 'qwen-code.interaction';
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/telemetry/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ export {
logMemoryExtract,
logMemoryDream,
logMemoryRecall,
logMemoryRecallDelivery,
} from './loggers.js';
export type { SlashCommandEvent, ChatCompressionEvent } from './types.js';
export {
Expand Down Expand Up @@ -91,6 +92,7 @@ export {
MemoryExtractEvent,
MemoryDreamEvent,
MemoryRecallEvent,
MemoryRecallDeliveryEvent,
} from './types.js';
export { makeSlashCommandEvent, makeChatCompressionEvent } from './types.js';
export type {
Expand Down Expand Up @@ -137,6 +139,7 @@ export {
recordMemoryDreamMetrics,
recordMemoryRecallMetrics,
recordChannelMemoryRecallMetrics,
recordMemoryRecallDeliveryMetrics,
// Performance monitoring types
PerformanceMetricType,
MemoryMetricType,
Expand Down
Loading
Loading