diff --git a/packages/server/src/server/agent/providers/omp/agent.test.ts b/packages/server/src/server/agent/providers/omp/agent.test.ts index 75034dc99e..640ead5d30 100644 --- a/packages/server/src/server/agent/providers/omp/agent.test.ts +++ b/packages/server/src/server/agent/providers/omp/agent.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, test } from "vitest"; +import { describe, expect, onTestFinished, test } from "vitest"; import { setImmediate as waitForImmediate } from "node:timers/promises"; import { mkdtemp, mkdir, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; @@ -69,6 +69,51 @@ test("OMP import uses the runtime's custom agent directory without a configured expect.objectContaining({ providerHandleId: sessionFile }), ]); }); +test("OMP resumes a session whose model was removed on the model OMP falls back to", async () => { + const runtime = new FakeOmp(); + runtime.removeModel("9router/deepseek-v4-flash"); + runtime.queueSessionSetup((session) => { + session.state = { ...session.state, model: { provider: "openrouter", id: "fallback" } }; + }); + const client = new OmpAgentClient({ logger: createTestLogger(), runtime }); + + const session = await client.resumeSession({ + provider: "omp", + sessionId: "omp-session-1", + nativeHandle: "/tmp/omp-session.jsonl", + metadata: { cwd: "/workspace/project", model: "9router/deepseek-v4-flash" }, + }); + onTestFinished(() => session.close()); + + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ model: "openrouter/fallback" }); + expect(session.describePersistence()?.metadata?.model).toBe("openrouter/fallback"); +}); + +test("OMP resumes a session on the requested model when it differs from the session's", async () => { + const runtime = new FakeOmp(); + const requestedModel = { provider: "openrouter", id: "requested" }; + runtime.queueSessionSetup((session) => { + session.state = { ...session.state, model: { provider: "openrouter", id: "recorded" } }; + session.models = [requestedModel]; + session.setModelResult = requestedModel; + }); + const client = new OmpAgentClient({ logger: createTestLogger(), runtime }); + + const session = await client.resumeSession({ + provider: "omp", + sessionId: "omp-session-1", + nativeHandle: "/tmp/omp-session.jsonl", + metadata: { cwd: "/workspace/project", model: "openrouter/requested" }, + }); + onTestFinished(() => session.close()); + + expect(runtime.latestSession().setModelRequests).toEqual([ + { provider: "openrouter", modelId: "requested" }, + ]); + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ model: "openrouter/requested" }); + expect(session.describePersistence()?.metadata?.model).toBe("openrouter/requested"); +}); + class ManualIdleScheduler implements OmpProviderIdleScheduler { private readonly retries: Array<() => void> = []; private readonly waiters: Array<{ count: number; resolve: () => void }> = []; diff --git a/packages/server/src/server/agent/providers/omp/agent.ts b/packages/server/src/server/agent/providers/omp/agent.ts index ce29ea6d29..a04b4c8fb0 100644 --- a/packages/server/src/server/agent/providers/omp/agent.ts +++ b/packages/server/src/server/agent/providers/omp/agent.ts @@ -445,7 +445,6 @@ function buildResumeStartInput(input: { protocolMode: "rpc-ui", env: input.launchContext?.env, session: input.sessionFile, - model: input.resumeConfig.model, thinkingOptionId: normalizeOmpThinkingOption(input.resumeConfig.thinkingOptionId) ?? undefined, ...(input.launchMode.modeId ? { modeId: input.launchMode.modeId } : {}), ...(input.launchMode.extraArgs ? { extraArgs: input.launchMode.extraArgs } : {}), @@ -2373,16 +2372,17 @@ export class OmpAgentClient implements AgentClient { resumeConfig.config, startInput.env, ); - const initialState = await this.restoreFastMode( - runtimeSession, - resumeConfig.config, - await runtimeSession.getState(), - ); + const resumedState = await this.applyResumeModel(runtimeSession, resumeConfig.model); + const config = { + ...resumeConfig.config, + model: modelToId(resumedState.model) ?? resumeConfig.config.model, + }; + const initialState = await this.restoreFastMode(runtimeSession, config, resumedState); return new OmpAgentSession({ runtimeSession, hostTools, - restartRuntime: this.buildRestartRuntime(startInput, resumeConfig.config, launchContext), - config: resumeConfig.config, + restartRuntime: this.buildRestartRuntime(startInput, config, launchContext), + config, initialState, currentModeId: launchMode.modeId, logger: this.logger, @@ -2400,6 +2400,35 @@ export class OmpAgentClient implements AgentClient { } } + // OMP resumes a session on the model it recorded, or on its default when that model + // is gone. Switching afterwards keeps a removed model from blocking the resume. + private async applyResumeModel( + runtimeSession: OmpRuntimeSession, + requestedModel: string | undefined, + ): Promise { + const state = await runtimeSession.getState(); + const reference = parseModelReference(requestedModel ?? null); + if (!reference?.provider) { + return state; + } + const { provider, id } = reference; + const isRequested = (model: OmpModel | null | undefined) => + model?.provider === provider && model.id === id; + if (isRequested(state.model)) { + return state; + } + const availableModels = await runtimeSession.getAvailableModels(); + if (!availableModels.some(isRequested)) { + this.logger.warn( + { requestedModel, sessionModel: modelToId(state.model) }, + "OMP resumed on the session's model because the requested model is unavailable", + ); + return state; + } + await runtimeSession.setModel(provider, id); + return runtimeSession.getState(); + } + private buildRestartRuntime( startInput: OmpStartSessionInput, config: AgentSessionConfig, diff --git a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts index c54defed6f..403a389647 100644 --- a/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts +++ b/packages/server/src/server/agent/providers/omp/test-utils/fake-omp.ts @@ -54,6 +54,8 @@ export class FakeOmp implements OmpRuntime { private readonly sessions: FakeOmpSession[] = []; private readonly command: [string, ...string[]]; private readonly queuedCommands: OmpRpcSlashCommand[][] = []; + private readonly queuedSessionSetups: Array<(session: FakeOmpSession) => void> = []; + private readonly removedModels = new Set(); private nextStartError: Error | null = null; private nextSessionId = 0; private readonly sessionIdsByFile = new Map(); @@ -75,16 +77,20 @@ export class FakeOmp implements OmpRuntime { session: input, }); this.recordedLaunches.push(launch); + if (launch.model && this.removedModels.has(launch.model)) { + throw new Error(`Model "${launch.model}" not found`); + } const sessionId = (launch.session ? this.sessionIdsByFile.get(launch.session) : undefined) ?? `omp-session-${++this.nextSessionId}`; - const session = new FakeOmpSession(launch, sessionId); + const session = new FakeOmpSession(launch, sessionId, this.removedModels); if (!launch.noSession) this.sessionIdsByFile.set(session.state.sessionFile, sessionId); session.commands = this.queuedCommands.shift() ?? []; for (const [level, error] of this.queuedSubagentSubscriptionErrors) { session.subagentSubscriptionErrors.set(level, error); } this.queuedSubagentSubscriptionErrors.clear(); + this.queuedSessionSetups.shift()?.(session); this.sessions.push(session); return session; } @@ -93,6 +99,15 @@ export class FakeOmp implements OmpRuntime { this.queuedCommands.push(commands); } + queueSessionSetup(setup: (session: FakeOmpSession) => void): void { + this.queuedSessionSetups.push(setup); + } + + // OMP refuses to launch with, or switch to, a model it no longer knows. + removeModel(modelId: string): void { + this.removedModels.add(modelId); + } + failNextStart(error: Error): void { this.nextStartError = error; } @@ -174,7 +189,11 @@ export class FakeOmpSession implements OmpRuntimeSession { private activeHeldPrompt: { promise: Promise; reject: (error: Error) => void } | null = null; - constructor(launch: OmpRuntimeLaunch, sessionId = "omp-session-1") { + constructor( + launch: OmpRuntimeLaunch, + sessionId = "omp-session-1", + private readonly removedModels: ReadonlySet = new Set(), + ) { this.environment = launch.env ?? {}; this.state = { model: null, @@ -333,9 +352,13 @@ export class FakeOmpSession implements OmpRuntimeSession { async setModel(provider: string, modelId: string): Promise { this.setModelRequests.push({ provider, modelId }); + if (this.removedModels.has(`${provider}/${modelId}`)) { + throw new Error(`Model not found: ${provider}/${modelId}`); + } if (!this.setModelResult) { throw new Error("FakeOmp setModel requires setModelResult to be scripted"); } + this.state = { ...this.state, model: this.setModelResult }; return this.setModelResult; } diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index b87e0bace0..8035d59830 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -1528,8 +1528,6 @@ describe("PiRpcAgentSession", () => { "pi", "--mode", "rpc", - "--model", - "openrouter/model-a", "--thinking", "high", "--session", @@ -1539,6 +1537,87 @@ describe("PiRpcAgentSession", () => { ]); }); + test("resumes a session whose model was removed on the model Pi falls back to", async () => { + const pi = new FakePi(); + pi.removeModel("9router/deepseek-v4-flash"); + pi.queueSessionSetup((session) => { + session.state = { ...session.state, model: RESTRICTED_THINKING_MODEL }; + }); + + const session = await createClient(pi).resumeSession({ + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project", model: "9router/deepseek-v4-flash" }, + }); + onTestFinished(() => session.close()); + + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ + model: "kimi-coding/kimi-k3", + }); + expect(session.describePersistence()?.metadata?.model).toBe("kimi-coding/kimi-k3"); + }); + + test("resumes a session on the requested model when it differs from the session's", async () => { + const pi = new FakePi(); + const requestedModel = { ...RESTRICTED_THINKING_MODEL, provider: "openrouter", id: "a" }; + pi.queueSessionSetup((session) => { + session.state = { ...session.state, model: RESTRICTED_THINKING_MODEL }; + session.models = [RESTRICTED_THINKING_MODEL, requestedModel]; + session.setModelResult = requestedModel; + }); + + const session = await createClient(pi).resumeSession({ + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project", model: "openrouter/a" }, + }); + onTestFinished(() => session.close()); + + expect(pi.latestSession().setModelRequests).toEqual([{ provider: "openrouter", modelId: "a" }]); + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ model: "openrouter/a" }); + expect(session.describePersistence()?.metadata?.model).toBe("openrouter/a"); + }); + + test("resumes on a requested model written as provider:id", async () => { + const pi = new FakePi(); + const requestedModel = { ...RESTRICTED_THINKING_MODEL, provider: "openrouter", id: "a" }; + pi.queueSessionSetup((session) => { + session.state = { ...session.state, model: RESTRICTED_THINKING_MODEL }; + session.models = [RESTRICTED_THINKING_MODEL, requestedModel]; + session.setModelResult = requestedModel; + }); + + const session = await createClient(pi).resumeSession({ + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project", model: "openrouter:a" }, + }); + onTestFinished(() => session.close()); + + await expect(session.getRuntimeInfo()).resolves.toMatchObject({ model: "openrouter/a" }); + }); + + test("fails the resume when switching to an available model fails", async () => { + const pi = new FakePi(); + const requestedModel = { ...RESTRICTED_THINKING_MODEL, provider: "openrouter", id: "a" }; + pi.queueSessionSetup((session) => { + session.state = { ...session.state, model: RESTRICTED_THINKING_MODEL }; + session.models = [RESTRICTED_THINKING_MODEL, requestedModel]; + }); + + await expect( + createClient(pi).resumeSession({ + provider: "pi", + sessionId: "pi-session-1", + nativeHandle: "/tmp/native-pi-session", + metadata: { cwd: "/workspace/project", model: "openrouter/a" }, + }), + ).rejects.toThrow("FakePi setModel requires setModelResult to be scripted"); + }); + test("adopts Pi's clamped thinking level when resuming a session", async () => { const pi = new FakePi(); pi.queueSessionSetup((session) => { @@ -1715,8 +1794,6 @@ describe("PiRpcAgentSession", () => { "pi", "--mode", "rpc", - "--model", - "openrouter/model-a", "--thinking", "high", "--session", @@ -2576,6 +2653,11 @@ describe("PiRpcAgentClient", () => { ); const pi = new FakePi(); pi.queueSessionSetup((session) => { + session.state.model = { + ...RESTRICTED_THINKING_MODEL, + provider: "openrouter", + id: "anthropic/claude-sonnet-4.5", + }; session.state.thinkingLevel = "high"; }); const client = new PiRpcAgentClient({ @@ -2594,8 +2676,6 @@ describe("PiRpcAgentClient", () => { "pi", "--mode", "rpc", - "--model", - "openrouter/anthropic/claude-sonnet-4.5", "--thinking", "high", "--session", diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 5a281ae0a6..5c9c7bb949 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -495,7 +495,6 @@ function buildResumeStartInput(input: { cwd: input.resumeConfig.cwd, env: input.launchContext?.env, session: input.sessionFile, - model: input.resumeConfig.model, thinkingOptionId: normalizePiThinkingOption(input.resumeConfig.thinkingOptionId) ?? undefined, mcpConfigPath: input.mcpConfigFile?.path, extensionPaths: input.paseoExtension ? [input.paseoExtension.path] : undefined, @@ -2590,10 +2589,14 @@ export class PiRpcAgentClient implements AgentClient { throw error; } try { + const initialState = await this.applyResumeModel(runtimeSession, resumeConfig.model); return new PiRpcAgentSession({ runtimeSession, - config: resumeConfig.config, - initialState: await runtimeSession.getState(), + config: { + ...resumeConfig.config, + model: modelToId(initialState.model) ?? resumeConfig.config.model, + }, + initialState, capabilities: capabilitiesForSession(mcp !== null), cleanup: combineCleanup([mcpConfigFile?.cleanup, paseoExtension?.cleanup]), extensionTimeoutMs: providerOptions.extensionTimeoutMs, @@ -2608,6 +2611,35 @@ export class PiRpcAgentClient implements AgentClient { } } + // Pi resumes a session on the model it recorded, or on its default when that model + // is gone. Switching afterwards keeps a removed model from blocking the resume. + private async applyResumeModel( + runtimeSession: PiRuntimeSession, + requestedModel: string | undefined, + ): Promise { + const state = await runtimeSession.getState(); + const reference = parseModelReference(requestedModel ?? null); + if (!reference?.provider) { + return state; + } + const { provider, id } = reference; + const isRequested = (model: PiModel | null | undefined) => + model?.provider === provider && model.id === id; + if (isRequested(state.model)) { + return state; + } + const availableModels = await runtimeSession.getAvailableModels(); + if (!availableModels.some(isRequested)) { + this.logger.warn( + { requestedModel, sessionModel: modelToId(state.model) }, + "Pi resumed on the session's model because the requested model is unavailable", + ); + return state; + } + await runtimeSession.setModel(provider, id); + return runtimeSession.getState(); + } + async fetchCatalog( options: FetchCatalogOptions, context?: ProviderRefreshContext, diff --git a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts index 5b52b1eda4..b8c7930bc4 100644 --- a/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts +++ b/packages/server/src/server/agent/providers/pi/test-utils/fake-pi.ts @@ -57,6 +57,7 @@ export class FakePi implements PiRuntime { private readonly command: [string, ...string[]]; private readonly queuedCommands: PiRpcSlashCommand[][] = []; private readonly queuedSessionSetups: Array<(session: FakePiSession) => void> = []; + private readonly removedModels = new Set(); constructor(command: [string, ...string[]] = ["pi"]) { this.command = command; @@ -68,7 +69,10 @@ export class FakePi implements PiRuntime { session: input, }); this.recordedLaunches.push(launch); - const session = new FakePiSession(launch); + if (launch.model && this.removedModels.has(launch.model)) { + throw new Error(`Model "${launch.model}" not found`); + } + const session = new FakePiSession(launch, this.removedModels); session.commands = this.queuedCommands.shift() ?? []; this.queuedSessionSetups.shift()?.(session); this.sessions.push(session); @@ -83,6 +87,11 @@ export class FakePi implements PiRuntime { this.queuedSessionSetups.push(setup); } + // Pi refuses to launch with, or switch to, a model it no longer knows. + removeModel(modelId: string): void { + this.removedModels.add(modelId); + } + latestSession(): FakePiSession { const session = this.sessions.at(-1); if (!session) { @@ -145,7 +154,10 @@ export class FakePiSession implements PiRuntimeSession { private activeHeldPrompt: { promise: Promise; reject: (error: Error) => void } | null = null; - constructor(launch: PiRuntimeLaunch) { + constructor( + launch: PiRuntimeLaunch, + private readonly removedModels: ReadonlySet = new Set(), + ) { this.environment = launch.env ?? {}; this.state = { model: null, @@ -267,6 +279,9 @@ export class FakePiSession implements PiRuntimeSession { async setModel(provider: string, modelId: string): Promise { this.setModelRequests.push({ provider, modelId }); + if (this.removedModels.has(`${provider}/${modelId}`)) { + throw new Error(`Model not found: ${provider}/${modelId}`); + } if (!this.setModelResult) { throw new Error("FakePi setModel requires setModelResult to be scripted"); }