From 7c667e100ef54c8996b8fa914fc7a35b470b9bbe Mon Sep 17 00:00:00 2001 From: shrutip90 Date: Wed, 3 Sep 2025 11:44:26 -0700 Subject: [PATCH] Override Gemini CLI trust with VScode workspace trust when in IDE (#7433) --- .../cli/src/config/trustedFolders.test.ts | 59 +++++++++++++++++++ packages/cli/src/config/trustedFolders.ts | 26 ++++++-- packages/cli/src/ui/App.test.tsx | 6 ++ packages/cli/src/ui/App.tsx | 30 ++++++++++ .../cli/src/ui/hooks/useIdeTrustListener.ts | 47 +++++++++++++++ packages/core/index.ts | 1 + packages/core/src/config/config.ts | 6 ++ packages/core/src/ide/ide-client.ts | 15 +++++ packages/core/src/ide/ideContext.ts | 1 + packages/core/src/index.ts | 1 + packages/core/src/utils/ide-trust.ts | 15 +++++ .../src/extension.test.ts | 6 ++ .../vscode-ide-companion/src/extension.ts | 5 +- .../src/ide-server.test.ts | 5 +- .../vscode-ide-companion/src/ide-server.ts | 54 ++++++++++------- .../src/open-files-manager.ts | 1 + 16 files changed, 248 insertions(+), 30 deletions(-) create mode 100644 packages/cli/src/ui/hooks/useIdeTrustListener.ts create mode 100644 packages/core/src/utils/ide-trust.ts diff --git a/packages/cli/src/config/trustedFolders.test.ts b/packages/cli/src/config/trustedFolders.test.ts index 23abee438f..c8b34b765b 100644 --- a/packages/cli/src/config/trustedFolders.test.ts +++ b/packages/cli/src/config/trustedFolders.test.ts @@ -255,3 +255,62 @@ describe('isWorkspaceTrusted', () => { expect(isWorkspaceTrusted(mockSettings)).toBe(true); }); }); + +import { getIdeTrust } from '@google/gemini-cli-core'; + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = await importOriginal>(); + return { + ...actual, + getIdeTrust: vi.fn(), + }; +}); + +describe('isWorkspaceTrusted with IDE override', () => { + const mockSettings: Settings = { + security: { + folderTrust: { + enabled: true, + }, + }, + }; + + it('should return true when ideTrust is true, ignoring config', () => { + vi.mocked(getIdeTrust).mockReturnValue(true); + // Even if config says don't trust, ideTrust should win. + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify({ [process.cwd()]: TrustLevel.DO_NOT_TRUST }), + ); + expect(isWorkspaceTrusted(mockSettings)).toBe(true); + }); + + it('should return false when ideTrust is false, ignoring config', () => { + vi.mocked(getIdeTrust).mockReturnValue(false); + // Even if config says trust, ideTrust should win. + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify({ [process.cwd()]: TrustLevel.TRUST_FOLDER }), + ); + expect(isWorkspaceTrusted(mockSettings)).toBe(false); + }); + + it('should fall back to config when ideTrust is undefined', () => { + vi.mocked(getIdeTrust).mockReturnValue(undefined); + vi.spyOn(fs, 'existsSync').mockReturnValue(true); + vi.spyOn(fs, 'readFileSync').mockReturnValue( + JSON.stringify({ [process.cwd()]: TrustLevel.TRUST_FOLDER }), + ); + expect(isWorkspaceTrusted(mockSettings)).toBe(true); + }); + + it('should always return true if folderTrust setting is disabled', () => { + const settings: Settings = { + security: { + folderTrust: { + enabled: false, + }, + }, + }; + vi.mocked(getIdeTrust).mockReturnValue(false); + expect(isWorkspaceTrusted(settings)).toBe(true); + }); +}); diff --git a/packages/cli/src/config/trustedFolders.ts b/packages/cli/src/config/trustedFolders.ts index 87271cf4be..c25ff5de06 100644 --- a/packages/cli/src/config/trustedFolders.ts +++ b/packages/cli/src/config/trustedFolders.ts @@ -7,7 +7,11 @@ import * as fs from 'node:fs'; import * as path from 'node:path'; import { homedir } from 'node:os'; -import { getErrorMessage, isWithinRoot } from '@google/gemini-cli-core'; +import { + getErrorMessage, + isWithinRoot, + getIdeTrust, +} from '@google/gemini-cli-core'; import type { Settings } from './settings.js'; import stripJsonComments from 'strip-json-comments'; @@ -159,11 +163,7 @@ export function isFolderTrustEnabled(settings: Settings): boolean { return folderTrustSetting; } -export function isWorkspaceTrusted(settings: Settings): boolean | undefined { - if (!isFolderTrustEnabled(settings)) { - return true; - } - +function getWorkspaceTrustFromLocalConfig(): boolean | undefined { const folders = loadTrustedFolders(); if (folders.errors.length > 0) { @@ -176,3 +176,17 @@ export function isWorkspaceTrusted(settings: Settings): boolean | undefined { return folders.isPathTrusted(process.cwd()); } + +export function isWorkspaceTrusted(settings: Settings): boolean | undefined { + if (!isFolderTrustEnabled(settings)) { + return true; + } + + const ideTrust = getIdeTrust(); + if (ideTrust !== undefined) { + return ideTrust; + } + + // Fall back to the local user configuration + return getWorkspaceTrustFromLocalConfig(); +} diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index bfc62322bb..2b7169a0d2 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -242,6 +242,12 @@ vi.mock('./hooks/useFolderTrust', () => ({ })), })); +vi.mock('./hooks/useIdeTrustListener', () => ({ + useIdeTrustListener: vi.fn(() => ({ + needsRestart: false, + })), +})); + vi.mock('./hooks/useLogger', () => ({ useLogger: vi.fn(() => ({ getPreviousUserMessages: vi.fn().mockResolvedValue([]), diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 3b5bb225f6..e11aefc2df 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -27,6 +27,7 @@ import { useLoadingIndicator } from './hooks/useLoadingIndicator.js'; import { useThemeCommand } from './hooks/useThemeCommand.js'; import { useAuthCommand } from './hooks/useAuthCommand.js'; import { useFolderTrust } from './hooks/useFolderTrust.js'; +import { useIdeTrustListener } from './hooks/useIdeTrustListener.js'; import { useEditorSettings } from './hooks/useEditorSettings.js'; import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js'; import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js'; @@ -230,6 +231,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { IdeContext | undefined >(); const [showEscapePrompt, setShowEscapePrompt] = useState(false); + const [showIdeRestartPrompt, setShowIdeRestartPrompt] = useState(false); const [isProcessing, setIsProcessing] = useState(false); const { @@ -304,6 +306,23 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const { isFolderTrustDialogOpen, handleFolderTrustSelect, isRestarting } = useFolderTrust(settings, setIsTrustedFolder); + const { needsRestart: ideNeedsRestart } = useIdeTrustListener(config); + useEffect(() => { + if (ideNeedsRestart) { + // IDE trust changed, force a restart. + setShowIdeRestartPrompt(true); + } + }, [ideNeedsRestart]); + + useKeypress( + (key) => { + if (key.name === 'r' || key.name === 'R') { + process.exit(0); + } + }, + { isActive: showIdeRestartPrompt }, + ); + const { isAuthDialogOpen, openAuthDialog, @@ -1102,6 +1121,17 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { } }} /> + ) : showIdeRestartPrompt ? ( + + + Workspace trust has changed. Press 'r' to restart + Gemini to apply the changes. + + ) : isFolderTrustDialogOpen ? ( void) => { + const ideClient = config.getIdeClient(); + ideClient.addTrustChangeListener(onStoreChange); + return () => { + ideClient.removeTrustChangeListener(onStoreChange); + }; + }, + [config], + ); + + const getSnapshot = () => + ideContext.getIdeContext()?.workspaceState?.isTrusted; + + const isIdeTrusted = useSyncExternalStore(subscribe, getSnapshot); + + const [needsRestart, setNeedsRestart] = useState(false); + const [initialTrustValue] = useState(isIdeTrusted); + + useEffect(() => { + if ( + !needsRestart && + initialTrustValue !== undefined && + initialTrustValue !== isIdeTrusted + ) { + setNeedsRestart(true); + } + }, [isIdeTrusted, initialTrustValue, needsRestart]); + + return { isIdeTrusted, needsRestart }; +} diff --git a/packages/core/index.ts b/packages/core/index.ts index 447560d489..8a05dc5778 100644 --- a/packages/core/index.ts +++ b/packages/core/index.ts @@ -17,5 +17,6 @@ export { IdeConnectionEvent, IdeConnectionType, } from './src/telemetry/types.js'; +export { getIdeTrust } from './src/utils/ide-trust.js'; export { makeFakeConfig } from './src/test-utils/config.js'; export * from './src/utils/pathReader.js'; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d1ebd718d2..3dd012e4ae 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -43,6 +43,7 @@ import { import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import type { MCPOAuthConfig } from '../mcp/oauth-provider.js'; import { IdeClient } from '../ide/ide-client.js'; +import { ideContext } from '../ide/ideContext.js'; import type { Content } from '@google/genai'; import type { FileSystemService } from '../services/fileSystemService.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; @@ -749,6 +750,11 @@ export class Config { // restarts in the more common path. If the user chooses to mark the folder // as untrusted, the CLI will restart and we will have the trust value // reloaded. + const context = ideContext.getIdeContext(); + if (context?.workspaceState?.isTrusted !== undefined) { + return context.workspaceState.isTrusted; + } + return this.trustedFolder ?? true; } diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index 90e2ee3b32..c13f68e83a 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -77,6 +77,7 @@ export class IdeClient { private ideProcessInfo: { pid: number; command: string } | undefined; private diffResponses = new Map void>(); private statusListeners = new Set<(state: IDEConnectionState) => void>(); + private trustChangeListeners = new Set<(isTrusted: boolean) => void>(); private constructor() {} @@ -103,6 +104,14 @@ export class IdeClient { this.statusListeners.delete(listener); } + addTrustChangeListener(listener: (isTrusted: boolean) => void) { + this.trustChangeListeners.add(listener); + } + + removeTrustChangeListener(listener: (isTrusted: boolean) => void) { + this.trustChangeListeners.delete(listener); + } + async connect(): Promise { if (!this.currentIde || !this.currentIdeDisplayName) { this.setState( @@ -422,6 +431,12 @@ export class IdeClient { IdeContextNotificationSchema, (notification) => { ideContext.setIdeContext(notification.params); + const isTrusted = notification.params.workspaceState?.isTrusted; + if (isTrusted !== undefined) { + for (const listener of this.trustChangeListeners) { + listener(isTrusted); + } + } }, ); this.client.onerror = (_error) => { diff --git a/packages/core/src/ide/ideContext.ts b/packages/core/src/ide/ideContext.ts index 3052c0295d..9689c6323e 100644 --- a/packages/core/src/ide/ideContext.ts +++ b/packages/core/src/ide/ideContext.ts @@ -27,6 +27,7 @@ export const IdeContextSchema = z.object({ workspaceState: z .object({ openFiles: z.array(FileSchema).optional(), + isTrusted: z.boolean().optional(), }) .optional(), }); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index be437ec1f9..4b8d3aa3e8 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -47,6 +47,7 @@ export * from './utils/errorParsing.js'; export * from './utils/workspaceContext.js'; export * from './utils/ignorePatterns.js'; export * from './utils/partUtils.js'; +export * from './utils/ide-trust.js'; // Export services export * from './services/fileDiscoveryService.js'; diff --git a/packages/core/src/utils/ide-trust.ts b/packages/core/src/utils/ide-trust.ts new file mode 100644 index 0000000000..1cfaa88b47 --- /dev/null +++ b/packages/core/src/utils/ide-trust.ts @@ -0,0 +1,15 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ideContext } from '../ide/ideContext.js'; + +/** + * Gets the workspace trust from the IDE if available. + * @returns A boolean if the IDE provides a trust value, otherwise undefined. + */ +export function getIdeTrust(): boolean | undefined { + return ideContext.getIdeContext()?.workspaceState?.isTrusted; +} diff --git a/packages/vscode-ide-companion/src/extension.test.ts b/packages/vscode-ide-companion/src/extension.test.ts index 5a738c2a4f..2223050bcf 100644 --- a/packages/vscode-ide-companion/src/extension.test.ts +++ b/packages/vscode-ide-companion/src/extension.test.ts @@ -32,6 +32,7 @@ vi.mock('vscode', () => ({ onDidCloseTextDocument: vi.fn(), registerTextDocumentContentProvider: vi.fn(), onDidChangeWorkspaceFolders: vi.fn(), + onDidGrantWorkspaceTrust: vi.fn(), }, commands: { registerCommand: vi.fn(), @@ -91,6 +92,11 @@ describe('activate', () => { expect(vscode.window.showInformationMessage).not.toHaveBeenCalled(); }); + it('should register a handler for onDidGrantWorkspaceTrust', async () => { + await activate(context); + expect(vscode.workspace.onDidGrantWorkspaceTrust).toHaveBeenCalled(); + }); + it('should launch the Gemini CLI when the user clicks the button', async () => { const showInformationMessageMock = vi .mocked(vscode.window.showInformationMessage) diff --git a/packages/vscode-ide-companion/src/extension.ts b/packages/vscode-ide-companion/src/extension.ts index 4e03a2909e..36fc4029d1 100644 --- a/packages/vscode-ide-companion/src/extension.ts +++ b/packages/vscode-ide-companion/src/extension.ts @@ -72,7 +72,10 @@ export async function activate(context: vscode.ExtensionContext) { context.subscriptions.push( vscode.workspace.onDidChangeWorkspaceFolders(() => { - ideServer.updateWorkspacePath(); + ideServer.syncEnvVars(); + }), + vscode.workspace.onDidGrantWorkspaceTrust(() => { + ideServer.syncEnvVars(); }), vscode.commands.registerCommand('gemini-cli.runGeminiCLI', async () => { const workspaceFolders = vscode.workspace.workspaceFolders; diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index 836da59705..5dbbd3e936 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -45,6 +45,7 @@ const vscodeMock = vi.hoisted(() => ({ }, }, ], + isTrusted: true, }, })); @@ -229,7 +230,7 @@ describe('IDEServer', () => { { uri: { fsPath: '/foo/bar' } }, { uri: { fsPath: '/baz/qux' } }, ]; - await ideServer.updateWorkspacePath(); + await ideServer.syncEnvVars(); const expectedWorkspacePaths = ['/foo/bar', '/baz/qux'].join( path.delimiter, @@ -264,7 +265,7 @@ describe('IDEServer', () => { // Simulate removing a folder vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/baz/qux' } }]; - await ideServer.updateWorkspacePath(); + await ideServer.syncEnvVars(); expect(replaceMock).toHaveBeenCalledWith( 'GEMINI_CLI_IDE_WORKSPACE_PATH', diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 6b4c1689ae..a7a0116322 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -91,6 +91,9 @@ export class IDEServer { private portFile: string | undefined; private ppidPortFile: string | undefined; private port: number | undefined; + private transports: { [sessionId: string]: StreamableHTTPServerTransport } = + {}; + private openFilesManager: OpenFilesManager | undefined; diffManager: DiffManager; constructor(log: (message: string) => void, diffManager: DiffManager) { @@ -102,27 +105,19 @@ export class IDEServer { return new Promise((resolve) => { this.context = context; const sessionsWithInitialNotification = new Set(); - const transports: { [sessionId: string]: StreamableHTTPServerTransport } = - {}; const app = express(); app.use(express.json()); const mcpServer = createMcpServer(this.diffManager); - const openFilesManager = new OpenFilesManager(context); - const onDidChangeSubscription = openFilesManager.onDidChange(() => { - for (const transport of Object.values(transports)) { - sendIdeContextUpdateNotification( - transport, - this.log.bind(this), - openFilesManager, - ); - } + this.openFilesManager = new OpenFilesManager(context); + const onDidChangeSubscription = this.openFilesManager.onDidChange(() => { + this.broadcastIdeContextUpdate(); }); context.subscriptions.push(onDidChangeSubscription); const onDidChangeDiffSubscription = this.diffManager.onDidChange( (notification) => { - for (const transport of Object.values(transports)) { + for (const transport of Object.values(this.transports)) { transport.send(notification); } }, @@ -135,14 +130,14 @@ export class IDEServer { | undefined; let transport: StreamableHTTPServerTransport; - if (sessionId && transports[sessionId]) { - transport = transports[sessionId]; + if (sessionId && this.transports[sessionId]) { + transport = this.transports[sessionId]; } else if (!sessionId && isInitializeRequest(req.body)) { transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), onsessioninitialized: (newSessionId) => { this.log(`New session initialized: ${newSessionId}`); - transports[newSessionId] = transport; + this.transports[newSessionId] = transport; }, }); const keepAlive = setInterval(() => { @@ -161,7 +156,7 @@ export class IDEServer { if (transport.sessionId) { this.log(`Session closed: ${transport.sessionId}`); sessionsWithInitialNotification.delete(transport.sessionId); - delete transports[transport.sessionId]; + delete this.transports[transport.sessionId]; } }; mcpServer.connect(transport); @@ -204,13 +199,13 @@ export class IDEServer { const sessionId = req.headers[MCP_SESSION_ID_HEADER] as | string | undefined; - if (!sessionId || !transports[sessionId]) { + if (!sessionId || !this.transports[sessionId]) { this.log('Invalid or missing session ID'); res.status(400).send('Invalid or missing session ID'); return; } - const transport = transports[sessionId]; + const transport = this.transports[sessionId]; try { await transport.handleRequest(req, res); } catch (error) { @@ -222,11 +217,14 @@ export class IDEServer { } } - if (!sessionsWithInitialNotification.has(sessionId)) { + if ( + this.openFilesManager && + !sessionsWithInitialNotification.has(sessionId) + ) { sendIdeContextUpdateNotification( transport, this.log.bind(this), - openFilesManager, + this.openFilesManager, ); sessionsWithInitialNotification.add(sessionId); } @@ -260,7 +258,20 @@ export class IDEServer { }); } - async updateWorkspacePath(): Promise { + broadcastIdeContextUpdate() { + if (!this.openFilesManager) { + return; + } + for (const transport of Object.values(this.transports)) { + sendIdeContextUpdateNotification( + transport, + this.log.bind(this), + this.openFilesManager, + ); + } + } + + async syncEnvVars(): Promise { if ( this.context && this.server && @@ -275,6 +286,7 @@ export class IDEServer { this.ppidPortFile, this.log, ); + this.broadcastIdeContextUpdate(); } } diff --git a/packages/vscode-ide-companion/src/open-files-manager.ts b/packages/vscode-ide-companion/src/open-files-manager.ts index ffd1a568ba..57bbdd9e84 100644 --- a/packages/vscode-ide-companion/src/open-files-manager.ts +++ b/packages/vscode-ide-companion/src/open-files-manager.ts @@ -172,6 +172,7 @@ export class OpenFilesManager { return { workspaceState: { openFiles: [...this.openFiles], + isTrusted: vscode.workspace.isTrusted, }, }; }