fix: skyfire parameter injection and missing tool registration

Fixed critical bug where 8 out of 11 skyfire-enabled tools were missing
the required `skyfire-pay-id` parameter in their input schemas. Refactored
Skyfire integration to eliminate ~180 lines of duplicated code.

Critical fixes:
- Schema injection now uses SKYFIRE_ENABLED_TOOLS Set instead of hardcoded list
- Added missing abortActorRun tool to runs category

Changes:
- Created src/utils/skyfire.ts with reusable helper functions
- Refactored 13 code locations to use helpers
- Added skyfire mode detection via ?payment=skyfire query parameter
- Added integration test for schema injection
- Updated test helpers and documentation

Tools affected: call-actor, get-actor-output, get-actor-run, get-actor-log,
abort-actor-run, get-dataset, get-dataset-items, get-dataset-schema,
get-key-value-store, get-key-value-store-keys, get-key-value-store-record

Files: 13 files changed, 286 insertions(+), 92 deletions(-)

Author: MQ37

AGENTS.md

@@ -114,6 +114,8 @@ After completing ANY code change (feature, fix, refactor), you MUST:
 
 - `tests/unit/` - Unit tests for individual modules
 - `tests/integration/` - Integration tests for MCP server functionality
+  - `tests/integration/suite.ts` - **Main integration test suite** where all test cases should be added
+  - Other files in this directory set up different transport modes (stdio, SSE, streamable-http) that all use suite.ts
 - `tests/helpers.ts` - Shared test utilities
 - `tests/const.ts` - Test constants
 
@@ -124,6 +126,33 @@ After completing ANY code change (feature, fix, refactor), you MUST:
 - Follow existing test patterns in the codebase
 - Ensure all tests pass before submitting a PR
 
+### Adding integration tests
+
+**IMPORTANT**: When adding integration test cases, add them to `tests/integration/suite.ts`, NOT as separate test files.
+
+The `suite.ts` file contains a test suite factory function `createIntegrationTestsSuite()` that is used by all transport modes (stdio, SSE, streamable-http). Adding tests here ensures they run across all transport types.
+
+**How to add a test case:**
+1. Open `tests/integration/suite.ts`
+2. Add your test case inside the `describe` block (before the closing braces at the end)
+3. Use `it()` or `it.runIf()` for conditional tests
+4. Follow the existing patterns for client creation and assertions
+5. Use `client = await createClientFn(options)` to create the test client
+6. Always call `await client.close()` when done
+
+**Example:**
+```typescript
+it('should do something awesome', async () => {
+    client = await createClientFn({ tools: ['actors'] });
+    const result = await client.callTool({
+        name: HelperTools.SOME_TOOL,
+        arguments: { /* ... */ },
+    });
+    expect(result.content).toBeDefined();
+    await client.close();
+});
+```
+
 ### Manual testing as an MCP client
 
 To test the MCP server, a human must first configure the MCP server.

src/actor/server.ts

