turnkit 0.2.9 → 0.3.0

data/lib/turnkit/result.rb

@@ -2,19 +2,44 @@
 
 module TurnKit
   class Result
-    attr_reader :text, :tool_calls, :usage, :model, :finish_reason, :output_data
+    attr_reader :parts, :usage, :model, :finish_reason, :output_data
 
-    def initialize(text: "", tool_calls: [], usage: Usage.new, model: nil, finish_reason: nil, output_data: nil)
-      @text = text.to_s
-      @tool_calls = Array(tool_calls)
+    def initialize(text: "", tool_calls: [], parts: nil, usage: Usage.new, model: nil, finish_reason: nil, output_data: nil)
+      @parts = parts ? normalize_parts(parts) : synthesize_parts(text: text, tool_calls: tool_calls)
       @usage = usage || Usage.new
       @model = model
       @finish_reason = finish_reason
       @output_data = output_data
     end
 
+    def text
+      parts.filter_map { |part| part["text"] if part["type"] == "text" }.join("\n")
+    end
+
+    def tool_calls
+      parts.filter_map do |part|
+        next unless part["type"] == "tool_call"
+
+        ToolCall.new(id: part.fetch("id"), name: part.fetch("name"), arguments: part["arguments"] || {}, arguments_error: part["arguments_error"])
+      end
+    end
+
     def tool_calls?
       tool_calls.any?
     end
+
+    private
+      def synthesize_parts(text:, tool_calls:)
+        parts = []
+        parts << { "type" => "text", "text" => text.to_s } unless text.to_s.empty?
+        Array(tool_calls).each do |call|
+          parts << { "type" => "tool_call", "id" => call.id, "name" => call.name, "arguments" => call.arguments, "arguments_error" => call.arguments_error }.compact
+        end
+        parts
+      end
+
+      def normalize_parts(value)
+        Array(value).map { |part| part.to_h.transform_keys(&:to_s) }
+      end
   end
 end

data/lib/turnkit/run.rb

@@ -14,6 +14,8 @@ module TurnKit
     def output = output_text
     def output_text = turn.output_text
     def output_data = turn.output_data
+    def policy_audit = turn.policy_audit
+    def policy_clean? = policy_audit.nil? || policy_audit.fetch("clean", false)
     def usage = Usage.from_records(turn_records)
     def cost = Cost.from_records(turn_records)
     def steps = turn_records.length
@@ -25,9 +27,8 @@ module TurnKit
     end
 
     def messages
-      turn_records.flat_map do |record|
-        conversation = turn.store.load_conversation(record.fetch("conversation_id"))
-        turn.store.list_messages(conversation.fetch("id"))
+      turn_records.map { |record| record.fetch("conversation_id") }.uniq.flat_map do |conversation_id|
+        turn.store.list_messages(conversation_id).map { |attrs| Message.new(attrs) }
       end
     end
 

