181
packages/coding-agent/test/sdk-openrouter-attribution.test.ts
Normal file
181
packages/coding-agent/test/sdk-openrouter-attribution.test.ts
Normal file
@@ -0,0 +1,181 @@
|
||||
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,
|
||||
type SimpleStreamOptions,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { AuthStorage } from "../src/core/auth-storage.js";
|
||||
import { ModelRegistry } from "../src/core/model-registry.js";
|
||||
import { createAgentSession } from "../src/core/sdk.js";
|
||||
import { SessionManager } from "../src/core/session-manager.js";
|
||||
import { SettingsManager } from "../src/core/settings-manager.js";
|
||||
|
||||
describe("createAgentSession OpenRouter 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)}`);
|
||||
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): Model<Api> {
|
||||
return {
|
||||
id: `${provider}-test-model`,
|
||||
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>;
|
||||
} = {},
|
||||
): Promise<Record<string, string> | undefined> {
|
||||
const settingsManager = SettingsManager.create(cwd, agentDir);
|
||||
if (options.telemetryEnabled === false) {
|
||||
settingsManager.setEnableInstallTelemetry(false);
|
||||
}
|
||||
|
||||
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
|
||||
authStorage.setRuntimeApiKey(model.provider, "test-api-key");
|
||||
const modelRegistry = ModelRegistry.create(authStorage, join(agentDir, "models.json"));
|
||||
const registeredProviders = ["capture-provider"];
|
||||
let capturedOptions: SimpleStreamOptions | undefined;
|
||||
|
||||
modelRegistry.registerProvider("capture-provider", {
|
||||
api: "openai-completions",
|
||||
streamSimple: (_model, _context, providerOptions) => {
|
||||
capturedOptions = providerOptions;
|
||||
return createDoneStream();
|
||||
},
|
||||
});
|
||||
|
||||
if (options.providerHeaders) {
|
||||
modelRegistry.registerProvider(model.provider, { headers: options.providerHeaders });
|
||||
registeredProviders.push(model.provider);
|
||||
}
|
||||
|
||||
const { session } = await createAgentSession({
|
||||
cwd,
|
||||
agentDir,
|
||||
model,
|
||||
authStorage,
|
||||
modelRegistry,
|
||||
settingsManager,
|
||||
sessionManager: SessionManager.inMemory(cwd),
|
||||
});
|
||||
|
||||
try {
|
||||
await session.agent.streamFn(
|
||||
model,
|
||||
{ messages: [] },
|
||||
options.requestHeaders ? { headers: options.requestHeaders } : undefined,
|
||||
);
|
||||
return capturedOptions?.headers;
|
||||
} finally {
|
||||
session.dispose();
|
||||
for (const provider of registeredProviders.reverse()) {
|
||||
modelRegistry.unregisterProvider(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("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");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user