diff --git a/crates/tinyagents-harness/src/multimodal/stash_test.rs b/crates/tinyagents-harness/src/multimodal/stash_test.rs index 7f7f382b8..6d9e02b1f 100644 --- a/crates/tinyagents-harness/src/multimodal/stash_test.rs +++ b/crates/tinyagents-harness/src/multimodal/stash_test.rs @@ -173,7 +173,7 @@ async fn concurrent_writes_of_one_id_both_succeed() { } #[tokio::test] -async fn concurrent_writes_under_a_tight_cap_only_return_paths_that_exist() { +async fn concurrent_writes_under_a_tight_cap_leave_the_stash_within_its_cap() { let dir = tmp("tight-cap"); let one = { let probe = tmp("probe"); @@ -191,12 +191,17 @@ async fn concurrent_writes_under_a_tight_cap_only_return_paths_that_exist() { let s = stash(&dir, one, Duration::from_secs(3600)); let other = s.clone(); let (a, b) = tokio::join!(s.write("a", TINY_PNG), other.write("b", TINY_PNG)); - for path in [a, b].into_iter().flatten() { - assert!( - path.exists(), - "returned an evicted path: {}", - path.display() - ); - } + // A path may legitimately be evicted by a later write after its own write + // returned, so check the end state: at least one write succeeded and the + // stash is back under its cap. + assert!(a.is_ok() || b.is_ok(), "both concurrent writes failed"); + let files: Vec<_> = std::fs::read_dir(&dir) + .unwrap() + .flatten() + .filter(|e| !e.file_name().to_string_lossy().starts_with('.')) + .collect(); + let total: u64 = files.iter().map(|e| e.metadata().unwrap().len()).sum(); + assert!(!files.is_empty()); + assert!(total <= one, "stash over its cap: {total} > {one}"); let _ = std::fs::remove_dir_all(&dir); } diff --git a/crates/tinyagents-harness/src/run_queue/forwarder.rs b/crates/tinyagents-harness/src/run_queue/forwarder.rs index 89479225e..f617fe2d4 100644 --- a/crates/tinyagents-harness/src/run_queue/forwarder.rs +++ b/crates/tinyagents-harness/src/run_queue/forwarder.rs @@ -156,9 +156,9 @@ async fn forward_lane( mode, "[run_queue] forwarder closed mid-delivery; returned drained message(s) to the queue" ); - for msg in items { - queue.push(lane, msg).await; - } + // Back at the front, in order: anything queued after the drain is + // newer and must stay behind these. + queue.requeue_front(lane, items).await; return false; } tracing::debug!( @@ -391,7 +391,7 @@ impl Drop for SteeringForwarderGuard { .map(|(text, _lane, id)| (Some(id.clone()), Some(text.clone()))) .unwrap_or((None, None)); - // `RunQueue::push` is async (tokio `Mutex`); `Drop` is synchronous. Push + // `RunQueue` is async (tokio `Mutex`); `Drop` is synchronous. Put // the recovered steers back on a detached task so they land in the // session queue and become the next turn's input. The forwarder poll // loop keeps re-draining the queue, so even if the requeue completes @@ -400,11 +400,19 @@ impl Drop for SteeringForwarderGuard { Ok(rt) => { let label = thread_label.clone(); rt.spawn(async move { + // These were delivered before anything still queued, so + // they go back at the front of their lane, in order. + let mut steers = Vec::new(); + let mut collects = Vec::new(); for (text, lane, id) in requeue_items { - queue - .push(lane, T::requeued(id, text, &label, now_ms())) - .await; + let item = T::requeued(id, text, &label, now_ms()); + match lane { + QueueLane::Collect => collects.push(item), + _ => steers.push(item), + } } + queue.requeue_front(QueueLane::Steer, steers).await; + queue.requeue_front(QueueLane::Collect, collects).await; }); } Err(_) => { diff --git a/crates/tinyagents-harness/src/run_queue/forwarder_test.rs b/crates/tinyagents-harness/src/run_queue/forwarder_test.rs index dd0c18da7..6ba6a2bb4 100644 --- a/crates/tinyagents-harness/src/run_queue/forwarder_test.rs +++ b/crates/tinyagents-harness/src/run_queue/forwarder_test.rs @@ -235,3 +235,22 @@ async fn collect_reaches_the_next_model_boundary_as_additional_context() { ); assert!(!collect.text().starts_with(STEER_PREFIX)); } + +#[tokio::test] +async fn requeue_front_keeps_returned_items_ahead_of_newer_ones() { + let queue = RunQueue::::new(); + queue.push(QueueLane::Steer, item("newer", "newer")).await; + queue + .requeue_front( + QueueLane::Steer, + vec![item("old-1", "old-1"), item("old-2", "old-2")], + ) + .await; + let order: Vec = queue + .drain(QueueLane::Steer) + .await + .iter() + .map(|msg| msg.id().to_string()) + .collect(); + assert_eq!(order, ["old-1", "old-2", "newer"]); +} diff --git a/crates/tinyagents-harness/src/run_queue/mod.rs b/crates/tinyagents-harness/src/run_queue/mod.rs index e23a21b02..7d37fca6f 100644 --- a/crates/tinyagents-harness/src/run_queue/mod.rs +++ b/crates/tinyagents-harness/src/run_queue/mod.rs @@ -85,6 +85,22 @@ impl RunQueue { } } + /// Puts `items` back at the **front** of `lane`, in their given order, as + /// one atomic step. Used to return items that were drained but never + /// delivered, so they stay ahead of anything queued after them. + pub async fn requeue_front(&self, lane: QueueLane, items: Vec) { + if items.is_empty() { + return; + } + let mut inner = self.inner.lock().await; + let target = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + target.splice(0..0, items); + } + /// Drains one lane in FIFO order. pub async fn drain(&self, lane: QueueLane) -> Vec { let mut inner = self.inner.lock().await; diff --git a/crates/tinyagents-session/src/transcript/spend.rs b/crates/tinyagents-session/src/transcript/spend.rs index 8f35844f9..35e06c937 100644 --- a/crates/tinyagents-session/src/transcript/spend.rs +++ b/crates/tinyagents-session/src/transcript/spend.rs @@ -9,8 +9,8 @@ use std::collections::{BTreeMap, HashSet}; use std::path::{Path, PathBuf}; use super::{ - DisplayRecord, SessionTranscript, TranscriptMeta, TurnUsage, find_root_transcripts_for_thread, - read_transcript_display, + DisplayRecord, SessionTranscript, TranscriptMessage, TranscriptMeta, TurnUsage, + find_root_transcripts_for_thread, read_transcript_display, }; /// One transcript's own spend, summed from the per-turn `turn_usage` records @@ -124,6 +124,9 @@ fn file_spend(path: &Path) -> Option<(TranscriptSpend, TranscriptMeta)> { }; let mut usages: Vec<&TurnUsage> = Vec::new(); let mut counted_requests: HashSet<&str> = HashSet::new(); + // The logical context as of the previous record: what a compaction's kept + // rows are carried over from. + let mut context: Vec<&TranscriptMessage> = Vec::new(); for record in &display.records { match record { DisplayRecord::Message(message) if !message.interrupted => { @@ -133,43 +136,46 @@ fn file_spend(path: &Path) -> Option<(TranscriptSpend, TranscriptMeta)> { counted_requests.insert(request_id); } } + context.push(&message.message); } DisplayRecord::Message(_) => {} DisplayRecord::Compaction(marker) => { - let Some(turn) = marker.request_id.as_deref() else { - // A compaction written without a `request_id` has no turn - // identity to match rows against. Its writing turn's - // spend is the usage on the replacement's final assistant - // row (where the writer attaches it), and that row carries - // no `request_id` either; kept history rows that carry - // usage were read back and keep their own ids. Count it - // unless the very same record was already counted. - if let Some(usage) = marker - .replacement - .iter() - .rev() - .find(|row| row.message.role == "assistant") - .filter(|row| row.request_id.is_none()) - .and_then(|row| row.message.turn_usage.as_ref()) - && !usages.contains(&usage) - { - usages.push(usage); + match marker.request_id.as_deref() { + Some(turn) if !counted_requests.contains(turn) => { + let fresh: Vec<&TurnUsage> = marker + .replacement + .iter() + .filter(|row| row.request_id.as_deref() == Some(turn)) + .filter_map(|row| row.message.turn_usage.as_ref()) + .collect(); + if !fresh.is_empty() { + counted_requests.insert(turn); + } + usages.extend(fresh); + } + Some(_) => {} + None => { + // A compaction written without a `request_id` has no + // turn identity. Its writing turn's spend sits on the + // replacement's final assistant row; that row is fresh + // unless it was carried over from the context this + // compaction replaced (same position-free row: + // role, content and usage). + if let Some(row) = marker + .replacement + .iter() + .rev() + .find(|row| row.message.role == "assistant") + && let Some(usage) = row.message.turn_usage.as_ref() + && !context + .iter() + .any(|prior| carried_from(prior, &row.message)) + { + usages.push(usage); + } } - continue; - }; - if counted_requests.contains(turn) { - continue; - } - let fresh: Vec<&TurnUsage> = marker - .replacement - .iter() - .filter(|row| row.request_id.as_deref() == Some(turn)) - .filter_map(|row| row.message.turn_usage.as_ref()) - .collect(); - if !fresh.is_empty() { - counted_requests.insert(turn); } - usages.extend(fresh); + context = marker.replacement.iter().map(|row| &row.message).collect(); } } } @@ -177,6 +183,12 @@ fn file_spend(path: &Path) -> Option<(TranscriptSpend, TranscriptMeta)> { Some((spend, display.meta)) } +/// Whether `row` in a compaction replacement is `prior` carried over: same +/// role, content and usage record. +fn carried_from(prior: &TranscriptMessage, row: &TranscriptMessage) -> bool { + prior.role == row.role && prior.content == row.content && prior.turn_usage == row.turn_usage +} + /// Every descendant transcript of `root`, at any delegation depth. /// /// Sub-agent transcripts are named `{root_stem}__{child}` (and a grandchild diff --git a/crates/tinyagents-session/src/transcript/spend_test.rs b/crates/tinyagents-session/src/transcript/spend_test.rs index ca14c8942..5359b1518 100644 --- a/crates/tinyagents-session/src/transcript/spend_test.rs +++ b/crates/tinyagents-session/src/transcript/spend_test.rs @@ -491,3 +491,36 @@ fn request_less_compaction_turn_still_counts() { assert_eq!(spend.root.input_tokens, 12_000); assert_eq!(spend.root.turns, 2); } + +/// An anonymous compaction turn whose usage equals an earlier turn's is still +/// its own spend: its answer row was not carried over from the prior context. +#[test] +fn request_less_compaction_with_usage_identical_to_history_still_counts() { + let tmp = tempfile::tempdir().expect("tempdir"); + let thread = "thread-compacted-anon-twin"; + let stem = "1790000008_orchestrator_anon_twin"; + let root_meta = meta("orchestrator", "root", Some(thread)); + write_transcript_turns(tmp.path(), stem, &root_meta, &[(5_000, 50, 0)]); + let path = tmp.path().join("session_raw").join(format!("{stem}.jsonl")); + let persisted = vec![ + TranscriptMessage::new("user", "q0"), + TranscriptMessage::assistant("a0"), + ]; + let reduced = vec![ + TranscriptMessage::new("user", "summary"), + TranscriptMessage::new("user", "q1"), + TranscriptMessage::assistant("a1"), + ]; + append_transcript_turn( + &path, + &persisted, + &reduced, + &root_meta, + Some(&turn_usage(5_000, 50, 0)), + None, + ) + .expect("append compaction turn"); + let spend = thread_spend(tmp.path(), thread); + assert_eq!(spend.root.input_tokens, 10_000); + assert_eq!(spend.root.turns, 2); +}