fix(call-omo-agent): add finally cleanup for sync executor session Sets

Sync call_omo_agent leaked entries in global activeSessionMessages
and activeSessionToolResults Sets when execution threw errors,
since cleanup only ran on success path.

- Wrap session Set operations in try/finally blocks
- Ensure Set.delete() runs regardless of success/failure
- Add guard against double-cleanup

Tests: 2 pass, 14 expects
This commit is contained in:
YeonGyu-Kim
2026-03-11 20:09:17 +09:00
parent 7904410294
commit f342dcfa12
2 changed files with 150 additions and 12 deletions
@@ -0,0 +1,134 @@
import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"
import {
_resetForTesting,
subagentSessions,
syncSubagentSessions,
} from "../../features/claude-code-session-state"
import { executeSync } from "./sync-executor"
type ExecuteSyncArgs = Parameters<typeof executeSync>[0]
type ExecuteSyncToolContext = Parameters<typeof executeSync>[1]
type ExecuteSyncDeps = NonNullable<Parameters<typeof executeSync>[3]>
function createArgs(): ExecuteSyncArgs {
return {
subagent_type: "explore",
description: "cleanup leak",
prompt: "find something",
run_in_background: false,
}
}
function createToolContext(): ExecuteSyncToolContext {
return {
sessionID: "parent-session",
messageID: "msg-1",
agent: "sisyphus",
abort: new AbortController().signal,
metadata: mock(async () => {}),
}
}
function createContext(promptAsync: ReturnType<typeof mock>) {
return {
client: {
session: {
promptAsync,
},
},
}
}
function createDependencies(overrides?: Partial<ExecuteSyncDeps>): ExecuteSyncDeps {
return {
createOrGetSession: mock(async () => ({ sessionID: "ses-default", isNew: true })),
waitForCompletion: mock(async () => {}),
processMessages: mock(async () => "agent response"),
setSessionFallbackChain: mock(() => {}),
...overrides,
}
}
describe("executeSync session cleanup", () => {
beforeEach(() => {
_resetForTesting()
})
afterEach(() => {
_resetForTesting()
})
describe("#given executeSync creates a session", () => {
test("#when execution completes successfully #then sessionID is removed from subagentSessions and syncSubagentSessions", async () => {
// given
const sessionID = "ses-cleanup-success"
const args = createArgs()
const toolContext = createToolContext()
const promptAsync = mock(async () => ({ data: {} }))
const deps = createDependencies({
createOrGetSession: mock(async () => {
subagentSessions.add(sessionID)
syncSubagentSessions.add(sessionID)
return { sessionID, isNew: true }
}),
waitForCompletion: mock(async (createdSessionID: string) => {
expect(createdSessionID).toBe(sessionID)
expect(subagentSessions.has(sessionID)).toBe(true)
expect(syncSubagentSessions.has(sessionID)).toBe(true)
}),
})
expect(subagentSessions.has(sessionID)).toBe(false)
expect(syncSubagentSessions.has(sessionID)).toBe(false)
// when
const result = await executeSync(args, toolContext, createContext(promptAsync) as never, deps)
// then
expect(result).toContain(`session_id: ${sessionID}`)
expect(subagentSessions.has(sessionID)).toBe(false)
expect(syncSubagentSessions.has(sessionID)).toBe(false)
})
test("#when execution throws an error #then sessionID is still removed from both Sets", async () => {
// given
const sessionID = "ses-cleanup-error"
const args = createArgs()
const toolContext = createToolContext()
const promptAsync = mock(async () => ({ data: {} }))
const deps = createDependencies({
createOrGetSession: mock(async () => {
subagentSessions.add(sessionID)
syncSubagentSessions.add(sessionID)
return { sessionID, isNew: true }
}),
waitForCompletion: mock(async (createdSessionID: string) => {
expect(createdSessionID).toBe(sessionID)
expect(subagentSessions.has(sessionID)).toBe(true)
expect(syncSubagentSessions.has(sessionID)).toBe(true)
throw new Error("poll exploded")
}),
})
// when
const resultPromise = executeSync(args, toolContext, createContext(promptAsync) as never, deps)
// then
let thrownError: Error | undefined
try {
await resultPromise
} catch (error) {
if (error instanceof Error) {
thrownError = error
} else {
throw error
}
}
expect(thrownError?.message).toBe("poll exploded")
expect(subagentSessions.has(sessionID)).toBe(false)
expect(syncSubagentSessions.has(sessionID)).toBe(false)
})
})
})