diff --git a/.env.example b/.env.example index 3b5e099..2e35024 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,8 @@ OPENROUTER_API_KEY=your_openrouter_api_key_here # are sent to the OpenAI API directly instead of through OpenRouter. OPENAI_API_KEY=your_openai_api_key_here -# Optional: point the OpenAI provider at an OpenAI-compatible endpoint -# (Azure, a gateway, local server...) -# OPENAI_BASE_URL=https://api.openai.com/v1 +# Optional: an OpenAI-compatible endpoint (Ollama, vLLM, LM Studio, Azure, a +# gateway...). Same as --base-url. With no OpenRouter key, every model goes +# here; with one, only OpenAI-style and `openai/`-prefixed ids do. Local +# servers don't need OPENAI_API_KEY. +# OPENAI_BASE_URL=http://localhost:11434/v1 diff --git a/AGENTS.md b/AGENTS.md index c823aea..6782b4b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -41,9 +41,9 @@ bun test Copy `.env.example` to `.env` and set: - `OPENROUTER_API_KEY` - default provider for all models - `OPENAI_API_KEY` - optional; routes `openai/*` and `gpt-*`/`o*` model ids to the OpenAI API -- `OPENAI_BASE_URL` - optional; override the OpenAI endpoint (Azure, gateway, local) +- `OPENAI_BASE_URL` - optional; an OpenAI-compatible endpoint (same as `--base-url`) -At least one key is required. +At least one key is required, or a base URL for a keyless local server. ## Project Architecture @@ -64,7 +64,7 @@ src/ ### Providers (`src/provider.ts`) -Agents never construct a provider directly — they call `resolveModel(modelId, { openrouterApiKey })`. It routes OpenAI-style ids (`openai/*`, `gpt-*`, `o1/o3/o4-*`) to the OpenAI API when `OPENAI_API_KEY` is set, and everything else (plus OpenAI ids with no OpenAI key) through OpenRouter. To add another provider, extend `resolveModel` — don't reintroduce `createOpenRouter` in the agents. +Agents never construct a provider directly — they call `resolveModel(modelId, { openrouterApiKey })`. Without `OPENAI_BASE_URL`, it routes OpenAI-style ids (`openai/*`, `gpt-*`, `o1/o3/o4-*`) to the OpenAI API when `OPENAI_API_KEY` is set, and everything else (plus OpenAI ids with no OpenAI key) through OpenRouter. With `OPENAI_BASE_URL`, OpenAI-style ids always go to that endpoint, and so does every other id when there's no OpenRouter key. A leading `openai/` is stripped before the request. To add another provider, extend `resolveModel` — don't reintroduce `createOpenRouter` in the agents. ### Agent System (`src/agents/`) diff --git a/README.md b/README.md index 05d3557..7c9e550 100644 --- a/README.md +++ b/README.md @@ -126,6 +126,7 @@ zeroleaks techniques | `--no-multi-turn` | Skip multi-turn grooming probes | | `-d, --duration ` | Time budget; 0 = no limit, otherwise more than 30000 (the last 30 s is kept for wrap-up) | | `--injection-model ` | Model for the compliance judge (defaults to the evaluator model) | +| `--base-url ` | Send models to an OpenAI-compatible endpoint (see [Providers](#providers)) | | `-o, --output ` | Write the full JSON result to a file | | `--json` | Print the result as JSON to stdout | | `--no-color` / `-q, --quiet` | Disable color / suppress the progress spinner | @@ -267,7 +268,7 @@ interface ScanResult { By default every model runs through [OpenRouter](https://openrouter.ai), so any OpenRouter slug works (`anthropic/...`, `x-ai/...`, `openai/...`, etc.). -If `OPENAI_API_KEY` is set, OpenAI-style ids such as `openai/gpt-5`, `gpt-5`, and `o3-mini` go straight to the OpenAI API instead. To use an OpenAI-compatible endpoint (Azure, a gateway, a local server), set `OPENAI_BASE_URL`. One scan can mix providers, for example an OpenAI target with an OpenRouter attacker: +If `OPENAI_API_KEY` is set, OpenAI-style ids such as `openai/gpt-5`, `gpt-5`, and `o3-mini` go straight to the OpenAI API instead. One scan can mix providers, for example an OpenAI target with an OpenRouter attacker: ```bash zeroleaks scan -f ./prompt.txt \ @@ -275,15 +276,35 @@ zeroleaks scan -f ./prompt.txt \ --attacker-model "anthropic/claude-opus-4.8" ``` +### OpenAI-compatible endpoints + +Any server that speaks the OpenAI chat completions API works: Ollama, vLLM, LM Studio, llama.cpp, LiteLLM, Azure, Together, Groq, and so on. Pass its URL with `--base-url` (or set `OPENAI_BASE_URL`). With no OpenRouter key, every model goes to that endpoint, whatever its id. Local servers don't need a key. + +```bash +zeroleaks scan -f ./prompt.txt --base-url http://localhost:11434/v1 \ + --attacker-model llama3.1:70b --target-model llama3.1:8b \ + --evaluator-model llama3.1:70b +``` + +For a hosted endpoint, add its key with `--openai-api-key`. + +If you also have an OpenRouter key, only OpenAI-style ids go to the endpoint and the rest go to OpenRouter. Prefix a model with `openai/` to send it to the endpoint anyway. The prefix is stripped, so `openai/llama3.1:8b` reaches the server as `llama3.1:8b`: + +```bash +zeroleaks scan -f ./prompt.txt --base-url http://localhost:11434/v1 \ + --target-model openai/llama3.1:8b \ + --attacker-model "anthropic/claude-opus-4.8" +``` + ## Environment variables | Variable | Description | |----------|-------------| | `OPENROUTER_API_KEY` | OpenRouter key; used for all models by default | | `OPENAI_API_KEY` | Optional. Routes `openai/*` and `gpt-*`/`o*` models to the OpenAI API | -| `OPENAI_BASE_URL` | Optional. Override the OpenAI endpoint (e.g. Azure) | +| `OPENAI_BASE_URL` | Optional. An OpenAI-compatible endpoint, same as `--base-url` | -At least one key is required. Get an OpenRouter key at [openrouter.ai](https://openrouter.ai). +Set at least one key, or a base URL for a local server. Get an OpenRouter key at [openrouter.ai](https://openrouter.ai). ## Research references diff --git a/src/bin/cli.ts b/src/bin/cli.ts index ed6dcb1..7e7002f 100644 --- a/src/bin/cli.ts +++ b/src/bin/cli.ts @@ -55,6 +55,14 @@ function fail(message: string): never { process.exit(EXIT.noVerdict); } +function isHttpUrl(value: string): boolean { + try { + return ["http:", "https:"].includes(new URL(value).protocol); + } catch { + return false; + } +} + function readPromptFile(path: string): string { try { return readFileSync(path, "utf-8"); @@ -173,6 +181,10 @@ program "--openai-api-key ", "OpenAI API key (or set OPENAI_API_KEY); openai/* and gpt-* models route to the OpenAI API", ) + .option( + "--base-url ", + "OpenAI-compatible endpoint (or set OPENAI_BASE_URL), e.g. http://localhost:11434/v1", + ) .option( "--attacker-model ", `Model for the attacker agent (default: ${DEFAULT_MODELS.attacker})`, @@ -228,13 +240,18 @@ program const apiKey = options.apiKey || process.env.OPENROUTER_API_KEY; const openaiApiKey = options.openaiApiKey || process.env.OPENAI_API_KEY; - if (!apiKey && !openaiApiKey) { + const baseUrl = options.baseUrl || process.env.OPENAI_BASE_URL; + if (baseUrl && !isHttpUrl(baseUrl)) { + fail(`invalid --base-url "${baseUrl}"; expected an http(s) URL`); + } + if (!apiKey && !openaiApiKey && !baseUrl) { fail( - "no API key. Set OPENROUTER_API_KEY (--api-key) and/or OPENAI_API_KEY (--openai-api-key).", + "no API key. Set OPENROUTER_API_KEY (--api-key), OPENAI_API_KEY (--openai-api-key), or an OpenAI-compatible endpoint (--base-url).", ); } if (apiKey) process.env.OPENROUTER_API_KEY = apiKey; if (openaiApiKey) process.env.OPENAI_API_KEY = openaiApiKey; + if (baseUrl) process.env.OPENAI_BASE_URL = baseUrl; const mode = (options.mode || "dual") as | "extraction" @@ -271,6 +288,7 @@ program `${c.gray("Attacker")} ${options.attackerModel || DEFAULT_MODELS.attacker}`, `${c.gray("Target")} ${options.targetModel || DEFAULT_MODELS.target}`, `${c.gray("Evaluator")} ${options.evaluatorModel || DEFAULT_MODELS.evaluator}`, + ...(baseUrl ? [`${c.gray("Endpoint")} ${baseUrl}`] : []), ...(mode !== "extraction" ? [ `${c.gray("Probes")} ${maxProbes === 0 ? "all" : maxProbes}${ diff --git a/src/provider.ts b/src/provider.ts index c7e02f8..0212961 100644 --- a/src/provider.ts +++ b/src/provider.ts @@ -5,6 +5,8 @@ import type { LanguageModelV1 } from "ai"; export interface ProviderKeys { openrouterApiKey?: string; openaiApiKey?: string; + /** An OpenAI-compatible endpoint; defaults to `OPENAI_BASE_URL`. */ + openaiBaseUrl?: string; } /** @@ -23,28 +25,40 @@ export function isOpenAiModel(model: string): boolean { /** * Resolve a model id to an AI SDK model. * - * An OpenAI-style id routes to the OpenAI API when an OpenAI key is available - * (via `opts.openaiApiKey` or `OPENAI_API_KEY`); everything else — and OpenAI - * ids with no OpenAI key — goes through OpenRouter, preserving prior behavior. - * Set `OPENAI_BASE_URL` to target an OpenAI-compatible endpoint (Azure, a - * gateway, a local server). + * With `OPENAI_BASE_URL` set, the model goes to that OpenAI-compatible + * endpoint (OpenAI, Azure, a gateway, Ollama, vLLM, LM Studio...) when either: + * - the id is OpenAI-style (`openai/*`, `gpt-*`, `o1/o3/o4-*`), or + * - there is no OpenRouter key, so the endpoint is the only provider. + * A leading `openai/` is stripped, so `openai/llama3.1:8b` sends `llama3.1:8b` + * to the endpoint even when OpenRouter serves the other roles. Local servers + * don't need a key. + * + * Without a base URL, OpenAI-style ids go to the OpenAI API when an OpenAI key + * is available, and everything else goes through OpenRouter. */ export function resolveModel( model: string, opts: ProviderKeys = {}, ): LanguageModelV1 { const openaiApiKey = opts.openaiApiKey ?? process.env.OPENAI_API_KEY; + const openrouterApiKey = + opts.openrouterApiKey ?? process.env.OPENROUTER_API_KEY; + const baseURL = + opts.openaiBaseUrl ?? (process.env.OPENAI_BASE_URL || undefined); + + const useOpenAi = baseURL + ? isOpenAiModel(model) || !openrouterApiKey + : Boolean(openaiApiKey) && isOpenAiModel(model); - if (openaiApiKey && isOpenAiModel(model)) { + if (useOpenAi) { const openai = createOpenAI({ - apiKey: openaiApiKey, - baseURL: process.env.OPENAI_BASE_URL || undefined, + // The SDK throws on a missing key; keyless local servers ignore it. + apiKey: openaiApiKey ?? "", + baseURL, }); return openai(model.replace(/^openai\//, "")); } - const openrouter = createOpenRouter({ - apiKey: opts.openrouterApiKey ?? process.env.OPENROUTER_API_KEY, - }); + const openrouter = createOpenRouter({ apiKey: openrouterApiKey }); return openrouter(model) as LanguageModelV1; } diff --git a/test/e2e/harness.ts b/test/e2e/harness.ts index 7b845d1..d831b5c 100644 --- a/test/e2e/harness.ts +++ b/test/e2e/harness.ts @@ -47,6 +47,14 @@ interface RunOptions { runtime?: "bun" | "node"; /** A script to run with Bun in place of the CLI, such as LIBRARY_SCAN. */ script?: string; + /** + * How the child finds the mock: "env" sets OPENAI_API_KEY and + * OPENAI_BASE_URL, "flag" passes only `--base-url`, "none" leaves it to + * the caller's args. + */ + endpoint?: "env" | "flag" | "none"; + /** Extra environment variables for the child. */ + env?: Record; } const runs: CliRun[] = []; @@ -91,15 +99,24 @@ export async function runCli( name: string, scenario: Scenario, args: string[], - { runtime = "bun", script = CLI_PATH }: RunOptions = {}, + { + runtime = "bun", + script = CLI_PATH, + endpoint = "env", + env, + }: RunOptions = {}, ): Promise { const runDir = join(ARTIFACTS_DIR, slug(name)); mkdirSync(runDir, { recursive: true }); const reportPath = join(runDir, "result.json"); + const mock = startMockLlm(scenario); const setsOutput = args.includes("-o") || args.includes("--output"); - const cliArgs = - args[0] === "scan" && !setsOutput ? [...args, "-o", reportPath] : args; + const cliArgs = [ + ...args, + ...(args[0] === "scan" && !setsOutput ? ["-o", reportPath] : []), + ...(endpoint === "flag" ? ["--base-url", mock.url] : []), + ]; const command = runtime === "node" ? [ @@ -111,7 +128,6 @@ export async function runCli( ] : [process.execPath, "--no-env-file", script, ...cliArgs]; - const mock = startMockLlm(scenario); try { const child = Bun.spawn(command, { cwd: runDir, @@ -119,8 +135,10 @@ export async function runCli( PATH: process.env.PATH, HOME: process.env.HOME, NO_COLOR: "1", - OPENAI_API_KEY: "sk-mock", - OPENAI_BASE_URL: mock.url, + ...(endpoint === "env" + ? { OPENAI_API_KEY: "sk-mock", OPENAI_BASE_URL: mock.url } + : {}), + ...env, }, stdout: "pipe", stderr: "pipe", diff --git a/test/e2e/mock-llm.ts b/test/e2e/mock-llm.ts index f9b5d32..bdb77df 100644 --- a/test/e2e/mock-llm.ts +++ b/test/e2e/mock-llm.ts @@ -10,6 +10,14 @@ export const MOCK_MODELS = { judge: "gpt-mock-judge", } as const; +/** Ids that don't look like OpenAI's, as a local server would serve. */ +export const LOCAL_MOCK_MODELS = { + attacker: "llama-mock-attacker", + target: "llama-mock-target", + evaluator: "llama-mock-evaluator", + judge: "llama-mock-judge", +} as const; + type MockRole = keyof typeof MOCK_MODELS; type TargetBehavior = "refuse" | "comply" | "leak" | "error"; @@ -117,7 +125,10 @@ export function startMockLlm(scenario: Scenario): MockLlm { } function roleFor(model: string): MockRole | "unknown" { - const entry = Object.entries(MOCK_MODELS).find(([, id]) => id === model); + const entry = [ + ...Object.entries(MOCK_MODELS), + ...Object.entries(LOCAL_MOCK_MODELS), + ].find(([, id]) => id === model); return entry ? (entry[0] as MockRole) : "unknown"; } diff --git a/test/e2e/scan.test.ts b/test/e2e/scan.test.ts index a102a94..ce917b2 100644 --- a/test/e2e/scan.test.ts +++ b/test/e2e/scan.test.ts @@ -17,7 +17,7 @@ import { SYSTEM_PROMPT, writeSummary, } from "./harness"; -import type { GraderBehavior } from "./mock-llm"; +import { type GraderBehavior, LOCAL_MOCK_MODELS } from "./mock-llm"; setDefaultTimeout(30_000); beforeAll(resetArtifacts); @@ -528,9 +528,9 @@ describe("library callbacks", () => { describe("models", () => { test("the models named on screen are the models tested", async () => { - // No model flags: every role falls back to its default. Those defaults - // route to OpenRouter, which fails locally without a key, so the scan - // makes no network calls and ends inconclusive. + // No model flags: every role falls back to its default. With no + // OpenRouter key the defaults go to the mock endpoint, which doesn't + // serve them, so the scan ends inconclusive. const run = await runCli("default models", { target: "refuse" }, [ ...SCAN, "--mode", @@ -547,6 +547,75 @@ describe("models", () => { target: shown("Target"), evaluator: shown("Evaluator"), }); + const requested = new Set(run.requests.map((r) => r.model)); + expect(requested.size).toBeGreaterThan(0); + for (const model of requested) { + expect([shown("Target"), shown("Evaluator")]).toContain(model); + } + expect(run.exitCode).toBe(EXIT.inconclusive); + }); +}); + +const LOCAL_MODEL_FLAGS = (prefix = "") => [ + "--attacker-model", + `${prefix}${LOCAL_MOCK_MODELS.attacker}`, + "--target-model", + `${prefix}${LOCAL_MOCK_MODELS.target}`, + "--evaluator-model", + `${prefix}${LOCAL_MOCK_MODELS.evaluator}`, + "--injection-model", + `${prefix}${LOCAL_MOCK_MODELS.judge}`, +]; +const LOCAL_INJECTION = [ + ...SCAN, + "--mode", + "injection", + "--no-multi-turn", + "--max-probes", + "4", +]; + +describe("OpenAI-compatible endpoint", () => { + test("--base-url alone sends every model to the endpoint, with no key", async () => { + const run = await runCli( + "base-url flag, no keys", + { target: "refuse", judge: JUDGE_REFUSED }, + [...LOCAL_INJECTION, ...LOCAL_MODEL_FLAGS()], + { endpoint: "flag" }, + ); + + expect(run.stdout).toMatch(/Endpoint\s+http:\/\/127\.0\.0\.1/); + expect(run.requests.length).toBeGreaterThan(0); + expect(run.requests.every((r) => !r.failed)).toBe(true); + expect(run.result?.coverage.injection?.completed).toBe(4); + expect(run.exitCode).toBe(EXIT.secure); + }); + + test("with an OpenRouter key, openai/ ids still reach the endpoint", async () => { + const run = await runCli( + "base-url flag, openai/ prefix", + { target: "refuse", judge: JUDGE_REFUSED }, + [...LOCAL_INJECTION, ...LOCAL_MODEL_FLAGS("openai/")], + { endpoint: "flag", env: { OPENROUTER_API_KEY: "sk-or-unused" } }, + ); + + // The prefix is stripped before the request goes out. + expect(run.requests.length).toBeGreaterThan(0); + for (const request of run.requests) { + expect(Object.values(LOCAL_MOCK_MODELS)).toContain(request.model); + } + expect(run.exitCode).toBe(EXIT.secure); + }); + + test("a --base-url that isn't an http(s) URL is rejected before scanning", async () => { + const run = await runCli( + "invalid base-url", + { target: "refuse" }, + [...LOCAL_INJECTION, "--base-url", "localhost:11434"], + { endpoint: "none" }, + ); + + expect(run.stderr).toContain("invalid --base-url"); expect(run.requests).toHaveLength(0); expect(run.exitCode).toBe(EXIT.inconclusive); });