From e162f622e41068a2bbc004edee15ca2a38bc2a6b Mon Sep 17 00:00:00 2001 From: mkorwel Date: Sat, 14 Mar 2026 12:14:50 -0700 Subject: [PATCH] fix(core): resolve scheduler hang and improve policy violation visibility This PR addresses three core issues with the Policy Engine and Scheduler: 1. Scheduler Hang: Removed a redundant MessageBus listener in Scheduler.ts that caused race conditions in TTY environments. 2. Policy Visibility: Updated ToolGroupMessage.tsx to always display policy violation errors, regardless of verbosity settings. 3. User Feedback: Added emitFeedback to MessageBus.ts to ensure blocked tool calls are reported to the UI. --- .../components/messages/ToolGroupMessage.tsx | 4 ++- packages/cli/src/ui/hooks/toolMapping.ts | 4 +++ packages/cli/src/ui/types.ts | 2 ++ .../core/src/confirmation-bus/message-bus.ts | 5 +++ packages/core/src/scheduler/scheduler.ts | 33 +------------------ 5 files changed, 15 insertions(+), 33 deletions(-) diff --git a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx index e22d3c6313..b71226b77d 100644 --- a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx @@ -21,6 +21,7 @@ import { isShellTool } from './ToolShared.js'; import { shouldHideToolCall, CoreToolCallStatus, + ToolErrorType, } from '@google/gemini-cli-core'; import { useUIState } from '../../contexts/UIStateContext.js'; import { getToolGroupBorderAppearance } from '../../utils/borderStyles.js'; @@ -59,7 +60,8 @@ export const ToolGroupMessage: React.FC = ({ if ( isLowErrorVerbosity && t.status === CoreToolCallStatus.Error && - !t.isClientInitiated + !t.isClientInitiated && + t.errorType !== ToolErrorType.POLICY_VIOLATION ) { return false; } diff --git a/packages/cli/src/ui/hooks/toolMapping.ts b/packages/cli/src/ui/hooks/toolMapping.ts index e06ebf5bb5..bd4c6a95df 100644 --- a/packages/cli/src/ui/hooks/toolMapping.ts +++ b/packages/cli/src/ui/hooks/toolMapping.ts @@ -10,6 +10,7 @@ import { type ToolResultDisplay, debugLogger, CoreToolCallStatus, + type ToolErrorType, } from '@google/gemini-cli-core'; import { type HistoryItemToolGroup, @@ -63,6 +64,7 @@ export function mapToDisplay( let progressMessage: string | undefined = undefined; let progress: number | undefined = undefined; let progressTotal: number | undefined = undefined; + let errorType: ToolErrorType | undefined = undefined; switch (call.status) { case CoreToolCallStatus.Success: @@ -72,6 +74,7 @@ export function mapToDisplay( case CoreToolCallStatus.Error: case CoreToolCallStatus.Cancelled: resultDisplay = call.response.resultDisplay; + errorType = call.response.errorType; break; case CoreToolCallStatus.AwaitingApproval: correlationId = call.correlationId; @@ -114,6 +117,7 @@ export function mapToDisplay( progressTotal, approvalMode: call.approvalMode, originalRequestName: call.request.originalRequestName, + errorType, }; }); diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 2f8e414a83..7025a6655c 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -16,6 +16,7 @@ import { type AgentDefinition, type ApprovalMode, type Kind, + type ToolErrorType, CoreToolCallStatus, checkExhaustive, } from '@google/gemini-cli-core'; @@ -117,6 +118,7 @@ export interface IndividualToolCallDisplay { originalRequestName?: string; progress?: number; progressTotal?: number; + errorType?: ToolErrorType; } export interface CompressionProps { diff --git a/packages/core/src/confirmation-bus/message-bus.ts b/packages/core/src/confirmation-bus/message-bus.ts index 33aa10355b..fe4485f09a 100644 --- a/packages/core/src/confirmation-bus/message-bus.ts +++ b/packages/core/src/confirmation-bus/message-bus.ts @@ -11,6 +11,7 @@ import { PolicyDecision } from '../policy/types.js'; import { MessageBusType, type Message } from './types.js'; import { safeJsonStringify } from '../utils/safeJsonStringify.js'; import { debugLogger } from '../utils/debugLogger.js'; +import { coreEvents } from '../utils/events.js'; export class MessageBus extends EventEmitter { constructor( @@ -70,6 +71,10 @@ export class MessageBus extends EventEmitter { break; case PolicyDecision.DENY: // Emit both rejection and response messages + coreEvents.emitFeedback( + 'error', + `Tool call "${message.toolCall.name}" was blocked by policy.`, + ); this.emitMessage({ type: MessageBusType.TOOL_POLICY_REJECTION, toolCall: message.toolCall, diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 4a92617e6d..62abb48fa3 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -35,11 +35,7 @@ import { runInDevTraceSpan } from '../telemetry/trace.js'; import { logToolCall } from '../telemetry/loggers.js'; import { ToolCallEvent } from '../telemetry/types.js'; import type { EditorType } from '../utils/editor.js'; -import { - MessageBusType, - type SerializableConfirmationDetails, - type ToolConfirmationRequest, -} from '../confirmation-bus/types.js'; +import { type SerializableConfirmationDetails } from '../confirmation-bus/types.js'; import { runWithToolCallContext } from '../utils/toolCallContext.js'; import { coreEvents, @@ -91,9 +87,6 @@ const createErrorResponse = ( * Coordinates execution via state updates and event listening. */ export class Scheduler { - // Tracks which MessageBus instances have the legacy listener attached to prevent duplicates. - private static subscribedMessageBuses = new WeakSet(); - private readonly state: SchedulerStateManager; private readonly executor: ToolExecutor; private readonly modifier: ToolModificationHandler; @@ -127,8 +120,6 @@ export class Scheduler { this.executor = new ToolExecutor(this.context); this.modifier = new ToolModificationHandler(); - this.setupMessageBusListener(this.messageBus); - coreEvents.on(CoreEvent.McpProgress, this.handleMcpProgress); } @@ -161,28 +152,6 @@ export class Scheduler { }); }; - private setupMessageBusListener(messageBus: MessageBus): void { - if (Scheduler.subscribedMessageBuses.has(messageBus)) { - return; - } - - // TODO: Optimize policy checks. Currently, tools check policy via - // MessageBus even though the Scheduler already checked it. - messageBus.subscribe( - MessageBusType.TOOL_CONFIRMATION_REQUEST, - async (request: ToolConfirmationRequest) => { - await messageBus.publish({ - type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, - correlationId: request.correlationId, - confirmed: false, - requiresUserConfirmation: true, - }); - }, - ); - - Scheduler.subscribedMessageBuses.add(messageBus); - } - /** * Schedules a batch of tool calls. * @returns A promise that resolves with the results of the completed batch.