Skip to content

think: align lifecycle hook contexts with AI SDK; fix tool-call hook timing and extension dispatch#1340

Merged
threepointone merged 3 commits into
mainfrom
think-lifecycle-hooks-ai-sdk-alignment
Apr 18, 2026
Merged

think: align lifecycle hook contexts with AI SDK; fix tool-call hook timing and extension dispatch#1340
threepointone merged 3 commits into
mainfrom
think-lifecycle-hooks-ai-sdk-alignment

Conversation

@threepointone

@threepointone threepointone commented Apr 18, 2026

Copy link
Copy Markdown
Contributor

Summary

Resolves #1339. Four logical changes, layered on top of each other:

  1. Lifecycle hook context types are derived from the AI SDK. StepContext, ChunkContext, ToolCallContext, ToolCallResultContext are now derived from StepResult, TextStreamPart, and TypedToolCall instead of being hand-typed subsets full of unknown. Users get full typed access to reasoning, sources, files, providerMetadata (where Anthropic cache tokens like cacheCreationInputTokens / cacheReadInputTokens live), request/response, warnings, and the full LanguageModelUsage (totalTokens, reasoningTokens, cachedInputTokens). The relevant AI SDK types are re-exported from @cloudflare/think.

  2. Tool-call hook timing is now correct. beforeToolCall runs before the tool's execute (Think wraps every tool's execute); afterToolCall runs after via the AI SDK's experimental_onToolCallFinish callback with accurate durationMs and a discriminated success/output/error outcome. 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 — a real latent bug.

  3. ToolCallDecision is now functional. Returning { action: "block", reason }, { action: "substitute", output }, or { action: "allow", input } from beforeToolCall actually intercepts execution. The decision types were declared but had no implementation before.

  4. Extension dispatch for the four observation hooks. ExtensionManifest.hooks claimed support for beforeToolCall / afterToolCall / onStepFinish / onChunk but Think only ever dispatched beforeTurn. All five hooks now dispatch to subscribed extensions with JSON-safe snapshots. Extension hook handlers also receive (snapshot, host) symmetric with tool execute — previously only tool executes got the host bridge.

Breaking renames

Per AI SDK conventions:

  • ToolCallContext.argsinput
  • ToolCallResultContext.argsinput, ToolCallResultContext.resultoutput
  • afterToolCall is now a discriminated union — read output only when ctx.success === true, error when ctx.success === false
  • Equivalent renames on ToolCallDecision

The package is @experimental and at 0.2.x; minor bump.

Test mocks brought up to spec

