Skip to content

refactor: Dedup the tool-loading mode-resolution gate #1078

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jirispilka wants to merge 1 commit into
base: master
Choose a base branch
from
+105 −33
Draft

refactor: Dedup the tool-loading mode-resolution gate #1078

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
56 changes: 25 additions & 31 deletions src/mcp/server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -441,6 +441,27 @@ export class ActorsMcpServer {
return [...this.listInternalToolNames(), ...this.listActorToolNames(), ...this.listActorMcpServerToolIds()];
}

/**
* Buffer-or-compose gate shared by the actor-tools loaders: if the server mode isn't
* resolved yet, queue `actorTools` for `updateToolsAfterServerModeResolved` and
* (if non-empty) upsert them immediately with the given `notify`; otherwise compose
* the mode-specific set via `getToolsForServerMode` and upsert that with `notify`.
*
* Callers intentionally pass different `notify` values: `loadToolsByName` notifies
* eagerly (actor-tools-only visible pre-init), while `loadToolsFromUrl` and
* `loadToolsFromInput` pass `false` and defer to the post-initialize reconcile.
* See `updateToolsAfterServerModeResolved` for the full rationale.
*/
private registerFetchedActorTools(input: Input, actorTools: ToolEntry[], notify: boolean): void {
if (!this.serverModeResolved) {
this.pendingToolsAfterModeResolved.push({ input, actorTools });
if (actorTools.length > 0) this.upsertTools(actorTools, notify);
return;
}
const tools = getToolsForServerMode(input, actorTools, this.serverMode);
if (tools.length > 0) this.upsertTools(tools, notify);
}