data/lib/turnkit/schema_check.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module TurnKit
+  module SchemaCheck
+    module_function
+
+    def validate!(value, schema, error_class: ToolValidationError, label: "input")
+      schema = stringify_schema(schema || {})
+      type = schema["type"] || "object"
+      validate_type!(value, type, schema, error_class: error_class, label: label)
+      validate_enum!(value, schema, error_class: error_class, label: label)
+
+      if type == "object" && schema["properties"]
+        attrs = value.respond_to?(:to_h) ? value.to_h.transform_keys(&:to_s) : {}
+        required = Array(schema["required"]).map(&:to_s)
+        missing = required.reject { |name| attrs.key?(name) }
+        raise error_class, "#{label} missing required field#{missing.length == 1 ? "" : "s"}: #{missing.join(", ")}" if missing.any?
+
+        schema.fetch("properties", {}).each do |name, child_schema|
+          next unless attrs.key?(name)
+
+          validate!(attrs[name], child_schema, error_class: error_class, label: "#{label}.#{name}")
+        end
+      elsif type == "array" && schema["items"] && value.is_a?(Array)
+        value.each_with_index do |item, index|
+          validate!(item, schema["items"], error_class: error_class, label: "#{label}[#{index}]")
+        end
+      end
+
+      true
+    end
+
+    def stringify_schema(value)
+      case value
+      when Hash
+        value.transform_keys(&:to_s).transform_values { |nested| stringify_schema(nested) }
+      when Array
+        value.map { |nested| stringify_schema(nested) }
+      when Symbol
+        value.to_s
+      else
+        value
+      end
+    end
+
+    def validate_type!(value, type, schema, error_class:, label:)
+      return if value.nil? && !Array(schema["required"]).include?(label.to_s)
+
+      valid = case type.to_s
+      when "string" then value.is_a?(String)
+      when "integer" then value.is_a?(Integer)
+      when "number" then value.is_a?(Numeric)
+      when "boolean" then value == true || value == false
+      when "array" then value.is_a?(Array)
+      when "object" then value.is_a?(Hash)
+      else true
+      end
+      raise error_class, "#{label} must be a #{type}" unless valid
+    end
+
+    def validate_enum!(value, schema, error_class:, label:)
+      enum = schema["enum"]
+      return unless enum && !Array(enum).include?(value)
+
+      raise error_class, "#{label} must be one of: #{Array(enum).join(", ")}"
+    end
+  end
+end

data/lib/turnkit/skill.rb

@@ -1,13 +1,15 @@
 # frozen_string_literal: true
 
+require "yaml"
+
 module TurnKit
   class Skill
     attr_reader :key, :name, :description, :content
 
     def self.from_file(path, key: nil, name: nil, description: "")
-      content = File.read(path)
+      content, metadata = parse_file(File.read(path))
       base = File.basename(path, File.extname(path))
-      new(key: key || base, name: name || base.tr("_-", " ").split.map(&:capitalize).join(" "), description: description, content: content)
+      new(key: key || base, name: name || metadata["name"] || base.tr("_-", " ").split.map(&:capitalize).join(" "), description: description.to_s.empty? ? metadata["description"].to_s : description, content: content)
     end
 
     def self.from_directory(path, pattern: "*.md")
@@ -23,5 +25,17 @@ module TurnKit
       raise ArgumentError, "name is required" if @name.empty?
       raise ArgumentError, "content is required" if @content.empty?
     end
+
+    def self.parse_file(content)
+      text = content.to_s
+      return [ text, {} ] unless text.start_with?("---\n")
+
+      _, frontmatter, body = text.split(/^---\s*$/, 3)
+      return [ text, {} ] unless body
+
+      [ body.sub(/\A\n/, ""), YAML.safe_load(frontmatter, permitted_classes: [ Symbol ], aliases: false) || {} ]
+    rescue Psych::SyntaxError
+      [ text, {} ]
+    end
   end
 end

data/lib/turnkit/store.rb

@@ -12,6 +12,12 @@ module TurnKit
     def create_turn(_attributes) = raise(NotImplementedError)
     def load_turn(_id) = raise(NotImplementedError)
     def update_turn(_id, _attributes) = raise(NotImplementedError)
+    def claim_turn(id, from: "pending", to: "running", **attributes)
+      turn = load_turn(id)
+      return nil unless turn["status"] == from
+
+      update_turn(id, attributes.merge(status: to))
+    end
     def list_turns(root_turn_id: nil, conversation_id: nil, agent_name: nil) = raise(NotImplementedError)
 
     def create_tool_execution(_attributes) = raise(NotImplementedError)

data/lib/turnkit/stores/active_record_store.rb

@@ -38,7 +38,6 @@ module TurnKit
           kind: message.fetch("kind"),
           sequence: message.fetch("sequence"),
           content: message.fetch("content"),
