diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index 954c8b286a..5555fa130e 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -108,6 +108,7 @@ import { mapOmpRpcUiPermissionRequest, } from "./rpc-ui-permission-mapper.js"; import { DEFAULT_OMP_THINKING_LEVEL, mapOmpModel } from "./map-omp-model.js"; +import { resolveOmpUsageReference } from "./usage-reference.js"; const OMP_PROVIDER = "omp"; const QUESTION_RESPONSE_HEADER = "Response"; @@ -186,6 +187,7 @@ interface OmpAgentSessionOptions { noTurnScheduler?: OmpNoTurnScheduler; usagePollScheduler?: OmpUsagePollScheduler; paseoTools?: PaseoToolCatalog; + usageEnv?: NodeJS.ProcessEnv; /** * When false (resumed sessions), replayed session events are dropped until * the first prompt or agent_start so history is not re-emitted as live @@ -916,6 +918,7 @@ export class OmpAgentSession implements AgentSession { this.currentModeId = options.currentModeId ?? null; this.logger = options.logger; this.paseoTools = options.paseoTools; + this.usageEnv = options.usageEnv ?? process.env; this.live = options.live ?? true; this.providerIdleScheduler = options.providerIdleScheduler ?? createOmpProviderIdleScheduler(); this.noTurnScheduler = options.noTurnScheduler ?? createOmpNoTurnScheduler(); @@ -964,6 +967,12 @@ export class OmpAgentSession implements AgentSession { private readonly config: AgentSessionConfig; private readonly logger: Logger; private readonly paseoTools?: PaseoToolCatalog; + private readonly usageEnv: NodeJS.ProcessEnv; + + async getUsageReference() { + const state = await this.runtimeSession.getState(); + return resolveOmpUsageReference(state.sessionId, state.model?.provider ?? "", this.usageEnv); + } get id(): string | null { return this.state.sessionId; @@ -2304,6 +2313,7 @@ export class OmpAgentClient implements AgentClient { noTurnScheduler: this.noTurnScheduler, usagePollScheduler: this.usagePollScheduler, paseoTools: launchContext?.paseoTools, + usageEnv: { ...process.env, ...this.runtimeSettings?.env, ...launchContext?.env }, }); } catch (error) { await runtimeSession.close().catch(() => undefined); @@ -2346,6 +2356,7 @@ export class OmpAgentClient implements AgentClient { noTurnScheduler: this.noTurnScheduler, usagePollScheduler: this.usagePollScheduler, paseoTools: launchContext?.paseoTools, + usageEnv: { ...process.env, ...this.runtimeSettings?.env, ...launchContext?.env }, live: false, }); } catch (error) { diff --git a/packages/server/src/server/agent/providers/omp/usage-reference.test.ts b/packages/server/src/server/agent/providers/omp/usage-reference.test.ts new file mode 100644 index 0000000000..6a9a6c70a3 --- /dev/null +++ b/packages/server/src/server/agent/providers/omp/usage-reference.test.ts @@ -0,0 +1,206 @@ +import { mkdtempSync, readFileSync, rmSync, mkdirSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { createRequire } from "node:module"; +import { afterEach, expect, test } from "vitest"; +import { resolveOmpUsageReference } from "./usage-reference.js"; + +const { DatabaseSync } = createRequire(import.meta.url)("node:sqlite") as { + DatabaseSync: new (path: string) => { + exec(sql: string): void; + prepare(sql: string): { run(...args: unknown[]): void }; + close(): void; + }; +}; +const dirs: string[] = []; +const future = Date.now() + 3_600_000; +function fixture() { + const home = mkdtempSync(join(tmpdir(), "omp-usage-")); + dirs.push(home); + const agentDir = join(home, "agent"); + mkdirSync(agentDir); + const dbFile = join(agentDir, "agent.db"); + const db = new DatabaseSync(dbFile); + db.exec( + "CREATE TABLE cache (key TEXT PRIMARY KEY, value TEXT, expires_at INTEGER); CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT)", + ); + const env = { OMP_PROFILE: "", PI_CODING_AGENT_DIR: agentDir, XDG_DATA_HOME: "" }; + const row = ( + id: number, + provider = "openai-codex", + options: { disabled?: boolean; expires?: number; type?: string } = {}, + ) => + db.prepare("INSERT INTO auth_credentials VALUES (?, ?, ?, ?, ?)").run( + id, + provider, + options.type ?? "oauth", + JSON.stringify({ + access: `fixture-token-${id}`, + expires: options.expires ?? future, + accountId: `account-${id}`, + }), + options.disabled ? "disabled" : null, + ); + const sticky = (id: number, expires = Math.floor(Date.now() / 1000) + 100) => + db + .prepare("INSERT INTO cache VALUES (?, ?, ?)") + .run( + "session:sticky:openai-codex:session-1", + JSON.stringify({ type: "oauth", index: 0, credentialId: id }), + expires, + ); + const read = (override = {}) => + resolveOmpUsageReference("session-1", "openai-codex", { ...env, ...override }); + return { db, dbFile, agentDir, row, sticky, read }; +} +afterEach(() => { + for (const dir of dirs.splice(0)) rmSync(dir, { recursive: true, force: true }); +}); +test("sticky id selects the right OAuth row and leaves SQLite byte-identical", () => { + const f = fixture(); + f.row(1); + f.row(2); + f.sticky(2); + f.db.close(); + const before = readFileSync(f.dbFile); + expect(f.read()).toEqual({ + source: "codex", + input: { accessToken: "fixture-token-2", accountId: "account-2" }, + }); + expect(f.read()).toEqual({ + source: "codex", + input: { accessToken: "fixture-token-2", accountId: "account-2" }, + }); + expect(readFileSync(f.dbFile)).toEqual(before); +}); +test("single OAuth row resolves without sticky", () => { + const f = fixture(); + f.row(1); + f.db.close(); + expect(f.read()).toEqual({ + source: "codex", + input: { accessToken: "fixture-token-1", accountId: "account-1" }, + }); +}); +test("multiple rows without sticky are ambiguous", () => { + const f = fixture(); + f.row(1); + f.row(2); + f.db.close(); + expect(f.read()).toBeNull(); +}); +test("disabled and expired rows do not resolve", () => { + const a = fixture(); + a.row(1, "openai-codex", { disabled: true }); + a.sticky(1); + a.db.close(); + expect(a.read()).toBeNull(); + const b = fixture(); + b.row(1, "openai-codex", { expires: Date.now() - 1 }); + b.db.close(); + expect(b.read()).toBeNull(); +}); +test("expired sticky falls back to a single row", () => { + const f = fixture(); + f.row(1); + f.sticky(1, 1); + f.db.close(); + expect(f.read()?.source).toBe("codex"); +}); +test("env key does not displace OAuth with or without sticky", () => { + const f = fixture(); + f.row(1); + f.db.close(); + expect(f.read({ OPENAI_CODEX_OAUTH_TOKEN: "fixture-key" })?.source).toBe("codex"); + const g = fixture(); + g.row(1); + g.sticky(1); + g.db.close(); + expect(g.read({ OPENAI_CODEX_OAUTH_TOKEN: "fixture-key" })?.source).toBe("codex"); +}); +test("broker environment and config URLs block attribution", () => { + const f = fixture(); + f.row(1); + f.db.close(); + expect(f.read({ OMP_AUTH_BROKER_URL: "https://broker.invalid" })).toBeNull(); + writeFileSync( + join(f.agentDir, "config.yml"), + "auth:\n broker:\n url: https://broker.invalid\n", + ); + expect(f.read()).toBeNull(); +}); +test("models config API key overrides sticky OAuth", () => { + const f = fixture(); + f.row(1); + f.sticky(1); + f.db.close(); + writeFileSync( + join(f.agentDir, "models.yml"), + "providers:\n openai-codex:\n apiKey: fixture-key\n", + ); + expect(f.read()).toBeNull(); +}); +test("provider mismatch and malformed sticky fail closed", () => { + const f = fixture(); + f.row(1, "anthropic"); + f.sticky(1); + f.db.close(); + expect(f.read()).toBeNull(); +}); + +test("Anthropic OAuth maps to Claude and changes are read on each request", () => { + const f = fixture(); + f.row(1, "anthropic"); + f.db.close(); + expect( + resolveOmpUsageReference("session-1", "anthropic", { + OMP_PROFILE: "", + PI_CODING_AGENT_DIR: f.agentDir, + }), + ).toEqual({ + source: "claude", + input: { accessToken: "fixture-token-1", accountId: "account-1" }, + }); + writeFileSync( + join(f.agentDir, "config.yml"), + "auth:\n broker:\n url: https://broker.invalid\n", + ); + expect( + resolveOmpUsageReference("session-1", "anthropic", { + OMP_PROFILE: "", + PI_CODING_AGENT_DIR: f.agentDir, + }), + ).toBeNull(); +}); + +test("named profile ignores agent override and uses existing XDG profile data", () => { + const f = fixture(); + f.db.close(); + const root = mkdtempSync(join(tmpdir(), "omp-profile-")); + dirs.push(root); + const data = join(root, "xdg", "omp", "profiles", "work"); + mkdirSync(data, { recursive: true }); + const profileDb = new DatabaseSync(join(data, "agent.db")); + profileDb.exec( + "CREATE TABLE cache (key TEXT PRIMARY KEY, value TEXT, expires_at INTEGER); CREATE TABLE auth_credentials (id INTEGER PRIMARY KEY, provider TEXT, credential_type TEXT, data TEXT, disabled_cause TEXT)", + ); + profileDb + .prepare("INSERT INTO auth_credentials VALUES (?, ?, ?, ?, ?)") + .run( + 1, + "openai-codex", + "oauth", + JSON.stringify({ access: "profile-token", expires: future }), + null, + ); + profileDb.close(); + expect( + resolveOmpUsageReference("session-1", "openai-codex", { + OMP_PROFILE: "work", + PI_PROFILE: "other", + PI_CODING_AGENT_DIR: f.agentDir, + PI_CONFIG_DIR: join(root, "config"), + XDG_DATA_HOME: join(root, "xdg"), + }), + ).toEqual({ source: "codex", input: { accessToken: "profile-token" } }); +}); diff --git a/packages/server/src/server/agent/providers/omp/usage-reference.ts b/packages/server/src/server/agent/providers/omp/usage-reference.ts new file mode 100644 index 0000000000..ebd2260e6f --- /dev/null +++ b/packages/server/src/server/agent/providers/omp/usage-reference.ts @@ -0,0 +1,135 @@ +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { join } from "node:path"; + +import { z } from "zod"; +import type { UsageReference } from "../../agent-sdk-types.js"; +import { resolveOmpDiagnosticPaths } from "./provider-config.js"; + +const stickySchema = z.object({ type: z.string(), credentialId: z.number().int().positive() }); +const rowSchema = z.object({ + id: z.number().int().positive(), + provider: z.string(), + credential_type: z.string(), + data: z.string(), + disabled_cause: z.string().nullable(), +}); +const oauthSchema = z.object({ + access: z.string().min(1), + expires: z.number(), + accountId: z.string().optional(), + orgId: z.string().optional(), +}); +const configSchema = z + .object({ + auth: z + .object({ broker: z.object({ url: z.string().optional() }).passthrough().optional() }) + .passthrough() + .optional(), + }) + .passthrough(); +const modelsSchema = z + .object({ + providers: z + .record(z.string(), z.object({ apiKey: z.unknown().optional() }).passthrough()) + .optional(), + }) + .passthrough(); + +interface UsageDb { + prepare(sql: string): { get(...args: unknown[]): unknown; all(...args: unknown[]): unknown[] }; + close(): void; +} +const localRequire = createRequire(import.meta.url); +const yaml = localRequire("js-yaml") as { load(value: string): unknown }; + +function readYaml(file: string): unknown { + try { + return yaml.load(readFileSync(file, "utf8")); + } catch { + return undefined; + } +} + +function hasKeyOverride(provider: string, env: NodeJS.ProcessEnv, agentDir: string): boolean { + // Upstream: packages/coding-agent/src/session/auth-broker-config.ts:9-16. + const config = configSchema.safeParse(readYaml(join(agentDir, "config.yml"))); + if (env.OMP_AUTH_BROKER_URL || (config.success && config.data.auth?.broker?.url)) return true; + // Upstream: packages/ai/src/auth/cascade.ts:22-84. + const models = modelsSchema.safeParse(readYaml(join(agentDir, "models.yml"))); + return models.success && models.data.providers?.[provider]?.apiKey !== undefined; +} + +function readCredential(db: UsageDb, provider: string, sessionId: string): unknown { + // Upstream: packages/ai/src/auth/affinity.ts:83-88,115-132; sqlite-credential-store.ts:447-448,711-722. + const cache = db + .prepare( + "SELECT value FROM cache WHERE key = ? AND expires_at > CAST(strftime('%s','now') AS INTEGER)", + ) + .get(`session:sticky:${provider}:${sessionId}`) as { value?: unknown } | undefined; + if (cache) { + const sticky = stickySchema.safeParse(JSON.parse(String(cache.value))); + if (!sticky.success || sticky.data.type !== "oauth") return null; + return db + .prepare( + "SELECT id, provider, credential_type, data, disabled_cause FROM auth_credentials WHERE id = ?", + ) + .get(sticky.data.credentialId); + } + // Upstream: packages/ai/src/auth/cascade.ts:276-345 selects OAuth before env or stored API keys. + const rows = db + .prepare( + "SELECT id, provider, credential_type, data, disabled_cause FROM auth_credentials WHERE provider = ? AND credential_type = 'oauth' AND disabled_cause IS NULL", + ) + .all(provider); + return rows.length === 1 ? rows[0] : null; +} + +/** Read OMP's current credential attribution without changing its store. */ +export function resolveOmpUsageReference( + sessionId: string, + provider: string, + env: NodeJS.ProcessEnv, +): UsageReference | null { + let source: string; + if (provider === "openai-codex") source = "codex"; + else if (provider === "anthropic") source = "claude"; + else return null; + if (!sessionId) return null; + // Upstream: packages/utils/src/dirs.ts:82-94,330-343,583-585 (profile, override, agent dir). + const paths = resolveOmpDiagnosticPaths(env); + if (hasKeyOverride(provider, env, paths.agentDir)) return null; + let db: UsageDb; + try { + const sqlite = localRequire("node:sqlite") as { + DatabaseSync: new (path: string, options: { readOnly: boolean }) => UsageDb; + }; + db = new sqlite.DatabaseSync(paths.agentDb, { readOnly: true }); + } catch { + return null; + } + try { + const row = rowSchema.safeParse(readCredential(db, provider, sessionId)); + if ( + !row.success || + row.data.provider !== provider || + row.data.credential_type !== "oauth" || + row.data.disabled_cause !== null + ) + return null; + // Upstream: packages/ai/src/auth/sqlite-credential-store.ts:117-155 stores OAuth data as JSON without type. + const oauth = oauthSchema.safeParse(JSON.parse(row.data.data)); + if (!oauth.success || oauth.data.expires <= Date.now()) return null; + return { + source, + input: { + accessToken: oauth.data.access, + ...(oauth.data.accountId ? { accountId: oauth.data.accountId } : {}), + }, + }; + } catch { + return null; + } finally { + db.close(); + } +}