magi-code 0.61.0

Repository-aware CLI coding agent for terminal work
Documentation
use super::super::tokens::{TokenEncodingFamily, token_encoding_family_for_model};
use super::*;

#[test]
fn token_estimator_uses_usage_or_chars_fallback() {
    assert_eq!(estimate_text_tokens("abcd"), 1);
    assert_eq!(estimate_text_tokens("abcde"), 2);
    assert_eq!(
        usage_input_tokens(&Usage {
            input: 42,
            ..Usage::default()
        }),
        42
    );
}

#[test]
fn embedded_tokenizers_count_known_strings_with_ordinary_encoding() {
    let cases = [
        ("gpt-4o", "hello world", 2),
        ("gpt-4", "hello world", 2),
        ("gpt-4o", "お誕生日おめでとう", 8),
        ("gpt-4", "お誕生日おめでとう", 9),
        ("gpt-4o", "<|endoftext|>", 7),
        ("gpt-4", "<|endoftext|>", 7),
    ];

    for (model, text, expected_tokens) in cases {
        let projected = project_text_tokens(crate::providers::OPENAI_CODEX_PROVIDER, model, text);
        assert_eq!(projected.source, ContextUsageSource::TokenizerProjection);
        assert_eq!(projected.tokens, expected_tokens, "{model} {text:?}");
    }
}

#[test]
fn embedded_tokenizers_count_empty_string_as_zero_tokens() {
    for model in ["gpt-4o", "gpt-4"] {
        let projected = project_text_tokens(crate::providers::OPENAI_CODEX_PROVIDER, model, "");
        assert_eq!(projected.source, ContextUsageSource::TokenizerProjection);
        assert_eq!(projected.tokens, 0, "{model}");
    }
}

#[test]
fn token_family_mapping_uses_o200k_precedence_before_cl100k() {
    assert_eq!(
        token_encoding_family_for_model("gpt-4o"),
        Some(TokenEncodingFamily::O200KBase)
    );
    assert_eq!(
        token_encoding_family_for_model("gpt-4"),
        Some(TokenEncodingFamily::Cl100KBase)
    );
    assert_eq!(
        token_encoding_family_for_model("GPT-5.5"),
        Some(TokenEncodingFamily::O200KBase)
    );
    assert_eq!(token_encoding_family_for_model("unknown-model"), None);
}

#[test]
fn context_tokenizer_openai_codex_model_uses_tiktoken() {
    let request = ProviderRequest::from_conversation(
        "gpt-5.5",
        vec![ProviderConversationItem::Message(ChatMessage::user(
            "hello world from tokenizer",
        ))],
    );

    let projected =
        project_provider_request_input_tokens(crate::providers::OPENAI_CODEX_PROVIDER, &request);

    assert_eq!(projected.source, ContextUsageSource::TokenizerEstimate);
    assert!(projected.tokens > 0);
    assert_ne!(
        projected.tokens,
        estimate_provider_request_input_tokens(&request)
    );
}

#[test]
fn context_tokenizer_unknown_provider_uses_fallback_estimate() {
    let request = ProviderRequest::from_conversation(
        "unknown-model",
        vec![ProviderConversationItem::Message(ChatMessage::user(
            "hello",
        ))],
    );

    let projected = project_provider_request_input_tokens("local-ai", &request);

    assert_eq!(projected.source, ContextUsageSource::FallbackEstimate);
    assert_eq!(
        projected.tokens,
        estimate_provider_request_input_tokens(&request)
    );
}

#[test]
fn context_tokenizer_counts_request_item_variants_without_panic() {
    let request = ProviderRequest::from_conversation(
        "gpt-5.5",
        vec![
            ProviderConversationItem::Message(ChatMessage::system("system")),
            ProviderConversationItem::ResponseItem(json!({
                "type": "function_call",
                "call_id": "call_1",
                "name": "read",
                "arguments": "{}"
            })),
            ProviderConversationItem::ToolResult(ProviderToolResult {
                call_id: "call_1".to_string(),
                tool_name: "read".to_string(),
                success: true,
                output: "file text".to_string(),
            }),
            ProviderConversationItem::LegacyReplayNote {
                event_type: "legacy".to_string(),
                content: "payload".to_string(),
            },
        ],
    );

    let projected =
        project_provider_request_input_tokens(crate::providers::OPENAI_CODEX_PROVIDER, &request);
    let text_projection = project_text_tokens(
        crate::providers::OPENAI_CODEX_PROVIDER,
        "gpt-5.5",
        "assistant output",
    );

    assert_eq!(projected.source, ContextUsageSource::TokenizerEstimate);
    assert!(projected.tokens > 0);
    assert_eq!(
        text_projection.source,
        ContextUsageSource::TokenizerProjection
    );
    assert!(text_projection.tokens > 0);
}

