fix(prompt-gate): harden internal prompt dispatch

2026-05-19 16:02:18 +09:00
parent 6c63372ef9
commit 1492bffd20
49 changed files with 1488 additions and 141 deletions
@@ -13,6 +13,7 @@ import {
 export const DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS = 250
 export const DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS = 30_000
 export const DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS = 5_000
+export const DEFAULT_PROMPT_QUEUE_RETRY_MS = 250

 type PromptAsyncInput = {
  path?: { id?: string }
@@ -44,11 +45,16 @@ type PromptClient<TInput> = {
 }

 export type InternalPromptDispatchMode = "async" | "sync"
+export type InternalPromptQueueBehavior = "enqueue" | "defer"

 type InternalPromptDispatchCommonArgs<TInput> = {
  sessionID: string
  input: TInput
  source: string
+  dedupeKey?: string
+  queueBehavior?: InternalPromptQueueBehavior
+  queue?: boolean
+  queueRetryMs?: number
  settleMs?: number
  postDispatchHoldMs?: number
  dispatchTimeoutMs?: number
@@ -63,6 +69,7 @@ export type InternalPromptDispatchArgs<TInput = PromptAsyncInput> = InternalProm

 type PromptAsyncReservation = {
  source: string
+  dedupeKey: string
  reservedAt: number
  token: symbol
  expiresAt?: number
@@ -75,6 +82,7 @@ let promptGateMessagesFetchTimeoutMsForTesting: number | undefined

 export type InternalPromptDispatchResult =
  | { status: "dispatched"; response: unknown }
+  | { status: "queued"; queuedBy: string; position: number }
  | { status: "active" }
  | { status: "reserved"; reservedBy: string }
  | { status: "unavailable" }
@@ -88,6 +96,35 @@ type PromptAsyncReservationReleaseOptions = {
 }

 const promptAsyncReservations = new Map<string, PromptAsyncReservation>()
+const promptQueues = new Map<string, QueuedInternalPrompt[]>()
+const promptQueueDraining = new Set<string>()
+const promptQueueInFlight = new Map<string, QueuedInternalPrompt>()
+const promptQueueTimers = new Map<string, unknown>()
+let promptQueueSequence = 0
+
+type PromptDispatchClient = {
+  session?: {
+    status?: () => Promise<unknown>
+    messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise<unknown>
+  }
+}
+
+type QueuedInternalPrompt = {
+  id: number
+  sessionID: string
+  sessionName: "promptAsync" | "prompt"
+  client: PromptDispatchClient
+  input: unknown
+  source: string
+  dedupeKey: string
+  settleMs: number
+  postDispatchHoldMs: number
+  dispatchTimeoutMs: number
+  queueRetryMs: number
+  checkStatus: boolean
+  checkToolState: boolean
+  dispatch: (input: unknown) => Promise<unknown>
+}

 export function _setPromptGateMessagesFetchTimeoutMsForTesting(value: number | undefined): void {
  promptGateMessagesFetchTimeoutMsForTesting = value
@@ -98,15 +135,20 @@ function getPromptGateMessagesFetchTimeoutMs(): number {
 }

 function pruneExpiredReservations(now = Date.now()): void {
+  const expiredSessionIDs: string[] = []
  for (const [sessionID, reservation] of promptAsyncReservations) {
    if (typeof reservation.expiresAt === "number" && reservation.expiresAt <= now) {
      promptAsyncReservations.delete(sessionID)
+      expiredSessionIDs.push(sessionID)
      log("[prompt-async-gate] expired reservation released", {
        sessionID,
        source: reservation.source,
      })
    }
  }
+  for (const sessionID of expiredSessionIDs) {
+    schedulePromptQueueDrain(sessionID, 0)
+  }
 }

 function getActiveReservation(sessionID: string): PromptAsyncReservation | undefined {
@@ -114,6 +156,102 @@ function getActiveReservation(sessionID: string): PromptAsyncReservation | undef
  return promptAsyncReservations.get(sessionID)
 }

+function getPromptQueue(sessionID: string): QueuedInternalPrompt[] {
+  const existing = promptQueues.get(sessionID)
+  if (existing) {
+    return existing
+  }
+
+  const queue: QueuedInternalPrompt[] = []
+  promptQueues.set(sessionID, queue)
+  return queue
+}
+
+function setPromptQueue(sessionID: string, queue: QueuedInternalPrompt[]): void {
+  if (queue.length === 0) {
+    promptQueues.delete(sessionID)
+    return
+  }
+  promptQueues.set(sessionID, queue)
+}
+
+function stringifyPromptInputForDedupe(input: unknown): string {
+  try {
+    const serialized = JSON.stringify(input, (key: string, value: unknown): unknown => {
+      if (key === "signal") {
+        return "[AbortSignal]"
+      }
+      if (typeof value === "function") {
+        return `[Function:${value.name}]`
+      }
+      return value
+    })
+    return serialized ?? String(input)
+  } catch {
+    return String(input)
+  }
+}
+
+function createDefaultDedupeKey(source: string, input: unknown): string {
+  const fingerprint = stringifyPromptInputForDedupe(input)
+  return `${source}:${fingerprint.length}:${fingerprint.slice(0, 8192)}`
+}
+
+function queuedResult(entry: QueuedInternalPrompt, position: number, queuedBy = entry.source): InternalPromptDispatchResult {
+  return {
+    status: "queued",
+    queuedBy,
+    position,
+  }
+}
+
+function clearPromptQueueTimer(sessionID: string): void {
+  const timer = promptQueueTimers.get(sessionID)
+  if (timer !== undefined) {
+    clearTimeout(timer)
+    promptQueueTimers.delete(sessionID)
+  }
+}
+
+function schedulePromptQueueDrain(sessionID: string, delayMs: number): void {
+  const queue = promptQueues.get(sessionID)
+  if (!queue || queue.length === 0) {
+    clearPromptQueueTimer(sessionID)
+    return
+  }
+
+  clearPromptQueueTimer(sessionID)
+  const timer = setTimeout(() => {
+    promptQueueTimers.delete(sessionID)
+    void drainPromptQueue(sessionID).catch((error: unknown) => {
+      log("[prompt-async-gate] queued prompt drain failed", {
+        sessionID,
+        error: String(error),
+      })
+    })
+  }, Math.max(0, delayMs))
+  promptQueueTimers.set(sessionID, timer)
+}
+
+function removePromptQueueEntry(sessionID: string, entry: QueuedInternalPrompt): void {
+  const queue = promptQueues.get(sessionID)
+  if (!queue) {
+    return
+  }
+  const nextQueue = queue.filter((queued) => queued.id !== entry.id)
+  setPromptQueue(sessionID, nextQueue)
+}
+
+function getQueuedPromptBlocker(sessionID: string): string | undefined {
+  const inFlight = promptQueueInFlight.get(sessionID)
+  if (inFlight) {
+    return inFlight.source
+  }
+
+  const queue = promptQueues.get(sessionID)
+  return queue?.[0]?.source
+}
+
 function reservationSourceMatches(
  reservationSource: string,
  expectedSource: string | readonly string[],
@@ -341,6 +479,7 @@ async function dispatchAfterSessionIdle<TInput>(args: {
  sessionID: string
  input: TInput
  source: string
+  dedupeKey: string
  settleMs: number
  postDispatchHoldMs: number
  dispatchTimeoutMs: number
@@ -354,6 +493,7 @@ async function dispatchAfterSessionIdle<TInput>(args: {
    sessionID,
    input,
    source,
+    dedupeKey,
    settleMs,
    postDispatchHoldMs,
    dispatchTimeoutMs,
@@ -375,6 +515,7 @@ async function dispatchAfterSessionIdle<TInput>(args: {

  const reservation: PromptAsyncReservation = {
    source,
+    dedupeKey,
    reservedAt: Date.now(),
    token: Symbol(source),
  }
@@ -447,6 +588,117 @@ async function dispatchAfterSessionIdle<TInput>(args: {
  }
 }

+async function drainPromptQueue(sessionID: string, awaitedEntry?: QueuedInternalPrompt): Promise<InternalPromptDispatchResult | undefined> {
+  if (promptQueueDraining.has(sessionID)) {
+    return awaitedEntry ? queuedResult(awaitedEntry, 1) : undefined
+  }
+
+  promptQueueDraining.add(sessionID)
+  clearPromptQueueTimer(sessionID)
+
+  let awaitedResult: InternalPromptDispatchResult | undefined
+  try {
+    while (true) {
+      const queue = promptQueues.get(sessionID)
+      const entry = queue?.[0]
+      if (!entry) {
+        break
+      }
+
+      promptQueueInFlight.set(sessionID, entry)
+      const result = await dispatchAfterSessionIdle({
+        sessionName: entry.sessionName,
+        client: entry.client,
+        sessionID: entry.sessionID,
+        input: entry.input,
+        source: entry.source,
+        dedupeKey: entry.dedupeKey,
+        settleMs: entry.settleMs,
+        postDispatchHoldMs: entry.postDispatchHoldMs,
+        dispatchTimeoutMs: entry.dispatchTimeoutMs,
+        checkStatus: entry.checkStatus,
+        checkToolState: entry.checkToolState,
+        dispatch: entry.dispatch,
+      })
+      if (promptQueueInFlight.get(sessionID)?.id === entry.id) {
+        promptQueueInFlight.delete(sessionID)
+      }
+
+      if (result.status === "active" || result.status === "reserved") {
+        const queued = queuedResult(
+          entry,
+          1,
+          result.status === "reserved" ? result.reservedBy : entry.source,
+        )
+        if (awaitedEntry?.id === entry.id) {
+          awaitedResult = queued
+        }
+        schedulePromptQueueDrain(sessionID, entry.queueRetryMs)
+        break
+      }
+
+      removePromptQueueEntry(sessionID, entry)
+      if (awaitedEntry?.id === entry.id) {
+        awaitedResult = result
+      }
+
+      const remainingQueue = promptQueues.get(sessionID)
+      if (!remainingQueue || remainingQueue.length === 0) {
+        break
+      }
+
+      schedulePromptQueueDrain(sessionID, entry.postDispatchHoldMs)
+      break
+    }
+  } finally {
+    promptQueueDraining.delete(sessionID)
+  }
+
+  return awaitedResult
+}
+
+async function enqueueInternalPrompt(entry: QueuedInternalPrompt): Promise<InternalPromptDispatchResult> {
+  const activeReservation = getActiveReservation(entry.sessionID)
+  if (activeReservation?.dedupeKey === entry.dedupeKey) {
+    log("[prompt-async-gate] queued prompt coalesced with recent dispatch", {
+      sessionID: entry.sessionID,
+      source: entry.source,
+      queuedBy: activeReservation.source,
+    })
+    return queuedResult(entry, 0, activeReservation.source)
+  }
+
+  const queue = getPromptQueue(entry.sessionID)
+  const existingIndex = queue.findIndex((queued) => queued.dedupeKey === entry.dedupeKey)
+  if (existingIndex >= 0) {
+    const existing = queue[existingIndex]
+    if (existing) {
+      log("[prompt-async-gate] queued prompt coalesced with pending dispatch", {
+        sessionID: entry.sessionID,
+        source: entry.source,
+        queuedBy: existing.source,
+        position: existingIndex + 1,
+      })
+      return queuedResult(existing, existingIndex + 1)
+    }
+  }
+
+  queue.push(entry)
+  log("[prompt-async-gate] queued prompt accepted", {
+    sessionID: entry.sessionID,
+    source: entry.source,
+    position: queue.length,
+  })
+
+  if (queue.length > 1 || promptQueueDraining.has(entry.sessionID)) {
+    schedulePromptQueueDrain(entry.sessionID, 0)
+    return queuedResult(entry, queue.length)
+  }
+
+  const result = await drainPromptQueue(entry.sessionID, entry)
+  return result ?? queuedResult(entry, 1)
+}
+
 export async function dispatchInternalPrompt<TInput = PromptAsyncInput>(
  args: InternalPromptDispatchArgs<TInput>,
 ): Promise<InternalPromptDispatchResult> {
@@ -457,6 +709,8 @@ export async function dispatchInternalPrompt<TInput = PromptAsyncInput>(
    source,
    settleMs = DEFAULT_SESSION_IDLE_SETTLE_MS,
  } = args
+  const dedupeKey = args.dedupeKey ?? createDefaultDedupeKey(source, input)
+  const queueRetryMs = args.queueRetryMs ?? DEFAULT_PROMPT_QUEUE_RETRY_MS
  const postDispatchHoldMs = args.postDispatchHoldMs ?? DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS
  const dispatchTimeoutMs = args.dispatchTimeoutMs ?? DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS
  const sessionName = args.mode === "async" ? "promptAsync" : "prompt"
@@ -483,12 +737,59 @@ export async function dispatchInternalPrompt<TInput = PromptAsyncInput>(
    return { status: "unavailable" }
  }

+  if (args.queueBehavior === "defer") {
+    const activeReservation = getActiveReservation(sessionID)
+    if (activeReservation) {
+      return { status: "reserved", reservedBy: activeReservation.source }
+    }
+
+    const queuedBy = getQueuedPromptBlocker(sessionID)
+    if (queuedBy !== undefined || promptQueueDraining.has(sessionID)) {
+      return { status: "reserved", reservedBy: queuedBy ?? source }
+    }
+
+    return dispatchAfterSessionIdle({
+      sessionName,
+      client,
+      sessionID,
+      input,
+      source,
+      dedupeKey,
+      settleMs,
+      postDispatchHoldMs,
+      dispatchTimeoutMs,
+      checkStatus: args.checkStatus !== false,
+      checkToolState: args.checkToolState !== false,
+      dispatch,
+    })
+  }
+
+  if (args.queue !== false) {
+    return enqueueInternalPrompt({
+      id: promptQueueSequence += 1,
+      sessionID,
+      sessionName,
+      client,
+      input,
+      source,
+      dedupeKey,
+      settleMs,
+      postDispatchHoldMs,
+      dispatchTimeoutMs,
+      queueRetryMs,
+      checkStatus: args.checkStatus !== false,
+      checkToolState: args.checkToolState !== false,
+      dispatch: dispatch as (dispatchInput: unknown) => Promise<unknown>,
+    })
+  }
+
  return dispatchAfterSessionIdle({
    sessionName,
    client,
    sessionID,
    input,
    source,
+    dedupeKey,
    settleMs,
    postDispatchHoldMs,
    dispatchTimeoutMs,
@@ -500,9 +801,20 @@ export async function dispatchInternalPrompt<TInput = PromptAsyncInput>(

 export function releaseAllPromptAsyncReservationsForTesting(): void {
  promptAsyncReservations.clear()
+  promptQueues.clear()
+  promptQueueDraining.clear()
+  promptQueueInFlight.clear()
+  for (const timer of promptQueueTimers.values()) {
+    clearTimeout(timer)
+  }
+  promptQueueTimers.clear()
  promptGateMessagesFetchTimeoutMsForTesting = undefined
 }

+export function isInternalPromptDispatchAccepted(result: InternalPromptDispatchResult): boolean {
+  return result.status === "dispatched" || result.status === "queued"
+}
+
 export function releasePromptAsyncReservation(
  sessionID: string,
  source: string,
@@ -524,6 +836,13 @@ export function releasePromptAsyncReservation(
  }

  promptAsyncReservations.delete(sessionID)
+  const inFlight = promptQueueInFlight.get(sessionID)
+  if (inFlight?.dedupeKey === existing.dedupeKey) {
+    removePromptQueueEntry(sessionID, inFlight)
+    promptQueueInFlight.delete(sessionID)
+    promptQueueDraining.delete(sessionID)
+  }
+  schedulePromptQueueDrain(sessionID, 0)
  log("[prompt-async-gate] promptAsync reservation released", {
    sessionID,
    source,