[codex] Expose Codex rollout context snapshots

.github/workflows/ci.yml

@@ -7,6 +7,17 @@ on:
     branches: [main]
 
 jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: "3.3"
+          bundler-cache: true
+      - run: bundle exec rubocop
+
   test:
     runs-on: ubuntu-latest
     strategy:

.rubocop.yml

@@ -0,0 +1,38 @@
+AllCops:
+  NewCops: enable
+  SuggestExtensions: false
+  TargetRubyVersion: 3.2
+
+Metrics/BlockLength:
+  Exclude:
+    - "spec/**/*"
+
+Metrics/AbcSize:
+  Enabled: false
+
+Metrics/ClassLength:
+  Enabled: false
+
+Metrics/CyclomaticComplexity:
+  Enabled: false
+
+Metrics/MethodLength:
+  Enabled: false
+
+Metrics/PerceivedComplexity:
+  Enabled: false
+
+Style/StringLiterals:
+  Enabled: false
+
+Style/ArgumentsForwarding:
+  Enabled: false
+
+Style/ConcatArrayLiterals:
+  Enabled: false
+
+Style/Documentation:
+  Enabled: false
+
+Style/StringLiteralsInInterpolation:
+  Enabled: false

Gemfile

@@ -3,3 +3,7 @@
 source "https://rubygems.org"
 
 gemspec
+
+gem "rake", "~> 13.0"
+gem "rspec", "~> 3.0"
+gem "rubocop", "~> 1.72"

README.md

