ask-agent 0.28.0 → 0.29.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 36e26d74ffa1c3e3e97b8e172f4fef2ada07c1c194488a8f550e8b77b7473ee1
-  data.tar.gz: 82568a4cc3ed19617293416d10f2f3194c1a18079a1426431a2d878afee75fa0
+  metadata.gz: 7b6eca46c976555990a13366053665b9340f4a97d578a332a8cbf9615cd5647d
+  data.tar.gz: 7c16e54008a44cc0eafb8ebfea51c752e0d6879749635e606142f84cdd03d682
 SHA512:
-  metadata.gz: cfd202f83bb95e3349282c3771e591f93062976b58c964d65e928e09a6408818ef969200c10b68c4a307e0c0d9fabde112aaf20b9776110a5040590370ecd317
-  data.tar.gz: 7d3ae82b872659e9e27217bc98cf7dfffee28f5209ee5f351dcc944564a63a777888fdad780de218956019d38f223996ebd6372c7be84cf33fdbf9a1eeb4e5a7
+  metadata.gz: 05eb5c72881fcb3730dab2f89d51444b3f86f41413f454090331d5c9bb240b6ca4751c40f0fdc3ef10f2ea718989aa7e68e716e13570e247a094f76ae0702051
+  data.tar.gz: 712139e5f89850fe706268b74a08854c2f1f0be7276e4717039c73192a872d6f4a3c2566954ae4133ce10917d51ea4ae5687995b828ef024bead95a8222ac68d

data/CHANGELOG.md

@@ -1,3 +1,36 @@
+## [0.29.1] — 2026-08-06
+
+### Fixed
+
+- **Streamed token accounting now counts real tokens.** Stream usage arrives
+  as OpenAI-style `prompt_tokens`/`completion_tokens` (deepseek, openai,
+  most OpenAI-compatible providers), but `accumulated_tokens` read only
+  `input_tokens`/`output_tokens` — every streamed call reported 0 input and
+  ~1 output token (the content-chunk fallback), so token billing, cost
+  calculation, and usage metrics under-counted by orders of magnitude. Both
+  key shapes are read now, and the content-chunk fallback applies only when
+  the stream carries no usage at all (no double counting).
+
+## [0.29.0] — 2026-08-06
+
+### Added
+
+- **Tool-call repair** — malformed tool calls get one internal LLM round-trip
+  to fix them before execution. When a model emits a call with unparseable
+  arguments or an unknown tool name, the loop asks the model to re-emit it
+  corrected and executes the corrected version instead of burning a turn on
+  the error. Enable with `Session.new(tool_call_repair: true)` (built-in
+  repair prompt) or pass a callable for full control:
+  `Session.new(tool_call_repair: ->(chat, calls, tools) { ... })`.
+  - Corrections are remapped to the original call ids, so tool results stay
+    consistent with the conversation history; the internal repair exchange
+    is stripped from history.
+  - Calls the model cannot correct are dropped (the model saw them in the
+    repair prompt); repair is best-effort — a failing round-trip drops the
+    malformed calls instead of failing the turn.
+  - `Events::ToolCallRepaired` fires with name, id, original and corrected
+    arguments.
+
 ## [0.28.0] — 2026-08-06
 
 ### Changed (breaking)

data/lib/ask/agent/chat.rb

@@ -353,13 +353,23 @@ module Ask
       def accumulated_tokens(stream)
         input = 0
         output = 0
+        content_chunks = 0
         stream.chunks.each do |chunk|
           if chunk.usage
-            input = chunk.usage[:input_tokens] || chunk.usage["input_tokens"] || input
-            output = chunk.usage[:output_tokens] || chunk.usage["output_tokens"] || output
+            # Streams carry OpenAI-style prompt_tokens/completion_tokens
+            # (some providers also send input_tokens/output_tokens); reading
+            # only the latter reported 0 in / ~1 out for every streamed call.
+            input = chunk.usage[:input_tokens] || chunk.usage["input_tokens"] ||
+              chunk.usage[:prompt_tokens] || chunk.usage["prompt_tokens"] || input
+            output = chunk.usage[:output_tokens] || chunk.usage["output_tokens"] ||
+              chunk.usage[:completion_tokens] || chunk.usage["completion_tokens"] || output
           end
