Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions packages/cli/src/test-utils/AppRig.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down
17 changes: 12 additions & 5 deletions packages/core/src/agents/local-invocation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

/**
Expand Down
185 changes: 174 additions & 11 deletions packages/core/src/confirmation-bus/message-bus.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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({
Expand All @@ -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();
});
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand All @@ -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({
Expand All @@ -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();
}
},
Expand Down Expand Up @@ -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();
}
},
Expand Down
17 changes: 13 additions & 4 deletions packages/core/src/confirmation-bus/message-bus.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<
Expand Down Expand Up @@ -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),
)}`,
);
}

Expand Down Expand Up @@ -144,6 +149,7 @@ export class MessageBus extends EventEmitter {
}
} catch (error) {
this.emit('error', error);
throw error;
Comment thread
Adib234 marked this conversation as resolved.
}
}

Expand Down Expand Up @@ -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);
});
});
}
}
17 changes: 11 additions & 6 deletions packages/core/src/scheduler/scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down
17 changes: 12 additions & 5 deletions packages/core/src/scheduler/state-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading
Loading