diff --git a/packages/app/e2e/browser/usage-composer-pill.spec.ts b/packages/app/e2e/browser/usage-composer-pill.spec.ts index 42df720db7..36e99f340c 100644 --- a/packages/app/e2e/browser/usage-composer-pill.spec.ts +++ b/packages/app/e2e/browser/usage-composer-pill.spec.ts @@ -1,7 +1,11 @@ import type { UsageReportEntry } from "@getpaseo/protocol/messages"; import { expect, test, type Page } from "../support/fixtures"; import { expectComposerVisible, submitMessage } from "../support/helpers/composer"; -import { openAgentRoute, seedMockAgentWorkspace } from "../support/helpers/mock-agent"; +import { + openAgentRoute, + seedMockAgentWorkspace, + seedRunningMockAgentWorkspace, +} from "../support/helpers/mock-agent"; import { installUsageReportsFixture } from "../support/helpers/usage-reports"; const REPORT_ID = "fixture:fixture-account"; @@ -42,6 +46,30 @@ async function openMockAgent(page: Page) { } test.describe("usage composer pill", () => { + test("shows the report when an existing running chat opens", async ({ page }) => { + test.setTimeout(180_000); + const usage = await installUsageReportsFixture(page, { + agentReportIds: [REPORT_ID], + lists: [[agentEntry(42)]], + }); + const session = await seedRunningMockAgentWorkspace({ + repoPrefix: "usage-running-", + title: "Running usage pill", + initialPrompt: "Show usage during the turn.", + model: "one-minute-stream", + }); + try { + await openAgentRoute(page, session); + await expectComposerVisible(page); + await expect(page.getByTestId("usage-composer-pill")).toContainText("42%", { + timeout: 20_000, + }); + expect(usage.agentRequests()[0]).toEqual({ agentId: session.agentId }); + } finally { + await session.cleanup(); + } + }); + test("opens the cached report without a request and shows when it was fetched", async ({ page, }) => { diff --git a/packages/app/src/usage/queries.ts b/packages/app/src/usage/queries.ts index a7a717b627..3c3a6b7f9a 100644 --- a/packages/app/src/usage/queries.ts +++ b/packages/app/src/usage/queries.ts @@ -31,8 +31,13 @@ function usageReportsQueryKey(serverId: string) { return ["usage", "reports", serverId] as const; } -function agentUsageQueryKey(serverId: string, agentId: string, model: string | null) { - return ["usage", "agent", serverId, agentId, model] as const; +function agentUsageQueryKey( + serverId: string, + agentId: string, + model: string | null, + isRunning: boolean, +) { + return ["usage", "agent", serverId, agentId, model, isRunning] as const; } function requireClient(serverId: string) { @@ -160,8 +165,7 @@ export function useUsageByHost(): { /** * The usage report for the account an agent is spending. The daemon resolves the * report ID, which is re-resolved when the model changes (a new key) and when a - * turn completes (the query is paused while the agent runs, and its data is always - * stale, so resuming refetches). The report itself is read from the shared cache. + * turn completes (a new query key). The report itself is read from the shared cache. */ export function useAgentUsage( serverId: string, @@ -176,10 +180,10 @@ export function useAgentUsage( }), ); const identity = useFetchQuery({ - queryKey: agentUsageQueryKey(serverId, agentId, model), + queryKey: agentUsageQueryKey(serverId, agentId, model, isRunning), queryFn: async () => (await requireClient(serverId).resolveAgentUsageReport({ agentId })).reportId, - enabled: isConnected && isSupported && !isRunning, + enabled: isConnected && isSupported, dataShape: "value", staleTimeMs: 0, }); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 87d3b290aa..a4b3543075 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -50,6 +50,7 @@ import { type AgentTimelineItem, type AgentUsage, type AgentRuntimeInfo, + type UsageReference, type ImportedTimelineEntry, type ImportableProviderSession, type ListImportableSessionsOptions, @@ -1167,6 +1168,22 @@ export class AgentManager { return agent ? { ...agent } : null; } + async resolveUsageReference( + 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 reference = + (await this.clients.get(provider)?.resolveUsageReference?.({ + config, + session: agent?.session ?? null, + })) ?? null; + return { found: true, reference }; + } + async waitForAgentClose(agentId: string): Promise { // Loading during reload must wait for the replacement, not resume another writer. await this.lifecycleMutationTails.get(agentId); diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index d42d98aaf7..11207836ad 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -673,8 +673,6 @@ export interface AgentSession { /** New provider-owned rows to commit on registration. streamHistory must also * replay them at their original timestamps; restored sessions omit old rows. */ readonly initialTimeline?: ImportedTimelineEntry[]; - /** Resolved at fetch time because model and credentials may change during a session. */ - getUsageReference?(): Promise; run(prompt: AgentPromptInput, options?: AgentRunOptions): Promise; startTurn(prompt: AgentPromptInput, options?: AgentRunOptions): Promise<{ turnId: string }>; steerActiveTurn?(prompt: AgentPromptInput, options: SteerActiveTurnOptions): Promise; @@ -750,6 +748,11 @@ export interface ResolveAgentDefaultModeInput { export interface AgentClient { readonly provider: AgentProvider; readonly capabilities: AgentCapabilityFlags; + /** Account used by the agent. Resolving this must not start or load a session. */ + resolveUsageReference?(input: { + config: AgentSessionConfig; + session: AgentSession | null; + }): Promise; createSession( config: AgentSessionConfig, launchContext?: AgentLaunchContext, diff --git a/packages/server/src/server/agent/plugin-provider.test.ts b/packages/server/src/server/agent/plugin-provider.test.ts index 8774b8784c..e1be6b8c8b 100644 --- a/packages/server/src/server/agent/plugin-provider.test.ts +++ b/packages/server/src/server/agent/plugin-provider.test.ts @@ -406,7 +406,12 @@ describe("PluginAgentClientRegistry", () => { cwd: "/workspace", }); try { - expect(await session.getUsageReference?.()).toBeNull(); + expect( + await client.resolveUsageReference?.({ + config: { provider: harness.registration.id, cwd: "/workspace" }, + session, + }), + ).toBeNull(); expect(harness.inputs.some((input) => input.type === "session.usage_reference")).toBe(false); } finally { await session.close(); diff --git a/packages/server/src/server/agent/plugin-provider.ts b/packages/server/src/server/agent/plugin-provider.ts index e6712d3c3b..1ba6bcce47 100644 --- a/packages/server/src/server/agent/plugin-provider.ts +++ b/packages/server/src/server/agent/plugin-provider.ts @@ -897,6 +897,15 @@ class PluginAgentClient implements AgentClient { return agentCapabilities(this.runtime.negotiatedCapabilities); } + async resolveUsageReference({ + session, + }: { + config: AgentSessionConfig; + session: AgentSession | null; + }) { + return session instanceof PluginAgentSession ? session.getPluginUsageReference() : null; + } + async createSession( config: AgentSessionConfig, launchContext?: AgentLaunchContext, @@ -1096,7 +1105,7 @@ class PluginAgentSession implements AgentSession { return agentCapabilities(this.bridge.negotiatedCapabilities); } - getUsageReference(): Promise { + getPluginUsageReference(): Promise { return this.bridge.getUsageReference(); } diff --git a/packages/server/src/server/agent/provider-registry.ts b/packages/server/src/server/agent/provider-registry.ts index 50ad1ae94d..38c4c24621 100644 --- a/packages/server/src/server/agent/provider-registry.ts +++ b/packages/server/src/server/agent/provider-registry.ts @@ -458,7 +458,6 @@ export function wrapSessionProvider(provider: AgentProvider, inner: AgentSession } }, getRuntimeInfo: async () => mapRuntimeInfo(provider, await inner.getRuntimeInfo()), - getUsageReference: inner.getUsageReference?.bind(inner), getAvailableModes: () => inner.getAvailableModes(), getCurrentMode: () => inner.getCurrentMode(), setMode: (modeId) => inner.setMode(modeId), @@ -490,13 +489,25 @@ function wrapClientProvider( const listFeatures = inner.listFeatures?.bind(inner); const archiveNativeSession = inner.archiveNativeSession?.bind(inner); const unarchiveNativeSession = inner.unarchiveNativeSession?.bind(inner); + const sessions = new WeakMap(); + const wrapSession = (session: AgentSession) => { + const wrapped = wrapSessionProvider(provider, session); + sessions.set(wrapped, session); + return wrapped; + }; return { provider, capabilities: inner.capabilities, + resolveUsageReference: inner.resolveUsageReference + ? ({ config, session }) => + inner.resolveUsageReference!({ + config: { ...config, provider: inner.provider }, + session: session ? (sessions.get(session) ?? null) : null, + }) + : undefined, createSession: async (config, launchContext) => - wrapSessionProvider( - provider, + wrapSession( await inner.createSession( { ...config, @@ -506,8 +517,7 @@ function wrapClientProvider( ), ), resumeSession: async (handle, overrides, launchContext, options) => - wrapSessionProvider( - provider, + wrapSession( await inner.resumeSession( { ...handle, @@ -569,7 +579,7 @@ function wrapClientProvider( } return { ...imported, - session: wrapSessionProvider(provider, imported.session), + session: wrapSession(imported.session), config: { ...imported.config, provider, diff --git a/packages/server/src/server/agent/providers/acp-agent.test.ts b/packages/server/src/server/agent/providers/acp-agent.test.ts index 0394900d3b..f39d8e723e 100644 --- a/packages/server/src/server/agent/providers/acp-agent.test.ts +++ b/packages/server/src/server/agent/providers/acp-agent.test.ts @@ -224,8 +224,17 @@ function createSessionWithConfig( test("ACP usage reference uses the provider ID", async () => { for (const provider of ["copilot", "cursor", "kimi", "custom-source"]) { - const session = createSessionWithConfig({ provider }); - expect(await session.getUsageReference()).toEqual({ source: provider, input: {} }); + const client = new ACPAgentClient({ + provider, + logger: createTestLogger(), + defaultCommand: ["unused"], + }); + expect( + await client.resolveUsageReference({ + config: { provider, cwd: "/tmp/paseo-acp-test" }, + session: null, + }), + ).toEqual({ source: provider, input: {} }); } }); diff --git a/packages/server/src/server/agent/providers/acp-agent.ts b/packages/server/src/server/agent/providers/acp-agent.ts index 019725ce3f..acad3df87a 100644 --- a/packages/server/src/server/agent/providers/acp-agent.ts +++ b/packages/server/src/server/agent/providers/acp-agent.ts @@ -882,6 +882,9 @@ function isACPCreateConfigUnattended(input: AgentCreateConfigUnattendedInput): b export class ACPAgentClient implements AgentClient { readonly provider: string; readonly capabilities: AgentCapabilityFlags; + async resolveUsageReference() { + return { source: this.provider, input: {} }; + } readonly resolveCreateConfig = resolveACPCreateConfig; readonly isCreateConfigUnattended = isACPCreateConfigUnattended; @@ -1633,10 +1636,6 @@ export class ACPAgentSession implements AgentSession, ACPClient { readonly provider: string; readonly capabilities: AgentCapabilityFlags; - async getUsageReference() { - return { source: this.provider, input: {} }; - } - private readonly logger: Logger; private readonly runtimeSettings?: ProviderRuntimeSettings; private readonly defaultCommand: [string, ...string[]]; diff --git a/packages/server/src/server/agent/providers/claude/agent.test.ts b/packages/server/src/server/agent/providers/claude/agent.test.ts index 16aa193cce..f148f53101 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -531,16 +531,16 @@ describe("ClaudeAgentClient binary resolution", () => { const logger = createTestLogger(); test("Claude usage reference follows CLAUDE_CONFIG_DIR and excludes API overrides", async () => { - const client = new ClaudeAgentClient({ logger, resolveBinary: async () => "/test/claude/bin" }); - const session = await client.createSession( - { provider: "claude", cwd: process.cwd() }, - { env: { CLAUDE_CONFIG_DIR: "/accounts/second" } }, - ); - expect(await session.getUsageReference?.()).toEqual({ + const config = { provider: "claude", cwd: process.cwd() }; + const client = new ClaudeAgentClient({ + logger, + runtimeSettings: { env: { CLAUDE_CONFIG_DIR: "/accounts/second" } }, + resolveBinary: async () => "/test/claude/bin", + }); + expect(await client.resolveUsageReference?.({ config, session: null })).toEqual({ source: "claude", input: { configDir: "/accounts/second" }, }); - await session.close(); }); test.each(["ANTHROPIC_BASE_URL", "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"])( @@ -548,14 +548,15 @@ describe("ClaudeAgentClient binary resolution", () => { async (name) => { const client = new ClaudeAgentClient({ logger, + runtimeSettings: { env: { [name]: "override" } }, resolveBinary: async () => "/test/claude/bin", }); - const session = await client.createSession( - { provider: "claude", cwd: process.cwd() }, - { env: { [name]: "override" } }, - ); - expect(await session.getUsageReference?.()).toBeNull(); - await session.close(); + expect( + await client.resolveUsageReference?.({ + config: { provider: "claude", cwd: process.cwd() }, + session: null, + }), + ).toBeNull(); }, ); @@ -569,15 +570,16 @@ describe("ClaudeAgentClient binary resolution", () => { }, }, }); - const session = await registry["work-claude"].createClient(logger).createSession({ - provider: "work-claude", - cwd: process.cwd(), - }); - expect(await session.getUsageReference?.()).toEqual({ + const client = registry["work-claude"].createClient(logger); + expect( + await client.resolveUsageReference?.({ + config: { provider: "work-claude", cwd: process.cwd() }, + session: null, + }), + ).toEqual({ source: "claude", input: { configDir: "/accounts/work" }, }); - await session.close(); }); test("resolves the installed Claude Code version", async () => { diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 38476490b3..8fdae8722e 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1519,6 +1519,15 @@ export class ClaudeAgentClient implements AgentClient { this.rewindSdk = options.rewindSdk ?? realClaudeRewindSdk; } + async resolveUsageReference() { + const env = createProviderEnv({ baseEnv: process.env, runtimeSettings: this.runtimeSettings }); + if (env.ANTHROPIC_BASE_URL || env.ANTHROPIC_API_KEY || env.ANTHROPIC_AUTH_TOKEN) return null; + return { + source: "claude", + input: { configDir: env.CLAUDE_CONFIG_DIR || path.join(env.HOME || homedir(), ".claude") }, + }; + } + resolveConfiguredModel(model: AgentModelDefinition): AgentModelDefinition { return resolveConfiguredClaudeModel(model); } @@ -2042,15 +2051,6 @@ class ClaudeAgentSession implements AgentSession { readonly provider = "claude" as const; readonly capabilities = CLAUDE_CAPABILITIES; - async getUsageReference() { - const env = this.buildSdkEnv(); - if (env.ANTHROPIC_BASE_URL || env.ANTHROPIC_API_KEY || env.ANTHROPIC_AUTH_TOKEN) return null; - return { - source: "claude", - input: { configDir: env.CLAUDE_CONFIG_DIR || path.join(env.HOME || homedir(), ".claude") }, - }; - } - private readonly config: ClaudeAgentConfig; private readonly launchEnv?: Record; private readonly agentId?: string; diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index d4a96b45e6..26f1227657 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts @@ -172,41 +172,18 @@ function createSession( } test("Codex usage reference follows CODEX_HOME and excludes custom base URLs", async () => { - const session = new CodexAppServerAgentSession( - createConfig(), - null, - createTestLogger(), - () => { - throw new Error("unused"); - }, - {}, - false, - false, - false, - undefined, - "interactive", - { CODEX_HOME: "/accounts/second" }, - ); - expect(await session.getUsageReference()).toEqual({ + const config = createConfig(); + const client = new CodexAppServerAgentClient(createTestLogger(), { + env: { CODEX_HOME: "/accounts/second" }, + }); + expect(await client.resolveUsageReference({ config, session: null })).toEqual({ source: "codex", input: { codexHome: "/accounts/second" }, }); - const custom = new CodexAppServerAgentSession( - createConfig(), - null, - createTestLogger(), - () => { - throw new Error("unused"); - }, - {}, - false, - false, - false, - undefined, - "interactive", - { CODEX_HOME: "/accounts/second", OPENAI_BASE_URL: "https://example.test" }, - ); - expect(await custom.getUsageReference()).toBeNull(); + const custom = new CodexAppServerAgentClient(createTestLogger(), { + env: { CODEX_HOME: "/accounts/second", OPENAI_BASE_URL: "https://example.test" }, + }); + expect(await custom.resolveUsageReference({ config, session: null })).toBeNull(); }); function createProviderWithFakeAppServer( diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index 55d7886df8..a52706e1ae 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -3456,7 +3456,6 @@ export class CodexAppServerAgentSession implements AgentSession { private readonly autoReviewEnabled: boolean = false, private readonly agentId?: string, private readonly initialResumePurpose: "interactive" | "history" = "interactive", - private readonly usageEnv: NodeJS.ProcessEnv = process.env, ) { this.logger = logger.child({ module: "agent", @@ -3486,17 +3485,6 @@ export class CodexAppServerAgentSession implements AgentSession { } } - async getUsageReference() { - if (this.usageEnv.OPENAI_BASE_URL) return null; - return { - source: "codex", - input: { - codexHome: - this.usageEnv.CODEX_HOME || path.join(this.usageEnv.HOME || os.homedir(), ".codex"), - }, - }; - } - get id(): string | null { return this.currentThreadId; } @@ -7040,6 +7028,15 @@ export class CodexAppServerAgentClient implements AgentClient { private readonly deps: CodexAppServerAgentDeps = {}, ) {} + async resolveUsageReference() { + const env = buildCodexAppServerEnv(this.runtimeSettings); + if (env.OPENAI_BASE_URL) return null; + return { + source: "codex", + input: { codexHome: env.CODEX_HOME || path.join(env.HOME || os.homedir(), ".codex") }, + }; + } + private sessionDeps(launchEnv: Record | undefined): CodexAppServerAgentDeps { return { ...this.deps, @@ -7161,7 +7158,6 @@ export class CodexAppServerAgentClient implements AgentClient { autoReviewEnabled, launchContext?.agentId, "interactive", - buildCodexAppServerEnv(this.runtimeSettings, launchContext?.env), ); await session.connect(); return session; @@ -7194,7 +7190,6 @@ export class CodexAppServerAgentClient implements AgentClient { autoReviewEnabled, launchContext?.agentId, options?.purpose ?? "interactive", - buildCodexAppServerEnv(this.runtimeSettings, launchContext?.env), ); await session.connect(); return session; diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index 5555fa130e..aab87ddea6 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -969,9 +969,13 @@ export class OmpAgentSession implements AgentSession { private readonly paseoTools?: PaseoToolCatalog; private readonly usageEnv: NodeJS.ProcessEnv; - async getUsageReference() { + async getUsageContext() { const state = await this.runtimeSession.getState(); - return resolveOmpUsageReference(state.sessionId, state.model?.provider ?? "", this.usageEnv); + return { + sessionId: state.sessionId, + provider: state.model?.provider ?? "", + env: this.usageEnv, + }; } get id(): string | null { @@ -2274,6 +2278,18 @@ export class OmpAgentClient implements AgentClient { options.runtime ?? createRuntime(options.logger, runtimeSettings, this.providerParams); } + async resolveUsageReference({ + session, + }: { + config: AgentSessionConfig; + 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); + } + private async configureNativePaseoTools( runtimeSession: OmpRuntimeSession, catalog: PaseoToolCatalog | undefined, 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 7c99fa772c..0390a5905f 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 @@ -484,7 +484,10 @@ export class OmpHarness { } async getUsageReference() { - return this.requireSession().getUsageReference(); + return this.client.resolveUsageReference({ + config: { provider: "omp", cwd: CWD }, + session: this.requireSession(), + }); } runningToolCallIds(): string[] { 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 0634ec25cb..191a87c160 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -420,19 +420,28 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { }, { serverManager: runtime, createClient: runtime.createClient }, ); + expect( + await client.resolveUsageReference({ + config: { ...buildConfig(cwd), model: "openai/gpt-5" }, + session: null, + }), + ).toEqual({ + source: "codex", + input: { accessToken: "oauth-token", accountId: "acct-1" }, + }); const session = await client.createSession(buildConfig(cwd)); await session.setModel?.("openai/gpt-5"); - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toEqual({ source: "codex", input: { accessToken: "oauth-token", accountId: "acct-1" }, }); await session.setModel?.("opencode-go/qwen"); - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toEqual({ source: "opencode-go", input: { apiKey: "go-key" }, }); await session.setModel?.("other/model"); - expect(await session.getUsageReference?.()).toBeNull(); + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toBeNull(); const events: AgentStreamEvent[] = []; session.subscribe((event) => events.push(event)); const sessionId = (await session.getRuntimeInfo()).sessionId; @@ -448,7 +457,7 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { }, } as OpenCodeEvent); await vi.waitFor(async () => - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toEqual({ source: "codex", input: { accessToken: "oauth-token", accountId: "acct-1" }, }), @@ -487,7 +496,7 @@ describe("OpenCodeAgentClient adapter smoke tests", () => { const session = await client.createSession(buildConfig(cwd)); try { await session.setModel?.("opencode-go/qwen"); - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: buildConfig(cwd), session })).toEqual({ source: "opencode-go", input: { apiKey: "go-key" }, }); diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index b2cab99c6e..1bbb908f2d 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -1445,6 +1445,22 @@ export class OpenCodeAgentClient implements AgentClient { this.resolveHomeDir = deps.resolveHomeDir ?? resolveOpenCodeHomeDir; } + async resolveUsageReference({ + config, + session, + }: { + config: AgentSessionConfig; + session: AgentSession | null; + }): Promise { + const model = session ? (await session.getRuntimeInfo()).model : config.model; + if (!model?.startsWith("openai/") && !model?.startsWith("opencode-go/")) return null; + const env = + session instanceof OpenCodeAgentSession + ? session.getUsageEnvironment() + : { ...process.env, ...this.runtimeSettings?.env }; + return resolveOpenCodeUsageReference(model, await readOpenCodeUsageAuth(env)); + } + async createSession( config: AgentSessionConfig, launchContext?: AgentLaunchContext, @@ -3620,11 +3636,8 @@ class OpenCodeAgentSession implements AgentSession { }; } - async getUsageReference(): Promise { - const model = this.config.model; - if (!model?.startsWith("openai/") && !model?.startsWith("opencode-go/")) return null; - const auth = await readOpenCodeUsageAuth(this.usageEnv); - return resolveOpenCodeUsageReference(model, auth); + getUsageEnvironment(): NodeJS.ProcessEnv { + return this.usageEnv; } async setModel(modelId: string | null): Promise { 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 7308185db0..99770c16e9 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -229,15 +229,16 @@ test("Pi usage reference follows the current OAuth model and agent directory", a runtime = sessionRuntime; runtime.state.model = { provider: "openai-codex", id: "gpt", name: "GPT" }; }); - const session = await createClient(pi).createSession(createConfig(), { + const client = createClient(pi); + const session = await client.createSession(createConfig(), { env: { PI_CODING_AGENT_DIR: agentDir }, }); - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: createConfig(), session })).toEqual({ source: "codex", input: { accessToken: "codex-token", accountId: "account-1" }, }); runtime.state.model = { provider: "anthropic", id: "claude", name: "Claude" }; - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: createConfig(), session })).toEqual({ source: "claude", input: { accessToken: "claude-token" }, }); @@ -249,7 +250,7 @@ test("Pi usage reference follows the current OAuth model and agent directory", a }), ); runtime.state.model = { provider: "openai-codex", id: "gpt", name: "GPT" }; - expect(await session.getUsageReference?.()).toEqual({ + expect(await client.resolveUsageReference({ config: createConfig(), session })).toEqual({ source: "codex", input: { accessToken: "codex-token" }, }); @@ -258,15 +259,38 @@ test("Pi usage reference follows the current OAuth model and agent directory", a path.join(agentDir, "auth.json"), JSON.stringify({ anthropic: { type: "api_key", key: "api-key" } }), ); - expect(await session.getUsageReference?.()).toBeNull(); + expect(await client.resolveUsageReference({ config: createConfig(), session })).toBeNull(); runtime.state.model = { provider: "other", id: "other", name: "Other" }; - expect(await session.getUsageReference?.()).toBeNull(); + expect(await client.resolveUsageReference({ config: createConfig(), session })).toBeNull(); await session.close(); } finally { rmSync(agentDir, { recursive: true, force: true }); } }); +test("Pi resolves a stored model without opening a session", async () => { + const agentDir = mkdtempSync(path.join(tmpdir(), "paseo-pi-stored-usage-")); + try { + writeFileSync( + path.join(agentDir, "auth.json"), + JSON.stringify({ anthropic: { type: "oauth", access: "claude-token", refresh: "refresh" } }), + ); + const client = new PiRpcAgentClient({ + logger: pino({ level: "silent" }), + runtime: new FakePi(), + runtimeSettings: { env: { PI_CODING_AGENT_DIR: agentDir } }, + }); + expect( + await client.resolveUsageReference({ + config: createConfig({ model: "anthropic/claude-sonnet-4" }), + session: null, + }), + ).toEqual({ source: "claude", input: { accessToken: "claude-token" } }); + } finally { + rmSync(agentDir, { recursive: true, force: true }); + } +}); + test("forwards launch-context env to the Pi process launch", async () => { const pi = new FakePi(); const client = createClient(pi); diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5d7966367e..4ccab930ba 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -1141,37 +1141,40 @@ function createRuntime( }); } +async function resolvePiUsageReference(provider: string | undefined, env?: Record) { + let source: string | null = null; + if (provider === "openai-codex") source = "codex"; + if (provider === "anthropic") source = "claude"; + if (!source || !provider) return null; + try { + const auth: unknown = JSON.parse( + await readFile(join(resolvePiAgentDir(env), "auth.json"), "utf8"), + ); + if (provider === "openai-codex") { + const credential = piCodexUsageAuthSchema.parse(auth)["openai-codex"]; + if (!credential) return null; + return { + source, + input: { + accessToken: credential.access, + ...(credential.accountId ? { accountId: credential.accountId } : {}), + }, + }; + } + const credential = piClaudeUsageAuthSchema.parse(auth).anthropic; + return credential ? { source, input: { accessToken: credential.access } } : null; + } catch { + return null; + } +} + export class PiRpcAgentSession implements AgentSession { readonly provider: AgentProvider; readonly capabilities: AgentCapabilityFlags; - async getUsageReference() { + async getUsageContext() { await this.refreshState(); - const provider = this.state.model?.provider; - let source: string | null = null; - if (provider === "openai-codex") source = "codex"; - if (provider === "anthropic") source = "claude"; - if (!source || !provider) return null; - try { - const auth: unknown = JSON.parse( - await readFile(join(resolvePiAgentDir(this.usageEnv), "auth.json"), "utf8"), - ); - if (provider === "openai-codex") { - const credential = piCodexUsageAuthSchema.parse(auth)["openai-codex"]; - if (!credential) return null; - return { - source, - input: { - accessToken: credential.access, - ...(credential.accountId ? { accountId: credential.accountId } : {}), - }, - }; - } - const credential = piClaudeUsageAuthSchema.parse(auth).anthropic; - return credential ? { source, input: { accessToken: credential.access } } : null; - } catch { - return null; - } + return { provider: this.state.model?.provider, env: this.usageEnv }; } private readonly subscribers = new Set<(event: AgentStreamEvent) => void>(); @@ -2477,6 +2480,18 @@ export class PiRpcAgentClient implements AgentClient { this.usagePollScheduler = options.usagePollScheduler; } + async resolveUsageReference({ + config, + session, + }: { + config: AgentSessionConfig; + 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); + } + async createSession( config: AgentSessionConfig, launchContext?: AgentLaunchContext, 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 c4f2287bbf..1d7c7a47f3 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,10 +1,8 @@ import { fileURLToPath } from "node:url"; +import { mkdtemp, rm } from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; import { expect, test } from "vitest"; -import { ACPAgentSession } from "../agent/providers/acp-agent.js"; -import type { UsageReference } from "../agent/agent-sdk-types.js"; -import { ClaudeAgentClient } from "../agent/providers/claude/agent.js"; -import { CodexAppServerAgentSession } from "../agent/providers/codex-app-server-agent.js"; -import { createTestLogger } from "../../test-utils/test-logger.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"; @@ -53,75 +51,21 @@ test("resolves a provider plugin session reference through its usage source", as } }, 60_000); -test("agent.resolve_usage_report resolves source IDs from default built-in and ACP sessions", async () => { +test("agent.resolve_usage_report resolves source IDs from provider clients", async () => { const directory = fileURLToPath( new URL("./test-fixtures/session-usage-reference/", import.meta.url), ); const providers = ["claude", "codex", "copilot", "cursor", "kimi", "generic-acp"] as const; - const logger = createTestLogger(); - const references = new Map(); - const claude = await new ClaudeAgentClient({ - logger, - resolveBinary: async () => "/test/claude/bin", - }).createSession( - { provider: "claude", cwd: directory }, - { - env: { - HOME: directory, - CLAUDE_CONFIG_DIR: "", - ANTHROPIC_BASE_URL: "", - ANTHROPIC_API_KEY: "", - ANTHROPIC_AUTH_TOKEN: "", - }, - }, - ); - references.set("claude", (await claude.getUsageReference?.()) ?? null); - await claude.close(); - const codex = new CodexAppServerAgentSession( - { provider: "codex", cwd: directory }, - null, - logger, - () => { - throw new Error("Codex runtime should not start"); - }, - {}, - false, - false, - false, - undefined, - "interactive", - { HOME: directory, OPENAI_BASE_URL: "" }, - ); - references.set("codex", await codex.getUsageReference()); - for (const provider of providers.slice(2)) { - const session = new ACPAgentSession( - { provider, cwd: directory }, - { - provider, - logger, - defaultCommand: ["unused"], - defaultModes: [], - capabilities: { - supportsStreaming: true, - supportsSessionPersistence: true, - supportsDynamicModes: true, - supportsMcpServers: true, - supportsReasoningStream: true, - supportsToolInvocations: true, - }, - }, - ); - references.set(provider, await session.getUsageReference()); - } const agentClients = Object.fromEntries( providers.map((provider) => { const client = createTestAgentClient(provider); - const createSession = client.createSession.bind(client); - client.createSession = async (...args) => { - const session = await createSession(...args); - session.getUsageReference = async () => references.get(provider) ?? null; - return session; - }; + let input: Record = {}; + if (provider === "claude") input = { configDir: directory }; + if (provider === "codex") input = { codexHome: directory }; + client.resolveUsageReference = async () => ({ + source: provider, + input, + }); return [provider, client]; }), ); @@ -148,3 +92,64 @@ test("agent.resolve_usage_report resolves source IDs from default built-in and A await daemon.close(); } }, 60_000); + +test("stored Claude and Codex agents resolve without opening a provider session", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "paseo-usage-closed-")); + const staticDir = await mkdtemp(path.join(os.tmpdir(), "paseo-usage-static-")); + const sessionsOpened = new Map(); + const agentClients = Object.fromEntries( + ["claude", "codex"].map((provider) => { + const providerClient = createTestAgentClient(provider, { + beforeCreateSession: async () => { + sessionsOpened.set(provider, (sessionsOpened.get(provider) ?? 0) + 1); + }, + }); + providerClient.resolveUsageReference = async ({ session }) => { + expect(session).toBeNull(); + return { source: "fixture-session-usage", input: { account: "from-session" } }; + }; + return [provider, providerClient]; + }), + ); + const options = { + paseoHomeRoot: root, + staticDir, + cleanup: false, + pluginsEnabled: false, + builtinPlugins: fixtureBuiltins(), + agentClients, + }; + let daemon = await createTestPaseoDaemon(options); + let client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + try { + await client.connect(); + const ids: string[] = []; + for (const provider of ["claude", "codex"]) { + const agent = await client.createAgent({ provider, cwd: root }); + ids.push(agent.id); + await client.sendMessage(agent.id, "Say OK."); + await client.waitForAgentUpsert(agent.id, (snapshot) => snapshot.status === "idle"); + expect(sessionsOpened.get(provider)).toBe(1); + } + await client.close(); + await daemon.close(); + + daemon = await createTestPaseoDaemon(options); + client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + await client.connect(); + await client.fetchAgents({ subscribe: {} }); + for (const [index, provider] of ["claude", "codex"].entries()) { + const result = await client.resolveAgentUsageReport({ agentId: ids[index]! }); + expect(result.reportId).toBe("fixture-session-usage:from-session"); + expect( + (await client.listUsageReports({ reportIds: [result.reportId!] })).reports, + ).toHaveLength(1); + expect(sessionsOpened.get(provider)).toBe(1); + } + } finally { + await client.close().catch(() => undefined); + await daemon.close().catch(() => undefined); + await rm(root, { recursive: true, force: true }); + await rm(staticDir, { recursive: true, force: true }); + } +}, 60_000); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index f83040c8c9..f2365b54ff 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -1000,7 +1000,7 @@ export class Session { this.usageSession = new UsageSession({ emit: (msg) => this.emit(msg), listAgents: () => this.agentManager.listAgents(), - getAgent: (agentId) => this.agentManager.getAgent(agentId), + resolveAgentReference: (agentId) => this.agentManager.resolveUsageReference(agentId), runtime: pluginRuntime, logger: this.sessionLogger, }); diff --git a/packages/server/src/server/session/usage/usage-session.test.ts b/packages/server/src/server/session/usage/usage-session.test.ts index 7914af3872..578747333a 100644 --- a/packages/server/src/server/session/usage/usage-session.test.ts +++ b/packages/server/src/server/session/usage/usage-session.test.ts @@ -1,6 +1,5 @@ import pino from "pino"; import { expect, test } from "vitest"; -import type { AgentSession } from "../../agent/agent-sdk-types.js"; import type { SessionOutboundMessage } from "../../messages.js"; import { UsageSession } from "./usage-session.js"; @@ -9,9 +8,8 @@ test("collects live references for usage reports and resolves an agent report", const references: unknown[] = []; const reference = { source: "fixture", input: { account: "one" } }; const agent = { - session: { - getUsageReference: async () => reference, - } as AgentSession, + id: "one", + session: {}, }; const entry = { id: "fixture:one", @@ -23,8 +21,8 @@ test("collects live references for usage reports and resolves an agent report", }; const usage = new UsageSession({ emit: (message) => emitted.push(message), - listAgents: () => [agent, { session: null }], - getAgent: (id) => (id === "one" ? agent : null), + listAgents: () => [agent, { id: "closed", session: null }], + resolveAgentReference: async (id) => ({ found: id === "one", reference }), runtime: { async listUsageReports(options) { references.push(options.references); @@ -59,7 +57,7 @@ test("surfaces a legacy usage-list failure as an rpc_error envelope", async () = const usage = new UsageSession({ emit: (message) => emitted.push(message), listAgents: () => [], - getAgent: () => null, + resolveAgentReference: async () => ({ found: false, reference: null }), runtime: { async listUsageReports() { return []; @@ -85,7 +83,7 @@ test("unknown agent returns agent_not_found", async () => { const usage = new UsageSession({ emit: (message) => emitted.push(message), listAgents: () => [], - getAgent: () => null, + resolveAgentReference: async () => ({ found: false, reference: null }), runtime: { listUsageReports: async () => [], resolveUsageReference: async () => null, diff --git a/packages/server/src/server/session/usage/usage-session.ts b/packages/server/src/server/session/usage/usage-session.ts index 6665a7454f..21c8f99801 100644 --- a/packages/server/src/server/session/usage/usage-session.ts +++ b/packages/server/src/server/session/usage/usage-session.ts @@ -1,16 +1,20 @@ import type pino from "pino"; import type { ProviderUsage, UsageReportEntry } from "@getpaseo/protocol/messages"; -import type { AgentSession, UsageReference } from "../../agent/agent-sdk-types.js"; +import type { UsageReference } from "../../agent/agent-sdk-types.js"; import type { SessionInboundMessage, SessionOutboundMessage } from "../../messages.js"; interface UsageAgent { - session: AgentSession | null; + id: string; + session: object | null; } export interface UsageSessionOptions { emit(message: SessionOutboundMessage): void; listAgents(): UsageAgent[]; - getAgent(agentId: string): UsageAgent | null; + resolveAgentReference(agentId: string): Promise<{ + found: boolean; + reference: UsageReference | null; + }>; runtime?: { listUsageReports(options: { forceRefresh?: boolean; @@ -37,8 +41,12 @@ export class UsageSession { await Promise.all( this.options .listAgents() - .map( - async (agent) => agent.session?.getUsageReference?.().catch(() => null) ?? null, + .filter((agent) => agent.session !== null) + .map(async (agent) => + this.options + .resolveAgentReference(agent.id) + .then((resolved) => resolved.reference) + .catch(() => null), ), ) ).filter((reference): reference is UsageReference => reference !== null) @@ -62,12 +70,12 @@ export class UsageSession { ): Promise { try { if (!this.options.runtime) throw new Error("Plugin runtime is unavailable"); - const agent = this.options.getAgent(msg.agentId); - if (!agent) { + const resolved = await this.options.resolveAgentReference(msg.agentId); + if (!resolved.found) { this.emitError(msg, new Error(`Agent not found: ${msg.agentId}`), "agent_not_found"); return; } - const reference = (await agent.session?.getUsageReference?.()) ?? null; + const reference = resolved.reference; const reportId = reference ? await this.options.runtime.resolveUsageReference(reference) : null;