返回 CodeWhale
completion_usage_tests.rs
根目录 / crates / tui / src / tools / subagent / completion_usage_tests.rs
1 use super::*;
2 use tempfile::tempdir;
3
4 fn child(
5 manager: &mut SubAgentManager,
6 name: &str,
7 parent: Option<&str>,
8 units: Option<u64>,
9 ) -> String {
10 let workspace = manager.workspace.clone();
11 let id = manager.insert_test_running_agent(name, &workspace);
12 let record = manager.worker_records.get_mut(&id).unwrap();
13 record.spec.parent_run_id = parent.map(str::to_string);
14 record.parent_run_id = parent.map(str::to_string);
15 record.usage.input_tokens = units.map(|units| units * 8);
16 record.usage.output_tokens = units.map(|units| units * 2);
17 record.usage.total_tokens = units.map(|units| units * 10);
18 id
19 }
20
21 fn resume(manager: &mut SubAgentManager, id: &str, source: &str) {
22 let record = manager.worker_records.get_mut(id).unwrap();
23 record.spec.launch_manifest = Some(
24 serde_json::from_value(json!({
25 "owner_session": record.spec.parent_run_id.as_deref().unwrap_or("root"), "child_id": id,
26 "profile": record.spec.runtime_profile, "prompt": "continue",
27 "cwd": null, "worktree": false, "writable_roots": [],
28 "writable_files": [], "coordination_contracts": [],
29 "resume_from_agent_id": source, "generation": 1
30 }))
31 .unwrap(),
32 );
33 // The same edge is present in both persisted representations.
34 manager.resume_targets.insert(source.into(), id.into());
35 }
36
37 fn sentinel(completion: &SubAgentCompletion) -> Value {
38 let opening = "<codewhale:subagent.done>";
39 let start = completion.payload.rfind(opening).unwrap() + opening.len();
40 let end = completion
41 .payload
42 .rfind("</codewhale:subagent.done>")
43 .unwrap();
44 serde_json::from_str(&completion.payload[start..end]).unwrap()
45 }
46
47 fn terminal_result(manager: &SubAgentManager, id: &str) -> SubAgentResult {
48 let mut result = manager.get_result(id).unwrap();
49 result.status = SubAgentStatus::Completed;
50 result.result = Some("Measured work is complete.".into());
51 result
52 }
53
54 fn receipt(manager: &SubAgentManager, id: &str) -> Value {
55 sentinel(&manager.completion_from_result_with_ref_for_session(
56 "workspace",
57 &terminal_result(manager, id),
58 None,
59 ))
60 }
61
62 fn family(manager: &mut SubAgentManager) -> (String, String, String, String) {
63 let root = child(manager, "root", None, Some(1));
64 let direct = child(manager, "direct", Some(&root), Some(2));
65 let grandchild = child(manager, "grandchild", Some(&direct), Some(3));
66 let continued = child(manager, "continued", Some(&direct), Some(4));
67 resume(manager, &continued, &grandchild);
68 let _sibling = child(manager, "outside", None, Some(90));
69 for id in [&direct, &grandchild, &continued] {
70 manager.worker_records.get_mut(id).unwrap().status = AgentWorkerStatus::Completed;
71 manager.agents.get_mut(id).unwrap().status = SubAgentStatus::Completed;
72 }
73 (root, direct, grandchild, continued)
74 }
75
76 #[tokio::test]
77 async fn completion_usage_live_terminal_counts_grandchildren_and_continuations_once() {
78 for status in [SubAgentStatus::Completed, SubAgentStatus::BudgetExhausted] {
79 let dir = tempdir().unwrap();
80 let mut manager = SubAgentManager::new(dir.path().to_path_buf(), 8);
81 let (root, _, _, _) = family(&mut manager);
82 let (tx, mut rx) = mpsc::channel(16);
83 let (event_tx, mut event_rx) = mpsc::channel(8);
84 manager.agents.get_mut(&root).unwrap().terminal_delivery =
85 Some(SubAgentTerminalDeliveryContext {
86 spawn_depth: 1,
87 parent_completion_tx: Some(tx),
88 mailbox: None,
89 event_tx: Some(event_tx),
90 session_id: "workspace".into(),
91 });
92 let mut result = terminal_result(&manager, &root);
93 result.status = status;
94 // The locked ledger wins over an earlier result snapshot.
95 result.usage.as_mut().unwrap().total_tokens = Some(123_456);
96 assert!(manager.finish_terminal_result(&root, result, false, false));
97 let completion = rx.try_recv().unwrap();
98 let payload = sentinel(&completion);
99 assert_eq!(payload["usage"]["own"]["total_tokens"], 10);
100 assert_eq!(payload["usage"]["descendants"]["workers"], 3);
101 assert_eq!(payload["usage"]["descendants"]["total_tokens"]["known"], 90);
102 assert_eq!(payload["usage"]["subtree"]["workers"], 4);
103 assert_eq!(payload["usage"]["subtree"]["active_workers"], 0);
104 assert_eq!(payload["usage"]["subtree"]["input_tokens"]["known"], 80);
105 assert_eq!(payload["usage"]["subtree"]["output_tokens"]["known"], 20);
106 assert_eq!(payload["usage"]["subtree"]["total_tokens"]["known"], 100);
107 assert_eq!(
108 payload["usage"]["subtree"]["total_tokens"]["reported_workers"],
109 4
110 );
111 assert!(payload.get("verification").is_some());
112 if manager.get_result(&root).unwrap().status == SubAgentStatus::BudgetExhausted {
113 assert_eq!(payload["event"], "subagent.failed");
114 }
115 let Event::AgentComplete {
116 result: event_result,
117 ..
118 } = event_rx.try_recv().unwrap()
119 else {
120 panic!("expected a terminal UI event");
121 };
122 assert_eq!(event_result, completion.payload);
123 assert!(
124 rx.try_recv().is_err(),
125 "terminal fan-in remains exactly once"
126 );
127 assert_eq!(manager.worker_records[&root].usage.total_tokens, Some(10));
128 assert!(serde_json::to_vec(&payload["usage"]).unwrap().len() <= 1600);
129 }
130 }
131
132 #[tokio::test]
133 async fn completion_usage_recovery_restores_measured_lineage_without_recounting() {
134 let dir = tempdir().unwrap();
135 let base = dir.path().to_path_buf();
136 let path = base.join(".codewhale/subagents/state.json");
137 let mut manager = SubAgentManager::new(base.clone(), 8).with_state_path(path.clone());
138 let (root, _, grandchild, continued) = family(&mut manager);
139 let result = terminal_result(&manager, &root);
140 assert!(manager.finish_terminal_result(&root, result, false, false));
141 let before = receipt(&manager, &root);
142 manager.persist_state_synchronously().unwrap();
143 let mut loaded = SubAgentManager::new(base, 8).with_state_path(path);
144 loaded.load_state().unwrap();
145 let after = receipt(&loaded, &root);
146 assert_eq!(after["usage"], before["usage"]);
147 assert_eq!(after["usage"]["subtree"]["total_tokens"]["known"], 100);
148 assert_eq!(loaded.continuation_target(&grandchild).unwrap(), continued);
149 // A resumed child's receipt covers its own forward subtree, not its
150 // predecessor's already-delivered spend or unrelated siblings.
151 assert_eq!(
152 receipt(&loaded, &continued)["usage"]["subtree"]["total_tokens"]["known"],
153 40
154 );
155 }
156
157 #[tokio::test]
158 async fn completion_usage_rejects_foreign_bridges_and_counts_cycles_once() {
159 let dir = tempdir().unwrap();
160 let mut manager = SubAgentManager::new(dir.path().to_path_buf(), 8);
161 let root = child(&mut manager, "root", None, Some(1));
162 let direct = child(&mut manager, "direct", Some(&root), Some(2));
163 manager
164 .worker_records
165 .get_mut(&root)
166 .unwrap()
167 .spec
168 .parent_run_id = Some(direct.clone());
169 let foreign = child(&mut manager, "foreign", Some(&root), Some(90));
170 manager.assign_test_session_owner(&foreign, "another-owner");
171 let bridged = child(&mut manager, "bridged", Some(&foreign), Some(80));
172 resume(&mut manager, &bridged, &foreign);
173 manager.resume_targets.insert(root.clone(), foreign.clone());
174 // A forged manifest cannot create a same-owner edge from foreign authority.
175 let manifest = manager
176 .worker_records
177 .get_mut(&bridged)
178 .unwrap()
179 .spec
180 .launch_manifest
181 .as_mut()
182 .unwrap();
183 manifest.owner_session = "another-owner".into();
184 manifest.resume_from_agent_id = Some(root.clone());
185 let payload = receipt(&manager, &root);
186 assert_eq!(payload["usage"]["subtree"]["workers"], 2);
187 assert_eq!(payload["usage"]["subtree"]["total_tokens"]["known"], 30);
188 assert_eq!(payload["usage"]["descendants"]["active_workers"], 1);
189 let foreign_projection = sentinel(&manager.completion_from_result_with_ref_for_session(
190 "another-owner",
191 &terminal_result(&manager, &root),
192 None,
193 ));
194 assert_eq!(foreign_projection["usage"]["scope"], "unavailable");
195 assert!(foreign_projection["usage"]["own"]["total_tokens"].is_null());
196 }
197
198 #[tokio::test]
199 async fn completion_usage_distinguishes_unknown_zero_partial_and_overflow() {
200 let dir = tempdir().unwrap();
201 let mut manager = SubAgentManager::new(dir.path().to_path_buf(), 4);
202 let root = child(&mut manager, "unknown", None, None);
203 let direct = child(&mut manager, "zero", Some(&root), Some(0));
204 let payload = receipt(&manager, &root);
205 assert!(payload["usage"]["own"]["total_tokens"].is_null());
206 assert_eq!(payload["usage"]["subtree"]["total_tokens"]["known"], 0);
207 assert_eq!(
208 payload["usage"]["subtree"]["total_tokens"]["reported_workers"],
209 1
210 );
211 assert_eq!(payload["usage"]["subtree"]["workers"], 2);
212 manager
213 .worker_records
214 .get_mut(&direct)
215 .unwrap()
216 .usage
217 .total_tokens = None;
218 let payload = receipt(&manager, &root);
219 assert!(payload["usage"]["subtree"]["total_tokens"]["known"].is_null());
220 assert_eq!(
221 payload["usage"]["subtree"]["total_tokens"]["reported_workers"],
222 0
223 );
224 assert_eq!(
225 receipt(&manager, &direct)["usage"]["descendants"]["total_tokens"]["known"],
226 0
227 );
228 manager
229 .worker_records
230 .get_mut(&root)
231 .unwrap()
232 .usage
233 .total_tokens = Some(u64::MAX);
234 manager
235 .worker_records
236 .get_mut(&direct)
237 .unwrap()
238 .usage
239 .total_tokens = Some(1);
240 let payload = receipt(&manager, &root);
241 assert!(payload["usage"]["subtree"]["total_tokens"]["known"].is_null());
242 assert_eq!(
243 payload["usage"]["subtree"]["total_tokens"]["overflow"],
244 true
245 );
246 assert_eq!(
247 payload["usage"]["subtree"]["total_tokens"]["reported_workers"],
248 2
249 );
250 }
251
252 #[tokio::test]
253 async fn completion_usage_keeps_unknown_then_known_response_subtotals_visible() {
254 let dir = tempdir().unwrap();
255 let mut manager = SubAgentManager::new(dir.path().to_path_buf(), 8);
256 let root = child(&mut manager, "partial-root", None, None);
257 let direct = child(&mut manager, "partial-child", Some(&root), None);
258 let zero = child(&mut manager, "known-zero", Some(&direct), None);
259 // The parent, manifest and successor map describe the same descendant.
260 resume(&mut manager, &zero, &direct);
261 let foreign = child(&mut manager, "foreign-unknown", Some(&root), None);
262 manager.assign_test_session_owner(&foreign, "another-owner");
263 manager.record_worker_usage(&foreign, "foreign-unknown", &Usage::default(), None);
264
265 for (id, input, output) in [(&root, 8, 2), (&direct, 16, 4)] {
266 manager.record_worker_usage(id, "unreported-response", &Usage::default(), None);
267 let before = receipt(&manager, id);
268 assert!(before["usage"]["own"]["total_tokens"].is_null());
269 assert_eq!(before["usage"]["own"]["has_unreported_usage"], true);
270 manager.record_worker_usage(
271 id,
272 "reported-response",
273 &Usage {
274 input_tokens: input,
275 output_tokens: output,
276 ..Usage::default()
277 },
278 None,
279 );
280 }
281 manager.record_worker_usage(
282 &zero,
283 "reported-zero-response",
284 &Usage {
285 prompt_cache_hit_tokens: Some(0),
286 ..Usage::default()
287 },
288 None,
289 );
290
291 let payload = receipt(&manager, &root);
292 let usage = &payload["usage"];
293 assert_eq!(usage["own"]["input_tokens"], 8);
294 assert_eq!(usage["own"]["output_tokens"], 2);
295 assert_eq!(usage["own"]["total_tokens"], 10);
296 assert_eq!(usage["own"]["has_unreported_usage"], true);
297 assert_eq!(usage["descendants"]["workers"], 2);
298 assert_eq!(usage["descendants"]["unreported_usage_workers"], 1);
299 assert_eq!(usage["descendants"]["total_tokens"]["known"], 20);
300 assert_eq!(usage["subtree"]["workers"], 3);
301 assert_eq!(usage["subtree"]["unreported_usage_workers"], 2);
302 assert_eq!(usage["subtree"]["input_tokens"]["known"], 24);
303 assert_eq!(usage["subtree"]["output_tokens"]["known"], 6);
304 assert_eq!(usage["subtree"]["total_tokens"]["known"], 30);
305 assert_eq!(usage["subtree"]["total_tokens"]["reported_workers"], 3);
306 assert!(serde_json::to_vec(usage).unwrap().len() <= 1600);
307
308 let zero_receipt = receipt(&manager, &zero);
309 assert_eq!(zero_receipt["usage"]["own"]["total_tokens"], 0);
310 assert_eq!(zero_receipt["usage"]["own"]["has_unreported_usage"], false);
311 assert_eq!(
312 zero_receipt["usage"]["subtree"]["unreported_usage_workers"],
313 0
314 );
315 assert_eq!(
316 zero_receipt["usage"]["descendants"]["unreported_usage_workers"],
317 0
318 );
319 assert_eq!(manager.worker_records[&root].usage.total_tokens, Some(10));
320 assert_eq!(manager.worker_records[&direct].usage.total_tokens, Some(20));
321 }
322
323 #[tokio::test]
324 async fn completion_usage_counts_real_manifest_only_root_fork() {
325 let dir = tempdir().unwrap();
326 let manager = new_shared_subagent_manager(dir.path().to_path_buf(), 4);
327 let mut runtime = tests::stub_runtime()
328 .with_max_spawn_depth(3)
329 .child_runtime();
330 runtime.context = ToolContext::new(dir.path());
331 runtime.manager = Arc::clone(&manager);
332 // Exercise real registration without allowing any provider request.
333 runtime.cancel_token.cancel();
334 let mut guard = manager.write().await;
335 let source = child(&mut guard, "source", None, Some(1));
336 let fork = guard
337 .spawn_background_with_assignment_options(
338 Arc::clone(&manager),
339 runtime,
340 FleetRole::Scout,
341 "Read the prior work.".into(),
342 SubAgentAssignment::new("Read the prior work.".into(), None),
343 Some(vec![]),
344 SubAgentSpawnOptions {
345 resume_from_agent_id: Some(source.clone()),
346 checkpoint_continuation: false,
347 ..Default::default()
348 },
349 None,
350 )
351 .unwrap();
352 let record = guard.worker_records.get_mut(&fork.agent_id).unwrap();
353 assert_eq!(record.owner_session_id, "workspace");
354 assert!(record.spec.parent_run_id.is_none());
355 let manifest = record.spec.launch_manifest.as_ref().unwrap();
356 assert_eq!(manifest.owner_session, "root");
357 assert_eq!(
358 manifest.resume_from_agent_id.as_deref(),
359 Some(source.as_str())
360 );
361 record.usage.input_tokens = Some(16);
362 record.usage.output_tokens = Some(4);
363 record.usage.total_tokens = Some(20);
364 assert!(!guard.resume_targets.contains_key(&source));
365 let payload = receipt(&guard, &source);
366 assert_eq!(payload["usage"]["descendants"]["workers"], 1);
367 assert_eq!(payload["usage"]["descendants"]["total_tokens"]["known"], 20);
368 assert_eq!(payload["usage"]["subtree"]["total_tokens"]["known"], 30);
369 }
370
371 #[tokio::test]
372 #[expect(
373 clippy::print_stderr,
374 reason = "libtest-only measurement; never the TUI"
375 )]
376 async fn completion_usage_dozen_child_status_measures_bytes_and_keeps_descendant_rows() {
377 let dir = tempdir().unwrap();
378 let manager = new_shared_subagent_manager(dir.path().to_path_buf(), 16);
379 let mut ids: Vec<String> = Vec::new();
380 {
381 let mut guard = manager.write().await;
382 for (index, name) in [
383 "a2f4095d", "846bd172", "ce819730", "918deb4a", "716a54bf", "f291ac63", "b5701e92",
384 "294e7cab", "ac967f81", "670d2ea9", "edf94328", "547ab013",
385 ]
386 .into_iter()
387 .enumerate()
388 {
389 let parent = (index > 0).then(|| ids[(index - 1) / 2].clone());
390 let depth = parent.as_ref().map_or(1, |parent| {
391 guard.worker_records[parent].spec.spawn_depth + 1
392 });
393 let id = child(&mut guard, name, parent.as_deref(), Some(index as u64 + 1));
394 let record = guard.worker_records.get_mut(&id).unwrap();
395 record.spec.spawn_depth = depth;
396 record.spec.max_spawn_depth = 4;
397 record.spec.runtime_profile.spawn_depth = depth;
398 record.spec.runtime_profile.max_spawn_depth = 4;
399 record.latest_message = Some("starting".into());
400 record.spec.child_route = Some(ChildRouteReceipt {
401 requested_type: "explore".into(),
402 requested_profile: Some("scout".into()),
403 resolved_profile_id: Some("scout".into()),
404 profile_origin: Some("workspace".into()),
405 canonical_role: "explore".into(),
406 provider_id: "deepseek".into(),
407 model_id: "deepseek-v4-flash".into(),
408 route_source: "profile.model".into(),
409 fallback_note: None,
410 requested_reasoning: "inherit".into(),
411 effective_reasoning: Some("medium".into()),
412 runtime_version: "0.9.13".into(),
413 runtime_build_sha: "a".repeat(40),
414 });
415 record.spec.runtime_profile.max_steps = 12;
416 record.spec.runtime_profile.wall_time_secs = Some(600);
417 record.spec.runtime_profile.wall_deadline_ms = Some(record.updated_at_ms + 600_000);
418 if index == 11 {
419 record.verification.status = "deliverable_missing".into();
420 record.verification.summary = "The claimed report.md is missing.".into();
421 record.verification.deliverables = vec![DeliverableVerdict {
422 path: "report.md".into(),
423 status: "missing".into(),
424 bytes: None,
425 }];
426 }
427 ids.push(id);
428 }
429 }
430 let mut offset = 0;
431 let mut bytes = 0;
432 let mut pages = 0;
433 let mut seen = HashSet::new();
434 loop {
435 let output = inspect_agent_from_input(
436 &json!({"action":"status", "offset":offset}),
437 Arc::clone(&manager),
438 &ToolContext::new(dir.path()),
439 false,
440 None,
441 )
442 .await
443 .unwrap();
444 bytes += output.content.len();
445 pages += 1;
446 eprintln!("DOZEN_CHILD_STATUS_JSON {}", output.content);
447 assert!(output.content.len() <= lifecycle::COMPACT_STATUS_BYTES);
448 let payload: Value = serde_json::from_str(&output.content).unwrap();
449 assert_eq!(payload["usage"]["total_tokens"], 780);
450 for row in lifecycle_tests::status_rows(&payload) {
451 let id = row["agent_id"].as_str().unwrap();
452 assert!(seen.insert(id.to_string()));
453 let index = ids.iter().position(|expected| expected == id).unwrap();
454 assert_eq!(row["total_tokens"], (index + 1) * 10);
455 for key in [
456 "compact",
457 "terminal",
458 "name",
459 "steps_taken",
460 "child_route",
461 "effective_limits",
462 "usage",
463 "max_spawn_depth",
464 ] {
465 assert!(row.get(key).is_none(), "{key}: {row}");
466 }
467 assert_eq!(row["needs_continuation"], false, "{row}");
468 assert_eq!(row["activity"], "starting");
469 for key in ["duration_ms", "last_activity_ms", "spawn_depth"] {
470 assert!(row[key].is_u64(), "{key}: {row}");
471 }
472 if index == 11 {
473 assert_eq!(row["verification"]["status"], "deliverable_missing");
474 assert_eq!(row["verification"]["deliverable_counts"]["missing"], 1);
475 assert_eq!(
476 row["verification"]["summary"],
477 "The claimed report.md is missing."
478 );
479 } else {
480 assert_eq!(row["verification"], json!({"status": "self_report_only"}));
481 }
482 if index > 0 {
483 assert_eq!(row["parent_agent_id"], ids[(index - 1) / 2]);
484 }
485 }
486 let Some(next) = payload["next_offset"].as_u64() else {
487 break;
488 };
489 assert!(next > offset);
490 offset = next;
491 }
492 assert_eq!(seen.len(), 12);
493 assert_eq!(pages, 1, "ordinary twelve-worker roster must fit one page");
494 assert!(bytes <= 3072, "twelve-worker roster used {bytes} bytes");
495 let addressed = inspect_agent_from_input(
496 &json!({"action":"status", "agent_id":ids[0]}),
497 Arc::clone(&manager),
498 &ToolContext::new(dir.path()),
499 false,
500 None,
501 )
502 .await
503 .unwrap();
504 let addressed: Value = serde_json::from_str(&addressed.content).unwrap();
505 assert_eq!(addressed["compact"], true);
506 assert_eq!(addressed["child_route"]["model_id"], "deepseek-v4-flash");
507 assert_eq!(addressed["effective_limits"]["max_steps"], 12);
508 assert_eq!(addressed["effective_limits"]["wall_time_secs"], 600);
509 assert_eq!(addressed["max_spawn_depth"], 4);
510 assert_eq!(addressed["usage"]["input_tokens"], 8);
511 assert_eq!(addressed["usage"]["output_tokens"], 2);
512 eprintln!(
513 "DOZEN_CHILD_STATUS_MEASUREMENT children=12 pages={pages} serialized_bytes={bytes}; token_count=unmeasured"
514 );
515 }
516
516 lines RUST