-          output += 1 if chunk.content.to_s.length > 0
+          content_chunks += 1 if chunk.content.to_s.length > 0
         end
+        # No usage in the stream at all: approximate output as content
+        # chunks rather than reporting nothing. Never mixed with real
+        # usage (that would double count).
+        output = content_chunks if input.zero? && output.zero? && content_chunks.positive?
         { input: input, output: output }
       end
 

data/lib/ask/agent/checkpoint_store.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module Ask
+  module Agent
+    # Versioned, durable session checkpoints on any {Ask::State::Adapter}.
+    #
+    # Every checkpoint is a full session snapshot (the same payload
+    # {Session#save} writes) stored under a sequential key. The adapter only
+    # needs the minimal KV contract — +get+, +set+, +delete+ — so this works
+    # with every state provider (SQLite, Redis, Postgres, MySQL) and with
+    # custom adapters that implement nothing else. No list primitives are
+    # required.
+    #
+    # Keys (session id + suffix):
+    #   "<id>:checkpoint:<seq>"  — one key per checkpoint
+    #   "<id>:checkpoint:head"   — the current seq (the active timeline)
+    #
+    # Rolling back moves the head pointer; later checkpoints are kept, so a
+    # session can roll forward again (time travel). Forking copies the
+    # checkpoints up to a point into a new session id — the branch diverges
+    # from there.
+    class CheckpointStore
+      CHECKPOINT_KEY = ":checkpoint:"
+      HEAD_KEY = ":checkpoint:head"
+
+      # @param state_adapter [Ask::State::Adapter] backing store
+      def initialize(state_adapter)
+        @state = state_adapter
+      end
+
+      # @return [Ask::State::Adapter] the underlying adapter
+      attr_reader :state
+
+      # Append a checkpoint. The new checkpoint becomes the head.
+      #
+      # @param session_id [String]
+      # @param data [Hash] session snapshot
+      # @return [Integer] the new checkpoint seq
+      def checkpoint(session_id, data)
+        seq = (head(session_id) || 0) + 1
+        @state.set(checkpoint_key(session_id, seq), data)
+        @state.set(head_key(session_id), seq)
+        seq
+      end
+
+      # @param session_id [String]
+      # @return [Integer, nil] current head seq, or nil when the session has
+      #   no checkpoints
+      def head(session_id)
+        @state.get(head_key(session_id))
+      end
+
+      # @param session_id [String]
+      # @return [Array<Integer>] all checkpoint seqs, oldest first
+      def history(session_id)
+        current = head(session_id)
+        current ? (1..current).to_a : []
+      end
+
+      # Load a checkpoint's data.
+      #
+      # @param session_id [String]
+      # @param seq [Integer, nil] checkpoint seq; defaults to the head
+      # @return [Hash, nil] the snapshot, or nil when it does not exist
+      def load(session_id, seq: nil)
+        seq ||= head(session_id)
+        return nil unless seq
+
+        @state.get(checkpoint_key(session_id, seq))
+      end
+
+      # Move the head pointer to an earlier (or later) checkpoint. Later
+      # checkpoints are kept so the session can roll forward again.
+      #
+      # @param session_id [String]
+      # @param seq [Integer]
+      # @return [Integer] the seq rolled back to
+      # @raise [ArgumentError] when the checkpoint does not exist
+      def rollback(session_id, seq)
+        raise ArgumentError, "no checkpoint #{seq}" unless @state.get(checkpoint_key(session_id, seq))
+
+        @state.set(head_key(session_id), seq)
+        seq
+      end
+
+      # Copy the checkpoints up to +at_seq+ into a new session id — a branch
+      # that diverges from that point.
+      #
+      # @param session_id [String]
+      # @param at_seq [Integer, nil] checkpoint to fork from; defaults to the
+      #   head
+      # @param new_id [String] id for the forked session (defaults to a new
+      #   uuid)
+      # @return [String] the forked session's id
+      # @raise [ArgumentError] when the checkpoint does not exist
+      def fork(session_id, at_seq: nil, new_id: SecureRandom.uuid)
+        at_seq ||= head(session_id)
+        raise ArgumentError, "no checkpoint #{at_seq}" unless at_seq && @state.get(checkpoint_key(session_id, at_seq))
+
+        (1..at_seq).each do |seq|
+          data = @state.get(checkpoint_key(session_id, seq))
+          @state.set(checkpoint_key(new_id, seq), data)
+        end
+        @state.set(head_key(new_id), at_seq)
+        new_id
+      end
+
+      # Remove every checkpoint for a session.
+      #
+      # @param session_id [String]
+      # @return [void]
+      def delete(session_id)
+        history(session_id).each do |seq|
+          @state.delete(checkpoint_key(session_id, seq))
+        end
+        @state.delete(head_key(session_id))
+        nil
+      end
+
+      private
+
+      def checkpoint_key(session_id, seq)
+        "#{session_id}#{CHECKPOINT_KEY}#{seq}"
+      end
+
+      def head_key(session_id)
+        "#{session_id}#{HEAD_KEY}"
+      end
+    end
+  end
+end