-          text: message.fetch("text"),
           tool_execution_uid: message["tool_execution_id"],
           provider_message_id: message["provider_message_id"],
           metadata: message.fetch("metadata")
@@ -93,6 +92,15 @@ module TurnKit
       turn_hash(record)
     end
 
+    def claim_turn(id, from: "pending", to: "running", **attributes)
+      attrs = Record.turn_update(attributes.merge(status: to))
+      attrs.delete("output_data") unless turn_has_attribute?("output_data")
+      affected = turn_class.where(uid: id, status: from).update_all(attrs.merge(updated_at: Clock.now))
+      return nil if affected.zero?
+
+      load_turn(id)
+    end
+
     def list_turns(root_turn_id: nil, conversation_id: nil, agent_name: nil)
       scope = turn_class.all
       scope = scope.where(root_turn_uid: root_turn_id) if root_turn_id
@@ -185,7 +193,7 @@ module TurnKit
         {
           "id" => record.uid, "conversation_id" => record.conversation_uid, "turn_id" => record.turn_uid,
           "role" => record.role, "kind" => record.kind, "sequence" => record.sequence, "content" => record.content,
-          "text" => record.text, "tool_execution_id" => record.tool_execution_uid,
+          "tool_execution_id" => record.tool_execution_uid,
           "provider_message_id" => record.provider_message_id, "metadata" => record.metadata || {}, "created_at" => record.created_at
         }
       end

data/lib/turnkit/sub_agent_tool.rb

@@ -42,7 +42,8 @@ module TurnKit
         agent: sub_agent,
         on_event: parent_turn.agent.effective_on_event
       )
-      { "conversation_id" => conversation.id, "turn_id" => child.id, "status" => child.status, "result" => child.output_text, "output_data" => child.output_data }.compact
+      error = child.store.load_turn(child.id)["error"] if child.failed?
+      { "conversation_id" => conversation.id, "turn_id" => child.id, "status" => child.status, "result" => child.output_text, "output_data" => child.output_data, "error" => error }.compact
     end
   end
 end

data/lib/turnkit/system_prompt.rb

@@ -194,7 +194,7 @@ module TurnKit
 
       tagged(
         "skills_available",
-        "Load or follow a skill when the task matches its description.\n\n#{entries.join("\n")}"
+        "These skills are listed but not loaded. When a task matches a skill description, call load_skill with the skill key before relying on it.\n\n#{entries.join("\n")}"
       )
     end
 

data/lib/turnkit/tool.rb

@@ -106,7 +106,7 @@ module TurnKit
             next
           end
 
-          validate_value!(value, param)
+          SchemaCheck.validate!(value, schema_for(param), error_class: ToolValidationError, label: name)
           normalized[name] = value
         end
         normalized
@@ -169,26 +169,7 @@ module TurnKit
         end
 
         def validate_value!(value, param)
-          return if value.nil? && !param.fetch(:required)
-
-          case param.fetch(:type)
-          when :string, :enum
-            raise ToolValidationError, "#{param.fetch(:name)} must be a string" unless value.is_a?(String)
-          when :integer
-            raise ToolValidationError, "#{param.fetch(:name)} must be an integer" unless value.is_a?(Integer)
-          when :number
-            raise ToolValidationError, "#{param.fetch(:name)} must be a number" unless value.is_a?(Numeric)
-          when :boolean
-            raise ToolValidationError, "#{param.fetch(:name)} must be a boolean" unless value == true || value == false
-          when :array
-            raise ToolValidationError, "#{param.fetch(:name)} must be an array" unless value.is_a?(Array)
-          when :object
-            raise ToolValidationError, "#{param.fetch(:name)} must be an object" unless value.is_a?(Hash)
-          end
-
-          if param[:enum] && !Array(param[:enum]).include?(value)
-            raise ToolValidationError, "#{param.fetch(:name)} must be one of: #{Array(param[:enum]).join(", ")}"
-          end
+          SchemaCheck.validate!(value, schema_for(param), error_class: ToolValidationError, label: param.fetch(:name))
         end
 
         def accepts_turnkit_context?(instance)

