From 89ccc5fac135ea0c23ff97e5c61f4220cee7487d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B6=85=E9=A2=84=E6=9C=9F?= Date: Tue, 10 Mar 2026 22:32:41 +0100 Subject: [PATCH 1/2] feat: per-agent auth profile binding for rate limit isolation Add optional `auth` field to agent config entries that maps provider names to specific auth profile IDs. This allows each agent to use a dedicated authentication profile per provider, enabling: - Rate limit isolation between agents - Per-agent billing visibility - Independent credential rotation Changes: - Schema: add `auth: Record` to AgentEntrySchema + AgentConfig - Runner: look up per-agent auth binding and inject as preferredProfile - Doctor: validate per-agent auth profile references exist in store - Tests: 5 new test cases for schema validation + preferredProfile ordering - Help/labels: add user-facing documentation for the new field --- src/agents/auth-profiles/order.test.ts | 92 ++++++++++++++++++++++++++ src/agents/pi-embedded-runner/run.ts | 13 +++- src/commands/doctor-auth.ts | 21 ++++++ src/config/schema.help.ts | 2 + src/config/schema.labels.ts | 1 + src/config/types.agents.ts | 2 + src/config/zod-schema.agent-runtime.ts | 1 + 7 files changed, 131 insertions(+), 1 deletion(-) diff --git a/src/agents/auth-profiles/order.test.ts b/src/agents/auth-profiles/order.test.ts index a1b15192e16..3ba355a122b 100644 --- a/src/agents/auth-profiles/order.test.ts +++ b/src/agents/auth-profiles/order.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it } from "vitest"; +import { AgentEntrySchema } from "../../config/zod-schema.agent-runtime.js"; import { resolveAuthProfileOrder } from "./order.js"; import type { AuthProfileStore } from "./types.js"; @@ -22,4 +23,95 @@ describe("resolveAuthProfileOrder", () => { expect(order).toEqual(["volcengine:default"]); }); + + it("places per-agent preferredProfile first in the order", () => { + const store: AuthProfileStore = { + version: 1, + profiles: { + "openai:account1": { + type: "api_key", + provider: "openai", + key: "sk-1", + }, + "openai:account2": { + type: "api_key", + provider: "openai", + key: "sk-2", + }, + }, + }; + + const order = resolveAuthProfileOrder({ + store, + provider: "openai", + preferredProfile: "openai:account2", + }); + + // The preferred profile should be first in the ordering + expect(order[0]).toBe("openai:account2"); + expect(order).toContain("openai:account1"); + }); + + it("ignores preferredProfile that does not exist in store", () => { + const store: AuthProfileStore = { + version: 1, + profiles: { + "openai:account1": { + type: "api_key", + provider: "openai", + key: "sk-1", + }, + }, + }; + + const order = resolveAuthProfileOrder({ + store, + provider: "openai", + preferredProfile: "openai:nonexistent", + }); + + expect(order).toEqual(["openai:account1"]); + }); +}); + +describe("AgentEntrySchema auth field", () => { + it("accepts an agent entry with per-agent auth bindings", () => { + const result = AgentEntrySchema.safeParse({ + id: "my-agent", + auth: { + openai: "openai:account2", + anthropic: "anthropic:team-key", + }, + }); + + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.auth).toEqual({ + openai: "openai:account2", + anthropic: "anthropic:team-key", + }); + } + }); + + it("accepts an agent entry without auth field (backward compat)", () => { + const result = AgentEntrySchema.safeParse({ + id: "my-agent", + }); + + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.auth).toBeUndefined(); + } + }); + + it("rejects auth field with non-string values", () => { + const result = AgentEntrySchema.safeParse({ + id: "my-agent", + auth: { + openai: 123, + }, + }); + + expect(result.success).toBe(false); + }); }); diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 7f5f4f525b7..3254ed434c6 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -409,6 +409,17 @@ export async function runEmbeddedPiAgent( const authStore = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false }); const preferredProfileId = params.authProfileId?.trim(); + + // Per-agent auth binding: if config specifies a dedicated profile for this + // agent+provider combination, use it as the preferred profile. This gives + // per-agent bindings priority over session-level overrides while + // maintaining full backward compatibility for agents without bindings. + const normalizedProvider = normalizeProviderId(provider); + const agentAuthBinding = params.config?.agents?.list?.find( + (a: { id: string }) => a.id === params.agentId, + )?.auth?.[normalizedProvider]; + const effectivePreferredProfile = agentAuthBinding?.trim() || preferredProfileId; + let lockedProfileId = params.authProfileIdSource === "user" ? preferredProfileId : undefined; if (lockedProfileId) { const lockedProfile = authStore.profiles[lockedProfileId]; @@ -423,7 +434,7 @@ export async function runEmbeddedPiAgent( cfg: params.config, store: authStore, provider, - preferredProfile: preferredProfileId, + preferredProfile: effectivePreferredProfile, }); if (lockedProfileId && !profileOrder.includes(lockedProfileId)) { throw new Error(`Auth profile "${lockedProfileId}" is not configured for ${provider}.`); diff --git a/src/commands/doctor-auth.ts b/src/commands/doctor-auth.ts index 56ba510f41d..52a4bb464d0 100644 --- a/src/commands/doctor-auth.ts +++ b/src/commands/doctor-auth.ts @@ -256,6 +256,27 @@ export async function noteAuthProfileHealth(params: { const store = ensureAuthProfileStore(undefined, { allowKeychainPrompt: params.allowKeychainPrompt, }); + + // Validate per-agent auth bindings reference existing profiles. + const agents = params.cfg.agents?.list ?? []; + const perAgentWarnings: string[] = []; + for (const agent of agents) { + if (agent.auth && typeof agent.auth === "object") { + for (const [provider, profileId] of Object.entries(agent.auth)) { + if (typeof profileId !== "string" || !profileId.trim()) { + continue; + } + if (!store.profiles[profileId]) { + perAgentWarnings.push( + `- agent "${agent.id}": auth profile "${profileId}" (provider: ${provider}) not found in auth store`, + ); + } + } + } + } + if (perAgentWarnings.length > 0) { + note(perAgentWarnings.join("\n"), "Per-agent auth binding warnings"); + } const unusable = (() => { const now = Date.now(); const out: string[] = []; diff --git a/src/config/schema.help.ts b/src/config/schema.help.ts index 908829cbf33..4a7d7042dfe 100644 --- a/src/config/schema.help.ts +++ b/src/config/schema.help.ts @@ -200,6 +200,8 @@ export const FIELD_HELP: Record = { "Idle runtime TTL in minutes for ACP session workers before eligible cleanup.", "acp.runtime.installCommand": "Optional operator install/setup command shown by `/acp install` and `/acp doctor` when ACP backend wiring is missing.", + "agents.list[].auth": + 'Per-agent auth profile bindings. Maps provider names to specific auth profile IDs (e.g. {"openai": "openai:account2"}) for rate limit isolation and per-agent billing.', "agents.list.*.skills": "Optional allowlist of skills for this agent (omit = all skills; empty = no skills).", "agents.list[].skills": diff --git a/src/config/schema.labels.ts b/src/config/schema.labels.ts index c643cf91cd9..662b935c076 100644 --- a/src/config/schema.labels.ts +++ b/src/config/schema.labels.ts @@ -56,6 +56,7 @@ export const FIELD_LABELS: Record = { "diagnostics.cacheTrace.includeSystem": "Cache Trace Include System", "agents.list.*.identity.avatar": "Identity Avatar", "agents.list.*.skills": "Agent Skill Filter", + "agents.list[].auth": "Per-Agent Auth", "agents.list[].runtime": "Agent Runtime", "agents.list[].runtime.type": "Agent Runtime Type", "agents.list[].runtime.acp": "Agent ACP Runtime", diff --git a/src/config/types.agents.ts b/src/config/types.agents.ts index a979506a2ab..3b6e3cd7c84 100644 --- a/src/config/types.agents.ts +++ b/src/config/types.agents.ts @@ -64,6 +64,8 @@ export type AgentConfig = { name?: string; workspace?: string; agentDir?: string; + /** Per-agent auth profile bindings. Maps provider names to specific auth profile IDs. */ + auth?: Record; model?: AgentModelConfig; /** Optional allowlist of skills for this agent (omit = all skills; empty = none). */ skills?: string[]; diff --git a/src/config/zod-schema.agent-runtime.ts b/src/config/zod-schema.agent-runtime.ts index 3ede7218b80..505b3205f27 100644 --- a/src/config/zod-schema.agent-runtime.ts +++ b/src/config/zod-schema.agent-runtime.ts @@ -719,6 +719,7 @@ export const AgentEntrySchema = z name: z.string().optional(), workspace: z.string().optional(), agentDir: z.string().optional(), + auth: z.record(z.string(), z.string()).optional(), model: AgentModelSchema.optional(), skills: z.array(z.string()).optional(), memorySearch: MemorySearchSchema, From ea1d5a30dbebd4fa59e1d443aa138230c1b0a448 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B6=85=E9=A2=84=E6=9C=9F?= Date: Wed, 11 Mar 2026 01:05:33 +0100 Subject: [PATCH 2/2] fix: use effectivePreferredProfile for lockedProfileId derivation Addresses Greptile review feedback: lockedProfileId now derives from effectivePreferredProfile instead of raw preferredProfileId, ensuring per-agent auth bindings take priority over session-level overrides. Added integration test for per-agent binding vs user-locked profile scenario. --- src/agents/auth-profiles/order.test.ts | 35 ++++++++++++++++++++++++++ src/agents/pi-embedded-runner/run.ts | 3 ++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/agents/auth-profiles/order.test.ts b/src/agents/auth-profiles/order.test.ts index 3ba355a122b..d4bcd660162 100644 --- a/src/agents/auth-profiles/order.test.ts +++ b/src/agents/auth-profiles/order.test.ts @@ -72,6 +72,41 @@ describe("resolveAuthProfileOrder", () => { expect(order).toEqual(["openai:account1"]); }); + + it("per-agent binding overrides session-level profile via effectivePreferredProfile", () => { + // Simulates the integration scenario where a per-agent binding (account2) + // should take priority over a session-level preferred profile (account1). + const store: AuthProfileStore = { + version: 1, + profiles: { + "openai:account1": { + type: "api_key", + provider: "openai", + key: "sk-1", + }, + "openai:account2": { + type: "api_key", + provider: "openai", + key: "sk-2", + }, + }, + }; + + // Session requests account1, but per-agent config binds to account2. + const sessionPreferred = "openai:account1"; + const agentAuthBinding = "openai:account2"; + const effectivePreferred = agentAuthBinding || sessionPreferred; + + const order = resolveAuthProfileOrder({ + store, + provider: "openai", + preferredProfile: effectivePreferred, + }); + + // Per-agent binding (account2) must come first, not the session override (account1) + expect(order[0]).toBe("openai:account2"); + expect(order).toContain("openai:account1"); + }); }); describe("AgentEntrySchema auth field", () => { diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 3254ed434c6..fc508f1986f 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -420,7 +420,8 @@ export async function runEmbeddedPiAgent( )?.auth?.[normalizedProvider]; const effectivePreferredProfile = agentAuthBinding?.trim() || preferredProfileId; - let lockedProfileId = params.authProfileIdSource === "user" ? preferredProfileId : undefined; + let lockedProfileId = + params.authProfileIdSource === "user" ? effectivePreferredProfile : undefined; if (lockedProfileId) { const lockedProfile = authStore.profiles[lockedProfileId]; if (