token-count 0.4.0

Count tokens for LLM models using exact tokenization
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160

//! Integration tests for Claude API mode
//!
//! Tests the --accurate flag behavior including error handling, consent requirements,
//! and fallback to estimation when API is unavailable.
//!
//! Note: These tests use a mutex to serialize access to environment variables,
//! preventing race conditions when tests run in parallel.

use assert_cmd::Command;
use predicates::prelude::*;
use std::sync::Mutex;

// Global mutex to serialize tests that manipulate environment variables
static ENV_MUTEX: Mutex<()> = Mutex::new(());

#[test]
fn test_api_mode_with_invalid_key() {
    let _lock = ENV_MUTEX.lock().unwrap();

    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.env("ANTHROPIC_API_KEY", "invalid-key-123")
        .arg("--model")
        .arg("claude")
        .arg("--accurate")
        .arg("-y")
        .write_stdin("test");

    // With invalid key, API call will fail and fall back to estimation
    // Should succeed with estimation and show warning in stderr
    let output = cmd.output().unwrap();

    // Should output token count (estimation fallback)
    let stdout = String::from_utf8(output.stdout).unwrap();
    assert!(!stdout.trim().is_empty(), "Should output token count");

    // Should have warning about API failure and fallback
    let stderr = String::from_utf8(output.stderr).unwrap();
    assert!(
        stderr.contains("falling back to estimation"),
        "Expected fallback warning, got: {}",
        stderr
    );
}

#[test]
fn test_api_mode_missing_key_verbose_error() {
    let _lock = ENV_MUTEX.lock().unwrap();

    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.env_remove("ANTHROPIC_API_KEY")
        .arg("--model")
        .arg("claude")
        .arg("--accurate")
        .arg("-y")
        .write_stdin("test");

    // Error should include setup instructions
    cmd.assert()
        .failure()
        .stderr(predicate::str::contains("export ANTHROPIC_API_KEY="))
        .stderr(predicate::str::contains("omit --accurate flag"));
}

#[test]
fn test_api_mode_non_interactive_without_yes() {
    let _lock = ENV_MUTEX.lock().unwrap();

    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.env_remove("ANTHROPIC_API_KEY")
        .arg("--model")
        .arg("claude")
        .arg("--accurate")
        .write_stdin("test");

    // Should fail in non-interactive mode without -y flag
    cmd.assert()
        .failure()
        .stderr(predicate::str::contains("non-interactive"))
        .stderr(predicate::str::contains("-y"));
}

#[test]
fn test_api_mode_only_applies_to_claude() {
    // OpenAI models should work with --accurate without needing consent or API key
    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.arg("--model").arg("gpt-4").arg("--accurate").write_stdin("test");

    // Should succeed (tiktoken is offline)
    cmd.assert().success();
}

#[test]
fn test_accurate_flag_without_yes_non_interactive() {
    let _lock = ENV_MUTEX.lock().unwrap();

    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.env("ANTHROPIC_API_KEY", "sk-test-key")
        .arg("--model")
        .arg("claude")
        .arg("--accurate")
        .write_stdin("test");

    // Non-interactive without -y should fail with clear message
    cmd.assert()
        .failure()
        .stderr(predicate::str::contains("consent"))
        .stderr(predicate::str::contains("non-interactive"));
}

#[test]
fn test_yes_flag_requires_accurate() {
    // -y flag should work but not change behavior without --accurate
    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.arg("--model").arg("claude").arg("-y").write_stdin("test");

    // Should succeed with estimation (no API call)
    cmd.assert().success();
}

#[test]
fn test_estimation_fallback_message() {
    let _lock = ENV_MUTEX.lock().unwrap();

    let mut cmd = Command::cargo_bin("token-count").unwrap();
    cmd.env("ANTHROPIC_API_KEY", "sk-invalid")
        .arg("--model")
        .arg("claude")
        .arg("--accurate")
        .arg("-y")
        .write_stdin("test");

    let output = cmd.output().unwrap();

    // Should produce output (estimation fallback)
    let stdout = String::from_utf8(output.stdout).unwrap();
    assert!(!stdout.trim().is_empty(), "Should output token count");

    // May have warning in stderr
    let stderr = String::from_utf8(output.stderr).unwrap();
    if !stderr.is_empty() {
        assert!(
            stderr.contains("Falling back") || stderr.contains("API"),
            "Unexpected stderr: {}",
            stderr
        );
    }
}

#[test]
fn test_all_claude_models_require_consent() {
    let models = vec!["claude-opus-4-6", "claude-sonnet-4-6", "claude-haiku-4-5"];

    for model in models {
        let mut cmd = Command::cargo_bin("token-count").unwrap();
        cmd.arg("--model").arg(model).arg("--accurate").write_stdin("test");

        // All Claude models in accurate mode require consent
        cmd.assert().failure().stderr(predicate::str::contains("non-interactive"));
    }
}