From d0c24f491536c9e5a5e6f281b9786b566673f1ff Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Sat, 26 Sep 2026 20:03:21 +0200
Subject: [PATCH] test(omp): verify live session usage references

---
 .../server/agent/providers/omp/agent.test.ts  | 77 +++++++++++++++++++
 .../providers/omp/test-utils/omp-harness.ts   |  7 +-
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/packages/server/src/server/agent/providers/omp/agent.test.ts b/packages/server/src/server/agent/providers/omp/agent.test.ts
index 4c75aafc29..665dfee5ad 100644
--- a/packages/server/src/server/agent/providers/omp/agent.test.ts
+++ b/packages/server/src/server/agent/providers/omp/agent.test.ts
@@ -1,5 +1,9 @@
 import { describe, expect, test } from "vitest";
 import { setImmediate as waitForImmediate } from "node:timers/promises";
+import { mkdtempSync, mkdirSync, rmSync } from "node:fs";
+import { createRequire } from "node:module";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 
 import type { AgentStreamEvent } from "../../agent-sdk-types.js";
 import type { PaseoToolCatalog } from "../../tools/types.js";
@@ -144,6 +148,79 @@ function createToolCatalog(): PaseoToolCatalog {
 }
 
 describe("OMP agent client and session", () => {
+  test("usage reference reads current get_state provider and account on each call", async () => {
+    const home = mkdtempSync(join(tmpdir(), "paseo-omp-reference-"));
+    const agentDir = join(home, "agent");
+    mkdirSync(agentDir);
+    // Upstream: packages/ai/src/auth/sqlite-credential-store.ts:576-581,711-724.
+    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 db = new DatabaseSync(join(agentDir, "agent.db"));
+    db.exec(`
+      CREATE TABLE cache (key TEXT PRIMARY KEY, value TEXT NOT NULL, expires_at INTEGER NOT NULL);
+      CREATE INDEX idx_cache_expires ON cache(expires_at);
+      CREATE TABLE auth_credentials (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        provider TEXT NOT NULL,
+        credential_type TEXT NOT NULL,
+        data TEXT NOT NULL,
+        disabled_cause TEXT DEFAULT NULL,
+        identity_key TEXT DEFAULT NULL,
+        created_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s','now') AS INTEGER)),
+        updated_at INTEGER NOT NULL DEFAULT (CAST(strftime('%s','now') AS INTEGER))
+      );
+    `);
+    const insert = db.prepare(
+      "INSERT INTO auth_credentials (provider, credential_type, data) VALUES (?, 'oauth', ?)",
+    );
+    insert.run(
+      "openai-codex",
+      JSON.stringify({
+        access: "codex-fixture",
+        accountId: "codex-account",
+        expires: Date.now() + 60_000,
+      }),
+    );
+    insert.run(
+      "anthropic",
+      JSON.stringify({
+        access: "claude-fixture",
+        accountId: "claude-account",
+        expires: Date.now() + 60_000,
+      }),
+    );
+    db.close();
+    try {
+      const omp = new OmpHarness();
+      await omp.start({}, undefined, {
+        OMP_PROFILE: "",
+        PI_CODING_AGENT_DIR: agentDir,
+        XDG_DATA_HOME: "",
+      });
+      const runtime = omp.runtime();
+      runtime.state = { ...runtime.state, model: { provider: "openai-codex", id: "codex-model" } };
+      const requestsBefore = runtime.getStateRequestCount;
+      expect(await omp.getUsageReference()).toEqual({
+        source: "codex",
+        input: { accessToken: "codex-fixture", accountId: "codex-account" },
+      });
+      expect(runtime.getStateRequestCount).toBe(requestsBefore + 1);
+      runtime.state = { ...runtime.state, model: { provider: "anthropic", id: "claude-model" } };
+      expect(await omp.getUsageReference()).toEqual({
+        source: "claude",
+        input: { accessToken: "claude-fixture" },
+      });
+      expect(runtime.getStateRequestCount).toBe(requestsBefore + 2);
+      await omp.close();
+    } finally {
+      rmSync(home, { recursive: true, force: true });
+    }
+  });
   test("owns launch configuration and registers native host tools", async () => {
     const omp = new OmpHarness();
     await omp.start({ modeId: "ask" }, createToolCatalog());
diff --git a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts
index 4f2bae8838..7c99fa772c 100644
--- a/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts
+++ b/packages/server/src/server/agent/providers/omp/test-utils/omp-harness.ts
@@ -93,10 +93,11 @@ export class OmpHarness {
   async start(
     config: Partial<AgentSessionConfig> = {},
     paseoTools?: PaseoToolCatalog,
+    env?: Record<string, string>,
   ): Promise<void> {
     const session = await this.client.createSession(
       { provider: "omp", cwd: CWD, ...config },
-      paseoTools ? { paseoTools } : undefined,
+      paseoTools || env ? { paseoTools, env } : undefined,
     );
     if (!(session instanceof OmpAgentSession)) {
       throw new Error("OMP client returned a non-OMP session");
@@ -482,6 +483,10 @@ export class OmpHarness {
     return this.omp.latestSession();
   }
 
+  async getUsageReference() {
+    return this.requireSession().getUsageReference();
+  }
+
   runningToolCallIds(): string[] {
     const statusByCall = new Map<string, string>();
     for (const item of this.timeline()) {