claude-agent-sdk 0.34.0 → 0.36.0

data/lib/claude_agent_sdk/fiber_boundary.rb

@@ -133,6 +133,58 @@ module ClaudeAgentSDK
       end
     end
 
+    # Carries a SystemExit / SignalException (Interrupt included) raised by
+    # a user callback out of the FiberBoundary hop — see .invoke_callback.
+    # A StandardError so the hop ends normally: a :thread worker that died
+    # with SystemExit would have it re-raised by Ruby on the MAIN thread,
+    # asynchronously, before the SDK could answer the pending request.
+    # #cause (and #original) is the exception it carries. A callback_wrapper
+    # sees this carrier, never the original.
+    # @api private
+    class ProcessExitCarrier < StandardError
+      attr_reader :original
+
+      def initialize(original)
+        @original = original
+        super(FiberBoundary.process_exit_message(original))
+      end
+    end
+
+    # Hands a callback's process-exit exception from the code running the
+    # callback to the SDK code waiting for it — or, once that waiter has
+    # stopped waiting (hook timeout, cancelled request), lets the callback
+    # side re-raise it where it is, as plain Ruby would. Either way it is
+    # never dropped.
+    # @api private
+    class ProcessExitHandoff
+      def initialize
+        @mutex = Mutex.new
+        @state = :waiting
+        @original = nil
+      end
+
+      # Callback side: true when the waiter will receive +original+.
+      def hand_off(original)
+        @mutex.synchronize do
+          next false unless @state == :waiting
+
+          @state = :handed_off
+          @original = original
+          true
+        end
+      end
+
+      # Waiter side, when it stops waiting: the exception handed off but not
+      # yet received, if any.
+      def abandon
+        @mutex.synchronize do
+          handed_off = @state == :handed_off
+          @state = :abandoned
+          handed_off ? @original : nil
+        end
+      end
+    end
+
     # Sentinel returned by .invoke_iteration when the user block attempted `break`.
     class Break
       attr_reader :value
@@ -144,6 +196,63 @@ module ClaudeAgentSDK
 
     module_function
 
+    # Invoke a user callback that answers a CLI control request (hook,
+    # can_use_tool, SDK MCP tool / resource / prompt handler) across the
+    # boundary, like .invoke. A SystemExit or SignalException (Interrupt
+    # included) raised while the callback runs is never swallowed: it is
+    # re-raised here, on the calling fiber, as the ORIGINAL exception, so
+    # Query#handle_control_request can answer the request first and then let
+    # it terminate the process as Ruby normally would.
+    #
+    # The conversion must sit INSIDE the hop, innermost around the user
+    # call: in :thread mode a worker dying with SystemExit has it re-raised
+    # by Ruby on the main thread at an arbitrary point, before any response
+    # is written. So the worker ends with a ProcessExitCarrier instead, and
+    # the carrier is unwrapped once control is back on the calling fiber. In
+    # :inline mode the callback runs on the reactor fiber — usually the main
+    # thread — so this also covers a real Ctrl-C / SIGTERM delivered while
+    # the callback runs. A callback_wrapper sees the carrier (a
+    # StandardError, #cause = the original); ensure-based wrappers still
+    # run their cleanup, and a wrapper that swallows the carrier cannot
+    # swallow the exit (the handoff re-raises it).
+    #
+    # If the caller stops waiting first (hook timeout, cancelled request),
+    # the exception is re-raised in the abandoned worker thread instead — a
+    # SystemExit from a non-main thread then ends the process, as in plain
+    # Ruby. Cancellation (Async::Stop, InlineCancellation) is not a
+    # SignalException and passes through untouched.
+    # @api private
+    def invoke_callback(scheduling: :thread, wrapper: nil, &callback)
+      handoff = ProcessExitHandoff.new
+      received = false
+      begin
+        invoke(scheduling: scheduling, wrapper: wrapper) do
+          callback.call
+        rescue SystemExit, SignalException => e
+          raise unless handoff.hand_off(e)
+
+          raise ProcessExitCarrier, e
+        end
+      rescue ProcessExitCarrier => e
+        received = true
+        raise e.original
+      ensure
+        unless received
+          pending = handoff.abandon
+          raise pending if pending
+        end
+      end
+    end
+
+    # Text reporting a process-exit exception to the CLI: its class, plus
+    # its message when that adds anything ("SystemExit: exit",
+    # "SignalException: SIGTERM", "Interrupt").
+    # @api private
+    def process_exit_message(error)
+      detail = error.message
+      detail.empty? || detail == error.class.name ? error.class.name : "#{error.class}: #{detail}"
+    end
+
     # Capture only the optional OTel context before crossing a fiber/thread
     # boundary. OTel keeps its current context fiber-local; copying generic
     # thread locals would also copy unsafe connection/request state. The
