fix: skyfire parameter injection and missing tool registration (#393)

* 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(-)

* 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>

* refactor: move createApifyClientWithSkyfireSupport to apify-clients.ts

* fix compiler and lint issues after merge

---------

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: MQ37 <29043708+MQ37@users.noreply.github.com>

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.

package-lock.json

@@ -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/actor/server.ts

@@ -3,6 +3,8 @@ import { ApifyClient as _ApifyClient } from 'apify-client';
 import type { AxiosRequestConfig } from 'axios';
 
 import { USER_AGENT_ORIGIN } from './const.js';
+import type { ActorsMcpServer } from './mcp/server.js';
+import type { ApifyToken } from './types.js';
 
 type ExtendedApifyClientOptions = Omit<ApifyClientOptions, 'token'> & {
     token?: string | null | undefined;
@@ -62,3 +64,22 @@ export class ApifyClient extends _ApifyClient {
         });
     }
 }
+
+/**
+ * Creates ApifyClient with appropriate credentials based on Skyfire mode.
+ * In Skyfire mode, uses skyfire-pay-id from args; otherwise uses apifyToken.
+ *
+ * @param apifyMcpServer - The MCP server instance with configuration options
+ * @param args - Tool arguments that may contain skyfire-pay-id
+ * @param apifyToken - Standard Apify token for non-Skyfire mode
+ * @returns ApifyClient instance configured for the appropriate mode
+ */
+export function createApifyClientWithSkyfireSupport(
+    apifyMcpServer: ActorsMcpServer,
+    args: Record<string, unknown>,
+    apifyToken: ApifyToken,
+): ApifyClient {
+    return apifyMcpServer.options.skyfireMode && typeof args['skyfire-pay-id'] === 'string'
+        ? new ApifyClient({ skyfirePayId: args['skyfire-pay-id'] })
+        : new ApifyClient({ token: apifyToken });
+}

src/apify-client.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