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
36 changes: 34 additions & 2 deletions packages/cli/src/acp-integration/session/Session.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33338,6 +33338,7 @@ describe('Session', () => {

it('fires prompt-suggestion extNotification after end_turn when enabled', async () => {
generateMock.mockResolvedValue({ suggestion: 'Run the tests next?' });
vi.mocked(mockChat.getHistoryTail).mockClear();

await session.prompt({
sessionId: 'test-session-id',
Expand All @@ -33356,13 +33357,44 @@ describe('Session', () => {
);
});

// Pin the curated-history tail: a revert to `chat.getHistory(true)`
// (full structuredClone per end_turn, the #4624 heap-peak shape) or a
// dropped curated argument (`getHistoryTail(40)` defaults
// curated=false) must not pass silently (#9233).
expect(vi.mocked(mockChat.getHistoryTail)).toHaveBeenCalledWith(40, true);

// The generator received an AbortSignal so the daemon can cancel
// mid-flight if the next prompt arrives first.
// mid-flight if the next prompt arrives first. `merged.ui` above leaves
// `enableCacheSharing` UNSET, so the gate must honour the schema's
// declared `default: true` — `mergeSettings` never applies schema
Comment thread
yiliang114 marked this conversation as resolved.
// defaults, and gating on `=== true` turned the cache-aware fork into
// dead code unless the user explicitly set the flag (#9230).
expect(generateMock).toHaveBeenCalledWith(
mockConfig,
expect.any(Array),
expect.any(AbortSignal),
expect.objectContaining({ enableCacheSharing: true }),
);
});

it('forwards an explicit enableCacheSharing=false opt-out', async () => {
(mockSettings as unknown as { merged: { ui: unknown } }).merged.ui = {
enableFollowupSuggestions: true,
enableCacheSharing: false,
};
generateMock.mockResolvedValue({ suggestion: null });

await session.prompt({
sessionId: 'test-session-id',
prompt: [{ type: 'text', text: 'hello' }],
});

await vi.waitFor(() => expect(generateMock).toHaveBeenCalled());
expect(generateMock).toHaveBeenCalledWith(
mockConfig,
expect.any(Array),
expect.any(AbortSignal),
expect.objectContaining({ enableCacheSharing: expect.any(Boolean) }),
expect.objectContaining({ enableCacheSharing: false }),
);
});

Expand Down
15 changes: 9 additions & 6 deletions packages/cli/src/acp-integration/session/Session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3889,7 +3889,7 @@ export class Session implements SessionContext {
* `qwen/notify/session/prompt-suggestion` extNotification. Mirrors
* the CLI's `AppContainer.tsx` integration: same `generatePromptSuggestion`
* call, same `enableCacheSharing` flag forwarding, same curated
* history slice (`getHistory(true).slice(-40)`).
* history tail (`getHistoryTail(40, true)`).
*
* Differences from the CLI:
* - Triggers only on `stopReason === 'end_turn'` (the daemon
Expand Down Expand Up @@ -3929,24 +3929,27 @@ export class Session implements SessionContext {

void (async () => {
try {
const fullHistory = chat.getHistory(true);
const lastEntry = fullHistory[fullHistory.length - 1];
const conversationHistory = chat.getHistoryTail(40, true);
Comment thread
yiliang114 marked this conversation as resolved.
const lastEntry = conversationHistory[conversationHistory.length - 1];
Comment thread
yiliang114 marked this conversation as resolved.
if (!lastEntry || lastEntry.role !== 'model') {
debugLogger.debug(
'Skipping followup suggestion: last history entry is not model',
);
return;
}
const conversationHistory =
fullHistory.length > 40 ? fullHistory.slice(-40) : fullHistory;

const r = await generatePromptSuggestion(
this.config,
conversationHistory,
ac.signal,
{
// On by default: the schema declares `default: true`, but
// `mergeSettings` doesn't apply schema defaults, so an unset value
// is `undefined` and a `=== true` gate left the cache-aware fork
// as dead code unless the flag was explicitly set (#9230). Mirrors
// AppContainer — only an explicit `false` opts out.
enableCacheSharing:
this.settings.merged.ui?.enableCacheSharing === true,
this.settings.merged.ui?.enableCacheSharing !== false,
Comment thread
yiliang114 marked this conversation as resolved.
Comment thread
yiliang114 marked this conversation as resolved.
Comment thread
yiliang114 marked this conversation as resolved.
},
);
if (ac.signal.aborted) return;
Expand Down
8 changes: 7 additions & 1 deletion packages/cli/src/ui/AppContainer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3153,7 +3153,13 @@ export const AppContainer = (props: AppContainerProps) => {
// causes transient heap peaks that trigger OOM (#4624).
const conversationHistory = geminiClient.getHistoryTail(40, true);
generatePromptSuggestion(config, conversationHistory, ac.signal, {
enableCacheSharing: settings.merged.ui?.enableCacheSharing === true,
// On by default: the schema declares `default: true`, but
// `mergeSettings` doesn't apply schema defaults, so an unset value is
// `undefined` and a `=== true` gate left the cache-aware fork as dead
// code unless the flag was explicitly set (#9230). Same treatment as
// `enableFollowupSuggestions` above — only an explicit `false` opts
// out.
enableCacheSharing: settings.merged.ui?.enableCacheSharing !== false,
Comment thread
yiliang114 marked this conversation as resolved.
Comment thread
yiliang114 marked this conversation as resolved.
Comment thread
yiliang114 marked this conversation as resolved.
})
.then((result) => {
if (ac.signal.aborted) return;
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/core/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4809,6 +4809,7 @@ describe('Gemini Client (client.ts)', () => {
const history = JSON.stringify(getCacheSafeParams()?.history);
expect(history).not.toContain('image-bytes');
expect(history).toContain('pdf-bytes');
expect(getCacheSafeParams()?.sessionId).toBe('test-session-id');
});

it.each([
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3871,6 +3871,7 @@ export class GeminiClient {
chat.getGenerationConfig(),
cachedHistory,
this.config.getModel(),
this.config.getSessionId(),
Comment thread
yiliang114 marked this conversation as resolved.
);
} catch {
// Best-effort — don't block the main flow
Expand Down
106 changes: 98 additions & 8 deletions packages/core/src/followup/suggestionGenerator.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,33 @@ import type { Content } from '@google/genai';
import { beforeEach, describe, it, expect, vi } from 'vitest';
import type { Config } from '../config/config.js';

const { mockGetCacheSafeParams, mockRunForkedAgent } = vi.hoisted(() => ({
const {
mockGetCacheSafeParams,
mockGetCacheSafeParamsSessionId,
mockRunForkedAgent,
mockRunSideQuery,
} = vi.hoisted(() => ({
mockGetCacheSafeParams: vi.fn(),
mockGetCacheSafeParamsSessionId: vi.fn(),
mockRunForkedAgent: vi.fn(),
mockRunSideQuery: vi.fn(),
}));

vi.mock('../utils/sideQuery.js', async (importOriginal) => {
const actual = await importOriginal<typeof import('../utils/sideQuery.js')>();
return {
...actual,
runSideQuery: mockRunSideQuery,
};
});

vi.mock('../utils/forkedAgent.js', async (importOriginal) => {
const actual =
await importOriginal<typeof import('../utils/forkedAgent.js')>();
return {
...actual,
getCacheSafeParams: mockGetCacheSafeParams,
getCacheSafeParamsSessionId: mockGetCacheSafeParamsSessionId,
runForkedAgent: mockRunForkedAgent,
};
});
Expand All @@ -39,7 +55,10 @@ const conversationHistory: Content[] = [
describe('generatePromptSuggestion', () => {
beforeEach(() => {
mockGetCacheSafeParams.mockReset();
mockGetCacheSafeParamsSessionId.mockReset();
mockGetCacheSafeParamsSessionId.mockReturnValue('test-session');
mockRunForkedAgent.mockReset();
mockRunSideQuery.mockReset();
});

it('passes cache-safe model in cache mode when no explicit or fast model exists', async () => {
Expand All @@ -48,6 +67,7 @@ describe('generatePromptSuggestion', () => {
history: conversationHistory,
model: 'main-model',
version: 1,
sessionId: 'test-session',
});
mockRunForkedAgent.mockResolvedValue({
text: null,
Expand All @@ -57,17 +77,16 @@ describe('generatePromptSuggestion', () => {
const config = {
getFastModel: vi.fn(() => undefined),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'test-session'),
} as unknown as Config;

await generatePromptSuggestion(
config,
conversationHistory,
new AbortController().signal,
{ enableCacheSharing: true },
);
const signal = new AbortController().signal;
await generatePromptSuggestion(config, conversationHistory, signal, {
enableCacheSharing: true,
});

expect(mockRunForkedAgent).toHaveBeenCalledWith(
expect.objectContaining({ model: 'main-model' }),
expect.objectContaining({ model: 'main-model', abortSignal: signal }),
);
});

Expand All @@ -77,6 +96,7 @@ describe('generatePromptSuggestion', () => {
history: conversationHistory,
model: 'main-model',
version: 1,
sessionId: 'test-session',
});
mockRunForkedAgent.mockResolvedValue({
text: null,
Expand All @@ -86,6 +106,7 @@ describe('generatePromptSuggestion', () => {
const config = {
getFastModel: vi.fn(() => 'openai:fast-model'),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'test-session'),
} as unknown as Config;

await generatePromptSuggestion(
Expand All @@ -105,6 +126,7 @@ describe('generatePromptSuggestion', () => {
history: conversationHistory,
model: 'main-model',
version: 1,
sessionId: 'test-session',
});
mockRunForkedAgent.mockResolvedValue({
text: null,
Expand All @@ -114,6 +136,7 @@ describe('generatePromptSuggestion', () => {
const config = {
getFastModel: vi.fn(() => undefined),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'test-session'),
} as unknown as Config;

await generatePromptSuggestion(
Expand All @@ -128,12 +151,78 @@ describe('generatePromptSuggestion', () => {
);
});

it('falls back to the base LLM when the cache-safe slot belongs to another session', async () => {
// The cache-safe slot is a process-global: in a multi-session daemon it
// can hold ANOTHER session's transcript + systemInstruction. The
// suggestion must NOT fork from a foreign session's params (cross-session
// content leak) — it falls back to the session-safe base-LLM path
// (#9233).
mockGetCacheSafeParamsSessionId.mockReturnValue('session-B');
mockRunSideQuery.mockResolvedValue({
text: '{"suggestion":"from base llm"}',
usage: { inputTokens: 1, outputTokens: 1 },
});
const config = {
getFastModel: vi.fn(() => undefined),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'session-A'), // this session is different
} as unknown as Config;

const result = await generatePromptSuggestion(
config,
conversationHistory,
new AbortController().signal,
{ enableCacheSharing: true },
);

// The foreign slot is rejected before cloning the full cached payload.
expect(mockGetCacheSafeParams).not.toHaveBeenCalled();
// 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.
expect(mockRunSideQuery).toHaveBeenCalled();
expect(result.suggestion).toBe('from base llm');
});

it('does not use the cache-safe fork when cache sharing is disabled', async () => {
mockGetCacheSafeParams.mockReturnValue({
generationConfig: {},
history: conversationHistory,
model: 'main-model',
version: 1,
sessionId: 'test-session',
});
mockRunSideQuery.mockResolvedValue({
text: '{"suggestion":"from base llm"}',
usage: { inputTokens: 1, outputTokens: 1 },
});
const config = {
getFastModel: vi.fn(() => undefined),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'test-session'),
} as unknown as Config;

const result = await generatePromptSuggestion(
config,
conversationHistory,
new AbortController().signal,
{ enableCacheSharing: false },
);

expect(mockGetCacheSafeParamsSessionId).not.toHaveBeenCalled();
expect(mockGetCacheSafeParams).not.toHaveBeenCalled();
expect(mockRunForkedAgent).not.toHaveBeenCalled();
expect(mockRunSideQuery).toHaveBeenCalled();
expect(result.suggestion).toBe('from base llm');
});

it('passes preserveTools: false when fast model differs from cache-safe model', async () => {
mockGetCacheSafeParams.mockReturnValue({
generationConfig: {},
history: conversationHistory,
model: 'main-model',
version: 1,
sessionId: 'test-session',
});
mockRunForkedAgent.mockResolvedValue({
text: null,
Expand All @@ -143,6 +232,7 @@ describe('generatePromptSuggestion', () => {
const config = {
getFastModel: vi.fn(() => 'different-fast-model'),
getModel: vi.fn(() => 'main-model'),
getSessionId: vi.fn(() => 'test-session'),
} as unknown as Config;

await generatePromptSuggestion(
Expand Down
Loading
Loading