mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(ai,agent,coding-agent): move client identity and provider attribution into pi-ai
pi-ai gets a ClientIdentity request option and createModels() default that choose the User-Agent (detailed with version, short, or custom) and provider attribution headers for OpenRouter, NVIDIA NIM, and OpenCode. Attribution is off by default in pi-ai and on by default in pi-agent-core and pi-durable (HarnessOptions.identity). pi-coding-agent sets its identity from the telemetry setting instead of adding headers itself.
This commit is contained in:
Armin Ronacher
Mario Zechner
parent
adfc14461a
commit
adc5fa9f80
55 changed files with 958 additions and 465 deletions
|
|
@ -8,6 +8,7 @@
|
|||
|
||||
### Changed
|
||||
|
||||
- Provider attribution headers now come from pi-ai's client identity. Extensions' `before_provider_headers` handlers no longer see the OpenRouter, NVIDIA NIM, and OpenCode attribution headers; returning `null` for one still removes it. Cloudflare requests send pi's detailed `User-Agent` instead of `pi-coding-agent`, as do GitHub release checks for tool downloads, and OpenCode requests send `x-opencode-client` whenever attribution is enabled, not only with a session.
|
||||
- `generateSummary()` and `generateSummaryWithUsage()` now ignore their `reserveTokens` parameter. Summary requests use the model's output limit instead.
|
||||
- Changed the default GitHub Copilot model from `gpt-5.4`, which Copilot retired, to `gpt-6.1-sol`.
|
||||
|
||||
|
|
|
|||
40
packages/coding-agent/src/core/client-identity.ts
Normal file
40
packages/coding-agent/src/core/client-identity.ts
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
import type { Api, ClientIdentity, Model, ProviderHeaders } from "@earendil-works/pi-ai";
|
||||
import { VERSION } from "../config.ts";
|
||||
import type { SettingsManager } from "./settings-manager.ts";
|
||||
import { isInstallTelemetryEnabled } from "./telemetry.ts";
|
||||
|
||||
const OPENCODE_HOST = "opencode.ai";
|
||||
|
||||
/**
|
||||
* How pi's requests identify it to providers: a detailed `User-Agent` with pi's version, and attribution to pi as a
|
||||
* CLI agent while install telemetry is enabled.
|
||||
*/
|
||||
export function getClientIdentity(settingsManager: SettingsManager): ClientIdentity {
|
||||
return {
|
||||
userAgent: "detailed",
|
||||
version: VERSION,
|
||||
attribution: isInstallTelemetryEnabled(settingsManager) ? { categories: "cli-agent" } : false,
|
||||
};
|
||||
}
|
||||
|
||||
function isOpenCodeModel(model: Model<Api>): boolean {
|
||||
if (model.provider === "opencode" || model.provider === "opencode-go") return true;
|
||||
try {
|
||||
return new URL(model.baseUrl).hostname === OPENCODE_HOST;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* OpenCode's per-conversation routing header. pi-ai's OpenCode providers add it themselves; this also covers custom
|
||||
* providers pointing at OpenCode. Configured and request headers override it.
|
||||
*/
|
||||
export function withOpenCodeSessionHeader(
|
||||
model: Model<Api>,
|
||||
sessionId: string | undefined,
|
||||
headers: ProviderHeaders,
|
||||
): ProviderHeaders {
|
||||
if (!sessionId || !isOpenCodeModel(model)) return headers;
|
||||
return { "x-opencode-session": sessionId, ...headers };
|
||||
}
|
||||
|
|
@ -1,97 +0,0 @@
|
|||
import type { Api, Model, ProviderHeaders } 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<Api>): boolean {
|
||||
return model.provider === "openrouter" || model.baseUrl.includes(OPENROUTER_HOST);
|
||||
}
|
||||
|
||||
function isNvidiaNimModel(model: Model<Api>): boolean {
|
||||
return model.provider === "nvidia" || matchesHost(model.baseUrl, NVIDIA_NIM_HOST);
|
||||
}
|
||||
|
||||
function isCloudflareModel(model: Model<Api>): 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<Api>,
|
||||
settingsManager: SettingsManager,
|
||||
): Record<string, string> | 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<Api>, sessionId: string | undefined): Record<string, string> | 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<Api>,
|
||||
settingsManager: SettingsManager,
|
||||
sessionId: string | undefined,
|
||||
...headerSources: Array<ProviderHeaders | undefined>
|
||||
): ProviderHeaders | undefined {
|
||||
const merged: ProviderHeaders = {
|
||||
...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;
|
||||
}
|
||||
|
|
@ -7,13 +7,13 @@ import { resolvePath } from "../utils/paths.ts";
|
|||
import { AgentSession } from "./agent-session.ts";
|
||||
import { formatNoModelsAvailableMessage } from "./auth-guidance.ts";
|
||||
import { CacheWarmer } from "./cache-warmer.ts";
|
||||
import { getClientIdentity, withOpenCodeSessionHeader } from "./client-identity.ts";
|
||||
import { DEFAULT_THINKING_LEVEL } from "./defaults.ts";
|
||||
import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.ts";
|
||||
import { createToolNameMatcher } from "./mcp-servers.ts";
|
||||
import { convertToLlm } from "./messages.ts";
|
||||
import { findInitialModel } from "./model-resolver.ts";
|
||||
import { ModelRuntime } from "./model-runtime.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";
|
||||
|
|
@ -351,16 +351,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
|||
websocketConnectTimeoutMs: options.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(),
|
||||
maxRetries: options.maxRetries ?? providerRetrySettings.maxRetries,
|
||||
maxRetryDelayMs: options.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
|
||||
// Replaces the agent loop's default identity; read per request so telemetry changes apply at once.
|
||||
identity: getClientIdentity(settingsManager),
|
||||
transformHeaders: async (requestHeaders) => {
|
||||
const headers = mergeProviderAttributionHeaders(
|
||||
requestModel,
|
||||
settingsManager,
|
||||
options.sessionId,
|
||||
requestHeaders,
|
||||
);
|
||||
const headers = withOpenCodeSessionHeader(requestModel, options.sessionId, requestHeaders);
|
||||
return headerRunner?.hasHandlers("before_provider_headers")
|
||||
? headerRunner.emitBeforeProviderHeaders(headers ?? {})
|
||||
: (headers ?? {});
|
||||
? headerRunner.emitBeforeProviderHeaders(headers)
|
||||
: headers;
|
||||
},
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import { formatUserAgent } from "@earendil-works/pi-ai";
|
||||
|
||||
/** pi's `User-Agent` for its own network requests; the same detailed form pi-ai sends to providers. */
|
||||
export function getPiUserAgent(version: string): string {
|
||||
const runtime = process.versions.bun ? `bun/${process.versions.bun}` : `node/${process.version}`;
|
||||
return `pi/${version} (${process.platform}; ${runtime}; ${process.arch})`;
|
||||
return formatUserAgent({ userAgent: "detailed", version });
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@ import { arch, platform } from "os";
|
|||
import { join } from "path";
|
||||
import { Readable } from "stream";
|
||||
import { pipeline } from "stream/promises";
|
||||
import { APP_NAME, getBinDir } from "../config.ts";
|
||||
import { getBinDir, VERSION } from "../config.ts";
|
||||
import { fetchWithRetry } from "./management-http.ts";
|
||||
import { getPiUserAgent } from "./pi-user-agent.ts";
|
||||
|
||||
const TOOLS_DIR = getBinDir();
|
||||
const NETWORK_TIMEOUT_MS = 10_000;
|
||||
|
|
@ -111,7 +112,7 @@ export async function getLatestVersion(repo: string): Promise<string> {
|
|||
const response = await fetchWithRetry(
|
||||
`https://github.com/${repo}/releases/latest`,
|
||||
{
|
||||
headers: { "User-Agent": `${APP_NAME}-coding-agent` },
|
||||
headers: { "User-Agent": getPiUserAgent(VERSION) },
|
||||
redirect: "manual",
|
||||
},
|
||||
{ timeoutMs: NETWORK_TIMEOUT_MS },
|
||||
|
|
|
|||
190
packages/coding-agent/test/sdk-client-identity.test.ts
Normal file
190
packages/coding-agent/test/sdk-client-identity.test.ts
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
import { existsSync, mkdirSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import {
|
||||
type Api,
|
||||
type AssistantMessage,
|
||||
createAssistantMessageEventStream,
|
||||
getAttributionHeaders,
|
||||
type Model,
|
||||
normalizeContext,
|
||||
type SimpleStreamOptions,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { VERSION } from "../src/config.ts";
|
||||
import { AuthStorage } from "../src/core/auth-storage.ts";
|
||||
import { createAgentSession } from "../src/core/sdk.ts";
|
||||
import { SessionManager } from "../src/core/session-manager.ts";
|
||||
import { SettingsManager } from "../src/core/settings-manager.ts";
|
||||
import { createInMemoryModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts";
|
||||
|
||||
describe("createAgentSession client identity", () => {
|
||||
let tempDir: string;
|
||||
let cwd: string;
|
||||
let agentDir: string;
|
||||
let originalTelemetryEnv: string | undefined;
|
||||
|
||||
beforeEach(() => {
|
||||
tempDir = join(tmpdir(), `pi-sdk-identity-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
|
||||
cwd = join(tempDir, "project");
|
||||
agentDir = join(tempDir, "agent");
|
||||
mkdirSync(cwd, { recursive: true });
|
||||
mkdirSync(agentDir, { recursive: true });
|
||||
originalTelemetryEnv = process.env.PI_TELEMETRY;
|
||||
delete process.env.PI_TELEMETRY;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalTelemetryEnv === undefined) {
|
||||
delete process.env.PI_TELEMETRY;
|
||||
} else {
|
||||
process.env.PI_TELEMETRY = originalTelemetryEnv;
|
||||
}
|
||||
if (tempDir && existsSync(tempDir)) {
|
||||
rmSync(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
function createModel(provider: string, baseUrl: string, id = `${provider}-test-model`): Model<Api> {
|
||||
return {
|
||||
id,
|
||||
name: `${provider} Test Model`,
|
||||
api: "openai-completions",
|
||||
provider,
|
||||
baseUrl,
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
};
|
||||
}
|
||||
|
||||
function createDoneStream() {
|
||||
const stream = createAssistantMessageEventStream();
|
||||
const message: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "ok" }],
|
||||
api: "openai-completions",
|
||||
provider: "capture-provider",
|
||||
model: "capture-model",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
stream.end(message);
|
||||
return stream;
|
||||
}
|
||||
|
||||
async function captureOptions(
|
||||
model: Model<Api>,
|
||||
options: {
|
||||
telemetryEnabled?: boolean;
|
||||
providerHeaders?: Record<string, string>;
|
||||
sessionId?: string;
|
||||
} = {},
|
||||
): Promise<SimpleStreamOptions | undefined> {
|
||||
const settingsManager = SettingsManager.create(cwd, agentDir);
|
||||
if (options.telemetryEnabled === false) {
|
||||
settingsManager.setEnableInstallTelemetry(false);
|
||||
}
|
||||
|
||||
const authStorage = AuthStorage.inMemory({
|
||||
[model.provider]: { type: "api_key", key: "test-api-key" },
|
||||
});
|
||||
const modelRegistry = await createInMemoryModelRegistry(authStorage);
|
||||
let capturedOptions: SimpleStreamOptions | undefined;
|
||||
|
||||
modelRegistry.registerProvider(model.provider, {
|
||||
api: model.api,
|
||||
headers: options.providerHeaders,
|
||||
streamSimple: (_model, _context, providerOptions) => {
|
||||
capturedOptions = providerOptions;
|
||||
return createDoneStream();
|
||||
},
|
||||
});
|
||||
|
||||
const modelRuntime = getModelRuntime(modelRegistry);
|
||||
const sessionManager = SessionManager.inMemory(cwd);
|
||||
if (options.sessionId) {
|
||||
sessionManager.newSession({ id: options.sessionId });
|
||||
}
|
||||
|
||||
const { session } = await createAgentSession({
|
||||
cwd,
|
||||
agentDir,
|
||||
model,
|
||||
modelRuntime,
|
||||
settingsManager,
|
||||
sessionManager,
|
||||
});
|
||||
|
||||
try {
|
||||
const stream = await session.agent.streamFunction(model, normalizeContext({ messages: [] }), {
|
||||
sessionId: session.sessionId,
|
||||
});
|
||||
await stream.result();
|
||||
return capturedOptions;
|
||||
} finally {
|
||||
session.dispose();
|
||||
modelRegistry.unregisterProvider(model.provider);
|
||||
}
|
||||
}
|
||||
|
||||
it("identifies pi with a detailed User-Agent and attribution as a CLI agent", async () => {
|
||||
const model = createModel("openrouter", "https://openrouter.ai/api/v1");
|
||||
const options = await captureOptions(model);
|
||||
|
||||
expect(options?.identity).toEqual({
|
||||
userAgent: "detailed",
|
||||
version: VERSION,
|
||||
attribution: { categories: "cli-agent" },
|
||||
});
|
||||
expect(getAttributionHeaders(model, options?.identity)).toEqual({
|
||||
"HTTP-Referer": "https://pi.dev",
|
||||
"X-OpenRouter-Title": "pi",
|
||||
"X-OpenRouter-Categories": "cli-agent",
|
||||
});
|
||||
});
|
||||
|
||||
it("turns attribution off when telemetry is disabled", async () => {
|
||||
const model = createModel("nvidia", "https://integrate.api.nvidia.com/v1");
|
||||
const options = await captureOptions(model, { telemetryEnabled: false });
|
||||
|
||||
expect(options?.identity).toEqual({ userAgent: "detailed", version: VERSION, attribution: false });
|
||||
expect(getAttributionHeaders(model, options?.identity)).toEqual({});
|
||||
});
|
||||
|
||||
it("turns attribution off with PI_TELEMETRY=0", async () => {
|
||||
process.env.PI_TELEMETRY = "0";
|
||||
const options = await captureOptions(createModel("openrouter", "https://openrouter.ai/api/v1"));
|
||||
|
||||
expect(options?.identity?.attribution).toBe(false);
|
||||
});
|
||||
|
||||
it("adds OpenCode session headers for custom OpenCode providers", async () => {
|
||||
const options = await captureOptions(createModel("custom-opencode", "https://opencode.ai/zen/v1"), {
|
||||
sessionId: "opencode-session",
|
||||
});
|
||||
|
||||
expect(options?.headers?.["x-opencode-session"]).toBe("opencode-session");
|
||||
});
|
||||
|
||||
it("lets configured OpenCode headers override the session header", async () => {
|
||||
const options = await captureOptions(createModel("opencode", "https://opencode.ai/zen/v1"), {
|
||||
sessionId: "opencode-session",
|
||||
providerHeaders: {
|
||||
"x-opencode-session": "configured-session",
|
||||
},
|
||||
});
|
||||
|
||||
expect(options?.headers?.["x-opencode-session"]).toBe("configured-session");
|
||||
});
|
||||
});
|
||||
|
|
@ -1,263 +0,0 @@
|
|||
import { existsSync, mkdirSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import {
|
||||
type Api,
|
||||
type AssistantMessage,
|
||||
createAssistantMessageEventStream,
|
||||
type Model,
|
||||
normalizeContext,
|
||||
type ProviderHeaders,
|
||||
type SimpleStreamOptions,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { AuthStorage } from "../src/core/auth-storage.ts";
|
||||
import { createAgentSession } from "../src/core/sdk.ts";
|
||||
import { SessionManager } from "../src/core/session-manager.ts";
|
||||
import { SettingsManager } from "../src/core/settings-manager.ts";
|
||||
import { createInMemoryModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts";
|
||||
|
||||
describe("createAgentSession provider attribution headers", () => {
|
||||
let tempDir: string;
|
||||
let cwd: string;
|
||||
let agentDir: string;
|
||||
let originalTelemetryEnv: string | undefined;
|
||||
|
||||
beforeEach(() => {
|
||||
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 });
|
||||
mkdirSync(agentDir, { recursive: true });
|
||||
originalTelemetryEnv = process.env.PI_TELEMETRY;
|
||||
delete process.env.PI_TELEMETRY;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalTelemetryEnv === undefined) {
|
||||
delete process.env.PI_TELEMETRY;
|
||||
} else {
|
||||
process.env.PI_TELEMETRY = originalTelemetryEnv;
|
||||
}
|
||||
if (tempDir && existsSync(tempDir)) {
|
||||
rmSync(tempDir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
function createModel(provider: string, baseUrl: string, id = `${provider}-test-model`): Model<Api> {
|
||||
return {
|
||||
id,
|
||||
name: `${provider} Test Model`,
|
||||
api: "openai-completions",
|
||||
provider,
|
||||
baseUrl,
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
};
|
||||
}
|
||||
|
||||
function createDoneStream() {
|
||||
const stream = createAssistantMessageEventStream();
|
||||
const message: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "ok" }],
|
||||
api: "openai-completions",
|
||||
provider: "capture-provider",
|
||||
model: "capture-model",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
stream.end(message);
|
||||
return stream;
|
||||
}
|
||||
|
||||
async function captureHeaders(
|
||||
model: Model<Api>,
|
||||
options: {
|
||||
telemetryEnabled?: boolean;
|
||||
providerHeaders?: Record<string, string>;
|
||||
requestHeaders?: Record<string, string>;
|
||||
sessionId?: string;
|
||||
} = {},
|
||||
): Promise<ProviderHeaders | undefined> {
|
||||
const settingsManager = SettingsManager.create(cwd, agentDir);
|
||||
if (options.telemetryEnabled === false) {
|
||||
settingsManager.setEnableInstallTelemetry(false);
|
||||
}
|
||||
|
||||
const authStorage = AuthStorage.inMemory({
|
||||
[model.provider]: { type: "api_key", key: "test-api-key" },
|
||||
});
|
||||
const modelRegistry = await createInMemoryModelRegistry(authStorage);
|
||||
let capturedOptions: SimpleStreamOptions | undefined;
|
||||
|
||||
modelRegistry.registerProvider(model.provider, {
|
||||
api: model.api,
|
||||
headers: options.providerHeaders,
|
||||
streamSimple: (_model, _context, providerOptions) => {
|
||||
capturedOptions = providerOptions;
|
||||
return createDoneStream();
|
||||
},
|
||||
});
|
||||
|
||||
const modelRuntime = getModelRuntime(modelRegistry);
|
||||
const sessionManager = SessionManager.inMemory(cwd);
|
||||
if (options.sessionId) {
|
||||
sessionManager.newSession({ id: options.sessionId });
|
||||
}
|
||||
|
||||
const { session } = await createAgentSession({
|
||||
cwd,
|
||||
agentDir,
|
||||
model,
|
||||
modelRuntime,
|
||||
settingsManager,
|
||||
sessionManager,
|
||||
});
|
||||
|
||||
try {
|
||||
const stream = await session.agent.streamFunction(model, normalizeContext({ messages: [] }), {
|
||||
sessionId: session.sessionId,
|
||||
...(options.requestHeaders ? { headers: options.requestHeaders } : {}),
|
||||
});
|
||||
await stream.result();
|
||||
return capturedOptions?.headers;
|
||||
} finally {
|
||||
session.dispose();
|
||||
modelRegistry.unregisterProvider(model.provider);
|
||||
}
|
||||
}
|
||||
|
||||
it("adds default attribution headers for OpenRouter models", async () => {
|
||||
const headers = await captureHeaders(createModel("openrouter", "https://openrouter.ai/api/v1"));
|
||||
|
||||
expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev");
|
||||
expect(headers?.["X-OpenRouter-Title"]).toBe("pi");
|
||||
expect(headers?.["X-OpenRouter-Categories"]).toBe("cli-agent");
|
||||
});
|
||||
|
||||
it("does not add attribution headers when telemetry is disabled", async () => {
|
||||
const headers = await captureHeaders(createModel("openrouter", "https://openrouter.ai/api/v1"), {
|
||||
telemetryEnabled: false,
|
||||
});
|
||||
|
||||
expect(headers?.["HTTP-Referer"]).toBeUndefined();
|
||||
expect(headers?.["X-OpenRouter-Title"]).toBeUndefined();
|
||||
expect(headers?.["X-OpenRouter-Categories"]).toBeUndefined();
|
||||
});
|
||||
|
||||
it("adds attribution headers for custom providers routed through OpenRouter", async () => {
|
||||
const headers = await captureHeaders(createModel("custom-openrouter", "https://openrouter.ai/api/v1"));
|
||||
|
||||
expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev");
|
||||
expect(headers?.["X-OpenRouter-Title"]).toBe("pi");
|
||||
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: {
|
||||
"HTTP-Referer": "https://provider.example",
|
||||
"X-OpenRouter-Categories": "provider-category",
|
||||
},
|
||||
requestHeaders: {
|
||||
"X-OpenRouter-Title": "request-title",
|
||||
},
|
||||
});
|
||||
|
||||
expect(headers?.["HTTP-Referer"]).toBe("https://provider.example");
|
||||
expect(headers?.["X-OpenRouter-Title"]).toBe("request-title");
|
||||
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",
|
||||
});
|
||||
|
||||
expect(headers?.["x-opencode-session"]).toBe("opencode-session");
|
||||
expect(headers?.["x-opencode-client"]).toBe("pi");
|
||||
});
|
||||
|
||||
it("lets configured OpenCode headers override the defaults", async () => {
|
||||
const headers = await captureHeaders(createModel("opencode", "https://opencode.ai/zen/v1"), {
|
||||
sessionId: "opencode-session",
|
||||
providerHeaders: {
|
||||
"x-opencode-session": "configured-session",
|
||||
"x-opencode-client": "configured-client",
|
||||
},
|
||||
});
|
||||
|
||||
expect(headers?.["x-opencode-session"]).toBe("configured-session");
|
||||
expect(headers?.["x-opencode-client"]).toBe("configured-client");
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue