Skip to content

Commit 7c74c49

Browse files
CopilotMQ37
Copilot
and
MQ37
authored
Centralize skyfire pay ID validation in tool dispatcher (#395)
* Initial plan * Centralize skyfire validation in MCP server dispatcher Co-authored-by: MQ37 <29043708+MQ37@users.noreply.github.com> * fix: enable task mode support for internal tools with taskSupport Previously, task mode execution (executeToolAndUpdateTask) rejected all internal tools, only allowing actor tools. This was inconsistent with tools like call-actor that declare taskSupport: 'optional'. Changes: - Remove hard rejection of internal tools in task mode - Add execution path for internal tools in task mode - Pass userRentedActorIds to task execution context - Add integration test for call-actor in task mode This fixes the architectural issue identified in PR #396 review where centralized skyfire validation worked but internal tools couldn't reach that code path due to early rejection. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: MQ37 <29043708+MQ37@users.noreply.github.com> Co-authored-by: MQ37 <themq37@gmail.com>
1 parent 3ea2ca1 commit 7c74c49

9 files changed

Lines changed: 130 additions & 80 deletions

File tree

‎package-lock.json‎

Lines changed: 0 additions & 14 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/mcp/server.ts‎

Lines changed: 55 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -818,6 +818,7 @@ Please remove the "task" parameter from the tool call request or use a different
818818
progressToken,
819819
extra,
820820
mcpSessionId,
821+
userRentedActorIds,
821822
});
822823
});
823824

@@ -831,6 +832,15 @@ Please remove the "task" parameter from the tool call request or use a different
831832
let toolStatus: ToolStatus = TOOL_STATUS.SUCCEEDED;
832833

