Skip to content
Merged
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
49 changes: 44 additions & 5 deletions packages/core/src/agents/agent-scheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,23 +19,24 @@ vi.mock('../scheduler/scheduler.js', () => ({
}));

describe('agent-scheduler', () => {
let mockConfig: Mocked<Config>;
let mockToolRegistry: Mocked<ToolRegistry>;
let mockMessageBus: Mocked<MessageBus>;

beforeEach(() => {
vi.mocked(Scheduler).mockClear();
mockMessageBus = {} as Mocked<MessageBus>;
mockToolRegistry = {
getTool: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
} as unknown as Mocked<ToolRegistry>;
mockConfig = {
});

it('should create a scheduler with agent-specific config', async () => {
const mockConfig = {
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
toolRegistry: mockToolRegistry,
} as unknown as Mocked<Config>;
});

it('should create a scheduler with agent-specific config', async () => {
const requests: ToolCallRequestInfo[] = [
{
callId: 'call-1',
Expand Down Expand Up @@ -68,8 +69,46 @@ describe('agent-scheduler', () => {
}),
);

// Verify that the scheduler's config has the overridden tool registry
const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
expect(schedulerConfig.toolRegistry).toBe(mockToolRegistry);
});

it('should override toolRegistry getter from prototype chain', async () => {
const mainRegistry = { _id: 'main' } as unknown as Mocked<ToolRegistry>;
const agentRegistry = {
_id: 'agent',
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
} as unknown as Mocked<ToolRegistry>;

const config = {
getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
} as unknown as Mocked<Config>;
Object.defineProperty(config, 'toolRegistry', {
get: () => mainRegistry,
configurable: true,
});

await scheduleAgentTools(
config as unknown as Config,
[
{
callId: 'c1',
name: 'new_page',
args: {},
isClientInitiated: false,
prompt_id: 'p1',
},
],
{
schedulerId: 'browser-1',
toolRegistry: agentRegistry as unknown as ToolRegistry,
signal: new AbortController().signal,
},
);

const schedulerConfig = vi.mocked(Scheduler).mock.calls[0][0].config;
expect(schedulerConfig.toolRegistry).toBe(agentRegistry);
expect(schedulerConfig.toolRegistry).not.toBe(mainRegistry);
expect(schedulerConfig.getToolRegistry()).toBe(agentRegistry);
});
});
5 changes: 5 additions & 0 deletions packages/core/src/agents/agent-scheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,11 @@ export async function scheduleAgentTools(
const agentConfig: Config = Object.create(config);
agentConfig.getToolRegistry = () => toolRegistry;
agentConfig.getMessageBus = () => toolRegistry.getMessageBus();
// Override toolRegistry property so AgentLoopContext reads the agent-specific registry.
Object.defineProperty(agentConfig, 'toolRegistry', {
get: () => toolRegistry,
configurable: true,
});
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious, if we're doing this for the AgentLoopContext interface, do we also need to override the messageBus in a similar way?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We will need to replace this Object.create hack with an explicit AgentLoopContext to avoid fragile prototype getter overrides altogether.


const scheduler = new Scheduler({
config: agentConfig,
Expand Down
39 changes: 39 additions & 0 deletions packages/core/src/agents/browser/browserAgentFactory.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,45 @@ describe('browserAgentFactory', () => {
.map((t) => t.name) ?? [];
expect(toolNames).toContain('analyze_screenshot');
});

it('should include all MCP navigation tools (new_page, navigate_page) in definition', async () => {
mockBrowserManager.getDiscoveredTools.mockResolvedValue([
{ name: 'take_snapshot', description: 'Take snapshot' },
{ name: 'click', description: 'Click element' },
{ name: 'fill', description: 'Fill form field' },
{ name: 'navigate_page', description: 'Navigate to URL' },
{ name: 'new_page', description: 'Open a new page/tab' },
{ name: 'close_page', description: 'Close page' },
{ name: 'select_page', description: 'Select page' },
{ name: 'press_key', description: 'Press key' },
{ name: 'hover', description: 'Hover element' },
]);

const { definition } = await createBrowserAgentDefinition(
mockConfig,
mockMessageBus,
);

const toolNames =
definition.toolConfig?.tools
?.filter(
(t): t is { name: string } => typeof t === 'object' && 'name' in t,
)
.map((t) => t.name) ?? [];

// All MCP tools must be present
expect(toolNames).toContain('new_page');
expect(toolNames).toContain('navigate_page');
expect(toolNames).toContain('close_page');
expect(toolNames).toContain('select_page');
expect(toolNames).toContain('click');
expect(toolNames).toContain('take_snapshot');
expect(toolNames).toContain('press_key');
// Custom composite tool must also be present
expect(toolNames).toContain('type_text');
// Total: 9 MCP + 1 type_text (no analyze_screenshot without visualModel)
expect(definition.toolConfig?.tools).toHaveLength(10);
});
});

describe('cleanupBrowserAgent', () => {
Expand Down
Loading