mistri 0.5.0 → 0.6.1

data/lib/mistri/sub_agent/runtime.rb

@@ -0,0 +1,297 @@
+# frozen_string_literal: true
+
+module Mistri
+  # Agent options a child may inherit without replacing lifecycle-owned
+  # provider, session, prompt, tools, task, signal, or event state.
+  module ChildAgentOptions
+    ALLOWED = %i[
+      budget max_concurrency transform_context compaction retries skills
+      before_tool after_tool context
+    ].freeze
+    private_constant :ALLOWED
+
+    module_function
+
+    def validate!(options)
+      unsupported = options.keys - ALLOWED
+      return if unsupported.empty?
+
+      raise ConfigurationError,
+            "unsupported sub-agent options: #{unsupported.sort.join(", ")}"
+    end
+  end
+  private_constant :ChildAgentOptions
+
+  class SubAgent
+    DISPATCH_SPEC_VERSION = 1
+    DISPATCH_SPEC_KEYS = %w[
+      spec_version name session_id parent_session_id type instructions task tool_names model
+    ].freeze
+    UNSET_RUNTIME_FIELD = Object.new.freeze
+    private_constant :DISPATCH_SPEC_KEYS, :UNSET_RUNTIME_FIELD
+
+    # The live dependencies a host constructs for one dispatched child.
+    # Mistri verifies the provider and tools against the durable spec; the
+    # host owns tenant scope, backend isolation, and the freshness of every
+    # object placed here.
+    class Runtime
+      attr_reader :provider, :system, :tools, :schema, :agent_options
+
+      def initialize(provider:, system: nil, tools: [], schema: nil, cleanup: nil,
+                     **agent_options)
+        raise ArgumentError, "runtime tools must be an Array" unless tools.is_a?(Array)
+        unless cleanup.nil? || cleanup.respond_to?(:call)
+          raise ArgumentError, "runtime cleanup must be callable"
+        end
+
+        @provider = provider
+        @system = system
+        @tools = Array.new(tools).freeze
+        @schema = schema
+        @cleanup = cleanup
+        @agent_options = agent_options.dup.freeze
+        freeze
+      end
+
+      def close = @cleanup&.call
+    end
+
+    # The fail-closed boundary between a durable dispatch spec and live Ruby
+    # dependencies. Kept separate from SubAgent's execution lifecycle so each
+    # concern stays auditable.
+    class RuntimeContract
+      Resolved = Data.define(:provider, :system, :tools, :schema, :agent_options)
+      private_constant :Resolved
+
+      class << self
+        def own_spec(spec)
+          owned, error = ToolArguments.canonicalize(spec)
+          unless error.nil? && owned.is_a?(Hash)
+            raise ConfigurationError, "dispatched child spec must be a bounded JSON object"
+          end
+
+          owned
+        end
+
+        def validate_identity!(spec)
+          %w[name session_id].each do |field|
+            value = spec[field]
+            unless value.is_a?(String) && !value.empty?
+              raise ConfigurationError, "dispatched child spec needs a non-empty #{field}"
+            end
+          end
+        end
+
+        def bind_spec(entries, supplied)
+          dispatched = entries.find { |entry| entry["type"] == Child::DISPATCHED }
+          raise DispatchGrantError, "child session has no durable dispatch grant" unless dispatched
+
+          stored = dispatched["spec"]
+          return supplied if stored.nil? && supplied["spec_version"].nil?
+          unless stored
+            raise DispatchGrantError, "versioned child is missing its durable dispatch grant"
+          end
+
+          authoritative = own_spec(stored)
+          unless supplied == authoritative
+            raise DispatchGrantError, "queue payload does not match the durable dispatch grant"
+          end
+
+          authoritative
+        end
+
+        def validate_resolution!(factory:, direct:)
+          supplied = direct.values.any? { |value| !value.equal?(UNSET_RUNTIME_FIELD) }
+          if factory
+            raise ArgumentError, "choose runtime_factory or direct runtime fields" if supplied
+            unless factory.respond_to?(:call)
+              raise ConfigurationError, "runtime_factory must be callable"
+            end
+          elsif direct.fetch(:provider).equal?(UNSET_RUNTIME_FIELD)
+            raise ConfigurationError, "run_dispatched needs runtime_factory or provider"
+          end
+        end
+
+        def validate_spec!(spec)
+          version = spec["spec_version"]
+          validate_version!(version)
+          validate_v1_shape!(spec) if version == DISPATCH_SPEC_VERSION
+          validate_task!(spec["task"])
+          validate_model!(spec["model"], version: version)
+          validate_tool_names!(spec["tool_names"])
+        end
+
+        def resolve(spec, factory:, direct:, &factory_failed)
+          return resolve_factory(spec, factory, &factory_failed) if factory
+
+          resolve_direct(direct)
+        end
+
+        def validate_runtime!(runtime, spec)
+          provider = runtime.provider
+          system = runtime.system
+          tools = runtime.tools
+          schema = runtime.schema
+          options = runtime.agent_options
+          validate_provider!(provider, spec["model"])
+          validate_agent_options!(options)
+          ordered = validate_tools!(tools, spec.fetch("tool_names"))
+          Resolved.new(provider: provider, system: system, tools: ordered,
+                       schema: schema, agent_options: options)
+        end
+
+        def cleanup(runtime)
+          runtime&.close
+          nil
+        rescue StandardError => e
+          e
+        end
+
+        private
+
+        def validate_version!(version)
+          return if version.nil? || version == DISPATCH_SPEC_VERSION
+
+          raise ConfigurationError,
+                "unsupported dispatched child spec version #{version.inspect}"
+        end
+
+        def validate_task!(task)
+          return if task.is_a?(String) && !task.empty?
+
+          raise ConfigurationError, "dispatched child spec needs a non-empty task"
+        end
+
+        def validate_v1_shape!(spec)
+          extra = spec.keys - DISPATCH_SPEC_KEYS
+          unless extra.empty?
+            raise ConfigurationError,
+                  "dispatched child spec has unknown fields: #{extra.join(", ")}"
+          end
+          validate_optional_string!(spec["parent_session_id"], "parent_session_id")
+          validate_required_string!(spec["type"], "type")
+          validate_optional_string!(spec["instructions"], "instructions")
+        end
+
+        def validate_model!(model, version:)
+          return if version.nil? && model.nil?
+          return if model.is_a?(String) && !model.empty?
+
+          raise ConfigurationError, "dispatched child spec model must be a non-empty string"
+        end
+
+        def validate_required_string!(value, field)
+          return if value.is_a?(String) && !value.empty?
+
+          raise ConfigurationError,
+                "dispatched child spec #{field} must be a non-empty string"
+        end
+
+        def validate_optional_string!(value, field)
+          return if value.nil? || value.is_a?(String)
+
+          raise ConfigurationError,
+                "dispatched child spec #{field} must be a string or null"
+        end
+
+        def validate_tool_names!(names)
+          valid = names.is_a?(Array) &&
+                  names.all? { |name| name.is_a?(String) && !name.empty? }
+          unless valid
+            raise ConfigurationError,
+                  "dispatched child spec tool_names must be non-empty strings"
+          end
+          return if names.uniq.length == names.length
+
+          raise ConfigurationError, "dispatched child spec has duplicate tool names"
+        end
+
+        def resolve_factory(spec, factory)
+          runtime = begin
+            factory.call(spec)
+          rescue DispatchGrantError
+            yield if block_given?
+            raise ConfigurationError,
+                  "runtime_factory must not raise Mistri::DispatchGrantError; " \
+                  "that class is reserved for queue grant verification"
+          rescue StandardError
+            yield if block_given?
+            raise
+          end
+          return runtime if runtime.instance_of?(Runtime)
+
+          raise ConfigurationError,
+                "runtime_factory must return Mistri::SubAgent::Runtime"
+        end
+
+        def resolve_direct(direct)
+          provider = direct.fetch(:provider)
+          Runtime.new(provider: provider,
+                      system: value_or_nil(direct.fetch(:system)),
+                      tools: value_or(direct.fetch(:tools), []),
+                      schema: value_or_nil(direct.fetch(:schema)),
+                      **value_or(direct.fetch(:agent_options), {}))
+        end
+
+        def value_or(value, fallback)
+          value.equal?(UNSET_RUNTIME_FIELD) ? fallback : value
+        end
+
+        def value_or_nil(value) = value_or(value, nil)
+
+        def validate_provider!(provider, expected_model)
+          unless provider.respond_to?(:stream)
+            raise ConfigurationError, "background runtime provider must respond to stream"
+          end
+          return unless expected_model
+
+          actual = provider.model.to_s if provider.respond_to?(:model)
+          return if actual == expected_model
+
+          raise ConfigurationError,
+                "background runtime model #{actual.inspect} does not match " \
+                "the granted model #{expected_model.inspect}"
+        end
+
+        def validate_agent_options!(options)
+          ChildAgentOptions.validate!(options)
+          return unless options.key?(:skills)
+
+          raise ConfigurationError,
+                "background runtime skills would add tools outside the durable grant"
+        end
+
+        def validate_tools!(tools, expected)
+          unless tools.all?(Tool)
+            raise ConfigurationError, "background runtime tools must be Mistri::Tool instances"
+          end
+
+          names = tools.map(&:name)
+          if names.uniq.length != names.length
+            raise ConfigurationError, "background runtime has duplicate tool names"
+          end
+
+          check_exact_grant!(expected, names)
+          by_name = tools.to_h { |tool| [tool.name, tool] }
+          ordered = expected.map { |name| by_name.fetch(name) }
+          SubAgent.forbid_gated!(ordered)
+          ordered
+        end
+
+        def check_exact_grant!(expected, names)
+          missing = expected - names
+          extra = names - expected
+          return if missing.empty? && extra.empty?
+
+          details = []
+          details << "missing: #{missing.join(", ")}" if missing.any?
+          details << "extra: #{extra.join(", ")}" if extra.any?
+          raise ConfigurationError,
+                "background runtime tools do not match the durable grant " \
+                "(#{details.join("; ")})"
+        end
+      end
+    end
+    private_constant :RuntimeContract
+  end
+end

