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