#[test]
fn token_estimators_traverse_all_conversation_item_variants() {
    let items = branch_complete_token_fixture();
    let request = ProviderRequest::from_conversation("gpt-5.5", items.clone());

    let fallback = project_provider_request_input_tokens("local-ai", &request);
    let tokenizer =
        project_provider_request_input_tokens(crate::providers::OPENAI_CODEX_PROVIDER, &request);

    assert_eq!(fallback.source, ContextUsageSource::FallbackEstimate);
    assert_eq!(fallback.tokens, manual_old_fallback_tokens(&items));
    assert_eq!(
        fallback.tokens,
        estimate_provider_request_input_tokens(&request)
    );
    assert_eq!(tokenizer.source, ContextUsageSource::TokenizerEstimate);
    assert!(tokenizer.tokens > 0);
}

fn branch_complete_token_fixture() -> Vec<ProviderConversationItem> {
    vec![
        ProviderConversationItem::Message(ChatMessage::user("message content")),
        ProviderConversationItem::ToolResult(ProviderToolResult {
            call_id: "call_tool".to_string(),
            tool_name: "read".to_string(),
            success: true,
            output: "tool output".to_string(),
        }),
        ProviderConversationItem::LegacyReplayNote {
            event_type: "legacy_note".to_string(),
            content: "legacy content".to_string(),
        },
        ProviderConversationItem::ResponseItem(json!({
            "type": "function_call",
            "call_id": "call_1",
            "name": "bash",
            "arguments": {"cmd": "echo hi", "flags": ["-n"]}
        })),
        ProviderConversationItem::ResponseItem(json!({
            "type": "function_call_output",
            "call_id": "call_1",
            "output": {"text": "done", "code": 0}
        })),
        ProviderConversationItem::ResponseItem(json!({
            "type": "reasoning",
            "summary": [{"text": "summary"}],
            "content": "reasoning content"
        })),
        ProviderConversationItem::ResponseItem(json!({
            "role": "assistant",
            "content": [{"type": "output_text", "text": "hello"}],
            "tool_calls": [{"id": "call_2", "name": "read"}],
            "tool_call_id": "call_2"
        })),
        ProviderConversationItem::ResponseItem(json!({
            "opaque": {"nested": [true, null, 7]}
        })),
    ]
}

fn manual_old_fallback_tokens(items: &[ProviderConversationItem]) -> usize {
    items
        .iter()
        .map(|item| match item {
            ProviderConversationItem::Message(message) => {
                estimate_text_tokens(&message.content) + 4
            }
            ProviderConversationItem::ResponseItem(value) => manual_response_item_tokens(value) + 4,
            ProviderConversationItem::ToolResult(result) => {
                estimate_text_tokens(&result.call_id)
                    + estimate_text_tokens(&result.tool_name)
                    + estimate_text_tokens(&result.output)
                    + 4
            }
            ProviderConversationItem::LegacyReplayNote {
                event_type,
                content,
            } => estimate_text_tokens(event_type) + estimate_text_tokens(content) + 4,
        })
        .sum()
}

fn manual_json_value_tokens(value: &Value) -> usize {
    match value {
        Value::String(text) => estimate_text_tokens(text),
        Value::Array(items) => items
            .iter()
            .map(manual_json_value_tokens)
            .sum::<usize>()
            .max(1),
        Value::Object(fields) => fields
            .iter()
            .map(|(key, value)| estimate_text_tokens(key) + manual_json_value_tokens(value))
            .sum::<usize>()
            .max(1),
        Value::Null => 1,
        other => estimate_text_tokens(&other.to_string()),
    }
}

fn manual_response_item_tokens(item: &Value) -> usize {
    match item.get("type").and_then(Value::as_str) {
        Some("function_call") => {
            estimate_text_tokens("function_call")
                + item
                    .get("call_id")
                    .and_then(Value::as_str)
                    .map(estimate_text_tokens)
                    .unwrap_or(0)
                + item
                    .get("name")
                    .and_then(Value::as_str)
                    .map(estimate_text_tokens)
                    .unwrap_or(0)
                + item
                    .get("arguments")
                    .map(manual_json_value_tokens)
                    .unwrap_or(0)
        }
        Some("function_call_output") => {
            estimate_text_tokens("function_call_output")
                + item
                    .get("call_id")
                    .and_then(Value::as_str)
                    .map(estimate_text_tokens)
                    .unwrap_or(0)
                + item
                    .get("output")
                    .map(manual_json_value_tokens)
                    .unwrap_or(0)
        }
        Some("reasoning") => {
            estimate_text_tokens("reasoning")
                + item
                    .get("summary")
                    .map(manual_json_value_tokens)
                    .unwrap_or(0)
                + item
                    .get("content")
                    .map(manual_json_value_tokens)
                    .unwrap_or(0)
        }
        _ => {
            if let Some(role) = item.get("role").and_then(Value::as_str) {
                estimate_text_tokens(role)
                    + item
                        .get("content")
                        .map(manual_json_value_tokens)
                        .unwrap_or(0)
                    + item
                        .get("tool_calls")
                        .map(manual_json_value_tokens)
                        .unwrap_or(0)
                    + item
                        .get("tool_call_id")
                        .and_then(Value::as_str)
                        .map(estimate_text_tokens)
                        .unwrap_or(0)
            } else {
                estimate_text_tokens(&item.to_string())
            }
        }
    }
}