ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(coding-agent): replace model registry with model runtime

Browse Source 
Move provider auth and OAuth flows onto pi-ai Models, compose models.json and extension overlays through ModelRuntime, and retain ModelRegistry as an extension compatibility facade.
This commit is contained in:
Mario Zechner
2026-07-14 17:48:45 +02:00
parent 6731a0ba9e
commit 9993c96907
133 changed files with 5103 additions and 4340 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/agent-session.ts
+68 -35
View File
@@ -24,7 +24,15 @@ import type {
PrepareNextTurnContext,
ThinkingLevel,
} from "@earendil-works/pi-agent-core";
import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@earendil-works/pi-ai/compat";
import type {
AssistantMessage,
AuthResult,
ImageContent,
Message,
Model,
ProviderHeaders,
TextContent,
} from "@earendil-works/pi-ai/compat";
import {
clampThinkingLevel,
cleanupSessionResources,
@@ -83,7 +91,8 @@ import {
} from "./extensions/index.ts";
import { emitSessionShutdownEvent } from "./extensions/runner.ts";
import type { BashExecutionMessage, CustomMessage } from "./messages.ts";
import type { ModelRegistry } from "./model-registry.ts";
import { ModelRegistry } from "./model-registry.ts";
import type { ModelRuntime } from "./model-runtime.ts";
import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts";
import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts";
import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts";
@@ -159,6 +168,12 @@ export type AgentSessionEventListener = (event: AgentSessionEvent) => void;
// Types
// ============================================================================
function withoutDeletedHeaders(headers: ProviderHeaders | undefined): Record<string, string> | undefined {
return headers
? Object.fromEntries(Object.entries(headers).filter((entry): entry is [string, string] => entry[1] !== null))
: undefined;
}
export interface AgentSessionConfig {
agent: Agent;
sessionManager: SessionManager;
@@ -170,8 +185,8 @@ export interface AgentSessionConfig {
resourceLoader: ResourceLoader;
/** SDK custom tools registered outside extensions */
customTools?: ToolDefinition[];
/** Model registry for API key resolution and model discovery */
modelRegistry: ModelRegistry;
/** Canonical model/auth runtime used by coding-agent internals. */
modelRuntime: ModelRuntime;
/** Initial active built-in tool names. Default: [read, bash, edit, write] */
initialActiveToolNames?: string[];
/** Optional allowlist of tool names. When provided, only these tool names are exposed. */
@@ -325,8 +340,7 @@ export class AgentSession {
private _extensionErrorListener?: ExtensionErrorListener;
private _extensionErrorUnsubscriber?: () => void;
// Model registry for API key resolution
private _modelRegistry: ModelRegistry;
private _modelRuntime: ModelRuntime;
// Tool registry for extension getTools/setTools
private _toolRegistry: Map<string, AgentTool> = new Map();
@@ -347,7 +361,7 @@ export class AgentSession {
this._resourceLoader = config.resourceLoader;
this._customTools = config.customTools ?? [];
this._cwd = config.cwd;
this._modelRegistry = config.modelRegistry;
this._modelRuntime = config.modelRuntime;
this._extensionRunnerRef = config.extensionRunnerRef;
this._initialActiveToolNames = config.initialActiveToolNames;
this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
@@ -367,9 +381,8 @@ export class AgentSession {
});
}
/** Model registry for API key resolution and model discovery */
get modelRegistry(): ModelRegistry {
return this._modelRegistry;
get modelRuntime(): ModelRuntime {
return this._modelRuntime;
}
private async _getRequiredRequestAuth(model: Model<any>): Promise<{
@@ -377,18 +390,25 @@ export class AgentSession {
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
if (!result.ok) {
if (result.error.startsWith("No API key found")) {
let result: AuthResult | undefined;
try {
result = await this._modelRuntime.getAuth(model);
} catch (error) {
const cause = error instanceof Error ? error.cause : undefined;
if (cause instanceof Error && cause.message === "authHeader requires a resolved API key") {
throw new Error(formatNoApiKeyFoundMessage(model.provider));
}
throw new Error(result.error);
throw error;
}
if (result.apiKey) {
return { apiKey: result.apiKey, headers: result.headers, env: result.env };
if (result?.auth.apiKey) {
return {
apiKey: result.auth.apiKey,
headers: withoutDeletedHeaders(result.auth.headers),
env: result.env,
};
}
const isOAuth = this._modelRegistry.isUsingOAuth(model);
const isOAuth = this._modelRuntime.isUsingOAuth(model.provider);
if (isOAuth) {
throw new Error(
`Authentication failed for "${model.provider}". ` +
@@ -408,8 +428,14 @@ export class AgentSession {
return this._getRequiredRequestAuth(model);
}
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {};
try {
const result = await this._modelRuntime.getAuth(model);
return result
? { apiKey: result.auth.apiKey, headers: withoutDeletedHeaders(result.auth.headers), env: result.env }
: {};
} catch {
return {};
}
}
/**
@@ -1141,8 +1167,11 @@ export class AgentSession {
throw new Error(formatNoModelSelectedMessage());
}
if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
const hasConfiguredAuth =
this._modelRuntime.hasConfiguredAuth(this.model.provider) ||
(await this._modelRuntime.checkAuth(this.model.provider)) !== undefined;
if (!hasConfiguredAuth) {
const isOAuth = this._modelRuntime.isUsingOAuth(this.model.provider);
if (isOAuth) {
throw new Error(
`Authentication failed for "${this.model.provider}". ` +
@@ -1535,7 +1564,7 @@ export class AgentSession {
* @throws Error if no auth is configured for the model
*/
async setModel(model: Model<any>): Promise<void> {
if (!this._modelRegistry.hasConfiguredAuth(model)) {
if (!(await this._modelRuntime.checkAuth(model.provider))) {
throw new Error(`No API key for ${model.provider}/${model.id}`);
}
@@ -1565,7 +1594,13 @@ export class AgentSession {
}
private async _cycleScopedModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
const scopedModels = this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
const checks = await Promise.all(
this._scopedModels.map(async (scoped) => ({
scoped,
auth: await this._modelRuntime.checkAuth(scoped.model.provider),
})),
);
const scopedModels = checks.filter(({ auth }) => auth !== undefined).map(({ scoped }) => scoped);
if (scopedModels.length <= 1) return undefined;
const currentModel = this.model;
@@ -1594,7 +1629,7 @@ export class AgentSession {
}
private async _cycleAvailableModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
const availableModels = await this._modelRegistry.getAvailable();
const availableModels = await this._modelRuntime.getAvailable();
if (availableModels.length <= 1) return undefined;
const currentModel = this.model;
@@ -2004,12 +2039,10 @@ export class AgentSession {
let headers: Record<string, string> | undefined;
let env: Record<string, string> | undefined;
if (this.agent.streamFn === streamSimple) {
const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
if (!authResult.ok || !authResult.apiKey) {
return false;
}
apiKey = authResult.apiKey;
headers = authResult.headers;
const authResult = await this._modelRuntime.getAuth(this.model);
if (!authResult?.auth.apiKey) return false;
apiKey = authResult.auth.apiKey;
headers = withoutDeletedHeaders(authResult.auth.headers);
env = authResult.env;
} else {
({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model));
@@ -2267,7 +2300,7 @@ export class AgentSession {
return;
}
const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id);
const refreshedModel = this._modelRuntime.getModel(currentModel.provider, currentModel.id);
if (!refreshedModel || refreshedModel === currentModel) {
return;
}
@@ -2343,7 +2376,7 @@ export class AgentSession {
refreshTools: () => this._refreshToolRegistry(),
getCommands,
setModel: async (model) => {
if (!this.modelRegistry.hasConfiguredAuth(model)) return false;
if (!this._modelRuntime.hasConfiguredAuth(model.provider)) return false;
await this.setModel(model);
return true;
},
@@ -2383,11 +2416,11 @@ export class AgentSession {
},
{
registerProvider: (name, config) => {
this._modelRegistry.registerProvider(name, config);
this._modelRuntime.registerProvider(name, config);
this._refreshCurrentModelFromRegistry();
},
unregisterProvider: (name) => {
this._modelRegistry.unregisterProvider(name);
this._modelRuntime.unregisterProvider(name);
this._refreshCurrentModelFromRegistry();
},
},
@@ -2523,7 +2556,7 @@ export class AgentSession {
extensionsResult.runtime,
this._cwd,
this.sessionManager,
this._modelRegistry,
new ModelRegistry(this._modelRuntime),
);
if (this._extensionRunnerRef) {
this._extensionRunnerRef.current = this._extensionRunner;