mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-08-02 05:01:04 -07:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2b9ce09b6e | |||
| a24eccbf57 | |||
| ab4c6461db |
@@ -609,13 +609,15 @@ export class AppRig {
|
||||
this.removeToolPolicy(pending.toolName);
|
||||
}
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
||||
messageBus.publish({
|
||||
const p = messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: pending.correlationId,
|
||||
confirmed: outcome !== ToolConfirmationOutcome.Cancel,
|
||||
outcome,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
|
||||
@@ -90,11 +90,18 @@ export class LocalSubagentInvocation extends BaseToolInvocation<
|
||||
}
|
||||
|
||||
private publishActivity(activity: SubagentActivityItem): void {
|
||||
void this.messageBus.publish({
|
||||
type: MessageBusType.SUBAGENT_ACTIVITY,
|
||||
subagentName: this.definition.displayName ?? this.definition.name,
|
||||
activity,
|
||||
});
|
||||
try {
|
||||
const p = this.messageBus.publish({
|
||||
type: MessageBusType.SUBAGENT_ACTIVITY,
|
||||
subagentName: this.definition.displayName ?? this.definition.name,
|
||||
activity,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors in fire-and-forget activity update
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -8,8 +8,10 @@ import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { MessageBus } from './message-bus.js';
|
||||
import { PolicyEngine } from '../policy/policy-engine.js';
|
||||
import { PolicyDecision } from '../policy/types.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import {
|
||||
MessageBusType,
|
||||
type Message,
|
||||
type ToolConfirmationRequest,
|
||||
type ToolConfirmationResponse,
|
||||
type ToolPolicyRejection,
|
||||
@@ -30,8 +32,9 @@ describe('MessageBus', () => {
|
||||
const errorHandler = vi.fn();
|
||||
messageBus.on('error', errorHandler);
|
||||
|
||||
// @ts-expect-error - Testing invalid message
|
||||
await messageBus.publish({ invalid: 'message' });
|
||||
await expect(
|
||||
messageBus.publish({ invalid: 'message' } as unknown as Message),
|
||||
).rejects.toThrow('Invalid message structure');
|
||||
|
||||
expect(errorHandler).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
@@ -44,11 +47,12 @@ describe('MessageBus', () => {
|
||||
const errorHandler = vi.fn();
|
||||
messageBus.on('error', errorHandler);
|
||||
|
||||
// @ts-expect-error - Testing missing correlationId
|
||||
await messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: { name: 'test' },
|
||||
});
|
||||
await expect(
|
||||
messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: { name: 'test' },
|
||||
} as unknown as Message),
|
||||
).rejects.toThrow('Invalid message structure');
|
||||
|
||||
expect(errorHandler).toHaveBeenCalled();
|
||||
});
|
||||
@@ -164,6 +168,62 @@ describe('MessageBus', () => {
|
||||
);
|
||||
});
|
||||
|
||||
it('should sanitize sensitive data in error messages', async () => {
|
||||
const errorHandler = vi.fn();
|
||||
messageBus.on('error', errorHandler);
|
||||
|
||||
const invalidMessage = {
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: {
|
||||
name: 'test-tool',
|
||||
args: { password: 'secret-password' },
|
||||
},
|
||||
// missing correlationId makes it invalid
|
||||
} as unknown as Message;
|
||||
|
||||
await expect(messageBus.publish(invalidMessage)).rejects.toThrow(
|
||||
/\[REDACTED\]/,
|
||||
);
|
||||
await expect(messageBus.publish(invalidMessage)).rejects.not.toThrow(
|
||||
/secret-password/,
|
||||
);
|
||||
|
||||
expect(errorHandler).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
message: expect.stringContaining('[REDACTED]'),
|
||||
}),
|
||||
);
|
||||
expect(errorHandler).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
message: expect.not.stringContaining('secret-password'),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should sanitize sensitive data in debug logs', async () => {
|
||||
const debugSpy = vi.spyOn(debugLogger, 'debug');
|
||||
// @ts-expect-error - access private member for test
|
||||
messageBus.debug = true;
|
||||
|
||||
const message: ToolExecutionSuccess<string> = {
|
||||
type: MessageBusType.TOOL_EXECUTION_SUCCESS as const,
|
||||
toolCall: {
|
||||
name: 'test-tool',
|
||||
args: { api_key: 'my-api-key' },
|
||||
},
|
||||
result: 'success',
|
||||
};
|
||||
|
||||
await messageBus.publish(message);
|
||||
|
||||
expect(debugSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('[REDACTED]'),
|
||||
);
|
||||
expect(debugSpy).toHaveBeenCalledWith(
|
||||
expect.not.stringContaining('my-api-key'),
|
||||
);
|
||||
});
|
||||
|
||||
it('should emit other message types directly', async () => {
|
||||
const successHandler = vi.fn();
|
||||
messageBus.subscribe(
|
||||
@@ -251,8 +311,10 @@ describe('MessageBus', () => {
|
||||
correlationId: '123',
|
||||
};
|
||||
|
||||
// Should not throw
|
||||
await expect(messageBus.publish(request)).resolves.not.toThrow();
|
||||
// Should throw
|
||||
await expect(messageBus.publish(request)).rejects.toThrow(
|
||||
'Policy check failed',
|
||||
);
|
||||
|
||||
// Should emit error
|
||||
expect(errorHandler).toHaveBeenCalledWith(
|
||||
@@ -263,6 +325,101 @@ describe('MessageBus', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('request', () => {
|
||||
it('should fail fast if publish fails', async () => {
|
||||
// Mock publish to throw
|
||||
vi.spyOn(messageBus, 'publish').mockRejectedValue(
|
||||
new Error('Publish failed'),
|
||||
);
|
||||
|
||||
const request: Omit<ToolConfirmationRequest, 'correlationId'> = {
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: { name: 'test-tool', args: {} },
|
||||
};
|
||||
|
||||
const start = Date.now();
|
||||
const requestPromise = messageBus.request(
|
||||
request,
|
||||
MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
60000,
|
||||
);
|
||||
|
||||
await expect(requestPromise).rejects.toThrow('Publish failed');
|
||||
const duration = Date.now() - start;
|
||||
|
||||
// Should have failed way before the 60s timeout
|
||||
expect(duration).toBeLessThan(1000);
|
||||
});
|
||||
|
||||
it('should handle timeout if no response is received', async () => {
|
||||
const request: Message = {
|
||||
type: MessageBusType.SUBAGENT_ACTIVITY,
|
||||
subagentName: 'test',
|
||||
activity: {
|
||||
id: '1',
|
||||
type: 'thought',
|
||||
status: 'running',
|
||||
content: 'thinking',
|
||||
},
|
||||
};
|
||||
|
||||
const requestPromise = messageBus.request(
|
||||
request,
|
||||
MessageBusType.ASK_USER_RESPONSE,
|
||||
10, // 10ms timeout
|
||||
);
|
||||
|
||||
await expect(requestPromise).rejects.toThrow(
|
||||
/Request timed out waiting for ask-user-response/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('pattern resiliency', () => {
|
||||
it('should handle publish returning non-promise (mock behavior)', () => {
|
||||
// @ts-expect-error - mock returning undefined
|
||||
vi.spyOn(messageBus, 'publish').mockReturnValue(undefined);
|
||||
|
||||
const publishAndCatch = () => {
|
||||
const p = messageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [],
|
||||
schedulerId: 'test',
|
||||
} as Message);
|
||||
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
};
|
||||
|
||||
expect(publishAndCatch).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle publish throwing synchronously (mock behavior)', () => {
|
||||
vi.spyOn(messageBus, 'publish').mockImplementation(() => {
|
||||
throw new Error('Sync throw');
|
||||
});
|
||||
|
||||
const publishAndCatch = () => {
|
||||
try {
|
||||
const p = messageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [],
|
||||
schedulerId: 'test',
|
||||
} as Message);
|
||||
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
} catch {
|
||||
// handled
|
||||
}
|
||||
};
|
||||
|
||||
expect(publishAndCatch).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('derive', () => {
|
||||
it('should receive responses from parent bus on derived bus', async () => {
|
||||
vi.spyOn(policyEngine, 'check').mockResolvedValue({
|
||||
@@ -288,11 +445,14 @@ describe('MessageBus', () => {
|
||||
MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
(msg) => {
|
||||
if (msg.subagent === subagentName) {
|
||||
void messageBus.publish({
|
||||
const p = messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: msg.correlationId,
|
||||
confirmed: true,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
resolve();
|
||||
}
|
||||
},
|
||||
@@ -330,11 +490,14 @@ describe('MessageBus', () => {
|
||||
MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
(msg) => {
|
||||
if (msg.subagent === 'agent1/agent2') {
|
||||
void messageBus.publish({
|
||||
const p = messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: msg.correlationId,
|
||||
confirmed: true,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
resolve();
|
||||
}
|
||||
},
|
||||
|
||||
@@ -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 { sanitizeToolArgs } from '../utils/agent-sanitization-utils.js';
|
||||
|
||||
export class MessageBus extends EventEmitter {
|
||||
private listenerToAbortCleanup = new WeakMap<
|
||||
@@ -78,12 +79,16 @@ export class MessageBus extends EventEmitter {
|
||||
|
||||
async publish(message: Message): Promise<void> {
|
||||
if (this.debug) {
|
||||
debugLogger.debug(`[MESSAGE_BUS] publish: ${safeJsonStringify(message)}`);
|
||||
debugLogger.debug(
|
||||
`[MESSAGE_BUS] publish: ${safeJsonStringify(sanitizeToolArgs(message))}`,
|
||||
);
|
||||
}
|
||||
try {
|
||||
if (!this.isValidMessage(message)) {
|
||||
throw new Error(
|
||||
`Invalid message structure: ${safeJsonStringify(message)}`,
|
||||
`Invalid message structure: ${safeJsonStringify(
|
||||
sanitizeToolArgs(message),
|
||||
)}`,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -144,6 +149,7 @@ export class MessageBus extends EventEmitter {
|
||||
}
|
||||
} catch (error) {
|
||||
this.emit('error', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -239,8 +245,11 @@ export class MessageBus extends EventEmitter {
|
||||
this.subscribe<TResponse>(responseType, responseHandler);
|
||||
|
||||
// Publish the request with correlation ID
|
||||
// eslint-disable-next-line @typescript-eslint/no-floating-promises, @typescript-eslint/no-unsafe-type-assertion
|
||||
this.publish({ ...request, correlationId } as TRequest);
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
this.publish({ ...request, correlationId } as TRequest).catch((error) => {
|
||||
cleanup();
|
||||
reject(error);
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import { checkPolicy, updatePolicy, getPolicyDenialError } from './policy.js';
|
||||
import { evaluateBeforeToolHook } from './hook-utils.js';
|
||||
import { ToolExecutor } from './tool-executor.js';
|
||||
import { ToolModificationHandler } from './tool-modifier.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import {
|
||||
type ToolCallRequestInfo,
|
||||
type ToolCall,
|
||||
@@ -166,12 +167,16 @@ export class Scheduler {
|
||||
private readonly handleToolConfirmationRequest = async (
|
||||
request: ToolConfirmationRequest,
|
||||
) => {
|
||||
await this.messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: request.correlationId,
|
||||
confirmed: false,
|
||||
requiresUserConfirmation: true,
|
||||
});
|
||||
try {
|
||||
await this.messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: request.correlationId,
|
||||
confirmed: false,
|
||||
requiresUserConfirmation: true,
|
||||
});
|
||||
} catch (error) {
|
||||
debugLogger.error('Failed to publish confirmation response', error);
|
||||
}
|
||||
};
|
||||
|
||||
private setupMessageBusListener(messageBus: MessageBus): void {
|
||||
|
||||
@@ -244,11 +244,18 @@ export class SchedulerStateManager {
|
||||
const snapshot = this.getSnapshot();
|
||||
|
||||
// Fire and forget - The message bus handles the publish and error handling.
|
||||
void this.messageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: snapshot,
|
||||
schedulerId: this.schedulerId,
|
||||
});
|
||||
try {
|
||||
const p = this.messageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: snapshot,
|
||||
schedulerId: this.schedulerId,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors in fire-and-forget update
|
||||
}
|
||||
}
|
||||
|
||||
private isTerminalCall(call: ToolCall): call is CompletedToolCall {
|
||||
|
||||
@@ -242,12 +242,19 @@ export abstract class BaseToolInvocation<
|
||||
) {
|
||||
if (this._toolName) {
|
||||
const options = this.getPolicyUpdateOptions(outcome);
|
||||
void this.messageBus.publish({
|
||||
type: MessageBusType.UPDATE_POLICY,
|
||||
toolName: this._toolName,
|
||||
persist: outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave,
|
||||
...options,
|
||||
});
|
||||
try {
|
||||
const p = this.messageBus.publish({
|
||||
type: MessageBusType.UPDATE_POLICY,
|
||||
toolName: this._toolName,
|
||||
persist: outcome === ToolConfirmationOutcome.ProceedAlwaysAndSave,
|
||||
...options,
|
||||
});
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {});
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors in fire-and-forget update
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -362,7 +369,13 @@ export abstract class BaseToolInvocation<
|
||||
};
|
||||
|
||||
try {
|
||||
void this.messageBus.publish(request);
|
||||
const p = this.messageBus.publish(request);
|
||||
if (p instanceof Promise) {
|
||||
p.catch(() => {
|
||||
cleanup();
|
||||
resolve('allow');
|
||||
});
|
||||
}
|
||||
} catch {
|
||||
cleanup();
|
||||
resolve('allow');
|
||||
|
||||
Reference in New Issue
Block a user