ask-agent 0.29.0 → 0.29.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 11cf486fa2d955d92cc6f8268f30e5a418b6f0567cdfb6fe6eac015a3352c74b
-  data.tar.gz: 5b4990efc0b66de943f6e5c3b7940d31cca6295c54c40fafef485c6004433601
+  metadata.gz: 7b6eca46c976555990a13366053665b9340f4a97d578a332a8cbf9615cd5647d
+  data.tar.gz: 7c16e54008a44cc0eafb8ebfea51c752e0d6879749635e606142f84cdd03d682
 SHA512:
-  metadata.gz: 3adf87ccd9b0cbb606133cb91b8baf9e24a328f96e4ed8f77266258883b13663762cbd16a3e1e34475ec0b93ab763518c4b661fa675fba37985144c01435a736
-  data.tar.gz: b6093801427cf4c8bc56fd9b6b0bf6f9fd6bc5252ce336731149413984b5f83862e4cfeae2e9705b2517bb73e3063091950293f74cdb32109624965e52ccc760
+  metadata.gz: 05eb5c72881fcb3730dab2f89d51444b3f86f41413f454090331d5c9bb240b6ca4751c40f0fdc3ef10f2ea718989aa7e68e716e13570e247a094f76ae0702051
+  data.tar.gz: 712139e5f89850fe706268b74a08854c2f1f0be7276e4717039c73192a872d6f4a3c2566954ae4133ce10917d51ea4ae5687995b828ef024bead95a8222ac68d

data/CHANGELOG.md

@@ -1,3 +1,16 @@
+## [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

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

@@ -29,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/session.rb

@@ -23,7 +23,7 @@ module Ask
                      reflector: nil, telemetry: true, meta_agent: nil,
                      agent_dir: nil, evaluator: nil, audit_log: nil,
                      skills_disclosure: true, approval: nil,
-                     tool_call_repair: nil, **chat_options)
+                     tool_call_repair: nil, checkpoints: false, **chat_options)
         @id = id || SecureRandom.uuid
         @agent_dir = agent_dir
         @max_turns = max_turns
@@ -61,6 +61,11 @@ module Ask
         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
@@ -309,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|
@@ -327,6 +345,7 @@ module Ask
 
       def delete
         @deleted = true
+        @checkpoint_store&.delete(@id)
         @state&.delete(@id)
       end
 
@@ -336,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
@@ -556,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|
             {
@@ -575,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/version.rb

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

data/lib/ask/agent.rb

@@ -38,6 +38,7 @@ module Ask
     end
 
     autoload :ToolCallRepair, "ask/agent/tool_call_repair"
+    autoload :CheckpointStore, "ask/agent/checkpoint_store"
 
     module Middleware
       autoload :Base, "ask/agent/middleware/base"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ask-agent
 version: !ruby/object:Gem::Version
-  version: 0.29.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