data/lib/turnkit/tool_call.rb

@@ -4,10 +4,10 @@ module TurnKit
   class ToolCall
     attr_reader :id, :name, :arguments, :arguments_error
 
-    def initialize(id:, name:, arguments: {})
+    def initialize(id:, name:, arguments: {}, arguments_error: nil)
       @id = id.to_s
       @name = name.to_s
-      @arguments_error = nil
+      @arguments_error = arguments_error
       @arguments = normalize_arguments(arguments)
     end
 
@@ -23,7 +23,7 @@ module TurnKit
           {}
         end
       rescue JSON::ParserError
-        @arguments_error = "invalid JSON arguments"
+        @arguments_error ||= "invalid JSON arguments"
         {}
       end
   end

data/lib/turnkit/tool_runner.rb

@@ -7,9 +7,12 @@ module TurnKit
     end
 
     def dispatch(tool_calls)
-      tool_calls.each do |tool_call|
+      tool_calls.each_with_index do |tool_call, index|
         execution = run(tool_call)
-        return execution if execution.completed? && tool_for(tool_call.name)&.ends_turn?
+        if execution.completed? && tool_for(tool_call.name)&.ends_turn?
+          skip_remaining(tool_calls.drop(index + 1), terminal: tool_call)
+          return execution
+        end
       end
       nil
     end
@@ -23,9 +26,10 @@ module TurnKit
       attr_reader :turn
 
       def run(tool_call)
-        turn.budget.count_tool_execution!
-        tool = tool_for(tool_call.name)
         execution = ToolExecution.new(create_execution(tool_call))
+        heartbeat!
+
+        tool = tool_for(tool_call.name)
 
         unless tool
           return finish_error(execution, tool_call, "unknown tool: #{tool_call.name}")
@@ -35,6 +39,13 @@ module TurnKit
           return finish_error(execution, tool_call, tool_call.arguments_error)
         end
 
+        begin
+          turn.budget.count_tool_execution!(tool_call.name)
+        rescue BudgetError => error
+          finish_error(execution, tool_call, error.message, details: { "class" => error.class.name, "budget_denied" => true })
+          raise
+        end
+
         context = ToolContext.new(turn: turn, execution: execution)
         payload = begin
           normalize_payload(call_tool(tool, tool_call.arguments, context: context))
@@ -56,32 +67,50 @@ module TurnKit
       end
 
       def finish_success(execution, tool_call, payload)
+        json = payload.to_json
         attrs = turn.store.update_tool_execution(execution.id, "status" => "completed", "result" => payload, "completed_at" => Clock.now)
-        append_result(execution, tool_call, payload)
-        turn.emit("tool_call.completed", id: tool_call.id, name: tool_call.name)
+        append_result(execution, tool_call, payload, json: json, error: false)
+        heartbeat!
+        turn.emit("tool_call.completed", id: tool_call.id, name: tool_call.name, result_chars: json.length)
         ToolExecution.new(attrs)
       end
 
       def finish_error(execution, tool_call, message, details: nil)
         error = { "message" => message.to_s, "details" => details }.compact
+        json = error.to_json
         attrs = turn.store.update_tool_execution(execution.id, "status" => "failed", "error" => error, "completed_at" => Clock.now)
-        append_result(execution, tool_call, error)
-        turn.emit("tool_call.failed", id: tool_call.id, name: tool_call.name, error: error)
+        append_result(execution, tool_call, error, json: json, error: true)
+        heartbeat!
+        turn.emit("tool_call.failed", id: tool_call.id, name: tool_call.name, error: error, result_chars: json.length)
         ToolExecution.new(attrs)
       end
 