data/lib/ask/agent/events.rb

@@ -16,6 +16,9 @@ module Ask
       MessageEnd = Data.define(:tool_calls)
 
       ToolExecutionStart = Data.define(:name, :arguments, :id)
+      # Emitted when a malformed tool call (unparseable arguments or unknown
+      # tool) was repaired by the model before execution.
+      ToolCallRepaired = Data.define(:name, :id, :original_arguments, :corrected_arguments)
       # Emitted when a tool returns Ask::Result.pending (async work started).
       ToolPending = Data.define(:name, :id)
       # Emitted when a pending (async) tool's background work completes.
@@ -26,6 +29,11 @@ module Ask
       CompactionStart = Data.define(:tokens_before, :reason)
       CompactionEnd = Data.define(:tokens_before, :tokens_after, :summary)
 
+      # Emitted when a session was rewound to an earlier checkpoint.
+      SessionRolledBack = Data.define(:session_id, :seq, :turn_count)
+      # Emitted when a session was forked from a checkpoint.
+      SessionForked = Data.define(:session_id, :forked_id, :seq)
+
       LoopDetected = Data.define(:tool_name, :repeated_count)
       MaxTurnsExceeded = Data.define(:max_turns)
 

data/lib/ask/agent/loop.rb

@@ -17,7 +17,7 @@ module Ask
       @max_consecutive_tool_turns = max_consecutive_tool_turns
       end
 
-      def run_turn(chat:, message:, tools:, tool_executor:, compactor:, hooks:, event_emitter:, session_id: nil, persist: nil)
+      def run_turn(chat:, message:, tools:, tool_executor:, compactor:, hooks:, event_emitter:, session_id: nil, persist: nil, tool_call_repair: nil)
         raise MaxTurnsExceeded if @turn_count >= @max_turns
 
         event_emitter.emit(Events::TurnStart.new)
@@ -86,6 +86,10 @@ module Ask
         user_tool_calls = response.tool_calls.reject { |id, _| provider_results.key?(id) }
 
         if user_tool_calls.any?
+          # Repair malformed calls (unparseable arguments, unknown tool
+          # names) with one internal LLM round-trip before executing.
+          user_tool_calls = repair_malformed_calls(user_tool_calls, tools, chat, event_emitter, tool_call_repair)
+
           # Execute user tool calls locally
           # Respect the session's parallel_tools setting: parallel tools run
           # in threads (with the caller's thread-local context inherited);
@@ -162,7 +166,8 @@ module Ask
           hooks: hooks,
           event_emitter: event_emitter,
           session_id: session_id,
-          persist: persist
+          persist: persist,
+          tool_call_repair: tool_call_repair
         )
       end
 
@@ -174,6 +179,43 @@ module Ask
 
       private
 
