fix(omp): match usage source input and valid credential fallback

packages/server/src/server/agent/providers/omp/usage-reference.test.ts

@@ -89,6 +89,20 @@ test("multiple rows without sticky are ambiguous", () => {
   f.db.close();
   expect(f.read()).toBeNull();
 });
+test("fallback counts only valid OAuth rows", () => {
+  const f = fixture();
+  f.row(1);
+  f.row(2, "openai-codex", { expires: Date.now() - 1 });
+  f.row(3, "openai-codex", { disabled: true });
+  f.db
+    .prepare("INSERT INTO auth_credentials VALUES (?, ?, ?, ?, ?)")
+    .run(4, "openai-codex", "oauth", "invalid-json", null);
+  f.db.close();
+  expect(f.read()).toEqual({
+    source: "codex",
+    input: { accessToken: "fixture-token-1", accountId: "account-1" },
+  });
+});
 test("disabled and expired rows do not resolve", () => {
   const a = fixture();
   a.row(1, "openai-codex", { disabled: true });
@@ -129,6 +143,20 @@ test("broker environment and config URLs block attribution", () => {
   );
   expect(f.read()).toBeNull();
 });
+test("broker config.yaml and flat keys block attribution with config.yml precedence", () => {
+  const f = fixture();
+  f.row(1);
+  f.db.close();
+  writeFileSync(join(f.agentDir, "config.yaml"), '"auth.broker.url": https://broker.invalid\n');
+  expect(f.read()).toBeNull();
+  writeFileSync(join(f.agentDir, "config.yml"), "auth:\n  broker: {}\n");
+  expect(f.read()).toEqual({
+    source: "codex",
+    input: { accessToken: "fixture-token-1", accountId: "account-1" },
+  });
+  writeFileSync(join(f.agentDir, "config.yml"), '"auth.broker.url": https://broker.invalid\n');
+  expect(f.read()).toBeNull();
+});
 test("models config API key overrides sticky OAuth", () => {
   const f = fixture();
   f.row(1);
@@ -159,7 +187,7 @@ test("Anthropic OAuth maps to Claude and changes are read on each request", () =
     }),
   ).toEqual({
     source: "claude",
-    input: { accessToken: "fixture-token-1", accountId: "account-1" },
+    input: { accessToken: "fixture-token-1" },
   });
   writeFileSync(
     join(f.agentDir, "config.yml"),

packages/server/src/server/agent/providers/omp/usage-reference.ts

@@ -1,4 +1,4 @@
-import { readFileSync } from "node:fs";
+import { existsSync, readFileSync } from "node:fs";
 import { createRequire } from "node:module";
 import { join } from "node:path";
 
@@ -22,6 +22,7 @@ const oauthSchema = z.object({
 });
 const configSchema = z
   .object({
+    "auth.broker.url": z.string().optional(),
     auth: z
       .object({ broker: z.object({ url: z.string().optional() }).passthrough().optional() })
       .passthrough()
@@ -52,9 +53,16 @@ function readYaml(file: string): unknown {
 }
 
 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-broker/discover.ts:87-110,197-225,321-329.
+  const configFile = ["config.yml", "config.yaml"]
+    .map((name) => join(agentDir, name))
+    .find((file) => existsSync(file));
+  const config = configSchema.safeParse(configFile ? readYaml(configFile) : undefined);
+  if (
+    env.OMP_AUTH_BROKER_URL ||
+    (config.success && (config.data.auth?.broker?.url ?? 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;
@@ -82,7 +90,25 @@ function readCredential(db: UsageDb, provider: string, sessionId: string): unkno
       "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;
+  const validRows = rows.filter((row) => validOAuthRow(row, provider));
+  return validRows.length === 1 ? validRows[0] : null;
+}
+
+function validOAuthRow(value: unknown, provider: string): boolean {
+  const row = rowSchema.safeParse(value);
+  if (
+    !row.success ||
+    row.data.provider !== provider ||
+    row.data.credential_type !== "oauth" ||
+    row.data.disabled_cause !== null
+  )
+    return false;
+  try {
+    const oauth = oauthSchema.safeParse(JSON.parse(row.data.data));
+    return oauth.success && oauth.data.expires > Date.now();
+  } catch {
+    return false;
+  }
 }
 
 /** Read OMP's current credential attribution without changing its store. */
@@ -120,6 +146,9 @@ export function resolveOmpUsageReference(
     // 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;
+    if (source === "claude") {
+      return { source, input: { accessToken: oauth.data.access } };
+    }
     return {
       source,
       input: {