| 1 | //! End-to-end shape lock for the per-model-call `turn_usage` event on the |
| 2 | //! `codewhale exec --output-format stream-json` stream (#52 / FINISH-0.9.4). |
| 3 | //! |
| 4 | //! A `wiremock` OpenAI-compatible endpoint stands in for the provider. Two |
| 5 | //! cases pin the contract: |
| 6 | //! |
| 7 | //! - usage reported by the provider -> exactly one `turn_usage` event per |
| 8 | //! model call, carrying the reported input/output/reasoning/cache fields, |
| 9 | //! and the pre-existing event sequence (`content` … `metadata` → `done`) |
| 10 | //! is unchanged for existing consumers; |
| 11 | //! - usage absent from the provider stream -> no `turn_usage` event at all |
| 12 | //! (honest absence, never fabricated zeros-as-data). |
| 13 | |
| 14 | #![cfg(unix)] |
| 15 | |
| 16 | use std::io::Read; |
| 17 | use std::path::PathBuf; |
| 18 | use std::process::{Command, Stdio}; |
| 19 | use std::time::Duration; |
| 20 | |
| 21 | use serde_json::{Value, json}; |
| 22 | use tempfile::TempDir; |
| 23 | use wait_timeout::ChildExt; |
| 24 | use wiremock::matchers::{method, path}; |
| 25 | use wiremock::{Mock, MockServer, ResponseTemplate}; |
| 26 | |
| 27 | const TEST_MODEL: &str = "turn-usage-model"; |
| 28 | const RUN_TIMEOUT: Duration = Duration::from_secs(60); |
| 29 | |
| 30 | fn sse_chunk(value: Value) -> String { |
| 31 | format!( |
| 32 | "data: {}\n\n", |
| 33 | serde_json::to_string(&value).expect("SSE JSON") |
| 34 | ) |
| 35 | } |
| 36 | |
| 37 | /// Final-answer SSE whose closing chunk reports usage with reasoning and |
| 38 | /// DeepSeek-style prompt-cache fields. |
| 39 | fn answer_sse_with_usage(answer: &str) -> String { |
| 40 | [ |
| 41 | sse_chunk(json!({ |
| 42 | "id": "chatcmpl-usage", |
| 43 | "object": "chat.completion.chunk", |
| 44 | "model": TEST_MODEL, |
| 45 | "choices": [{"index": 0, "delta": {"content": answer}, "finish_reason": null}] |
| 46 | })), |
| 47 | sse_chunk(json!({ |
| 48 | "id": "chatcmpl-usage", |
| 49 | "object": "chat.completion.chunk", |
| 50 | "model": TEST_MODEL, |
| 51 | "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], |
| 52 | "usage": { |
| 53 | "prompt_tokens": 20, |
| 54 | "completion_tokens": 8, |
| 55 | "total_tokens": 28, |
| 56 | "completion_tokens_details": {"reasoning_tokens": 5}, |
| 57 | "prompt_cache_hit_tokens": 12, |
| 58 | "prompt_cache_miss_tokens": 8 |
| 59 | } |
| 60 | })), |
| 61 | "data: [DONE]\n\n".to_string(), |
| 62 | ] |
| 63 | .join("") |
| 64 | } |
| 65 | |
| 66 | /// Final-answer SSE whose provider never reports usage. |
| 67 | fn answer_sse_without_usage(answer: &str) -> String { |
| 68 | [ |
| 69 | sse_chunk(json!({ |
| 70 | "id": "chatcmpl-no-usage", |
| 71 | "object": "chat.completion.chunk", |
| 72 | "model": TEST_MODEL, |
| 73 | "choices": [{"index": 0, "delta": {"content": answer}, "finish_reason": null}] |
| 74 | })), |
| 75 | sse_chunk(json!({ |
| 76 | "id": "chatcmpl-no-usage", |
| 77 | "object": "chat.completion.chunk", |
| 78 | "model": TEST_MODEL, |
| 79 | "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}] |
| 80 | })), |
| 81 | "data: [DONE]\n\n".to_string(), |
| 82 | ] |
| 83 | .join("") |
| 84 | } |
| 85 | |
| 86 | fn sse_response(body: String) -> ResponseTemplate { |
| 87 | ResponseTemplate::new(200) |
| 88 | .insert_header("content-type", "text/event-stream") |
| 89 | .insert_header("cache-control", "no-cache") |
| 90 | .set_body_string(body) |
| 91 | } |
| 92 | |
| 93 | fn json_response(value: Value) -> ResponseTemplate { |
| 94 | ResponseTemplate::new(200) |
| 95 | .insert_header("content-type", "application/json") |
| 96 | .set_body_json(value) |
| 97 | } |
| 98 | |
| 99 | async fn start_mock_llm(answer_sse: String) -> MockServer { |
| 100 | let server = MockServer::start().await; |
| 101 | |
| 102 | Mock::given(method("GET")) |
| 103 | .and(path("/v1/models")) |
| 104 | .respond_with(json_response(json!({ |
| 105 | "object": "list", |
| 106 | "data": [{ "id": TEST_MODEL, "object": "model" }] |
| 107 | }))) |
| 108 | .mount(&server) |
| 109 | .await; |
| 110 | |
| 111 | Mock::given(method("POST")) |
| 112 | .and(path("/v1/chat/completions")) |
| 113 | .respond_with(sse_response(answer_sse)) |
| 114 | .mount(&server) |
| 115 | .await; |
| 116 | |
| 117 | server |
| 118 | } |
| 119 | |
| 120 | fn preserve_host_env(command: &mut Command) { |
| 121 | command.env_clear(); |
| 122 | for key in [ |
| 123 | "PATH", |
| 124 | "PATHEXT", |
| 125 | "SystemRoot", |
| 126 | "SystemDrive", |
| 127 | "WINDIR", |
| 128 | "COMSPEC", |
| 129 | "TEMP", |
| 130 | "TMP", |
| 131 | "TERM", |
| 132 | "COLORTERM", |
| 133 | "LANG", |
| 134 | "LC_ALL", |
| 135 | ] { |
| 136 | if let Some(value) = std::env::var_os(key) { |
| 137 | command.env(key, value); |
| 138 | } |
| 139 | } |
| 140 | } |
| 141 | |
| 142 | fn run_exec_stream_json(server: &MockServer) -> Vec<Value> { |
| 143 | let workspace = TempDir::new().expect("workspace tempdir"); |
| 144 | let home = TempDir::new().expect("home tempdir"); |
| 145 | |
| 146 | let mut command = Command::new(codewhale_tui_binary()); |
| 147 | preserve_host_env(&mut command); |
| 148 | command |
| 149 | .current_dir(workspace.path()) |
| 150 | .arg("--workspace") |
| 151 | .arg(workspace.path()) |
| 152 | .arg("--no-project-config") |
| 153 | .arg("exec") |
| 154 | .arg("--auto") |
| 155 | .arg("--model") |
| 156 | .arg(TEST_MODEL) |
| 157 | .arg("--output-format") |
| 158 | .arg("stream-json") |
| 159 | .arg("answer briefly") |
| 160 | .env("HOME", home.path()) |
| 161 | .env("USERPROFILE", home.path()) |
| 162 | .env("XDG_CONFIG_HOME", home.path().join(".config")) |
| 163 | .env("XDG_DATA_HOME", home.path().join(".local").join("share")) |
| 164 | .env("XDG_CACHE_HOME", home.path().join(".cache")) |
| 165 | .env( |
| 166 | "CODEWHALE_CONFIG_PATH", |
| 167 | home.path().join(".codewhale").join("config.toml"), |
| 168 | ) |
| 169 | .env( |
| 170 | "DEEPSEEK_CONFIG_PATH", |
| 171 | home.path().join(".deepseek").join("config.toml"), |
| 172 | ) |
| 173 | .env("DEEPSEEK_API_KEY", "ci-test-key-not-real") |
| 174 | .env("DEEPSEEK_BASE_URL", server.uri()) |
| 175 | .env("CODEWHALE_BASE_URL", server.uri()) |
| 176 | .env("DEEPSEEK_MODEL", TEST_MODEL) |
| 177 | .env("CODEWHALE_MODEL", TEST_MODEL) |
| 178 | .env("RUST_LOG", "warn") |
| 179 | .stdout(Stdio::piped()) |
| 180 | .stderr(Stdio::piped()); |
| 181 | |
| 182 | std::fs::create_dir_all(home.path().join(".codewhale")).expect("create codewhale config dir"); |
| 183 | std::fs::create_dir_all(home.path().join(".deepseek")).expect("create deepseek config dir"); |
| 184 | |
| 185 | let mut child = command.spawn().expect("spawn codewhale-tui exec"); |
| 186 | let stdout_reader = read_pipe_in_background(child.stdout.take().expect("stdout pipe")); |
| 187 | let stderr_reader = read_pipe_in_background(child.stderr.take().expect("stderr pipe")); |
| 188 | |
| 189 | let status = match child |
| 190 | .wait_timeout(RUN_TIMEOUT) |
| 191 | .expect("wait for codewhale-tui") |
| 192 | { |
| 193 | Some(status) => status, |
| 194 | None => { |
| 195 | let _ = child.kill(); |
| 196 | let _ = child.wait(); |
| 197 | let stdout = join_pipe_reader(stdout_reader, "stdout"); |
| 198 | let stderr = join_pipe_reader(stderr_reader, "stderr"); |
| 199 | panic!( |
| 200 | "codewhale-tui exec timed out after {RUN_TIMEOUT:?}\nstdout:\n{}\nstderr:\n{}", |
| 201 | String::from_utf8_lossy(&stdout), |
| 202 | String::from_utf8_lossy(&stderr) |
| 203 | ); |
| 204 | } |
| 205 | }; |
| 206 | |
| 207 | let stdout = join_pipe_reader(stdout_reader, "stdout"); |
| 208 | let stderr = join_pipe_reader(stderr_reader, "stderr"); |
| 209 | assert!( |
| 210 | status.success(), |
| 211 | "codewhale-tui exec failed\nstdout:\n{}\nstderr:\n{}", |
| 212 | String::from_utf8_lossy(&stdout), |
| 213 | String::from_utf8_lossy(&stderr) |
| 214 | ); |
| 215 | |
| 216 | let stdout = String::from_utf8_lossy(&stdout).into_owned(); |
| 217 | stdout |
| 218 | .lines() |
| 219 | .filter(|line| !line.trim().is_empty()) |
| 220 | .map(|line| { |
| 221 | serde_json::from_str(line).unwrap_or_else(|err| { |
| 222 | panic!("stream-json line should parse: {err}\nline: {line}\nstdout:\n{stdout}") |
| 223 | }) |
| 224 | }) |
| 225 | .collect() |
| 226 | } |
| 227 | |
| 228 | fn read_pipe_in_background<R>(mut reader: R) -> std::thread::JoinHandle<std::io::Result<Vec<u8>>> |
| 229 | where |
| 230 | R: Read + Send + 'static, |
| 231 | { |
| 232 | std::thread::spawn(move || { |
| 233 | let mut output = Vec::new(); |
| 234 | reader.read_to_end(&mut output).map(|_| output) |
| 235 | }) |
| 236 | } |
| 237 | |
| 238 | fn join_pipe_reader( |
| 239 | handle: std::thread::JoinHandle<std::io::Result<Vec<u8>>>, |
| 240 | stream_name: &str, |
| 241 | ) -> Vec<u8> { |
| 242 | handle |
| 243 | .join() |
| 244 | .unwrap_or_else(|_| panic!("{stream_name} reader thread panicked")) |
| 245 | .unwrap_or_else(|err| panic!("failed to read {stream_name}: {err}")) |
| 246 | } |
| 247 | |
| 248 | fn codewhale_tui_binary() -> PathBuf { |
| 249 | if let Some(path) = option_env!("CARGO_BIN_EXE_codewhale-tui") { |
| 250 | return PathBuf::from(path); |
| 251 | } |
| 252 | if let Ok(path) = std::env::var("CARGO_BIN_EXE_codewhale-tui") { |
| 253 | return PathBuf::from(path); |
| 254 | } |
| 255 | |
| 256 | let mut path = std::env::current_exe().expect("current test executable path"); |
| 257 | path.pop(); |
| 258 | if path.ends_with("deps") { |
| 259 | path.pop(); |
| 260 | } |
| 261 | path.push(format!("codewhale-tui{}", std::env::consts::EXE_SUFFIX)); |
| 262 | path |
| 263 | } |
| 264 | |
| 265 | fn events_of_type<'a>(events: &'a [Value], event_type: &str) -> Vec<&'a Value> { |
| 266 | events |
| 267 | .iter() |
| 268 | .filter(|event| event.get("type").and_then(Value::as_str) == Some(event_type)) |
| 269 | .collect() |
| 270 | } |
| 271 | |
| 272 | #[tokio::test(flavor = "multi_thread")] |
| 273 | async fn turn_usage_event_is_emitted_with_reported_fields_and_stream_contract_holds() { |
| 274 | let server = start_mock_llm(answer_sse_with_usage("done in one step")).await; |
| 275 | let events = run_exec_stream_json(&server); |
| 276 | |
| 277 | // Every event carries the stream schema envelope. |
| 278 | for event in &events { |
| 279 | assert_eq!(event["schema"], "codewhale.exec-stream"); |
| 280 | assert_eq!(event["schema_version"], 1); |
| 281 | } |
| 282 | |
| 283 | // Exactly one per-call usage receipt, numbered from 1. |
| 284 | let usage_events = events_of_type(&events, "turn_usage"); |
| 285 | assert_eq!( |
| 286 | usage_events.len(), |
| 287 | 1, |
| 288 | "expected one turn_usage event: {events:#?}" |
| 289 | ); |
| 290 | let usage = usage_events[0]; |
| 291 | assert_eq!(usage["turn"], 1); |
| 292 | assert_eq!(usage["input_tokens"], 20); |
| 293 | assert_eq!(usage["output_tokens"], 8); |
| 294 | assert_eq!(usage["reasoning_tokens"], 5); |
| 295 | assert_eq!(usage["prompt_cache_hit_tokens"], 12); |
| 296 | assert_eq!(usage["prompt_cache_miss_tokens"], 8); |
| 297 | assert!( |
| 298 | usage["duration_ms"].as_u64().is_some(), |
| 299 | "duration_ms must be a non-negative integer: {usage}" |
| 300 | ); |
| 301 | // Fields the provider did not report are omitted, not zero-filled. |
| 302 | let usage_object = usage.as_object().expect("turn_usage object"); |
| 303 | for absent in ["prompt_cache_write_tokens", "reasoning_replay_tokens"] { |
| 304 | assert!( |
| 305 | !usage_object.contains_key(absent), |
| 306 | "{absent} must be omitted when unreported: {usage}" |
| 307 | ); |
| 308 | } |
| 309 | |
| 310 | // The usage receipt lands after the model output it accounts for and |
| 311 | // before the terminal receipts. |
| 312 | let types: Vec<&str> = events |
| 313 | .iter() |
| 314 | .filter_map(|event| event.get("type").and_then(Value::as_str)) |
| 315 | .collect(); |
| 316 | let content_pos = types.iter().position(|t| *t == "content"); |
| 317 | let usage_pos = types.iter().position(|t| *t == "turn_usage"); |
| 318 | assert!( |
| 319 | content_pos.is_some_and(|c| usage_pos.is_some_and(|u| c < u)), |
| 320 | "turn_usage must follow the content it accounts for: {types:?}" |
| 321 | ); |
| 322 | |
| 323 | // Existing consumers' terminal contract is unchanged: `metadata` |
| 324 | // immediately precedes exactly one trailing `done`. |
| 325 | assert_eq!(types.last(), Some(&"done"), "stream must end with done"); |
| 326 | assert_eq!( |
| 327 | types.get(types.len() - 2), |
| 328 | Some(&"metadata"), |
| 329 | "metadata must immediately precede done: {types:?}" |
| 330 | ); |
| 331 | assert_eq!( |
| 332 | events_of_type(&events, "done").len(), |
| 333 | 1, |
| 334 | "exactly one done event" |
| 335 | ); |
| 336 | let metadata = events_of_type(&events, "metadata"); |
| 337 | assert_eq!(metadata.len(), 1, "exactly one metadata event"); |
| 338 | // The terminal receipt still carries the cumulative usage. |
| 339 | assert_eq!(metadata[0]["meta"]["input_tokens"], 20); |
| 340 | assert_eq!(metadata[0]["meta"]["output_tokens"], 8); |
| 341 | assert_eq!(metadata[0]["meta"]["reasoning_tokens"], 5); |
| 342 | } |
| 343 | |
| 344 | #[tokio::test(flavor = "multi_thread")] |
| 345 | async fn turn_usage_event_is_skipped_when_provider_reports_no_usage() { |
| 346 | let server = start_mock_llm(answer_sse_without_usage("quiet answer")).await; |
| 347 | let events = run_exec_stream_json(&server); |
| 348 | |
| 349 | assert!( |
| 350 | events_of_type(&events, "turn_usage").is_empty(), |
| 351 | "no turn_usage event without provider-reported usage: {events:#?}" |
| 352 | ); |
| 353 | |
| 354 | // The rest of the stream contract still holds. |
| 355 | let types: Vec<&str> = events |
| 356 | .iter() |
| 357 | .filter_map(|event| event.get("type").and_then(Value::as_str)) |
| 358 | .collect(); |
| 359 | assert!(types.contains(&"content"), "content missing: {types:?}"); |
| 360 | assert_eq!(types.last(), Some(&"done"), "stream must end with done"); |
| 361 | assert_eq!(types.get(types.len() - 2), Some(&"metadata")); |
| 362 | } |
| 363 |