+      # Repair malformed tool calls before execution. Calls with corrected
+      # versions execute in place of the originals (same ids); calls the
+      # model could not correct are dropped — the model saw them in the
+      # repair prompt. Best-effort: a failing repair round-trip drops the
+      # malformed calls instead of failing the turn.
+      def repair_malformed_calls(user_tool_calls, tools, chat, event_emitter, tool_call_repair)
+        return user_tool_calls if tool_call_repair.nil? || user_tool_calls.empty?
+
+        repairer = if tool_call_repair == true
+          @tool_call_repair ||= ToolCallRepair.new
+        elsif tool_call_repair.respond_to?(:call)
+          ToolCallRepair.new(tool_call_repair)
+        end
+        return user_tool_calls unless repairer
+
+        malformed, _valid = user_tool_calls.partition do |_id, tc|
+          ToolCallRepair.repair_info(tc, tools)
+        end
+        return user_tool_calls if malformed.empty?
+
+        corrections = repairer.call(chat: chat, calls: malformed.to_h, tools: tools)
+
+        corrections.each do |id, corrected|
+          event_emitter.emit(Events::ToolCallRepaired.new(
+            name: corrected.name,
+            id: id,
+            original_arguments: user_tool_calls[id].arguments,
+            corrected_arguments: corrected.arguments
+          ))
+          user_tool_calls[id] = corrected
+        end
+
+        malformed_ids = malformed.map(&:first)
+        user_tool_calls.reject! { |id, _| malformed_ids.include?(id) && !corrections.key?(id) }
+        user_tool_calls
+      end
+
       # Whether the session asked the loop to stop (barge-in). Emitters that
       # don't support aborting (plain stubs) are treated as never aborted.
       def aborted?(event_emitter)

data/lib/ask/agent/session.rb

@@ -22,7 +22,8 @@ module Ask
                      id: nil, system_prompt: nil, parallel_tools: true,
                      reflector: nil, telemetry: true, meta_agent: nil,
                      agent_dir: nil, evaluator: nil, audit_log: nil,
-                     skills_disclosure: true, approval: nil, **chat_options)
+                     skills_disclosure: true, approval: nil,
+                     tool_call_repair: nil, checkpoints: false, **chat_options)
         @id = id || SecureRandom.uuid
         @agent_dir = agent_dir
         @max_turns = max_turns
@@ -54,11 +55,17 @@ module Ask
         @hooks = Hooks.new(hooks)
         @audit_log = build_audit_log(audit_log)
         @approval_queue = build_approval(approval)
+        @tool_call_repair = tool_call_repair
 
         @system_context = build_system_context(system_prompt)
         apply_system_context
 
         @state = state || persistence
+        if checkpoints && !@state
+          raise ArgumentError, "checkpoints: requires a state: adapter"
+        end
+        @checkpoints = !!checkpoints
+        @checkpoint_store = CheckpointStore.new(@state) if @checkpoints
 
         reflector_opts = reflector.is_a?(Hash) ? reflector : {}
         @reflector = if reflector
@@ -129,6 +136,7 @@ module Ask
 	            hooks: @hooks,
 	            event_emitter: self,
 	            session_id: @id,
+            tool_call_repair: @tool_call_repair,
 	            persist: @state ? method(:persist!) : nil
 	          )
 
@@ -202,7 +210,8 @@ module Ask
               compactor: @compactor,
               hooks: @hooks,
               event_emitter: self,
-              session_id: @id
+              session_id: @id,
+              tool_call_repair: @tool_call_repair
             )
 
             @total_input_tokens += @loop.last_input_tokens.to_i
@@ -241,7 +250,8 @@ module Ask
               compactor: @compactor,
               hooks: @hooks,
               event_emitter: self,
-              session_id: @id
+              session_id: @id,
+              tool_call_repair: @tool_call_repair
             )
 
             @total_input_tokens += @loop.last_input_tokens.to_i
@@ -304,8 +314,21 @@ module Ask
         session = new(
           id: data[:id],
           model: data.dig(:metadata, :model),
-          tools: data.dig(:metadata, :tools)&.map(&:constantize) || [],
-          state: adapter
+          # Instantiate saved tool classes defensively: tools that cannot be
+          # auto-constructed (e.g. the built-in LoadSkillTool, which needs a
+          # registry) are skipped — resolve_tools re-adds them with a proper
+          # registry. Callers can also pass their own `tools:` after load.
+          tools: data.dig(:metadata, :tools).to_a.filter_map do |name|
+            begin
+              name.constantize.new
+            rescue StandardError
+              nil
+            end
+          end,
+          state: adapter,
+          # Checkpointing is restored automatically when the session has
+          # checkpoints in the store.
+          checkpoints: !adapter.get("#{id}#{CheckpointStore::HEAD_KEY}").nil?
         )
 
 	        data[:messages].each do |msg|