@@ -267,10 +376,10 @@ module ClaudeAgentSDK
     # need a bounded wait use :thread scheduling (hard Thread#join bound);
     # inline callbacks keep their cleanup fiber-aware.
     # @api private
-    def with_cooperative_timeout(task, timeout, on_timeout:, &block)
+    def with_cooperative_timeout(task, timeout, on_timeout:, &)
       cancellation = Class.new(InlineCancellation)
       begin
-        task.with_timeout(timeout, cancellation, &block)
+        task.with_timeout(timeout, cancellation, &)
       rescue cancellation
         raise on_timeout.call
       end

data/lib/claude_agent_sdk/option_warnings.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'set'
-
 module ClaudeAgentSDK
   # Advisory warnings for option combinations that silently change behavior.
   module OptionWarnings

data/lib/claude_agent_sdk/query.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require 'json'
-require 'set'
 require 'async'
 require 'async/queue'
 require 'async/condition'
@@ -271,11 +270,11 @@ module ClaudeAgentSDK
     # Fine for the current one-shot call sites (max two tasks per Query); do
     # not route per-request work (control handlers, per-turn streams) through
     # this without adding completion-based removal.
-    def spawn_task(&block)
+    def spawn_task(&)
       parent = Async::Task.current?
       raise CLIConnectionError, 'Query#spawn_task must be called inside an Async{} block' unless parent
 
-      task = parent.async(&FiberBoundary.capture_otel_context(&block))
+      task = parent.async(&FiberBoundary.capture_otel_context(&))
       @child_tasks << task
       task
     end
@@ -575,13 +574,50 @@ module ClaudeAgentSDK
         }
       }
       writeln(JSON.generate(success_response))
+      responded = true
     rescue Async::Stop
       # Cancellation requested; respond with an error so the CLI can unblock.
       send_control_error(request_id, 'Cancelled')
+    rescue SystemExit, SignalException => e
+      # exit / Interrupt / a signal raised while a user callback ran
+      # (FiberBoundary.invoke_callback re-raises it here, on the reactor) —
+      # or a real signal landing on this fiber. Never swallowed: answer the
+      # request the way an ordinary callback failure is answered, so the CLI
+      # is not left waiting, then let it terminate the process as Ruby
+      # normally would. The transport flushes every write.
+      respond_to_process_exit(request_id, request_data, e) unless responded
+      raise
     rescue StandardError => e
       send_control_error(request_id, e.message)
     end
 
