From 4ccd4f81fce8acd34388383a96a7d538d58e29cf Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Fri, 17 Jul 2026 21:32:02 +0200 Subject: [PATCH] Migrate piSessionService to ModelRuntime (slice 4) Pass modelRuntime to createAgentSessionServices instead of authStorage + modelRegistry; carry ModelRuntime on PiAgentSession; make modelRuntime a required PiSessionService dependency (sessiond already injects auth.runtime). Switch anthropicSubscriptionWarning to readStoredCredential, and rederive model reads (availableModels/setModel/syncCurrentModelAuthWarning) via the runtime (getAvailableSnapshot/getModel/hasConfiguredAuth). sessiond.ts and piSessionService.ts now typecheck; only slice-5 test/support files remain. --- src/server/sessions/piSessionService.ts | 50 ++++++++++++------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index 6b26a43..5899e3c 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -1,20 +1,21 @@ import { statSync } from "node:fs"; +import { join } from "node:path"; import { open, readFile, writeFile } from "node:fs/promises"; import type { ImageContent } from "@earendil-works/pi-ai"; import type { StreamFn } from "@earendil-works/pi-agent-core"; import { - AuthStorage, createAgentSessionFromServices, createAgentSessionRuntime, createAgentSessionServices, createEditToolDefinition, defineTool, - ModelRegistry, + readStoredCredential, SessionManager, type AgentSessionRuntimeDiagnostic, type AgentSessionServices, type CreateAgentSessionRuntimeFactory, type EditToolDetails, + type ModelRuntime, type ResourceDiagnostic, } from "@earendil-works/pi-coding-agent"; import type { ClientArchiveSessionsResponse, ClientCommand, ClientCommandResult, ClientMessagePage, ClientSession, ClientSessionCleanupExecuteResponse, ClientSessionCleanupPreviewResponse, ClientSessionModel, ClientSessionStatus, ClientThinkingLevel, SessionStreamSnapshot, SessionUiEvent } from "../types.js"; @@ -26,7 +27,6 @@ import { SessionCommandService } from "./sessionCommandService.js"; import { SessionArchiveStore, type ArchivedSessionRecord, type ArchiveSessionInput } from "./sessionArchiveStore.js"; import { findArchiveCandidateByIdOrPrefix, planSessionArchiveTree, type SessionArchiveTreeCandidate } from "./sessionArchiveTree.js"; import type { ActiveSession } from "./sessionRuntimeStore.js"; -import { createModelRegistryForAgentDir, type AuthChange } from "./authService.js"; import { deterministicSessionName, fallbackSessionName, generateShortSessionName } from "./sessionNameGenerator.js"; import { computeEditPreview, type EditPreviewResult } from "./editPreview.js"; import { attachmentsToInlineImages, saveAttachmentsToWorkspace } from "./attachmentService.js"; @@ -34,6 +34,7 @@ import { parsePromptAttachments } from "../../shared/promptAttachments.js"; import type { SavedPromptAttachment, SessionBulkArchiveResponse, SessionBulkDeleteArchivedResponse, SessionBulkFailure, SessionBulkMutationRef, SessionWarning } from "../../shared/apiTypes.js"; import type { SessionRouteLookup, SessionRouteRef, SessionRouteService } from "./sessionService.js"; +import { type AuthChange } from "./authService.js"; import { canonicalizeStoredCwd, cwdPathsEqual } from "../workingDirectory.js"; import type { WorkspaceActivityService } from "../activity/workspaceActivityService.js"; import { createSpawnSessionToolDefinition, type SpawnSessionInvocation, type SpawnSessionResult } from "./spawnSessionTool.js"; @@ -171,7 +172,6 @@ interface BulkDeletePlanItem { } type AgentModel = NonNullable; -type ModelRegistryInstance = ReturnType; export interface PiSessionManager { getCwd(): string; @@ -208,7 +208,7 @@ interface PiExtensionBindings { } export interface PiAgentSession { - modelRegistry: ModelRegistryInstance; + modelRuntime: ModelRuntime; /** * Narrow read/write of the SDK `SettingsManager`, exposing only the warning * suppression flags consumed here (e.g. `anthropicExtraUsage`). Used to gate @@ -383,11 +383,12 @@ const ANTHROPIC_EXTRA_USAGE_DISMISS_ID = "anthropicExtraUsage"; * synchronous live status computation. */ export function anthropicSubscriptionWarning( - session: Pick, + session: Pick, + authPath?: string, ): SessionWarning | undefined { if (session.settingsManager.getWarnings().anthropicExtraUsage === false) return undefined; if (session.model?.provider !== "anthropic") return undefined; - const credential = session.modelRegistry.authStorage.get("anthropic"); + const credential = readStoredCredential("anthropic", authPath); if (credential === undefined) return undefined; const isSubscriptionAuth = credential.type === "oauth" ? true @@ -487,14 +488,13 @@ export function createPiWebCustomToolDefinitions( } function createDefaultRuntimeFactory( - authStorage: AuthStorage, - modelRegistry: ModelRegistryInstance, + modelRuntime: ModelRuntime, sessionManagers: Pick, spawn?: SpawnSessionFn, subsessions?: SubsessionToolDeps, ): PiWebCreateAgentSessionRuntimeFactory { return async ({ cwd, agentDir, sessionManager, sessionStartEvent, initialModel, delegationToolsEnabled }) => { - const services = await createAgentSessionServices({ cwd, agentDir, authStorage, modelRegistry }); + const services = await createAgentSessionServices({ cwd, agentDir, modelRuntime }); const resolvedDelegationToolsEnabled = delegationToolsEnabled ?? await sessionAllowsDelegationTools(sessionManager, sessionManagers); const customTools = createPiWebCustomToolDefinitions(cwd, resolvedDelegationToolsEnabled, spawn, subsessions); @@ -539,7 +539,7 @@ export interface PiSessionServiceDependencies { archiveStore?: SessionArchiveRepository; createRuntime?: PiWebCreateAgentSessionRuntimeFactory; createAgentRuntime?: CreateAgentRuntime; - modelRegistry?: ModelRegistryInstance; + modelRuntime: ModelRuntime; heartbeatIntervalMs?: number; workspaceActivity?: Pick; /** @@ -589,7 +589,7 @@ export class PiSessionService implements SessionRouteService { private readonly sessionManager: PiSessionManagerGateway; private readonly createRuntime: PiWebCreateAgentSessionRuntimeFactory; private readonly createAgentRuntime: CreateAgentRuntime; - private readonly modelRegistry: ModelRegistryInstance; + private readonly modelRuntime: ModelRuntime; private readonly workspaceActivity: Pick | undefined; private readonly spawnTargets: SpawnTargetResolver | undefined; private readonly logger: PiSessionLogger; @@ -599,7 +599,7 @@ export class PiSessionService implements SessionRouteService { this.archiveStore = deps.archiveStore ?? new SessionArchiveStore(); this.agentDir = deps.agentDir; this.sessionManager = deps.sessionManager; - this.modelRegistry = deps.modelRegistry ?? createModelRegistryForAgentDir(this.agentDir); + this.modelRuntime = deps.modelRuntime; this.spawnTargets = deps.spawnTargets; this.logger = deps.logger ?? noopLogger; this.now = deps.now ?? (() => new Date()); @@ -607,8 +607,7 @@ export class PiSessionService implements SessionRouteService { // also require the spawn capability (they share its project-scope resolver). const subsessionsActive = this.spawnTargets !== undefined && deps.subsessionsEnabled === true; this.createRuntime = deps.createRuntime ?? createDefaultRuntimeFactory( - this.modelRegistry.authStorage, - this.modelRegistry, + this.modelRuntime, this.sessionManager, this.spawnTargets === undefined ? undefined : (input) => this.spawnSession(input), !subsessionsActive ? undefined : { @@ -1159,22 +1158,22 @@ export class PiSessionService implements SessionRouteService { async availableModels(ref: PiSessionLookup): Promise { const session = await this.getOrOpen(ref); - session.modelRegistry.refresh(); + await session.modelRuntime.refresh(); const models = session.scopedModels.length > 0 ? session.scopedModels.map((scoped) => scoped.model) - : session.modelRegistry.getAvailable(); + : session.modelRuntime.getAvailableSnapshot(); return models.map(modelToClientModel); } async setModel(ref: PiSessionLookup, provider: string, modelId: string): Promise { await this.assertWritable(ref); const session = await this.getOrOpen(ref); - session.modelRegistry.refresh(); + await session.modelRuntime.refresh(); const candidates = session.scopedModels.length > 0 ? session.scopedModels.map((scoped) => scoped.model) - : session.modelRegistry.getAvailable(); + : session.modelRuntime.getAvailableSnapshot(); const model = candidates.find((candidate) => candidate.provider === provider && candidate.id === modelId) - ?? session.modelRegistry.find(provider, modelId); + ?? session.modelRuntime.getModel(provider, modelId); if (model === undefined) throw new Error(`Model not found: ${provider}/${modelId}`); await session.setModel(model); this.publishActivity(session, `model: ${model.id}`, "idle", model.provider); @@ -2020,10 +2019,11 @@ export class PiSessionService implements SessionRouteService { } applyAuthChange(change: AuthChange = {}): void { - this.modelRegistry.refresh(); + // The shared model runtime is refreshed by AuthService before it emits the + // change (and every session shares that runtime), so no refresh is needed + // here — this keeps the subscribe callback synchronous. for (const active of this.active.values()) { const { session } = active.runtime; - session.modelRegistry.refresh(); this.syncCurrentModelAuthWarning(session, change.removedProviderId); this.publishStatus(session); } @@ -2034,9 +2034,9 @@ export class PiSessionService implements SessionRouteService { if (model === undefined) return; if (model.provider === "unknown" && model.id === "unknown") return; const warningKey = authLossWarningKey(session.sessionId, model.provider, model.id); - const registered = session.modelRegistry.find(model.provider, model.id); + const registered = session.modelRuntime.getModel(model.provider, model.id); if (registered === undefined) return; - if (session.modelRegistry.hasConfiguredAuth(registered)) { + if (session.modelRuntime.hasConfiguredAuth(model.provider)) { this.authLossWarnings.delete(warningKey); return; } @@ -2182,7 +2182,7 @@ export class PiSessionService implements SessionRouteService { private warningsForSession(session: PiAgentSession): SessionWarning[] { const runtime = this.active.get(session.sessionId)?.runtime; const warnings = runtime === undefined ? [] : collectRuntimeWarnings(runtime); - const anthropic = anthropicSubscriptionWarning(session); + const anthropic = anthropicSubscriptionWarning(session, join(this.agentDir, "auth.json")); if (anthropic !== undefined) warnings.push(anthropic); return warnings; }