@@ -322,6 +345,7 @@ module Ask
 
       def delete
         @deleted = true
+        @checkpoint_store&.delete(@id)
         @state&.delete(@id)
       end
 
@@ -331,6 +355,81 @@ module Ask
 
       def abort_requested? = @abort_requested
 
+      # --- Checkpoints (fork, rollback, resume) ---
+
+      # @return [Array<Integer>] checkpoint seqs, oldest first
+      # @raise [RuntimeError] when checkpointing is not enabled
+      def checkpoint_history
+        require_checkpoints!
+        @checkpoint_store.history(@id)
+      end
+
+      # Load a checkpoint's snapshot.
+      #
+      # @param seq [Integer, nil] checkpoint seq; defaults to the head
+      # @return [Hash, nil] the snapshot with symbol keys
+      # @raise [RuntimeError] when checkpointing is not enabled
+      def load_checkpoint(seq: nil)
+        require_checkpoints!
+        data = @checkpoint_store.load(@id, seq: seq)
+        data && self.class.deep_symbolize_keys(data)
+      end
+
+      # Rewind the session to an earlier checkpoint: messages and turn count
+      # are restored from the snapshot, and the store's head moves back.
+      # Later checkpoints are kept, so the session can roll forward again.
+      #
+      # @param seq [Integer, nil] checkpoint seq (xor +turn:)
+      # @param turn [Integer, nil] roll back to the last checkpoint whose
+      #   turn count equals +turn+ (xor +seq:)
+      # @return [self]
+      # @raise [ArgumentError] when the checkpoint does not exist
+      # @raise [RuntimeError] when checkpointing is not enabled or the
+      #   session is running
+      def rollback!(seq: nil, turn: nil)
+        require_checkpoints!
+        raise "cannot roll back a running session" if @running
+
+        seq = resolve_checkpoint_seq(seq, turn)
+        data = load_checkpoint(seq: seq)
+        raise ArgumentError, "no checkpoint #{seq}" unless data
+
+        @checkpoint_store.rollback(@id, seq)
+        restore_from_snapshot(data)
+        emit(Events::SessionRolledBack.new(session_id: @id, seq: seq, turn_count: @turn_count))
+        self
+      end
+
+      # Fork the session at a checkpoint: a new session (new id, same model
+      # and tools) whose history is everything up to that point, backed by
+      # its own checkpoint chain. Continue the branch with +run+.
+      #
+      # @param at_seq [Integer, nil] checkpoint to fork from (xor +at_turn:)
+      # @param at_turn [Integer, nil] fork at the last checkpoint whose turn
+      #   count equals +at_turn+ (xor +at_seq:)
+      # @return [Ask::Agent::Session] the forked session
+      # @raise [ArgumentError] when the checkpoint does not exist
+      # @raise [RuntimeError] when checkpointing is not enabled
+      def fork(at_seq: nil, at_turn: nil)
+        require_checkpoints!
+
+        seq = resolve_checkpoint_seq(at_seq, at_turn)
+        data = load_checkpoint(seq: seq)
+        raise ArgumentError, "no checkpoint #{seq}" unless data
+
+        forked_id = @checkpoint_store.fork(@id, at_seq: seq)
+        forked = self.class.new(
+          id: forked_id,
+          model: data[:metadata][:model],
+          tools: @tools,
+          state: @state,
+          checkpoints: true
+        )
+        restore_into(forked, data)
+        emit(Events::SessionForked.new(session_id: @id, forked_id: forked_id, seq: seq))
+        forked
+      end
+
       # --- Async (pending) tools ---
 
       # Registers a pending tool call (called by the loop when a tool
@@ -551,8 +650,52 @@ module Ask
         compactor
       end
 