+    # The response an ordinary exception from the callback would have
+    # produced, with the process-exit exception named by class: an error
+    # control response for hooks / can_use_tool; for SDK MCP requests an
+    # in-band isError result (tools/call) or a JSON-RPC internal error
+    # (resources/read, prompts/get), inside a successful control response.
+    def respond_to_process_exit(request_id, request_data, error)
+      message = FiberBoundary.process_exit_message(error)
+      mcp_message = request_data[:message] if request_data.is_a?(Hash) && request_data[:subtype] == 'mcp_message'
+      return send_control_error(request_id, message) unless mcp_message.is_a?(Hash)
+
+      mcp_response = { jsonrpc: '2.0', id: mcp_message[:id] }
+      if mcp_message[:method] == 'tools/call'
+        mcp_response[:result] = { content: [{ type: 'text', text: message }], isError: true }
+      else
+        mcp_response[:error] = { code: -32_603, message: message }
+      end
+      writeln(JSON.generate({
+                              type: 'control_response',
+                              response: {
+                                subtype: 'success', request_id: request_id, requestId: request_id,
+                                response: { mcp_response: mcp_response }
+                              }
+                            }))
+    rescue CLIConnectionError
+      nil # the CLI is already gone; nothing is waiting for the answer
+    end
+
     def send_control_error(request_id, message)
       error_response = {
         type: 'control_response',
@@ -628,8 +664,10 @@ module ClaudeAgentSDK
       # so AR/PG calls inside it aren't intercepted by the Fiber scheduler;
       # with callback_scheduling: :inline it runs in place on this control-
       # request task, where control_cancel_request (task.stop) can actually
-      # cancel it at suspension points.
-      response = FiberBoundary.invoke(scheduling: @callback_scheduling, wrapper: @callback_wrapper) do
+      # cancel it at suspension points. exit / Interrupt from the callback
+      # re-raise here after the hop; handle_control_request answers the
+      # request before letting them propagate (FiberBoundary.invoke_callback).
+      response = FiberBoundary.invoke_callback(scheduling: @callback_scheduling, wrapper: @callback_wrapper) do
         @can_use_tool.call(request_data[:tool_name], request_data[:input], context)
       end
       # A worker may return a decision after the read loop invalidated the
@@ -684,8 +722,11 @@ module ClaudeAgentSDK
       # genuine cooperative cancellation: the hook is interrupted at its next
       # suspension point and its ensure blocks run (Python parity — anyio
       # cancels the coroutine). A CPU-stuck inline hook cannot be timed out.
+      # All three variants go through FiberBoundary.invoke_callback, so exit
+      # / Interrupt from the hook reach handle_control_request, which answers
+      # the request before letting them propagate.
       unless @hook_callback_timeouts[callback_id]
-        hook_output = FiberBoundary.invoke(scheduling: @callback_scheduling, wrapper: @callback_wrapper) do
+        hook_output = FiberBoundary.invoke_callback(scheduling: @callback_scheduling, wrapper: @callback_wrapper) do
           callback.call(hook_input, request_data[:tool_use_id], context)
         end
       end
@@ -709,13 +750,13 @@ module ClaudeAgentSDK
               Async::Task.current, timeout,
               on_timeout: -> { Async::TimeoutError.new('execution expired') }
             ) do
-              FiberBoundary.invoke(scheduling: :inline, wrapper: @callback_wrapper) do
+              FiberBoundary.invoke_callback(scheduling: :inline, wrapper: @callback_wrapper) do
                 callback.call(hook_input, request_data[:tool_use_id], context)
               end
             end
           else
             Async::Task.current.with_timeout(timeout) do
-              FiberBoundary.invoke(wrapper: @callback_wrapper) do
+              FiberBoundary.invoke_callback(wrapper: @callback_wrapper) do
                 callback.call(hook_input, request_data[:tool_use_id], context)
               end
             end

