fix(prompt-gate): block prompts into pending tool turns

This commit is contained in:
YeonGyu-Kim
2026-05-17 15:42:58 +09:00
parent 6eb88a0545
commit a7b7ace7ed
9 changed files with 454 additions and 23 deletions
+25 -1
View File
@@ -1,5 +1,6 @@
import { afterEach, describe, expect, test } from "bun:test"
import { createSessionRecoveryHook } from "./hook"
import { _setInterruptedIdleMessagesFetchTimeoutMsForTesting } from "./interrupted-idle-message-fetch-timeout"
import { releaseAllPromptAsyncReservationsForTesting } from "../../shared/prompt-async-gate"
type RecoverableInfo = Parameters<ReturnType<typeof createSessionRecoveryHook>["handleSessionRecovery"]>[0]
@@ -19,6 +20,7 @@ type PromptAsyncCall = {
afterEach(() => {
releaseAllPromptAsyncReservationsForTesting()
_setInterruptedIdleMessagesFetchTimeoutMsForTesting(undefined)
})
function createPrefillErrorInfo(): RecoverableInfo {
@@ -127,7 +129,8 @@ describe("session-recovery hook interrupted idle recovery", () => {
id: "msg_assistant_unfinished",
role: "assistant",
sessionID: "ses_idle_interrupted",
time: { created: 1778995446058 },
finish: "tool-calls",
time: { created: 1778995446058, completed: 1778995447058 },
},
parts: [
{
@@ -186,4 +189,25 @@ describe("session-recovery hook interrupted idle recovery", () => {
expect(promptAsyncCalls[0]?.body.model).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-7" })
expect(promptAsyncCalls[0]?.body.variant).toBe("max")
})
test("#given session.messages hangs during idle recovery #when timeout elapses #then idle recovery returns false", async () => {
// given
_setInterruptedIdleMessagesFetchTimeoutMsForTesting(5)
const ctx = {
client: {
session: {
messages: async () => new Promise(() => {}),
promptAsync: async () => ({}),
},
},
directory: "/tmp/session-recovery-timeout-test",
}
const hook = createSessionRecoveryHook(ctx as never)
// when
const result = await hook.handleInterruptedToolResultsOnIdle("ses_messages_hangs")
// then
expect(result).toBe(false)
})
})
+14 -4
View File
@@ -5,6 +5,10 @@ import { detectErrorType } from "./detect-error-type"
import type { RecoveryErrorType } from "./detect-error-type"
import type { MessageData } from "./types"
import { normalizeSDKResponse } from "../../shared"
import {
getInterruptedIdleMessagesFetchTimeoutMs,
withInterruptedIdleMessagesFetchTimeout,
} from "./interrupted-idle-message-fetch-timeout"
import { recoverToolResultMissing } from "./recover-tool-result-missing"
import { recoverUnavailableTool } from "./recover-unavailable-tool"
import { recoverThinkingBlockOrder } from "./recover-thinking-block-order"
@@ -56,6 +60,9 @@ export function createSessionRecoveryHook(ctx: PluginInput, options?: SessionRec
}
const finish = message.info?.finish
if (finish === "tool-calls") {
return false
}
if ((typeof finish === "string" && finish.length > 0) || finish === true) {
return true
}
@@ -99,10 +106,13 @@ export function createSessionRecoveryHook(ctx: PluginInput, options?: SessionRec
let recoveryStarted = false
let assistantMessageIDForRecovery: string | undefined
try {
const messagesResp = await ctx.client.session.messages({
path: { id: sessionID },
query: { directory: ctx.directory },
})
const messagesResp = await withInterruptedIdleMessagesFetchTimeout(
ctx.client.session.messages({
path: { id: sessionID },
query: { directory: ctx.directory },
}),
getInterruptedIdleMessagesFetchTimeoutMs(),
)
const messages = normalizeSDKResponse(messagesResp, [] as MessageData[])
const latestAssistant = findLatestAssistantMessage(messages)
if (!latestAssistant?.info?.id) {
@@ -0,0 +1,38 @@
export const DEFAULT_INTERRUPTED_IDLE_MESSAGES_FETCH_TIMEOUT_MS = 5_000
let interruptedIdleMessagesFetchTimeoutMsForTesting: number | undefined
export function _setInterruptedIdleMessagesFetchTimeoutMsForTesting(value: number | undefined): void {
interruptedIdleMessagesFetchTimeoutMsForTesting = value
}
export function getInterruptedIdleMessagesFetchTimeoutMs(): number {
return interruptedIdleMessagesFetchTimeoutMsForTesting ?? DEFAULT_INTERRUPTED_IDLE_MESSAGES_FETCH_TIMEOUT_MS
}
export class InterruptedIdleMessagesFetchTimeoutError extends Error {
constructor(timeoutMs: number) {
super(`[session-recovery] session.messages timed out after ${timeoutMs}ms while checking interrupted idle tools`)
this.name = "InterruptedIdleMessagesFetchTimeoutError"
}
}
export function withInterruptedIdleMessagesFetchTimeout<T>(operation: Promise<T>, timeoutMs: number): Promise<T> {
if (timeoutMs <= 0) {
return operation
}
let timeoutID: ReturnType<typeof globalThis.setTimeout> | undefined
const timeoutPromise = new Promise<never>((_, reject) => {
timeoutID = globalThis.setTimeout(
() => reject(new InterruptedIdleMessagesFetchTimeoutError(timeoutMs)),
timeoutMs,
)
})
return Promise.race([operation, timeoutPromise]).finally(() => {
if (timeoutID !== undefined) {
globalThis.clearTimeout(timeoutID)
}
})
}
@@ -174,6 +174,7 @@ export async function recoverToolResultMissing(
sessionID,
source: options?.source ?? "session-recovery-tool-result-missing",
input: promptInput,
checkToolState: false,
})
return promptResult.status === "dispatched"