From 601480122970ebc3a9b5ce8059f6b20dd8390e78 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 2 Jun 2026 15:09:36 +0200 Subject: [PATCH] Add NVIDIA NIM provider --- packages/ai/CHANGELOG.md | 4 + packages/ai/README.md | 6 +- packages/ai/scripts/generate-models.ts | 97 +++++ packages/ai/src/env-api-keys.ts | 1 + packages/ai/src/models.generated.ts | 363 ++++++++++++++++++ .../ai/src/providers/openai-completions.ts | 10 +- packages/ai/src/types.ts | 1 + packages/ai/test/stream.test.ts | 24 ++ packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/README.md | 5 +- packages/coding-agent/docs/providers.md | 2 + packages/coding-agent/docs/usage.md | 2 +- packages/coding-agent/src/cli/args.ts | 1 + .../coding-agent/src/core/agent-session.ts | 1 + .../core/compaction/branch-summarization.ts | 31 +- .../coding-agent/src/core/model-resolver.ts | 1 + .../src/core/provider-attribution.ts | 97 +++++ .../src/core/provider-display-names.ts | 1 + packages/coding-agent/src/core/sdk.ts | 52 +-- .../test/sdk-openrouter-attribution.test.ts | 66 +++- 20 files changed, 700 insertions(+), 66 deletions(-) create mode 100644 packages/coding-agent/src/core/provider-attribution.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index af63cdc4..4a616f89 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added NVIDIA NIM as a built-in OpenAI-compatible provider, exposing public NIM models that support tool use. + ### Fixed - Fixed Amazon Bedrock requests to replace blank required user/tool-result text with a placeholder and skip blank replay text blocks ([#4975](https://github.com/earendil-works/pi/issues/4975)). diff --git a/packages/ai/README.md b/packages/ai/README.md index e6624d3b..76a98769 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -54,6 +54,7 @@ Unified LLM API with automatic model discovery, provider configuration, token an - **Azure OpenAI (Responses)** - **OpenAI Codex** (ChatGPT Plus/Pro subscription, requires OAuth, see below) - **DeepSeek** +- **NVIDIA NIM** - **Anthropic** - **Google** - **Vertex AI** (Gemini via Vertex AI) @@ -801,7 +802,7 @@ A **provider** offers models through a specific API. For example: - **Google** models use the `google-generative-ai` API - **OpenAI** models use the `openai-responses` API - **Mistral** models use the `mistral-conversations` API -- **xAI, Cerebras, Groq, Together AI, etc.** models use the `openai-completions` API (OpenAI-compatible) +- **xAI, Cerebras, Groq, NVIDIA NIM, Together AI, etc.** models use the `openai-completions` API (OpenAI-compatible) ### Querying Providers and Models @@ -923,7 +924,7 @@ const ollamaReasoningModel: Model<'openai-completions'> = { ### OpenAI Compatibility Settings -The `openai-completions` API is implemented by many providers with minor differences. By default, the library auto-detects compatibility settings based on `baseUrl` for a small set of known OpenAI-compatible providers (Cerebras, xAI, Chutes, DeepSeek, Together AI, zAi, OpenCode, Cloudflare Workers AI, etc.). For custom proxies or unknown endpoints, you can override these settings via the `compat` field. For `openai-responses` models, the compat field only supports Responses-specific flags. +The `openai-completions` API is implemented by many providers with minor differences. By default, the library auto-detects compatibility settings based on `baseUrl` for a small set of known OpenAI-compatible providers (Cerebras, xAI, Chutes, DeepSeek, NVIDIA NIM, Together AI, zAi, OpenCode, Cloudflare Workers AI, etc.). For custom proxies or unknown endpoints, you can override these settings via the `compat` field. For `openai-responses` models, the compat field only supports Responses-specific flags. ```typescript interface OpenAICompletionsCompat { @@ -1102,6 +1103,7 @@ In Node.js environments, you can set environment variables to avoid passing API | Azure OpenAI | `AZURE_OPENAI_API_KEY` + `AZURE_OPENAI_BASE_URL` (e.g. `https://{resource}.openai.azure.com`) or `AZURE_OPENAI_RESOURCE_NAME`. Supports `*.openai.azure.com` and `*.cognitiveservices.azure.com`; root endpoints auto-normalize to `/openai/v1`. Optional: `AZURE_OPENAI_API_VERSION` (default `v1`), `AZURE_OPENAI_DEPLOYMENT_NAME_MAP`. | | Anthropic | `ANTHROPIC_API_KEY` or `ANTHROPIC_OAUTH_TOKEN` | | DeepSeek | `DEEPSEEK_API_KEY` | +| NVIDIA NIM | `NVIDIA_API_KEY` | | Google | `GEMINI_API_KEY` | | Vertex AI | `GOOGLE_CLOUD_API_KEY` or `GOOGLE_CLOUD_PROJECT` (or `GCLOUD_PROJECT`) + `GOOGLE_CLOUD_LOCATION` + ADC | | Mistral | `MISTRAL_API_KEY` | diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 50df88f2..8fb62d14 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -32,12 +32,17 @@ interface ModelsDevModel { }; modalities?: { input?: string[]; + output?: string[]; }; provider?: { npm?: string; }; } +interface NvidiaNimModelListItem { + id: string; +} + interface AiGatewayModel { id: string; name?: string; @@ -117,6 +122,39 @@ const TOGETHER_TOGGLE_REASONING_LEVEL_MAP = { const AI_GATEWAY_MODELS_URL = "https://ai-gateway.vercel.sh/v1"; const AI_GATEWAY_BASE_URL = "https://ai-gateway.vercel.sh"; +const NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"; +const NVIDIA_HEADERS = { + "NVCF-POLL-SECONDS": "3600", +} as const; +const NVIDIA_OPENAI_COMPAT: OpenAICompletionsCompat = { + supportsStore: false, + supportsDeveloperRole: false, + supportsReasoningEffort: false, + maxTokensField: "max_tokens", + supportsStrictMode: false, + supportsLongCacheRetention: false, +}; +const NVIDIA_NIM_UNSUPPORTED_MODELS = new Set([ + "abacusai/dracarys-llama-3.1-70b-instruct", + "bytedance/seed-oss-36b-instruct", + "deepseek-ai/deepseek-v4-flash", + "deepseek-ai/deepseek-v4-pro", + "google/gemma-2-2b-it", + "google/gemma-3n-e2b-it", + "google/gemma-3n-e4b-it", + "google/gemma-4-31b-it", + "meta/llama-3.2-1b-instruct", + "meta/llama-4-maverick-17b-128e-instruct", + "microsoft/phi-4-mini-instruct", + "minimaxai/minimax-m2.7", + "mistralai/mistral-nemotron", + "nvidia/nemotron-mini-4b-instruct", + "qwen/qwen3-next-80b-a3b-instruct", + "qwen/qwen3.5-122b-a10b", + "qwen/qwen3.5-397b-a17b", + "sarvamai/sarvam-m", + "upstage/solar-10.7b-instruct", +]); const ZAI_TOOL_STREAM_UNSUPPORTED_MODELS = new Set(["glm-4.5", "glm-4.5-air", "glm-4.5-flash", "glm-4.5v"]); const EAGER_TOOL_INPUT_STREAMING_UNSUPPORTED_ANTHROPIC_MODELS = new Set([ "github-copilot:claude-haiku-4.5", @@ -312,6 +350,30 @@ function getBedrockBaseUrl(modelId: string): string { : "https://bedrock-runtime.us-east-1.amazonaws.com"; } +function normalizeNvidiaModelId(modelId: string): string { + return modelId.toLowerCase().replaceAll("_", "."); +} + +async function fetchNvidiaNimModelIds(): Promise> { + try { + console.log("Fetching models from NVIDIA NIM API..."); + const response = await fetch(`${NVIDIA_BASE_URL}/models`); + const data = (await response.json()) as { data?: NvidiaNimModelListItem[] }; + const modelIds = new Map(); + + for (const model of data.data ?? []) { + modelIds.set(model.id, model.id); + modelIds.set(normalizeNvidiaModelId(model.id), model.id); + } + + console.log(`Fetched ${data.data?.length ?? 0} model IDs from NVIDIA NIM`); + return modelIds; + } catch (error) { + console.error("Failed to fetch NVIDIA NIM models:", error); + return new Map(); + } +} + async function fetchOpenRouterModels(): Promise[]> { try { console.log("Fetching models from OpenRouter API..."); @@ -435,6 +497,7 @@ async function loadModelsDevData(): Promise[]> { const data = await response.json(); const models: Model[] = []; + const nvidiaNimModelIds = data.nvidia?.models ? await fetchNvidiaNimModelIds() : new Map(); // Process Amazon Bedrock models if (data["amazon-bedrock"]?.models) { @@ -836,6 +899,40 @@ async function loadModelsDevData(): Promise[]> { } } + // Process NVIDIA NIM models + if (data.nvidia?.models) { + for (const [modelId, model] of Object.entries(data.nvidia.models)) { + const m = model as ModelsDevModel; + if (m.tool_call !== true) continue; + if (!m.modalities?.input?.includes("text")) continue; + if (!m.modalities?.output?.includes("text")) continue; + + const liveModelId = nvidiaNimModelIds.get(modelId) ?? nvidiaNimModelIds.get(normalizeNvidiaModelId(modelId)); + if (!liveModelId) continue; + if (NVIDIA_NIM_UNSUPPORTED_MODELS.has(liveModelId)) continue; + + models.push({ + id: liveModelId, + name: m.name || liveModelId, + api: "openai-completions", + provider: "nvidia", + baseUrl: NVIDIA_BASE_URL, + headers: { ...NVIDIA_HEADERS }, + reasoning: m.reasoning === true, + input: m.modalities.input.includes("image") ? ["text", "image"] : ["text"], + cost: { + input: m.cost?.input || 0, + output: m.cost?.output || 0, + cacheRead: m.cost?.cache_read || 0, + cacheWrite: m.cost?.cache_write || 0, + }, + compat: NVIDIA_OPENAI_COMPAT, + contextWindow: m.limit?.context || 4096, + maxTokens: m.limit?.output || 4096, + }); + } + } + // Process Together AI models const togetherProvider = data.together ?? data.togetherai ?? data["together-ai"]; if (togetherProvider?.models) { diff --git a/packages/ai/src/env-api-keys.ts b/packages/ai/src/env-api-keys.ts index a00e8470..6321a03a 100644 --- a/packages/ai/src/env-api-keys.ts +++ b/packages/ai/src/env-api-keys.ts @@ -101,6 +101,7 @@ function getApiKeyEnvVars(provider: string): readonly string[] | undefined { const envMap: Record = { openai: "OPENAI_API_KEY", "azure-openai-responses": "AZURE_OPENAI_API_KEY", + nvidia: "NVIDIA_API_KEY", deepseek: "DEEPSEEK_API_KEY", google: "GEMINI_API_KEY", "google-vertex": "GOOGLE_CLOUD_API_KEY", diff --git a/packages/ai/src/models.generated.ts b/packages/ai/src/models.generated.ts index 8e83c580..cd022b79 100644 --- a/packages/ai/src/models.generated.ts +++ b/packages/ai/src/models.generated.ts @@ -6335,6 +6335,369 @@ export const MODELS = { maxTokens: 262144, } satisfies Model<"openai-completions">, }, + "nvidia": { + "meta/llama-3.1-70b-instruct": { + id: "meta/llama-3.1-70b-instruct", + name: "Llama 3.1 70b Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 4096, + } satisfies Model<"openai-completions">, + "meta/llama-3.1-8b-instruct": { + id: "meta/llama-3.1-8b-instruct", + name: "Llama 3.1 8B Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 16000, + maxTokens: 4096, + } satisfies Model<"openai-completions">, + "meta/llama-3.2-11b-vision-instruct": { + id: "meta/llama-3.2-11b-vision-instruct", + name: "Llama 3.2 11b Vision Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 4096, + } satisfies Model<"openai-completions">, + "meta/llama-3.2-90b-vision-instruct": { + id: "meta/llama-3.2-90b-vision-instruct", + name: "Llama-3.2-90B-Vision-Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 8192, + } satisfies Model<"openai-completions">, + "meta/llama-3.3-70b-instruct": { + id: "meta/llama-3.3-70b-instruct", + name: "Llama 3.3 70b Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 4096, + } satisfies Model<"openai-completions">, + "mistralai/mistral-large-3-675b-instruct-2512": { + id: "mistralai/mistral-large-3-675b-instruct-2512", + name: "Mistral Large 3 675B Instruct 2512", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 262144, + maxTokens: 262144, + } satisfies Model<"openai-completions">, + "mistralai/mistral-small-4-119b-2603": { + id: "mistralai/mistral-small-4-119b-2603", + name: "mistral-small-4-119b-2603", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 8192, + } satisfies Model<"openai-completions">, + "moonshotai/kimi-k2.6": { + id: "moonshotai/kimi-k2.6", + name: "Kimi K2.6", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 262144, + maxTokens: 262144, + } satisfies Model<"openai-completions">, + "nvidia/llama-3.3-nemotron-super-49b-v1": { + id: "nvidia/llama-3.3-nemotron-super-49b-v1", + name: "Llama 3.3 Nemotron Super 49B v1", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 131072, + } satisfies Model<"openai-completions">, + "nvidia/llama-3.3-nemotron-super-49b-v1.5": { + id: "nvidia/llama-3.3-nemotron-super-49b-v1.5", + name: "Llama 3.3 Nemotron Super 49B v1.5", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 131072, + } satisfies Model<"openai-completions">, + "nvidia/nemotron-3-nano-30b-a3b": { + id: "nvidia/nemotron-3-nano-30b-a3b", + name: "nemotron-3-nano-30b-a3b", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 131072, + } satisfies Model<"openai-completions">, + "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning": { + id: "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning", + name: "Nemotron 3 Nano Omni", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 256000, + maxTokens: 65536, + } satisfies Model<"openai-completions">, + "nvidia/nemotron-3-super-120b-a12b": { + id: "nvidia/nemotron-3-super-120b-a12b", + name: "Nemotron 3 Super", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0.2, + output: 0.8, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 262144, + maxTokens: 262144, + } satisfies Model<"openai-completions">, + "nvidia/nvidia-nemotron-nano-9b-v2": { + id: "nvidia/nvidia-nemotron-nano-9b-v2", + name: "nvidia-nemotron-nano-9b-v2", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 131072, + } satisfies Model<"openai-completions">, + "openai/gpt-oss-20b": { + id: "openai/gpt-oss-20b", + name: "GPT OSS 20B", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 32768, + } satisfies Model<"openai-completions">, + "qwen/qwen3-coder-480b-a35b-instruct": { + id: "qwen/qwen3-coder-480b-a35b-instruct", + name: "Qwen3 Coder 480B A35B Instruct", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: false, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 262144, + maxTokens: 66536, + } satisfies Model<"openai-completions">, + "stepfun-ai/step-3.5-flash": { + id: "stepfun-ai/step-3.5-flash", + name: "Step 3.5 Flash", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 256000, + maxTokens: 16384, + } satisfies Model<"openai-completions">, + "stepfun-ai/step-3.7-flash": { + id: "stepfun-ai/step-3.7-flash", + name: "Step 3.7 Flash", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text", "image"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 256000, + maxTokens: 16384, + } satisfies Model<"openai-completions">, + "z-ai/glm-5.1": { + id: "z-ai/glm-5.1", + name: "GLM-5.1", + api: "openai-completions", + provider: "nvidia", + baseUrl: "https://integrate.api.nvidia.com/v1", + headers: {"NVCF-POLL-SECONDS":"3600"}, + compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false}, + reasoning: true, + input: ["text"], + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 131072, + } satisfies Model<"openai-completions">, + }, "openai": { "gpt-4": { id: "gpt-4", diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index ae06326e..c58cbf46 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -1078,8 +1078,10 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet const isOpenRouter = provider === "openrouter" || baseUrl.includes("openrouter.ai"); const isCloudflareWorkersAI = provider === "cloudflare-workers-ai" || baseUrl.includes("api.cloudflare.com"); const isCloudflareAiGateway = provider === "cloudflare-ai-gateway" || baseUrl.includes("gateway.ai.cloudflare.com"); + const isNvidia = provider === "nvidia" || baseUrl.includes("integrate.api.nvidia.com"); const isNonStandard = + isNvidia || provider === "cerebras" || baseUrl.includes("cerebras.ai") || provider === "xai" || @@ -1094,7 +1096,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet isCloudflareWorkersAI || isCloudflareAiGateway; - const useMaxTokens = baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway || isTogether; + const useMaxTokens = baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway || isTogether || isNvidia; const isGrok = provider === "xai" || baseUrl.includes("api.x.ai"); const isDeepSeek = provider === "deepseek" || baseUrl.includes("deepseek.com"); @@ -1103,7 +1105,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet return { supportsStore: !isNonStandard, supportsDeveloperRole: !isNonStandard && !isOpenRouter, - supportsReasoningEffort: !isGrok && !isZai && !isMoonshot && !isTogether && !isCloudflareAiGateway, + supportsReasoningEffort: !isGrok && !isZai && !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, supportsUsageInStreaming: true, maxTokensField: useMaxTokens ? "max_tokens" : "max_completion_tokens", requiresToolResultName: false, @@ -1122,10 +1124,10 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet openRouterRouting: {}, vercelGatewayRouting: {}, zaiToolStream: false, - supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway, + supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, cacheControlFormat, sendSessionAffinityHeaders: false, - supportsLongCacheRetention: !(isTogether || isCloudflareWorkersAI || isCloudflareAiGateway), + supportsLongCacheRetention: !(isTogether || isCloudflareWorkersAI || isCloudflareAiGateway || isNvidia), }; } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index d0b053ea..7a448772 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -28,6 +28,7 @@ export type KnownProvider = | "openai" | "azure-openai-responses" | "openai-codex" + | "nvidia" | "deepseek" | "github-copilot" | "xai" diff --git a/packages/ai/test/stream.test.ts b/packages/ai/test/stream.test.ts index 298104dd..1693f406 100644 --- a/packages/ai/test/stream.test.ts +++ b/packages/ai/test/stream.test.ts @@ -787,6 +787,30 @@ describe("Generate E2E Tests", () => { }); }); + describe.skipIf(!process.env.NVIDIA_API_KEY)("NVIDIA NIM Provider (Nemotron 3 Super via OpenAI Completions)", () => { + const llm = getModel("nvidia", "nvidia/nemotron-3-super-120b-a12b"); + + it("should complete basic text generation", { retry: 3 }, async () => { + await basicTextGeneration(llm); + }); + + it("should handle tool calling", { retry: 3 }, async () => { + await handleToolCall(llm); + }); + + it("should handle streaming", { retry: 3 }, async () => { + await handleStreaming(llm); + }); + + it("should handle thinking mode", { retry: 3 }, async () => { + await handleThinking(llm, { reasoningEffort: "high" }); + }); + + it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => { + await multiTurn(llm, { reasoningEffort: "high" }); + }); + }); + describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter Provider (glm-4.5v via OpenAI Completions)", () => { const llm = getModel("openrouter", "z-ai/glm-4.5v"); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 85b26f8f..28401d53 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- Added NVIDIA NIM provider selection, setup documentation, and direct NIM request attribution headers. - Added `ctx.mode` to extension contexts so extensions can distinguish TUI, RPC, JSON, and print mode. - Added `ctx.getSystemPromptOptions()` for extension commands to inspect the current base system prompt inputs. diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index 0cf35255..585a93a0 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -113,6 +113,7 @@ For each built-in provider, pi maintains a list of tool-capable models, updated - OpenAI - Azure OpenAI - DeepSeek +- NVIDIA NIM - Google Gemini - Google Vertex - Amazon Bedrock @@ -290,7 +291,7 @@ See [docs/settings.md](docs/settings.md) for all options. Pi has two separate startup features: - **Update check:** fetches `https://pi.dev/api/latest-version` to check whether a newer Pi version exists. Disable it with `PI_SKIP_VERSION_CHECK=1`. Disabling update checks only turns off this check. -- **Install/update telemetry:** after first install or a changelog-detected update, sends an anonymous version ping to `https://pi.dev/api/report-install`. Opt out by setting `enableInstallTelemetry` to `false` in `settings.json`, or by setting `PI_TELEMETRY=0`. This does not disable update checks; Pi may still contact `pi.dev` for the latest version unless update checks are disabled or offline mode is enabled. +- **Install/update telemetry:** after first install or a changelog-detected update, sends an anonymous version ping to `https://pi.dev/api/report-install`. This setting also controls optional provider attribution headers for OpenRouter, Cloudflare, and direct NVIDIA NIM requests. Opt out by setting `enableInstallTelemetry` to `false` in `settings.json`, or by setting `PI_TELEMETRY=0`. This does not disable update checks; Pi may still contact `pi.dev` for the latest version unless update checks are disabled or offline mode is enabled. Use `--offline` or `PI_OFFLINE=1` to disable all startup network operations described here, including update checks, package update checks, and install/update telemetry. @@ -641,7 +642,7 @@ pi --thinking high "Solve this complex problem" | `PI_PACKAGE_DIR` | Override package directory (useful for Nix/Guix where store paths tokenize poorly) | | `PI_OFFLINE` | Disable startup network operations, including update checks, package update checks, and install/update telemetry | | `PI_SKIP_VERSION_CHECK` | Skip the Pi version update check at startup. This prevents the `pi.dev` latest-version request | -| `PI_TELEMETRY` | Override install/update telemetry. Use `1`/`true`/`yes` to enable or `0`/`false`/`no` to disable. This does not disable update checks | +| `PI_TELEMETRY` | Override install/update telemetry and provider attribution headers. Use `1`/`true`/`yes` to enable or `0`/`false`/`no` to disable. This does not disable update checks | | `PI_CACHE_RETENTION` | Set to `long` for extended prompt cache (Anthropic: 1h, OpenAI: 24h) | | `VISUAL`, `EDITOR` | External editor for Ctrl+G | diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index 07b6ba08..e7a77c5a 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -52,6 +52,7 @@ pi | Azure OpenAI Responses | `AZURE_OPENAI_API_KEY` | `azure-openai-responses` | | OpenAI | `OPENAI_API_KEY` | `openai` | | DeepSeek | `DEEPSEEK_API_KEY` | `deepseek` | +| NVIDIA NIM | `NVIDIA_API_KEY` | `nvidia` | | Google Gemini | `GEMINI_API_KEY` | `google` | | Mistral | `MISTRAL_API_KEY` | `mistral` | | Groq | `GROQ_API_KEY` | `groq` | @@ -86,6 +87,7 @@ Store credentials in `~/.pi/agent/auth.json`: "anthropic": { "type": "api_key", "key": "sk-ant-..." }, "openai": { "type": "api_key", "key": "sk-..." }, "deepseek": { "type": "api_key", "key": "sk-..." }, + "nvidia": { "type": "api_key", "key": "nvapi-..." }, "google": { "type": "api_key", "key": "..." }, "opencode": { "type": "api_key", "key": "..." }, "opencode-go": { "type": "api_key", "key": "..." }, diff --git a/packages/coding-agent/docs/usage.md b/packages/coding-agent/docs/usage.md index 47d08769..2f2a6449 100644 --- a/packages/coding-agent/docs/usage.md +++ b/packages/coding-agent/docs/usage.md @@ -275,7 +275,7 @@ pi --exclude-tools ask_question | `PI_PACKAGE_DIR` | Override package directory, useful for Nix/Guix store paths | | `PI_OFFLINE` | Disable startup network operations, including update checks, package update checks, and install/update telemetry | | `PI_SKIP_VERSION_CHECK` | Skip the Pi version update check at startup. This prevents the `pi.dev` latest-version request | -| `PI_TELEMETRY` | Override install/update telemetry: `1`/`true`/`yes` or `0`/`false`/`no`. This does not disable update checks | +| `PI_TELEMETRY` | Override install/update telemetry and provider attribution headers: `1`/`true`/`yes` or `0`/`false`/`no`. This does not disable update checks | | `PI_CACHE_RETENTION` | Set to `long` for extended prompt cache where supported | | `VISUAL`, `EDITOR` | External editor for Ctrl+G | diff --git a/packages/coding-agent/src/cli/args.ts b/packages/coding-agent/src/cli/args.ts index 2decae24..8156bd03 100644 --- a/packages/coding-agent/src/cli/args.ts +++ b/packages/coding-agent/src/cli/args.ts @@ -335,6 +335,7 @@ ${chalk.bold("Environment Variables:")} AZURE_OPENAI_API_VERSION - Azure OpenAI API version (default: v1) AZURE_OPENAI_DEPLOYMENT_NAME_MAP - Azure OpenAI model=deployment map (comma-separated) DEEPSEEK_API_KEY - DeepSeek API key + NVIDIA_API_KEY - NVIDIA NIM API key GEMINI_API_KEY - Google Gemini API key GROQ_API_KEY - Groq API key CEREBRAS_API_KEY - Cerebras API key diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 556ca20e..de76a619 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -2787,6 +2787,7 @@ export class AgentSession { customInstructions, replaceInstructions, reserveTokens: branchSummarySettings.reserveTokens, + streamFn: this.agent.streamFn, }); if (result.aborted) { return { cancelled: true, aborted: true }; diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 0039bc79..f6ce3d67 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -5,8 +5,8 @@ * a summary of the branch being left so context isn't lost. */ -import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import type { Model } from "@earendil-works/pi-ai"; +import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; +import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai"; import { completeSimple } from "@earendil-works/pi-ai"; import { convertToLlm, @@ -77,6 +77,8 @@ export interface GenerateBranchSummaryOptions { replaceInstructions?: boolean; /** Tokens reserved for prompt + LLM response (default 16384) */ reserveTokens?: number; + /** Optional session stream function. Used to preserve SDK request behavior without mutating agent state. */ + streamFn?: StreamFn; } // ============================================================================ @@ -284,7 +286,16 @@ export async function generateBranchSummary( entries: SessionEntry[], options: GenerateBranchSummaryOptions, ): Promise { - const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; + const { + model, + apiKey, + headers, + signal, + customInstructions, + replaceInstructions, + reserveTokens = 16384, + streamFn, + } = options; // Token budget = context window minus reserved space for prompt + response const contextWindow = model.contextWindow || 128000; @@ -320,12 +331,14 @@ export async function generateBranchSummary( }, ]; - // Call LLM for summarization - const response = await completeSimple( - model, - { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - { apiKey, headers, signal, maxTokens: 2048 }, - ); + // Call LLM for summarization. Prefer the session stream function so SDK + // request behavior (timeouts, retries, attribution headers) stays consistent + // without running through agent state/events. + const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }; + const requestOptions: SimpleStreamOptions = { apiKey, headers, signal, maxTokens: 2048 }; + const response = streamFn + ? await (await streamFn(model, context, requestOptions)).result() + : await completeSimple(model, context, requestOptions); // Check if aborted or errored if (response.stopReason === "aborted") { diff --git a/packages/coding-agent/src/core/model-resolver.ts b/packages/coding-agent/src/core/model-resolver.ts index 87f7354e..e3af076d 100644 --- a/packages/coding-agent/src/core/model-resolver.ts +++ b/packages/coding-agent/src/core/model-resolver.ts @@ -17,6 +17,7 @@ export const defaultModelPerProvider: Record = { openai: "gpt-5.4", "azure-openai-responses": "gpt-5.4", "openai-codex": "gpt-5.5", + nvidia: "nvidia/nemotron-3-super-120b-a12b", deepseek: "deepseek-v4-pro", google: "gemini-3.1-pro-preview", "google-vertex": "gemini-3.1-pro-preview", diff --git a/packages/coding-agent/src/core/provider-attribution.ts b/packages/coding-agent/src/core/provider-attribution.ts new file mode 100644 index 00000000..91265c92 --- /dev/null +++ b/packages/coding-agent/src/core/provider-attribution.ts @@ -0,0 +1,97 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { SettingsManager } from "./settings-manager.ts"; +import { isInstallTelemetryEnabled } from "./telemetry.ts"; + +const OPENROUTER_HOST = "openrouter.ai"; +const NVIDIA_NIM_HOST = "integrate.api.nvidia.com"; +const CLOUDFLARE_API_HOST = "api.cloudflare.com"; +const CLOUDFLARE_AI_GATEWAY_HOST = "gateway.ai.cloudflare.com"; +const OPENCODE_HOST = "opencode.ai"; + +function matchesHost(baseUrl: string, expectedHost: string): boolean { + try { + return new URL(baseUrl).hostname === expectedHost; + } catch { + return false; + } +} + +function isOpenRouterModel(model: Model): boolean { + return model.provider === "openrouter" || model.baseUrl.includes(OPENROUTER_HOST); +} + +function isNvidiaNimModel(model: Model): boolean { + return model.provider === "nvidia" || matchesHost(model.baseUrl, NVIDIA_NIM_HOST); +} + +function isCloudflareModel(model: Model): boolean { + return ( + model.provider === "cloudflare-workers-ai" || + model.provider === "cloudflare-ai-gateway" || + matchesHost(model.baseUrl, CLOUDFLARE_API_HOST) || + matchesHost(model.baseUrl, CLOUDFLARE_AI_GATEWAY_HOST) + ); +} + +function getDefaultAttributionHeaders( + model: Model, + settingsManager: SettingsManager, +): Record | undefined { + if (!isInstallTelemetryEnabled(settingsManager)) { + return undefined; + } + + if (isOpenRouterModel(model)) { + return { + "HTTP-Referer": "https://pi.dev", + "X-OpenRouter-Title": "pi", + "X-OpenRouter-Categories": "cli-agent", + }; + } + + if (isNvidiaNimModel(model)) { + return { + "X-BILLING-INVOKE-ORIGIN": "Pi", + }; + } + + if (isCloudflareModel(model)) { + return { + "User-Agent": "pi-coding-agent", + }; + } + + return undefined; +} + +function getSessionHeaders(model: Model, sessionId: string | undefined): Record | undefined { + if (!sessionId) return undefined; + if ( + model.provider !== "opencode" && + model.provider !== "opencode-go" && + !matchesHost(model.baseUrl, OPENCODE_HOST) + ) { + return undefined; + } + return { "x-opencode-session": sessionId, "x-opencode-client": "pi" }; +} + +export function mergeProviderAttributionHeaders( + model: Model, + settingsManager: SettingsManager, + sessionId: string | undefined, + ...headerSources: Array | undefined> +): Record | undefined { + const merged = { + ...getSessionHeaders(model, sessionId), + ...getDefaultAttributionHeaders(model, settingsManager), + }; + + for (const headers of headerSources) { + if (headers) { + Object.assign(merged, headers); + } + } + + return Object.keys(merged).length > 0 ? merged : undefined; +} diff --git a/packages/coding-agent/src/core/provider-display-names.ts b/packages/coding-agent/src/core/provider-display-names.ts index 24f7b3ce..789f3153 100644 --- a/packages/coding-agent/src/core/provider-display-names.ts +++ b/packages/coding-agent/src/core/provider-display-names.ts @@ -17,6 +17,7 @@ export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record = { "minimax-cn": "MiniMax (China)", moonshotai: "Moonshot AI", "moonshotai-cn": "Moonshot AI (China)", + nvidia: "NVIDIA NIM", opencode: "OpenCode Zen", "opencode-go": "OpenCode Go", openai: "OpenAI", diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 7a1c922b..e6db747b 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -11,11 +11,11 @@ import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefi import { convertToLlm } from "./messages.ts"; import { ModelRegistry } from "./model-registry.ts"; import { findInitialModel } from "./model-resolver.ts"; +import { mergeProviderAttributionHeaders } from "./provider-attribution.ts"; import type { ResourceLoader } from "./resource-loader.ts"; import { DefaultResourceLoader } from "./resource-loader.ts"; import { getDefaultSessionDir, SessionManager } from "./session-manager.ts"; import { SettingsManager } from "./settings-manager.ts"; -import { isInstallTelemetryEnabled } from "./telemetry.ts"; import { time } from "./timings.ts"; import { createBashTool, @@ -128,44 +128,6 @@ function getDefaultAgentDir(): string { return getAgentDir(); } -function getAttributionHeaders( - model: Model, - settingsManager: SettingsManager, - sessionId?: string, -): Record | undefined { - if ( - sessionId && - (model.provider === "opencode" || model.provider === "opencode-go" || model.baseUrl.includes("opencode.ai")) - ) { - return { "x-opencode-session": sessionId, "x-opencode-client": "pi" }; - } - - if (!isInstallTelemetryEnabled(settingsManager)) { - return undefined; - } - - if (model.provider === "openrouter" || model.baseUrl.includes("openrouter.ai")) { - return { - "HTTP-Referer": "https://pi.dev", - "X-OpenRouter-Title": "pi", - "X-OpenRouter-Categories": "cli-agent", - }; - } - - if ( - model.provider === "cloudflare-workers-ai" || - model.provider === "cloudflare-ai-gateway" || - model.baseUrl.includes("api.cloudflare.com") || - model.baseUrl.includes("gateway.ai.cloudflare.com") - ) { - return { - "User-Agent": "pi-coding-agent", - }; - } - - return undefined; -} - /** * Create an AgentSession with the specified options. * @@ -349,7 +311,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs; const websocketConnectTimeoutMs = options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(); - const attributionHeaders = getAttributionHeaders(model, settingsManager, options?.sessionId); return streamSimple(model, context, { ...options, apiKey: auth.apiKey, @@ -357,10 +318,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} websocketConnectTimeoutMs, maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries, maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs, - headers: - attributionHeaders || auth.headers || options?.headers - ? { ...attributionHeaders, ...auth.headers, ...options?.headers } - : undefined, + headers: mergeProviderAttributionHeaders( + model, + settingsManager, + options?.sessionId, + auth.headers, + options?.headers, + ), }); }, onPayload: async (payload, _model) => { diff --git a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts index c7032103..9f035189 100644 --- a/packages/coding-agent/test/sdk-openrouter-attribution.test.ts +++ b/packages/coding-agent/test/sdk-openrouter-attribution.test.ts @@ -15,14 +15,14 @@ import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; -describe("createAgentSession OpenRouter attribution headers", () => { +describe("createAgentSession provider attribution headers", () => { let tempDir: string; let cwd: string; let agentDir: string; let originalTelemetryEnv: string | undefined; beforeEach(() => { - tempDir = join(tmpdir(), `pi-sdk-openrouter-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); + tempDir = join(tmpdir(), `pi-sdk-attribution-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); cwd = join(tempDir, "project"); agentDir = join(tempDir, "agent"); mkdirSync(cwd, { recursive: true }); @@ -42,9 +42,9 @@ describe("createAgentSession OpenRouter attribution headers", () => { } }); - function createModel(provider: string, baseUrl: string): Model { + function createModel(provider: string, baseUrl: string, id = `${provider}-test-model`): Model { return { - id: `${provider}-test-model`, + id, name: `${provider} Test Model`, api: "openai-completions", provider, @@ -172,6 +172,14 @@ describe("createAgentSession OpenRouter attribution headers", () => { expect(headers?.["X-OpenRouter-Categories"]).toBe("cli-agent"); }); + it("preserves legacy OpenRouter base URL substring attribution matching", async () => { + const headers = await captureHeaders(createModel("custom-openrouter", "not-a-url-openrouter.ai")); + + expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev"); + expect(headers?.["X-OpenRouter-Title"]).toBe("pi"); + expect(headers?.["X-OpenRouter-Categories"]).toBe("cli-agent"); + }); + it("lets provider and request headers override the defaults", async () => { const headers = await captureHeaders(createModel("openrouter", "https://openrouter.ai/api/v1"), { providerHeaders: { @@ -188,6 +196,56 @@ describe("createAgentSession OpenRouter attribution headers", () => { expect(headers?.["X-OpenRouter-Categories"]).toBe("provider-category"); }); + it("adds default attribution headers for direct NVIDIA NIM endpoints", async () => { + const headers = await captureHeaders(createModel("custom-nim", "https://integrate.api.nvidia.com/v1")); + + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Pi"); + }); + + it("adds default attribution headers for the NVIDIA provider", async () => { + const headers = await captureHeaders(createModel("nvidia", "https://example.test/v1")); + + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Pi"); + }); + + it("does not add NVIDIA NIM attribution headers when telemetry is disabled", async () => { + const headers = await captureHeaders(createModel("nvidia", "https://integrate.api.nvidia.com/v1"), { + telemetryEnabled: false, + }); + + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined(); + }); + + it("lets provider and request headers override NVIDIA NIM defaults", async () => { + const headers = await captureHeaders(createModel("nvidia", "https://integrate.api.nvidia.com/v1"), { + providerHeaders: { + "X-BILLING-INVOKE-ORIGIN": "Provider", + }, + requestHeaders: { + "X-BILLING-INVOKE-ORIGIN": "Request", + }, + }); + + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Request"); + }); + + it("does not add NVIDIA NIM attribution headers for NVIDIA models routed through OpenRouter", async () => { + const headers = await captureHeaders( + createModel("openrouter", "https://openrouter.ai/api/v1", "nvidia/nemotron-3-super-120b-a12b"), + ); + + expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev"); + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined(); + }); + + it("does not add NVIDIA NIM attribution headers for NVIDIA models routed through Vercel AI Gateway", async () => { + const headers = await captureHeaders( + createModel("vercel-ai-gateway", "https://ai-gateway.vercel.sh/v1", "nvidia/nemotron-3-super-120b-a12b"), + ); + + expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined(); + }); + it("adds OpenCode session headers", async () => { const headers = await captureHeaders(createModel("opencode", "https://opencode.ai/zen/v1"), { sessionId: "opencode-session",