prompt_objects 0.5.0 → 0.7.0

data/lib/prompt_objects/execution/runner.rb

@@ -0,0 +1,498 @@
+# frozen_string_literal: true
+
+module PromptObjects
+  module Execution
+    # Executes one persisted run to completion without owning scheduler policy.
+    # The scheduler may call this concurrently for different leased threads.
+    class Runner
+      class CapabilityError < StandardError; end
+      PreparedTool = Data.define(:tool_call, :execution, :policy, :policy_error)
+
+      def initialize(runtime, event_publisher: nil, llm_resolver: nil, provider_permits: nil,
+                     tool_permits: nil, cpu_permits: nil, resource_lock_manager: nil)
+        @runtime = runtime
+        @event_publisher = event_publisher
+        @llm_resolver = llm_resolver || ->(_definition) { runtime.llm }
+        @provider_permits = provider_permits
+        @tool_permits = tool_permits
+        @cpu_permits = cpu_permits
+        @resource_lock_manager = resource_lock_manager
+        store = runtime.session_store
+        raise ArgumentError, "run execution requires a session store" unless store
+
+        @threads = Repositories::ThreadRepository.new(store)
+        @runs = Repositories::RunRepository.new(store)
+        @definitions = Repositories::DefinitionRepository.new(store)
+        @messages = Repositories::MessageRepository.new(store)
+        @events = Repositories::EventRepository.new(store)
+        @tools = Repositories::ToolExecutionRepository.new(store)
+      end
+
+      def execute(run_id, cancellation_token: CancellationToken.new, tui_mode: false)
+        run = required_run(run_id)
+        definition = required_definition(run_id)
+        prompt_object = required_prompt_object(run.po_name)
+        context = build_context(run, cancellation_token, tui_mode)
+
+        start_run(run, context)
+        history = history_for(run)
+        loop do
+          context.check_cancelled!
+          response = call_model(prompt_object, definition, history, run)
+          context.check_cancelled!
+          if response.tool_calls?
+            persist_assistant_tool_message(run, response)
+            results = execute_tool_batch(run, definition, response.tool_calls, context)
+            persist_tool_results(run, results)
+            history = history_for(run)
+          else
+            return complete_run(run, response)
+          end
+        end
+      rescue CancelledError => e
+        cancel_run(run_id, e.message)
+        raise
+      rescue => e
+        fail_run(run_id, e)
+        raise
+      end
+
+      private
+
+      def required_run(run_id)
+        run = @runs.get(run_id)
+        raise ArgumentError, "unknown run: #{run_id}" unless run
+        raise ArgumentError, "run #{run_id} is #{run.status}, expected queued" unless run.status == "queued"
+
+        run
+      end
+
+      def required_definition(run_id)
+        definition = @definitions.get(run_id)
+        raise ArgumentError, "run #{run_id} has no definition snapshot" unless definition
+
+        definition
+      end
+
+      def required_prompt_object(name)
+        prompt_object = @runtime.registry.get(name)
+        raise ArgumentError, "unknown prompt object: #{name}" unless prompt_object.is_a?(PromptObject)
+
+        prompt_object
+      end
+
+      def build_context(run, cancellation_token, tui_mode)
+        Context.new(
+          env: @runtime,
+          workspace_session_id: run.workspace_session_id,
+          thread_id: run.thread_id,
+          run_id: run.id,
+          target_po: run.po_name,
+          caller_po: caller_for(run),
+          source_type: run.source_type,
+          source_name: run.source_name,
+          cancellation_token: cancellation_token,
+          deadline: run_deadline(run),
+          resource_lock_manager: @resource_lock_manager,
+          event_publisher: @event_publisher,
+          tui_mode: tui_mode
+        )
+      end
+
+      def caller_for(run)
+        return nil unless run.parent_run_id
+
+        @runs.get(run.parent_run_id)&.po_name
+      end
+
+      def start_run(run, context)
+        context.check_cancelled!
+        running = @runs.transition(run.id, status: "running")
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "run_started",
+          data: { source_type: run.source_type, turn_ordinal: run.turn_ordinal }
+        ))
+        input = @messages.append(
+          thread_id: run.thread_id,
+          run_id: run.id,
+          turn_ordinal: run.turn_ordinal,
+          role: :user,
+          content: run.input_content,
+          from_po: run.source_name || run.source_type
+        )
+        @runs.transition(
+          run.id,
+          status: running.status,
+          input_message_id: input[:id],
+          history_cutoff_message_id: input[:id]
+        )
+        # The input is durable now, so expose it immediately. Waiting for the
+        # first tool result or final model response makes a freshly submitted
+        # message appear to vanish while the run is thinking.
+        @runtime.publish_thread_snapshot(run.thread_id, po_name: run.po_name)
+      end
+
+      def call_model(prompt_object, definition, history, run)
+        publish_event(@events.append_run_event(run_id: run.id, kind: "model_started"))
+        request = lambda do
+          @llm_resolver.call(definition).chat(
+            system: prompt_object.system_prompt_for(definition),
+            messages: history,
+            tools: prompt_object.tool_descriptors_for(definition)
+          )
+        end
+        response = @provider_permits ? @provider_permits.with(&request) : request.call
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "model_completed",
+          data: { tool_call_count: response.tool_calls.size, usage: response.usage }
+        ))
+        response
+      end
+
+      def persist_assistant_tool_message(run, response)
+        @messages.append(
+          thread_id: run.thread_id,
+          run_id: run.id,
+          turn_ordinal: run.turn_ordinal,
+          role: :assistant,
+          content: response.content,
+          tool_calls: response.tool_calls.map(&:to_h),
+          usage: response.usage
+        )
+        @runtime.publish_thread_snapshot(run.thread_id, po_name: run.po_name)
+      end
+
+      def execute_tool_batch(run, definition, tool_calls, context)
+        batch_id = SecureRandom.uuid
+        prepared = tool_calls.each_with_index.map do |tool_call, position|
+          capability = @runtime.registry.get(tool_call.name)
+          policy, policy_error = resolve_policy(capability, tool_call, context)
+          execution = @tools.create(
+            run_id: run.id,
+            provider_tool_call_id: tool_call.id,
+            batch_id: batch_id,
+            batch_position: position,
+            capability_name: tool_call.name,
+            capability_policy_version: definition.capability_policy_version,
+            arguments: tool_call.arguments,
+            resource_keys: policy.resource_keys
+          )
+          PreparedTool.new(
+            tool_call: tool_call,
+            execution: execution,
+            policy: policy,
+            policy_error: policy_error
+          )
+        end
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "tool_batch_started",
+          data: { batch_id: batch_id, tool_execution_ids: prepared.map { |item| item.execution.id } }
+        ))
+
+        results = Array.new(prepared.size)
+        parallel_wave = []
+        prepared.each_with_index do |item, position|
+          if item.policy.parallel_safe
+            parallel_wave << [item, position]
+          else
+            execute_parallel_wave(parallel_wave, results, run, definition, context)
+            parallel_wave = []
+            results[position] = execute_tool(run, definition, item, context)
+          end
+        end
+        execute_parallel_wave(parallel_wave, results, run, definition, context)
+        results
+      end
+
+      def execute_parallel_wave(wave, results, run, definition, context)
+        return if wave.empty?
+
+        threads = wave.map do |item, position|
+          Thread.new do
+            Thread.current.report_on_exception = false
+            [position, execute_tool(run, definition, item, context)]
+          end
+        end
+        errors = []
+        threads.each do |thread|
+          begin
+            position, result = thread.value
+            results[position] = result
+          rescue => e
+            errors << e
+          end
+        end
+        raise errors.first if errors.any?
+      end
+
+      def execute_tool(run, definition, prepared, context)
+        tool_call = prepared.tool_call
+        execution = prepared.execution
+        context.check_cancelled!
+        tool_context = context.for_tool_call(tool_call)
+        raise prepared.policy_error if prepared.policy_error
+
+        validate_capability!(definition, tool_call.name)
+        capability = @runtime.registry.get(tool_call.name)
+        raise CapabilityError, "Unknown capability: #{tool_call.name}" unless capability
+
+        invoke = lambda do
+          with_resource_locks(prepared.policy, tool_context) do
+            @tools.transition(execution.id, status: "running")
+            publish_event(@events.append_run_event(
+              run_id: run.id,
+              kind: "tool_started",
+              data: { tool_execution_id: execution.id, name: tool_call.name }
+            ))
+            @runtime.bus.publish(from: run.po_name, to: tool_call.name, message: tool_call.arguments)
+            if capability.is_a?(PromptObject)
+              execute_delegation(run, capability, tool_call, tool_context)
+            else
+              capability.receive(tool_call.arguments, context: tool_context)
+            end
+          end
+        end
+        result = if capability.is_a?(PromptObject)
+                   invoke.call
+                 else
+                   with_workload_permit(prepared.policy, &invoke)
+                 end
+        tool_context.check_cancelled!
+        content = serialize_tool_content(result)
+        @tools.transition(execution.id, status: "completed", result: result)
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "tool_completed",
+          data: { tool_execution_id: execution.id, name: tool_call.name }
+        ))
+        @runtime.bus.publish(from: tool_call.name, to: run.po_name, message: result)
+        { tool_call_id: tool_call.id, name: tool_call.name, content: content }
+      rescue CancelledError => e
+        @tools.transition(execution.id, status: "cancelled", cancellation_reason: e.message)
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "tool_cancelled",
+          data: { tool_execution_id: execution.id, name: tool_call.name }
+        ))
+        raise
+      rescue ResourceLockManager::LockTimeoutError => e
+        @tools.transition(execution.id, status: "timed_out", error: e)
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "tool_failed",
+          data: { tool_execution_id: execution.id, name: tool_call.name, error: error_data(e) }
+        ))
+        { tool_call_id: tool_call.id, name: tool_call.name, content: "Error: #{e.message}" }
+      rescue => e
+        @tools.transition(execution.id, status: "failed", error: e)
+        publish_event(@events.append_run_event(
+          run_id: run.id,
+          kind: "tool_failed",
+          data: { tool_execution_id: execution.id, name: tool_call.name, error: error_data(e) }
+        ))
+        { tool_call_id: tool_call.id, name: tool_call.name, content: "Error: #{e.message}" }
+      end
+
+      def resolve_policy(capability, tool_call, context)
+        unless capability
+          fallback = Policy.new(parallel_safe: false, side_effect: :exclusive)
+          return [fallback.resolve(arguments: tool_call.arguments, context: context, capability: nil), nil]
+        end
+
+        [capability.resolved_execution_policy(arguments: tool_call.arguments, context: context), nil]
+      rescue => e
+        fallback = Policy.new(parallel_safe: false, side_effect: :exclusive)
+        [fallback.resolve(arguments: tool_call.arguments, context: context, capability: capability), e]
+      end
+
+      def with_workload_permit(policy, &block)
+        permits = policy.workload == :cpu ? @cpu_permits : @tool_permits
+        permits ? permits.with(&block) : block.call
+      end
+
+      def with_resource_locks(policy, context, &block)
+        return block.call unless @resource_lock_manager
+
+        @resource_lock_manager.with(
+          policy.resource_keys,
+          cancellation_token: context.cancellation_token,
+          &block
+        )
+      end
+
+      def execute_delegation(parent_run, target_po, tool_call, context)
+        definition = @runtime.execution_definition_for(target_po)
+        child_thread = nil
+        child_run = nil
+        @runtime.with_child_run_admission(parent_run) do |depth|
+          child_thread = @threads.create(
+            workspace_session_id: parent_run.workspace_session_id,
+            po_name: target_po.name,
+            thread_type: "delegation",
+            parent_thread_id: parent_run.thread_id,
+            parent_run_id: parent_run.id,
+            retention_policy: "full"
+          )
+          child_run = @runs.enqueue(
+            workspace_session_id: parent_run.workspace_session_id,
+            thread_id: child_thread.id,
+            po_name: target_po.name,
+            input_content: delegation_content(parent_run, target_po, tool_call.arguments),
+            definition_version: definition.version,
+            source_type: "delegation",
+            source_name: parent_run.po_name,
+            parent_run_id: parent_run.id,
+            priority: parent_run.priority,
+            metadata: {
+              delegation_depth: depth,
+              deadline: context.deadline&.iso8601(6)
+            }.compact
+          )
+        end
+        @definitions.attach(run_id: child_run.id, definition: definition)
+        @runtime.publish_workspace_snapshot(parent_run.workspace_session_id)
+        publish_event(@events.append_run_event(
+          run_id: child_run.id,
+          kind: "run_queued",
+          data: { source_type: child_run.source_type, turn_ordinal: child_run.turn_ordinal }
+        ))
+        publish_event(@events.append_run_event(
+          run_id: parent_run.id,
+          kind: "delegation_started",
+          data: { child_run_id: child_run.id, child_thread_id: child_thread.id, target_po: target_po.name }
+        ))
+        @runs.transition(parent_run.id, status: "waiting")
+        publish_event(@events.append_run_event(
+          run_id: parent_run.id,
+          kind: "run_waiting",
+          data: { child_run_id: child_run.id, reason: "awaited delegation" }
+        ))
+        response = execute(child_run.id, cancellation_token: context.cancellation_token, tui_mode: context.tui_mode)
+        publish_event(@events.append_run_event(
+          run_id: parent_run.id,
+          kind: "delegation_completed",
+          data: { child_run_id: child_run.id, child_thread_id: child_thread.id, target_po: target_po.name }
+        ))
+        response
+      ensure
+        @runtime.publish_thread_snapshot(child_thread.id, po_name: target_po.name) if child_thread
+        @runtime.publish_workspace_snapshot(parent_run.workspace_session_id)
+        current = @runs.get(parent_run.id)
+        @runs.transition(parent_run.id, status: "running") if current&.status == "waiting"
+      end
+
+      def delegation_content(parent_run, target_po, arguments)
+        content = if arguments.is_a?(Hash)
+                    arguments[:message] || arguments["message"] || arguments.to_s
+                  else
+                    arguments.to_s
+                  end
+        <<~MESSAGE.strip
+          ---
+          [Delegation Context]
+          Called by: #{parent_run.po_name}
+          Target: #{target_po.name}
+          ---
+
+          #{content}
+        MESSAGE
+      end
+
+      def run_deadline(run)
+        value = run.metadata[:deadline] || run.metadata["deadline"]
+        value && Time.iso8601(value)
+      end
+
+      def persist_tool_results(run, results)
+        @messages.append(
+          thread_id: run.thread_id,
+          run_id: run.id,
+          turn_ordinal: run.turn_ordinal,
+          role: :tool,
+          tool_results: results
+        )
+        @runtime.publish_thread_snapshot(run.thread_id, po_name: run.po_name)
+      end
+
+      def complete_run(run, response)
+        @messages.append(
+          thread_id: run.thread_id,
+          run_id: run.id,
+          turn_ordinal: run.turn_ordinal,
+          role: :assistant,
+          content: response.content,
+          usage: response.usage
+        )
+        @runs.transition(
+          run.id,
+          status: "completed",
+          final_response: response.content,
+          usage: response.usage
+        )
+        publish_event(@events.append_run_event(run_id: run.id, kind: "run_completed"))
+        @runtime.enforce_retention
+        response.content
+      end
+
+      def cancel_run(run_id, reason)
+        run = @runs.get(run_id)
+        return unless run && %w[queued running waiting].include?(run.status)
+
+        @runs.transition(run_id, status: "cancelled", cancellation_reason: reason)
+        publish_event(@events.append_run_event(run_id: run_id, kind: "run_cancelled", data: { reason: reason }))
+        @runtime.enforce_retention
+      end
+
+      def fail_run(run_id, error)
+        run = @runs.get(run_id)
+        return unless run && %w[queued running waiting].include?(run.status)
+
+        @runs.transition(run_id, status: "failed", error: error)
+        publish_event(@events.append_run_event(
+          run_id: run_id,
+          kind: "run_failed",
+          data: { error: error_data(error) }
+        ))
+        @runtime.enforce_retention
+      end
+
+      def history_for(run)
+        @messages.list(thread_id: run.thread_id, through_turn: run.turn_ordinal, limit: 10_000).map do |message|
+          history = { role: message[:role], content: message[:content] }
+          history[:from] = message[:from_po] if message[:from_po]
+          history[:tool_calls] = message[:tool_calls] if message[:tool_calls]
+          history[:results] = message[:tool_results] if message[:tool_results]
+          history
+        end
+      end
+
+      def validate_capability!(definition, capability_name)
+        declared = definition.config["capabilities"] || definition.config[:capabilities] || []
+        allowed = declared + UNIVERSAL_CAPABILITIES
+        return if allowed.include?(capability_name)
+
+        raise CapabilityError, "Capability '#{capability_name}' is not available"
+      end
+
+      def serialize_tool_content(result)
+        return result if result.is_a?(String)
+
+        JSON.generate(result)
+      rescue JSON::GeneratorError
+        result.to_s
+      end
+
+      def error_data(error)
+        { class: error.class.name, message: error.message }
+      end
+
+      def publish_event(event)
+        @event_publisher&.call(event)
+        event
+      end
+    end
+  end
+end