+      def require_checkpoints!
+        raise "checkpointing is not enabled (pass state: and checkpoints: true)" unless @checkpoints
+      end
+
+      # Resolve a checkpoint seq from either an explicit seq or a turn
+      # count (the last checkpoint whose metadata turn_count matches).
+      def resolve_checkpoint_seq(seq, turn)
+        raise ArgumentError, "pass either seq: or turn:, not both" if seq && turn
+
+        if seq
+          seq
+        elsif turn
+          found = checkpoint_history.reverse_each.find do |s|
+            data = self.class.deep_symbolize_keys(@checkpoint_store.load(@id, seq: s) || {})
+            data.dig(:metadata, :turn_count) == turn
+          end
+          raise ArgumentError, "no checkpoint at turn #{turn}" unless found
+
+          found
+        else
+          raise ArgumentError, "pass either seq: or turn:"
+        end
+      end
+
+      # Replace the session's in-memory state with a snapshot (symbol keys)
+      # and keep the legacy blob consistent.
+      def restore_from_snapshot(data)
+        restore_into(self, data)
+        @state.set(@id, data)
+      end
+
+      def restore_into(target, data)
+        target.chat.reset_messages!
+        data[:messages].each do |msg|
+          target.chat.add_message(
+            role: msg[:role].to_sym,
+            content: msg[:content],
+            tool_call_id: msg[:tool_call_id]
+          )
+        end
+        target.instance_variable_set(:@messages, target.chat.messages.dup)
+        target.instance_variable_set(:@turn_count, data.dig(:metadata, :turn_count) || 0)
+      end
+
       def persist!
