diff --git a/src/tools/call-omo-agent/sync-executor-leak.test.ts b/src/tools/call-omo-agent/sync-executor-leak.test.ts index 778af537e..eb236e973 100644 --- a/src/tools/call-omo-agent/sync-executor-leak.test.ts +++ b/src/tools/call-omo-agent/sync-executor-leak.test.ts @@ -45,6 +45,7 @@ function createDependencies(overrides?: Partial): ExecuteSyncDe waitForCompletion: mock(async () => {}), processMessages: mock(async () => "agent response"), setSessionFallbackChain: mock(() => {}), + clearSessionFallbackChain: mock(() => {}), ...overrides, } } @@ -133,7 +134,7 @@ describe("executeSync session cleanup", () => { }) describe("#given executeSync reuses an existing session", () => { - test("#when execution completes successfully #then the reused session stays tracked in both Sets", async () => { + test("#when execution completes successfully #then the reused session is tracked in both Sets", async () => { // given const sessionID = "ses-reused" const args = { ...createArgs(), session_id: sessionID } @@ -141,10 +142,15 @@ describe("executeSync session cleanup", () => { const promptAsync = mock(async () => ({ data: {} })) const deps = createDependencies({ createOrGetSession: mock(async () => ({ sessionID, isNew: false })), + waitForCompletion: mock(async (createdSessionID: string) => { + expect(createdSessionID).toBe(sessionID) + expect(subagentSessions.has(sessionID)).toBe(true) + expect(syncSubagentSessions.has(sessionID)).toBe(true) + }), }) - subagentSessions.add(sessionID) - syncSubagentSessions.add(sessionID) + expect(subagentSessions.has(sessionID)).toBe(false) + expect(syncSubagentSessions.has(sessionID)).toBe(false) // when const result = await executeSync(args, toolContext, createContext(promptAsync) as never, deps) @@ -154,5 +160,25 @@ describe("executeSync session cleanup", () => { expect(subagentSessions.has(sessionID)).toBe(true) expect(syncSubagentSessions.has(sessionID)).toBe(true) }) + + test("#when execution applies a fallback chain #then it clears that chain in finally", async () => { + // given + const sessionID = "ses-reused-fallback" + const args = { ...createArgs(), session_id: sessionID } + const toolContext = createToolContext() + const promptAsync = mock(async () => ({ data: {} })) + const clearSessionFallbackChain = mock(() => {}) + const deps = createDependencies({ + createOrGetSession: mock(async () => ({ sessionID, isNew: false })), + clearSessionFallbackChain, + }) + const fallbackChain = [{ providers: ["openai"], model: "gpt-5.4" }] + + // when + await executeSync(args, toolContext, createContext(promptAsync) as never, deps, fallbackChain) + + // then + expect(clearSessionFallbackChain).toHaveBeenCalledWith(sessionID) + }) }) }) diff --git a/src/tools/call-omo-agent/sync-executor.test.ts b/src/tools/call-omo-agent/sync-executor.test.ts index 513588676..d1b1dd3f6 100644 --- a/src/tools/call-omo-agent/sync-executor.test.ts +++ b/src/tools/call-omo-agent/sync-executor.test.ts @@ -24,6 +24,7 @@ type Dependencies = { waitForCompletion: ReturnType processMessages: ReturnType setSessionFallbackChain: ReturnType + clearSessionFallbackChain: ReturnType } async function importExecuteSync(): Promise { @@ -37,6 +38,7 @@ function createDependencies(overrides?: Partial): Dependencies { waitForCompletion: mock(async () => {}), processMessages: mock(async () => "agent response"), setSessionFallbackChain: mock(() => {}), + clearSessionFallbackChain: mock(() => {}), ...overrides, } } @@ -259,6 +261,7 @@ describe("executeSync", () => { waitForCompletion: mock(async () => {}), processMessages: mock(async () => "agent response"), setSessionFallbackChain: mock(() => {}), + clearSessionFallbackChain: mock(() => {}), } const spawnReservation = { diff --git a/src/tools/call-omo-agent/sync-executor.ts b/src/tools/call-omo-agent/sync-executor.ts index ce3d5089b..015b81b5d 100644 --- a/src/tools/call-omo-agent/sync-executor.ts +++ b/src/tools/call-omo-agent/sync-executor.ts @@ -1,7 +1,7 @@ import type { CallOmoAgentArgs } from "./types" import type { PluginInput } from "@opencode-ai/plugin" import { subagentSessions, syncSubagentSessions } from "../../features/claude-code-session-state" -import { setSessionFallbackChain } from "../../hooks/model-fallback/hook" +import { clearSessionFallbackChain, setSessionFallbackChain } from "../../hooks/model-fallback/hook" import { getAgentToolRestrictions, log } from "../../shared" import type { FallbackEntry } from "../../shared/model-requirements" import { waitForCompletion } from "./completion-poller" @@ -17,6 +17,7 @@ type ExecuteSyncDeps = { waitForCompletion: typeof waitForCompletion processMessages: typeof processMessages setSessionFallbackChain: typeof setSessionFallbackChain + clearSessionFallbackChain: typeof clearSessionFallbackChain } type SpawnReservation = { @@ -29,6 +30,7 @@ const defaultDeps: ExecuteSyncDeps = { waitForCompletion, processMessages, setSessionFallbackChain, + clearSessionFallbackChain, } export async function executeSync( @@ -47,11 +49,14 @@ export async function executeSync( ): Promise { let sessionID: string | undefined let createdSessionForExecution = false + let appliedFallbackChain = false try { const session = await deps.createOrGetSession(args, toolContext, ctx) sessionID = session.sessionID createdSessionForExecution = session.isNew + subagentSessions.add(sessionID) + syncSubagentSessions.add(sessionID) if (session.isNew) { spawnReservation?.commit() @@ -59,6 +64,7 @@ export async function executeSync( if (fallbackChain && fallbackChain.length > 0) { deps.setSessionFallbackChain(sessionID, fallbackChain) + appliedFallbackChain = true } await Promise.resolve( @@ -102,6 +108,10 @@ export async function executeSync( spawnReservation?.rollback() throw error } finally { + if (sessionID && appliedFallbackChain) { + deps.clearSessionFallbackChain(sessionID) + } + if (sessionID && createdSessionForExecution) { subagentSessions.delete(sessionID) syncSubagentSessions.delete(sessionID)