@@ -43,6 +43,7 @@ thread = client.start_thread(
 turn = thread.run("Explain this codebase")
 puts turn.final_response
 puts "Tokens used: #{turn.usage.input_tokens} in, #{turn.usage.output_tokens} out"
+puts "Final context: #{turn.context_snapshot.context_tokens} / #{turn.context_snapshot.model_context_window}"
 
 # Streaming run - yields events as they arrive
 thread.run_streamed("Fix the failing tests") do |event|
@@ -62,6 +63,8 @@ thread.run_streamed("Fix the failing tests") do |event|
     puts "Error: #{event.error_message}"
   end
 end
+
+puts "Final context: #{thread.context_snapshot.context_tokens} / #{thread.context_snapshot.model_context_window}"
 ```
 
 ### Resume a thread
@@ -120,6 +123,20 @@ client = CodexSDK::Client.new(
 | `Events::ItemCompleted` | Item finished, provides typed `item` |
 | `Events::Error` | Stream-level error, provides `message` |
 
+## Context snapshots
+
+Codex CLI writes richer rollout logs under `~/.codex/sessions` (or `CODEX_HOME/sessions`). After a run completes, `Thread#run` and `Thread#run_streamed` expose a final `context_snapshot` derived from the latest `token_count` entry in those rollout files.
+
+```ruby
+snapshot = thread.context_snapshot
+snapshot.context_tokens          # => current prompt/context footprint
+snapshot.model_context_window    # => model max context window
+snapshot.last_token_usage.total_tokens
+snapshot.total_token_usage.total_tokens
+```
+
+This is separate from `turn.completed.usage`, which is still the per-turn API usage reported by the JSON event stream.
+
 ## Item types
 
 | Item | Fields |

codex-ruby.gemspec

@@ -18,10 +18,8 @@ Gem::Specification.new do |spec|
   spec.metadata["homepage_uri"] = spec.homepage
   spec.metadata["source_code_uri"] = spec.homepage
   spec.metadata["changelog_uri"] = "#{spec.homepage}/blob/main/CHANGELOG.md"
+  spec.metadata["rubygems_mfa_required"] = "true"
 
   spec.files = Dir["lib/**/*.rb", "LICENSE.txt", "README.md", "CHANGELOG.md"]
   spec.require_paths = ["lib"]
-
-  spec.add_development_dependency "rake", "~> 13.0"
-  spec.add_development_dependency "rspec", "~> 3.0"
 end

lib/codex_sdk.rb

@@ -29,6 +29,7 @@ def initialize(message, line: nil)
 require_relative "codex_sdk/config_serializer"
 require_relative "codex_sdk/items"
 require_relative "codex_sdk/events"
+require_relative "codex_sdk/rollout_context_snapshot_reader"
 require_relative "codex_sdk/exec"
 require_relative "codex_sdk/agent_thread"
 require_relative "codex_sdk/client"

lib/codex_sdk/agent_thread.rb

@@ -4,7 +4,7 @@
 
 module CodexSDK
   class AgentThread
-    attr_reader :id
+    attr_reader :id, :context_snapshot
 
     def initialize(options, thread_options:, resume_id: nil)
       @options = options
@@ -33,17 +33,15 @@ def run(input, turn_options: TurnOptions.new)
         end
       end
 
-      Turn.new(items: items, final_response: final_response, usage: usage)
+      Turn.new(items: items, final_response: final_response, usage: usage, context_snapshot: @context_snapshot)
     end
 
     # Streaming run: yields each event to the block as it arrives.
     def run_streamed(input, turn_options: TurnOptions.new, &block)
       prompt = normalize_input(input)
 
       output_schema_path = nil
-      if turn_options.output_schema
-        output_schema_path = write_output_schema(turn_options.output_schema)
-      end
+      output_schema_path = write_output_schema(turn_options.output_schema) if turn_options.output_schema
 
       @exec = Exec.new(
         @options,
@@ -61,6 +59,7 @@ def run_streamed(input, turn_options: TurnOptions.new, &block)
         block.call(event)
       end
     ensure
+      @context_snapshot = @exec&.context_snapshot
       cleanup_output_schema(output_schema_path)
     end
 
@@ -76,9 +75,9 @@ def normalize_input(input)
       when String
         input
       when Array
-        input.filter_map { |entry|
+        input.filter_map do |entry|
           entry[:text] if entry[:type] == "text"
-        }.join("\n\n")
+        end.join("\n\n")
       else
         input.to_s
       end
@@ -93,6 +92,7 @@ def write_output_schema(schema)
 
     def cleanup_output_schema(path)
       return unless path
+
       dir = File.dirname(path)
       FileUtils.rm_rf(dir)
     rescue StandardError

lib/codex_sdk/config_serializer.rb

@@ -38,6 +38,7 @@ def to_toml_value(value)
         value.to_json
       when Integer, Float
         raise ArgumentError, "cannot serialize non-finite number" unless value.to_f.finite?
+
         value.to_s
       when true, false
         value.to_s

lib/codex_sdk/events.rb

@@ -19,9 +19,7 @@ def self.parse(data)
 
     ThreadStarted = Data.define(:thread_id)
 
-    TurnStarted = Data.define do
-      def initialize; super(); end
-    end
+    TurnStarted = Data.define
 
     TurnCompleted = Data.define(:usage) do
       def self.from_json(data)

lib/codex_sdk/exec.rb

@@ -10,7 +10,7 @@ module CodexSDK
   class Exec
     SHUTDOWN_TIMEOUT = 10 # seconds to wait after SIGTERM before SIGKILL
 
-    attr_reader :pid
+    attr_reader :pid, :context_snapshot
 
     def initialize(options, thread_options: ThreadOptions.new)
       @options = options
@@ -27,6 +27,9 @@ def initialize(options, thread_options: ThreadOptions.new)
     def run(prompt, resume_thread_id: nil, images: [], output_schema_path: nil, &block)
       args = build_args(resume_thread_id: resume_thread_id, images: images, output_schema_path: output_schema_path)
       env = build_env
+      sessions_root = codex_sessions_root(env)
+      started_at = Time.now
+      @context_snapshot = nil
 
       @stdin, @stdout, @stderr, @wait_thread = Open3.popen3(env, *args)
 
@@ -35,7 +38,11 @@ def run(prompt, resume_thread_id: nil, images: [], output_schema_path: nil, &blo
       @stdin.close
 
       # Read stderr in background thread
-      stderr_reader = ::Thread.new { @stderr.read rescue "" }
+      stderr_reader = ::Thread.new do
+        @stderr.read
+      rescue StandardError
+        ""
+      end
 
       # Read JSONL from stdout line by line
       @stdout.each_line do |line|
@@ -64,6 +71,11 @@ def run(prompt, resume_thread_id: nil, images: [], output_schema_path: nil, &blo
           stderr: stderr_buf
         )
       end
+
+      @context_snapshot = read_context_snapshot(
+        sessions_root: sessions_root,
+        started_at: started_at
+      )
     ensure
       cleanup
     end
@@ -122,9 +134,7 @@ def build_args(resume_thread_id: nil, images: [], output_schema_path: nil)
         args.concat(["--config", "sandbox_workspace_write.network_access=#{to.network_access}"])
       end
 
-      if to.web_search
-        args.concat(["--config", "web_search=#{ConfigSerializer.to_toml_value(to.web_search)}"])
-      end
+      args.concat(["--config", "web_search=#{ConfigSerializer.to_toml_value(to.web_search)}"]) if to.web_search
 
       if to.approval_policy
         args.concat(["--config", "approval_policy=#{ConfigSerializer.to_toml_value(to.approval_policy)}"])
@@ -153,15 +163,37 @@ def build_env
     def find_codex_path
       path = `which codex 2>/dev/null`.strip
       raise Error, "codex binary not found in PATH" if path.empty?
+
       path
     end
 
+    def codex_sessions_root(env)
+      return File.join(env["CODEX_HOME"], "sessions") if env["CODEX_HOME"] && !env["CODEX_HOME"].empty?
+
+      return unless env["HOME"] && !env["HOME"].empty?
+
+      File.join(env["HOME"], ".codex", "sessions")
+    end
+
+    def read_context_snapshot(sessions_root:, started_at:)
+      return unless sessions_root
+
+      RolloutContextSnapshotReader.new(
+        sessions_root: sessions_root,
+        started_at: started_at
+      ).read
+    rescue StandardError
+      nil
+    end
+
     def wait_for_exit(timeout)
       deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout
       loop do
         return true unless @wait_thread&.alive?
+
         remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
         return false if remaining <= 0
+
         sleep([0.1, remaining].min)
       end
     end

lib/codex_sdk/options.rb

@@ -63,9 +63,43 @@ def initialize(input_tokens: 0, cached_input_tokens: 0, output_tokens: 0)
     end
   end
 
+  # Detailed token usage from rollout token_count snapshots.
+  TokenUsage = Data.define(
+    :input_tokens,
+    :cached_input_tokens,
+    :output_tokens,
+    :reasoning_output_tokens,
+    :total_tokens
+  ) do
+    def initialize(
+      input_tokens: 0,
+      cached_input_tokens: 0,
+      output_tokens: 0,
+      reasoning_output_tokens: 0,
+      total_tokens: 0
+    )
+      super
+    end
+  end
+
+  # Final context snapshot derived from Codex rollout logs.
+  ContextSnapshot = Data.define(:model_context_window, :last_token_usage, :total_token_usage) do
+    def initialize(
+      model_context_window: 0,
+      last_token_usage: TokenUsage.new,
+      total_token_usage: TokenUsage.new
+    )
+      super
+    end
+
+    def context_tokens
+      last_token_usage.total_tokens
+    end
+  end
+
   # Result of a blocking Thread#run call.
-  Turn = Data.define(:items, :final_response, :usage) do
-    def initialize(items: [], final_response: "", usage: nil)
+  Turn = Data.define(:items, :final_response, :usage, :context_snapshot) do
+    def initialize(items: [], final_response: "", usage: nil, context_snapshot: nil)
       super
     end
   end