From a134873d470db9b0e810c1c1c93f870aebec08d5 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 29 Sep 2026 13:47:57 +0200 Subject: [PATCH] fix(usage): resolve saved provider accounts without sessions Default-model chats retain their effective model in saved runtime info, and the persistence handle retains the session ID needed for credential attribution. Pass both through the provider client so closed chats resolve usage after restart. --- .../server/src/server/agent/agent-manager.ts | 12 ++- .../src/server/agent/agent-sdk-types.ts | 2 + .../src/server/agent/provider-registry.ts | 5 +- .../server/agent/providers/omp/agent.test.ts | 26 ++++- .../src/server/agent/providers/omp/agent.ts | 26 +++-- .../providers/omp/test-utils/omp-harness.ts | 16 +++ .../agent/providers/opencode-agent.test.ts | 45 ++++++++ .../server/agent/providers/opencode-agent.ts | 6 +- .../server/agent/providers/pi/agent.test.ts | 29 ++++- .../src/server/agent/providers/pi/agent.ts | 18 ++-- .../session-usage-reference.e2e.test.ts | 101 +++++++++++++++++- 11 files changed, 256 insertions(+), 30 deletions(-) diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index a4b3543075..f9990954a2 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -1172,13 +1172,15 @@ export class AgentManager { id: string, ): Promise<{ found: boolean; reference: UsageReference | null }> { const agent = this.getAgent(id); - const stored = agent ? null : await this.registry?.get(id); - if (!agent && !stored) return { found: false, reference: null }; - const provider = agent?.provider ?? stored!.provider; - const config = agent?.config ?? buildStoredAgentConfig(stored!); + const stored = agent?.session ? null : await this.registry?.get(id); + const source = stored ?? agent; + if (!source) return { found: false, reference: null }; + const config = stored ? buildStoredAgentConfig(stored) : agent!.config; const reference = - (await this.clients.get(provider)?.resolveUsageReference?.({ + (await this.clients.get(source.provider)?.resolveUsageReference?.({ config, + runtimeInfo: source.runtimeInfo ?? null, + persistence: source.persistence ?? null, session: agent?.session ?? null, })) ?? null; return { found: true, reference }; diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index 11207836ad..b6387ac749 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -751,6 +751,8 @@ export interface AgentClient { /** Account used by the agent. Resolving this must not start or load a session. */ resolveUsageReference?(input: { config: AgentSessionConfig; + runtimeInfo: AgentRuntimeInfo | null; + persistence: AgentPersistenceHandle | null; session: AgentSession | null; }): Promise; createSession( diff --git a/packages/server/src/server/agent/provider-registry.ts b/packages/server/src/server/agent/provider-registry.ts index 38c4c24621..ff49b77afe 100644 --- a/packages/server/src/server/agent/provider-registry.ts +++ b/packages/server/src/server/agent/provider-registry.ts @@ -489,6 +489,7 @@ function wrapClientProvider( const listFeatures = inner.listFeatures?.bind(inner); const archiveNativeSession = inner.archiveNativeSession?.bind(inner); const unarchiveNativeSession = inner.unarchiveNativeSession?.bind(inner); + // An alias's live session can carry launch env that is fresher than stored config. const sessions = new WeakMap(); const wrapSession = (session: AgentSession) => { const wrapped = wrapSessionProvider(provider, session); @@ -500,9 +501,11 @@ function wrapClientProvider( provider, capabilities: inner.capabilities, resolveUsageReference: inner.resolveUsageReference - ? ({ config, session }) => + ? ({ config, runtimeInfo, persistence, session }) => inner.resolveUsageReference!({ config: { ...config, provider: inner.provider }, + runtimeInfo, + persistence, session: session ? (sessions.get(session) ?? null) : null, }) : undefined, 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 665dfee5ad..cfe8ba0bd9 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -194,14 +194,20 @@ describe("OMP agent client and session", () => { expires: Date.now() + 60_000, }), ); + db.prepare("INSERT INTO cache (key, value, expires_at) VALUES (?, ?, ?)").run( + "session:sticky:openai-codex:omp-session-1", + JSON.stringify({ type: "oauth", credentialId: 1 }), + Math.floor(Date.now() / 1000) + 100, + ); db.close(); try { - const omp = new OmpHarness(); - await omp.start({}, undefined, { + const env = { OMP_PROFILE: "", PI_CODING_AGENT_DIR: agentDir, XDG_DATA_HOME: "", - }); + }; + const omp = new OmpHarness({ env }); + await omp.start({}, undefined, env); const runtime = omp.runtime(); runtime.state = { ...runtime.state, model: { provider: "openai-codex", id: "codex-model" } }; const requestsBefore = runtime.getStateRequestCount; @@ -216,6 +222,20 @@ describe("OMP agent client and session", () => { input: { accessToken: "claude-fixture" }, }); expect(runtime.getStateRequestCount).toBe(requestsBefore + 2); + const beforeClosed = runtime.getStateRequestCount; + expect( + await omp.getStoredUsageReference("openai-codex/codex-model", "omp-session-1"), + ).toEqual({ + source: "codex", + input: { accessToken: "codex-fixture", accountId: "codex-account" }, + }); + expect( + await omp.getStoredUsageReference("openai-codex/codex-model", "omp-session-1", null), + ).toEqual({ + source: "codex", + input: { accessToken: "codex-fixture", accountId: "codex-account" }, + }); + expect(runtime.getStateRequestCount).toBe(beforeClosed); await omp.close(); } finally { rmSync(home, { recursive: true, force: true }); diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index aab87ddea6..d0d06ec93d 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -969,13 +969,8 @@ export class OmpAgentSession implements AgentSession { private readonly paseoTools?: PaseoToolCatalog; private readonly usageEnv: NodeJS.ProcessEnv; - async getUsageContext() { - const state = await this.runtimeSession.getState(); - return { - sessionId: state.sessionId, - provider: state.model?.provider ?? "", - env: this.usageEnv, - }; + getUsageEnvironment() { + return this.usageEnv; } get id(): string | null { @@ -2279,15 +2274,24 @@ export class OmpAgentClient implements AgentClient { } async resolveUsageReference({ + config, + runtimeInfo, + persistence, session, }: { config: AgentSessionConfig; + runtimeInfo: AgentRuntimeInfo | null; + persistence: AgentPersistenceHandle | null; session: AgentSession | null; }) { - // OMP's credential lookup needs the runtime session ID, which stored config does not contain. - if (!(session instanceof OmpAgentSession)) return null; - const { sessionId, provider, env } = await session.getUsageContext(); - return resolveOmpUsageReference(sessionId, provider, env); + const liveInfo = session ? await session.getRuntimeInfo() : null; + const model = liveInfo?.model ?? config.model ?? runtimeInfo?.model; + const sessionId = liveInfo?.sessionId ?? persistence?.sessionId; + const env = + session instanceof OmpAgentSession + ? session.getUsageEnvironment() + : { ...process.env, ...this.runtimeSettings?.env }; + return resolveOmpUsageReference(sessionId ?? "", model?.split("/")[0] ?? "", env); } private async configureNativePaseoTools( 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 0390a5905f..96f3badb13 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 @@ -71,11 +71,13 @@ export class OmpHarness { providerIdleScheduler?: OmpProviderIdleScheduler; noTurnScheduler?: OmpNoTurnScheduler; usagePollScheduler?: OmpUsagePollScheduler; + env?: Record; } = {}, ) { this.client = new OmpAgentClient({ logger: pino({ level: "silent" }), runtime: this.omp, + runtimeSettings: options.env ? { env: options.env } : undefined, providerIdleScheduler: options.providerIdleScheduler, noTurnScheduler: options.noTurnScheduler, usagePollScheduler: options.usagePollScheduler, @@ -490,6 +492,20 @@ export class OmpHarness { }); } + async getStoredUsageReference( + model: string, + sessionId: string, + configuredModel: string | null = model, + ) { + const input = { + config: { provider: "omp", cwd: CWD, model: configuredModel }, + runtimeInfo: { provider: "omp", sessionId, model }, + persistence: { provider: "omp", sessionId }, + session: null, + }; + return this.client.resolveUsageReference(input); + } + runningToolCallIds(): string[] { const statusByCall = new Map(); for (const item of this.timeline()) { diff --git a/packages/server/src/server/agent/providers/opencode-agent.test.ts b/packages/server/src/server/agent/providers/opencode-agent.test.ts index 191a87c160..ba9007732f 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -440,6 +440,11 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { source: "opencode-go", input: { apiKey: "go-key" }, }); + await session.setModel?.("openai/gpt-5"); + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }); await session.setModel?.("other/model"); expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toBeNull(); const events: AgentStreamEvent[] = []; @@ -477,6 +482,46 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { rmSync(cwd, { recursive: true, force: true }); }); + test("usage reference survives a model switch and default-model restart", async () => { + const cwd = tmpCwd(); + const client = new OpenCodeAgentClient(logger, { + env: { + OPENCODE_AUTH_CONTENT: JSON.stringify({ + openai: { type: "oauth", access: "oauth-token", accountId: "acct-1" }, + "opencode-go": { type: "api", key: "go-key" }, + }), + }, + }); + try { + const runtimeInfo = { + provider: "opencode", + sessionId: "saved-session", + model: "openai/gpt-5", + }; + const closed = { + config: { ...buildConfig(cwd), model: "openai/gpt-5" }, + runtimeInfo, + persistence: null, + session: null, + }; + expect(await client.resolveUsageReference(closed)).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }); + expect( + await client.resolveUsageReference({ + ...closed, + config: { ...buildConfig(cwd), model: null }, + }), + ).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + } + }); + test("usage reference ignores another auth entry with a different type", async () => { const cwd = tmpCwd(); const runtime = new TestOpenCodeHarness(); diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index 1bbb908f2d..f06f15b30d 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -1447,12 +1447,16 @@ export class OpenCodeAgentClient implements AgentClient { async resolveUsageReference({ config, + runtimeInfo, session, }: { config: AgentSessionConfig; + runtimeInfo: AgentRuntimeInfo | null; session: AgentSession | null; }): Promise { - const model = session ? (await session.getRuntimeInfo()).model : config.model; + const model = session + ? (await session.getRuntimeInfo()).model + : (config.model ?? runtimeInfo?.model); if (!model?.startsWith("openai/") && !model?.startsWith("opencode-go/")) return null; const env = session instanceof OpenCodeAgentSession diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 99770c16e9..917da68a79 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -273,7 +273,15 @@ test("Pi resolves a stored model without opening a session", async () => { try { writeFileSync( path.join(agentDir, "auth.json"), - JSON.stringify({ anthropic: { type: "oauth", access: "claude-token", refresh: "refresh" } }), + JSON.stringify({ + anthropic: { type: "oauth", access: "claude-token", refresh: "refresh" }, + "openai-codex": { + type: "oauth", + access: "codex-token", + accountId: "account-1", + refresh: "refresh", + }, + }), ); const client = new PiRpcAgentClient({ logger: pino({ level: "silent" }), @@ -286,6 +294,25 @@ test("Pi resolves a stored model without opening a session", async () => { session: null, }), ).toEqual({ source: "claude", input: { accessToken: "claude-token" } }); + const runtimeInfo = { + provider: "pi", + sessionId: "saved-session", + model: "openai-codex/gpt-5", + }; + const closed = { + config: createConfig({ model: "openai-codex/gpt-5" }), + runtimeInfo, + persistence: null, + session: null, + }; + expect(await client.resolveUsageReference(closed)).toEqual({ + source: "codex", + input: { accessToken: "codex-token", accountId: "account-1" }, + }); + expect(await client.resolveUsageReference({ ...closed, config: createConfig() })).toEqual({ + source: "codex", + input: { accessToken: "codex-token", accountId: "account-1" }, + }); } finally { rmSync(agentDir, { recursive: true, force: true }); } diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 4ccab930ba..a7b4d98df0 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1172,9 +1172,8 @@ export class PiRpcAgentSession implements AgentSession { readonly provider: AgentProvider; readonly capabilities: AgentCapabilityFlags; - async getUsageContext() { - await this.refreshState(); - return { provider: this.state.model?.provider, env: this.usageEnv }; + getUsageEnvironment() { + return this.usageEnv; } private readonly subscribers = new Set<(event: AgentStreamEvent) => void>(); @@ -2482,14 +2481,21 @@ export class PiRpcAgentClient implements AgentClient { async resolveUsageReference({ config, + runtimeInfo, session, }: { config: AgentSessionConfig; + runtimeInfo: AgentRuntimeInfo | null; session: AgentSession | null; }) { - const context = session instanceof PiRpcAgentSession ? await session.getUsageContext() : null; - const provider = context?.provider ?? config.model?.split("/")[0]; - return resolvePiUsageReference(provider, context?.env ?? this.runtimeSettings?.env); + const model = session + ? (await session.getRuntimeInfo()).model + : (config.model ?? runtimeInfo?.model); + const env = + session instanceof PiRpcAgentSession + ? session.getUsageEnvironment() + : this.runtimeSettings?.env; + return resolvePiUsageReference(model?.split("/")[0], env); } async createSession( diff --git a/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts b/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts index 1d7c7a47f3..eea23bf5bd 100644 --- a/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts +++ b/packages/server/src/server/plugins/session-usage-reference.e2e.test.ts @@ -1,8 +1,13 @@ import { fileURLToPath } from "node:url"; -import { mkdtemp, rm } from "node:fs/promises"; +import { randomUUID } from "node:crypto"; +import { mkdtemp, mkdir, rm, writeFile } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; +import pino from "pino"; import { expect, test } from "vitest"; +import { AgentStorage } from "../agent/agent-storage.js"; +import { OpenCodeAgentClient } from "../agent/providers/opencode-agent.js"; +import { PiRpcAgentClient } from "../agent/providers/pi/agent.js"; import { DaemonClient } from "../test-utils/daemon-client.js"; import { createTestAgentClient } from "../test-utils/fake-agent-client.js"; import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; @@ -104,8 +109,11 @@ test("stored Claude and Codex agents resolve without opening a provider session" sessionsOpened.set(provider, (sessionsOpened.get(provider) ?? 0) + 1); }, }); - providerClient.resolveUsageReference = async ({ session }) => { + providerClient.resolveUsageReference = async ({ session, runtimeInfo, persistence }) => { expect(session).toBeNull(); + expect(runtimeInfo?.sessionId).toBeTruthy(); + expect(runtimeInfo?.model).toBe(`${provider}/saved-model`); + expect(persistence?.sessionId).toBe(runtimeInfo?.sessionId); return { source: "fixture-session-usage", input: { account: "from-session" } }; }; return [provider, providerClient]; @@ -134,6 +142,19 @@ test("stored Claude and Codex agents resolve without opening a provider session" await client.close(); await daemon.close(); + const storage = new AgentStorage( + path.join(root, ".paseo", "agents"), + pino({ level: "silent" }), + ); + for (const [index, provider] of ["claude", "codex"].entries()) { + const record = await storage.get(ids[index]!); + expect(record).not.toBeNull(); + await storage.upsert({ + ...record!, + runtimeInfo: { ...record!.runtimeInfo!, model: `${provider}/saved-model` }, + }); + } + daemon = await createTestPaseoDaemon(options); client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); await client.connect(); @@ -153,3 +174,79 @@ test("stored Claude and Codex agents resolve without opening a provider session" await rm(staticDir, { recursive: true, force: true }); } }, 60_000); + +test("stored default-model OpenCode and Pi agents resolve their saved runtime model", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "paseo-usage-saved-model-")); + const staticDir = await mkdtemp(path.join(os.tmpdir(), "paseo-usage-static-")); + const piDir = path.join(root, "pi-agent"); + await mkdir(piDir); + await writeFile( + path.join(piDir, "auth.json"), + JSON.stringify({ + "openai-codex": { + type: "oauth", + access: "codex-fixture-token", + accountId: "fixture-account", + refresh: "fixture-refresh", + }, + }), + ); + const storage = new AgentStorage(path.join(root, ".paseo", "agents"), pino({ level: "silent" })); + const ids: string[] = []; + const now = new Date().toISOString(); + for (const provider of ["opencode", "pi"]) { + const id = randomUUID(); + ids.push(id); + await storage.upsert({ + id, + provider, + cwd: root, + createdAt: now, + updatedAt: now, + labels: {}, + lastStatus: "closed", + config: { model: null }, + runtimeInfo: { + provider, + sessionId: `${provider}-saved-session`, + model: provider === "opencode" ? "openai/gpt-5" : "openai-codex/gpt-5", + }, + persistence: { provider, sessionId: `${provider}-saved-session` }, + }); + } + const daemon = await createTestPaseoDaemon({ + paseoHomeRoot: root, + staticDir, + cleanup: false, + agentClients: { + opencode: new OpenCodeAgentClient(pino({ level: "silent" }), { + env: { + OPENCODE_AUTH_CONTENT: JSON.stringify({ + openai: { type: "oauth", access: "codex-fixture-token", accountId: "fixture-account" }, + }), + }, + }), + pi: new PiRpcAgentClient({ + logger: pino({ level: "silent" }), + runtimeSettings: { env: { PI_CODING_AGENT_DIR: piDir } }, + }), + }, + builtinPlugins: new BuiltinPluginLoader(undefined, ["codex-usage-source"]), + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + try { + await client.connect(); + await client.fetchAgents({ subscribe: {} }); + for (const [index, id] of ids.entries()) { + expect( + (await client.resolveAgentUsageReport({ agentId: id })).reportId, + ["opencode", "pi"][index], + ).not.toBeNull(); + } + } finally { + await client.close(); + await daemon.close(); + await rm(root, { recursive: true, force: true }); + await rm(staticDir, { recursive: true, force: true }); + } +}, 60_000);