mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-28 05:55:17 -07:00
fix(core): dynamic session ID injection to resolve resume bugs (#24972)
This commit is contained in:
committed by
GitHub
parent
80764c8bb5
commit
d06dba3538
@@ -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
@@ -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.
|
||||
|
||||
@@ -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'),
|
||||
}),
|
||||
);
|
||||
|
||||
|
||||
@@ -107,7 +107,7 @@ export async function startInteractiveUI(
|
||||
<TerminalProvider>
|
||||
<ScrollProvider>
|
||||
<OverflowProvider>
|
||||
<SessionStatsProvider>
|
||||
<SessionStatsProvider sessionId={config.getSessionId()}>
|
||||
<VimModeProvider>
|
||||
<AppContainer
|
||||
config={config}
|
||||
|
||||
@@ -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}
|
||||
>
|
||||
|
||||
@@ -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: () => [],
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -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;
|
||||
};
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
@@ -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';
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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';
|
||||
},
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -6,8 +6,6 @@
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
export const sessionId = randomUUID();
|
||||
|
||||
export function createSessionId(): string {
|
||||
return randomUUID();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user