getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-16 21:03:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix expired OAuth tokens in long-running agent loops (#223)

Browse source 
Add getApiKey hook to AgentLoopConfig that resolves API keys dynamically
before each LLM call. This allows short-lived OAuth tokens (e.g. GitHub
Copilot, Anthropic OAuth) to be refreshed between turns when tool
execution takes a long time.

Previously, the API key was resolved once when ProviderTransport.run()
was called and passed as a static string to the agent loop. If the loop
ran for longer than the token lifetime (e.g. 30 minutes for Copilot),
subsequent LLM calls would fail with expired token errors.

Changes:
- Add getApiKey hook to AgentLoopConfig (packages/ai)
- Call getApiKey before each LLM call in streamAssistantResponse
- Update ProviderTransport to pass getApiKey instead of static apiKey
- Update web-ui ProviderTransport with same pattern
This commit is contained in:
Ahmed Kamal 2025-12-19 02:36:25 +02:00 • committed by GitHub
parent 139af12b37
commit 1167e84453
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 42 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

26
packages/agent/src/transports/ProviderTransport.ts
View file

@ -34,15 +34,7 @@ export class ProviderTransport implements AgentTransport {
this.options = options;
}
private async getModelAndKey(cfg: AgentRunConfig) {
let apiKey: string | undefined;
if (this.options.getApiKey) {
apiKey = await this.options.getApiKey(cfg.model.provider);
}
if (!apiKey) {
throw new Error(`No API key found for provider: ${cfg.model.provider}`);
}
private getModel(cfg: AgentRunConfig) {
let model = cfg.model;
if (this.options.corsProxyUrl && cfg.model.baseUrl) {
model = {
@ -50,8 +42,7 @@ export class ProviderTransport implements AgentTransport {
baseUrl: `${this.options.corsProxyUrl}/?url=${encodeURIComponent(cfg.model.baseUrl)}`,
};
}
return { model, apiKey };
return model;
}
private buildContext(messages: Message[], cfg: AgentRunConfig): AgentContext {
@ -62,19 +53,20 @@ export class ProviderTransport implements AgentTransport {
};
}
private buildLoopConfig(model: typeof cfg.model, apiKey: string, cfg: AgentRunConfig): AgentLoopConfig {
private buildLoopConfig(model: AgentRunConfig["model"], cfg: AgentRunConfig): AgentLoopConfig {
return {
model,
reasoning: cfg.reasoning,
apiKey,
// Resolve API key per assistant response (important for expiring OAuth tokens)
getApiKey: this.options.getApiKey ? (provider) => this.options.getApiKey?.(provider) : undefined,
getQueuedMessages: cfg.getQueuedMessages,
};
}
async *run(messages: Message[], userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
const { model, apiKey } = await this.getModelAndKey(cfg);
const model = this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, apiKey, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal)) {
yield ev;
@ -82,9 +74,9 @@ export class ProviderTransport implements AgentTransport {
}
async *continue(messages: Message[], cfg: AgentRunConfig, signal?: AbortSignal) {
const { model, apiKey } = await this.getModelAndKey(cfg);
const model = this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, apiKey, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoopContinue(context, pc, signal)) {
yield ev;

7
packages/ai/src/agent/agent-loop.ts
View file

@ -174,7 +174,12 @@ async function streamAssistantResponse(
// Use custom stream function if provided, otherwise use default streamSimple
const streamFunction = streamFn || streamSimple;
const response = await streamFunction(config.model, processedContext, { ...config, signal });
// Resolve API key for every assistant response (important for expiring tokens)
const resolvedApiKey =
(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
const response = await streamFunction(config.model, processedContext, { ...config, apiKey: resolvedApiKey, signal });
let partialMessage: AssistantMessage | null = null;
let addedPartial = false;

15
packages/ai/src/agent/types.ts
View file

@ -85,6 +85,21 @@ export interface QueuedMessage<TApp = Message> {
// Configuration for agent loop execution
export interface AgentLoopConfig extends SimpleStreamOptions {
model: Model<any>;
/**
* Optional hook to resolve an API key dynamically for each LLM call.
*
* This is useful for short-lived OAuth tokens (e.g. GitHub Copilot) that may
* expire during long-running tool execution phases.
*
* The agent loop will call this before each assistant response and pass the
* returned value as `apiKey` to `streamSimple()` (or a custom `streamFn`).
*
* If it returns `undefined`, the loop falls back to `config.apiKey`, and then
* to `streamSimple()`'s own provider key lookup (setApiKey/env vars).
*/
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
preprocessor?: (messages: AgentContext["messages"], abortSignal?: AbortSignal) => Promise<AgentContext["messages"]>;
getQueuedMessages?: <T>() => Promise<QueuedMessage<T>[]>;
}

20
packages/web-ui/src/agent/transports/ProviderTransport.ts
View file

@ -15,7 +15,7 @@ import type { AgentRunConfig, AgentTransport } from "./types.js";
* Uses CORS proxy only for providers that require it (Anthropic OAuth, Z-AI).
*/
export class ProviderTransport implements AgentTransport {
private async getModelAndKey(cfg: AgentRunConfig) {
private async getModel(cfg: AgentRunConfig) {
const apiKey = await getAppStorage().providerKeys.get(cfg.model.provider);
if (!apiKey) {
throw new Error("no-api-key");
@ -25,7 +25,7 @@ export class ProviderTransport implements AgentTransport {
const proxyUrl = await getAppStorage().settings.get<string>("proxy.url");
const model = applyProxyIfNeeded(cfg.model, apiKey, proxyEnabled ? proxyUrl || undefined : undefined);
return { model, apiKey };
return model;
}
private buildContext(messages: Message[], cfg: AgentRunConfig): AgentContext {
@ -36,19 +36,23 @@ export class ProviderTransport implements AgentTransport {
};
}
private buildLoopConfig(model: typeof cfg.model, apiKey: string, cfg: AgentRunConfig): AgentLoopConfig {
private buildLoopConfig(model: AgentRunConfig["model"], cfg: AgentRunConfig): AgentLoopConfig {
return {
model,
reasoning: cfg.reasoning,
apiKey,
// Resolve API key per assistant response (important for expiring OAuth tokens)
getApiKey: async (provider: string) => {
const key = await getAppStorage().providerKeys.get(provider);
return key ?? undefined; // Convert null to undefined for type compatibility
},
getQueuedMessages: cfg.getQueuedMessages,
};
}
async *run(messages: Message[], userMessage: Message, cfg: AgentRunConfig, signal?: AbortSignal) {
const { model, apiKey } = await this.getModelAndKey(cfg);
const model = await this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, apiKey, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoop(userMessage as unknown as UserMessage, context, pc, signal)) {
yield ev;
@ -56,9 +60,9 @@ export class ProviderTransport implements AgentTransport {
}
async *continue(messages: Message[], cfg: AgentRunConfig, signal?: AbortSignal) {
const { model, apiKey } = await this.getModelAndKey(cfg);
const model = await this.getModel(cfg);
const context = this.buildContext(messages, cfg);
const pc = this.buildLoopConfig(model, apiKey, cfg);
const pc = this.buildLoopConfig(model, cfg);
for await (const ev of agentLoopContinue(context, pc, signal)) {
yield ev;