feat: sticky auth profile rotation + usage headers

This commit is contained in:
Peter Steinberger
2026-01-16 00:24:31 +00:00
parent e274b5a040
commit 8c3cdba21c
16 changed files with 334 additions and 31 deletions

View File

@@ -1,5 +1,6 @@
import { Type } from "@sinclair/typebox";
import { resolveAgentDir } from "../../agents/agent-scope.js";
import { buildAuthHealthSummary, formatRemainingShort } from "../../agents/auth-health.js";
import {
ensureAuthProfileStore,
resolveAuthProfileDisplayLabel,
@@ -28,9 +29,10 @@ import {
updateSessionStore,
} from "../../config/sessions.js";
import {
formatUsageSummaryLine,
formatUsageWindowSummary,
loadProviderUsageSummary,
resolveUsageProviderId,
type UsageProviderId,
} from "../../infra/provider-usage.js";
import {
buildAgentMainSessionKey,
@@ -257,10 +259,14 @@ export function createSessionStatusTool(opts?: {
delete nextEntry.providerOverride;
delete nextEntry.modelOverride;
delete nextEntry.authProfileOverride;
delete nextEntry.authProfileOverrideSource;
delete nextEntry.authProfileOverrideCompactionCount;
} else {
nextEntry.providerOverride = selection.provider;
nextEntry.modelOverride = selection.model;
delete nextEntry.authProfileOverride;
delete nextEntry.authProfileOverrideSource;
delete nextEntry.authProfileOverrideCompactionCount;
}
store[resolved.key] = nextEntry;
await updateSessionStore(storePath, (nextStore) => {
@@ -277,24 +283,48 @@ export function createSessionStatusTool(opts?: {
defaultModel: DEFAULT_MODEL,
});
const providerForCard = resolved.entry.providerOverride?.trim() || configured.provider;
const usageProvider = resolveUsageProviderId(providerForCard);
let usageLine: string | undefined;
if (usageProvider) {
const authStore = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false });
const authHealth = buildAuthHealthSummary({
store: authStore,
cfg,
});
const oauthProfiles = authHealth.profiles.filter(
(profile) => profile.type === "oauth" || profile.type === "token",
);
const usageProviders = Array.from(
new Set(
oauthProfiles
.map((profile) => resolveUsageProviderId(profile.provider))
.filter((provider): provider is UsageProviderId => Boolean(provider)),
),
);
const usageByProvider = new Map<string, string>();
if (usageProviders.length > 0) {
try {
const usageSummary = await loadProviderUsageSummary({
timeoutMs: 3500,
providers: [usageProvider],
providers: usageProviders,
agentDir,
});
const formatted = formatUsageSummaryLine(usageSummary, {
now: Date.now(),
});
if (formatted) usageLine = formatted;
for (const snapshot of usageSummary.providers) {
const formatted = formatUsageWindowSummary(snapshot, {
now: Date.now(),
maxWindows: 2,
});
if (formatted) usageByProvider.set(snapshot.provider, formatted);
}
} catch {
// ignore
}
}
const usageProvider = resolveUsageProviderId(providerForCard);
const usageLine =
oauthProfiles.length === 0 && usageProvider && usageByProvider.has(usageProvider)
? `📊 Usage: ${usageByProvider.get(usageProvider)}`
: undefined;
const isGroup =
resolved.entry.chatType === "group" ||
resolved.entry.chatType === "room" ||
@@ -340,13 +370,53 @@ export function createSessionStatusTool(opts?: {
includeTranscriptUsage: false,
});
const authStatusLines = (() => {
if (oauthProfiles.length === 0) return [];
const formatStatus = (status: string) => {
if (status === "ok") return "ok";
if (status === "static") return "static";
if (status === "expiring") return "expiring";
if (status === "missing") return "unknown";
return "expired";
};
const profilesByProvider = new Map<string, typeof oauthProfiles>();
for (const profile of oauthProfiles) {
const current = profilesByProvider.get(profile.provider);
if (current) current.push(profile);
else profilesByProvider.set(profile.provider, [profile]);
}
const lines: string[] = ["OAuth/token status"];
for (const [provider, profiles] of profilesByProvider) {
const usageKey = resolveUsageProviderId(provider);
const usage = usageKey ? usageByProvider.get(usageKey) : undefined;
const usageSuffix = usage ? ` — usage: ${usage}` : "";
lines.push(`- ${provider}${usageSuffix}`);
for (const profile of profiles) {
const labelText = profile.label || profile.profileId;
const status = formatStatus(profile.status);
const expiry =
profile.status === "static"
? ""
: profile.expiresAt
? ` expires in ${formatRemainingShort(profile.remainingMs)}`
: " expires unknown";
const source = profile.source !== "store" ? ` (${profile.source})` : "";
lines.push(` - ${labelText} ${status}${expiry}${source}`);
}
}
return lines;
})();
const fullStatusText =
authStatusLines.length > 0 ? `${statusText}\n\n${authStatusLines.join("\n")}` : statusText;
return {
content: [{ type: "text", text: statusText }],
content: [{ type: "text", text: fullStatusText }],
details: {
ok: true,
sessionKey: resolved.key,
changedModel,
statusText,
statusText: fullStatusText,
},
};
},

View File

@@ -264,8 +264,12 @@ export async function handleDirectiveOnly(params: {
}
if (profileOverride) {
sessionEntry.authProfileOverride = profileOverride;
sessionEntry.authProfileOverrideSource = "user";
delete sessionEntry.authProfileOverrideCompactionCount;
} else if (directives.hasModelDirective) {
delete sessionEntry.authProfileOverride;
delete sessionEntry.authProfileOverrideSource;
delete sessionEntry.authProfileOverrideCompactionCount;
}
}
if (directives.hasQueueDirective && directives.queueReset) {

View File

@@ -156,8 +156,12 @@ export async function persistInlineDirectives(params: {
}
if (profileOverride) {
sessionEntry.authProfileOverride = profileOverride;
sessionEntry.authProfileOverrideSource = "user";
delete sessionEntry.authProfileOverrideCompactionCount;
} else if (directives.hasModelDirective) {
delete sessionEntry.authProfileOverride;
delete sessionEntry.authProfileOverrideSource;
delete sessionEntry.authProfileOverrideCompactionCount;
}
provider = resolved.ref.provider;
model = resolved.ref.model;

View File

@@ -5,6 +5,11 @@ import {
isEmbeddedPiRunStreaming,
resolveEmbeddedSessionLane,
} from "../../agents/pi-embedded.js";
import {
ensureAuthProfileStore,
isProfileInCooldown,
resolveAuthProfileOrder,
} from "../../agents/auth-profiles.js";
import type { ClawdbotConfig } from "../../config/config.js";
import {
resolveSessionFilePath,
@@ -97,6 +102,86 @@ type RunPreparedReplyParams = {
abortedLastRun: boolean;
};
async function resolveSessionAuthProfileOverride(params: {
cfg: ClawdbotConfig;
provider: string;
agentDir: string;
sessionEntry?: SessionEntry;
sessionStore?: Record<string, SessionEntry>;
sessionKey?: string;
storePath?: string;
isNewSession: boolean;
}): Promise<string | undefined> {
const {
cfg,
provider,
agentDir,
sessionEntry,
sessionStore,
sessionKey,
storePath,
isNewSession,
} = params;
if (!sessionEntry || !sessionStore || !sessionKey) return sessionEntry?.authProfileOverride;
const store = ensureAuthProfileStore(agentDir, { allowKeychainPrompt: false });
const order = resolveAuthProfileOrder({ cfg, store, provider });
if (order.length === 0) return sessionEntry.authProfileOverride;
const pickFirstAvailable = () =>
order.find((profileId) => !isProfileInCooldown(store, profileId)) ?? order[0];
const pickNextAvailable = (current: string) => {
const startIndex = order.indexOf(current);
if (startIndex < 0) return pickFirstAvailable();
for (let offset = 1; offset <= order.length; offset += 1) {
const candidate = order[(startIndex + offset) % order.length];
if (!isProfileInCooldown(store, candidate)) return candidate;
}
return order[startIndex] ?? order[0];
};
const compactionCount = sessionEntry.compactionCount ?? 0;
const storedCompaction =
typeof sessionEntry.authProfileOverrideCompactionCount === "number"
? sessionEntry.authProfileOverrideCompactionCount
: compactionCount;
let current = sessionEntry.authProfileOverride?.trim();
if (current && !order.includes(current)) current = undefined;
const source = sessionEntry.authProfileOverrideSource ?? (current ? "user" : undefined);
if (source === "user" && current && !isNewSession) {
return current;
}
let next = current;
if (isNewSession) {
next = current ? pickNextAvailable(current) : pickFirstAvailable();
} else if (current && compactionCount > storedCompaction) {
next = pickNextAvailable(current);
} else if (!current || isProfileInCooldown(store, current)) {
next = pickFirstAvailable();
}
if (!next) return current;
const shouldPersist =
next !== sessionEntry.authProfileOverride ||
sessionEntry.authProfileOverrideSource !== "auto" ||
sessionEntry.authProfileOverrideCompactionCount !== compactionCount;
if (shouldPersist) {
sessionEntry.authProfileOverride = next;
sessionEntry.authProfileOverrideSource = "auto";
sessionEntry.authProfileOverrideCompactionCount = compactionCount;
sessionEntry.updatedAt = Date.now();
sessionStore[sessionKey] = sessionEntry;
if (storePath) {
await saveSessionStore(storePath, sessionStore);
}
}
return next;
}
export async function runPreparedReply(
params: RunPreparedReplyParams,
): Promise<ReplyPayload | ReplyPayload[] | undefined> {
@@ -314,7 +399,16 @@ export async function runPreparedReply(
resolvedQueue.mode === "followup" ||
resolvedQueue.mode === "collect" ||
resolvedQueue.mode === "steer-backlog";
const authProfileId = sessionEntry?.authProfileOverride;
const authProfileId = await resolveSessionAuthProfileOverride({
cfg,
provider,
agentDir,
sessionEntry,
sessionStore,
sessionKey,
storePath,
isNewSession,
});
const followupRun = {
prompt: queuedBody,
messageId: sessionCtx.MessageSid,

View File

@@ -217,6 +217,8 @@ export async function createModelSelectionState(params: {
const profile = store.profiles[sessionEntry.authProfileOverride];
if (!profile || profile.provider !== provider) {
delete sessionEntry.authProfileOverride;
delete sessionEntry.authProfileOverrideSource;
delete sessionEntry.authProfileOverrideCompactionCount;
sessionEntry.updatedAt = Date.now();
sessionStore[sessionKey] = sessionEntry;
if (storePath) {

View File

@@ -286,6 +286,8 @@ export async function agentCommand(
const profile = store.profiles[authProfileId];
if (!profile || profile.provider !== provider) {
delete entry.authProfileOverride;
delete entry.authProfileOverrideSource;
delete entry.authProfileOverrideCompactionCount;
entry.updatedAt = Date.now();
if (sessionStore && sessionKey) {
sessionStore[sessionKey] = entry;

View File

@@ -14,6 +14,12 @@ import { resolveEnvApiKey } from "../../agents/model-auth.js";
import { parseModelRef, resolveConfiguredModelRef } from "../../agents/model-selection.js";
import { CONFIG_PATH_CLAWDBOT, loadConfig } from "../../config/config.js";
import { getShellEnvAppliedKeys, shouldEnableShellEnvFallback } from "../../infra/shell-env.js";
import {
formatUsageWindowSummary,
loadProviderUsageSummary,
resolveUsageProviderId,
type UsageProviderId,
} from "../../infra/provider-usage.js";
import type { RuntimeEnv } from "../../runtime.js";
import { colorize, theme } from "../../terminal/theme.js";
import { shortenHomePath } from "../../utils.js";
@@ -402,6 +408,32 @@ export async function modelsStatusCommand(
return;
}
const usageByProvider = new Map<string, string>();
const usageProviders = Array.from(
new Set(
oauthProfiles
.map((profile) => resolveUsageProviderId(profile.provider))
.filter((provider): provider is UsageProviderId => Boolean(provider)),
),
);
if (usageProviders.length > 0) {
try {
const usageSummary = await loadProviderUsageSummary({
providers: usageProviders,
agentDir,
timeoutMs: 3500,
});
for (const snapshot of usageSummary.providers) {
const formatted = formatUsageWindowSummary(snapshot, { now: Date.now(), maxWindows: 2 });
if (formatted) {
usageByProvider.set(snapshot.provider, formatted);
}
}
} catch {
// ignore usage failures
}
}
const formatStatus = (status: string) => {
if (status === "ok") return colorize(rich, theme.success, "ok");
if (status === "static") return colorize(rich, theme.muted, "static");
@@ -410,19 +442,32 @@ export async function modelsStatusCommand(
return colorize(rich, theme.error, "expired");
};
const profilesByProvider = new Map<string, typeof oauthProfiles>();
for (const profile of oauthProfiles) {
const labelText = profile.label || profile.profileId;
const label = colorize(rich, theme.accent, labelText);
const status = formatStatus(profile.status);
const expiry =
profile.status === "static"
? ""
: profile.expiresAt
? ` expires in ${formatRemainingShort(profile.remainingMs)}`
: " expires unknown";
const source =
profile.source !== "store" ? colorize(rich, theme.muted, ` (${profile.source})`) : "";
runtime.log(`- ${label} ${status}${expiry}${source}`);
const current = profilesByProvider.get(profile.provider);
if (current) current.push(profile);
else profilesByProvider.set(profile.provider, [profile]);
}
for (const [provider, profiles] of profilesByProvider) {
const usageKey = resolveUsageProviderId(provider);
const usage = usageKey ? usageByProvider.get(usageKey) : undefined;
const usageSuffix = usage ? colorize(rich, theme.muted, ` usage: ${usage}`) : "";
runtime.log(`- ${colorize(rich, theme.heading, provider)}${usageSuffix}`);
for (const profile of profiles) {
const labelText = profile.label || profile.profileId;
const label = colorize(rich, theme.accent, labelText);
const status = formatStatus(profile.status);
const expiry =
profile.status === "static"
? ""
: profile.expiresAt
? ` expires in ${formatRemainingShort(profile.remainingMs)}`
: " expires unknown";
const source =
profile.source !== "store" ? colorize(rich, theme.muted, ` (${profile.source})`) : "";
runtime.log(` - ${label} ${status}${expiry}${source}`);
}
}
if (opts.check) runtime.exit(checkStatus);

View File

@@ -33,6 +33,8 @@ export type SessionEntry = {
providerOverride?: string;
modelOverride?: string;
authProfileOverride?: string;
authProfileOverrideSource?: "auto" | "user";
authProfileOverrideCompactionCount?: number;
groupActivation?: "mention" | "always";
groupActivationNeedsSystemIntro?: boolean;
sendPolicy?: "allow" | "deny";

View File

@@ -1,5 +1,5 @@
import { clampPercent } from "./provider-usage.shared.js";
import type { UsageSummary, UsageWindow } from "./provider-usage.types.js";
import type { ProviderUsageSnapshot, UsageSummary, UsageWindow } from "./provider-usage.types.js";
function formatResetRemaining(targetMs?: number, now?: number): string | null {
if (!targetMs) return null;
@@ -35,6 +35,28 @@ function formatWindowShort(window: UsageWindow, now?: number): string {
return `${remaining.toFixed(0)}% left (${window.label}${resetSuffix})`;
}
export function formatUsageWindowSummary(
snapshot: ProviderUsageSnapshot,
opts?: { now?: number; maxWindows?: number; includeResets?: boolean },
): string | null {
if (snapshot.error) return `error: ${snapshot.error}`;
if (snapshot.windows.length === 0) return null;
const now = opts?.now ?? Date.now();
const maxWindows =
typeof opts?.maxWindows === "number" && opts.maxWindows > 0
? Math.min(opts.maxWindows, snapshot.windows.length)
: snapshot.windows.length;
const includeResets = opts?.includeResets ?? false;
const windows = snapshot.windows.slice(0, maxWindows);
const parts = windows.map((window) => {
const remaining = clampPercent(100 - window.usedPercent);
const reset = includeResets ? formatResetRemaining(window.resetAt, now) : null;
const resetSuffix = reset ? `${reset}` : "";
return `${window.label} ${remaining.toFixed(0)}% left${resetSuffix}`;
});
return parts.join(" · ");
}
export function formatUsageSummaryLine(
summary: UsageSummary,
opts?: { now?: number; maxProviders?: number },

View File

@@ -1,4 +1,8 @@
export { formatUsageReportLines, formatUsageSummaryLine } from "./provider-usage.format.js";
export {
formatUsageReportLines,
formatUsageSummaryLine,
formatUsageWindowSummary,
} from "./provider-usage.format.js";
export { loadProviderUsageSummary } from "./provider-usage.load.js";
export { resolveUsageProviderId } from "./provider-usage.shared.js";
export type {