From ebb4eaf5ee75c3b4accb55fcd07cb56c6e039631 Mon Sep 17 00:00:00 2001 From: Shakker Date: Sun, 2 Aug 2026 23:30:03 +0100 Subject: [PATCH] fix: retire model state across owners --- ui/src/lib/sessions/index.test.ts | 10 +- ui/src/lib/sessions/session-mutations.ts | 33 ++++--- ui/src/pages/chat/chat-commands.ts | 7 +- ui/src/pages/chat/chat-send.test.ts | 65 ++++++------- ui/src/pages/chat/chat-session.ts | 5 +- ui/src/pages/chat/chat-view.test.ts | 113 ++++++++++++----------- 6 files changed, 121 insertions(+), 112 deletions(-) diff --git a/ui/src/lib/sessions/index.test.ts b/ui/src/lib/sessions/index.test.ts index a298c1fcb27f..f6f07207f32f 100644 --- a/ui/src/lib/sessions/index.test.ts +++ b/ui/src/lib/sessions/index.test.ts @@ -625,7 +625,7 @@ describe("createSessionCapability", () => { sessions.dispose(); }); - it("rolls back an optimistic model patch when its connection epoch retires", async () => { + it("clears optimistic and settled model overrides when its connection epoch retires", async () => { const stalePatch = deferred(); const request = vi.fn(async (method: string) => { if (method === "sessions.patch") { @@ -643,18 +643,20 @@ describe("createSessionCapability", () => { const { gateway, publish } = createGatewayHarness(client); const sessions = createSessionCapability(gateway); const key = "agent:main:main"; + const inactiveKey = "agent:main:inactive"; sessions.setModelOverride(key, "openai/gpt-old"); + sessions.setModelOverride(inactiveKey, "openai/gpt-old-account"); const operation = sessions.patch(key, { model: "openai/gpt-new" }); expect(sessions.state.modelOverrides[key]).toBe("openai/gpt-new"); publish(false); - expect(sessions.state.modelOverrides[key]).toBe("openai/gpt-old"); + expect(sessions.state.modelOverrides).toEqual({}); publish(true); stalePatch.resolve({}); await expect(operation).resolves.toBeNull(); - expect(sessions.state.modelOverrides[key]).toBe("openai/gpt-old"); + expect(sessions.state.modelOverrides).toEqual({}); sessions.dispose(); }); @@ -692,7 +694,7 @@ describe("createSessionCapability", () => { await expect(operation).resolves.toBeNull(); expect(request).not.toHaveBeenCalledWith("sessions.patch", expect.anything()); - expect(sessions.state.modelOverrides[key]).toBe("openai/gpt-old"); + expect(sessions.state.modelOverrides[key]).toBeUndefined(); sessions.dispose(); }); diff --git a/ui/src/lib/sessions/session-mutations.ts b/ui/src/lib/sessions/session-mutations.ts index 471a97ea9c1d..89d919402fa2 100644 --- a/ui/src/lib/sessions/session-mutations.ts +++ b/ui/src/lib/sessions/session-mutations.ts @@ -39,9 +39,8 @@ type SessionMutationsHost = { export function createSessionMutations(host: SessionMutationsHost) { const pendingModelPatches = new Map< string, - { token: symbol; previous: string | null | undefined } + { token: symbol; previous: string | null | undefined; revision: number } >(); - const modelOverrideRevisions = new Map(); const preparedWorkSessionKeys = new Set(); const setModelOverride = (key: string, value: string | null | undefined) => { @@ -49,8 +48,11 @@ export function createSessionMutations(host: SessionMutationsHost) { if (!normalizedKey) { return; } - // Equal-value writes still transfer ownership between agent-scoped queues. - modelOverrideRevisions.set(normalizedKey, (modelOverrideRevisions.get(normalizedKey) ?? 0) + 1); + // Equal-value writes still transfer ownership while a patch is pending. + const pendingModelPatch = pendingModelPatches.get(normalizedKey); + if (pendingModelPatch) { + pendingModelPatch.revision += 1; + } const state = host.readState(); const modelOverrides = { ...state.modelOverrides }; if (value === undefined) { @@ -90,14 +92,6 @@ export function createSessionMutations(host: SessionMutationsHost) { } }; - const rollback = () => { - const pending = [...pendingModelPatches]; - pendingModelPatches.clear(); - for (const [key, operation] of pending) { - setModelOverride(key, operation.previous); - } - }; - const retireModelOverride = (key: string) => { const normalizedKey = key.trim(); if (!normalizedKey) { @@ -192,19 +186,21 @@ export function createSessionMutations(host: SessionMutationsHost) { pendingModelPatches.set(normalizedKey, { token: modelPatchToken, previous: previousModelOverride, + revision: 0, }); setModelOverride(key, patchParams.model); - modelPatchRevision = modelOverrideRevisions.get(normalizedKey) ?? 0; + modelPatchRevision = pendingModelPatches.get(normalizedKey)?.revision ?? 0; }; if (!options.waitFor) { startModelPatch(); } const settleModelOverride = (completed: boolean) => { - if (modelPatchStarted && pendingModelPatches.get(normalizedKey)?.token === modelPatchToken) { + const pendingModelPatch = pendingModelPatches.get(normalizedKey); + if (modelPatchStarted && pendingModelPatch?.token === modelPatchToken) { pendingModelPatches.delete(normalizedKey); if (host.connection.isCurrent(scope) && ownsModelOverride()) { setModelOverride(key, completed ? patchParams.model : previousModelOverride); - } else if ((modelOverrideRevisions.get(normalizedKey) ?? 0) === modelPatchRevision) { + } else if (pendingModelPatch.revision === modelPatchRevision) { // The shared key now belongs to another agent/connection. Remove only // this operation's untouched optimistic value; preserve newer claims. setModelOverride(key, undefined); @@ -364,12 +360,15 @@ export function createSessionMutations(host: SessionMutationsHost) { } }, retireConnection() { - rollback(); + pendingModelPatches.clear(); preparedWorkSessionKeys.clear(); + const state = host.readState(); + if (Object.keys(state.modelOverrides).length > 0) { + host.publish({ ...state, modelOverrides: {} }); + } }, dispose() { pendingModelPatches.clear(); - modelOverrideRevisions.clear(); preparedWorkSessionKeys.clear(); }, }; diff --git a/ui/src/pages/chat/chat-commands.ts b/ui/src/pages/chat/chat-commands.ts index f2f3c94b2467..22b7475e778a 100644 --- a/ui/src/pages/chat/chat-commands.ts +++ b/ui/src/pages/chat/chat-commands.ts @@ -20,8 +20,9 @@ import { } from "../../lib/sessions/index.ts"; import { areUiSessionKeysEquivalent, - isUiGlobalSessionKey, + isUiSelectedGlobalSessionKey, resolveUiDefaultAgentId, + resolveUiSelectedGlobalAgentId, type UiSessionDefaultsHost, } from "../../lib/sessions/session-key.ts"; import { executeSlashCommand } from "./chat-command-executor.ts"; @@ -148,8 +149,8 @@ function isChatCommandModelCacheOwnerCurrent( // The selected-agent global session shares one UI cache key across agents. // Keep delayed results out when that selection changes on the same Gateway. return ( - !isUiGlobalSessionKey(target.sessionKey) || - scopedAgentIdForSession(host, target.sessionKey) === target.agentId + !isUiSelectedGlobalSessionKey(target.sessionKey) || + resolveUiSelectedGlobalAgentId(host) === target.agentId ); } diff --git a/ui/src/pages/chat/chat-send.test.ts b/ui/src/pages/chat/chat-send.test.ts index 56c775fc4a54..0fa526d4110f 100644 --- a/ui/src/pages/chat/chat-send.test.ts +++ b/ui/src/pages/chat/chat-send.test.ts @@ -4560,41 +4560,44 @@ describe("handleSendChat", () => { expect(host.sessions.state.modelOverrides[item.sessionKey]).toBeUndefined(); }); - it("does not apply a late global model result after the selected agent changes", async () => { - const command = createDeferred>>(); - executeSlashCommandMock.mockImplementationOnce(() => command.promise); + it.each(["global", "agent:work:main"])( + "does not apply a late selected-global model result after the selected agent changes for %s", + async (sessionKey) => { + const command = createDeferred>>(); + executeSlashCommandMock.mockImplementationOnce(() => command.promise); - const item = createQueuedLocalCommand("switched-global-model-command", "/model gpt-5-mini", { - sessionKey: "global", - }); - const host = makeHost({ - requestHandlers: { - "chat.history": () => idleChatHistory(item.sessionKey), - }, - assistantAgentId: "work", - agentsList: { defaultId: "main" }, - chatQueue: [item], - sessionKey: item.sessionKey, - }); - const setModelOverride = vi.spyOn(host.sessions, "setModelOverride"); - expect(admitQueuedMessageForSession(host, item.sessionKey, item)).toBe(true); + const item = createQueuedLocalCommand("switched-global-model-command", "/model gpt-5-mini", { + sessionKey, + }); + const host = makeHost({ + requestHandlers: { + "chat.history": () => idleChatHistory(item.sessionKey), + }, + assistantAgentId: "work", + agentsList: { defaultId: "main" }, + chatQueue: [item], + sessionKey: item.sessionKey, + }); + const setModelOverride = vi.spyOn(host.sessions, "setModelOverride"); + expect(admitQueuedMessageForSession(host, item.sessionKey, item)).toBe(true); - const draining = retryReconnectableQueuedChatSends(host); - await waitForFast(() => expect(executeSlashCommandMock).toHaveBeenCalledTimes(1)); + const draining = retryReconnectableQueuedChatSends(host); + await waitForFast(() => expect(executeSlashCommandMock).toHaveBeenCalledTimes(1)); - host.assistantAgentId = "main"; - command.resolve({ - action: "refresh", - content: "Model set to `gpt-5-mini`.", - sessionPatch: { - modelOverride: { kind: "qualified", value: "openai/gpt-5-mini" }, - }, - }); - await draining; + host.assistantAgentId = "main"; + command.resolve({ + action: "refresh", + content: "Model set to `gpt-5-mini`.", + sessionPatch: { + modelOverride: { kind: "qualified", value: "openai/gpt-5-mini" }, + }, + }); + await draining; - expect(setModelOverride).not.toHaveBeenCalled(); - expect(host.sessions.state.modelOverrides[item.sessionKey]).toBeUndefined(); - }); + expect(setModelOverride).not.toHaveBeenCalled(); + expect(host.sessions.state.modelOverrides[item.sessionKey]).toBeUndefined(); + }, + ); it("does not borrow a replacement connection error for a stale queued command", async () => { const command = createDeferred>>(); diff --git a/ui/src/pages/chat/chat-session.ts b/ui/src/pages/chat/chat-session.ts index a1d6b3e7d284..4e5914122363 100644 --- a/ui/src/pages/chat/chat-session.ts +++ b/ui/src/pages/chat/chat-session.ts @@ -18,6 +18,7 @@ import { isUiGlobalSessionKey, isUiSelectedGlobalSessionKey, resolveUiGlobalAliasAgentId, + resolveUiSelectedGlobalAgentId, } from "../../lib/sessions/session-key.ts"; import { normalizeOptionalString } from "../../lib/string-coerce.ts"; import type { ChatHistoryResult } from "./chat-history.ts"; @@ -415,8 +416,8 @@ export async function switchChatModel( } const modelOwnerAgentId = scopedAgentParamsForSession(host, targetSessionKey).agentId; const ownsModelOverride = () => - !isUiGlobalSessionKey(targetSessionKey) || - scopedAgentParamsForSession(host, targetSessionKey).agentId === modelOwnerAgentId; + !isUiSelectedGlobalSessionKey(targetSessionKey) || + resolveUiSelectedGlobalAgentId(host) === modelOwnerAgentId; setChatError(host, null, true); const switchPromiseRef: { current?: Promise } = {}; const clearPendingSwitch = () => { diff --git a/ui/src/pages/chat/chat-view.test.ts b/ui/src/pages/chat/chat-view.test.ts index a0d474f5b996..23e7f5b95df2 100644 --- a/ui/src/pages/chat/chat-view.test.ts +++ b/ui/src/pages/chat/chat-view.test.ts @@ -6031,64 +6031,67 @@ describe("chat model controls", () => { expect(host.chatThinkingLevel).toBe("high"); }); - it("does not report a failed global model switch after the selected agent changes", async () => { - const modelPatch = createDeferred(); - const modelOverrides: Record = { - global: "openai/gpt-agent-a-old", - }; - let patchOptions: SessionPatchOptions | undefined; - const sessions = { - state: { modelOverrides }, - patch: vi.fn( - async (_key: string, _patch: Record, options?: SessionPatchOptions) => { - patchOptions = options; - return await modelPatch.promise; - }, - ), - refresh: async () => {}, - setModelOverride: vi.fn((key: string, value: string | null | undefined) => { - if (value === undefined) { - delete modelOverrides[key]; - } else { - modelOverrides[key] = value; - } - }), - patchRowLocal: vi.fn(), - }; - const host = { - assistantAgentId: "work", - agentsList: { defaultId: "main", scope: "global" }, - client: {}, - connected: true, - sessionKey: "global", - chatModelCatalog: [], - chatModelSwitchPromises: {}, - chatThinkingLevel: null, - sessions, - sessionsResult: createSessionsResultFromRows([ - { - key: "global", - kind: "direct", - updatedAt: 1, - model: "gpt-agent-a-old", - modelProvider: "openai", - }, - ]), - } as unknown as Parameters[0]; + it.each(["global", "agent:work:main"])( + "does not report a failed selected-global model switch after the selected agent changes for %s", + async (sessionKey) => { + const modelPatch = createDeferred(); + const modelOverrides: Record = { + [sessionKey]: "openai/gpt-agent-a-old", + }; + let patchOptions: SessionPatchOptions | undefined; + const sessions = { + state: { modelOverrides }, + patch: vi.fn( + async (_key: string, _patch: Record, options?: SessionPatchOptions) => { + patchOptions = options; + return await modelPatch.promise; + }, + ), + refresh: async () => {}, + setModelOverride: vi.fn((key: string, value: string | null | undefined) => { + if (value === undefined) { + delete modelOverrides[key]; + } else { + modelOverrides[key] = value; + } + }), + patchRowLocal: vi.fn(), + }; + const host = { + assistantAgentId: "work", + agentsList: { defaultId: "main", scope: "global" }, + client: {}, + connected: true, + sessionKey, + chatModelCatalog: [], + chatModelSwitchPromises: {}, + chatThinkingLevel: null, + sessions, + sessionsResult: createSessionsResultFromRows([ + { + key: sessionKey, + kind: "direct", + updatedAt: 1, + model: "gpt-agent-a-old", + modelProvider: "openai", + }, + ]), + } as unknown as Parameters[0]; - const switching = switchChatModel(host, "openai/gpt-agent-a-new"); - await waitForFast(() => expect(patchOptions).toBeDefined()); - expect(patchOptions?.ownsModelOverride?.()).toBe(true); + const switching = switchChatModel(host, "openai/gpt-agent-a-new"); + await waitForFast(() => expect(patchOptions).toBeDefined()); + expect(patchOptions?.ownsModelOverride?.()).toBe(true); - host.assistantAgentId = "main"; - modelPatch.reject(new Error("agent A patch failed")); + host.assistantAgentId = "main"; + modelPatch.reject(new Error("agent A patch failed")); - await expect(switching).resolves.toBe(false); - expect(patchOptions?.ownsModelOverride?.()).toBe(false); - expect(modelOverrides.global).toBe("openai/gpt-agent-a-old"); - expect(host.lastError ?? null).toBeNull(); - expect(host.chatError ?? null).toBeNull(); - }); + await expect(switching).resolves.toBe(false); + expect(patchOptions?.ownsModelOverride?.()).toBe(false); + expect(modelOverrides[sessionKey]).toBe("openai/gpt-agent-a-old"); + expect(host.lastError ?? null).toBeNull(); + expect(host.chatError ?? null).toBeNull(); + }, + ); it("keeps the newest speed selection when an older patch fails late", async () => { const pendingPatches: Array<{ resolve: () => void; reject: (error: Error) => void }> = [];