Support Token Counter for custom models

What version of Codex is running?

OpenAI Codex TUI (v0.80.0)

What subscription do you have?

plus

Which model were you using?

local nemotron30b

What platform is your computer?

linux mint

What issue are you seeing?

Codex Local Token Counter Fix - Instructional Guide

Author: Claude (Anthropic) with Jeff Bulger
Date: 2026-01-14
For: OpenAI Codex TUI (v0.80.0) running with local llama.cpp backend
GitHub: Feel free to use and share - MIT License

---

🎯 What This Fixes

The stock OpenAI Codex TUI doesn't display token usage when using local llama.cpp. After this fix:

╭───────────────────────────────────────────────────────────────╮
│  Token usage:      34.5K total  (33K input + 1.45K output)    │
│  Context window:   93% left (28.2K used / 249K)               │
╰───────────────────────────────────────────────────────────────╯

Before: Always shows 0 total (0 input + 0 output)
After: Real-time token tracking + auto-compact when you hit your limit

---

📋 Prerequisites

• OpenAI Codex TUI (Rust version, v0.80.0 or similar)
• llama.cpp server running locally
• Using the Chat wire API (not Responses API)
• Rust toolchain for rebuilding

---

🔧 The Fix (4 Files to Modify)

File 1: codex-api/src/requests/chat.rs

What: Tell llama.cpp to include usage data in the streaming response.

Find the payload construction (around line 303-310):

let payload = json!({
    "model": self.model,
    "messages": messages,
    "stream": true,
    "tools": self.tools,
});

Replace with:

let payload = json!({
    "model": self.model,
    "messages": messages,
    "stream": true,
    "stream_options": {"include_usage": true},
    "tools": self.tools,
});

---

File 2: codex-api/src/sse/chat.rs

What: Extract the token usage from llama.cpp's response and pass it to the TUI.

Step 2a: Add Import

At the top of the file, add:

use codex_protocol::protocol::TokenUsage;

Step 2b: Add Tracking Variable

In the spawn_stream_handler function, find these lines:

let mut assistant_item: Option<ResponseItem> = None;
let mut reasoning_item: Option<ResponseItem> = None;
let mut completed_sent = false;

Add below them:

let mut captured_usage: Option<TokenUsage> = None;

Step 2c: Modify flush_and_complete Function

Find the flush_and_complete function and change it to accept and use token_usage:

async fn flush_and_complete(
    tx_event: &mpsc::Sender<Result<ResponseEvent, ApiError>>,
    reasoning_item: &mut Option<ResponseItem>,
    assistant_item: &mut Option<ResponseItem>,
    token_usage: Option<TokenUsage>,  // NEW PARAMETER
) {
    if let Some(reasoning) = reasoning_item.take() {
        let _ = tx_event
            .send(Ok(ResponseEvent::OutputItemDone(reasoning)))
            .await;
    }

    if let Some(assistant) = assistant_item.take() {
        let _ = tx_event
            .send(Ok(ResponseEvent::OutputItemDone(assistant)))
            .await;
    }

    let _ = tx_event
        .send(Ok(ResponseEvent::Completed {
            response_id: String::new(),
            token_usage,  // WAS: None - NOW: actual usage
        }))
        .await;
}

Step 2d: Extract Usage from JSON

After the JSON is parsed (let value: serde_json::Value = serde_json::from_str(&data)?;), add:

// Extract usage from final chunk if present
if let Some(usage) = value.get("usage") {
    if let (Some(prompt), Some(completion), Some(total)) = (
        usage.get("prompt_tokens").and_then(|v| v.as_i64()),
        usage.get("completion_tokens").and_then(|v| v.as_i64()),
        usage.get("total_tokens").and_then(|v| v.as_i64()),
    ) {
        captured_usage = Some(TokenUsage {
            input_tokens: prompt,
            output_tokens: completion,
            total_tokens: total,
            cached_input_tokens: 0,
            reasoning_output_tokens: 0,
        });
    }
}

Step 2e: Update All flush_and_complete Calls

Find all calls to flush_and_complete and add the captured_usage parameter:

Before:

flush_and_complete(&tx_event, &mut reasoning_item, &mut assistant_item).await;

After:

flush_and_complete(&tx_event, &mut reasoning_item, &mut assistant_item, captured_usage.clone()).await;

---

File 3: core/src/models_manager/model_info.rs

What: Set the correct context window for your model so the percentage display is accurate.

Find your model's definition (or add one for local models). Example for Nemotron:

} else if slug.contains("nemotron") || slug.contains("Nemotron") || slug.contains("30b") {
    model_info!(
        slug,
        base_instructions: NEMOTRON_INSTRUCTIONS.to_string(),
        apply_patch_tool_type: Some(ApplyPatchToolType::Function),
        context_window: Some(262_144),  // MATCH YOUR llama.cpp -c VALUE
        shell_type: ConfigShellToolType::ShellCommand,
    )
}

Important: Set context_window to match your llama.cpp server's -c flag!

---

File 4: config.toml

What: Enable auto-compaction to prevent OOM when you hit your context limit.

Add at the top level of your config:

model_auto_compact_token_limit = 250000

This triggers automatic context compaction at 250k tokens (adjust based on your context window).

---

🔨 Build & Test

cd codex-rs
cargo build

Then launch your TUI and run /status - you should see real token counts!

---

🔍 How It Works

┌─────────────────────────────────────────────────────────────────┐
│  Your Prompt                                                    │
│       ↓                                                         │
│  Codex TUI builds request with stream_options: {include_usage}  │
│       ↓                                                         │
│  llama.cpp processes and streams response                       │
│       ↓                                                         │
│  Final SSE chunk includes: {"usage": {prompt_tokens, ...}}      │
│       ↓                                                         │
│  Our fix extracts usage → TokenUsage struct                     │
│       ↓                                                         │
│  TUI displays: "Token usage: 34.5K total (33K in + 1.45K out)"  │
└─────────────────────────────────────────────────────────────────┘

---

⚠️ Troubleshooting

Still showing 0 tokens?

1. Make sure you rebuilt after changes: cargo build
2. Verify you're using Chat wire API (check wire_api in config)
3. Check llama.cpp logs - should show POST /v1/chat/completions

Context percentage wrong?

• Update context_window in model_info.rs to match your llama.cpp -c value

Auto-compact not triggering?

• Check model_auto_compact_token_limit is set in config.toml
• Value should be slightly below your context window

---

📊 Token Count Breakdown (For Reference)

Typical baseline overhead:
| Component | Tokens | Notes |
|-----------|--------|-------|
| System prompt | 1,400-5,200 | Depends on prompt file |
| Tools (MCP) | 2,000-3,000 | Depends on tools loaded |
| Grammar overhead | 1,200-1,500 | llama.cpp peg-constructed |
| Request formatting | 500-700 | JSON structure |

Optimal baseline: ~5,600 tokens with lightweight prompt
Heavy baseline: ~9,800 tokens with full prompt + features

---

🙏 Credits

Fix developed by Claude (Anthropic) working with Jeff Bulger on 2026-01-14.

This was a multi-hour debugging session tracing token flow from llama.cpp → SSE parser → TUI. The key insight: the Chat wire API (for local models) never implemented usage extraction, even though Responses API (cloud) had it.

The Problem: Hardcoded token_usage: None in the completion handler
The Fix: Extract usage from llama.cpp's JSON response and pass it through

---

📜 License

MIT - Use freely, attribution appreciated.

---

"When Nemo runs well, he runs GREAT. This config makes it happen." - Jeff

What steps can reproduce the bug?

Load and local model.

What is the expected behavior?

_No response_

Additional information

_No response_