The previous LanguageModelV3 mocks were emitting finishReason: "stop" (string) and usage: { inputTokens: 10, outputTokens: 5 } (flat). The current AI SDK expects finishReason: { unified, raw } and structured usage: { 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 — clean
  • tsc -p src/tests/tsconfig.json --noEmit — clean
  • tsc --noEmit in examples/assistant — clean
  • npm test in packages/think247/247 passing (16 new this PR, 0 regressions on the 231 baseline)
  • New regression coverage:
    • 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 (the tc.args bug)
    • afterToolCall receives typed output (the tr.result bug)
    • 5 ToolCallDecision cases (void, allow+input, block+reason, block-default, substitute)
    • async beforeToolCall (Promise<ToolCallDecision>) is awaited correctly
    • a throwing beforeToolCall surfaces as a tool error in afterToolCall
    • 4 extension dispatch tests (one per observation hook), end-to-end via marker files written through the host bridge

Documentation

  • docs/think/lifecycle-hooks.md: rewrote the four context tables to match the SDK shapes; added typed-TOOLS examples; documented ToolCallDecision with 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, and stepNumber undefined in start ctx; new "Extension hook subscriptions" section with snapshot shapes per hook and the onChunk perf footgun.
  • packages/think/README.md: contexts table, discriminated afterToolCall example, beforeToolCall input-clamping example, extension subscriptions subsection, migration note for the args/resultinput/output rename.
  • examples/assistant/src/server.ts: updated to use ctx.input / ctx.success discriminator with durationMs.

Follow-ups

Filed as separate issues:

  • Streaming preliminary tool results (AsyncIterable returns) currently collapse when beforeToolCall is in play — preserving both is non-trivial and not blocking.
  • Per-tool typing on output / ToolCallDecision.input would benefit from generic distribution over TOOLS.
  • Skip _wrapToolsWithDecision when no beforeToolCall override + no extensions subscribe (perf nicety).
  • Verify ordering with needsApproval tools.
  • Counter-based test for "block prevents execute from running" (currently verified indirectly).

Made with Cursor


Open in Devin Review

…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-bot

changeset-bot Bot commented Apr 18, 2026

Copy link
Copy Markdown

🦋 Changeset detected

Latest commit: 35578bc

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 2 packages
Name Type
@cloudflare/ai-chat Patch
@cloudflare/think Minor

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

@devin-ai-integration devin-ai-integration Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

Open in Devin Review

`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
@threepointone

Copy link
Copy Markdown
Contributor Author

Pushed an additional commit (3bf23ae) covering an unrelated-but-tangential @cloudflare/ai-chat fix that came out of investigating a flaky test in this PR's run:

  • waitForIdle / waitUntilStable now drain in-flight submits past the concurrency decision but not yet enqueued (i.e. mid-persistMessages). They previously only awaited _turnQueue.waitForIdle(), which could return while a submit was still tracked in _pendingEnqueueCount — racing with anything that called them (tests, recovery code, etc.).
  • Bumps the merge concatenates... test stream durations (10×100ms → 15×150ms) to give the WS dispatch enough headroom under CI load to bump _latestOverlappingSubmitSequence before the first turn finishes — waitForOverlappingSubmits only helps if the third submit's handler has actually run.

Includes a patch changeset (@cloudflare/ai-chat). 459/459 ai-chat tests passing, 247/247 think tests passing.

devin-ai-integration[bot]

This comment was marked as resolved.

@pkg-pr-new

pkg-pr-new Bot commented Apr 18, 2026

Copy link
Copy Markdown

Open in StackBlitz

agents

npm i https://pkg.pr.new/agents@1340

@cloudflare/ai-chat

npm i https://pkg.pr.new/@cloudflare/ai-chat@1340

@cloudflare/codemode

npm i https://pkg.pr.new/@cloudflare/codemode@1340

hono-agents

npm i https://pkg.pr.new/hono-agents@1340

@cloudflare/shell

npm i https://pkg.pr.new/@cloudflare/shell@1340

@cloudflare/think

npm i https://pkg.pr.new/@cloudflare/think@1340

@cloudflare/voice

npm i https://pkg.pr.new/@cloudflare/voice@1340

@cloudflare/worker-bundler

npm i https://pkg.pr.new/@cloudflare/worker-bundler@1340

commit: 35578bc

…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
@threepointone

Copy link
Copy Markdown
Contributor Author

Pushed 35578bc addressing the review feedback on _wrapToolsWithDecision.

The reviewer was right — the collapse logic was effectively dead code for the most common shape of streaming tool. Trace:

  • Original: async function execute(...) { return makeAsyncIterable(); }
  • Wrapper: const result = originalExecute(...); ← no awaitresult is Promise<AsyncIterable>
  • Symbol.asyncIterator in Promisefalse → collapse skipped
  • Wrapper is async, auto-unwraps one layer → returns Promise<AsyncIterable>
  • AI SDK's executeTool awaits it, gets the AsyncIterable, yields { type: "final", output: AsyncIterable } — exactly what the wrapper's own comment said would be broken.

Fix: const result = await originalExecute(finalInput, options);. This handles all three shapes uniformly (sync return, async function*, async () => makeIter()).

Verification:

  • New test collapses Promise<AsyncIterable> returns to the last yielded value — fails on the pre-fix code (expect("echo: hello").toBe(JSON.parse("{}"))), passes after the await.
  • New test collapses sync AsyncIterable returns to the last yielded value — belt-and-suspenders for the case that already worked.
  • 249/249 think tests passing. Docs and JSDoc updated to spell out that all three iterator-return shapes are handled.

Patch changeset added (@cloudflare/think).

@threepointone threepointone merged commit 3cbe776 into main Apr 18, 2026
3 checks passed
@threepointone threepointone deleted the think-lifecycle-hooks-ai-sdk-alignment branch April 18, 2026 23:27
@github-actions github-actions Bot mentioned this pull request Apr 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Quality of live improvements for Think

1 participant