833834
try {
835+
// Centralized skyfire validation for tools that require it
836+
if (tool.requiresSkyfirePayId) {
837+
const skyfireError = validateSkyfirePayId(this, args);
838+
if (skyfireError) {
839+
toolStatus = TOOL_STATUS.SOFT_FAIL;
840+
return skyfireError;
841+
}
842+
}
843+
834844
// Handle internal tool
835845
if (tool.type === 'internal') {
836846
// Only create progress tracker for call-actor tool
@@ -917,9 +927,6 @@ Please verify the server URL is correct and accessible, and ensure you have a va
917927

918928
// Handle actor tool
919929
if (tool.type === 'actor') {
920-
const skyfireError = validateSkyfirePayId(this, args);
921-
if (skyfireError) return skyfireError;
922-
923930
// Create progress tracker if progressToken is available
924931
const progressTracker = createProgressTracker(progressToken, extra.sendNotification);
925932

@@ -1035,8 +1042,9 @@ Please verify the tool name and ensure the tool is properly registered.`;
10351042
progressToken: string | number | undefined;
10361043
extra: RequestHandlerExtra<Request, Notification>;
10371044
mcpSessionId: string | undefined;
1045+
userRentedActorIds?: string[];
10381046
}): Promise<void> {
1039-
const { taskId, tool, args, apifyToken, progressToken, extra, mcpSessionId } = params;
1047+
const { taskId, tool, args, apifyToken, progressToken, extra, mcpSessionId, userRentedActorIds } = params;
10401048
let toolStatus: ToolStatus = TOOL_STATUS.SUCCEEDED;
10411049
const startTime = Date.now();
10421050

@@ -1051,24 +1059,6 @@ Please verify the tool name and ensure the tool is properly registered.`;
10511059
const { telemetryData, userId } = await this.prepareTelemetryData(tool, mcpSessionId, apifyToken);
10521060

10531061
try {
1054-
// Only support actor tool types for long-running tasks
1055-
if (tool.type !== 'actor') {
1056-
const msg = `Tool "${tool.name}" does not support long-running task execution. Only actor tools are supported in task mode.`;
1057-
log.softFail(msg, { statusCode: 400 });
1058-
await this.server.sendLoggingMessage({ level: 'error', data: msg });
1059-
toolStatus = TOOL_STATUS.SOFT_FAIL;
1060-
await this.taskStore.storeTaskResult(taskId, 'failed', {
1061-
content: [{
1062-
type: 'text' as const,
1063-
text: msg,
1064-
}],
1065-
isError: true,
1066-
internalToolStatus: toolStatus,
1067-
}, mcpSessionId);
1068-
this.finalizeAndTrackTelemetry(telemetryData, userId, startTime, toolStatus);
1069-
return;
1070-
}
1071-
10721062
// Check if task was already cancelled before we start execution.
10731063
// Critical: if a client cancels the task immediately after creation (race condition),
10741064
// attempting to transition from 'cancelled' (terminal state) to 'working' will fail in the SDK
@@ -1091,12 +1081,49 @@ Please verify the tool name and ensure the tool is properly registered.`;
10911081
// Execute the tool and get the result
10921082
let result: Record<string, unknown> = {};
10931083

1094-
// Handle actor tool
1095-
const skyfireError = validateSkyfirePayId(this, args);
1096-
if (skyfireError) {
1097-
result = skyfireError;
1098-
toolStatus = TOOL_STATUS.SOFT_FAIL;
1099-
} else {
1084+
// Centralized skyfire validation for tools that require it
1085+
if (tool.requiresSkyfirePayId) {
1086+
const skyfireError = validateSkyfirePayId(this, args);
1087+
if (skyfireError) {
1088+
result = skyfireError;
1089+
toolStatus = TOOL_STATUS.SOFT_FAIL;
1090+
}
1091+
}
1092+
1093+
// Handle internal tool execution in task mode
1094+
if (toolStatus === TOOL_STATUS.SUCCEEDED && tool.type === 'internal') {
1095+
const progressTracker = createProgressTracker(progressToken, extra.sendNotification, taskId);
1096+
1097+
log.info('Calling internal tool for task', { taskId, name: tool.name, input: args });
1098+
const res = await tool.call({
1099+
args,
1100+
extra,
1101+
apifyMcpServer: this,
1102+
mcpServer: this.server,
1103+
apifyToken,
1104+
userRentedActorIds,
1105+
progressTracker,
1106+
}) as object;
1107+
1108+
if (progressTracker) {
1109+
progressTracker.stop();
1110+
}
1111+
1112+
// If tool provided internal status, use it; otherwise infer from isError flag
1113+
const { internalToolStatus, ...rest } = res as { internalToolStatus?: ToolStatus; isError?: boolean };
1114+
if (internalToolStatus !== undefined) {
1115+
toolStatus = internalToolStatus;
1116+
} else if ('isError' in rest && rest.isError) {
1117+
toolStatus = TOOL_STATUS.FAILED;
1118+
} else {
1119+
toolStatus = TOOL_STATUS.SUCCEEDED;
1120+
}
1121+
1122+
result = rest;
1123+
}
1124+
1125+
// Handle actor tool execution in task mode
1126+
if (toolStatus === TOOL_STATUS.SUCCEEDED && tool.type === 'actor') {
11001127
const progressTracker = createProgressTracker(progressToken, extra.sendNotification, taskId);
11011128
const callOptions: ActorCallOptions = { memory: tool.memoryMbytes };
11021129
const { 'skyfire-pay-id': _skyfirePayId, ...actorArgs } = args as Record<string, unknown>;

‎src/tools/actor.ts‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import { buildMCPResponse } from '../utils/mcp.js';
2727
import type { ProgressTracker } from '../utils/progress.js';
2828
import type { JsonSchemaProperty } from '../utils/schema-generation.js';
2929
import { generateSchemaFromItems } from '../utils/schema-generation.js';
30-
import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
30+
import { createApifyClientWithSkyfireSupport } from '../utils/skyfire.js';
3131
import { getWidgetConfig, WIDGET_URIS } from '../utils/widgets.js';
3232
import { getActorDefinition } from './build.js';
3333
import { actorNameToToolName, buildActorInputSchema, fixedAjvCompile, isActorInfoMcpServer } from './utils.js';
@@ -198,6 +198,7 @@ Actor description: ${definition.description}`;
198198
description,
199199
inputSchema: inputSchema as ToolInputSchema,
200200
ajvValidate,
201+
requiresSkyfirePayId: true,
201202
memoryMbytes,
202203
icons: definition.pictureUrl
203204
? [{ src: definition.pictureUrl, mimeType: 'image/png' }]
@@ -398,6 +399,7 @@ EXAMPLES:
398399
// Additional props true to allow skyfire-pay-id
399400
additionalProperties: true,
400401
}),
402+
requiresSkyfirePayId: true,
401403
_meta: {
402404
...getWidgetConfig(WIDGET_URIS.ACTOR_RUN)?.meta,
403405
},
@@ -442,9 +444,6 @@ EXAMPLES:
442444
}
443445

444446
try {
445-
const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
446-
if (skyfireError) return skyfireError;
447-
448447
const apifyClient = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
449448

450449
// Determine execution mode: always async when UI mode is enabled, otherwise respect the parameter

‎src/tools/dataset.ts‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { compileSchema } from '../utils/ajv.js';
66
import { parseCommaSeparatedList } from '../utils/generic.js';
77
import { buildMCPResponse } from '../utils/mcp.js';
88
import { generateSchemaFromItems } from '../utils/schema-generation.js';
9-
import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
9+
import { createApifyClientWithSkyfireSupport } from '../utils/skyfire.js';
1010

1111
const getDatasetArgs = z.object({
1212
datasetId: z.string()
@@ -60,6 +60,7 @@ USAGE EXAMPLES:
6060
* Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
6161
*/
6262
ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetArgs), additionalProperties: true }),
63+
requiresSkyfirePayId: true,
6364
annotations: {
6465
title: 'Get dataset',
6566
readOnlyHint: true,
@@ -70,9 +71,6 @@ USAGE EXAMPLES:
7071
const { args, apifyToken, apifyMcpServer } = toolArgs;
7172
const parsed = getDatasetArgs.parse(args);
7273

73-
const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
74-
if (skyfireError) return skyfireError;
75-
7674
const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
7775
const v = await client.dataset(parsed.datasetId).get();
7876
if (!v) {
@@ -109,6 +107,7 @@ USAGE EXAMPLES:
109107
* Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
110108
*/
111109
ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetItemsArgs), additionalProperties: true }),
110+
requiresSkyfirePayId: true,
112111
annotations: {
113112
title: 'Get dataset items',
114113
readOnlyHint: true,
@@ -119,9 +118,6 @@ USAGE EXAMPLES:
119118
const { args, apifyToken, apifyMcpServer } = toolArgs;
120119
const parsed = getDatasetItemsArgs.parse(args);
121120

122-
const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
123-
if (skyfireError) return skyfireError;
124-
125121
const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
126122

127123
// Convert comma-separated strings to arrays
@@ -185,6 +181,7 @@ USAGE EXAMPLES:
185181
* Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
186182
*/
187183
ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetSchemaArgs), additionalProperties: true }),
184+
requiresSkyfirePayId: true,
188185
annotations: {
189186
title: 'Get dataset schema',
190187
readOnlyHint: true,
@@ -195,9 +192,6 @@ USAGE EXAMPLES:
195192
const { args, apifyToken, apifyMcpServer } = toolArgs;
196193
const parsed = getDatasetSchemaArgs.parse(args);
197194

198-
const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
199-
if (skyfireError) return skyfireError;
200-
201195
const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
202196

203197
// Get dataset items

‎src/tools/get-actor-output.ts‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import type { InternalToolArgs, ToolEntry, ToolInputSchema } from '../types.js';
55
import { compileSchema } from '../utils/ajv.js';
66
import { getValuesByDotKeys, parseCommaSeparatedList } from '../utils/generic.js';
77
import { buildMCPResponse } from '../utils/mcp.js';
8-
import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
8+
import { createApifyClientWithSkyfireSupport } from '../utils/skyfire.js';
99

1010
/**
1111
* Zod schema for get-actor-output tool arguments
@@ -88,6 +88,7 @@ Note: This tool is automatically included if the Apify MCP Server is configured
8888
* Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
8989
*/
9090
ajvValidate: compileSchema({ ...z.toJSONSchema(getActorOutputArgs), additionalProperties: true }),
91+
requiresSkyfirePayId: true,
9192
annotations: {
9293
title: 'Get Actor output',
9394
readOnlyHint: true,
@@ -97,9 +98,6 @@ Note: This tool is automatically included if the Apify MCP Server is configured
9798
call: async (toolArgs: InternalToolArgs) => {
9899
const { args, apifyToken, apifyMcpServer } = toolArgs;
99100

100-
const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
101-
if (skyfireError) return skyfireError;
102-
103101
const apifyClient = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
104102
const parsed = getActorOutputArgs.parse(args);
105103

0 commit comments

Comments
 (0)