fix(mothership): parallel tool calls

This commit is contained in:
Siddharth Ganesan
2026-03-24 02:45:31 -07:00
parent 092525e8aa
commit 41a7d247ea
2 changed files with 80 additions and 0 deletions
@@ -278,4 +278,41 @@ describe('orchestrateCopilotStream async continuation', () => {
expect(releaseCompletedAsyncToolClaim).toHaveBeenCalledWith('tool-1', 'run-1')
expect(markAsyncToolDelivered).not.toHaveBeenCalled()
})
it('does not send a partial resume payload when only some pending tool calls are claimable', async () => {
claimCompletedAsyncToolCall
.mockResolvedValueOnce({ toolCallId: 'tool-1' })
.mockResolvedValueOnce(null)
.mockResolvedValueOnce({ toolCallId: 'tool-1' })
.mockResolvedValueOnce(null)
.mockResolvedValueOnce({ toolCallId: 'tool-1' })
.mockResolvedValueOnce(null)
.mockResolvedValueOnce({ toolCallId: 'tool-1' })
.mockResolvedValueOnce(null)
getAsyncToolCall.mockResolvedValue(null)
runStreamLoop.mockImplementationOnce(async (_url: string, _opts: RequestInit, context: any) => {
context.awaitingAsyncContinuation = {
checkpointId: 'checkpoint-1',
runId: 'run-1',
pendingToolCallIds: ['tool-1', 'tool-2'],
}
})
const result = await orchestrateCopilotStream(
{ message: 'hello' },
{
userId: 'user-1',
workflowId: 'workflow-1',
chatId: 'chat-1',
executionId: 'exec-1',
runId: 'run-1',
}
)
expect(result.success).toBe(true)
expect(runStreamLoop).toHaveBeenCalledTimes(1)
expect(releaseCompletedAsyncToolClaim).toHaveBeenCalledWith('tool-1', 'run-1')
expect(markAsyncToolDelivered).not.toHaveBeenCalled()
})
})
@@ -246,6 +246,49 @@ export async function orchestrateCopilotStream(
continue
}
const missingToolCallIds = continuation.pendingToolCallIds.filter(
(toolCallId) => !claimableToolCallIds.includes(toolCallId)
)
if (missingToolCallIds.length > 0) {
if (claimedToolCallIds.length > 0 && claimedByWorkerId) {
logger.info('Releasing partial async tool claims before retrying resume', {
checkpointId: continuation.checkpointId,
runId: continuation.runId,
claimedToolCallIds,
missingToolCallIds,
})
await Promise.all(
claimedToolCallIds.map((toolCallId) =>
releaseCompletedAsyncToolClaim(toolCallId, claimedByWorkerId!).catch(() => null)
)
)
claimedToolCallIds = []
claimedByWorkerId = null
}
if (emptyClaimRetries < 3) {
emptyClaimRetries++
logger.info(
'Retrying async resume claim after only a subset of tool calls were claimable',
{
checkpointId: continuation.checkpointId,
runId: continuation.runId,
retry: emptyClaimRetries,
missingToolCallIds,
}
)
await new Promise((resolve) => setTimeout(resolve, 250 * emptyClaimRetries))
continue
}
logger.warn('Skipping async resume because not all tool calls were claimable', {
checkpointId: continuation.checkpointId,
runId: continuation.runId,
claimableToolCallIds,
missingToolCallIds,
})
context.awaitingAsyncContinuation = undefined
break
}
if (claimableToolCallIds.length === 0) {
if (emptyClaimRetries < 3 && continuation.pendingToolCallIds.length > 0) {
emptyClaimRetries++