mirror of
https://github.com/getpaseo/paseo.git
synced 2026-11-01 14:12:51 +00:00
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.
This commit is contained in:
@@ -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 };
|
||||
|
||||
@@ -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<UsageReference | null>;
|
||||
createSession(
|
||||
|
||||
@@ -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<AgentSession, AgentSession>();
|
||||
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,
|
||||
|
||||
@@ -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 });
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -71,11 +71,13 @@ export class OmpHarness {
|
||||
providerIdleScheduler?: OmpProviderIdleScheduler;
|
||||
noTurnScheduler?: OmpNoTurnScheduler;
|
||||
usagePollScheduler?: OmpUsagePollScheduler;
|
||||
env?: Record<string, string>;
|
||||
} = {},
|
||||
) {
|
||||
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<string, string>();
|
||||
for (const item of this.timeline()) {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -1447,12 +1447,16 @@ export class OpenCodeAgentClient implements AgentClient {
|
||||
|
||||
async resolveUsageReference({
|
||||
config,
|
||||
runtimeInfo,
|
||||
session,
|
||||
}: {
|
||||
config: AgentSessionConfig;
|
||||
runtimeInfo: AgentRuntimeInfo | null;
|
||||
session: AgentSession | null;
|
||||
}): Promise<UsageReference | null> {
|
||||
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
|
||||
|
||||
@@ -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 });
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user