basalt-ui 1.10.0 → 1.11.0

package/src/agent/use-agent-thread-runs.test.tsx

@@ -1,10 +1,12 @@
 /**
- * useAgentThreadRuns — concurrency, busy-thread no-op, finalize ordering, failure status, retry.
+ * useAgentThreadRuns — concurrency, busy-thread no-op, finalize ordering, failure status, retry,
+ * and (as of B2) the stop()/stopAll() lifecycle (partial-turn preservation, the idle no-op
+ * guard, and the resolveOutcome-race fix).
  *
- * Scope: the multi-thread run manager AS THE PUBLIC API EXISTS TODAY (B1). No mergePart/part ids/
+ * Scope: the multi-thread run manager AS THE PUBLIC API EXISTS TODAY. No mergePart/part ids/
  * offsets, no seven-state ToolCallPart, no ResumableAgentTransport/isResumable/idempotentReplay
- * beyond the plain optional `resume` already on AgentTransport, no stop-preserves-partial-turn, no
- * ThreadsStoreAdapter — all later releases, not exercised here.
+ * beyond the plain optional `resume` already on AgentTransport, no ThreadsStoreAdapter — later
+ * releases, not exercised here.
  *
  * The mount-reconcile/resume path (StartPart handling, orphaned-thread resumption on mount) has
  * its own file, `use-agent-thread-runs.resume.test.tsx`; the F3 wedge has its own file,
@@ -345,15 +347,12 @@ describe('useAgentThreadRuns — concurrency, no-op guard, finalize order, failu
     expect(store.threads.find((t) => t.id === threadId)?.messages).toHaveLength(4)
   })
 
-  // Flagged for scheduling, not fixed here (out of this task's scope — see the brief's "while you
-  // are in there, report, do not fix" note (i)): consumeAndFinalize's success path awaits
-  // resolveOutcome (:219) and then writes setOutcome/setStatus (:221-222) WITHOUT re-checking the
-  // supersede/abort guards it checked right before the await (:205-207). A slow resolveOutcome
-  // racing a stop() call is reachable: stop() sets the thread to 'done' immediately, but once the
-  // delayed resolveOutcome finally resolves, the unconditional setStatus(threadId, outcome.status)
-  // silently overwrites that 'done' with whatever the outcome resolved to. Verified failing against
-  // the current implementation before being marked .skip here.
-  test.skip('stop() during a slow resolveOutcome must not be clobbered by the outcome settling later', async () => {
+  // Pins the fix for a race in consumeAndFinalize's success path: it awaits resolveOutcome and
+  // then writes setOutcome/setStatus, and now RE-CHECKS the supersede/abort guards after that
+  // await (not just before it) before performing those writes. Without the recheck, a slow
+  // resolveOutcome racing a stop() call would clobber stop()'s synchronous 'done' with whatever
+  // the outcome resolved to once it finally settled.
+  test('stop() during a slow resolveOutcome must not be clobbered by the outcome settling later', async () => {
     const store = createTestThreadsStore()
     const threadId = store.create()
 
@@ -391,8 +390,392 @@ describe('useAgentThreadRuns — concurrency, no-op guard, finalize order, failu
     releaseOutcome?.()
 
     await waitFor(() => {
-      // Currently fails: the late setStatus overwrites 'done' with 'attention'.
+      // The late resolveOutcome settling must not resurrect the thread past the 'done' stop()
+      // already committed.
       expect(store.threads.find((t) => t.id === threadId)?.status).toBe('done')
     })
   })
 })
+
+describe('useAgentThreadRuns — stop()/stopAll() lifecycle', () => {
+  test('stop() mid-stream persists the accumulated parts as a finish: "stopped" assistant message', async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    const gate = deferred<void>()
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream() {
+        yield { type: 'text', text: 'partial' }
+        await gate.promise
+        yield { type: 'text', text: 'never arrives' }
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    act(() => {
+      result.current.start(threadId, 'hi')
+    })
+
+    await waitFor(() => {
+      expect(result.current.runs.get(threadId)?.parts).toHaveLength(1)
+    })
+
+    act(() => {
+      result.current.stop(threadId)
+    })
+
+    const thread = store.threads.find((t) => t.id === threadId)
+    expect(thread?.status).toBe('done')
+    expect(thread?.messages).toHaveLength(2)
+    const assistantMessage = thread?.messages[1]
+    expect(assistantMessage?.role).toBe('assistant')
+    expect(assistantMessage?.finish).toBe('stopped')
+    // withPartIds stamps the missing id as `${threadId}#0` — the first (and only) id-less part
+    // this run's generator yielded.
+    expect(assistantMessage?.parts).toEqual([
+      { id: `${threadId}#0`, type: 'text', text: 'partial' },
+    ])
+    // The run entry is torn down synchronously — no lingering 'streaming' indicator.
+    expect(result.current.runs.has(threadId)).toBe(false)
+
+    // Release the still-suspended (aborted, guard-caught) generator so it doesn't leak into the
+    // next test.
+    gate.resolve(undefined)
+  })
+
+  test('stop() does not double-append when the stream already fully arrived before stop() was called', async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    let releaseOutcome: (() => void) | undefined
+    const slowResolveOutcome: OutcomeResolver<AgentPart> = () =>
+      new Promise<AgentOutcome>((resolve) => {
+        releaseOutcome = () => resolve({ title: 't', summary: 's', status: 'done' })
+      })
+
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream() {
+        yield { type: 'text', text: 'hi' }
+      },
+    }
+
+    const { result } = renderHook(() =>
+      useAgentThreadRuns({ transport, store, resolveOutcome: slowResolveOutcome }),
+    )
+
+    act(() => {
+      result.current.start(threadId, 'hi')
+    })
+
+    // Wait until consumeAndFinalize's own appendMessage has already landed (it happens BEFORE
+    // its await on resolveOutcome) — this is the narrow window stop() can race without ever
+    // being able to observe "not yet appended" (see finalizeStop's doc comment).
+    await waitFor(() => {
+      expect(store.threads.find((t) => t.id === threadId)?.messages.length).toBe(2)
+    })
+
+    act(() => {
+      result.current.stop(threadId)
+    })
+
+    // Still exactly 2 messages (user + the ORIGINAL 'complete' assistant message) — stop() must
+    // recognize the terminal message already landed and skip its own append rather than
+    // duplicating it as a second 'stopped' message.
+    const thread = store.threads.find((t) => t.id === threadId)
+    expect(thread?.messages).toHaveLength(2)
+    expect(thread?.messages[1]?.finish).toBe('complete')
+
+    releaseOutcome?.()
+  })
+
+  test('stop() is a true no-op when threadId has no in-flight run', () => {
+    const store = createTestThreadsStore()
+    // No store.create() here on purpose — mirrors the "retry() is a no-op" test above: isolates
+    // stop()'s own "nothing in flight" guard from the mount-reconcile effect's unrelated handling
+    // of a freshly-created 'pending' thread.
+    const threadId = crypto.randomUUID()
+
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream() {
+        yield { type: 'text', text: 'x' }
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    act(() => {
+      result.current.stop(threadId)
+    })
+
+    expect(store.threads).toEqual([])
+  })
+
+  // Pins the fix for the reference-equality hole: runsRef mirrors `runs` only as of the last
+  // FLUSHED render (assigned during render, `:~355`), so a stop() landing inside the narrow
+  // window between consumeAndFinalize's post-loop appendMessage and its `await resolveOutcome`
+  // resolving could observe a STALE parts array (one delta behind) via runsRef — defeating the
+  // old `lastMessage.parts === parts` guard and double-appending. This test lands stop()
+  // deterministically inside that exact window (fired synchronously from inside the injected
+  // resolveOutcome, itself called synchronously right after appendMessage, before any of the
+  // hook's pending state updates can flush — flushing needs a macrotask, and nothing here crosses
+  // one) rather than relying on timing luck, so it is a reliable regression guard for the fixed
+  // marker-based mechanism, not just a one-off repro.
+  test("stop() fired synchronously from inside resolveOutcome (before the final delta's render flushes) must not double-append", async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    const gate = deferred<void>()
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream() {
+        // Distinct ids: mergePart addresses by id, and two parts sharing the (undefined) id would
+        // splice into ONE entry instead of accumulating two — this test needs the SECOND delta to
+        // produce a genuinely NEW array entry (and therefore a new `parts` array reference) to
+        // exercise the race at all.
+        yield { id: 'p1', type: 'text', text: 'first' }
+        await gate.promise
+        yield { id: 'p2', type: 'text', text: 'second' }
+      },
+    }
+
+    const hookRef: { stop?: (id: string) => void } = {}
+    let resolveCalls = 0
+    const raceResolveOutcome: OutcomeResolver<AgentPart> = async () => {
+      resolveCalls += 1
+      // Fire stop() from inside the FIRST resolveOutcome call — synchronously, before this
+      // function even returns its promise. Everything from the second delta's setRuns dispatch to
+      // this point runs purely through microtasks (async-generator machinery + this direct call),
+      // so no macrotask has elapsed and React's scheduler cannot have flushed that dispatch into
+      // runsRef yet — this reliably lands stop() inside the race window.
+      if (resolveCalls === 1) {
+        hookRef.stop?.(threadId)
+      }
+      return { title: 't', summary: 's', status: 'done' }
+    }
+
+    const { result } = renderHook(() =>
+      useAgentThreadRuns({ transport, store, resolveOutcome: raceResolveOutcome }),
+    )
+    hookRef.stop = result.current.stop
+
+    act(() => {
+      result.current.start(threadId, 'hi')
+    })
+
+    // Let the FIRST delta's render flush — runsRef now mirrors a non-empty parts array
+    // (['first']), which is what makes the eventual mismatch observable rather than vacuously
+    // guarded by parts.length === 0.
+    await waitFor(() => {
+      expect(result.current.runs.get(threadId)?.parts).toHaveLength(1)
+    })
+
+    // Resolve the gate outside act() on purpose: this must not synchronously flush React state —
+    // the whole point is to let the second delta, the loop exit, and appendMessage run purely as
+    // microtask continuations, with stop() firing from inside resolveOutcome before any render for
+    // the second delta occurs.
+    gate.resolve(undefined)
+
+    await waitFor(() => {
+      expect(store.threads.find((t) => t.id === threadId)?.status).toBe('done')
+    })
+
+    const thread = store.threads.find((t) => t.id === threadId)
+    const assistantMessages = thread?.messages.filter((m) => m.role === 'assistant') ?? []
+    expect(assistantMessages).toHaveLength(1)
+    // The surviving message must be the ORIGINAL 'complete' one with the full (['first','second'])
+    // content — not a truncated 'stopped' message built from the stale ['first']-only parts.
+    expect(assistantMessages[0]?.finish).toBe('complete')
+    expect(assistantMessages[0]?.parts).toEqual([
+      { id: 'p1', type: 'text', text: 'first' },
+      { id: 'p2', type: 'text', text: 'second' },
+    ])
+  })
+
+  // Pins the flip side of the marker fix: the marker must be cleared once a run concludes, or a
+  // stale entry from an EARLIER finished run on the same thread would wrongly suppress a
+  // legitimate append on a LATER run — a worse bug than the double-append this mechanism exists to
+  // prevent.
+  test('a stale marker from an earlier finished run must not suppress a legitimate later stop() append', async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    const inputs: string[] = []
+    let gate = deferred<void>()
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream(input) {
+        inputs.push(input)
+        // Distinct ids: two id-less parts sharing the (undefined) id would silently splice into
+        // ONE entry (the F1 bug) — turn 1's own "did these two deltas stay separate?" assertion
+        // below is invisible otherwise, since it only checked assistant MESSAGE count, not parts.
+        yield { id: `${input}-partial`, type: 'text', text: `${input}-partial` }
+        await gate.promise
+        yield { id: `${input}-final`, type: 'text', text: `${input}-final` }
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    // Turn 1: runs to full completion (no stop()) so consumeAndFinalize's own success path sets
+    // and then clears the marker for this thread.
+    act(() => {
+      result.current.start(threadId, 'first-turn')
+    })
+    gate.resolve(undefined)
+    await waitFor(() => {
+      expect(store.threads.find((t) => t.id === threadId)?.status).toBe('done')
+    })
+    const turn1Assistant = store.threads
+      .find((t) => t.id === threadId)
+      ?.messages.filter((m) => m.role === 'assistant')
+    expect(turn1Assistant).toHaveLength(1)
+    // The two deltas landed as two SEPARATE parts, not spliced into one.
+    expect(turn1Assistant?.[0]?.parts).toEqual([
+      { id: 'first-turn-partial', type: 'text', text: 'first-turn-partial' },
+      { id: 'first-turn-final', type: 'text', text: 'first-turn-final' },
+    ])
+
+    // Turn 2: a genuinely new run on the SAME thread, stopped mid-stream. If a stale marker from
+    // turn 1 leaked, this stop() would be wrongly treated as "already persisted" and silently drop
+    // the partial content instead of appending it.
+    gate = deferred<void>()
+    act(() => {
+      result.current.retry(threadId)
+    })
+    expect(inputs).toEqual(['first-turn', 'first-turn'])
+
+    await waitFor(() => {
+      expect(result.current.runs.get(threadId)?.parts).toHaveLength(1)
+    })
+
+    act(() => {
+      result.current.stop(threadId)
+    })
+
+    const thread = store.threads.find((t) => t.id === threadId)
+    const assistantMessages = thread?.messages.filter((m) => m.role === 'assistant') ?? []
+    expect(assistantMessages).toHaveLength(2)
+    expect(assistantMessages[1]?.finish).toBe('stopped')
+    expect(assistantMessages[1]?.parts).toEqual([
+      { id: 'first-turn-partial', type: 'text', text: 'first-turn-partial' },
+    ])
+
+    gate.resolve(undefined)
+  })
+
+  test('stopAll() resolves every aborted thread instead of stranding them in "streaming"', async () => {
+    const store = createTestThreadsStore()
+    const threadA = store.create()
+    const threadB = store.create()
+
+    const gateA = deferred<void>()
+    const gateB = deferred<void>()
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream(input) {
+        if (input === 'a') {
+          yield { type: 'text', text: 'partial-a' }
+          await gateA.promise
+        } else {
+          yield { type: 'text', text: 'partial-b' }
+          await gateB.promise
+        }
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    act(() => {
+      result.current.start(threadA, 'a')
+      result.current.start(threadB, 'b')
+    })
+
+    await waitFor(() => {
+      expect(result.current.runs.get(threadA)?.parts).toHaveLength(1)
+      expect(result.current.runs.get(threadB)?.parts).toHaveLength(1)
+    })
+
+    act(() => {
+      result.current.stopAll()
+    })
+
+    expect(store.threads.find((t) => t.id === threadA)?.status).toBe('done')
+    expect(store.threads.find((t) => t.id === threadB)?.status).toBe('done')
+    expect(result.current.runs.size).toBe(0)
+
+    // Release the still-suspended (aborted, guard-caught) generators so they don't leak into the
+    // next test.
+    gateA.resolve(undefined)
+    gateB.resolve(undefined)
+  })
+
+  // F4 regression: stop()'s no-op guard used to read runsRef, which only mirrors `runs` as of the
+  // last FLUSHED RENDER. start() then stop() called in the SAME synchronous act() batch — e.g.
+  // `onClick={() => { start(id, text); stop(id) }}` — never lets a render land in between, so
+  // runsRef was still empty when stop() ran and it wrongly no-op'd: the controller was never
+  // aborted, and a transport that only resolves on its AbortSignal firing would strand the thread
+  // in 'streaming' forever (this is exactly what `stopAll()`'s own regression test above cannot
+  // see — it deliberately `waitFor`s a render to flush before calling stopAll()).
+  test('F4 regression: start() then stop() in the same synchronous tick must actually abort, not strand the thread in "streaming"', async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    // Hangs until its AbortSignal fires — never resolves on its own. If stop()'s no-op guard skips
+    // the abort, this generator never settles and the thread is stuck 'streaming' forever.
+    const transport: AgentTransport<AgentPart, string> = {
+      // oxlint-disable-next-line require-yield -- hangs until aborted; never has anything to yield
+      async *stream(_input, signal) {
+        await new Promise((resolve) => {
+          signal?.addEventListener('abort', () => resolve(undefined))
+        })
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    act(() => {
+      result.current.start(threadId, 'hi')
+      result.current.stop(threadId)
+    })
+
+    await waitFor(() => {
+      expect(store.threads.find((t) => t.id === threadId)?.status).toBe('done')
+    })
+    expect(result.current.runs.has(threadId)).toBe(false)
+  })
+})
+
+describe('useAgentThreadRuns — F1 regression: id-less drafts within one run', () => {
+  test('drafts of different types do not destroy each other — mergePart no longer sees undefined === undefined', async () => {
+    const store = createTestThreadsStore()
+    const threadId = store.create()
+
+    const transport: AgentTransport<AgentPart, string> = {
+      async *stream() {
+        yield { type: 'text', text: 'Hi' }
+        yield {
+          type: 'tool',
+          toolCallId: 'c1',
+          toolName: 'search',
+          state: 'input-available',
+          input: { q: 'x' },
+        }
+        yield { type: 'text', text: 'Bye' }
+      },
+    }
+
+    const { result } = renderHook(() => useAgentThreadRuns({ transport, store, resolveOutcome }))
+
+    act(() => {
+      result.current.start(threadId, 'go')
+    })
+
+    await waitFor(() => {
+      expect(store.threads.find((t) => t.id === threadId)?.status).toBe('done')
+    })
+
+    const assistantMessage = store.threads
+      .find((t) => t.id === threadId)
+      ?.messages.find((m) => m.role === 'assistant')
+    expect(assistantMessage?.parts).toHaveLength(3)
+    expect(assistantMessage?.parts.map((part) => part.type)).toEqual(['text', 'tool', 'text'])
+  })
+})