Skip to content

callback_wrapper middleware around every user-callback dispatch (#47 phase 2) #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ya-luotao merged 2 commits into from
Jul 31, 2026
Merged

callback_wrapper middleware around every user-callback dispatch (#47 phase 2) #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
db8b55e
Add callback_wrapper middleware around every user-callback dispatch (…
ya-luotao Jul 31, 2026
14a6fa9
Address codex review: atomic dispatch-pair resolution, break invisibl…
ya-luotao Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## [Unreleased]

### Added
- **`ClaudeAgentOptions#callback_wrapper`** (#47 phase 2): optional middleware wrapped around every user-callback dispatch (message blocks, observers, hooks, permission callbacks, SDK MCP handlers). A callable receiving a zero-arg invocation that it must call and return: `callback_wrapper: ->(inv) { Rails.application.executor.wrap { inv.call } }`. The wrapper runs on the same execution context as the callback — inside the worker thread in the default `:thread` mode (so `executor.wrap` checks ActiveRecord connections back in when the callback ends, retiring the stranded-connection workaround without adopting `:inline`), in place on the reactor fiber in `:inline` mode. Exceptions propagate through it unchanged. Also settable per SDK MCP server for direct calls (`server.callback_wrapper=`); session dispatches carry the session's wrapper via the same fiber-storage scope as `callback_scheduling`. See "Rails executor around callbacks" in docs/rails.md.

## [0.25.0] - 2026-07-31

### Added
Expand Down
2 changes: 1 addition & 1 deletion CLAUDE.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ Optional adapter for mirroring session transcripts to external storage (the subp

`async` installs a Fiber scheduler, but most Ruby libraries (pg, mysql2, ActiveRecord pools) key state on `Thread.current` and are thread-safe, not fiber-safe. `FiberBoundary.invoke` (`lib/claude_agent_sdk/fiber_boundary.rb`) hops every user-supplied callback (tool handlers, hooks, permission callbacks, message blocks, observers) to a plain thread before invoking it. Consequence: the thread hop severs `break`/`return`/`next` from the surrounding method — SDK loops yielding user callbacks must keep loop control outside the invoked block (see `Client#receive_response`; user `break` is bridged via `.invoke_iteration`).

Opt-in `ClaudeAgentOptions#callback_scheduling: :inline` (issue #47) skips the hop and runs callbacks in place on the reactor fiber — for hosts that are fiber-isolated end to end (solid_queue fiber workers with `IsolatedExecutionState.isolation_level = :fiber`). The mode is plumbed per-call (Query kwarg; the SDK-MCP dispatch path carries it via fiber storage — `Fiber[FiberBoundary::SCHEDULING_KEY]` — so a server shared by concurrent sessions is never mutated), never via a thread-local. Timeout-bounded store-adapter calls always hop regardless; inline hook timeouts are cooperative (`with_timeout`). Public `ClaudeAgentSDK.offload { }` hops one heavy block manually.
Opt-in `ClaudeAgentOptions#callback_scheduling: :inline` (issue #47) skips the hop and runs callbacks in place on the reactor fiber — for hosts that are fiber-isolated end to end (solid_queue fiber workers with `IsolatedExecutionState.isolation_level = :fiber`). The mode is plumbed per-call (Query kwarg; the SDK-MCP dispatch path carries it via fiber storage — `Fiber[FiberBoundary::SCHEDULING_KEY]` — so a server shared by concurrent sessions is never mutated), never via a thread-local. Timeout-bounded store-adapter calls always hop regardless; inline hook timeouts are cooperative (`with_timeout`). Public `ClaudeAgentSDK.offload { }` hops one heavy block manually. `ClaudeAgentOptions#callback_wrapper` (issue #47 phase 2) is optional middleware composed around every user-callback dispatch BEFORE the hop, so it runs on the callback's own execution context (worker thread in `:thread`, reactor fiber in `:inline`) — e.g. `->(inv) { Rails.application.executor.wrap { inv.call } }`; it travels the SDK-MCP dispatch path inside the same fiber-storage scope as the mode.

### Global Configuration

Expand Down
16 changes: 16 additions & 0 deletions docs/rails.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,22 @@ end

The trade-off: because callbacks run on a plain thread rather than inside an `Async::Task`, fiber-specific primitives aren't available to them — `Async::Task.current` will raise "No async task available". If a callback wants cooperative concurrency it should open its own `Async { }` block. In practice, callbacks typically do some Ruby work, call external services, and return — so this rarely matters. If you wrap your own call site in an outer `Async { }` block, the scheduler is visible to your code again; you've opted in, and whatever fiber-safety rules your app uses apply there.

### Rails executor around callbacks: `callback_wrapper`

One consequence of the thread hop: an ActiveRecord connection implicitly checked out inside a callback belongs to that throwaway thread and stays stranded until the pool reaper reclaims it. Rails' own answer to "code running on a thread Rails didn't create" is the executor — and `callback_wrapper` lets you install it around every user-callback dispatch:

```ruby
ClaudeAgentSDK.configure do |config|
config.default_options = {
callback_wrapper: ->(invocation) { Rails.application.executor.wrap { invocation.call } }
}
end
```

The wrapper is a callable receiving a zero-arg `invocation`; it must call it and return its value. It runs on the **same execution context as the callback** — inside the worker thread in `:thread` mode, which is the whole point: `executor.wrap` runs on the thread that touches ActiveRecord, so connections check back in when the callback ends. Exceptions from the callback propagate through the wrapper unchanged (don't rescue them); `ensure`-based wrappers like `executor.wrap` are safe, including around a `break` from a message block. Beyond the executor, this is a generic hook for APM span propagation, `CurrentAttributes`/logging context, etc.

When do you want this vs `callback_scheduling: :inline`? `callback_wrapper` + default `:thread` mode is the right choice for ordinary threaded hosts (Puma, threaded Sidekiq/solid_queue): it fixes connection hygiene without any fiber-isolation precondition. `:inline` is only for hosts that are fiber-isolated end to end (solid_queue fiber workers with `isolation_level = :fiber`); there the wrapper still applies — it simply runs in place on the reactor fiber.

## Fiber workers (solid_queue) and `callback_scheduling: :inline`

[solid_queue 728](https://github.com/rails/solid_queue/pull/728) added a fiber-based worker mode: workers configured with `fibers: N` run claimed jobs as fibers on one async reactor thread — built for exactly the long-lived, I/O-bound "LLM streaming" jobs this SDK produces. It requires the app to be fiber-isolated end to end:
Expand Down
69 changes: 46 additions & 23 deletions lib/claude_agent_sdk.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,9 +84,9 @@ def self.extract_exclude_dynamic_sections(system_prompt)
# Each observer is invoked through FiberBoundary so that user code runs
# on a plain thread (no Fiber scheduler) even when called from inside
# the SDK's Async reactor — or in place when scheduling is :inline.
def self.notify_observers(observers, method, *args, scheduling: :thread)
def self.notify_observers(observers, method, *args, scheduling: :thread, wrapper: nil)
observers.each do |obs|
FiberBoundary.invoke(scheduling: scheduling) { obs.send(method, *args) }
FiberBoundary.invoke(scheduling: scheduling, wrapper: wrapper) { obs.send(method, *args) }
rescue StandardError, ScriptError
# ScriptError too: NotImplementedError < ScriptError (not
# StandardError), and a stubbed observer must never mask the original
Expand Down Expand Up @@ -187,13 +187,13 @@ def self.prompt_text_from_content(content)
# Wrap a streaming-input enumerable so observers get on_user_prompt for
# each user message before it is written to stdin. Identity when no
# observers are configured.
def self.observing_prompt_stream(prompt, observers, scheduling: :thread)
def self.observing_prompt_stream(prompt, observers, scheduling: :thread, wrapper: nil)
return prompt if observers.empty?

Enumerator.new do |yielder|
prompt.each do |message|
text = extract_user_prompt_text(message)
notify_observers(observers, :on_user_prompt, text, scheduling: scheduling) if text
notify_observers(observers, :on_user_prompt, text, scheduling: scheduling, wrapper: wrapper) if text
yielder << message
end
end
Expand Down Expand Up @@ -461,8 +461,10 @@ def self.query(prompt:, options: nil, transport: nil, &block)
# Resolve callable observers into fresh instances (thread-safe for global defaults)
resolved_observers = ClaudeAgentSDK.resolve_observers(configured_options.observers)

# Where user callbacks run (see ClaudeAgentOptions#callback_scheduling).
# Where user callbacks run (see ClaudeAgentOptions#callback_scheduling)
# and the middleware wrapped around them (#callback_wrapper).
callback_scheduling = configured_options.callback_scheduling || :thread
callback_wrapper = configured_options.callback_wrapper
ClaudeAgentSDK.check_inline_isolation(callback_scheduling)

raise ArgumentError, 'transport must respond to #connect (see ClaudeAgentSDK::Transport)' if transport && !transport.respond_to?(:connect)
Expand Down Expand Up @@ -525,7 +527,8 @@ def self.query(prompt:, options: nil, transport: nil, &block)
sdk_mcp_servers: sdk_mcp_servers,
exclude_dynamic_sections: ClaudeAgentSDK.extract_exclude_dynamic_sections(configured_options.system_prompt),
skills: configured_options.skills,
callback_scheduling: callback_scheduling
callback_scheduling: callback_scheduling,
callback_wrapper: callback_wrapper
)

# Mirror transcripts to the session_store, if configured. Installed
Expand All @@ -549,7 +552,8 @@ def self.query(prompt:, options: nil, transport: nil, &block)

# Send prompt(s) as user messages, then close stdin
if prompt.is_a?(String)
ClaudeAgentSDK.notify_observers(resolved_observers, :on_user_prompt, prompt, scheduling: callback_scheduling)
ClaudeAgentSDK.notify_observers(resolved_observers, :on_user_prompt, prompt,
scheduling: callback_scheduling, wrapper: callback_wrapper)
message = {
type: 'user',
message: { role: 'user', content: prompt },
Expand All @@ -567,7 +571,8 @@ def self.query(prompt:, options: nil, transport: nil, &block)
# here kept the root reactor alive forever when the read loop died
# while the user enumerator was still blocked (matches Python's
# query.spawn_task(query.stream_input(prompt))).
observed_prompt = ClaudeAgentSDK.observing_prompt_stream(prompt, resolved_observers, scheduling: callback_scheduling)
observed_prompt = ClaudeAgentSDK.observing_prompt_stream(prompt, resolved_observers,
scheduling: callback_scheduling, wrapper: callback_wrapper)
query_handler.spawn_task { query_handler.stream_input(observed_prompt) }
end

Expand All @@ -579,8 +584,10 @@ def self.query(prompt:, options: nil, transport: nil, &block)
message = MessageParser.parse(data)
next unless message

ClaudeAgentSDK.notify_observers(resolved_observers, :on_message, message, scheduling: callback_scheduling)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: callback_scheduling)
ClaudeAgentSDK.notify_observers(resolved_observers, :on_message, message,
scheduling: callback_scheduling, wrapper: callback_wrapper)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: callback_scheduling,
wrapper: callback_wrapper)
break signal.value if signal.is_a?(FiberBoundary::Break)
end
rescue StandardError => e
Expand All @@ -589,10 +596,12 @@ def self.query(prompt:, options: nil, transport: nil, &block)
# parse errors, and user-block errors. StandardError only: Async::Stop
# is cancellation, not an error. Bare raise preserves the backtrace;
# the ensure below still fires on_close after on_error.
ClaudeAgentSDK.notify_observers(resolved_observers, :on_error, e, scheduling: callback_scheduling)
ClaudeAgentSDK.notify_observers(resolved_observers, :on_error, e,
scheduling: callback_scheduling, wrapper: callback_wrapper)
raise
ensure
ClaudeAgentSDK.notify_observers(resolved_observers, :on_close, scheduling: callback_scheduling)
ClaudeAgentSDK.notify_observers(resolved_observers, :on_close,
scheduling: callback_scheduling, wrapper: callback_wrapper)
# query_handler.close stops the background read task and closes the
# transport (flushing the mirror batcher first). Fall back to a bare
# transport close when the handler was never built.
Expand Down Expand Up @@ -663,6 +672,7 @@ class Client
def initialize(options: nil, transport_class: SubprocessCLITransport, transport_args: {})
@options = options || ClaudeAgentOptions.new
@callback_scheduling = @options.callback_scheduling || :thread
@callback_wrapper = @options.callback_wrapper
@transport_class = transport_class
@transport_args = transport_args
@transport = nil
Expand Down Expand Up @@ -799,7 +809,8 @@ def query(prompt, session_id: 'default')

begin
if prompt.is_a?(String)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, prompt, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, prompt,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
message = {
type: 'user',
message: { role: 'user', content: prompt },
Expand Down Expand Up @@ -840,8 +851,10 @@ def receive_messages(&block)
message = MessageParser.parse(data)
next unless message

ClaudeAgentSDK.notify_observers(@resolved_observers, :on_message, message, scheduling: @callback_scheduling)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_message, message,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: @callback_scheduling,
wrapper: @callback_wrapper)
break signal.value if signal.is_a?(FiberBoundary::Break)
end
rescue StandardError => e
Expand All @@ -866,8 +879,10 @@ def receive_response(&block)
message = MessageParser.parse(data)
next unless message

ClaudeAgentSDK.notify_observers(@resolved_observers, :on_message, message, scheduling: @callback_scheduling)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_message, message,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
signal = FiberBoundary.invoke_iteration(block, message, scheduling: @callback_scheduling,
wrapper: @callback_wrapper)
break signal.value if signal.is_a?(FiberBoundary::Break)
break if message.is_a?(ResultMessage)
end
Expand Down Expand Up @@ -959,7 +974,10 @@ def get_server_info

# Disconnect from Claude
def disconnect
ClaudeAgentSDK.notify_observers(@resolved_observers || [], :on_close, scheduling: @callback_scheduling) if @connected
if @connected
ClaudeAgentSDK.notify_observers(@resolved_observers || [], :on_close,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
end
# Tear down whatever exists — robust to a partial/failed connect, where
# @connected is still false but a transport and/or materialized temp dir
# were already created. #close on the query handler also closes the
Expand Down Expand Up @@ -1046,7 +1064,8 @@ def connect_inner(configured_options, prompt)
agents: configured_options.agents,
exclude_dynamic_sections: exclude_dynamic_sections,
skills: configured_options.skills,
callback_scheduling: @callback_scheduling
callback_scheduling: @callback_scheduling,
callback_wrapper: @callback_wrapper
)

# Mirror transcripts to the session_store, if configured.
Expand Down Expand Up @@ -1077,7 +1096,8 @@ def connect_inner(configured_options, prompt)
# Observer#on_error contract; notifying a swallowed error would mark
# a still-live OTel trace as failed). Same behavior as query()'s
# streaming path.
observed = ClaudeAgentSDK.observing_prompt_stream(prompt, @resolved_observers, scheduling: @callback_scheduling)
observed = ClaudeAgentSDK.observing_prompt_stream(prompt, @resolved_observers,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
@query_handler.spawn_task { @query_handler.stream_input(observed) }
end
end
Expand All @@ -1094,12 +1114,14 @@ def stream_query_messages(prompt, session_id)
when Hash
msg = msg.merge(session_id: session_id) unless msg.key?(:session_id) || msg.key?('session_id')
if (text = ClaudeAgentSDK.extract_user_prompt_text(msg))
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, text, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, text,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
end
writeln(JSON.generate(msg))
when String
if (text = ClaudeAgentSDK.extract_user_prompt_text(msg))
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, text, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers, :on_user_prompt, text,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
end
writeln(msg)
else
Expand All @@ -1113,7 +1135,8 @@ def stream_query_messages(prompt, session_id)
# Notify observers of an error surfacing to the consumer. `|| []` keeps a
# mis-scoped call before connect harmless instead of NoMethodError on nil.
def notify_error(error)
ClaudeAgentSDK.notify_observers(@resolved_observers || [], :on_error, error, scheduling: @callback_scheduling)
ClaudeAgentSDK.notify_observers(@resolved_observers || [], :on_error, error,
scheduling: @callback_scheduling, wrapper: @callback_wrapper)
end

# Build and install the transcript-mirror batcher on the query handler when
Expand Down
Loading
Loading