-      def append_result(execution, tool_call, payload)
+      def append_result(execution, tool_call, payload, json: payload.to_json, error: false)
         message = turn.conversation.append_message(
           role: "tool",
           kind: "tool_result",
-          text: payload.to_json,
+          content: [ { "type" => "tool_result", "tool_call_id" => tool_call.id, "text" => json, "error" => error } ],
           turn_id: turn.id,
           tool_execution_id: execution.id,
-          metadata: { "tool_call_id" => tool_call.id, "tool_name" => tool_call.name }
+          metadata: { "tool_name" => tool_call.name }
         )
         turn.emit("message.created", message_id: message.id, role: message.role, kind: message.kind)
       end
 
+      def skip_remaining(calls, terminal:)
+        calls.each do |call|
+          payload = { "skipped" => true, "message" => "not executed: turn ended by #{terminal.name}" }
+          execution = ToolExecution.new(create_execution(call))
+          attrs = turn.store.update_tool_execution(execution.id, "status" => "cancelled", "result" => payload, "completed_at" => Clock.now)
+          append_result(ToolExecution.new(attrs), call, payload)
+          turn.emit("tool_call.skipped", id: call.id, name: call.name)
+        end
+      end
+
+      def heartbeat!
+        turn.send(:heartbeat!)
+      end
+
       def tool_for(name)
         turn.agent.effective_tools.find { |tool| tool.tool_name == name.to_s }
       end

data/lib/turnkit/turn.rb

@@ -36,36 +36,52 @@ module TurnKit
       @on_event = block if block
       return self unless status == "pending"
 
-      update!(status: "running", started_at: Clock.now, heartbeat_at: Clock.now)
+      claimed = store.claim_turn(id, from: "pending", to: "running", started_at: Clock.now, heartbeat_at: Clock.now)
+      return self unless claimed
+
+      @record = claimed
+      @started_at = @record["started_at"]
       emit("turn.started", status: status, model: model)
       agent.effective_client.validate!(model: model)
+      @budget = Budget.resume(store: store, root_turn_id: root_turn_id, limits: budget_limits)
+      revisions_used = 0
       loop do
         budget.check!(depth: depth)
-        budget.count_iteration!
+        count_iteration!
         TurnKit::Compaction.maybe_compact!(self)
 
         request = model_request
-        emit("model.requested", model: request.model, tool_names: request.tool_names)
+        emit_model_requested("model.requested", request)
         result = call_client(request)
-        emit("model.completed", model: result.model || model, tool_call_count: result.tool_calls.length)
         result_cost = Cost.from_usage(result.usage, model: result.model || model)
 
-        budget.add_cost!(result_cost.total)
         add_usage!(result.usage, cost: result_cost)
+        emit_model_completed("model.completed", result, result_cost, model: model)
+        budget.add_cost!(result_cost.total)
         persist_assistant_message(result)
 
         if result.tool_calls?
           runner = ToolRunner.new(self)
           terminal = runner.dispatch(result.tool_calls)
           if terminal
-            complete_from_terminal_tool(runner, terminal)
-            break
+            candidate = append_terminal_completion(runner, terminal)
+          else
+            next
           end
         else
-          update!(status: "completed", output_text: result.text, output_data: result.output_data, completed_at: Clock.now)
-          emit("turn.completed", status: status, output_text: result.text)
-          break
+          candidate = result.text
+        end
+
+        audit = check_policy(candidate, output_data: result.output_data)
+        if should_revise?(audit, revisions_used)
+          revisions_used += 1
+          append_revision_message(audit, attempt: revisions_used, terminal_tool_name: terminal&.tool_name)
+          emit("output_policy.revision", violation_count: audit.violations.length, attempt: revisions_used)
+          next
         end
+
+        complete_with_output(candidate, output_data: result.output_data, audit: audit)
+        break
       end
       reload
       self
@@ -96,6 +112,10 @@ module TurnKit
       @record["output_data"]
     end
 
