fix(core): dynamic session ID injection to resolve resume bugs (#24972)

This commit is contained in:
Tommaso Sciortino
2026-04-08 23:27:24 +00:00
committed by GitHub
parent 80764c8bb5
commit d06dba3538
33 changed files with 165 additions and 189 deletions
+1 -1
View File
@@ -372,7 +372,7 @@ export class GeminiAgent {
mcpServers,
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(config.storage);
const { sessionData, sessionPath } =
await sessionSelector.resolveSession(sessionId);
+37 -35
View File
@@ -13,7 +13,7 @@ import {
type OutputPayload,
type ConsoleLogPayload,
type UserFeedbackPayload,
sessionId,
createSessionId,
logUserPrompt,
AuthType,
UserPromptEvent,
@@ -33,6 +33,7 @@ import {
type AdminControlsSettings,
debugLogger,
isHeadlessMode,
Storage,
} from '@google/gemini-cli-core';
import { loadCliConfig, parseArguments } from './config/config.js';
@@ -185,6 +186,39 @@ ${reason.stack}`
});
}
export async function resolveSessionId(resumeArg: string | undefined): Promise<{
sessionId: string;
resumedSessionData?: ResumedSessionData;
}> {
if (!resumeArg) {
return { sessionId: createSessionId() };
}
const storage = new Storage(process.cwd());
await storage.initialize();
try {
const { sessionData, sessionPath } = await new SessionSelector(
storage,
).resolveSession(resumeArg);
return {
sessionId: sessionData.sessionId,
resumedSessionData: { conversation: sessionData, filePath: sessionPath },
};
} catch (error) {
if (error instanceof SessionError && error.code === 'NO_SESSIONS_FOUND') {
coreEvents.emitFeedback('warning', error.message);
return { sessionId: createSessionId() };
}
coreEvents.emitFeedback(
'error',
`Error resuming session: ${error instanceof Error ? error.message : 'Unknown error'}`,
);
await runExitCleanup();
process.exit(ExitCodes.FATAL_INPUT_ERROR);
}
}
export async function startInteractiveUI(
config: Config,
settings: LoadedSettings,
@@ -280,6 +314,8 @@ export async function main() {
const argv = await argvPromise;
const { sessionId, resumedSessionData } = await resolveSessionId(argv.resume);
if (
(argv.allowedTools && argv.allowedTools.length > 0) ||
(settings.merged.tools?.allowed && settings.merged.tools.allowed.length > 0)
@@ -599,40 +635,6 @@ export async function main() {
})),
];
// Handle --resume flag
let resumedSessionData: ResumedSessionData | undefined = undefined;
if (argv.resume) {
const sessionSelector = new SessionSelector(config);
try {
const result = await sessionSelector.resolveSession(argv.resume);
resumedSessionData = {
conversation: result.sessionData,
filePath: result.sessionPath,
};
// Use the existing session ID to continue recording to the same session
config.setSessionId(resumedSessionData.conversation.sessionId);
} catch (error) {
if (
error instanceof SessionError &&
error.code === 'NO_SESSIONS_FOUND'
) {
// No sessions to resume — start a fresh session with a warning
startupWarnings.push({
id: 'resume-no-sessions',
message: error.message,
priority: WarningPriority.High,
});
} else {
coreEvents.emitFeedback(
'error',
`Error resuming session: ${error instanceof Error ? error.message : 'Unknown error'}`,
);
await runExitCleanup();
process.exit(ExitCodes.FATAL_INPUT_ERROR);
}
}
}
cliStartupHandle?.end();
// Render UI, passing necessary config values. Check that there is no command line question.
+3
View File
@@ -73,6 +73,7 @@ vi.mock('./config/config.js', () => ({
getSandbox: vi.fn(() => false),
getQuestion: vi.fn(() => ''),
isInteractive: () => false,
getSessionId: vi.fn().mockReturnValue('test-session-id'),
storage: { initialize: vi.fn().mockResolvedValue(undefined) },
} as unknown as Config),
parseArguments: vi.fn().mockResolvedValue({}),
@@ -213,6 +214,7 @@ describe('gemini.tsx main function cleanup', () => {
getSandbox: vi.fn(() => false),
getDebugMode: vi.fn(() => false),
getPolicyEngine: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
getMessageBus: () => ({ subscribe: vi.fn() }),
getEnableHooks: vi.fn(() => false),
getHookSystem: () => undefined,
@@ -273,6 +275,7 @@ describe('gemini.tsx main function cleanup', () => {
vi.mocked(loadCliConfig).mockResolvedValue(
buildMockConfig({
getHookSystem: vi.fn(() => mockHookSystem),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
}),
);
+1 -1
View File
@@ -107,7 +107,7 @@ export async function startInteractiveUI(
<TerminalProvider>
<ScrollProvider>
<OverflowProvider>
<SessionStatsProvider>
<SessionStatsProvider sessionId={config.getSessionId()}>
<VimModeProvider>
<AppContainer
config={config}
+1 -1
View File
@@ -731,7 +731,7 @@ export const renderWithProviders = async (
<UIStateContext.Provider value={finalUiState}>
<VimModeProvider>
<ShellFocusContext.Provider value={shellFocus}>
<SessionStatsProvider>
<SessionStatsProvider sessionId={config.getSessionId()}>
<StreamingContext.Provider
value={finalUiState.streamingState}
>
+1 -1
View File
@@ -444,7 +444,7 @@ export const AppContainer = (props: AppContainerProps) => {
const [isConfigInitialized, setConfigInitialized] = useState(false);
const logger = useLogger(config.storage);
const logger = useLogger(config);
const { inputHistory, addInput, initializeFromLogger } =
useInputHistoryStore();
@@ -9,7 +9,7 @@ import open from 'open';
import path from 'node:path';
import { bugCommand } from './bugCommand.js';
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
import { getVersion } from '@google/gemini-cli-core';
import { getVersion, type Config } from '@google/gemini-cli-core';
import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
import { formatBytes } from '../utils/formatters.js';
@@ -89,7 +89,8 @@ describe('bugCommand', () => {
getBugCommand: () => undefined,
getIdeMode: () => true,
getContentGeneratorConfig: () => ({ authType: 'oauth-personal' }),
},
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config,
geminiClient: {
getChat: () => ({
getHistory: () => [],
@@ -137,7 +138,8 @@ describe('bugCommand', () => {
storage: {
getProjectTempDir: () => '/tmp/gemini',
},
},
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config,
geminiClient: {
getChat: () => ({
getHistory: () => history,
@@ -182,7 +184,8 @@ describe('bugCommand', () => {
getBugCommand: () => ({ urlTemplate: customTemplate }),
getIdeMode: () => true,
getContentGeneratorConfig: () => ({ authType: 'vertex-ai' }),
},
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config,
geminiClient: {
getChat: () => ({
getHistory: () => [],
+1 -2
View File
@@ -16,7 +16,6 @@ import { GIT_COMMIT_INFO } from '../../generated/git-commit.js';
import { formatBytes } from '../utils/formatters.js';
import {
IdeClient,
sessionId,
getVersion,
INITIAL_HISTORY_LENGTH,
debugLogger,
@@ -59,7 +58,7 @@ export const bugCommand: SlashCommand = {
let info = `
* **CLI Version:** ${cliVersion}
* **Git Commit:** ${GIT_COMMIT_INFO}
* **Session ID:** ${sessionId}
* **Session ID:** ${config?.getSessionId() || 'Unknown'}
* **Operating System:** ${osVersion}
* **Sandbox Environment:** ${sandboxEnv}
* **Model Version:** ${modelVersion}
@@ -158,6 +158,7 @@ Implement a comprehensive authentication system with multiple providers.
getIdeMode: () => false,
isTrustedFolder: () => true,
getPreferredEditor: () => undefined,
getSessionId: () => 'test-session-id',
storage: {
getPlansDir: () => mockPlansDir,
},
@@ -464,6 +465,7 @@ Implement a comprehensive authentication system with multiple providers.
getTargetDir: () => mockTargetDir,
getIdeMode: () => false,
isTrustedFolder: () => true,
getSessionId: () => 'test-session-id',
storage: {
getPlansDir: () => mockPlansDir,
},
@@ -82,6 +82,7 @@ const mockConfigPlain = {
getExtensionRegistryURI: () => undefined,
getContentGeneratorConfig: () => ({ authType: undefined }),
getSandboxEnabled: () => false,
getSessionId: () => 'test-session-id',
};
const mockConfig = mockConfigPlain as unknown as Config;
@@ -124,7 +124,7 @@ describe('<HistoryItemDisplay />', () => {
duration: '1s',
};
const { lastFrame, unmount } = await renderWithProviders(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<HistoryItemDisplay {...baseItem} item={item} />
</SessionStatsProvider>,
);
@@ -157,7 +157,7 @@ describe('<HistoryItemDisplay />', () => {
type: 'model_stats',
};
const { lastFrame, unmount } = await renderWithProviders(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<HistoryItemDisplay {...baseItem} item={item} />
</SessionStatsProvider>,
);
@@ -173,7 +173,7 @@ describe('<HistoryItemDisplay />', () => {
type: 'tool_stats',
};
const { lastFrame, unmount } = await renderWithProviders(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<HistoryItemDisplay {...baseItem} item={item} />
</SessionStatsProvider>,
);
@@ -190,7 +190,7 @@ describe('<HistoryItemDisplay />', () => {
duration: '1s',
};
const { lastFrame, unmount } = await renderWithProviders(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<HistoryItemDisplay {...baseItem} item={item} />
</SessionStatsProvider>,
);
@@ -86,6 +86,7 @@ describe('<ModelDialog />', () => {
getProModelNoAccess: mockGetProModelNoAccess,
getProModelNoAccessSync: mockGetProModelNoAccessSync,
getLastRetrievedQuota: () => ({ buckets: [] }),
getSessionId: () => 'test-session-id',
};
beforeEach(() => {
@@ -55,6 +55,7 @@ describe('ToolConfirmationQueue', () => {
getFileSystemService: () => ({
readFile: vi.fn().mockResolvedValue('Plan content'),
}),
getSessionId: () => 'test-session-id',
storage: {
getPlansDir: () => '/mock/temp/plans',
},
@@ -60,7 +60,7 @@ describe('SessionStatsContext', () => {
> = { current: undefined };
const { unmount } = await render(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<TestHarness contextRef={contextRef} />
</SessionStatsProvider>,
);
@@ -79,7 +79,7 @@ describe('SessionStatsContext', () => {
> = { current: undefined };
const { unmount } = await render(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<TestHarness contextRef={contextRef} />
</SessionStatsProvider>,
);
@@ -162,7 +162,7 @@ describe('SessionStatsContext', () => {
};
const { unmount } = await render(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<CountingTestHarness />
</SessionStatsProvider>,
);
@@ -245,7 +245,7 @@ describe('SessionStatsContext', () => {
> = { current: undefined };
const { unmount } = await render(
<SessionStatsProvider>
<SessionStatsProvider sessionId="test-session-id">
<TestHarness contextRef={contextRef} />
</SessionStatsProvider>,
);
@@ -13,14 +13,13 @@ import {
useMemo,
useEffect,
} from 'react';
import type {
SessionMetrics,
ModelMetrics,
RoleMetrics,
ToolCallStats,
} from '@google/gemini-cli-core';
import { uiTelemetryService, sessionId } from '@google/gemini-cli-core';
import { uiTelemetryService } from '@google/gemini-cli-core';
export enum ToolCallDecision {
ACCEPT = 'accept',
@@ -183,9 +182,10 @@ const SessionStatsContext = createContext<SessionStatsContextValue | undefined>(
// --- Provider Component ---
export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
children,
}) => {
export const SessionStatsProvider: React.FC<{
children: React.ReactNode;
sessionId: string;
}> = ({ children, sessionId }) => {
const [stats, setStats] = useState<SessionStatsState>({
sessionId,
sessionStartTime: new Date(),
+6 -14
View File
@@ -262,14 +262,13 @@ export const useGeminiStream = (
useStateAndRef<boolean>(true);
const processedMemoryToolsRef = useRef<Set<string>>(new Set());
const { startNewPrompt, getPromptCount } = useSessionStats();
const storage = config.storage;
const logger = useLogger(storage);
const logger = useLogger(config);
const gitService = useMemo(() => {
if (!config.getProjectRoot()) {
return;
}
return new GitService(config.getProjectRoot(), storage);
}, [config, storage]);
return new GitService(config.getProjectRoot(), config.storage);
}, [config]);
useEffect(() => {
const handleRetryAttempt = (payload: RetryAttemptPayload) => {
@@ -1580,6 +1579,7 @@ export const useGeminiStream = (
operation: options?.isContinuation
? GeminiCliOperation.SystemPrompt
: GeminiCliOperation.UserPrompt,
sessionId: config.getSessionId(),
},
async ({ metadata: spanMetadata }) => {
spanMetadata.input = query;
@@ -2105,7 +2105,7 @@ export const useGeminiStream = (
}
if (checkpointsToWrite.size > 0) {
const checkpointDir = storage.getProjectTempCheckpointsDir();
const checkpointDir = config.storage.getProjectTempCheckpointsDir();
try {
await fs.mkdir(checkpointDir, { recursive: true });
for (const [fileName, content] of checkpointsToWrite) {
@@ -2122,15 +2122,7 @@ export const useGeminiStream = (
};
// eslint-disable-next-line @typescript-eslint/no-floating-promises
saveRestorableToolCalls();
}, [
toolCalls,
config,
onDebugMessage,
gitService,
history,
geminiClient,
storage,
]);
}, [toolCalls, config, onDebugMessage, gitService, history, geminiClient]);
const lastOutputTime = Math.max(
lastToolOutputTime,
+4 -31
View File
@@ -8,14 +8,7 @@ import { act } from 'react';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { renderHook } from '../../test-utils/render.js';
import { useLogger } from './useLogger.js';
import {
sessionId as globalSessionId,
Logger,
type Storage,
type Config,
} from '@google/gemini-cli-core';
import { ConfigContext } from '../contexts/ConfigContext.js';
import type React from 'react';
import { Logger, type Storage, type Config } from '@google/gemini-cli-core';
let deferredInit: { resolve: (val?: unknown) => void };
@@ -41,35 +34,15 @@ describe('useLogger', () => {
const mockStorage = {} as Storage;
const mockConfig = {
getSessionId: vi.fn().mockReturnValue('active-session-id'),
storage: mockStorage,
} as unknown as Config;
beforeEach(() => {
vi.clearAllMocks();
});
it('should initialize with the global sessionId by default', async () => {
const { result } = await renderHook(() => useLogger(mockStorage));
expect(result.current).toBeNull();
await act(async () => {
deferredInit.resolve();
});
expect(result.current).not.toBeNull();
expect(Logger).toHaveBeenCalledWith(globalSessionId, mockStorage);
});
it('should initialize with the active sessionId from ConfigContext when available', async () => {
const wrapper = ({ children }: { children: React.ReactNode }) => (
<ConfigContext.Provider value={mockConfig}>
{children}
</ConfigContext.Provider>
);
const { result } = await renderHook(() => useLogger(mockStorage), {
wrapper,
});
it('should initialize with the sessionId from config', async () => {
const { result } = await renderHook(() => useLogger(mockConfig));
expect(result.current).toBeNull();
+6 -15
View File
@@ -4,24 +4,17 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { useState, useEffect, useContext } from 'react';
import {
sessionId as globalSessionId,
Logger,
type Storage,
} from '@google/gemini-cli-core';
import { ConfigContext } from '../contexts/ConfigContext.js';
import { useState, useEffect } from 'react';
import { Logger, type Config } from '@google/gemini-cli-core';
/**
* Hook to manage the logger instance.
*/
export const useLogger = (storage: Storage): Logger | null => {
export const useLogger = (config: Config): Logger | null => {
const [logger, setLogger] = useState<Logger | null>(null);
const config = useContext(ConfigContext);
useEffect(() => {
const activeSessionId = config?.getSessionId() ?? globalSessionId;
const newLogger = new Logger(activeSessionId, storage);
const newLogger = new Logger(config.getSessionId(), config.storage);
/**
* Start async initialization, no need to await. Using await slows down the
@@ -30,11 +23,9 @@ export const useLogger = (storage: Storage): Logger | null => {
*/
newLogger
.initialize()
.then(() => {
setLogger(newLogger);
})
.then(() => setLogger(newLogger))
.catch(() => {});
}, [storage, config]);
}, [config]);
return logger;
};
+19 -25
View File
@@ -15,7 +15,7 @@ import {
} from './sessionUtils.js';
import {
SESSION_FILE_PREFIX,
type Config,
type Storage,
type MessageRecord,
CoreToolCallStatus,
} from '@google/gemini-cli-core';
@@ -25,20 +25,17 @@ import { randomUUID } from 'node:crypto';
describe('SessionSelector', () => {
let tmpDir: string;
let config: Config;
let storage: Storage;
beforeEach(async () => {
// Create a temporary directory for testing
tmpDir = path.join(process.cwd(), '.tmp-test-sessions');
await fs.mkdir(tmpDir, { recursive: true });
// Mock config
config = {
storage: {
getProjectTempDir: () => tmpDir,
},
getSessionId: () => 'current-session-id',
} as Partial<Config> as Config;
// Mock storage
storage = {
getProjectTempDir: () => tmpDir,
} as Partial<Storage> as Storage;
});
afterEach(async () => {
@@ -104,7 +101,7 @@ describe('SessionSelector', () => {
JSON.stringify(session2, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
// Test resolving by UUID
const result1 = await sessionSelector.resolveSession(sessionId1);
@@ -170,7 +167,7 @@ describe('SessionSelector', () => {
JSON.stringify(session2, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
// Test resolving by index (1-based)
const result1 = await sessionSelector.resolveSession('1');
@@ -234,7 +231,7 @@ describe('SessionSelector', () => {
JSON.stringify(session2, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
// Test resolving latest
const result = await sessionSelector.resolveSession('latest');
@@ -271,7 +268,7 @@ describe('SessionSelector', () => {
JSON.stringify(session, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
// Test resolving by UUID with leading/trailing spaces
const result = await sessionSelector.resolveSession(` ${sessionId} `);
@@ -334,7 +331,7 @@ describe('SessionSelector', () => {
JSON.stringify(sessionDuplicate, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
const sessions = await sessionSelector.listSessions();
expect(sessions.length).toBe(1);
@@ -373,7 +370,7 @@ describe('SessionSelector', () => {
JSON.stringify(session1, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
await expect(
sessionSelector.resolveSession('invalid-uuid'),
@@ -389,14 +386,11 @@ describe('SessionSelector', () => {
const chatsDir = path.join(tmpDir, 'chats');
await fs.mkdir(chatsDir, { recursive: true });
const emptyConfig = {
storage: {
getProjectTempDir: () => tmpDir,
},
getSessionId: () => 'current-session-id',
} as Partial<Config> as Config;
const emptyStorage = {
getProjectTempDir: () => tmpDir,
} as Partial<Storage> as Storage;
const sessionSelector = new SessionSelector(emptyConfig);
const sessionSelector = new SessionSelector(emptyStorage);
await expect(sessionSelector.resolveSession('latest')).rejects.toSatisfy(
(error) => {
@@ -469,7 +463,7 @@ describe('SessionSelector', () => {
JSON.stringify(sessionSystemOnly, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
const sessions = await sessionSelector.listSessions();
// Should only list the session with user message
@@ -508,7 +502,7 @@ describe('SessionSelector', () => {
JSON.stringify(sessionGeminiOnly, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
const sessions = await sessionSelector.listSessions();
// Should list the session with gemini message
@@ -574,7 +568,7 @@ describe('SessionSelector', () => {
JSON.stringify(subagentSession, null, 2),
);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(storage);
const sessions = await sessionSelector.listSessions();
// Should only list the main session
+6 -15
View File
@@ -9,7 +9,7 @@ import {
partListUnionToString,
SESSION_FILE_PREFIX,
CoreToolCallStatus,
type Config,
type Storage,
type ConversationRecord,
type MessageRecord,
} from '@google/gemini-cli-core';
@@ -399,17 +399,14 @@ export const getSessionFiles = async (
* Utility class for session discovery and selection.
*/
export class SessionSelector {
constructor(private config: Config) {}
constructor(private storage: Storage) {}
/**
* Lists all available sessions for the current project.
*/
async listSessions(): Promise<SessionInfo[]> {
const chatsDir = path.join(
this.config.storage.getProjectTempDir(),
'chats',
);
return getSessionFiles(chatsDir, this.config.getSessionId());
const chatsDir = path.join(this.storage.getProjectTempDir(), 'chats');
return getSessionFiles(chatsDir);
}
/**
@@ -452,10 +449,7 @@ export class SessionSelector {
return sortedSessions[index - 1];
}
const chatsDir = path.join(
this.config.storage.getProjectTempDir(),
'chats',
);
const chatsDir = path.join(this.storage.getProjectTempDir(), 'chats');
throw SessionError.invalidSessionIdentifier(trimmedIdentifier, chatsDir);
}
@@ -507,10 +501,7 @@ export class SessionSelector {
private async selectSession(
sessionInfo: SessionInfo,
): Promise<SessionSelectionResult> {
const chatsDir = path.join(
this.config.storage.getProjectTempDir(),
'chats',
);
const chatsDir = path.join(this.storage.getProjectTempDir(), 'chats');
const sessionPath = path.join(chatsDir, sessionInfo.fileName);
try {
+2 -2
View File
@@ -21,7 +21,7 @@ export async function listSessions(config: Config): Promise<void> {
// Generate summary for most recent session if needed
await generateSummary(config);
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(config.storage);
const sessions = await sessionSelector.listSessions();
if (sessions.length === 0) {
@@ -55,7 +55,7 @@ export async function deleteSession(
config: Config,
sessionIndex: string,
): Promise<void> {
const sessionSelector = new SessionSelector(config);
const sessionSelector = new SessionSelector(config.storage);
const sessions = await sessionSelector.listSessions();
if (sessions.length === 0) {
@@ -182,6 +182,7 @@ class SubAgentInvocation extends BaseToolInvocation<AgentInputs, ToolResult> {
{
operation: GeminiCliOperation.AgentCall,
logPrompts: this.context.config.getTelemetryLogPromptsEnabled(),
sessionId: this.context.config.getSessionId(),
attributes: {
[GEN_AI_AGENT_NAME]: this.definition.name,
[GEN_AI_AGENT_DESCRIPTION]: this.definition.description,
@@ -74,6 +74,7 @@ describe('LoggingContentGenerator', () => {
}),
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(true),
refreshUserQuotaIfStale: vi.fn().mockResolvedValue(undefined),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config;
loggingContentGenerator = new LoggingContentGenerator(wrapped, config);
vi.useFakeTimers();
@@ -350,6 +350,7 @@ export class LoggingContentGenerator implements ContentGenerator {
{
operation: GeminiCliOperation.LLMCall,
logPrompts: this.config.getTelemetryLogPromptsEnabled(),
sessionId: this.config.getSessionId(),
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
[GEN_AI_PROMPT_NAME]: userPromptId,
@@ -440,6 +441,7 @@ export class LoggingContentGenerator implements ContentGenerator {
{
operation: GeminiCliOperation.LLMCall,
logPrompts: this.config.getTelemetryLogPromptsEnabled(),
sessionId: this.config.getSessionId(),
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
[GEN_AI_PROMPT_NAME]: userPromptId,
@@ -594,6 +596,7 @@ export class LoggingContentGenerator implements ContentGenerator {
{
operation: GeminiCliOperation.LLMCall,
logPrompts: this.config.getTelemetryLogPromptsEnabled(),
sessionId: this.config.getSessionId(),
attributes: {
[GEN_AI_REQUEST_MODEL]: req.model,
},
+1 -1
View File
@@ -252,7 +252,7 @@ export * from './telemetry/index.js';
export * from './telemetry/billingEvents.js';
export { logBillingEvent } from './telemetry/loggers.js';
export * from './telemetry/constants.js';
export { sessionId, createSessionId } from './utils/session.js';
export { createSessionId } from './utils/session.js';
export * from './utils/compatibility.js';
export * from './utils/browser.js';
export { Storage } from './config/storage.js';
+23 -14
View File
@@ -51,8 +51,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
@@ -79,8 +79,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
@@ -161,8 +161,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
@@ -226,8 +226,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const toolCall = {
request: { name: 'test-tool', args: {}, isClientInitiated: true },
tool: { name: 'test-tool' },
@@ -243,8 +243,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -273,8 +273,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -307,6 +307,7 @@ describe('policy.ts', () => {
isTrustedFolder: vi.fn().mockReturnValue(false),
getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
@@ -339,8 +340,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -379,8 +380,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -420,8 +421,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -447,8 +448,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -473,8 +474,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -499,8 +500,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -540,8 +541,8 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
mockConfig as Config;
const mockMessageBus = {
@@ -583,6 +584,7 @@ describe('policy.ts', () => {
isTrustedFolder: vi.fn().mockReturnValue(false),
getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config =
@@ -628,6 +630,7 @@ describe('policy.ts', () => {
.fn()
.mockReturnValue('/mock/project/policies'),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const mockMessageBus = {
publish: vi.fn(),
@@ -659,6 +662,7 @@ describe('policy.ts', () => {
.fn()
.mockReturnValue('/mock/project/policies'),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const mockMessageBus = {
publish: vi.fn(),
@@ -689,6 +693,7 @@ describe('policy.ts', () => {
getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
getTargetDir: vi.fn().mockReturnValue('/mock/dir'),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const mockMessageBus = {
publish: vi.fn(),
@@ -727,6 +732,7 @@ describe('policy.ts', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const mockMessageBus = {
publish: vi.fn(),
@@ -766,6 +772,7 @@ describe('policy.ts', () => {
it('should return default denial message when no rule provided', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config;
(mockConfig as unknown as { config: Config }).config = mockConfig;
@@ -779,6 +786,7 @@ describe('policy.ts', () => {
it('should return custom deny message if provided', () => {
const mockConfig = {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Config;
(mockConfig as unknown as { config: Config }).config = mockConfig;
@@ -840,7 +848,6 @@ describe('Plan Mode Denial Consistency', () => {
publish: vi.fn(),
subscribe: vi.fn(),
} as unknown as Mocked<MessageBus>;
mockConfig = {
getPolicyEngine: vi.fn().mockReturnValue(mockPolicyEngine),
toolRegistry: mockToolRegistry,
@@ -852,6 +859,7 @@ describe('Plan Mode Denial Consistency', () => {
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.PLAN), // Key: Plan Mode
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false),
setApprovalMode: vi.fn(),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config = mockConfig as Config;
@@ -933,6 +941,7 @@ describe('Plan Mode Denial Consistency', () => {
getApprovalMode: vi.fn().mockReturnValue(currentMode),
isTrustedFolder: vi.fn().mockReturnValue(false),
getWorkspacePoliciesDir: vi.fn().mockReturnValue(undefined),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
const mockMessageBus = {
@@ -177,6 +177,7 @@ describe('Scheduler (Orchestrator)', () => {
setApprovalMode: vi.fn(),
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config = mockConfig as Config;
@@ -1423,6 +1424,7 @@ describe('Scheduler MCP Progress', () => {
setApprovalMode: vi.fn(),
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config = mockConfig as Config;
+1
View File
@@ -197,6 +197,7 @@ export class Scheduler {
{
operation: GeminiCliOperation.ScheduleToolCalls,
logPrompts: this.context.config.getTelemetryLogPromptsEnabled(),
sessionId: this.context.config.getSessionId(),
},
async ({ metadata: spanMetadata }) => {
const requests = Array.isArray(request) ? request : [request];
@@ -218,6 +218,7 @@ describe('Scheduler Parallel Execution', () => {
setApprovalMode: vi.fn(),
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
} as unknown as Mocked<Config>;
(mockConfig as unknown as { config: Config }).config = mockConfig as Config;
@@ -84,6 +84,7 @@ export class ToolExecutor {
{
operation: GeminiCliOperation.ToolCall,
logPrompts: this.config.getTelemetryLogPromptsEnabled(),
sessionId: this.config.getSessionId(),
attributes: {
[GEN_AI_TOOL_NAME]: toolName,
[GEN_AI_TOOL_CALL_ID]: callId,
+12 -9
View File
@@ -110,7 +110,7 @@ describe('runInDevTraceSpan', () => {
const fn = vi.fn(async () => 'result');
const result = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
fn,
);
@@ -125,7 +125,7 @@ describe('runInDevTraceSpan', () => {
it('should set default attributes on the span metadata', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
expect(metadata.attributes[GEN_AI_OPERATION_NAME]).toBe(
GeminiCliOperation.LLMCall,
@@ -143,7 +143,7 @@ describe('runInDevTraceSpan', () => {
it('should set span attributes from metadata on completion', async () => {
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
metadata.input = { query: 'hello' };
metadata.output = { response: 'world' };
@@ -169,9 +169,12 @@ describe('runInDevTraceSpan', () => {
it('should handle errors in the wrapped function', async () => {
const error = new Error('test error');
await expect(
runInDevTraceSpan({ operation: GeminiCliOperation.LLMCall }, async () => {
throw error;
}),
runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => {
throw error;
},
),
).rejects.toThrow(error);
expect(mockSpan.setStatus).toHaveBeenCalledWith({
@@ -189,7 +192,7 @@ describe('runInDevTraceSpan', () => {
}
const resultStream = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => testStream(),
);
@@ -212,7 +215,7 @@ describe('runInDevTraceSpan', () => {
}
const resultStream = await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async () => errorStream(),
);
@@ -231,7 +234,7 @@ describe('runInDevTraceSpan', () => {
});
await runInDevTraceSpan(
{ operation: GeminiCliOperation.LLMCall },
{ operation: GeminiCliOperation.LLMCall, sessionId: 'test-session-id' },
async ({ metadata }) => {
metadata.input = 'trigger error';
},
+6 -3
View File
@@ -23,7 +23,6 @@ import {
SERVICE_DESCRIPTION,
SERVICE_NAME,
} from './constants.js';
import { sessionId } from '../utils/session.js';
import { truncateString } from '../utils/textUtils.js';
@@ -96,10 +95,14 @@ export interface SpanMetadata {
* @returns The result of the function.
*/
export async function runInDevTraceSpan<R>(
opts: SpanOptions & { operation: GeminiCliOperation; logPrompts?: boolean },
opts: SpanOptions & {
operation: GeminiCliOperation;
logPrompts?: boolean;
sessionId: string;
},
fn: ({ metadata }: { metadata: SpanMetadata }) => Promise<R>,
): Promise<R> {
const { operation, logPrompts, ...restOfSpanOpts } = opts;
const { operation, logPrompts, sessionId, ...restOfSpanOpts } = opts;
const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION);
return tracer.startActiveSpan(operation, restOfSpanOpts, async (span) => {
-2
View File
@@ -6,8 +6,6 @@
import { randomUUID } from 'node:crypto';
export const sessionId = randomUUID();
export function createSessionId(): string {
return randomUUID();
}