-        @state.set(@id, {
+        payload = {
           id: @id,
           messages: @chat.messages.map { |m|
             {
@@ -570,7 +713,9 @@ module Ask
             created_at: @created_at.iso8601,
             updated_at: Time.now.iso8601
           }
-        })
+        }
+        @state.set(@id, payload)
+        @checkpoint_store.checkpoint(@id, payload) if @checkpoints
       end
 
       def try_auto_meta_agent

data/lib/ask/agent/tool_call_repair.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Ask
+  module Agent
+    # Repairs malformed tool calls before they execute.
+    #
+    # When a model emits a tool call with unparseable arguments or an unknown
+    # tool name, execution currently fails and the model burns a turn seeing
+    # the error. Repair intercepts those calls before execution, asks the
+    # model to re-emit them corrected (one internal LLM round-trip), and
+    # executes the corrected versions — remapped to the original call ids so
+    # tool results stay consistent with the conversation history.
+    #
+    # Enable with `Session.new(tool_call_repair: true)` for the built-in
+    # repair prompt, or pass a callable for full control:
+    #
+    #   Session.new(tool_call_repair: ->(chat, calls, tools) {
+    #     { calls.keys.first => ToolCallInfo.new(
+    #         id: calls.keys.first, name: "bash", arguments: "{}"
+    #       ) }
+    #   })
+    #
+    # The callable receives the chat, the malformed calls hash (id =>
+    # ToolCallInfo), and the available tools; it returns a hash of
+    # corrections keyed by the ORIGINAL call ids. Calls without a correction
+    # are dropped — the model saw them in the prompt and declined to fix
+    # them.
+    class ToolCallRepair
+      # @param tool_call [ToolCallInfo]
+      # @param tools [Array<Object>]
+      # @return [String, nil] why the call is repairable, or nil when it is
+      #   well-formed and executable
+      def self.repair_info(tool_call, tools)
+        unless tools.any? { |t| t.respond_to?(:name) && t.name == tool_call.name }
+          return "unknown tool '#{tool_call.name}'"
+        end
+
+        args = tool_call.arguments
+        return nil if args.is_a?(Hash)
+
+        parsed = JSON.parse(args.to_s)
+        return nil if parsed.is_a?(Hash)
+
+        "arguments must be a JSON object, got #{parsed.class}"
+      rescue JSON::ParserError => e
+        "arguments are not valid JSON: #{e.message}"
+      end
+
+      # @param callable [Proc, nil] custom repair function; nil uses the
+      #   built-in repair prompt
+      def initialize(callable = nil)
+        @callable = callable
+      end
+
+      # Ask the model to correct the malformed calls.
+      #
+      # @param chat [Ask::Agent::Chat] the session chat (history is restored
+      #   after the internal round-trip)
+      # @param calls [Hash{String => ToolCallInfo}] malformed calls by id
+      # @param tools [Array<Object>]
+      # @return [Hash{String => ToolCallInfo}] corrections keyed by the
+      #   original call ids
+      def call(chat:, calls:, tools:)
+        if @callable
+          normalize(@callable.call(chat, calls, tools))
+        else
+          built_in(chat, calls, tools)
+        end
+      end
+
+      private
+
+      def built_in(chat, calls, tools)
+        size = chat.messages.size
+        response = chat.ask(repair_prompt(calls, tools))
+        # Remove the internal repair exchange from the conversation so the
+        # history stays clean and the model never sees it.
+        chat.messages.slice!(size..)
+
+        corrections = {}
+        response.tool_calls.values.each_with_index do |tc, index|
+          original_id = calls.keys[index]
+          break unless original_id
+
+          corrections[original_id] = ToolCallInfo.new(
+            id: original_id, name: tc.name, arguments: tc.arguments
+          )
+        end
+        corrections
+      rescue StandardError
+        # Repair is best-effort: on any failure, drop the malformed calls
+        # rather than failing the turn. Restore history either way.
+        chat.messages.slice!(size..) rescue nil
+        {}
+      end
+
+      # Normalize a custom callable's result to the corrections contract
+      # (original id => ToolCallInfo), dropping anything malformed.
+      def normalize(result)
+        return {} unless result.respond_to?(:each)
+
+        result.each_with_object({}) do |(id, tc), acc|
+          next unless tc.respond_to?(:name)
+
+          acc[id] = ToolCallInfo.new(id: id, name: tc.name, arguments: tc.arguments)
+        end
+      end
+
+      def repair_prompt(calls, tools)
+        lines = calls.map.with_index do |(id, tc), i|
+          reason = self.class.repair_info(tc, tools) || "invalid"
+          "#{i + 1}. call id \"#{id}\", tool \"#{tc.name}\", " \
+            "arguments: #{tc.arguments.inspect} — #{reason}"
+        end
+        tool_list = tools.map { |t| t.respond_to?(:name) ? t.name : t.to_s }.join(", ")
+
+        <<~PROMPT.strip
+          Some tool calls from your last message were invalid and could not be executed:
+          #{lines.join("\n")}
+
+          Available tools: #{tool_list.empty? ? "(none)" : tool_list}
+
+          Reply by calling the tools again with corrected arguments, in the same order as listed above. Call ONLY the tools listed above that you can correct; if a tool truly does not exist or cannot be corrected, do not call it.
+        PROMPT
+      end
+    end
+  end
+end

data/lib/ask/agent/version.rb

@@ -2,6 +2,6 @@
 
 module Ask
   module Agent
-    VERSION = "0.28.0"
+    VERSION = "0.29.1"
   end
 end

data/lib/ask/agent.rb

@@ -37,6 +37,9 @@ module Ask
       autoload :ApprovalPolicy, "ask/agent/policies/approval_policy"
     end
 
+    autoload :ToolCallRepair, "ask/agent/tool_call_repair"
+    autoload :CheckpointStore, "ask/agent/checkpoint_store"
+
     module Middleware
       autoload :Base, "ask/agent/middleware/base"
       autoload :Pipeline, "ask/agent/middleware/pipeline"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ask-agent
 version: !ruby/object:Gem::Version
-  version: 0.28.0
+  version: 0.29.1
 platform: ruby
 authors:
 - Kaka Ruto
@@ -166,6 +166,7 @@ files:
 - lib/ask/agent.rb
 - lib/ask/agent/approval_queue.rb
 - lib/ask/agent/chat.rb
+- lib/ask/agent/checkpoint_store.rb
 - lib/ask/agent/cli.rb
 - lib/ask/agent/compactor.rb
 - lib/ask/agent/configuration.rb
@@ -204,6 +205,7 @@ files:
 - lib/ask/agent/telemetry.rb
 - lib/ask/agent/test.rb
 - lib/ask/agent/tool_abort_controller.rb
+- lib/ask/agent/tool_call_repair.rb
 - lib/ask/agent/tool_executor.rb
 - lib/ask/agent/version.rb
 homepage: https://github.com/ask-rb/ask-agent