@@ -93,6 +93,10 @@ export function createExpressApp(
             const uiModeParam = urlParams.get('ui') as UiMode | undefined;
             const uiMode = uiModeParam ?? process.env.UI_MODE as UiMode | undefined;
 
+            // Extract payment mode parameter - if payment=skyfire, enable skyfire mode
+            const paymentParam = urlParams.get('payment');
+            const skyfireMode = paymentParam === 'skyfire';
+
             const mcpServer = new ActorsMcpServer({
                 taskStore,
                 setupSigintHandler: false,
@@ -101,6 +105,7 @@ export function createExpressApp(
                     enabled: telemetryEnabled,
                 },
                 uiMode,
+                skyfireMode,
             });
             const transport = new SSEServerTransport(Routes.MESSAGE, res);
 
@@ -221,6 +226,10 @@ export function createExpressApp(
                 const uiModeParam = urlParams.get('ui') as UiMode | undefined;
                 const uiMode = uiModeParam ?? process.env.UI_MODE as UiMode | undefined;
 
+                // Extract payment mode parameter - if payment=skyfire, enable skyfire mode
+                const paymentParam = urlParams.get('payment');
+                const skyfireMode = paymentParam === 'skyfire';
+
                 const mcpServer = new ActorsMcpServer({
                     taskStore,
                     setupSigintHandler: false,
@@ -230,6 +239,7 @@ export function createExpressApp(
                         enabled: telemetryEnabled,
                     },
                     uiMode,
+                    skyfireMode,
                 });
 
                 // Load MCP server tools

src/const.ts

@@ -73,6 +73,24 @@ export const SKYFIRE_PAY_ID_PROPERTY_DESCRIPTION = `Skyfire PAY token ID used fo
 
 export const SKYFIRE_README_CONTENT = `The Apify MCP Server allows clients to interact with Apify Actors, storage, and documentation. Before running any Apify Actor, you need to set up a Skyfire PAY token for billing purposes. To set up the PAY token you need to use the create-pay-token tool from the Skyfire MCP server. The PAY token must be charged with a minimum of ${SKYFIRE_MIN_CHARGE_USD} USD to be valid.${SKYFIRE_SELLER_ID ? ` The Skyfire seller ID for Apify MCP server is "${SKYFIRE_SELLER_ID}".` : ''} When running an Actor using the call-actor tool, you must provide the Skyfire PAY token ID in the \`skyfire-pay-id\` input property. Similarly, when retrieving Actor output using the get-actor-output tool, you must also provide the same Skyfire PAY token ID in the \`skyfire-pay-id\` input property.`;
 
+/**
+ * Set of internal tool names that require Skyfire PAY token ID in Skyfire mode.
+ * These tools interact with Actor runs, datasets, or key-value stores and need billing support.
+ */
+export const SKYFIRE_ENABLED_TOOLS = new Set([
+    HelperTools.ACTOR_CALL,
+    HelperTools.ACTOR_OUTPUT_GET,
+    HelperTools.ACTOR_RUNS_GET,
+    HelperTools.ACTOR_RUNS_LOG,
+    HelperTools.ACTOR_RUNS_ABORT,
+    HelperTools.DATASET_GET,
+    HelperTools.DATASET_GET_ITEMS,
+    HelperTools.DATASET_SCHEMA_GET,
+    HelperTools.KEY_VALUE_STORE_GET,
+    HelperTools.KEY_VALUE_STORE_KEYS_GET,
+    HelperTools.KEY_VALUE_STORE_RECORD_GET,
+]);
+
 export const CALL_ACTOR_MCP_MISSING_TOOL_NAME_MSG = `When calling an MCP server Actor, you must specify the tool name in the actor parameter as "{actorName}:{toolName}" in the "actor" input property.`;
 
 // Cache

src/mcp/server.ts

@@ -44,6 +44,7 @@ import {
     SERVER_INSTRUCTIONS,
     SERVER_NAME,
     SERVER_VERSION,
+    SKYFIRE_ENABLED_TOOLS,
     SKYFIRE_PAY_ID_PROPERTY_DESCRIPTION,
     SKYFIRE_README_CONTENT,
     SKYFIRE_TOOL_INSTRUCTIONS,
@@ -69,6 +70,7 @@ import { parseBooleanFromString } from '../utils/generic.js';
 import { logHttpError } from '../utils/logging.js';
 import { buildMCPResponse } from '../utils/mcp.js';
 import { createProgressTracker } from '../utils/progress.js';
+import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
 import { getToolStatusFromError } from '../utils/tool-status.js';
 import { cloneToolEntry, getToolPublicFieldOnly } from '../utils/tools.js';
 import { getUserIdFromTokenCached } from '../utils/userid-cache.js';
@@ -346,8 +348,7 @@ export class ActorsMcpServer {
 
                 // Handle Skyfire mode modifications
                 if (this.options.skyfireMode && (wrap.type === 'actor'
-                    || (wrap.type === 'internal' && wrap.name === HelperTools.ACTOR_CALL)
-                    || (wrap.type === 'internal' && wrap.name === HelperTools.ACTOR_OUTPUT_GET))) {
+                    || (wrap.type === 'internal' && SKYFIRE_ENABLED_TOOLS.has(wrap.name as HelperTools)))) {
                     // Add Skyfire instructions to description if not already present
                     if (clonedWrap.description && !clonedWrap.description.includes(SKYFIRE_TOOL_INSTRUCTIONS)) {
                         clonedWrap.description += `\n\n${SKYFIRE_TOOL_INSTRUCTIONS}`;
@@ -916,22 +917,16 @@ Please verify the server URL is correct and accessible, and ensure you have a va
 
                 // Handle actor tool
                 if (tool.type === 'actor') {
-                    if (this.options.skyfireMode && args['skyfire-pay-id'] === undefined) {
-                        return buildMCPResponse({ texts: [SKYFIRE_TOOL_INSTRUCTIONS] });
-                    }
+                    const skyfireError = validateSkyfirePayId(this, args);
+                    if (skyfireError) return skyfireError;
 
                     // Create progress tracker if progressToken is available
                     const progressTracker = createProgressTracker(progressToken, extra.sendNotification);
 
                     const callOptions: ActorCallOptions = { memory: tool.memoryMbytes };
 
-                    /**
-                     * Create Apify token, for Skyfire mode use `skyfire-pay-id` and for normal mode use `apifyToken`.
-                     */
-                    const { 'skyfire-pay-id': skyfirePayId, ...actorArgs } = args as Record<string, unknown>;
-                    const apifyClient = this.options.skyfireMode && typeof skyfirePayId === 'string'
-                        ? new ApifyClient({ skyfirePayId })
-                        : new ApifyClient({ token: apifyToken });
+                    const { 'skyfire-pay-id': _skyfirePayId, ...actorArgs } = args as Record<string, unknown>;
+                    const apifyClient = createApifyClientWithSkyfireSupport(this, args, apifyToken);
 
                     try {
                         log.info('Calling Actor', { actorName: tool.actorFullName, input: actorArgs });
@@ -1097,16 +1092,15 @@ Please verify the tool name and ensure the tool is properly registered.`;
             let result: Record<string, unknown> = {};
 
             // Handle actor tool
-            if (this.options.skyfireMode && args['skyfire-pay-id'] === undefined) {
-                result = buildMCPResponse({ texts: [SKYFIRE_TOOL_INSTRUCTIONS], isError: true });
+            const skyfireError = validateSkyfirePayId(this, args);
+            if (skyfireError) {
+                result = skyfireError;
                 toolStatus = TOOL_STATUS.SOFT_FAIL;
             } else {
                 const progressTracker = createProgressTracker(progressToken, extra.sendNotification, taskId);
                 const callOptions: ActorCallOptions = { memory: tool.memoryMbytes };
-                const { 'skyfire-pay-id': skyfirePayId, ...actorArgs } = args as Record<string, unknown>;
-                const apifyClient = this.options.skyfireMode && typeof skyfirePayId === 'string'
-                    ? new ApifyClient({ skyfirePayId })
-                    : new ApifyClient({ token: apifyToken });
+                const { 'skyfire-pay-id': _skyfirePayId, ...actorArgs } = args as Record<string, unknown>;
+                const apifyClient = createApifyClientWithSkyfireSupport(this, args, apifyToken);
 
                 log.info('Calling Actor for task', { taskId, actorName: tool.actorFullName, input: actorArgs });
                 const callResult = await callActorGetDataset(

src/tools/actor.ts

@@ -11,7 +11,6 @@ import {
     HelperTools,
     RAG_WEB_BROWSER,
     RAG_WEB_BROWSER_ADDITIONAL_DESC,
-    SKYFIRE_TOOL_INSTRUCTIONS,
     TOOL_MAX_OUTPUT_CHARS,
     TOOL_STATUS,
 } from '../const.js';
@@ -28,6 +27,7 @@ import { buildMCPResponse } from '../utils/mcp.js';
 import type { ProgressTracker } from '../utils/progress.js';
 import type { JsonSchemaProperty } from '../utils/schema-generation.js';
 import { generateSchemaFromItems } from '../utils/schema-generation.js';
+import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
 import { getWidgetConfig, WIDGET_URIS } from '../utils/widgets.js';
 import { getActorDefinition } from './build.js';
 import { actorNameToToolName, buildActorInputSchema, fixedAjvCompile, isActorInfoMcpServer } from './utils.js';
@@ -442,25 +442,10 @@ EXAMPLES:
         }
 
         try {
-            /**
-            * In Skyfire mode, we check for the presence of `skyfire-pay-id`.
-            * If it is missing, we return instructions to the LLM on how to create it and pass it to the tool.
-            */
-            if (apifyMcpServer.options.skyfireMode && args['skyfire-pay-id'] === undefined) {
-                return {
-                    content: [{
-                        type: 'text',
-                        text: SKYFIRE_TOOL_INSTRUCTIONS,
-                    }],
-                };
-            }
+            const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
+            if (skyfireError) return skyfireError;
 
-            /**
-            * Create Apify token, for Skyfire mode use `skyfire-pay-id` and for normal mode use `apifyToken`.
-            */
-            const apifyClient = apifyMcpServer.options.skyfireMode && typeof args['skyfire-pay-id'] === 'string'
-                ? new ApifyClient({ skyfirePayId: args['skyfire-pay-id'] })
-                : new ApifyClient({ token: apifyToken });
+            const apifyClient = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
 
             // Determine execution mode: always async when UI mode is enabled, otherwise respect the parameter
             const isAsync = apifyMcpServer.options.uiMode === 'openai'

src/tools/categories.ts

@@ -13,7 +13,7 @@ import { getHtmlSkeleton } from './get-html-skeleton.js';
 import { addTool } from './helpers.js';
 import { getKeyValueStore, getKeyValueStoreKeys, getKeyValueStoreRecord } from './key_value_store.js';
 import { getUserKeyValueStoresList } from './key_value_store_collection.js';
-import { getActorRun, getActorRunLog } from './run.js';
+import { abortActorRun, getActorRun, getActorRunLog } from './run.js';
 import { getUserRunsList } from './run_collection.js';
 import { searchApifyDocsTool } from './search-apify-docs.js';
 import { searchActors } from './store_collection.js';
@@ -35,6 +35,7 @@ export const toolCategories = {
         getActorRun,
         getUserRunsList,
         getActorRunLog,
+        abortActorRun,
     ],
     storage: [
         getDataset,

src/tools/dataset.ts

@@ -1,12 +1,12 @@
 import { z } from 'zod';
 
-import { ApifyClient } from '../apify-client.js';
 import { HelperTools, TOOL_STATUS } from '../const.js';
 import type { InternalToolArgs, ToolEntry, ToolInputSchema } from '../types.js';
 import { compileSchema } from '../utils/ajv.js';
 import { parseCommaSeparatedList } from '../utils/generic.js';
 import { buildMCPResponse } from '../utils/mcp.js';
 import { generateSchemaFromItems } from '../utils/schema-generation.js';
+import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
 
 const getDatasetArgs = z.object({
     datasetId: z.string()
@@ -56,17 +56,24 @@ USAGE EXAMPLES:
 - user_input: Show info for dataset xyz123
 - user_input: What fields does username~my-dataset have?`,
     inputSchema: z.toJSONSchema(getDatasetArgs) as ToolInputSchema,
-    ajvValidate: compileSchema(z.toJSONSchema(getDatasetArgs)),
+    /**
+     * Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
+     */
+    ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetArgs), additionalProperties: true }),
     annotations: {
         title: 'Get dataset',
         readOnlyHint: true,
         idempotentHint: true,
         openWorldHint: false,
     },
     call: async (toolArgs: InternalToolArgs) => {
-        const { args, apifyToken } = toolArgs;
+        const { args, apifyToken, apifyMcpServer } = toolArgs;
         const parsed = getDatasetArgs.parse(args);
-        const client = new ApifyClient({ token: apifyToken });
+
+        const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
+        if (skyfireError) return skyfireError;
+
+        const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
         const v = await client.dataset(parsed.datasetId).get();
         if (!v) {
             return buildMCPResponse({
@@ -98,17 +105,24 @@ USAGE EXAMPLES:
 - user_input: Get first 100 items from dataset abd123
 - user_input: Get only metadata.url and title from dataset username~my-dataset (flatten metadata)`,
     inputSchema: z.toJSONSchema(getDatasetItemsArgs) as ToolInputSchema,
-    ajvValidate: compileSchema(z.toJSONSchema(getDatasetItemsArgs)),
+    /**
+     * Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
+     */
+    ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetItemsArgs), additionalProperties: true }),
     annotations: {
         title: 'Get dataset items',
         readOnlyHint: true,
         idempotentHint: true,
         openWorldHint: false,
     },
     call: async (toolArgs: InternalToolArgs) => {
-        const { args, apifyToken } = toolArgs;
+        const { args, apifyToken, apifyMcpServer } = toolArgs;
         const parsed = getDatasetItemsArgs.parse(args);
-        const client = new ApifyClient({ token: apifyToken });
+
+        const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
+        if (skyfireError) return skyfireError;
+
+        const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
 
         // Convert comma-separated strings to arrays
         const fields = parseCommaSeparatedList(parsed.fields);
@@ -167,17 +181,24 @@ USAGE EXAMPLES:
 - user_input: Generate schema for dataset 34das2 using 10 items
 - user_input: Show schema of username~my-dataset (clean items only)`,
     inputSchema: z.toJSONSchema(getDatasetSchemaArgs) as ToolInputSchema,
-    ajvValidate: compileSchema(z.toJSONSchema(getDatasetSchemaArgs)),
+    /**
+     * Allow additional properties for Skyfire mode to pass `skyfire-pay-id`.
+     */
+    ajvValidate: compileSchema({ ...z.toJSONSchema(getDatasetSchemaArgs), additionalProperties: true }),
     annotations: {
         title: 'Get dataset schema',
         readOnlyHint: true,
         idempotentHint: true,
         openWorldHint: false,
     },
     call: async (toolArgs: InternalToolArgs) => {
-        const { args, apifyToken } = toolArgs;
+        const { args, apifyToken, apifyMcpServer } = toolArgs;
         const parsed = getDatasetSchemaArgs.parse(args);
-        const client = new ApifyClient({ token: apifyToken });
+
+        const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
+        if (skyfireError) return skyfireError;
+
+        const client = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
 
         // Get dataset items
         const datasetResponse = await client.dataset(parsed.datasetId).listItems({

src/tools/get-actor-output.ts

@@ -1,11 +1,11 @@
 import { z } from 'zod';
 
-import { ApifyClient } from '../apify-client.js';
-import { HelperTools, SKYFIRE_TOOL_INSTRUCTIONS, TOOL_MAX_OUTPUT_CHARS, TOOL_STATUS } from '../const.js';
+import { HelperTools, TOOL_MAX_OUTPUT_CHARS, TOOL_STATUS } from '../const.js';
 import type { InternalToolArgs, ToolEntry, ToolInputSchema } from '../types.js';
 import { compileSchema } from '../utils/ajv.js';
 import { getValuesByDotKeys, parseCommaSeparatedList } from '../utils/generic.js';
 import { buildMCPResponse } from '../utils/mcp.js';
+import { createApifyClientWithSkyfireSupport, validateSkyfirePayId } from '../utils/skyfire.js';
 
 /**
  * Zod schema for get-actor-output tool arguments
@@ -97,27 +97,10 @@ Note: This tool is automatically included if the Apify MCP Server is configured
     call: async (toolArgs: InternalToolArgs) => {
         const { args, apifyToken, apifyMcpServer } = toolArgs;
 
-        /**
-             * In Skyfire mode, we check for the presence of `skyfire-pay-id`.
-             * If it is missing, we return instructions to the LLM on how to create it and pass it to the tool.
-             */
-        if (apifyMcpServer.options.skyfireMode
-                && args['skyfire-pay-id'] === undefined
-        ) {
-            return {
-                content: [{
-                    type: 'text',
-                    text: SKYFIRE_TOOL_INSTRUCTIONS,
-                }],
-            };
-        }
+        const skyfireError = validateSkyfirePayId(apifyMcpServer, args);
+        if (skyfireError) return skyfireError;
 
-        /**
-             * Create Apify token, for Skyfire mode use `skyfire-pay-id` and for normal mode use `apifyToken`.
-             */
-        const apifyClient = apifyMcpServer.options.skyfireMode && typeof args['skyfire-pay-id'] === 'string'
-            ? new ApifyClient({ skyfirePayId: args['skyfire-pay-id'] })
-            : new ApifyClient({ token: apifyToken });
+        const apifyClient = createApifyClientWithSkyfireSupport(apifyMcpServer, args, apifyToken);
         const parsed = getActorOutputArgs.parse(args);
 
         // Parse fields into array