diff --git a/docs/aimlapi-setup.md b/docs/aimlapi-setup.md index 9197d7c378..b55c087e48 100644 --- a/docs/aimlapi-setup.md +++ b/docs/aimlapi-setup.md @@ -86,10 +86,53 @@ for the request and response format. - `/model` lists chat-capable models discovered from the catalog. - Send any prompt to confirm responses come back from the selected model. +## AI/ML API MCP server (Toolbox and media generation) + +AI/ML API also runs a remote MCP server at `https://mcp.aimlapi.com/mcp`. It adds 19 tools to a session: the Toolbox (`tools_search`, `tools_inspect`, `tools_run`, `tools_run_status`), `image_generate`, `video_generate`, `audio_generate`, `llm_chat`, `embeddings_create`, `job_status`, `job_result`, model lookup (`models_list`, `models_search`, `models_find_best`, `models_compare`, `model_get`), `account_balance`, `account_usage`, and `ping`. + +> **Billing:** Toolbox runs (`tools_run`), generations, `llm_chat` and embeddings are charged to your AI/ML API balance. `ping`, `tools_search` and `tools_inspect` are free. The server is never connected unless you turn it on. + +### One step: the built-in `aimlapi` plugin + +If AI/ML API is already your provider, turn the server on with: + +```bash +openclaude plugin enable aimlapi@builtin +``` + +You can also open `/plugin` in a session, find the built-in **aimlapi** plugin and enable it there. To turn it off again: + +```bash +openclaude plugin disable aimlapi@builtin +``` + +- The plugin signs in with the AI/ML API key you already configured. It reads the key from `AIMLAPI_API_KEY`, which an active aimlapi.com provider profile sets for the session. You can also export the variable yourself. The settings file only stores `"aimlapi@builtin": true`. The key is never written there. +- The plugin is off by default, so having a key does not connect the server. It is also hidden when no `AIMLAPI_API_KEY` is set. It stays hidden when your aimlapi profile points at a proxy base URL (`OPENAI_BASE_URL` or `OPENAI_API_BASE`), so a proxy credential is never sent to aimlapi.com. The check runs again whenever the MCP configuration is rebuilt and whenever the server connects, so this also holds after you switch `/provider` to such a profile during a session. +- If `AIMLAPI_API_KEY` holds a comma-separated pool, only its first usable key is sent to the MCP server, with no rotation across keys. +- The server appears as `plugin:aimlapi:aimlapi`, and its tools appear as `mcp__plugin_aimlapi_aimlapi__`, for example `mcp__plugin_aimlapi_aimlapi__tools_search`. +- To check it, run `openclaude mcp list`. You should see `plugin:aimlapi:aimlapi: https://mcp.aimlapi.com/mcp (HTTP) - ✓ Connected`. In a session, `/mcp` shows the same server. + +### Without the plugin: `openclaude mcp add` + +The standard MCP commands work too. To use your API key, keep the single quotes so that your shell does not expand the variable. OpenClaude expands `${AIMLAPI_API_KEY}` when it connects, so the config file stores the reference and not the key: + +```bash +openclaude mcp add --transport http -s user aimlapi https://mcp.aimlapi.com/mcp \ + --header 'Authorization: Bearer ${AIMLAPI_API_KEY}' +``` + +To sign in with your aimlapi.com account (OAuth) instead, leave out the header. `openclaude mcp list` reports **Needs authentication**. In a session, open `/mcp`, select the server and authenticate in the browser: + +```bash +openclaude mcp add --transport http -s user aimlapi https://mcp.aimlapi.com/mcp +``` + +If you configure the server both ways, the hand-added entry wins and the plugin copy is skipped as a duplicate. Note that `openclaude mcp get aimlapi` prints header values after expansion, which includes your key. + ## Notes - Model discovery uses the public, unauthenticated `GET /models` endpoint and surfaces only chat-completions models; image, audio, embeddings, and other modalities are intentionally not routed through the coding workflow. -- Every request to AI/ML API (inference, catalog, sign-in/account, checkout, and Toolbox web search) carries the two mandatory attribution headers `X-AIMLAPI-Source: agent/openclaude` and `X-AIMLAPI-Partner-ID`. On the canonical `api.aimlapi.com` endpoint the inference/catalog requests additionally send `X-AIMLAPI-Integration-*` (owner/repo/version) plus `HTTP-Referer: OpenClaude` and `X-Title: OpenClaude`; all attribution headers are stripped when a non-canonical (proxy) base URL is configured, so a third-party host never receives OpenClaude's partner identity. +- Every request to AI/ML API (inference, catalog, sign-in/account, checkout, and Toolbox web search) carries the two mandatory attribution headers `X-AIMLAPI-Source: agent/openclaude` and `X-AIMLAPI-Partner-ID`. On the canonical `api.aimlapi.com` endpoint the inference/catalog requests additionally send `X-AIMLAPI-Integration-*` (owner/repo/version) plus `HTTP-Referer: OpenClaude` and `X-Title: OpenClaude`; all attribution headers are stripped when a non-canonical (proxy) base URL is configured, so a third-party host never receives OpenClaude's partner identity. MCP requests to `mcp.aimlapi.com` (matched by exact hostname, from the plugin or a hand-added server) carry `X-AIMLAPI-Source` and `X-AIMLAPI-Partner-ID`. Other MCP servers never receive them. - `AIMLAPI_AUTH_URL`, `AIMLAPI_APP_URL`, `AIMLAPI_PAY_URL`, and `AIMLAPI_VERIFICATION_BASE_URL` can point the complete flow at another environment. Guided **new-account** onboarding and key provisioning require the canonical `api.aimlapi.com` inference endpoint, so `AIMLAPI_INFERENCE_URL` must be unset (or left at its default) to create a new account and mint its key — a non-canonical value is rejected there. The existing-key top-up path (paste an aimlapi.com key) instead runs against whatever `AIMLAPI_INFERENCE_URL` is configured. `AIMLAPI_RETURN_URL` overrides only the browser landing page. - Checkout success is detected by polling. The browser return target is an HTTPS page; OpenClaude does not install a custom URL-scheme handler. - Usage (`/usage`) reporting is not supported for this provider. diff --git a/src/cli/handlers/mcp.tsx b/src/cli/handlers/mcp.tsx index 9a2dc89751..030b5b2011 100644 --- a/src/cli/handlers/mcp.tsx +++ b/src/cli/handlers/mcp.tsx @@ -248,6 +248,13 @@ export async function mcpRemoveHandler(name: string, options: { // mcp list (lines 4641–4688) export async function mcpListHandler(): Promise { logEvent('tengu_mcp_list', {}); + // The mcp subcommands skip the session startup that registers built-in + // plugins; register them so an enabled built-in plugin's MCP server (e.g. + // aimlapi@builtin) is listed and health-checked the way a session sees it. + const { + initBuiltinPlugins + } = await import('../../plugins/bundled/index.js'); + initBuiltinPlugins(); const { servers: configs } = await getAllMcpConfigs(); diff --git a/src/integrations/aimlapi/config.test.ts b/src/integrations/aimlapi/config.test.ts index 2b7dec338a..25efa78729 100644 --- a/src/integrations/aimlapi/config.test.ts +++ b/src/integrations/aimlapi/config.test.ts @@ -2,12 +2,15 @@ import { afterEach, beforeEach, expect, test } from 'bun:test' import { AIMLAPI_MANAGED_ATTRIBUTION_HEADER_NAMES, + AIMLAPI_MCP_URL, DEFAULT_PARTNER_ID, buildPartnerCheckoutReturnUrls, buildPartnerReturnUrl, + isAimlapiMcpUrl, isCanonicalAimlapiInferenceBaseUrl, isTrustedAimlapiRequestUrl, resolveAimlapiAttributionHeaders, + resolveAimlapiMcpHeaders, resolvePartnerId, resolveEndpoints, withResolvedPartnerHeader, @@ -264,3 +267,117 @@ test('proxy attribution strips every managed name spelling and keeps X-Tenant', expect(headers.get('x-tenant')).toBe('acme') expect(Object.keys(resolved)).toEqual(['X-Tenant']) }) + +test('MCP host gate is an exact https hostname match, never a suffix', () => { + expect(isAimlapiMcpUrl(AIMLAPI_MCP_URL)).toBe(true) + expect(isAimlapiMcpUrl('https://MCP.aimlapi.com/mcp')).toBe(true) + expect(isAimlapiMcpUrl('https://mcp.aimlapi.com:443/other')).toBe(true) + // Sibling aimlapi hosts are not the MCP server. + expect(isAimlapiMcpUrl('https://api.aimlapi.com/v1')).toBe(false) + expect(isAimlapiMcpUrl('https://staging.mcp.aimlapi.com/mcp')).toBe(false) + // Look-alikes, cleartext and garbage are rejected. + expect(isAimlapiMcpUrl('https://mcp.aimlapi.com.attacker.test/mcp')).toBe(false) + expect(isAimlapiMcpUrl('https://evilmcp.aimlapi.com/mcp')).toBe(false) + expect(isAimlapiMcpUrl('https://attacker.test/mcp.aimlapi.com/mcp')).toBe(false) + expect(isAimlapiMcpUrl('http://mcp.aimlapi.com/mcp')).toBe(false) + expect(isAimlapiMcpUrl('not-a-url')).toBe(false) +}) + +test('MCP headers gain both mandatory attribution headers on the MCP host only', () => { + const resolved = resolveAimlapiMcpHeaders( + { Authorization: 'Bearer test-key' }, + AIMLAPI_MCP_URL, + ) + expectSingleManagedAttribution(resolved, { + 'X-AIMLAPI-Source': 'agent/openclaude', + 'X-AIMLAPI-Partner-ID': DEFAULT_PARTNER_ID, + }) + expect(resolved.Authorization).toBe('Bearer test-key') + + // Any other MCP server keeps exactly the headers it was configured with. + for (const url of [ + 'https://mcp.example.test/mcp', + 'https://mcp.aimlapi.com.attacker.test/mcp', + 'http://mcp.aimlapi.com/mcp', + ]) { + expect(resolveAimlapiMcpHeaders({ 'X-Tenant': 'acme' }, url)).toEqual({ + 'X-Tenant': 'acme', + }) + } +}) + +test('MCP attribution restamps a configured partner id instead of combining it', () => { + const resolved = resolveAimlapiMcpHeaders( + { + 'x-aimlapi-partner-id': 'part_attacker', + 'X-AIMLAPI-SOURCE': 'agent/attacker', + 'X-Tenant': 'acme', + }, + AIMLAPI_MCP_URL, + ) + expectSingleManagedAttribution(resolved, { + 'X-AIMLAPI-Source': 'agent/openclaude', + 'X-AIMLAPI-Partner-ID': DEFAULT_PARTNER_ID, + }) + expect(new Headers(resolved).get('x-aimlapi-partner-id')).not.toContain(',') + expect(resolved['X-Tenant']).toBe('acme') +}) + +test('MCP Authorization keeps only the first key of a pooled credential', () => { + expect( + resolveAimlapiMcpHeaders( + { Authorization: 'Bearer aiml-first-key, aiml-second-key' }, + AIMLAPI_MCP_URL, + ).Authorization, + ).toBe('Bearer aiml-first-key') + // A single key and a non-bearer value pass through unchanged. + expect( + resolveAimlapiMcpHeaders({ Authorization: 'Bearer aiml-only-key' }, AIMLAPI_MCP_URL) + .Authorization, + ).toBe('Bearer aiml-only-key') + expect( + resolveAimlapiMcpHeaders({ authorization: 'Basic abc,def' }, AIMLAPI_MCP_URL) + .authorization, + ).toBe('Basic abc,def') + // Other MCP servers keep exactly what they were configured with. + expect( + resolveAimlapiMcpHeaders( + { Authorization: 'Bearer first, second' }, + 'https://mcp.example.test/mcp', + ), + ).toEqual({ Authorization: 'Bearer first, second' }) +}) + +test('MCP Authorization reduces a line-broken pool to one header-safe key', () => { + // An expanded `Bearer ${AIMLAPI_API_KEY}` keeps whatever whitespace the pool + // was written with. firstUsableCredential accepts these pools, so the MCP + // header must reduce them too instead of handing Headers a line break. + for (const pool of [ + 'aiml-first-key,\naiml-second-key', + 'aiml-first-key,\r\naiml-second-key', + 'aiml-first-key\n, aiml-second-key', + ',\naiml-first-key, aiml-second-key', + '\r\n,aiml-first-key,\r\naiml-second-key', + ]) { + const authorization = resolveAimlapiMcpHeaders( + { Authorization: `Bearer ${pool}` }, + AIMLAPI_MCP_URL, + ).Authorization + expect(authorization).toBe('Bearer aiml-first-key') + expect(() => new Headers({ Authorization: authorization! })).not.toThrow() + } + // The ordinary comma-and-space pool is unchanged by the multiline handling. + expect( + resolveAimlapiMcpHeaders( + { Authorization: 'Bearer aiml-first-key, aiml-second-key' }, + AIMLAPI_MCP_URL, + ).Authorization, + ).toBe('Bearer aiml-first-key') + // Other hosts still keep exactly what they were configured with. + expect( + resolveAimlapiMcpHeaders( + { Authorization: 'Bearer first,\nsecond' }, + 'https://mcp.example.test/mcp', + ), + ).toEqual({ Authorization: 'Bearer first,\nsecond' }) +}) diff --git a/src/integrations/aimlapi/config.ts b/src/integrations/aimlapi/config.ts index a71b69db9d..f3537ecd53 100644 --- a/src/integrations/aimlapi/config.ts +++ b/src/integrations/aimlapi/config.ts @@ -9,6 +9,8 @@ * Override any single URL via the corresponding `AIMLAPI_*_URL` env var. */ +import { firstUsableCredential } from '../../services/api/credentialPool.js' + export type AimlapiEndpoints = { /** app/auth service - mints the user access (Bearer) token. */ authBaseUrl: string @@ -251,6 +253,119 @@ export function resolveAimlapiAttributionHeaders( return rest } +/** + * AI/ML API MCP server (streamable HTTP). It accepts the same + * `Authorization: Bearer ` as the inference API and falls back to + * OAuth when no Authorization header is sent. + */ +export const AIMLAPI_MCP_URL = 'https://mcp.aimlapi.com/mcp' +const AIMLAPI_MCP_HOSTNAME = new URL(AIMLAPI_MCP_URL).hostname + +/** + * True only for an https URL whose hostname is exactly the AI/ML API MCP host. + * Deliberately not a suffix match: attribution must never reach a look-alike + * (`mcp.aimlapi.com.evil.test`) or a sibling host the user pointed elsewhere. + */ +export function isAimlapiMcpUrl(url: string): boolean { + try { + const parsed = new URL(url) + return ( + parsed.protocol === 'https:' && + parsed.hostname.toLowerCase() === AIMLAPI_MCP_HOSTNAME + ) + } catch { + return false + } +} + +/** + * Whether the session holds an AI/ML API key that may be sent to the MCP host. + * When the aimlapi route is retargeted to a non-canonical base URL (a user + * proxy), the profile exports that proxy's credential under the same variable + * name, so it is withheld rather than forwarded to aimlapi.com. + * + * Read from the live environment on every call: a provider switch can replace + * AIMLAPI_API_KEY at any point in a session. + */ +export function hasAimlapiMcpCredential( + env: NodeJS.ProcessEnv = process.env, +): boolean { + if (!firstUsableCredential(env.AIMLAPI_API_KEY)) return false + const baseUrl = env.OPENAI_BASE_URL?.trim() || env.OPENAI_API_BASE?.trim() + if ( + env.CLAUDE_CODE_PROVIDER_ROUTE_ID === 'aimlapi' && + baseUrl && + !isCanonicalAimlapiInferenceBaseUrl(baseUrl) + ) { + return false + } + return true +} + +/** + * `pluginSource` stamped on the MCP server of the built-in `aimlapi` plugin + * (`{name}@builtin`, see src/plugins/bundled/aimlapiMcp.ts). + */ +export const AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE = 'aimlapi@builtin' + +/** + * Connection-time guard for the built-in plugin's server. Its config is only + * resolved while hasAimlapiMcpCredential() holds, but a provider switch can + * replace the session's AIMLAPI_API_KEY afterwards. The connection is then + * refused rather than authenticated with a key the session no longer routes + * to aimlapi.com. Servers from any other source are not affected. + */ +export function assertAimlapiMcpBuiltinCredential( + pluginSource: string | undefined, + env: NodeJS.ProcessEnv = process.env, +): void { + if (pluginSource !== AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE) return + if (hasAimlapiMcpCredential(env)) return + throw new Error( + 'The aimlapi@builtin MCP server was not connected: AIMLAPI_API_KEY is not set to an aimlapi.com key in this session.', + ) +} + +/** + * An aimlapi provider profile can export a comma-separated credential pool as + * AIMLAPI_API_KEY, so `Bearer ${AIMLAPI_API_KEY}` can expand to several keys. + * The MCP server takes one key: keep the first usable one, as the web search + * backend does. Any other Authorization value is returned unchanged. + * + * The payload is matched with `[\s\S]` so a pool whose entries are separated + * by line breaks (LF or CRLF) reaches the shared selector whole; it trims every + * entry, so the selected key never carries a line break into the header. + */ +function singleBearerCredential(value: string): string { + const match = /^Bearer\s+([\s\S]+)$/i.exec(value.trim()) + if (!match?.[1]?.includes(',')) return value + const key = firstUsableCredential(match[1]) + return key ? `Bearer ${key}` : value +} + +/** + * Stamp the mandatory attribution headers onto an MCP server's request headers + * when (and only when) the server is the AI/ML API MCP host. This covers both + * the built-in `aimlapi` plugin server and a server the user added by hand + * with `openclaude mcp add`. Any other MCP server's headers are returned + * untouched. Managed names are restamped through the same helper the + * inference path uses, so a configured header cannot replace the partner id. + */ +export function resolveAimlapiMcpHeaders( + headers: Readonly>, + url: string, +): Record { + if (!isAimlapiMcpUrl(url)) return { ...headers } + const single: Record = {} + for (const [name, value] of Object.entries(headers)) { + single[name] = name.trim().toLowerCase() === 'authorization' + ? singleBearerCredential(value) + : value + } + // An undefined base URL selects the canonical (attributed) branch. + return resolveAimlapiAttributionHeaders(single, undefined) +} + /** * Build the co-branded checkout return URLs the hosted payment page redirects * to after the user pays or cancels. Carrying `sessionToken` + `partnerCheckout=1` diff --git a/src/main.tsx b/src/main.tsx index 1d22f19fcb..833ff6e70d 100644 --- a/src/main.tsx +++ b/src/main.tsx @@ -1726,6 +1726,15 @@ async function run(): Promise { return allowed; }) : Promise.resolve({}); + // Register built-in plugins before the first plugin/MCP discovery below. + // loadAllPluginsCacheOnly() memoizes its result, and getInputPrompt() can + // await piped stdin long enough for that discovery to finish, so a later + // registration would leave an enabled built-in plugin (and its MCP + // servers) out of the session. In-memory map writes only, no I/O. + if (process.env.CLAUDE_CODE_ENTRYPOINT !== 'local-agent') { + initBuiltinPlugins(); + } + // Kick off MCP config loading early (safe - just reads files, no execution). // Both interactive and -p use getClaudeCodeMcpConfigs (local file reads only). // The local promise is awaited later (before prefetchAllMcpResources) to @@ -1854,12 +1863,12 @@ async function run(): Promise { // since --worktree makes setup() process.chdir() (setup.ts:203), and // commands/agents need the post-chdir cwd. const preSetupCwd = getCwd(); - // Register bundled skills/plugins before kicking getCommands() — they're - // pure in-memory array pushes (<1ms, zero I/O) that getBundledSkills() - // reads synchronously. Previously ran inside setup() after ~20ms of - // await points, so the parallel getCommands() memoized an empty list. + // Register bundled skills before kicking getCommands() — they're pure + // in-memory array pushes (<1ms, zero I/O) that getBundledSkills() reads + // synchronously. Previously ran inside setup() after ~20ms of await + // points, so the parallel getCommands() memoized an empty list. Built-in + // plugins are registered earlier, before the first MCP config load. if (process.env.CLAUDE_CODE_ENTRYPOINT !== 'local-agent') { - initBuiltinPlugins(); initBundledSkills(); } const setupPromise = setup(preSetupCwd, permissionMode, allowDangerouslySkipPermissions, worktreeEnabled, worktreeName, tmuxEnabled, sessionId ? validateUuid(sessionId) : undefined, worktreePRNumber, messagingSocketPath); diff --git a/src/plugins/bundled/aimlapiMcp.test.ts b/src/plugins/bundled/aimlapiMcp.test.ts new file mode 100644 index 0000000000..b3f5d4ccee --- /dev/null +++ b/src/plugins/bundled/aimlapiMcp.test.ts @@ -0,0 +1,377 @@ +import { afterEach, beforeEach, expect, mock, test } from 'bun:test' + +import { + AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE, + AIMLAPI_MCP_URL, + DEFAULT_PARTNER_ID, + hasAimlapiMcpCredential, +} from '../../integrations/aimlapi/config.js' +import { getMcpServerHeaders } from '../../services/mcp/headersHelper.js' +import type { + McpHTTPServerConfig, + ScopedMcpServerConfig, +} from '../../services/mcp/types.js' +import { + acquireSharedMutationLock, + releaseSharedMutationLock, +} from '../../test/sharedMutationLock.js' +import { getPluginMcpServers } from '../../utils/plugins/mcpPluginIntegration.js' +import { + clearPluginCache, + loadAllPluginsCacheOnly, +} from '../../utils/plugins/pluginLoader.js' +import { applyProviderProfileToProcessEnv } from '../../utils/providerProfiles.js' +import { + resetSettingsCache, + setSessionSettingsCache, +} from '../../utils/settings/settingsCache.js' +import type { SettingsJson } from '../../utils/settings/types.js' +import { clearBuiltinPlugins, getBuiltinPlugins } from '../builtinPlugins.js' +import { + AIMLAPI_MCP_PLUGIN_NAME, + registerAimlapiMcpPlugin, +} from './aimlapiMcp.js' + +const PLUGIN_ID = `${AIMLAPI_MCP_PLUGIN_NAME}@builtin` +const SERVER_KEY = `plugin:${AIMLAPI_MCP_PLUGIN_NAME}:aimlapi` +const TEST_KEY = 'aimlapi-test-key-not-real' +const PROXY_KEY = 'proxy-test-credential-not-real' +const envNames = [ + 'AIMLAPI_API_KEY', + 'OPENAI_BASE_URL', + 'OPENAI_API_BASE', + 'CLAUDE_CODE_PROVIDER_ROUTE_ID', +] as const +const originalEnv = Object.fromEntries( + envNames.map(name => [name, process.env[name]]), +) + +// Test files share a process, and an earlier file can leave its own +// settings.js mock registered. Built-in plugins read their enabled state +// through getSettings_DEPRECATED, so pin it to this file's settings and put +// the real module back after each test. +const realSettings = { + ...await import(`../../utils/settings/settings.js?real=${Date.now()}-${Math.random()}`), +} as typeof import('../../utils/settings/settings.js') + +let settingsForTest: SettingsJson = {} + +function useEnabledPlugins(enabledPlugins?: Record): void { + settingsForTest = enabledPlugins ? { enabledPlugins } : {} + setSessionSettingsCache({ settings: settingsForTest, errors: [] }) +} + +beforeEach(async () => { + await acquireSharedMutationLock('plugins/bundled/aimlapiMcp.test.ts') + settingsForTest = {} + mock.module('../../utils/settings/settings.js', () => ({ + ...realSettings, + getSettings_DEPRECATED: () => settingsForTest, + })) + for (const name of envNames) delete process.env[name] + registerAimlapiMcpPlugin() +}) + +afterEach(() => { + try { + // A test that ran plugin discovery leaves its memoized result behind; + // drop it so no later test or file sees this file's built-in plugin. + clearPluginCache() + clearBuiltinPlugins() + mock.module('../../utils/settings/settings.js', () => realSettings) + resetSettingsCache() + settingsForTest = {} + for (const name of envNames) { + const value = originalEnv[name] + if (value === undefined) delete process.env[name] + else process.env[name] = value + } + } finally { + releaseSharedMutationLock() + } +}) + +test('stays disabled when an AIMLAPI key is present but the user has not opted in', () => { + process.env.AIMLAPI_API_KEY = TEST_KEY + useEnabledPlugins() + + const { enabled, disabled } = getBuiltinPlugins() + + expect(enabled.map(plugin => plugin.name)).not.toContain( + AIMLAPI_MCP_PLUGIN_NAME, + ) + expect(disabled.map(plugin => plugin.name)).toContain(AIMLAPI_MCP_PLUGIN_NAME) +}) + +test('enabling the plugin exposes the MCP server with the configured key and attribution', async () => { + process.env.AIMLAPI_API_KEY = TEST_KEY + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const plugin = getBuiltinPlugins().enabled.find( + candidate => candidate.name === AIMLAPI_MCP_PLUGIN_NAME, + ) + expect(plugin).toBeDefined() + + const servers = await getPluginMcpServers(plugin!) + const server = servers?.[`plugin:${AIMLAPI_MCP_PLUGIN_NAME}:aimlapi`] as + | McpHTTPServerConfig + | undefined + expect(server?.type).toBe('http') + expect(server?.url).toBe(AIMLAPI_MCP_URL) + + expect(await getMcpServerHeaders('plugin:aimlapi:aimlapi', server!)).toEqual({ + Authorization: `Bearer ${TEST_KEY}`, + 'X-AIMLAPI-Partner-ID': DEFAULT_PARTNER_ID, + 'X-AIMLAPI-Source': 'agent/openclaude', + }) +}) + +test('a pooled AIMLAPI_API_KEY sends only its first key to the MCP server', async () => { + process.env.AIMLAPI_API_KEY = `${TEST_KEY}, aimlapi-second-key-not-real` + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const plugin = getBuiltinPlugins().enabled.find( + candidate => candidate.name === AIMLAPI_MCP_PLUGIN_NAME, + ) + const servers = await getPluginMcpServers(plugin!) + const server = servers?.[`plugin:${AIMLAPI_MCP_PLUGIN_NAME}:aimlapi`] as + | McpHTTPServerConfig + | undefined + + expect( + (await getMcpServerHeaders('plugin:aimlapi:aimlapi', server!)).Authorization, + ).toBe(`Bearer ${TEST_KEY}`) +}) + +test('the stored server config references the key instead of embedding it', () => { + process.env.AIMLAPI_API_KEY = TEST_KEY + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const plugin = getBuiltinPlugins().enabled.find( + candidate => candidate.name === AIMLAPI_MCP_PLUGIN_NAME, + ) + const serialized = JSON.stringify(plugin?.mcpServers) + + expect(serialized).toContain('${AIMLAPI_API_KEY}') + expect(serialized).not.toContain(TEST_KEY) +}) + +test('is hidden without an AIMLAPI key even when enabled', () => { + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const { enabled, disabled } = getBuiltinPlugins() + + expect( + [...enabled, ...disabled].map(plugin => plugin.name), + ).not.toContain(AIMLAPI_MCP_PLUGIN_NAME) +}) + +test('withholds a credential the aimlapi route exported for a proxy base URL', () => { + expect(hasAimlapiMcpCredential({ AIMLAPI_API_KEY: TEST_KEY })).toBe(true) + expect(hasAimlapiMcpCredential({ AIMLAPI_API_KEY: ' ' })).toBe(false) + expect(hasAimlapiMcpCredential({ AIMLAPI_API_KEY: 'null' })).toBe(false) + expect( + hasAimlapiMcpCredential({ + AIMLAPI_API_KEY: TEST_KEY, + CLAUDE_CODE_PROVIDER_ROUTE_ID: 'aimlapi', + OPENAI_API_BASE: 'https://proxy.example.test/v1', + }), + ).toBe(false) + expect( + hasAimlapiMcpCredential({ + AIMLAPI_API_KEY: TEST_KEY, + CLAUDE_CODE_PROVIDER_ROUTE_ID: 'aimlapi', + OPENAI_BASE_URL: 'https://api.aimlapi.com/v1', + }), + ).toBe(true) + expect( + hasAimlapiMcpCredential({ + AIMLAPI_API_KEY: TEST_KEY, + CLAUDE_CODE_PROVIDER_ROUTE_ID: 'aimlapi', + OPENAI_BASE_URL: 'https://proxy.example.test/v1', + }), + ).toBe(false) + // Another provider's base URL does not make an exported AIMLAPI key a proxy one. + expect( + hasAimlapiMcpCredential({ + AIMLAPI_API_KEY: TEST_KEY, + CLAUDE_CODE_PROVIDER_ROUTE_ID: 'openrouter', + OPENAI_BASE_URL: 'https://openrouter.ai/api/v1', + }), + ).toBe(true) +}) + +async function withRestoredProcessEnv( + run: () => void | Promise, +): Promise { + const snapshot = { ...process.env } + try { + await run() + } finally { + for (const name of Object.keys(process.env)) { + if (!(name in snapshot)) delete process.env[name] + } + Object.assign(process.env, snapshot) + } +} + +test('a saved aimlapi.com provider profile supplies the key without exporting it', async () => { + await withRestoredProcessEnv(() => { + applyProviderProfileToProcessEnv({ + id: 'provider_aimlapi', + name: 'AI/ML API', + provider: 'aimlapi', + baseUrl: 'https://api.aimlapi.com/v1', + model: 'gpt-4o', + apiKey: TEST_KEY, + }) + + expect(process.env.AIMLAPI_API_KEY).toBe(TEST_KEY) + expect(hasAimlapiMcpCredential()).toBe(true) + }) +}) + +test('a saved aimlapi profile retargeted to a proxy never offers its credential', async () => { + await withRestoredProcessEnv(() => { + applyProviderProfileToProcessEnv({ + id: 'provider_aimlapi_proxy', + name: 'AI/ML API Proxy', + provider: 'aimlapi', + baseUrl: 'https://proxy.example.test/v1', + model: 'gpt-4o', + apiKey: 'proxy-credential', + }) + + expect(process.env.AIMLAPI_API_KEY).toBe('proxy-credential') + expect(hasAimlapiMcpCredential()).toBe(false) + }) +}) + +const canonicalProfile = { + id: 'provider_aimlapi', + name: 'AI/ML API', + provider: 'aimlapi', + baseUrl: 'https://api.aimlapi.com/v1', + model: 'gpt-4o', + apiKey: TEST_KEY, +} as const + +const proxyProfile = { + id: 'provider_aimlapi_proxy', + name: 'AI/ML API Proxy', + provider: 'aimlapi', + baseUrl: 'https://proxy.example.test/v1', + model: 'gpt-4o', + apiKey: PROXY_KEY, +} as const + +/** + * The plugin part of getClaudeCodeMcpConfigs(): memoized discovery, then a + * fresh resolution of every enabled plugin's MCP servers. + */ +async function rebuildPluginMcpServer(): Promise< + McpHTTPServerConfig | undefined +> { + const { enabled } = await loadAllPluginsCacheOnly() + const servers: Record = {} + for (const resolved of await Promise.all( + enabled.map(plugin => getPluginMcpServers(plugin)), + )) { + if (resolved) Object.assign(servers, resolved) + } + return servers[SERVER_KEY] as McpHTTPServerConfig | undefined +} + +test('the plugin server carries the builtin plugin source the connection guard checks', async () => { + process.env.AIMLAPI_API_KEY = TEST_KEY + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const plugin = getBuiltinPlugins().enabled.find( + candidate => candidate.name === AIMLAPI_MCP_PLUGIN_NAME, + ) + const servers = await getPluginMcpServers(plugin!) + + expect(PLUGIN_ID).toBe(AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE) + expect(servers?.[SERVER_KEY]?.pluginSource).toBe( + AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE, + ) +}) + +test('a switch to a proxy profile after cached discovery never resolves or sends the proxy credential', async () => { + // Start from a cold discovery cache; nothing below clears it again. + clearPluginCache() + useEnabledPlugins({ [PLUGIN_ID]: true }) + + await withRestoredProcessEnv(async () => { + applyProviderProfileToProcessEnv(canonicalProfile) + expect(hasAimlapiMcpCredential()).toBe(true) + + // Positive control: the canonical key reaches the MCP host. + const canonical = await rebuildPluginMcpServer() + expect(canonical?.url).toBe(AIMLAPI_MCP_URL) + expect(canonical?.headers?.Authorization).toBe(`Bearer ${TEST_KEY}`) + expect( + (await getMcpServerHeaders(SERVER_KEY, canonical!)).Authorization, + ).toBe(`Bearer ${TEST_KEY}`) + + // /provider switches to an aimlapi profile whose base URL and key belong + // to a proxy. The memoized discovery still lists the enabled plugin... + applyProviderProfileToProcessEnv(proxyProfile) + expect(process.env.AIMLAPI_API_KEY).toBe(PROXY_KEY) + expect( + (await loadAllPluginsCacheOnly()).enabled.map(plugin => plugin.name), + ).toContain(AIMLAPI_MCP_PLUGIN_NAME) + + // ...but a rebuilt MCP configuration (as after /clear) has no server for + // it, so the proxy credential is never resolved into one. + const rebuilt = await rebuildPluginMcpServer() + expect(rebuilt).toBeUndefined() + + // A config resolved before the switch is not connected either while the + // session's key belongs to the proxy. + await expect(getMcpServerHeaders(SERVER_KEY, canonical!)).rejects.toThrow( + AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE, + ) + + // Back on the canonical profile the same cached plugin resolves again. + applyProviderProfileToProcessEnv(canonicalProfile) + const restored = await rebuildPluginMcpServer() + expect(restored?.headers?.Authorization).toBe(`Bearer ${TEST_KEY}`) + expect( + (await getMcpServerHeaders(SERVER_KEY, restored!)).Authorization, + ).toBe(`Bearer ${TEST_KEY}`) + }) +}) + +test('the connection guard leaves a hand-added server on the MCP host alone', async () => { + process.env.AIMLAPI_API_KEY = PROXY_KEY + process.env.CLAUDE_CODE_PROVIDER_ROUTE_ID = 'aimlapi' + process.env.OPENAI_BASE_URL = 'https://proxy.example.test/v1' + expect(hasAimlapiMcpCredential()).toBe(false) + + // Explicit authentication on a server the user added is their choice. + const headers = await getMcpServerHeaders('aimlapi', { + type: 'http', + url: AIMLAPI_MCP_URL, + headers: { Authorization: 'Bearer user-chosen-key' }, + }) + expect(headers.Authorization).toBe('Bearer user-chosen-key') +}) + +test('a line-broken AIMLAPI_API_KEY pool reaches the MCP server as one header-safe key', async () => { + process.env.AIMLAPI_API_KEY = `${TEST_KEY},\r\naimlapi-second-key-not-real` + useEnabledPlugins({ [PLUGIN_ID]: true }) + + const plugin = getBuiltinPlugins().enabled.find( + candidate => candidate.name === AIMLAPI_MCP_PLUGIN_NAME, + ) + expect(plugin).toBeDefined() + const servers = await getPluginMcpServers(plugin!) + const headers = await getMcpServerHeaders( + SERVER_KEY, + servers![SERVER_KEY] as McpHTTPServerConfig, + ) + + expect(headers.Authorization).toBe(`Bearer ${TEST_KEY}`) + expect(() => new Headers(headers)).not.toThrow() +}) diff --git a/src/plugins/bundled/aimlapiMcp.ts b/src/plugins/bundled/aimlapiMcp.ts new file mode 100644 index 0000000000..3e25926145 --- /dev/null +++ b/src/plugins/bundled/aimlapiMcp.ts @@ -0,0 +1,43 @@ +import { + AIMLAPI_MCP_URL, + hasAimlapiMcpCredential, +} from '../../integrations/aimlapi/config.js' +import { registerBuiltinPlugin } from '../builtinPlugins.js' + +export const AIMLAPI_MCP_PLUGIN_NAME = 'aimlapi' +export const AIMLAPI_MCP_SERVER_NAME = 'aimlapi' + +/** + * The plugin authenticates with the AI/ML API key the user already configured. + * It is an environment reference, expanded when the server config is loaded, + * so the key itself is never copied into settings or an MCP config file. An + * active aimlapi.com provider profile exports its saved key as + * AIMLAPI_API_KEY for the session, and a user can export it directly. + */ +const AIMLAPI_MCP_AUTHORIZATION = 'Bearer ${AIMLAPI_API_KEY}' + +/** + * Opt-in built-in plugin exposing the AI/ML API MCP server (Toolbox `tools_*`, + * image/video/audio generation, model and account tools). Disabled by default: + * Toolbox runs and generations are billed to the user's AI/ML API balance, so + * a key being present must never connect it on its own. Enable it with + * `openclaude plugin enable aimlapi@builtin` or from the /plugin UI. Attribution + * headers are added for the MCP host by `getMcpServerHeaders`. + */ +export function registerAimlapiMcpPlugin(): void { + registerBuiltinPlugin({ + name: AIMLAPI_MCP_PLUGIN_NAME, + description: + 'AI/ML API MCP server: Toolbox tools, image/video/audio generation, model and account tools. Uses your AIMLAPI_API_KEY; tool runs and generations are billed to your AI/ML API balance.', + version: '1.0.0', + defaultEnabled: false, + isAvailable: () => hasAimlapiMcpCredential(), + mcpServers: { + [AIMLAPI_MCP_SERVER_NAME]: { + type: 'http', + url: AIMLAPI_MCP_URL, + headers: { Authorization: AIMLAPI_MCP_AUTHORIZATION }, + }, + }, + }) +} diff --git a/src/plugins/bundled/index.ts b/src/plugins/bundled/index.ts index c27876ceb1..daae3d3b26 100644 --- a/src/plugins/bundled/index.ts +++ b/src/plugins/bundled/index.ts @@ -14,6 +14,7 @@ * 2. Call that helper from initBuiltinPlugins() */ +import { registerAimlapiMcpPlugin } from './aimlapiMcp.js' import { registerKarpathyGuidelinesPlugin } from './karpathyGuidelines.js' /** @@ -21,4 +22,5 @@ import { registerKarpathyGuidelinesPlugin } from './karpathyGuidelines.js' */ export function initBuiltinPlugins(): void { registerKarpathyGuidelinesPlugin() + registerAimlapiMcpPlugin() } diff --git a/src/plugins/bundled/startupRegistration.test.ts b/src/plugins/bundled/startupRegistration.test.ts new file mode 100644 index 0000000000..9f3453ba9b --- /dev/null +++ b/src/plugins/bundled/startupRegistration.test.ts @@ -0,0 +1,103 @@ +import { expect, test } from 'bun:test' +import { spawn } from 'node:child_process' +import { existsSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +import { AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE } from '../../integrations/aimlapi/config.js' + +const cliPath = resolve(import.meta.dir, '../../../dist/cli.mjs') +const CHILD_TIMEOUT_MS = 30_000 +const DISCOVERY_LINE = /Found \d+ plugins \((\d+) enabled/ + +/** + * Start a fresh print-mode session whose stdin is a pipe that never delivers + * input, and return the enabled count of the session's FIRST plugin discovery. + * That discovery is memoized, and the pending stdin keeps getInputPrompt() + * waiting while it completes, so a built-in plugin registered after it is + * missing from the whole session. + */ +function firstDiscoveryEnabledCount(withAimlapiKey: boolean): Promise { + const configDir = mkdtempSync(join(tmpdir(), 'openclaude-builtin-startup-')) + writeFileSync( + join(configDir, 'settings.json'), + JSON.stringify({ + enabledPlugins: { [AIMLAPI_MCP_BUILTIN_PLUGIN_SOURCE]: true }, + }), + ) + const env: Record = { + ...process.env, + OPENCLAUDE_CONFIG_DIR: configDir, + OPENCLAUDE_DISABLE_TELEMETRY: '1', + } + // The child must run its entrypoint as a plain print-mode session. + for (const name of [ + 'OPENCLAUDE_DISABLE_CLI_ENTRYPOINT_AUTO_RUN', + 'CLAUDE_CODE_ENTRYPOINT', + 'CLAUDE_CODE_SIMPLE', + 'CLAUDE_CODE_SYNC_PLUGIN_INSTALL', + 'CLAUDE_CODE_PROVIDER_ROUTE_ID', + 'CLAUDE_CODE_USE_OPENAI', + 'OPENAI_BASE_URL', + 'OPENAI_API_BASE', + 'AIMLAPI_API_KEY', + ]) { + delete env[name] + } + if (withAimlapiKey) env.AIMLAPI_API_KEY = 'aimlapi-test-key-not-real' + + return new Promise((resolveCount, reject) => { + const child = spawn('node', [cliPath, '-p', '--debug-to-stderr'], { + cwd: configDir, + env, + // stdin stays open and silent: the delayed, non-TTY input case. + stdio: ['pipe', 'ignore', 'pipe'], + }) + let stderr = '' + let settled = false + const finish = (result: number | Error) => { + if (settled) return + settled = true + clearTimeout(timer) + child.kill('SIGKILL') + rmSync(configDir, { recursive: true, force: true }) + if (result instanceof Error) reject(result) + else resolveCount(result) + } + const timer = setTimeout( + () => + finish(new Error(`No plugin discovery was logged:\n${stderr}`)), + CHILD_TIMEOUT_MS, + ) + child.stderr.setEncoding('utf8') + child.stderr.on('data', (chunk: string) => { + stderr += chunk + const match = DISCOVERY_LINE.exec(stderr) + if (match) finish(Number(match[1])) + }) + child.once('error', finish) + // 'close' fires after stderr is drained, so a line logged just before an + // exit is still matched above first. + child.once('close', code => + finish( + new Error( + `The CLI exited (${code}) before logging plugin discovery:\n${stderr}`, + ), + ), + ) + }) +} + +test('an enabled built-in plugin is in the first discovery of a session with pending piped input', async () => { + if (!existsSync(cliPath)) return // needs `bun run build`; always present in CI + + // The same settings enable aimlapi@builtin in both runs; it is available + // only when AIMLAPI_API_KEY is set. Comparing the two keeps the assertion + // independent of any other built-in plugin's default state. + const [withKey, withoutKey] = await Promise.all([ + firstDiscoveryEnabledCount(true), + firstDiscoveryEnabledCount(false), + ]) + + expect(withKey).toBe(withoutKey + 1) +}, 2 * CHILD_TIMEOUT_MS) diff --git a/src/services/mcp/headersHelper.test.ts b/src/services/mcp/headersHelper.test.ts new file mode 100644 index 0000000000..f4bedd5a37 --- /dev/null +++ b/src/services/mcp/headersHelper.test.ts @@ -0,0 +1,35 @@ +import { expect, test } from 'bun:test' + +import { + AIMLAPI_MCP_URL, + DEFAULT_PARTNER_ID, +} from '../../integrations/aimlapi/config.js' +import { getMcpServerHeaders } from './headersHelper.js' + +test('requests to the AI/ML API MCP host carry the attribution headers', async () => { + const headers = await getMcpServerHeaders('aimlapi', { + type: 'http', + url: AIMLAPI_MCP_URL, + headers: { Authorization: 'Bearer test-key' }, + }) + + expect(headers).toEqual({ + Authorization: 'Bearer test-key', + 'X-AIMLAPI-Partner-ID': DEFAULT_PARTNER_ID, + 'X-AIMLAPI-Source': 'agent/openclaude', + }) +}) + +test('other MCP servers never receive AI/ML API attribution', async () => { + for (const url of [ + 'https://mcp.example.test/mcp', + 'https://mcp.aimlapi.com.attacker.test/mcp', + ]) { + const headers = await getMcpServerHeaders('other', { + type: 'http', + url, + headers: { Authorization: 'Bearer other-key' }, + }) + expect(headers).toEqual({ Authorization: 'Bearer other-key' }) + } +}) diff --git a/src/services/mcp/headersHelper.ts b/src/services/mcp/headersHelper.ts index 7fea3ccdc6..44860e436e 100644 --- a/src/services/mcp/headersHelper.ts +++ b/src/services/mcp/headersHelper.ts @@ -1,4 +1,8 @@ import { getIsNonInteractiveSession } from '../../bootstrap/state.js' +import { + assertAimlapiMcpBuiltinCredential, + resolveAimlapiMcpHeaders, +} from '../../integrations/aimlapi/config.js' import { checkHasTrustDialogAccepted } from '../../utils/config.js' import { logAntError } from '../../utils/debug.js' import { errorMessage } from '../../utils/errors.js' @@ -125,15 +129,26 @@ export async function getMcpHeadersFromHelper( */ export async function getMcpServerHeaders( serverName: string, - config: McpSSEServerConfig | McpHTTPServerConfig | McpWebSocketServerConfig, + config: ( + | McpSSEServerConfig + | McpHTTPServerConfig + | McpWebSocketServerConfig + ) & { pluginSource?: string }, ): Promise> { + // The built-in aimlapi plugin's server authenticates with AIMLAPI_API_KEY; + // recheck the live session before its headers are used to connect. + assertAimlapiMcpBuiltinCredential(config.pluginSource) const staticHeaders = config.headers || {} const dynamicHeaders = (await getMcpHeadersFromHelper(serverName, config)) || {} - // Dynamic headers override static headers if both are present - return { - ...staticHeaders, - ...dynamicHeaders, - } + // Dynamic headers override static headers if both are present. Requests to + // the AI/ML API MCP host (exact hostname) also carry its attribution headers. + return resolveAimlapiMcpHeaders( + { + ...staticHeaders, + ...dynamicHeaders, + }, + config.url, + ) } diff --git a/src/utils/plugins/mcpPluginIntegration.ts b/src/utils/plugins/mcpPluginIntegration.ts index f13a12c48c..994254d865 100644 --- a/src/utils/plugins/mcpPluginIntegration.ts +++ b/src/utils/plugins/mcpPluginIntegration.ts @@ -1,4 +1,5 @@ import { join } from 'path' +import { getBuiltinPluginDefinition } from '../../plugins/builtinPlugins.js' import { expandEnvVarsInString } from '../../services/mcp/envExpansion.js' import { type McpServerConfig, @@ -334,6 +335,20 @@ function loadChannelUserConfig( return loadMcpServerUserConfig(plugin.repository, serverName) ?? undefined } +/** + * A built-in plugin's isAvailable() runs during plugin discovery, and that + * result is memoized (loadAllPluginsCacheOnly). Its MCP server configs are + * resolved from the environment afresh on every load, so check availability + * again here: a server whose environment stopped qualifying after discovery + * (e.g. a provider switch replaced the credential it references) is not + * resolved at all. + */ +function isBuiltinPluginUnavailable(plugin: LoadedPlugin): boolean { + if (!plugin.isBuiltin) return false + const definition = getBuiltinPluginDefinition(plugin.name) + return definition?.isAvailable ? !definition.isAvailable() : false +} + /** * Add plugin scope to MCP server configs * This adds a prefix to server names to avoid conflicts between plugins @@ -371,7 +386,7 @@ export async function extractMcpServersFromPlugins( const scopedResults = await Promise.all( plugins.map(async plugin => { - if (!plugin.enabled) return null + if (!plugin.enabled || isBuiltinPluginUnavailable(plugin)) return null const servers = await loadPluginMcpServers(plugin, errors) if (!servers) return null @@ -590,7 +605,7 @@ export async function getPluginMcpServers( plugin: LoadedPlugin, errors: PluginError[] = [], ): Promise | undefined> { - if (!plugin.enabled) { + if (!plugin.enabled || isBuiltinPluginUnavailable(plugin)) { return undefined }