+    def policy_audit
+      (@record["options"] || {})["policy_audit"]
+    end
+
     def usage
       Usage.from_h(@record["usage"] || {})
     end
@@ -125,6 +145,28 @@ module TurnKit
       emit_event(Event.new(type: type, turn_id: id, conversation_id: conversation.id, payload: payload))
     end
 
+    def internal_model_call(model:, messages:, instructions:, tools: [], thinking: nil, output_schema: nil, metadata: {}, purpose:, client: nil)
+      request = ModelRequest.new(
+        model: model,
+        messages: messages,
+        tools: tools,
+        instructions: instructions,
+        thinking: thinking,
+        output_schema: output_schema,
+        metadata: { purpose: purpose.to_s, turn_id: id, conversation_id: conversation.id }.merge(metadata || {})
+      )
+      model_client = client || agent.effective_client
+      model_client.validate!(model: request.model)
+
+      emit_model_requested("#{purpose}.model.requested", request)
+      result = call_client(request, client: model_client)
+      result_cost = Cost.from_usage(result.usage, model: result.model || request.model)
+      add_usage!(result.usage, cost: result_cost)
+      emit_model_completed("#{purpose}.model.completed", result, result_cost, model: request.model)
+      budget.add_cost!(result_cost.total)
+      result
+    end
+
     private
       def model_request
         prompt = SystemPrompt.new(agent: agent, turn: self, conversation: conversation, mode: prompt_mode || agent.effective_prompt_mode(turn: self))
@@ -148,7 +190,7 @@ module TurnKit
         )
       end
 
-      def call_client(request)
+      def call_client(request, client: agent.effective_client)
         kwargs = {
           model: request.model,
           messages: request.messages,
@@ -159,9 +201,9 @@ module TurnKit
           metadata: request.metadata,
           on_event: ->(event) { emit_event(event) }
         }
-        accepted = chat_keyword_names(agent.effective_client)
+        accepted = chat_keyword_names(client)
         kwargs = kwargs.slice(*accepted) unless accepted.include?(:keyrest)
-        agent.effective_client.chat(**kwargs)
+        client.chat(**kwargs)
       end
 
       def chat_keyword_names(client)
@@ -176,6 +218,26 @@ module TurnKit
         MessageProjection.for(TurnKit::Compaction.project(conversation.messages_for_turn(self)))
       end
 
+      def emit_model_requested(type, request)
+        emit(
+          type,
+          model: request.model,
+          tool_names: request.tool_names,
+          message_count: request.messages.length,
+          prompt: request.report
+        )
+      end
+
+      def emit_model_completed(type, result, cost, model: self.model)
+        emit(
+          type,
+          model: result.model || model,
+          tool_call_count: result.tool_calls.length,
+          usage: result.usage.to_h,
+          cost: cost.to_h
+        )
+      end
+
       def thinking_from_options
         options = (@record["options"] || {}).transform_keys(&:to_s)
         return Agent.normalize_thinking(options["thinking"]) if options.key?("thinking")
@@ -203,9 +265,9 @@ module TurnKit
           message = conversation.append_message(
             role: "assistant",
             kind: "tool_call",
-            text: result.text,
+            content: result.parts,
             turn_id: id,
-            metadata: { "tool_calls" => result.tool_calls.map { |call| { "id" => call.id, "name" => call.name, "arguments" => call.arguments } } }
+            metadata: {}
           )
           emit("message.created", message_id: message.id, role: message.role, kind: message.kind)
           result.tool_calls.each { |call| emit("tool_call.created", id: call.id, name: call.name) }
@@ -215,12 +277,73 @@ module TurnKit
         end
       end
 
-      def complete_from_terminal_tool(runner, execution)
+      def append_terminal_completion(runner, execution)
         message = runner.completion_message(execution)
         assistant = conversation.append_message(role: "assistant", kind: "text", text: message, turn_id: id)
         emit("message.created", message_id: assistant.id, role: assistant.role, kind: assistant.kind)
