From 26d274ffb4d65435e2c5401f6d811acde0619c24 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 3 Feb 2026 23:07:23 +0000 Subject: [PATCH] initial redirection changes --- .../messages/ToolConfirmationMessage.tsx | 21 ++++++- .../RedirectionConfirmation.test.tsx.snap | 5 +- .../cli/src/zed-integration/zedIntegration.ts | 1 + packages/core/src/confirmation-bus/types.ts | 8 ++- packages/core/src/policy/config.ts | 4 ++ packages/core/src/policy/policy-engine.ts | 9 +++ .../core/src/policy/policy-updater.test.ts | 11 ++++ .../src/policy/redirection-session.test.ts | 61 +++++++++++++++++++ .../src/tools/base-tool-invocation.test.ts | 21 ++++++- packages/core/src/tools/tools.ts | 8 +++ 10 files changed, 142 insertions(+), 7 deletions(-) create mode 100644 packages/core/src/policy/redirection-session.test.ts diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx index a527c13314..2fd0085776 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx @@ -63,6 +63,15 @@ export const ToolConfirmationMessage: React.FC< const allowPermanentApproval = settings.merged.security.enablePermanentToolApproval; + const containsRedirection = useMemo(() => { + if (confirmationDetails.type !== 'exec') return false; + const commandsToDisplay = + confirmationDetails.commands && confirmationDetails.commands.length > 1 + ? confirmationDetails.commands + : [confirmationDetails.command]; + return commandsToDisplay.some((cmd) => hasRedirection(cmd)); + }, [confirmationDetails]); + const handlesOwnUI = confirmationDetails.type === 'ask_user' || confirmationDetails.type === 'exit_plan_mode'; @@ -149,6 +158,13 @@ export const ToolConfirmationMessage: React.FC< key: 'Allow once', }); if (isTrustedFolder) { + if (containsRedirection) { + options.push({ + label: 'Allow redirection for this session', + value: ToolConfirmationOutcome.ProceedAlwaysRedirection, + key: 'Allow redirection for this session', + }); + } options.push({ label: `Allow for this session`, value: ToolConfirmationOutcome.ProceedAlways, @@ -231,6 +247,7 @@ export const ToolConfirmationMessage: React.FC< allowPermanentApproval, config, isDiffingEnabled, + containsRedirection, ]); const availableBodyContentHeight = useCallback(() => { @@ -344,9 +361,6 @@ export const ToolConfirmationMessage: React.FC< executionProps.commands && executionProps.commands.length > 1 ? executionProps.commands : [executionProps.command]; - const containsRedirection = commandsToDisplay.some((cmd) => - hasRedirection(cmd), - ); let bodyContentHeight = availableBodyContentHeight(); let warnings: React.ReactNode = null; @@ -461,6 +475,7 @@ export const ToolConfirmationMessage: React.FC< availableBodyContentHeight, terminalWidth, handleConfirm, + containsRedirection, ]); if (confirmationDetails.type === 'edit') { diff --git a/packages/cli/src/ui/components/messages/__snapshots__/RedirectionConfirmation.test.tsx.snap b/packages/cli/src/ui/components/messages/__snapshots__/RedirectionConfirmation.test.tsx.snap index 4f89811121..e9dccbc494 100644 --- a/packages/cli/src/ui/components/messages/__snapshots__/RedirectionConfirmation.test.tsx.snap +++ b/packages/cli/src/ui/components/messages/__snapshots__/RedirectionConfirmation.test.tsx.snap @@ -8,7 +8,8 @@ Tip: Toggle auto-edit (Shift+Tab) to allow redirection in the future. Allow execution of: 'echo, redirection (>)'? ● 1. Allow once - 2. Allow for this session - 3. No, suggest changes (esc) + 2. Allow redirection for this session + 3. Allow for this session + 4. No, suggest changes (esc) " `; diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index 634c20a1a0..6681ccd00a 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -683,6 +683,7 @@ export class Session { case ToolConfirmationOutcome.ProceedAlwaysAndSave: case ToolConfirmationOutcome.ProceedAlwaysServer: case ToolConfirmationOutcome.ProceedAlwaysTool: + case ToolConfirmationOutcome.ProceedAlwaysRedirection: case ToolConfirmationOutcome.ModifyWithEditor: break; default: { diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts index 8aa21f8ca1..61032929a5 100644 --- a/packages/core/src/confirmation-bus/types.ts +++ b/packages/core/src/confirmation-bus/types.ts @@ -21,6 +21,7 @@ export enum MessageBusType { TOOL_CALLS_UPDATE = 'tool-calls-update', ASK_USER_REQUEST = 'ask-user-request', ASK_USER_RESPONSE = 'ask-user-response', + ALLOW_SESSION_REDIRECTION = 'allow-session-redirection', } export interface ToolCallsUpdateMessage { @@ -171,6 +172,10 @@ export interface AskUserResponse { cancelled?: boolean; } +export interface AllowSessionRedirection { + type: MessageBusType.ALLOW_SESSION_REDIRECTION; +} + export type Message = | ToolConfirmationRequest | ToolConfirmationResponse @@ -180,4 +185,5 @@ export type Message = | UpdatePolicy | AskUserRequest | AskUserResponse - | ToolCallsUpdateMessage; + | ToolCallsUpdateMessage + | AllowSessionRedirection; diff --git a/packages/core/src/policy/config.ts b/packages/core/src/policy/config.ts index 7f6f4d9f3d..a1363234cc 100644 --- a/packages/core/src/policy/config.ts +++ b/packages/core/src/policy/config.ts @@ -329,6 +329,10 @@ export function createPolicyUpdater( policyEngine: PolicyEngine, messageBus: MessageBus, ) { + messageBus.subscribe(MessageBusType.ALLOW_SESSION_REDIRECTION, () => { + policyEngine.setAllowSessionRedirection(true); + }); + messageBus.subscribe( MessageBusType.UPDATE_POLICY, async (message: UpdatePolicy) => { diff --git a/packages/core/src/policy/policy-engine.ts b/packages/core/src/policy/policy-engine.ts index c0baf3e5c7..eb7ffddcd1 100644 --- a/packages/core/src/policy/policy-engine.ts +++ b/packages/core/src/policy/policy-engine.ts @@ -87,6 +87,7 @@ export class PolicyEngine { private readonly nonInteractive: boolean; private readonly checkerRunner?: CheckerRunner; private approvalMode: ApprovalMode; + private allowSessionRedirection = false; constructor(config: PolicyEngineConfig = {}, checkerRunner?: CheckerRunner) { this.rules = (config.rules ?? []).sort( @@ -118,11 +119,19 @@ export class PolicyEngine { return this.approvalMode; } + /** + * Set whether redirection is allowed for the current session. + */ + setAllowSessionRedirection(allowed: boolean): void { + this.allowSessionRedirection = allowed; + } + private shouldDowngradeForRedirection( command: string, allowRedirection?: boolean, ): boolean { return ( + !this.allowSessionRedirection && !allowRedirection && hasRedirection(command) && this.approvalMode !== ApprovalMode.AUTO_EDIT && diff --git a/packages/core/src/policy/policy-updater.test.ts b/packages/core/src/policy/policy-updater.test.ts index aa6b7ac887..7567396e05 100644 --- a/packages/core/src/policy/policy-updater.test.ts +++ b/packages/core/src/policy/policy-updater.test.ts @@ -130,6 +130,17 @@ describe('createPolicyUpdater', () => { expect(parsed.rule).toHaveLength(1); expect(parsed.rule![0].commandPrefix).toEqual(['echo', 'ls']); }); + + it('should call setAllowSessionRedirection on ALLOW_SESSION_REDIRECTION message', async () => { + vi.spyOn(policyEngine, 'setAllowSessionRedirection'); + createPolicyUpdater(policyEngine, messageBus); + + await messageBus.publish({ + type: MessageBusType.ALLOW_SESSION_REDIRECTION, + }); + + expect(policyEngine.setAllowSessionRedirection).toHaveBeenCalledWith(true); + }); }); describe('ShellToolInvocation Policy Update', () => { diff --git a/packages/core/src/policy/redirection-session.test.ts b/packages/core/src/policy/redirection-session.test.ts new file mode 100644 index 0000000000..020a672767 --- /dev/null +++ b/packages/core/src/policy/redirection-session.test.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach } from 'vitest'; +import { PolicyEngine } from './policy-engine.js'; +import { ApprovalMode, PolicyDecision } from './types.js'; +import { MessageBus } from '../confirmation-bus/message-bus.js'; +import { MessageBusType } from '../confirmation-bus/types.js'; +import { createPolicyUpdater } from './config.js'; + +describe('PolicyEngine Redirection Session', () => { + let policyEngine: PolicyEngine; + let messageBus: MessageBus; + + beforeEach(() => { + policyEngine = new PolicyEngine({ + defaultDecision: PolicyDecision.ALLOW, + approvalMode: ApprovalMode.DEFAULT, + }); + messageBus = new MessageBus(policyEngine); + createPolicyUpdater(policyEngine, messageBus); + }); + + it('should downgrade for redirection by default', async () => { + const result = await policyEngine.check( + { name: 'run_shell_command', args: { command: 'echo test > file.txt' } }, + undefined, + ); + expect(result.decision).toBe(PolicyDecision.ASK_USER); + }); + + it('should NOT downgrade for redirection after ALLOW_SESSION_REDIRECTION message', async () => { + // Publish the allow session redirection message + await messageBus.publish({ + type: MessageBusType.ALLOW_SESSION_REDIRECTION, + }); + + const result = await policyEngine.check( + { name: 'run_shell_command', args: { command: 'echo test > file.txt' } }, + undefined, + ); + expect(result.decision).toBe(PolicyDecision.ALLOW); + }); + + it('should still downgrade for redirection if another session starts (new PolicyEngine)', async () => { + // Simulate new session/engine + const newEngine = new PolicyEngine({ + defaultDecision: PolicyDecision.ALLOW, + approvalMode: ApprovalMode.DEFAULT, + }); + + const result = await newEngine.check( + { name: 'run_shell_command', args: { command: 'echo test > file.txt' } }, + undefined, + ); + expect(result.decision).toBe(PolicyDecision.ASK_USER); + }); +}); diff --git a/packages/core/src/tools/base-tool-invocation.test.ts b/packages/core/src/tools/base-tool-invocation.test.ts index f25bc8828f..1088daf981 100644 --- a/packages/core/src/tools/base-tool-invocation.test.ts +++ b/packages/core/src/tools/base-tool-invocation.test.ts @@ -5,7 +5,11 @@ */ import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { BaseToolInvocation, type ToolResult } from './tools.js'; +import { + BaseToolInvocation, + type ToolResult, + ToolConfirmationOutcome, +} from './tools.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { type Message, @@ -132,4 +136,19 @@ describe('BaseToolInvocation', () => { // ignore abort error } }); + + it('should publish ALLOW_SESSION_REDIRECTION message when outcome is ProceedAlwaysRedirection', async () => { + const tool = new TestBaseToolInvocation({}, messageBus, 'test-tool'); + + // Access protected method for testing + await ( + tool as unknown as { + publishPolicyUpdate(outcome: ToolConfirmationOutcome): Promise; + } + ).publishPolicyUpdate(ToolConfirmationOutcome.ProceedAlwaysRedirection); + + expect(messageBus.publish).toHaveBeenCalledWith({ + type: MessageBusType.ALLOW_SESSION_REDIRECTION, + }); + }); }); diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 65aeb0884f..30a0f9bb1b 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -141,6 +141,13 @@ export abstract class BaseToolInvocation< protected async publishPolicyUpdate( outcome: ToolConfirmationOutcome, ): Promise { + if (outcome === ToolConfirmationOutcome.ProceedAlwaysRedirection) { + void this.messageBus.publish({ + type: MessageBusType.ALLOW_SESSION_REDIRECTION, + }); + return; + } + if ( outcome === ToolConfirmationOutcome.ProceedAlways || outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave @@ -777,6 +784,7 @@ export enum ToolConfirmationOutcome { ProceedAlwaysAndSave = 'proceed_always_and_save', ProceedAlwaysServer = 'proceed_always_server', ProceedAlwaysTool = 'proceed_always_tool', + ProceedAlwaysRedirection = 'proceed_always_redirection', ModifyWithEditor = 'modify_with_editor', Cancel = 'cancel', }