data/lib/mistri/sub_agent.rb

@@ -23,9 +23,10 @@ module Mistri
   #
   #   spawn = Mistri::SubAgent.spawner(provider:, tools: [fetch_page, search])
   #
-  # Children never receive a spawn tool: delegation is one level deep by
-  # construction. Parallel fan-out costs nothing: several spawn calls in
-  # one turn run concurrently on the executor pool.
+  # The open spawner never grants itself, preventing accidental recursion.
+  # Hosts may deliberately nest fixed specialists. Several spawn calls in
+  # one turn are scheduled concurrently; provider instances decide whether
+  # their network requests overlap.
   #
   # Child events forward into the parent's stream tagged with origin
   # ("researcher#ab12cd34"; nesting of named specialists joins with ">").
@@ -35,6 +36,8 @@ module Mistri
   # runtime is denied and reported in band. Gate the delegation itself
   # instead (needs_approval: on the definition or spawner).
   class SubAgent
+    require_relative "sub_agent/runtime"
+
     SPAWNER_DESCRIPTION =
       "Delegate a self-contained task to a focused child agent with a clean " \
       "context. The child starts blank: give it complete instructions and " \
@@ -50,6 +53,8 @@ module Mistri
     def initialize(name:, description:, provider:, system: nil, tools: [], schema: nil,
                    **agent_options)
       SubAgent.forbid_gated!(tools)