data/lib/claude_agent_sdk/railtie.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+module ClaudeAgentSDK
+  # Rails integration. Loaded by lib/claude_agent_sdk.rb only when
+  # Rails::Railtie is already defined (Bundler.require runs after
+  # `require 'rails'`), so non-Rails processes never see it.
+  #
+  # Deliberately minimal: it contributes the `claude_agent_sdk:*` rake tasks
+  # and anchors CLI discovery to the app root, nothing else. It installs
+  # nothing into callback dispatch — the generated initializer
+  # (`bin/rails g claude_agent_sdk:install`) opts in to {.callback_wrapper}
+  # explicitly, where it is visible and removable.
+  class Railtie < ::Rails::Railtie
+    rake_tasks do
+      load File.expand_path('tasks/claude_agent_sdk.rake', __dir__)
+    end
+
+    # Find the vendored CLI under Rails.root/vendor/claude whatever the
+    # process cwd — a daemonized worker or a job runner started elsewhere
+    # would otherwise look under its own cwd and fall through to PATH.
+    # Runs before config/initializers, so an app initializer can still set
+    # CLIInstaller.root (or nil, for the cwd) itself; a root set earlier,
+    # e.g. in config/application.rb, is left alone.
+    initializer 'claude_agent_sdk.cli_installer_root', before: :load_config_initializers do |app|
+      ClaudeAgentSDK::CLIInstaller.root ||= app.root
+    end
+
+    # A `callback_wrapper` (see ClaudeAgentOptions#callback_wrapper) that
+    # gives SDK callbacks Rails' connection hygiene without deadlocking
+    # development code reloading.
+    #
+    # The obvious wrapper, `->(inv) { Rails.application.executor.wrap { inv.call } }`,
+    # deadlocks whenever the executor carries a process-wide lock:
+    #
+    # - With code reloading enabled, every executor holds a share of the
+    #   ActiveSupport::Dependencies interlock. In the default `:thread`
+    #   scheduling the request/job thread (already inside the executor, so
+    #   already holding a share) blocks on the FiberBoundary thread running
+    #   the callback. If a reload starts meanwhile (another request after the
+    #   agent edited an app file), the reloader queues for the exclusive
+    #   unload lock, and the callback thread's own `executor.wrap` then waits
+    #   behind it for a new share — which the reloader can never get while
+    #   the parent's share is held. Three-way deadlock.
+    # - With `config.allow_concurrency = false`, the executor holds a
+    #   process-wide monitor that the parent thread already owns, so the
+    #   callback thread's `executor.wrap` blocks every time.
+    #
+    # So, per invocation:
+    #
+    # 1. Executor already active on this execution context (`:inline`
+    #    scheduling, or a callback running on the caller's own thread) —
+    #    call straight through; the enclosing executor already owns cleanup.
+    # 2. Executor carries a lock (the two cases above, mirroring railties'
+    #    `configure_executor_for_concurrency`) — call WITHOUT entering the
+    #    executor, then return this thread's ActiveRecord connections to the
+    #    pool. The caller's executor still covers the callback: its share of
+    #    the interlock keeps a reload from unloading code under it until the
+    #    whole SDK call returns.
+    # 3. Otherwise (production: no reloading, concurrency allowed) — run the
+    #    callback inside `Rails.application.executor.wrap`.
+    #
+    # The configuration is read on every call, so the same wrapper is correct
+    # in every environment.
+    #
+    # @return [Proc] a callable suitable for `callback_wrapper:`
+    # @example config/initializers/claude_agent_sdk.rb
+    #   ClaudeAgentSDK.configure do |config|
+    #     config.default_options = { callback_wrapper: ClaudeAgentSDK::Railtie.callback_wrapper }
+    #   end
+    def self.callback_wrapper
+      lambda do |invocation|
+        app = ::Rails.application
+        next invocation.call if app.nil? || app.executor.active?
+        next app.executor.wrap { invocation.call } unless executor_locks?(app.config)
+
+        begin
+          invocation.call
+        ensure
+          release_active_record_connections
+        end
+      end
+    end
+
+    # Whether railties registered a process-wide lock hook on the executor
+    # (Rails::Application::Finisher, initializer
+    # :configure_executor_for_concurrency).
+    def self.executor_locks?(config)
+      return true if config.allow_concurrency == false
+      return false if config.allow_concurrency == :unsafe
+
+      config.reloading_enabled?
+    end
+    private_class_method :executor_locks?
+
+    # What the executor's ActiveRecord completion hook does for a callback
+    # that ran on a thread of its own. Explicit :all — the no-argument form
+    # is deprecated on Rails 7.1.
+    def self.release_active_record_connections
+      return unless defined?(::ActiveRecord::Base)
+
+      ::ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
+    end
+    private_class_method :release_active_record_connections
+  end
+end

