Skip to content

fix: Fix tool-bridge slot leak when session.abort() throws inside runAgentTurn (#73) #77

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
KochC wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: Fix tool-bridge slot leak when session.abort() throws inside runAgentTurn (#73) #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 81 additions & 78 deletions index.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -752,59 +752,62 @@ async function runAgentTurn(client, model, messages, system, callerTools, onChun
toolsMap = buildToolsMap(baseTools, bridge)
}

const session = await client.session.create({ body: { title: `Proxy: ${model.id}` } })
const sessionID = session.data.id
const prompt = buildPrompt(messages)
const toolIDSet = bridge ? new Set(bridge.toolIDs) : null
// Wrap all work that happens after bridge acquisition in a try/finally so the
// slot is always returned to the pool even if session.create, event.subscribe,
// or session.promptAsync throw before the event-loop runs.
try {
const session = await client.session.create({ body: { title: `Proxy: ${model.id}` } })
const sessionID = session.data.id
const prompt = buildPrompt(messages)
const toolIDSet = bridge ? new Set(bridge.toolIDs) : null

// Subscribe to the event stream before sending the prompt so we don't miss events.
const { stream } = await client.event.subscribe()
// Subscribe to the event stream before sending the prompt so we don't miss events.
const { stream } = await client.event.subscribe()

await client.session.promptAsync({
path: { id: sessionID },
body: {
model: { providerID: model.providerID, modelID: model.modelID },
system,
tools: toolsMap,
parts: [{ type: "text", text: prompt }],
},
})
await client.session.promptAsync({
path: { id: sessionID },
body: {
model: { providerID: model.providerID, modelID: model.modelID },
system,
tools: toolsMap,
parts: [{ type: "text", text: prompt }],
},
})

let errorMessage = null
let content = ""
// Tool calls collected live off the event stream, keyed by callID so parallel calls
// and the pending -> running -> completed status updates for each collapse into one
// entry. session.messages() does NOT carry the tool parts in current OpenCode, so the
// live stream is the authoritative source here.
const toolCallsByID = new Map()
// The assistant message that carries this turn's tool calls. Once set, we only accept
// tool parts (and the terminating step-finish) from this same message, so a follow-up
// agent step can never leak spurious calls into the result.
let toolMessageID = null

const recordToolPart = (part) => {
const callID = part.callID
if (!callID) return
const input = part.state?.input
const hasInput =
input && typeof input === "object" && !Array.isArray(input) && Object.keys(input).length > 0
const existing = toolCallsByID.get(callID)
if (!existing) {
toolCallsByID.set(callID, {
id: callID,
name: bridge.nameMap.get(part.tool) ?? part.tool,
arguments: hasInput ? input : {},
hasInput: Boolean(hasInput),
})
} else if (hasInput && !existing.hasInput) {
// Upgrade from the empty-input "pending" snapshot to the populated one that
// arrives with "running"/"completed".
existing.arguments = input
existing.hasInput = true
let errorMessage = null
let content = ""
// Tool calls collected live off the event stream, keyed by callID so parallel calls
// and the pending -> running -> completed status updates for each collapse into one
// entry. session.messages() does NOT carry the tool parts in current OpenCode, so the
// live stream is the authoritative source here.
const toolCallsByID = new Map()
// The assistant message that carries this turn's tool calls. Once set, we only accept
// tool parts (and the terminating step-finish) from this same message, so a follow-up
// agent step can never leak spurious calls into the result.
let toolMessageID = null

const recordToolPart = (part) => {
const callID = part.callID
if (!callID) return
const input = part.state?.input
const hasInput =
input && typeof input === "object" && !Array.isArray(input) && Object.keys(input).length > 0
const existing = toolCallsByID.get(callID)
if (!existing) {
toolCallsByID.set(callID, {
id: callID,
name: bridge.nameMap.get(part.tool) ?? part.tool,
arguments: hasInput ? input : {},
hasInput: Boolean(hasInput),
})
} else if (hasInput && !existing.hasInput) {
// Upgrade from the empty-input "pending" snapshot to the populated one that
// arrives with "running"/"completed".
existing.arguments = input
existing.hasInput = true
}
}
}

try {
for await (const event of stream) {
if (event.type === "message.part.delta") {
// Real incremental token deltas arrive here, as flat properties (sessionID,
Expand Down Expand Up @@ -863,40 +866,40 @@ async function runAgentTurn(client, model, messages, system, callerTools, onChun
}
}
}
} finally {
releaseToolBridge(bridge)
}

const toolCalls = [...toolCallsByID.values()].map((call) => ({
id: call.id,
name: call.name,
arguments: call.arguments ?? {},
}))
const toolCalls = [...toolCallsByID.values()].map((call) => ({
id: call.id,
name: call.name,
arguments: call.arguments ?? {},
}))

if (errorMessage && toolCalls.length === 0) {
throw new Error(errorMessage)
}
if (errorMessage && toolCalls.length === 0) {
throw new Error(errorMessage)
}

// Each list item is { info: Message, parts: Part[] } - matching the shape
// client.session.prompt() (the non-tool-calling path) already returns directly.
const messagesResult = await client.session.messages({ path: { id: sessionID } })
const assistantEntry = (messagesResult.data ?? []).filter((m) => m.info?.role === "assistant").at(-1)
const assistantInfo = assistantEntry?.info

// Fallback for turns where message.part.delta never fired (observed for some
// multi-step turns, e.g. continuing a conversation with prior tool calls/results in
// history): use the authoritative final text from the fetched message's parts
// instead of leaving content empty.
if (!content && toolCalls.length === 0) {
content = extractAssistantText(assistantEntry?.parts ?? [])
}
// Each list item is { info: Message, parts: Part[] } - matching the shape
// client.session.prompt() (the non-tool-calling path) already returns directly.
const messagesResult = await client.session.messages({ path: { id: sessionID } })
const assistantEntry = (messagesResult.data ?? []).filter((m) => m.info?.role === "assistant").at(-1)
const assistantInfo = assistantEntry?.info

// Fallback for turns where message.part.delta never fired (observed for some
// multi-step turns, e.g. continuing a conversation with prior tool calls/results in
// history): use the authoritative final text from the fetched message's parts
// instead of leaving content empty.
if (!content && toolCalls.length === 0) {
content = extractAssistantText(assistantEntry?.parts ?? [])
}

return {
sessionID,
content,
toolCalls,
tokens: assistantInfo?.tokens ?? { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
finish: toolCalls.length > 0 ? "tool_calls" : assistantInfo?.finish,
return {
sessionID,
content,
toolCalls,
tokens: assistantInfo?.tokens ?? { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
finish: toolCalls.length > 0 ? "tool_calls" : assistantInfo?.finish,
}
} finally {
releaseToolBridge(bridge)
}
}

Expand Down
60 changes: 60 additions & 0 deletions index.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2511,6 +2511,66 @@ describe("buildToolsMap / registerToolBridge slot isolation", () => {
assert.ok(bridge.slotName)
after(() => releaseToolBridge(bridge))
})

// Regression test for issue #73: if session.create / event.subscribe / session.promptAsync
// throws *after* registerToolBridge() has already handed out a slot, the slot must still
// be returned to the pool by the outer try/finally in runAgentTurn(). Without the fix, the
// slot was only released inside the inner event-loop try/finally that was never entered,
// so any throw between bridge acquisition and the event loop would permanently lose the slot.
it("releases the bridge slot when session.create throws after bridge acquisition", async () => {
// Build a minimal client whose session.create always fails, simulating any error in the
// window between bridge acquisition and the start of the event-loop try/finally.
let callCount = 0
const client = {
tool: { ids: async () => ({ data: [] }) },
config: {
providers: async () => ({
data: {
providers: [{ id: "p", models: { m: { id: "m" } } }],
},
}),
},
mcp: {
disconnect: async () => { throw new Error("not connected") },
add: async () => ({ data: {} }),
},
session: {
create: async () => {
callCount++
throw new Error("session.create failed")
},
},
event: { subscribe: async () => ({ stream: (async function* () {})() }) },
}

const handler = createProxyFetchHandler(client)
const makeRequest = () => new Request("http://127.0.0.1:4010/v1/chat/completions", {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
model: "p/m",
messages: [{ role: "user", content: "hi" }],
tools: [{ type: "function", function: { name: "my_tool", description: "A tool" } }],
}),
})

const timeout = delay(3000).then(() => {
throw new Error("timed out - bridge slot was leaked after session.create threw")
})

// Fire many more requests than the pool size. If any slot leaks, the pool is exhausted
// and acquireBridgeSlot() hangs forever - caught by the timeout race.
await Promise.race([
(async () => {
for (let i = 0; i < 20; i++) {
await handler(makeRequest())
}
})(),
timeout,
])

assert.ok(callCount >= 20, "session.create should have been called for each request")
})
})

test("POST /v1beta/models/:model:generateContent returns a functionCall part", async () => {
Expand Down