+      @gate = agent_options.delete(:needs_approval) || false
+      ChildAgentOptions.validate!(agent_options)
       @name = name.to_s
       @description = description
       @provider = provider
@@ -57,7 +62,6 @@ module Mistri
       @tools = tools
       @schema = schema
       @agent_options = agent_options
-      @gate = agent_options.delete(:needs_approval) || false
     end
 
     # The delegate tool: each call runs a fresh child and answers with its
@@ -83,8 +87,8 @@ module Mistri
     def run_child(task, context, name: nil)
       SubAgent.run_child(label: SubAgent.sanitize_label(name, fallback: @name),
                          provider: @provider, system: @system,
-                         tools: @tools, task: task, context: context, schema: @schema,
-                         **@agent_options)
+                         tools: @tools, task: task, parent_context: context, schema: @schema,
+                         agent_options: @agent_options)
     end
 
     class << self
@@ -110,79 +114,141 @@ module Mistri
         label.empty? ? fallback : label
       end
 
-      def run_child(label:, provider:, system:, tools:, task:, context:, schema: nil,
-                    **agent_options)
-        store = context.session ? context.session.store : Stores::Memory.new
+      def run_child(label:, provider:, system:, tools:, task:, parent_context:, schema: nil,
+                    agent_options: {})
+        ChildAgentOptions.validate!(agent_options)
+        store = parent_context.session ? parent_context.session.store : Stores::Memory.new
         child = Session.new(store: store)