data/lib/claude_agent_sdk/sdk_mcp_server.rb

@@ -4,6 +4,7 @@ require 'mcp'
 
 module ClaudeAgentSDK
   # Recursively convert all hash keys to symbols
+  # @api private
   def self.deep_symbolize_keys(obj)
     case obj
     when Hash then obj.transform_keys(&:to_sym).transform_values { |v| deep_symbolize_keys(v) }
@@ -15,6 +16,7 @@ module ClaudeAgentSDK
   # Like deep_symbolize_keys, but also converts Symbol VALUES to strings so a
   # prebuilt schema written with symbols ({ type: :object, ... }) emits clean
   # wire-format JSON Schema.
+  # @api private
   def self.deep_normalize_schema(obj)
     case obj
     when Hash then obj.transform_keys(&:to_sym).transform_values { |v| deep_normalize_schema(v) }
@@ -34,6 +36,7 @@ module ClaudeAgentSDK
   # mangled into nonsense parameter lists ("additionalProperties" as a
   # required string param). A $ref-only schema without type: 'object' remains
   # indistinguishable from a params hash — declare the type alongside $ref.
+  # @api private
   def self.prebuilt_json_schema?(schema)
     return false unless schema.is_a?(Hash)
 
@@ -47,6 +50,7 @@ module ClaudeAgentSDK
   # Single source of truth for tool input schemas: prebuilt schemas are
   # normalized (symbol keys, string values); simple { name: :type } hashes
   # become a full JSON Schema with every param required (string keys).
+  # @api private
   def self.normalize_tool_schema(schema)
     return deep_normalize_schema(schema) if prebuilt_json_schema?(schema)
 
@@ -60,6 +64,7 @@ module ClaudeAgentSDK
     { type: 'object', properties: {} }
   end
 
+  # @api private
   def self.ruby_type_to_json_schema(type)
     # Class#=== matches instances, not the class object used in { id: Integer }.
     type = { String => :string, Integer => :integer, Float => :float, TrueClass => :boolean, FalseClass => :boolean }.fetch(type, type)
@@ -72,22 +77,15 @@ module ClaudeAgentSDK
     end
   end
 
-  # Internal: call a tool handler, reporting SystemExit / SignalException
-  # (Interrupt included) as an ordinary handler failure — re-raised as a
-  # RuntimeError (#cause holds the original) that both tools/call dispatch
-  # boundaries turn into an in-band isError result, so the pending control
-  # response is always written. Must run INSIDE the FiberBoundary.invoke
-  # block: a worker thread that dies with SystemExit has it re-raised by
-  # Ruby on the MAIN thread, tearing down the reactor, while the dispatcher
-  # only sees Async::Stop — a rescue after the hop cannot catch it in
-  # :thread mode. A callback_wrapper therefore observes the RuntimeError.
-  # Deliberately not `rescue Exception`: cancellation (Async::Stop, and
-  # InlineCancellation at an :inline suspension point) must propagate.
+  # Internal: expand a tool handler's String shorthand into a single text
+  # block. Every other value passes through untouched — Hash results behave
+  # exactly as before, and any other non-Hash value still gets the "must
+  # return a hash" diagnostic from the caller. Applied inside the callback
+  # dispatch at both tools/call paths, so a callback_wrapper sees the
+  # expanded Hash.
   # @api private
-  def self.call_tool_handler(handler, arguments)
-    handler.call(arguments)
-  rescue SystemExit, SignalException => e
-    raise e.message
+  def self.normalize_tool_result(result)
+    result.is_a?(String) ? { content: [{ type: 'text', text: result }] } : result
   end
 
   # SDK MCP Server - wraps official MCP::Server with block-based API
