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
4 changes: 2 additions & 2 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4037,8 +4037,8 @@ export class GeminiClient {

if (!turn.pendingToolCalls.length && signal && !signal.aborted) {
// Save cache-safe params here — before any early return — so that
// background extract/dream agents calling getCacheSafeParams() always
// see the current turn's history regardless of which path exits below.
// background readers calling getCacheSafeParams(sessionId) can see the
// current turn's history regardless of which path exits below.
try {
const chat = this.getChat();
const maxHistoryForCache = 40;
Expand Down
92 changes: 83 additions & 9 deletions packages/core/src/followup/speculation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,19 @@ import {
} from './speculation.js';
import type { Content } from '@google/genai';
import { ApprovalMode, type Config } from '../config/config.js';
import type { CacheSafeParams } from '../utils/forkedAgent.js';
import type { ToolResultBoundaryObservation } from '../utils/tool-result-boundary-diagnostics.js';

const forkedAgentMocks = vi.hoisted(() => ({
getCacheSafeParams: vi.fn<
(expectedSessionId?: string) => CacheSafeParams | null
>(() => ({
generationConfig: {},
history: [],
model: 'qwen-fast',
version: 1,
})),
createForkedChat: vi.fn(),
runForkedAgent: vi.fn(),
sendMessageStream: vi.fn(),
}));
Expand All @@ -33,15 +43,12 @@ vi.mock(
);

vi.mock('../utils/forkedAgent.js', () => ({
getCacheSafeParams: vi.fn(() => ({
generationConfig: {},
history: [],
model: 'qwen-fast',
version: 1,
})),
createForkedChat: vi.fn(() => ({
sendMessageStream: forkedAgentMocks.sendMessageStream,
})),
getCacheSafeParams: forkedAgentMocks.getCacheSafeParams,
createForkedChat: forkedAgentMocks.createForkedChat.mockImplementation(
() => ({
sendMessageStream: forkedAgentMocks.sendMessageStream,
}),
),
runForkedAgent: forkedAgentMocks.runForkedAgent,
runWithForkedChatModel: vi.fn(
async (
Expand All @@ -57,6 +64,20 @@ afterEach(() => {
});

describe('startSpeculation', () => {
it('does not start when the session-scoped lookup returns null', async () => {
const config = {
getSessionId: vi.fn().mockReturnValue('spec-session'),
} as unknown as Config;
forkedAgentMocks.getCacheSafeParams.mockReturnValueOnce(null);

await expect(startSpeculation(config, 'read a.ts')).rejects.toThrow(
'CacheSafeParams not available for speculation',
);

expect(forkedAgentMocks.createForkedChat).not.toHaveBeenCalled();
expect(forkedAgentMocks.runForkedAgent).not.toHaveBeenCalled();
});

it('stops at a boundary when the host guard denies a speculative invocation', async () => {
const execute = vi.fn();
const guard = vi.fn().mockResolvedValue({
Expand Down Expand Up @@ -112,6 +133,9 @@ describe('startSpeculation', () => {
const state = await startSpeculation(config, 'read a.ts');
await vi.waitFor(() => expect(state.status).toBe('boundary'));

expect(forkedAgentMocks.getCacheSafeParams).toHaveBeenCalledWith(
'spec-session',
);
expect(guard).toHaveBeenCalledWith({
callId: 'call-speculation-guard',
toolName: 'read_file',
Expand Down Expand Up @@ -179,7 +203,14 @@ describe('startSpeculation', () => {

const state = await startSpeculation(config, 'read a.ts');
await vi.waitFor(() => expect(state.status).toBe('completed'));
await vi.waitFor(() =>
expect(forkedAgentMocks.getCacheSafeParams).toHaveBeenCalledTimes(2),
);

expect(forkedAgentMocks.getCacheSafeParams).toHaveBeenNthCalledWith(
2,
'spec-session',
);
expect(guard).toHaveBeenCalledWith({
callId: 'call-speculation-guard-allow',
toolName: 'read_file',
Expand Down Expand Up @@ -212,6 +243,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
} as unknown as Config;

Expand Down Expand Up @@ -295,6 +327,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
} as unknown as Config;

Expand Down Expand Up @@ -357,6 +390,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
} as unknown as Config;

Expand Down Expand Up @@ -418,6 +452,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
} as unknown as Config;

Expand Down Expand Up @@ -481,6 +516,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
getToolOutputBatchBudget: vi.fn().mockReturnValue(10_000),
} as unknown as Config;
Expand Down Expand Up @@ -544,6 +580,7 @@ describe('startSpeculation', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue(toolRegistry),
} as unknown as Config;
forkedAgentMocks.runForkedAgent.mockResolvedValue({
Expand Down Expand Up @@ -585,6 +622,42 @@ describe('startSpeculation', () => {

await abortSpeculation(state);
});

it('does not generate a pipelined suggestion when its scoped lookup returns null', async () => {
const config = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('spec-session'),
} as unknown as Config;
forkedAgentMocks.getCacheSafeParams
.mockReturnValueOnce({
generationConfig: {},
history: [],
model: 'qwen-fast',
version: 1,
})
.mockReturnValueOnce(null);
forkedAgentMocks.sendMessageStream.mockImplementation(async function* () {
yield {
type: 'chunk',
value: {
candidates: [{ content: { parts: [{ text: 'done' }] } }],
},
};
});

const state = await startSpeculation(config, 'do something');
await vi.waitFor(() => expect(state.status).toBe('completed'));
await vi.waitFor(() =>
expect(forkedAgentMocks.getCacheSafeParams).toHaveBeenCalledTimes(2),
);

expect(forkedAgentMocks.runForkedAgent).not.toHaveBeenCalled();
expect(state.pipelinedSuggestion).toBeUndefined();

await abortSpeculation(state);
});
});

describe.each([
Expand All @@ -606,6 +679,7 @@ describe.each([
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getCwd: vi.fn().mockReturnValue(process.cwd()),
getFastModel: vi.fn().mockReturnValue(fastModel),
getSessionId: vi.fn().mockReturnValue('spec-session'),
getToolRegistry: vi.fn().mockReturnValue({
ensureTool: vi.fn().mockResolvedValue({
build: vi.fn().mockReturnValue({
Expand Down
4 changes: 2 additions & 2 deletions packages/core/src/followup/speculation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ export async function startSpeculation(
parentSignal?: AbortSignal,
options?: { model?: string },
): Promise<SpeculationState> {
const cacheSafe = getCacheSafeParams();
const cacheSafe = getCacheSafeParams(config.getSessionId());

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 fail-closed outcome of the session-scoped lookup is untested for both speculation readers. Every test in this file uses a mock that returns params unconditionally, and the new assertions only pin that the session id is passed — there is no case where getCacheSafeParams returns null, so neither the throw in startSpeculation nor the return null in generatePipelinedSuggestion is exercised. The sibling readers pin the equivalent path (extractionAgentPlanner.test.ts "throws when getCacheSafeParams returns null", suggestionGenerator.test.ts foreign-slot fallback), but speculation does not.

This matters in exactly the multi-session scenario this PR fixes: a future edit relaxing the null guard — for example falling back to an unscoped read to "fix" a speculation-not-starting report — would pass every test in this file green and silently reintroduce the cross-session fork. A mutation probe confirmed it: applying getCacheSafeParams(config.getSessionId()) ?? getCacheSafeParams() here kept the suite 18/18 green while a contention probe showed the mutated code consuming the foreign snapshot past the guard; the same probe passes on the unmodified source.

Add a test setting the mock to return null and asserting startSpeculation rejects with 'CacheSafeParams not available for speculation' and neither createForkedChat nor runForkedAgent is called, plus the analogous null case for the pipelined suggestion path:

it('does not start speculation when the session-scoped lookup returns null', async () => {
  forkedAgentMocks.getCacheSafeParams.mockReturnValue(null);
  await expect(startSpeculation(config, 'read a.ts')).rejects.toThrow(
    'CacheSafeParams not available for speculation',
  );
  expect(createForkedChat).not.toHaveBeenCalled();
  expect(forkedAgentMocks.runForkedAgent).not.toHaveBeenCalled();
});
中文说明

两个推测(speculation)读取方的 fail-closed(会话校验失败返回 null)结果都没有测试覆盖。本文件中所有测试都使用无条件返回 params 的 mock,新增断言也只钉住了「传入了 session id」这一行为——没有任何用例让 getCacheSafeParams 返回 null,因此 startSpeculation 的 throw 和 generatePipelinedSuggestionreturn null 都不会被执行到。相邻的读取方钉住了等价路径(extractionAgentPlanner.test.ts 的 "throws when getCacheSafeParams returns null"、suggestionGenerator.test.ts 的外部 slot 回退),唯独 speculation 没有。

这正是本 PR 修复的多 session 场景下的风险所在:未来若有人放宽 null 守卫——例如为了「修复」推测不启动的问题而回退到无 session 限定的读取——本文件的所有测试仍会全绿,跨 session fork 会被悄悄重新引入。变异探针证实了这一点:在此处应用 getCacheSafeParams(config.getSessionId()) ?? getCacheSafeParams() 后,测试套件仍 18/18 全绿,而竞争探针显示变异后的代码越过守卫消费了外来快照;同一探针在未修改的源码上通过。

建议新增一个测试:将 mock 设为返回 null,断言 startSpeculation 以 'CacheSafeParams not available for speculation' 拒绝,且 createForkedChatrunForkedAgent 均未被调用;流水线推测路径也加类似的 null 用例。

— qwen3.8-max via Qwen Code /review (v0.21.14)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Addressed in bd1babe. Added fail-closed coverage for both speculation readers: the initial null lookup now asserts rejection with zero forked-chat/agent calls, and a null second lookup completes speculation without producing a pipelined suggestion or calling the forked agent.

if (!cacheSafe) {
throw new Error('CacheSafeParams not available for speculation');
}
Expand Down Expand Up @@ -718,7 +718,7 @@ The assistant responded: ${speculatedSummary || '(tool calls executed)'}

${SUGGESTION_PROMPT}`;

const cacheSafeParams = getCacheSafeParams();
const cacheSafeParams = getCacheSafeParams(config.getSessionId());
if (!cacheSafeParams) return null;
const model = modelOverride ?? config.getFastModel();
const resolvedModel = model ?? cacheSafeParams.model;
Expand Down
5 changes: 3 additions & 2 deletions packages/core/src/followup/suggestionGenerator.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ describe('generatePromptSuggestion', () => {
enableCacheSharing: true,
});

expect(mockGetCacheSafeParams).toHaveBeenCalledWith('test-session');
expect(mockRunForkedAgent).toHaveBeenCalledWith(
expect.objectContaining({ model: 'main-model', abortSignal: signal }),
);
Expand Down Expand Up @@ -175,8 +176,8 @@ describe('generatePromptSuggestion', () => {
{ enableCacheSharing: true },
);

// The foreign slot is rejected before cloning the full cached payload.
expect(mockGetCacheSafeParams).not.toHaveBeenCalled();
// The cache API rejects the foreign slot before cloning its payload.
expect(mockGetCacheSafeParams).toHaveBeenCalledWith('session-A');
// The fork must NOT be used for a foreign session's params.
expect(mockRunForkedAgent).not.toHaveBeenCalled();
// The session-safe base-LLM path is used instead.
Expand Down
15 changes: 6 additions & 9 deletions packages/core/src/followup/suggestionGenerator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -110,29 +110,26 @@ export async function generatePromptSuggestion(
const cacheSafeSessionId = options?.enableCacheSharing
? getCacheSafeParamsSessionId()
: undefined;
const cacheSafe =
cacheSafeSessionId === sessionId ? getCacheSafeParams() : null;
const cacheSafe = options?.enableCacheSharing
? getCacheSafeParams(sessionId)
: null;
// The cache-safe slot is a process-global: in a multi-session daemon it
// can hold ANOTHER session's transcript + systemInstruction. Only use it
// when it belongs to THIS session; otherwise fall back to the
// session-safe base-LLM path (#9233).
const sessionCacheSafe =
cacheSafe && cacheSafe.sessionId === config.getSessionId()
? cacheSafe
: null;
const modelOverride = options?.model;
const cacheSharingState = sessionCacheSafe
const cacheSharingState = cacheSafe
? 'true'
: cacheSafeSessionId
? 'session_mismatch'
: 'false';
debugLogger.debug(
`Generating suggestion: cacheSharing=${cacheSharingState}, model=${modelOverride || '(default)'}`,
);
const raw = sessionCacheSafe
const raw = cacheSafe
? await generateViaForkedQuery(
config,
sessionCacheSafe,
cacheSafe,
abortSignal,
modelOverride,
)
Expand Down
66 changes: 66 additions & 0 deletions packages/core/src/memory/extract.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,17 @@ import {
rebuildUserAutoMemoryIndex,
} from './indexer.js';
import { refreshMemoryInstruction } from './refresh.js';
import { getCacheSafeParamsSessionId } from '../utils/forkedAgent.js';

vi.mock('./extractionAgentPlanner.js', () => ({
runAutoMemoryExtractionByAgent: vi.fn(),
}));

vi.mock('../utils/forkedAgent.js', async (importOriginal) => ({
...(await importOriginal<typeof import('../utils/forkedAgent.js')>()),
getCacheSafeParamsSessionId: vi.fn(),
}));

vi.mock('./indexer.js', () => ({
rebuildManagedAutoMemoryIndex: vi.fn().mockResolvedValue(''),
rebuildUserAutoMemoryIndex: vi.fn().mockResolvedValue(''),
Expand Down Expand Up @@ -75,6 +81,7 @@ describe('auto-memory extraction', () => {
getModel: vi.fn().mockReturnValue('qwen3-coder-plus'),
} as unknown as Config;
vi.clearAllMocks();
vi.mocked(getCacheSafeParamsSessionId).mockReturnValue('session-1');
});

afterEach(async () => {
Expand Down Expand Up @@ -125,6 +132,65 @@ describe('auto-memory extraction', () => {
expect(cursor.processedOffset).toBe(2);
});

it('skips a session mismatch without advancing the cursor', async () => {
vi.mocked(getCacheSafeParamsSessionId)
.mockReturnValueOnce('session-1')
.mockReturnValueOnce('session-2');
const cursorBefore = await fs.readFile(
getAutoMemoryExtractCursorPath(projectRoot),
'utf-8',
);

const result = await runAutoMemoryExtract({
projectRoot,
sessionId: 'session-1',
config: mockConfig,
history: [{ role: 'user', parts: [{ text: 'Remember this.' }] }],
});

expect(result.skippedReason).toBe('session_mismatch');
expect(result.cursor.processedOffset).toBeUndefined();
expect(runAutoMemoryExtractionByAgent).not.toHaveBeenCalled();
expect(
await fs.readFile(getAutoMemoryExtractCursorPath(projectRoot), 'utf-8'),
).toBe(cursorBefore);
});

it('skips an existing session mismatch before scaffold IO', async () => {
vi.mocked(getCacheSafeParamsSessionId).mockReturnValue('session-2');
const uncreatedProjectRoot = path.join(tempDir, 'not-created');

const result = await runAutoMemoryExtract({
projectRoot: uncreatedProjectRoot,
sessionId: 'session-1',
config: mockConfig,
history: [{ role: 'user', parts: [{ text: 'Remember this.' }] }],
});

expect(result.skippedReason).toBe('session_mismatch');
await expect(fs.stat(uncreatedProjectRoot)).rejects.toMatchObject({
code: 'ENOENT',
});
expect(runAutoMemoryExtractionByAgent).not.toHaveBeenCalled();
});

it('preserves the empty-cache failure path', async () => {
vi.mocked(getCacheSafeParamsSessionId).mockReturnValue(undefined);
vi.mocked(runAutoMemoryExtractionByAgent).mockRejectedValueOnce(
new Error('no cache-safe params'),
);

await expect(
runAutoMemoryExtract({
projectRoot,
sessionId: 'session-1',
config: mockConfig,
history: [{ role: 'user', parts: [{ text: 'Remember this.' }] }],
}),
).rejects.toThrow('no cache-safe params');
expect(runAutoMemoryExtractionByAgent).toHaveBeenCalledOnce();
});

it('throws when config is missing because heuristic fallback was removed', async () => {
await expect(
runAutoMemoryExtract({
Expand Down
Loading
Loading