From eb0931ed6d2aeb879b77f0514cf8b4bdbff7d19d Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Thu, 19 Feb 2026 16:26:06 +0900 Subject: [PATCH] fix(ultrawork): use session agent fallback and skip same-model override --- src/plugin/chat-message.ts | 2 +- src/plugin/ultrawork-model-override.test.ts | 41 ++++++++++++++++++++- src/plugin/ultrawork-model-override.ts | 30 +++++++++++++-- 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/src/plugin/chat-message.ts b/src/plugin/chat-message.ts index 612b7122b..69340fa1f 100644 --- a/src/plugin/chat-message.ts +++ b/src/plugin/chat-message.ts @@ -140,6 +140,6 @@ export function createChatMessageHandler(args: { } } - applyUltraworkModelOverrideOnMessage(pluginConfig, input.agent, output, ctx.client.tui) + applyUltraworkModelOverrideOnMessage(pluginConfig, input.agent, output, ctx.client.tui, input.sessionID) } } diff --git a/src/plugin/ultrawork-model-override.test.ts b/src/plugin/ultrawork-model-override.test.ts index 1e14c5fa5..4f167e963 100644 --- a/src/plugin/ultrawork-model-override.test.ts +++ b/src/plugin/ultrawork-model-override.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, test, beforeEach, afterEach, spyOn, mock } from "bun:test" +import { describe, expect, test, beforeEach, afterEach, spyOn } from "bun:test" import { applyUltraworkModelOverrideOnMessage, resolveUltraworkOverride, @@ -6,6 +6,7 @@ import { } from "./ultrawork-model-override" import * as sharedModule from "../shared" import * as dbOverrideModule from "./ultrawork-db-model-override" +import * as sessionStateModule from "../features/claude-code-session-state" describe("detectUltrawork", () => { test("should detect ultrawork keyword", () => { @@ -197,6 +198,22 @@ describe("resolveUltraworkOverride", () => { //#then expect(result).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-6", variant: undefined }) }) + + test("should use session agent when input and message agents are undefined", () => { + //#given + const config = createConfig("sisyphus", { model: "anthropic/claude-opus-4-6", variant: "max" }) + const output = createOutput("ultrawork do something") + const getSessionAgentSpy = spyOn(sessionStateModule, "getSessionAgent").mockReturnValue("sisyphus") + + //#when + const result = resolveUltraworkOverride(config, undefined, output, "ses_test") + + //#then + expect(getSessionAgentSpy).toHaveBeenCalledWith("ses_test") + expect(result).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-6", variant: "max" }) + + getSessionAgentSpy.mockRestore() + }) }) describe("applyUltraworkModelOverrideOnMessage", () => { @@ -378,4 +395,26 @@ describe("applyUltraworkModelOverrideOnMessage", () => { "max", ) }) + + test("should skip override trigger when current model already matches ultrawork model", () => { + //#given + const config = createConfig("sisyphus", { model: "anthropic/claude-opus-4-6", variant: "max" }) + const output = createOutput("ultrawork do something", { + existingModel: { providerID: "anthropic", modelID: "claude-opus-4-6" }, + messageId: "msg_123", + }) + let toastCalled = false + const tui = { + showToast: async () => { + toastCalled = true + }, + } + + //#when + applyUltraworkModelOverrideOnMessage(config, "sisyphus", output, tui) + + //#then + expect(dbOverrideSpy).not.toHaveBeenCalled() + expect(toastCalled).toBe(false) + }) }) diff --git a/src/plugin/ultrawork-model-override.ts b/src/plugin/ultrawork-model-override.ts index 83c3d632a..f6aa87bd2 100644 --- a/src/plugin/ultrawork-model-override.ts +++ b/src/plugin/ultrawork-model-override.ts @@ -1,5 +1,6 @@ import type { OhMyOpenCodeConfig } from "../config" import type { AgentOverrides } from "../config/schema/agent-overrides" +import { getSessionAgent } from "../features/claude-code-session-state" import { log } from "../shared" import { getAgentConfigKey } from "../shared/agent-display-names" import { scheduleDeferredModelOverride } from "./ultrawork-db-model-override" @@ -35,6 +36,18 @@ export type UltraworkOverrideResult = { variant?: string } +function isSameModel( + current: unknown, + target: { providerID: string; modelID: string }, +): boolean { + if (typeof current !== "object" || current === null) return false + const currentRecord = current as Record + return ( + currentRecord["providerID"] === target.providerID + && currentRecord["modelID"] === target.modelID + ) +} + /** * Resolves the ultrawork model override config for the given agent and prompt text. * Returns null if no override should be applied. @@ -46,13 +59,15 @@ export function resolveUltraworkOverride( message: Record parts: Array<{ type: string; text?: string; [key: string]: unknown }> }, + sessionID?: string, ): UltraworkOverrideResult | null { const promptText = extractPromptText(output.parts) if (!detectUltrawork(promptText)) return null const messageAgentName = typeof output.message["agent"] === "string" ? (output.message["agent"] as string) : undefined - const rawAgentName = inputAgentName ?? messageAgentName + const sessionAgentName = sessionID ? getSessionAgent(sessionID) : undefined + const rawAgentName = inputAgentName ?? messageAgentName ?? sessionAgentName if (!rawAgentName || !pluginConfig.agents) return null const agentConfigKey = getAgentConfigKey(rawAgentName) @@ -94,8 +109,9 @@ export function applyUltraworkModelOverrideOnMessage( parts: Array<{ type: string; text?: string; [key: string]: unknown }> }, tui: unknown, + sessionID?: string, ): void { - const override = resolveUltraworkOverride(pluginConfig, inputAgentName, output) + const override = resolveUltraworkOverride(pluginConfig, inputAgentName, output, sessionID) if (!override) return if (!override.providerID || !override.modelID) { @@ -106,10 +122,16 @@ export function applyUltraworkModelOverrideOnMessage( return } + const targetModel = { providerID: override.providerID, modelID: override.modelID } + if (isSameModel(output.message.model, targetModel)) { + log(`[ultrawork-model-override] Skip override; target model already active: ${override.modelID}`) + return + } + const messageId = output.message["id"] as string | undefined if (!messageId) { log("[ultrawork-model-override] No message ID found, falling back to direct mutation") - output.message.model = { providerID: override.providerID, modelID: override.modelID } + output.message.model = targetModel if (override.variant) { output.message["variant"] = override.variant output.message["thinking"] = override.variant @@ -125,7 +147,7 @@ export function applyUltraworkModelOverrideOnMessage( scheduleDeferredModelOverride( messageId, - { providerID: override.providerID, modelID: override.modelID }, + targetModel, override.variant, )