think: align lifecycle hook contexts with AI SDK; fix tool-call hook timing and extension dispatch#1340
Conversation
…timing and extension dispatch Resolves #1339. Lifecycle hook context types are now derived from the AI SDK so users get full typed access instead of `unknown`: - `StepContext<TOOLS>` = `StepResult<TOOLS>` — `reasoning`, `sources`, `files`, `providerMetadata` (cache tokens), `request`/`response`, `warnings`, full `LanguageModelUsage` (incl. `cachedInputTokens`, `reasoningTokens`). - `ChunkContext<TOOLS>` = AI SDK's `StreamTextOnChunkCallback` event (discriminated `TextStreamPart`). - `ToolCallContext<TOOLS>` spreads `TypedToolCall<TOOLS>` plus `stepNumber`, `messages`, `abortSignal`. - `ToolCallResultContext<TOOLS>` spreads `TypedToolCall<TOOLS>` plus `durationMs`, `messages`, `stepNumber`, and a discriminated `success`/`output`/`error` outcome. The relevant AI SDK types (`StepResult`, `TextStreamPart`, `TypedToolCall`, `TypedToolResult`) are re-exported from the package. Tool-call hook timing is now correct: `beforeToolCall` fires before `execute` (Think wraps every tool's `execute`), and `afterToolCall` fires after via `experimental_onToolCallFinish` with accurate `durationMs`. Previously both fired post-execution from `onStepFinish` data, and the wrapper was reading `tc.args` / `tr.result` (the AI SDK uses `input` / `output`) so subclass hooks always saw `{}` and `undefined`. `ToolCallDecision` is now functional. Returning `{ action: "block", reason }`, `{ action: "substitute", output }`, or `{ action: "allow", input }` from `beforeToolCall` actually intercepts execution. Extension dispatch for `beforeToolCall` / `afterToolCall` / `onStepFinish` / `onChunk`. The `ExtensionManifest.hooks` array claimed support for these but Think only ever dispatched `beforeTurn`. All five hooks now dispatch with JSON-safe snapshots. Extension hook handlers also receive `(snapshot, host)` symmetric with tool `execute`; previously only tool executes got the host bridge. Test mocks updated to the AI SDK v3 LanguageModel spec (`finishReason: { unified, raw }`, structured `usage` with `{ inputTokens: { total, ... }, outputTokens: { total, text, reasoning } }`, explicit `tool-call` chunks). The previous mocks were quietly relying on fallbacks in the wrapper that masked these shape mismatches. Breaking renames per AI SDK conventions: `ToolCallContext.args` -> `input`, `ToolCallResultContext.args` -> `input`, `ToolCallResultContext.result` -> `output`. `afterToolCall` is now a discriminated union — read `output` only when `ctx.success === true`, and `error` when `ctx.success === false`. Equivalent renames on `ToolCallDecision`. 247/247 tests passing. Docs and the assistant example are updated. Made-with: Cursor
🦋 Changeset detectedLatest commit: 35578bc The changes in this PR will be included in the next version bump. This PR includes changesets to release 2 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
`waitForIdle` and `waitUntilStable` previously only awaited `_turnQueue.waitForIdle()`. They missed submits that had passed `_getSubmitConcurrencyDecision` (so already bumped `_latestOverlappingSubmitSequence` and incremented `_pendingEnqueueCount`) but hadn't yet reached `_runExclusiveChatTurn` because they were mid-`persistMessages`. Anything calling these helpers (tests, recovery code, callers waiting for quiescence) could race the in-flight submit and observe the agent as "idle" when it really wasn't. Both helpers now poll on `_pendingEnqueueCount === 0` after the queue drains, with a brief setTimeout yield between checks so in-flight handlers can reach the enqueue call. Fixes the long-standing flake in `merge concatenates overlapping queued user messages into one follow-up turn`. Also bumps that test's stream durations (10×100ms → 15×150ms) to give the WS dispatch enough headroom under CI load to bump `_latestOverlappingSubmitSequence` *before* the first turn finishes — the supersession check on the second submit needs to see the latest counter value, and `waitForOverlappingSubmits` only helps if the third submit's handler has actually run. Made-with: Cursor
|
Pushed an additional commit (3bf23ae) covering an unrelated-but-tangential
Includes a patch changeset ( |
agents
@cloudflare/ai-chat
@cloudflare/codemode
hono-agents
@cloudflare/shell
@cloudflare/think
@cloudflare/voice
@cloudflare/worker-bundler
commit: |
…nwrapped
`_wrapToolsWithDecision` previously did:
const result = originalExecute(finalInput, options);
if (Symbol.asyncIterator in (result as object)) { ...collapse... }
For an `async function execute(...) { return makeAsyncIterable(); }`
(the common shape if you build the iterator inside the async body),
`originalExecute(...)` returns `Promise<AsyncIterable>` — and
`Symbol.asyncIterator in Promise` is always false. The collapse logic
was skipped. The wrapper is async, so it auto-unwrapped one Promise
layer and effectively returned `Promise<AsyncIterable>`. The AI SDK's
`executeTool` then `await`s that, gets the AsyncIterable, and yields
`{ type: "final", output: AsyncIterable }` — exactly the broken case
the wrapper's own comment warned against.
Fix: await `originalExecute(finalInput, options)` before inspecting,
so the symbol check sees the resolved value (an AsyncIterable in any
of: direct return, `async function*`, or `async () => makeIter()`).
Verified by adding two regression tests:
- `collapses Promise<AsyncIterable> returns to the last yielded value`
(the buggy case — fails on the pre-fix code with the iterator
serializing to `{}`)
- `collapses sync AsyncIterable returns to the last yielded value`
(the case that already worked — belt-and-suspenders)
Updated docs and JSDoc to spell out that all three iterator-return
shapes are handled identically.
Spotted by code review on #1340.
Made-with: Cursor
|
Pushed The reviewer was right — the collapse logic was effectively dead code for the most common shape of streaming tool. Trace:
Fix: Verification:
Patch changeset added ( |
Summary
Resolves #1339. Four logical changes, layered on top of each other:
Lifecycle hook context types are derived from the AI SDK.
StepContext,ChunkContext,ToolCallContext,ToolCallResultContextare now derived fromStepResult,TextStreamPart, andTypedToolCallinstead of being hand-typed subsets full ofunknown. Users get full typed access toreasoning,sources,files,providerMetadata(where Anthropic cache tokens likecacheCreationInputTokens/cacheReadInputTokenslive),request/response,warnings, and the fullLanguageModelUsage(totalTokens,reasoningTokens,cachedInputTokens). The relevant AI SDK types are re-exported from@cloudflare/think.Tool-call hook timing is now correct.
beforeToolCallruns before the tool'sexecute(Think wraps every tool'sexecute);afterToolCallruns after via the AI SDK'sexperimental_onToolCallFinishcallback with accuratedurationMsand a discriminatedsuccess/output/erroroutcome. Previously both fired post-execution fromonStepFinishdata, and the wrapper was readingtc.args/tr.result(the AI SDK usesinput/output) so subclass hooks always saw{}andundefined— a real latent bug.ToolCallDecisionis now functional. Returning{ action: "block", reason },{ action: "substitute", output }, or{ action: "allow", input }frombeforeToolCallactually intercepts execution. The decision types were declared but had no implementation before.Extension dispatch for the four observation hooks.
ExtensionManifest.hooksclaimed support forbeforeToolCall/afterToolCall/onStepFinish/onChunkbut Think only ever dispatchedbeforeTurn. All five hooks now dispatch to subscribed extensions with JSON-safe snapshots. Extension hook handlers also receive(snapshot, host)symmetric with toolexecute— previously only tool executes got the host bridge.Breaking renames
Per AI SDK conventions:
ToolCallContext.args→inputToolCallResultContext.args→input,ToolCallResultContext.result→outputafterToolCallis now a discriminated union — readoutputonly whenctx.success === true,errorwhenctx.success === falseToolCallDecisionThe package is
@experimentaland at0.2.x; minor bump.Test mocks brought up to spec
The previous
LanguageModelV3mocks were emittingfinishReason: "stop"(string) andusage: { inputTokens: 10, outputTokens: 5 }(flat). The current AI SDK expectsfinishReason: { unified, raw }and structuredusage: { inputTokens: { total, ... }, outputTokens: { total, text, reasoning } }. The wrapper had?? "stop"fallbacks that were quietly masking this. Mocks now use the correct shape.Test plan
tsc -p tsconfig.json --noEmit— cleantsc -p src/tests/tsconfig.json --noEmit— cleantsc --noEmitinexamples/assistant— cleannpm testinpackages/think— 247/247 passing (16 new this PR, 0 regressions on the 231 baseline)forwards the AI SDK's full StepResult — text and usage(Quality of live improvements for Think #1339 latent bug)forwards typed toolCalls/toolResults arrays(Quality of live improvements for Think #1339 latent bug)beforeToolCall receives toolName and typed input(thetc.argsbug)afterToolCall receives typed output(thetr.resultbug)ToolCallDecisioncases (void,allow+input,block+reason,block-default,substitute)async beforeToolCall (Promise<ToolCallDecision>) is awaited correctlya throwing beforeToolCall surfaces as a tool error in afterToolCallDocumentation
docs/think/lifecycle-hooks.md: rewrote the four context tables to match the SDK shapes; added typed-TOOLSexamples; documentedToolCallDecisionwith three concrete examples (block, substitute from cache, clamp model-emitted limit); added "Notes & limitations" section covering substituted-input not re-validated, AsyncIterable collapse, hook order, andstepNumberundefined in start ctx; new "Extension hook subscriptions" section with snapshot shapes per hook and theonChunkperf footgun.packages/think/README.md: contexts table, discriminatedafterToolCallexample,beforeToolCallinput-clamping example, extension subscriptions subsection, migration note for theargs/result→input/outputrename.examples/assistant/src/server.ts: updated to usectx.input/ctx.successdiscriminator withdurationMs.Follow-ups
Filed as separate issues:
AsyncIterablereturns) currently collapse whenbeforeToolCallis in play — preserving both is non-trivial and not blocking.output/ToolCallDecision.inputwould benefit from generic distribution overTOOLS._wrapToolsWithDecisionwhen nobeforeToolCalloverride + no extensions subscribe (perf nicety).needsApprovaltools.Made with Cursor