-        context.session&.append("subagent", "name" => label, "session_id" => child.id)
+        parent_context.session&.append("subagent", "name" => label, "session_id" => child.id)
         # An inline child runs on a signal derived from the parent's: the
         # parent's abort cascades down through the handle, while stopping
         # the child alone leaves the parent running.
-        signal, cascade = context.signal ? context.signal.derive : [AbortSignal.new, nil]
+        signal, cascade = if parent_context.signal
+                            parent_context.signal.derive
+                          else
+                            [AbortSignal.new, nil]
+                          end
         result = begin
           execute_child(child: child, label: label, provider: provider, system: system,
                         tools: tools, task: task, schema: schema, signal: signal,
-                        emit: context.emit, **agent_options)
+                        emit: parent_context.emit, agent_options: agent_options)
+        rescue StandardError => e
+          unless child.entries.any? { |entry| entry["type"] == Child::TERMINAL }
+            child.append(Child::TERMINAL,
+                         "status" => "failed", "error" => "#{e.class}: #{e.message}")
+          end
+          raise
         ensure
-          context.signal&.remove_callback(cascade) if cascade
+          parent_context.signal&.remove_callback(cascade) if cascade
         end
         outcome = answer(result, label, child)
         child.append(Child::TERMINAL, terminal(result))
         outcome
       end
 
-      # The host job's way back in: reconstruct provider, system, and tools
-      # from the spec through host registries, then hand them here. Reopens
-      # the child session the spawn created, runs it exactly like an inline
-      # child (started entry, lease, stop watching, terminals), streams
-      # origin-tagged events to the emit the job supplies, and reports the
-      # outcome back to the parent (see report_back). A background child
-      # runs on its own signal: the parent's turn is long over, so only
-      # stop_agent and the stop flag end it early.
+      # The host job's way back in: a runtime factory turns the durable spec
+      # into live dependencies inside the worker. The compatible direct
+      # provider/system/tools form remains for hosts that already reconstruct
+      # those values in their job. Either form is checked against the spec
+      # before execution. The child then runs exactly like an inline child,
+      # streams origin-tagged events, and reports back to its parent.
       #