-        update!(status: "completed", output_text: message, completed_at: Clock.now)
-        emit("turn.completed", status: status, output_text: message)
+        message
+      end
+
+      def complete_with_output(text, output_data: nil, audit: nil)
+        attrs = { output_text: text, output_data: output_data, completed_at: Clock.now }
+        if audit && !audit.clean? && agent.output_policy_mode == :fail
+          attrs[:status] = "failed"
+          attrs[:error] = { "class" => "TurnKit::OutputAudit", "message" => audit.messages.join("; "), "policy_audit" => audit.to_h }
+        else
+          attrs[:status] = "completed"
+        end
+        update!(attrs)
+        persist_policy_audit(audit) if audit
+
+        if failed?
+          emit("turn.failed", error: @record["error"])
+        else
+          emit("turn.completed", status: status, output_text: text)
+        end
+      end
+
+      def check_policy(text, output_data: nil)
+        constraints = agent.effective_output_policy
+        return nil if constraints.empty?
+
+        output = output_data.nil? ? text : output_data
+        TurnKit.check_output_policy(output, constraints: constraints, context: { turn: self, output_text: text, output_data: output_data })
+      end
+
+      def persist_policy_audit(audit)
+        options = (@record["options"] || {}).merge("policy_audit" => audit.to_h)
+        update!(options: options)
+        emit("output_policy.completed", clean: audit.clean?, violation_count: audit.violations.length)
+      end
+
+      def should_revise?(audit, revisions_used)
+        audit && !audit.clean? && revisions_used < agent.output_retries
+      end
+
+      def append_revision_message(audit, attempt:, terminal_tool_name: nil)
+        text = <<~TEXT.strip
+          The previous output failed policy checks.
+
+          Revise the previous output. Do not introduce new claims.
+          Do not deviate from the skill or policy below.
+
+          #{revision_policy_blocks}
+
+          Violations:
+          #{audit.violations.each_with_index.map { |violation, index| "#{index + 1}. #{violation.rule}: #{violation.message}" }.join("\n")}
+          #{terminal_tool_name ? "\nResubmit via #{terminal_tool_name}." : ""}
+        TEXT
+        message = conversation.append_message(role: "user", kind: "text", text: text, turn_id: id, metadata: { "source" => "output_policy", "attempt" => attempt })
+        emit("message.created", message_id: message.id, role: message.role, kind: message.kind)
+      end
+
+      def revision_policy_blocks
+        agent.effective_output_policy.filter_map do |policy|
+          next unless policy.respond_to?(:content)
+
+          key = policy.respond_to?(:name) ? policy.name : "output_policy"
+          "<skill key=\"#{key}\">\n#{policy.content}\n</skill>"
+        end.join("\n\n")
       end
 
       def add_usage!(usage, cost: nil)
@@ -239,6 +362,27 @@ module TurnKit
         update!(attributes)
       end
 
+      def count_iteration!
+        budget.count_iteration!
+        options = (@record["options"] || {}).merge("iterations" => (@record.dig("options", "iterations").to_i + 1))
+        update!(options: options)
+      end
+
+      def heartbeat!
+        update!(heartbeat_at: Clock.now)
+      end
+
+      def budget_limits
+        {
+          max_iterations: agent.max_iterations || TurnKit.max_iterations,
+          timeout: agent.timeout || TurnKit.timeout,
+          max_depth: agent.max_depth || TurnKit.max_depth,
+          max_tool_executions: agent.max_tool_executions || TurnKit.max_tool_executions,
+          max_tool_executions_by_name: agent.max_tool_executions_by_name || TurnKit.max_tool_executions_by_name,
+          max_spend: agent.max_spend || TurnKit.max_spend
+        }
+      end
+
       def aggregate_cost(current, cost)
         return cost unless current