feat(omp): resolve session usage from active OAuth credential

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) {

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" } });
+});

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();
+  }
+}