-      # The child's lease is the exactly-once fence, so it is taken before
-      # anything else: refused means another process is running this child
-      # right now (a queue redelivered a live job), so leave its owner alone.
-      # Holding it, a terminal decides: present means the child was
-      # cancelled while queued or the queue retried a finished job, so
-      # there is nothing to run; absent means run, and that includes the
-      # child a crashed process left mid-run, which is exactly what queue
-      # retries are for. Either kind of no-op returns nil.
-      def run_dispatched(spec, provider:, system:, tools:, store:, emit: nil, schema: nil,
-                         **agent_options)
+      # The child lease is duplicate suppression, not an exactly-once claim:
+      # a refused delivery leaves the current holder alone. A terminal means
+      # a queued cancellation or finished retry and returns nil. The runner
+      # retains an acquired lease through terminal persistence and reporting,
+      # suppressing ordinary redelivery while that lease remains live.
+      # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/MethodLength, Metrics/PerceivedComplexity -- ordered dispatch transitions are the safety contract
+      def run_dispatched(spec, store:, emit: nil, runtime_factory: nil,
+                         provider: UNSET_RUNTIME_FIELD, system: UNSET_RUNTIME_FIELD,
+                         tools: UNSET_RUNTIME_FIELD, schema: UNSET_RUNTIME_FIELD,
+                         retry_factory_errors: true, **agent_options)
+        runtime = nil
+        resolved = nil
+        authorized = false
+        terminalize = false
+        factory_failed = false
+        retryable_exit = false
+        primary_error = nil
+        delivery_owned = false
+        spec = RuntimeContract.own_spec(spec)
+        RuntimeContract.validate_identity!(spec)
         child = Session.new(store: store, id: spec.fetch("session_id"))
+        spec = RuntimeContract.bind_spec(child.entries, spec)
+        authorized = true
+        direct = { provider: provider, system: system, tools: tools, schema: schema,
+                   agent_options: agent_options.empty? ? UNSET_RUNTIME_FIELD : agent_options }
+        RuntimeContract.validate_resolution!(factory: runtime_factory, direct: direct)
         signal = AbortSignal.new
+        terminalize = true
+        delivery_owned = true
         lease = Locks.hold(Child.lease_key(child.id),
                            stop_key: Child.stop_key(child.id), signal: signal)
-        return nil if Mistri.locks && lease.nil?
+        if Mistri.locks && lease.nil?
+          delivery_owned = false
+          return nil
+        end
 
         if Child.new(name: spec.fetch("name"), session_id: child.id, store: store).finished?
-          lease&.release
           return nil
         end
 
-        result = execute_child(child: child, label: spec.fetch("name"), provider: provider,
-                               system: system, tools: tools, task: spec.fetch("task"),
-                               schema: schema, signal: signal, emit: emit, lease: lease,
-                               **agent_options)
+        child.append(Child::STARTED, {})
+        RuntimeContract.validate_spec!(spec)
+        signal.abort!("stopped by user") if Mistri.locks&.flag?(Child.stop_key(child.id))
+        unless signal.aborted?
+          runtime = RuntimeContract.resolve(spec, factory: runtime_factory, direct: direct) do
+            factory_failed = true
+          end
+          resolved = RuntimeContract.validate_runtime!(runtime, spec)
+          signal.abort!("stopped by user") if Mistri.locks&.flag?(Child.stop_key(child.id))
+        end
+        result = if signal.aborted?
+                   Result.new(message: nil, status: :aborted)
+                 else
+                   execute_child(child: child, label: spec.fetch("name"),
+                                 provider: resolved.provider, system: resolved.system,
+                                 tools: resolved.tools, task: spec.fetch("task"),
+                                 schema: resolved.schema, signal: signal, emit: emit,
+                                 lease: lease, started: true,
+                                 agent_options: resolved.agent_options)
+                 end
         deny_pending(result, child)
         child.append(Child::TERMINAL, terminal(result))
         result
       rescue StandardError => e
+        primary_error = e
         # Completion is a contract even when the runner dies in the
         # preamble: without this, a raise before the started entry (a lock
         # backend down, say) would leave the child reading :queued forever,
         # with nothing to report and nothing for a retry to heal.