@@ -289,8 +287,10 @@ module ClaudeAgentSDK
       # gem's Fiber scheduler is not visible to user code (which may hit
       # AR/PG); in :inline mode it runs in place on the reactor fiber.
       scheduling, wrapper = effective_callback_dispatch
-      result = FiberBoundary.invoke(scheduling: scheduling, wrapper: wrapper) do
-        ClaudeAgentSDK.call_tool_handler(tool.handler, arguments)
+      # exit / Interrupt from the handler propagate (never an isError
+      # result): see FiberBoundary.invoke_callback.
+      result = FiberBoundary.invoke_callback(scheduling: scheduling, wrapper: wrapper) do
+        ClaudeAgentSDK.normalize_tool_result(tool.handler.call(arguments))
       end
 
       # Guard before flexible_fetch: it raises on non-Hash inputs.
@@ -327,7 +327,7 @@ module ClaudeAgentSDK
       # as `call_tool` above: reader blocks may touch Thread.current-keyed
       # libraries (ActiveRecord, pg, ...) and must run on a plain thread.
       scheduling, wrapper = effective_callback_dispatch
-      content = FiberBoundary.invoke(scheduling: scheduling, wrapper: wrapper) do
+      content = FiberBoundary.invoke_callback(scheduling: scheduling, wrapper: wrapper) do
         resource.reader.call
       end
 
@@ -362,7 +362,7 @@ module ClaudeAgentSDK
       # Hop off the Fiber scheduler before invoking user code — same reason
       # as `call_tool` above.
       scheduling, wrapper = effective_callback_dispatch
-      result = FiberBoundary.invoke(scheduling: scheduling, wrapper: wrapper) do
+      result = FiberBoundary.invoke_callback(scheduling: scheduling, wrapper: wrapper) do
         prompt.generator.call(arguments)
       end
 
@@ -475,8 +475,11 @@ module ClaudeAgentSDK
               # Hop to a plain thread (default) so user handlers don't see
               # the Fiber scheduler; :inline runs in place on the reactor.
               scheduling, wrapper = @sdk_server.effective_callback_dispatch
-              result = FiberBoundary.invoke(scheduling: scheduling, wrapper: wrapper) do
-                ClaudeAgentSDK.call_tool_handler(@tool_def.handler, args)
+              # exit / Interrupt propagate past the gem (it rescues only
+              # StandardError) to Query#handle_control_request, which
+              # answers with an isError result and then re-raises them.
+              result = FiberBoundary.invoke_callback(scheduling: scheduling, wrapper: wrapper) do
+                ClaudeAgentSDK.normalize_tool_result(@tool_def.handler.call(args))
               end
 
               # Guard BEFORE flexible_fetch: on a non-Hash it raises
@@ -595,18 +598,26 @@ module ClaudeAgentSDK
   # @param name [String] Unique identifier for the tool
   # @param description [String] Human-readable description
   # @param input_schema [Hash] Schema defining input parameters
-  # @param handler [Proc] Block that implements the tool logic
+  # @param handler [Proc] Block that implements the tool logic. It returns a
+  #   String, sent to Claude as a single text block, or a Hash with a
+  #   +:content+ Array of MCP content blocks plus optional +:is_error+ /
+  #   +:structured_content+. Use the Hash form for error results, structured
+  #   output, images, or several blocks.
   # @return [SdkMcpTool] Tool definition
   #
-  # @example Simple tool
+  # @example Simple tool (a String return becomes one text block)
+  #   tool = create_tool('greet', 'Greet a user', { name: :string }) do |args|
+  #     "Hello, #{args[:name]}!"
+  #   end
+  #
+  # @example The same tool in the Hash form
   #   tool = create_tool('greet', 'Greet a user', { name: :string }) do |args|
   #     { content: [{ type: 'text', text: "Hello, #{args[:name]}!" }] }
   #   end
   #
   # @example Tool with multiple parameters
   #   tool = create_tool('add', 'Add two numbers', { a: :number, b: :number }) do |args|