/**
* Loads missing toolNames from a provided list of tool names.
* Skips toolNames that are already loaded and loads only the missing ones.
Expand All @@ -458,16 +479,7 @@ export class ActorsMcpServer {
paymentProvider: this.options.paymentProvider,
});

if (!this.serverModeResolved) {
this.pendingToolsAfterModeResolved.push({ input: restoreInput, actorTools });
if (actorTools.length > 0) this.upsertTools(actorTools, true);
return;
}

const toolsToLoad = getToolsForServerMode(restoreInput, actorTools, this.serverMode);
if (toolsToLoad.length > 0) {
this.upsertTools(toolsToLoad, actorTools.length > 0);
}
this.registerFetchedActorTools(restoreInput, actorTools, actorTools.length > 0);
}

/**
Expand Down Expand Up @@ -495,20 +507,8 @@ export class ActorsMcpServer {
paymentProvider: this.options.paymentProvider,
});

if (!this.serverModeResolved) {
this.pendingToolsAfterModeResolved.push({ input, actorTools });
if (actorTools.length > 0) {
log.debug('Loading actor tools from query parameters before mode resolution');
this.upsertTools(actorTools, false);
}
return;
}

const tools = getToolsForServerMode(input, actorTools, this.serverMode);
if (tools.length > 0) {
log.debug('Loading tools from query parameters');
this.upsertTools(tools, false);
}
log.debug('Loading tools from query parameters');
this.registerFetchedActorTools(input, actorTools, false);
}

/**
Expand All @@ -525,13 +525,7 @@ export class ActorsMcpServer {
actorStore: this.actorStore,
paymentProvider: this.options.paymentProvider,
});
if (!this.serverModeResolved) {
this.pendingToolsAfterModeResolved.push({ input, actorTools });
if (actorTools.length > 0) this.upsertTools(actorTools);
return;
}
const tools = getToolsForServerMode(input, actorTools, this.serverMode);
if (tools.length > 0) this.upsertTools(tools);
this.registerFetchedActorTools(input, actorTools, false);
}

/** Delete tools from the server and notify the handler.
Expand Down
82 changes: 80 additions & 2 deletions tests/unit/mcp.server.capability_gating.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,26 @@ import { RESOURCE_MIME_TYPE } from '../../src/resources/widgets.js';
import { callActorApps } from '../../src/tools/actors/call_actor.js';
import { searchActors } from '../../src/tools/actors/search_actors.js';
import { searchActorsWidget } from '../../src/tools/widgets/search_actors_widget.js';
import type { ServerModeOption } from '../../src/types.js';
import { SERVER_MODE } from '../../src/types.js';
import type { ServerModeOption, ToolEntry } from '../../src/types.js';
import { SERVER_MODE, TOOL_TYPE } from '../../src/types.js';
import type * as ToolsLoaderModule from '../../src/utils/tools_loader.js';
import { getActors } from '../../src/utils/tools_loader.js';

// Stub getActors so tests can produce a non-empty actorTools array without a network
// fetch to the Apify platform. getToolsForServerMode / toolNamesToInput stay real.
vi.mock('../../src/utils/tools_loader.js', async (importOriginal) => {
const actual = await importOriginal<typeof ToolsLoaderModule>();
return { ...actual, getActors: vi.fn() };
});

const getActorsMock = vi.mocked(getActors);

const FAKE_ACTOR_FULL_NAME = 'apify/fake-actor';
const FAKE_ACTOR_TOOL: ToolEntry = {
type: TOOL_TYPE.ACTOR,
name: 'apify--fake-actor',
actorFullName: FAKE_ACTOR_FULL_NAME,
} as unknown as ToolEntry;

type InitHandler = (req: InitializeRequest, ctx: unknown) => Promise<unknown>;

Expand Down Expand Up @@ -53,12 +71,18 @@ async function dispatchInitialize(server: ActorsMcpServer, request: InitializeRe
describe('ActorsMcpServer initialize handler', () => {
const servers: ActorsMcpServer[] = [];

// Default: no actor tools resolved, matching the real getActors() behavior for
// inputs that only reference HELPER_TOOLS names (used by the pre-existing tests below).
getActorsMock.mockResolvedValue([]);

afterEach(async () => {
while (servers.length > 0) {
const server = servers.pop();
server?.tools.clear();
await server?.close();
}
getActorsMock.mockReset();
getActorsMock.mockResolvedValue([]);
});

const track = (server: ActorsMcpServer): ActorsMcpServer => {
Expand Down Expand Up @@ -179,4 +203,58 @@ describe('ActorsMcpServer initialize handler', () => {
expect(server.tools.get(HELPER_TOOLS.STORE_SEARCH_WIDGET)).toBe(searchActorsWidget);
},
);

describe('registerFetchedActorTools notify behavior', () => {
it('loadToolsByName notifies eagerly pre-mode-resolution when actor tools are non-empty', async () => {
const server = track(makeServer('auto'));
const apifyClient = new ApifyClient({ token: 'test-token' });
const notifyHandler = vi.fn();
server.registerToolsChangedHandler(notifyHandler);
getActorsMock.mockResolvedValueOnce([FAKE_ACTOR_TOOL]);

await server.loadToolsByName(['apify/fake-actor'], apifyClient);

expect(notifyHandler).toHaveBeenCalledTimes(1);
});

it('loadToolsFromInput does not notify pre-mode-resolution even when actor tools are non-empty', async () => {
const server = track(makeServer('auto'));
const apifyClient = new ApifyClient({ token: 'test-token' });
const notifyHandler = vi.fn();
server.registerToolsChangedHandler(notifyHandler);
getActorsMock.mockResolvedValueOnce([FAKE_ACTOR_TOOL]);

await server.loadToolsFromInput({ actors: ['apify/fake-actor'] }, apifyClient);

expect(notifyHandler).not.toHaveBeenCalled();
});

it('loadToolsFromUrl does not notify pre-mode-resolution even when actor tools are non-empty', async () => {
const server = track(makeServer('auto'));
const apifyClient = new ApifyClient({ token: 'test-token' });
const notifyHandler = vi.fn();
server.registerToolsChangedHandler(notifyHandler);
getActorsMock.mockResolvedValueOnce([FAKE_ACTOR_TOOL]);

await server.loadToolsFromUrl('http://localhost?actors=apify/fake-actor', apifyClient);

expect(notifyHandler).not.toHaveBeenCalled();
});

it('notifies with the full set once initialize runs the pending-tools reconcile', async () => {
const server = track(makeServer('auto'));
const apifyClient = new ApifyClient({ token: 'test-token' });
getActorsMock.mockResolvedValueOnce([FAKE_ACTOR_TOOL]);

await server.loadToolsFromInput({ actors: ['apify/fake-actor'] }, apifyClient);

const notifyHandler = vi.fn();
server.registerToolsChangedHandler(notifyHandler);

await dispatchInitialize(server, makeInitializeRequest(true));

expect(notifyHandler).toHaveBeenCalledTimes(1);
expect(notifyHandler.mock.calls[0][0]).toContain(FAKE_ACTOR_FULL_NAME);
});
});
});
Loading