-        if child && !Child.new(name: spec.fetch("name"), session_id: child.id,
-                               store: store).finished?
-          child.append(Child::TERMINAL,
-                       "status" => "failed", "error" => "#{e.class}: #{e.message}")
+        stopped = factory_failed && child_stop_requested?(child, signal)
+        retryable = factory_failed && retry_factory_errors && !stopped
+        retryable_exit = retryable
+        if child && terminalize && !retryable &&
+           !Child.new(name: spec.fetch("name"), session_id: child.id, store: store).finished?
+          terminal = if stopped
+                       { "status" => "stopped" }
+                     else
+                       { "status" => "failed", "error" => "#{e.class}: #{e.message}" }
+                     end
+          child.append(Child::TERMINAL, terminal)
         end
-        lease&.release
         raise
       ensure
-        report_back(spec, store, emit) if child
+        cleanup_error = RuntimeContract.cleanup(runtime)
+        begin
+          report_back(spec, store, emit) if child && authorized && delivery_owned
+        ensure
+          begin
+            Mistri.locks&.clear_flag(Child.stop_key(child.id)) if child && lease && !retryable_exit
+          ensure
+            lease&.release
+          end
+        end
+        raise cleanup_error if cleanup_error && primary_error.nil?
       end
+      # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/MethodLength, Metrics/PerceivedComplexity
 
       # A child cannot wait for a human, whichever door it entered by: any
       # calls parked for approval are denied AND settled with the denial as
@@ -195,7 +261,7 @@ module Mistri
           child.deny(call.id, note: "sub-agents cannot pause for human approval")
           child.append_message(Message.tool(
                                  content: "Denied: sub-agents cannot pause for human approval.",
-                                 tool_call_id: call.id, tool_name: call.name
+                                 tool_call_id: call.id, tool_name: call.name, tool_error: true
                                ))
         end
       end
@@ -225,51 +291,18 @@ module Mistri
 
       private
 
-      # The one hardened execution path every child goes through, whichever
-      # door it entered by. From the started entry on, every exit writes a
-      # terminal, setup failures included, or the child would read as
-      # running forever. The lease says "alive right now" to other
-      # processes and its thread watches the child's stop flag; a
-      # dispatched run hands in the lease it already holds (its run-or-not
-      # decision happened under the fence), an inline child acquires its
-      # own here. Either way, this path releases it.
-      def execute_child(child:, label:, provider:, system:, tools:, task:, schema:,
-                        signal:, emit:, lease: nil, **agent_options)
-        child.append(Child::STARTED, {})
-        lease ||= Locks.hold(Child.lease_key(child.id),
-                             stop_key: Child.stop_key(child.id), signal: signal)
-        begin
-          agent = Agent.new(provider: provider, session: child, system: system,
-                            tools: tools, **agent_options)
-          origin = "#{label}##{child.id[0, 8]}"
-          tagged = ->(event) { forward(event, origin, emit) }
-          if schema
-            agent.task(task, schema: schema, signal: signal, &tagged)
-          else
-            agent.run(task, signal: signal, &tagged)
-          end
-        rescue StandardError => e
-          child.append(Child::TERMINAL, "status" => "failed", "error" => "#{e.class}: #{e.message}")
-          raise
-        ensure
-          lease&.release
-          Mistri.locks&.clear_flag(Child.stop_key(child.id))
-        end
-      end
-
-      def forward(event, origin, emit)
-        return unless emit
+      def child_stop_requested?(child, signal)
+        return false unless child
 
-        tagged = event.origin ? "#{origin}>#{event.origin}" : origin
-        emit.call(event.with(origin: tagged))
+        signal&.aborted? || Mistri.locks&.flag?(Child.stop_key(child.id))
       end
 
