Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 13 additions & 8 deletions crates/tinyagents-harness/src/multimodal/stash_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand All @@ -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);
}
22 changes: 15 additions & 7 deletions crates/tinyagents-harness/src/run_queue/forwarder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,9 +156,9 @@ async fn forward_lane<T: QueuedMessage>(
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!(
Expand Down Expand Up @@ -391,7 +391,7 @@ impl<T: QueuedMessage> Drop for SteeringForwarderGuard<T> {
.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
Expand All @@ -400,11 +400,19 @@ impl<T: QueuedMessage> Drop for SteeringForwarderGuard<T> {
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;
Comment thread
senamakel marked this conversation as resolved.
queue.requeue_front(QueueLane::Collect, collects).await;
});
}
Err(_) => {
Expand Down
19 changes: 19 additions & 0 deletions crates/tinyagents-harness/src/run_queue/forwarder_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<Item>::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<String> = queue
.drain(QueueLane::Steer)
.await
.iter()
.map(|msg| msg.id().to_string())
.collect();
assert_eq!(order, ["old-1", "old-2", "newer"]);
}
16 changes: 16 additions & 0 deletions crates/tinyagents-harness/src/run_queue/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,22 @@ impl<T> RunQueue<T> {
}
}

/// 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<T>) {
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<T> {
let mut inner = self.inner.lock().await;
Expand Down
80 changes: 46 additions & 34 deletions crates/tinyagents-session/src/transcript/spend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 => {
Expand All @@ -133,50 +136,59 @@ 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();
}
}
}
let spend = spend_from_usages(usages, display.meta.model.as_ref());
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
Expand Down
33 changes: 33 additions & 0 deletions crates/tinyagents-session/src/transcript/spend_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Loading