-  #     result = args[:a] + args[:b]
-  #     { content: [{ type: 'text', text: "Result: #{result}" }] }
+  #     "Result: #{args[:a] + args[:b]}"
   #   end
   #
   # @example Tool with error handling

data/lib/claude_agent_sdk/session_mutations.rb

@@ -34,7 +34,7 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id is invalid or title is empty
     # @raise [Errno::ENOENT] if the session file cannot be found
     def rename_session(session_id:, title:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
       stripped = title.strip
       raise ArgumentError, 'title must be non-empty' if stripped.empty?
@@ -55,7 +55,7 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id is invalid or tag is empty after sanitization
     # @raise [Errno::ENOENT] if the session file cannot be found
     def tag_session(session_id:, tag:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
       if tag
         sanitized = sanitize_unicode(tag).strip
@@ -79,7 +79,7 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id is invalid
     # @raise [Errno::ENOENT] if the session file cannot be found
     def delete_session(session_id:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
       result = find_session_file_with_dir(session_id, directory)
       raise Errno::ENOENT, "Session #{session_id} not found#{" in project directory for #{directory}" if directory}" unless result
@@ -115,9 +115,9 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id or up_to_message_id is invalid
     # @raise [Errno::ENOENT] if the session file cannot be found
     def fork_session(session_id:, directory: nil, up_to_message_id: nil, title: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
-      raise ArgumentError, "Invalid up_to_message_id: #{up_to_message_id}" if up_to_message_id && !up_to_message_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid up_to_message_id: #{up_to_message_id}" if up_to_message_id && !Sessions.valid_session_id?(up_to_message_id)
 
       result = find_session_file_with_dir(session_id, directory)
       raise Errno::ENOENT, "Session #{session_id} not found#{" in project directory for #{directory}" if directory}" unless result
@@ -159,7 +159,7 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id is invalid or title is empty
     # @raise [Errno::ENOENT] if the session is not found in the store
     def rename_session_via_store(session_store:, session_id:, title:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
       stripped = title.strip
       raise ArgumentError, 'title must be non-empty' if stripped.empty?
@@ -183,7 +183,7 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id is invalid or tag is empty after sanitization
     # @raise [Errno::ENOENT] if the session is not found in the store
     def tag_session_via_store(session_store:, session_id:, tag:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
 
       if tag
         sanitized = sanitize_unicode(tag).strip
@@ -213,7 +213,7 @@ module ClaudeAgentSDK
     #
     # @raise [ArgumentError] if session_id is invalid
     def delete_session_via_store(session_store:, session_id:, directory: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
       return unless SessionStore.implements?(session_store, :delete)
 
       key = { 'project_key' => Sessions.project_key_for_directory(directory), 'session_id' => session_id }
@@ -231,8 +231,8 @@ module ClaudeAgentSDK
     # @raise [ArgumentError] if session_id/up_to_message_id is invalid or the session has no messages
     # @raise [Errno::ENOENT] if the source session is not found in the store
     def fork_session_via_store(session_store:, session_id:, directory: nil, up_to_message_id: nil, title: nil)
-      raise ArgumentError, "Invalid session_id: #{session_id}" unless session_id.match?(Sessions::UUID_RE)
-      raise ArgumentError, "Invalid up_to_message_id: #{up_to_message_id}" if up_to_message_id && !up_to_message_id.match?(Sessions::UUID_RE)
+      raise ArgumentError, "Invalid session_id: #{session_id}" unless Sessions.valid_session_id?(session_id)
+      raise ArgumentError, "Invalid up_to_message_id: #{up_to_message_id}" if up_to_message_id && !Sessions.valid_session_id?(up_to_message_id)
 
       project_key = Sessions.project_key_for_directory(directory)
       raw = session_store.load('project_key' => project_key, 'session_id' => session_id)