-      # Every terminal outcome reports back, exactly once. The report joins
-      # the parent's inbox (a typed entry that folds at its next turn
-      # boundary, exactly like a steer) and a :subagent_report event
+      # A lease-backed runner reports before releasing the child lease. The
+      # report joins the parent's inbox (a typed entry that folds at its next
+      # turn boundary, exactly like a steer) and a :subagent_report event
       # closes the child's lane in whatever UI watched the spawn. A child
       # that never ran has nothing to say, and the parent session drops a
-      # duplicate delivery, so a redelivered job cannot repeat one.
+      # sequential duplicate delivery. Without locks, the host must serialize.
       def report_back(spec, store, emit)
         facade = Child.new(name: spec.fetch("name"), session_id: spec.fetch("session_id"),
                            store: store)
@@ -301,12 +334,14 @@ module Mistri
         when :awaiting_approval
           deny_pending(result, child)
           ToolResult.new(content: "The #{label} sub-agent stopped: it needed human " \
-                                  "approval, which sub-agents cannot wait for.", ui: link)
+                                  "approval, which sub-agents cannot wait for.", ui: link,
+                         error: true)
         when :aborted
-          ToolResult.new(content: "[the #{label} sub-agent was stopped]", ui: link)
+          ToolResult.new(content: "[the #{label} sub-agent was stopped]", ui: link, error: true)
         else
           reason = result.error_message || result.status
-          ToolResult.new(content: "The #{label} sub-agent failed: #{reason}", ui: link)
+          ToolResult.new(content: "The #{label} sub-agent failed: #{reason}", ui: link,
+                         error: true)
         end
       end
 
@@ -320,3 +355,5 @@ module Mistri
     end
   end
 end
+
+require_relative "sub_agent/execution"

data/lib/mistri/task_output.rb

@@ -9,26 +9,37 @@ module Mistri
   module TaskOutput
     # Distinguishable from a parsed nil: JSON "null" is a valid value.
     PARSE_FAILED = Object.new.freeze
+    OUTPUT_TOO_LARGE = Object.new.freeze
+    OUTPUT_TOO_COMPLEX = Object.new.freeze
 
     module_function
 
     def prompt(input, schema)
+      plan = plan_for(schema)
       "#{input}\n\nAnswer with ONLY a JSON value matching this schema:\n" \
-        "#{JSON.generate(Schema.strict(schema))}"
+        "#{JSON.generate(plan.schema)}"
     end
 
     def parse(text)
-      body = text.to_s.strip
+      body = text.to_s
+      return OUTPUT_TOO_LARGE if body.bytesize > ToolArguments::MAX_BYTES
+
+      body = body.strip
       body = body[/\A```(?:json)?\s*(.*?)```\z/m, 1] || body
-      JSON.parse(body)
-    rescue JSON::ParserError
-      PARSE_FAILED
+      value, error = ToolArguments.parse_json(body)
+      return OUTPUT_TOO_LARGE if error == "too_large"
+      return OUTPUT_TOO_COMPLEX if %w[too_deep too_many_nodes number_too_large].include?(error)
+      return PARSE_FAILED if error
+
+      value
     end
 
     def errors(value, schema)
+      return ["the answer exceeds the output byte limit"] if value.equal?(OUTPUT_TOO_LARGE)
+      return ["the answer exceeds the output complexity limit"] if value.equal?(OUTPUT_TOO_COMPLEX)
       return ["the answer is not valid JSON"] if value.equal?(PARSE_FAILED)
 
-      Schema.violations(value, schema)
+      plan_for(schema).violations(value)
     end
 
     def fix_prompt(errors)
@@ -36,5 +47,12 @@ module Mistri
       "Your answer did not satisfy the required output schema. Problems:\n" \
         "#{lines}\nReply with ONLY the corrected JSON."
     end
+
+    def plan_for(schema)
+      return schema if schema.respond_to?(:schema) && schema.respond_to?(:violations)
+
+      Schema.task_plan(schema)
+    end
+    private_class_method :plan_for
   end
 end