diff --git a/crates/cli/src/adapters/mod.rs b/crates/cli/src/adapters/mod.rs index ce60db69d..1a5349455 100644 --- a/crates/cli/src/adapters/mod.rs +++ b/crates/cli/src/adapters/mod.rs @@ -6,6 +6,7 @@ pub(crate) mod codex; pub(crate) mod hermes; use axum::http::HeaderMap; +use nemo_relay::api::scope::COMPACTION_EVENT_NAME; use serde_json::{Map, Value, json}; use uuid::Uuid; @@ -895,7 +896,7 @@ fn classify_primary( fallback_session_id, )) } - "precompact" | "compaction" => { + "precompact" | "postcompact" | COMPACTION_EVENT_NAME => { NormalizedEvent::Compaction(common_session_event_with_fallback( payload, headers, diff --git a/crates/cli/tests/coverage/adapters_tests.rs b/crates/cli/tests/coverage/adapters_tests.rs index 685f3c23f..b7a70fdcc 100644 --- a/crates/cli/tests/coverage/adapters_tests.rs +++ b/crates/cli/tests/coverage/adapters_tests.rs @@ -930,6 +930,8 @@ fn normalizes_mark_style_events_and_header_session_ids() { ("UserPromptSubmit", "prompt"), ("afterAgentResponse", "response"), ("PreCompact", "compact"), + ("PostCompact", "compact"), + (COMPACTION_EVENT_NAME, "compact"), ("Notification", "notification"), ("Unrecognized.Event", "hook"), ] { diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 0dd76ef85..7c587bf5e 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -29,7 +29,7 @@ use crate::api::shared::{ resolve_parent_uuid, run_request_intercepts_with_codec_and_recorder, sanitize_event_with_scope_stack, snapshot_event_subscribers, }; -use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::request::{AnnotatedLlmRequest, Message}; use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider}; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use crate::error::{FlowError, Result}; @@ -298,6 +298,89 @@ fn create_llm_handle(params: CreateLlmHandleParams<'_>) -> Result { Ok(state.create_llm_handle(params)) } +fn request_turn_projection_needed( + items: &[T], + is_user: &impl Fn(&T) -> bool, + is_instruction: &impl Fn(&T) -> bool, +) -> bool { + let Some(last_index) = items.len().checked_sub(1) else { + return false; + }; + match items.iter().rposition(is_user) { + Some(start) => items[..start].iter().any(|item| !is_instruction(item)), + None => items + .iter() + .enumerate() + .any(|(index, item)| index != last_index && !is_instruction(item)), + } +} + +fn retain_current_request_turn( + items: &mut Vec, + is_user: impl Fn(&T) -> bool, + is_instruction: impl Fn(&T) -> bool, +) -> bool { + if !request_turn_projection_needed(items, &is_user, &is_instruction) { + return false; + } + let last_index = items.len() - 1; + let Some(start) = items.iter().rposition(is_user) else { + let mut index = 0; + items.retain(|item| { + let retain = index == last_index || is_instruction(item); + index += 1; + retain + }); + return true; + }; + let mut current_turn = items.split_off(start); + items.retain(is_instruction); + items.append(&mut current_turn); + true +} + +fn project_llm_request_to_current_user_turn( + request: &mut LlmRequest, + annotated_request: &mut Option>, + request_codec: Option<&dyn LlmCodec>, +) { + let Some(annotation) = annotated_request.as_mut() else { + return; + }; + if !request_turn_projection_needed( + &annotation.messages, + &|message| matches!(message, Message::User { .. }), + &|message| matches!(message, Message::System { .. }), + ) { + return; + } + let original_annotation = request_codec.map(|_| Arc::clone(annotation)); + let projected = limit_annotated_request_history_to_current_user_turn(Arc::make_mut(annotation)); + debug_assert!(projected); + if let Some(codec) = request_codec { + match codec.encode(annotation, request) { + Ok(encoded) => *request = encoded, + Err(error) => { + eprintln!( + "nemo_relay: LLM request projection encode failed; preserving full event history: {error}" + ); + *annotation = original_annotation + .expect("codec-backed projection should preserve the original annotation") + } + } + } +} + +fn limit_annotated_request_history_to_current_user_turn( + annotated_request: &mut AnnotatedLlmRequest, +) -> bool { + retain_current_request_turn( + &mut annotated_request.messages, + |message| matches!(message, Message::User { .. }), + |message| matches!(message, Message::System { .. }), + ) +} + fn emit_llm_start( handle: &LlmHandle, request: &LlmRequest, @@ -339,9 +422,9 @@ fn emit_llm_start_with_subscribers( .map_err(|error| FlowError::Internal(error.to_string()))?; state.llm_sanitize_request_entries(&scope_locals) }; - let sanitized_request = + let mut sanitized_request = NemoRelayContextState::llm_sanitize_request_snapshot_chain(request.clone(), &entries); - let annotated_request = match request_codec { + let mut annotated_request = match request_codec { Some(codec) if sanitized_request.headers != request.headers || sanitized_request.content != request.content => @@ -350,6 +433,18 @@ fn emit_llm_start_with_subscribers( } _ => annotated_request, }; + let scope_stack = handle.captured_scope_stack(); + let agent_is_fresh = { + let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); + scope_guard.take_agent_freshness(handle.parent_uuid) + }; + if !agent_is_fresh { + project_llm_request_to_current_user_turn( + &mut sanitized_request, + &mut annotated_request, + request_codec, + ); + } let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null); let event = { let context = global_context(); @@ -358,7 +453,7 @@ fn emit_llm_start_with_subscribers( .map_err(|error| FlowError::Internal(error.to_string()))?; state.build_llm_start_event(handle, Some(input), annotated_request) }; - if let Some(event) = sanitize_event_with_scope_stack(event, handle.captured_scope_stack()) { + if let Some(event) = sanitize_event_with_scope_stack(event, scope_stack) { NemoRelayContextState::emit_event(&event, subscribers); } Ok(()) @@ -486,7 +581,10 @@ fn emit_optimization_marks_with( /// /// # Notes /// Sanitize-request guardrails affect only the emitted start-event payload, not -/// the caller-owned [`LlmRequest`]. +/// the caller-owned [`LlmRequest`]. When the owning agent is not fresh, the +/// emitted request annotation is limited to the current user turn. Managed +/// calls with a request codec also apply that projection to the event input, +/// without changing the request used for provider execution. pub fn llm_call(params: LlmCallParams<'_>) -> Result { let handle_params = CreateLlmHandleParams::builder() .name(params.name) diff --git a/crates/core/src/api/runtime/scope_stack.rs b/crates/core/src/api/runtime/scope_stack.rs index ca70e93cb..f62f7191b 100644 --- a/crates/core/src/api/runtime/scope_stack.rs +++ b/crates/core/src/api/runtime/scope_stack.rs @@ -9,6 +9,7 @@ //! context into worker threads. use std::cell::RefCell; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, RwLock}; use uuid::Uuid; @@ -21,12 +22,15 @@ use crate::registry::{RegistryEntry, SortedRegistry}; /// Mutable stack of active scopes plus their scope-local registries. /// -/// The stack always contains an implicit root scope. Additional scopes are -/// pushed as the public API opens lifecycle spans and removed when those spans -/// close. +/// The stack always contains an implicit root agent scope. It owns freshness +/// for work that is not nested under an explicit agent; non-agent scopes inherit +/// their nearest agent's freshness instead of creating a separate budget. +/// Additional scopes are pushed as the public API opens lifecycle spans and +/// removed when those spans close. pub struct ScopeStack { stack: Vec, - scope_registries: std::collections::HashMap, + scope_registries: HashMap, + fresh_agents: HashSet, } impl ScopeStack { @@ -40,9 +44,11 @@ impl ScopeStack { .name("root") .scope_type(ScopeType::Agent) .build(); + let root_uuid = root.uuid; Self { stack: vec![root], - scope_registries: std::collections::HashMap::new(), + scope_registries: HashMap::new(), + fresh_agents: HashSet::from([root_uuid]), } } @@ -51,6 +57,9 @@ impl ScopeStack { /// # Parameters /// - `handle`: Scope handle to make the new top-most active scope. pub fn push(&mut self, handle: ScopeHandle) { + if matches!(handle.scope_type, ScopeType::Agent) { + self.fresh_agents.insert(handle.uuid); + } self.stack.push(handle); } @@ -134,6 +143,7 @@ impl ScopeStack { )); } self.scope_registries.remove(uuid); + self.fresh_agents.remove(uuid); return Ok(self .stack .pop() @@ -149,6 +159,34 @@ impl ScopeStack { Err(FlowError::NotFound("scope handle not found".into())) } + fn owning_agent_uuid(&self, parent_uuid: Option) -> Uuid { + let search_end = parent_uuid + .and_then(|parent_uuid| { + self.stack + .iter() + .position(|scope| scope.uuid == parent_uuid) + }) + .map_or(self.stack.len(), |index| index + 1); + self.stack[..search_end] + .iter() + .rev() + .find(|scope| matches!(scope.scope_type, ScopeType::Agent)) + .map(|scope| scope.uuid) + .expect("scope stack should always contain an owning agent") + } + + /// Return whether the owning agent is fresh, then mark it stale. + pub(crate) fn take_agent_freshness(&mut self, parent_uuid: Option) -> bool { + let uuid = self.owning_agent_uuid(parent_uuid); + self.fresh_agents.remove(&uuid) + } + + /// Mark the agent that owns a compaction event as fresh. + pub(crate) fn mark_agent_fresh(&mut self, parent_uuid: Option) { + let uuid = self.owning_agent_uuid(parent_uuid); + self.fresh_agents.insert(uuid); + } + /// Get or create the scope-local registries for an active scope. /// /// # Parameters @@ -210,6 +248,7 @@ impl std::fmt::Debug for ScopeStack { f.debug_struct("ScopeStack") .field("stack", &self.stack) .field("scope_registries_count", &self.scope_registries.len()) + .field("fresh_agent_count", &self.fresh_agents.len()) .finish() } } diff --git a/crates/core/src/api/scope.rs b/crates/core/src/api/scope.rs index b00a97fbf..7546dc92e 100644 --- a/crates/core/src/api/scope.rs +++ b/crates/core/src/api/scope.rs @@ -19,6 +19,9 @@ use uuid::Uuid; pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; +/// Canonical mark-event name used to indicate agent context compaction. +pub const COMPACTION_EVENT_NAME: &str = "compaction"; + /// Runtime-owned handle identifying an active or completed scope. #[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)] #[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] @@ -328,11 +331,18 @@ pub fn pop_scope(params: PopScopeParams<'_>) -> Result<()> { pub fn event(params: EmitMarkEventParams<'_>) -> Result<()> { ensure_runtime_owner()?; let parent_uuid = resolve_parent_uuid(params.parent); + let scope_stack = current_scope_stack(); let (event, subscribers) = { - let scope_stack = current_scope_stack(); - let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); - let scope_subscribers = scope_guard.collect_scope_local_subscribers(); - let subscribers = snapshot_event_subscribers(scope_subscribers)?; + let subscribers = if params.name == COMPACTION_EVENT_NAME { + let mut scope_guard = scope_stack.write().expect("scope stack lock poisoned"); + let subscribers = + snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())?; + scope_guard.mark_agent_fresh(parent_uuid); + subscribers + } else { + let scope_guard = scope_stack.read().expect("scope stack lock poisoned"); + snapshot_event_subscribers(scope_guard.collect_scope_local_subscribers())? + }; let context = global_context(); let state = context .read() diff --git a/crates/core/tests/unit/context_tests.rs b/crates/core/tests/unit/context_tests.rs index 7b15dda3a..07e4ea6b6 100644 --- a/crates/core/tests/unit/context_tests.rs +++ b/crates/core/tests/unit/context_tests.rs @@ -95,6 +95,46 @@ fn scope_stack_rejects_removing_non_top_or_root_scopes() { assert!(err.to_string().contains("root scope cannot be removed")); } +#[test] +fn scope_stack_tracks_freshness_per_owning_agent() { + let mut stack = ScopeStack::new(); + let root_uuid = stack.root_uuid(); + assert!(stack.take_agent_freshness(Some(root_uuid))); + assert!(!stack.take_agent_freshness(Some(root_uuid))); + + let parent = ScopeHandle::builder() + .name("parent-agent") + .scope_type(ScopeType::Agent) + .parent_uuid(root_uuid) + .build(); + let parent_uuid = parent.uuid; + stack.push(parent); + let turn = ScopeHandle::builder() + .name("turn") + .scope_type(ScopeType::Custom) + .parent_uuid(parent_uuid) + .build(); + let turn_uuid = turn.uuid; + stack.push(turn); + assert!(stack.take_agent_freshness(Some(turn_uuid))); + + let child = ScopeHandle::builder() + .name("child-agent") + .scope_type(ScopeType::Agent) + .parent_uuid(turn_uuid) + .build(); + let child_uuid = child.uuid; + stack.push(child); + assert!(stack.take_agent_freshness(Some(child_uuid))); + stack.mark_agent_fresh(Some(turn_uuid)); + assert!(!stack.take_agent_freshness(Some(child_uuid))); + + stack.remove(&child_uuid).unwrap(); + assert!(stack.take_agent_freshness(Some(turn_uuid))); + stack.remove(&turn_uuid).unwrap(); + stack.remove(&parent_uuid).unwrap(); +} + #[test] fn merge_helpers_preserve_global_and_scope_local_priority_order() { let mut global_guardrails = SortedRegistry::new(); diff --git a/crates/core/tests/unit/llm_api_tests.rs b/crates/core/tests/unit/llm_api_tests.rs index 1e5c94350..625d5bf7c 100644 --- a/crates/core/tests/unit/llm_api_tests.rs +++ b/crates/core/tests/unit/llm_api_tests.rs @@ -5,20 +5,29 @@ #![allow(clippy::await_holding_lock)] -use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Barrier, Mutex}; use serde_json::json; use tokio_stream::StreamExt; use super::{ - LlmCallExecuteParams, LlmHandle, LlmRequest, LlmStreamCallExecuteParams, - emit_optimization_marks_with, llm_call_execute, llm_stream_call_execute, + LlmCallExecuteParams, LlmCallParams, LlmHandle, LlmRequest, LlmStreamCallExecuteParams, + emit_optimization_marks_with, llm_call, llm_call_execute, llm_stream_call_execute, + project_llm_request_to_current_user_turn, }; -use crate::api::event::ScopeCategory; +use crate::api::event::{Event, ScopeCategory}; use crate::api::optimization::finalize_optimization_summary; use crate::api::runtime::LlmJsonStream; -use crate::api::runtime::{NemoRelayContextState, global_context}; +use crate::api::runtime::{ + NemoRelayContextState, create_scope_stack, global_context, set_thread_scope_stack, +}; +use crate::api::scope::{COMPACTION_EVENT_NAME, EmitMarkEventParams, event}; +use crate::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use crate::codec::openai_chat::OpenAIChatCodec; +use crate::codec::request::{AnnotatedLlmRequest, Message}; +use crate::codec::traits::LlmCodec; use crate::error::FlowError; use crate::json::Json; use crate::{codec::optimization::LlmOptimizationContribution, codec::response::PricingResolver}; @@ -42,6 +51,580 @@ fn request() -> LlmRequest { } } +fn multi_turn_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ + "model": "demo", + "messages": [ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "earlier question"}, + {"role": "assistant", "content": "earlier answer"}, + {"role": "user", "content": "latest question"} + ] + }), + } +} + +fn multi_turn_annotation() -> Arc { + Arc::new(OpenAIChatCodec.decode(&multi_turn_request()).unwrap()) +} + +struct ProjectionFailingCodec { + projection_attempts: Arc, +} + +impl LlmCodec for ProjectionFailingCodec { + fn decode(&self, request: &LlmRequest) -> crate::error::Result { + OpenAIChatCodec.decode(request) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> crate::error::Result { + let original_messages = original.content["messages"].as_array().map_or(0, Vec::len); + if annotated.messages.len() < original_messages { + self.projection_attempts.fetch_add(1, Ordering::Relaxed); + return Err(FlowError::Internal("projection encode failed".into())); + } + OpenAIChatCodec.encode(annotated, original) + } +} + +fn emit_compaction() { + event( + EmitMarkEventParams::builder() + .name(COMPACTION_EVENT_NAME) + .build(), + ) + .unwrap(); +} + +#[test] +fn freshness_culls_annotations_and_repeated_compactions_are_idempotent() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let raw_request = multi_turn_request(); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "freshness-culling", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let emit = |name| { + llm_call( + LlmCallParams::builder() + .name(name) + .request(&raw_request) + .annotated_request(multi_turn_annotation()) + .build(), + ) + .unwrap(); + }; + emit("fresh-start"); + emit("stale-start"); + // PreCompact and PostCompact both normalize to this canonical mark. + emit_compaction(); + emit_compaction(); + emit("post-compaction-start"); + emit("post-compaction-stale"); + + flush_subscribers().unwrap(); + assert!(deregister_subscriber("freshness-culling").unwrap()); + let starts = events + .lock() + .unwrap() + .iter() + .filter(|event| event.scope_category() == Some(ScopeCategory::Start)) + .map(|event| { + assert_eq!(event.input().unwrap()["content"], raw_request.content); + ( + event.name().to_string(), + event.annotated_request().unwrap().messages.clone(), + ) + }) + .collect::>(); + assert_eq!(starts[0].0, "fresh-start"); + assert_eq!(starts[0].1.len(), 4); + assert_eq!(starts[1].0, "stale-start"); + assert_eq!(starts[1].1.len(), 2); + assert!(matches!(starts[1].1[0], Message::System { .. })); + assert!(matches!(starts[1].1[1], Message::User { .. })); + assert_eq!(starts[2].0, "post-compaction-start"); + assert_eq!(starts[2].1.len(), 4); + assert_eq!(starts[3].0, "post-compaction-stale"); + assert_eq!(starts[3].1.len(), 2); +} + +#[test] +fn request_projection_handles_history_edge_cases() { + let project = |messages: Json| { + let annotated: AnnotatedLlmRequest = + serde_json::from_value(json!({"messages": messages})).unwrap(); + let mut annotated = Some(Arc::new(annotated)); + let mut request = LlmRequest { + headers: serde_json::Map::new(), + content: Json::Null, + }; + project_llm_request_to_current_user_turn(&mut request, &mut annotated, None); + annotated.unwrap().messages.clone() + }; + + assert!(project(json!([])).is_empty()); + + let no_user = project(json!([ + {"role": "assistant", "content": "tool request"}, + {"role": "tool", "tool_call_id": "call-1", "content": "tool result"} + ])); + assert_eq!(no_user.len(), 1); + assert!(matches!(no_user[0], Message::Tool { .. })); + + let instructions_only = project(json!([ + {"role": "system", "content": "first instruction"}, + {"role": "system", "content": "second instruction"} + ])); + assert_eq!(instructions_only.len(), 2); + assert!( + instructions_only + .iter() + .all(|message| matches!(message, Message::System { .. })) + ); + + let current_turn = project(json!([ + {"role": "system", "content": "instructions"}, + {"role": "user", "content": "earlier question"}, + {"role": "assistant", "content": "earlier answer"}, + {"role": "user", "content": "latest question"}, + {"role": "assistant", "content": null, "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": {"name": "search", "arguments": "{}"} + }]}, + {"role": "tool", "tool_call_id": "call-1", "content": "latest result"} + ])); + assert_eq!(current_turn.len(), 4); + assert!(matches!(current_turn[0], Message::System { .. })); + assert!(matches!(current_turn[1], Message::User { .. })); + assert!(matches!(current_turn[2], Message::Assistant { .. })); + assert!(matches!(current_turn[3], Message::Tool { .. })); + + let original = Arc::new( + serde_json::from_value::(json!({ + "messages": [{"role": "system", "content": "instructions"}] + })) + .unwrap(), + ); + let mut annotated = Some(original.clone()); + let mut request = LlmRequest { + headers: serde_json::Map::new(), + content: Json::Null, + }; + project_llm_request_to_current_user_turn(&mut request, &mut annotated, Some(&OpenAIChatCodec)); + assert!(Arc::ptr_eq(annotated.as_ref().unwrap(), &original)); +} + +#[test] +fn managed_and_streaming_calls_cull_event_inputs_and_annotations_with_real_codec() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "managed-streaming-freshness", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let request = multi_turn_request(); + let codec: Arc = Arc::new(OpenAIChatCodec); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + for name in ["managed-fresh", "managed-stale"] { + let response = llm_call_execute( + LlmCallExecuteParams::builder() + .name(name) + .request(request.clone()) + .func(Arc::new(|request| { + Box::pin(async move { + Ok(json!({ + "provider_message_count": request.content["messages"] + .as_array() + .unwrap() + .len() + })) + }) + })) + .codec(codec.clone()) + .build(), + ) + .await + .unwrap(); + assert_eq!(response["provider_message_count"], 4); + } + + emit_compaction(); + for name in ["stream-fresh", "stream-stale"] { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name(name) + .request(request.clone()) + .func(Arc::new(|request| { + Box::pin(async move { + assert_eq!(request.content["messages"].as_array().unwrap().len(), 4); + Ok( + Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) + as LlmJsonStream, + ) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .codec(codec.clone()) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + } + }); + + flush_subscribers().unwrap(); + assert!(deregister_subscriber("managed-streaming-freshness").unwrap()); + let events = events.lock().unwrap(); + for (name, expected_messages) in [ + ("managed-fresh", 4), + ("managed-stale", 2), + ("stream-fresh", 4), + ("stream-stale", 2), + ] { + let start = events + .iter() + .find(|event| { + event.name() == name && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap_or_else(|| panic!("missing LLM start event {name}")); + assert_eq!( + start.input().unwrap()["content"]["messages"] + .as_array() + .unwrap() + .len(), + expected_messages, + "unexpected raw event history for {name}" + ); + assert_eq!( + start.annotated_request().unwrap().messages.len(), + expected_messages, + "unexpected annotation history for {name}" + ); + } +} + +#[test] +fn projection_encode_failures_do_not_block_managed_or_streaming_calls() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "projection-encode-failure", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let projection_attempts = Arc::new(AtomicUsize::new(0)); + let codec: Arc = Arc::new(ProjectionFailingCodec { + projection_attempts: projection_attempts.clone(), + }); + let request = multi_turn_request(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + for name in ["managed-encode-fresh", "managed-encode-stale"] { + let response = llm_call_execute( + LlmCallExecuteParams::builder() + .name(name) + .request(request.clone()) + .func(Arc::new(|request| { + Box::pin(async move { + Ok(json!({ + "provider_message_count": request.content["messages"] + .as_array() + .unwrap() + .len() + })) + }) + })) + .codec(codec.clone()) + .build(), + ) + .await + .unwrap(); + assert_eq!(response["provider_message_count"], 4); + } + + emit_compaction(); + for name in ["stream-encode-fresh", "stream-encode-stale"] { + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name(name) + .request(request.clone()) + .func(Arc::new(|request| { + Box::pin(async move { + assert_eq!(request.content["messages"].as_array().unwrap().len(), 4); + Ok( + Box::pin(tokio_stream::iter(vec![Ok(json!({"chunk": true}))])) + as LlmJsonStream, + ) + }) + })) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .codec(codec.clone()) + .build(), + ) + .await + .unwrap(); + while let Some(chunk) = stream.next().await { + chunk.unwrap(); + } + } + }); + + assert_eq!(projection_attempts.load(Ordering::Relaxed), 2); + flush_subscribers().unwrap(); + assert!(deregister_subscriber("projection-encode-failure").unwrap()); + let events = events.lock().unwrap(); + for name in [ + "managed-encode-fresh", + "managed-encode-stale", + "stream-encode-fresh", + "stream-encode-stale", + ] { + let start = events + .iter() + .find(|event| { + event.name() == name && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap_or_else(|| panic!("missing LLM start event {name}")); + assert_eq!( + start.input().unwrap()["content"]["messages"] + .as_array() + .unwrap() + .len(), + 4, + "encode failure should preserve the full event input for {name}" + ); + assert_eq!( + start.annotated_request().unwrap().messages.len(), + 4, + "encode failure should preserve the full annotation for {name}" + ); + } +} + +#[test] +fn nested_agents_track_freshness_independently_end_to_end() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "nested-agent-freshness", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let request = multi_turn_request(); + let emit = |name| { + llm_call( + LlmCallParams::builder() + .name(name) + .request(&request) + .annotated_request(multi_turn_annotation()) + .build(), + ) + .unwrap(); + }; + + let parent = push_scope( + PushScopeParams::builder() + .name("parent-agent") + .scope_type(ScopeType::Agent) + .build(), + ) + .unwrap(); + emit("parent-fresh"); + emit("parent-stale"); + + let child = push_scope( + PushScopeParams::builder() + .name("child-agent") + .scope_type(ScopeType::Agent) + .build(), + ) + .unwrap(); + emit("child-fresh"); + emit("child-stale"); + emit_compaction(); + emit_compaction(); + emit("child-after-compaction"); + emit("child-after-compaction-stale"); + pop_scope(PopScopeParams::builder().handle_uuid(&child.uuid).build()).unwrap(); + + emit("parent-after-child-compaction"); + emit_compaction(); + emit("parent-after-compaction"); + pop_scope(PopScopeParams::builder().handle_uuid(&parent.uuid).build()).unwrap(); + + flush_subscribers().unwrap(); + assert!(deregister_subscriber("nested-agent-freshness").unwrap()); + let events = events.lock().unwrap(); + for (name, expected_messages) in [ + ("parent-fresh", 4), + ("parent-stale", 2), + ("child-fresh", 4), + ("child-stale", 2), + ("child-after-compaction", 4), + ("child-after-compaction-stale", 2), + ("parent-after-child-compaction", 2), + ("parent-after-compaction", 4), + ] { + let start = events + .iter() + .find(|event| { + event.name() == name && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap_or_else(|| panic!("missing LLM start event {name}")); + assert_eq!( + start.annotated_request().unwrap().messages.len(), + expected_messages, + "unexpected annotation history for {name}" + ); + } +} + +#[test] +fn non_agent_scopes_share_the_implicit_root_freshness_budget() { + let _guard = lock_global_runtime(); + reset_global(); + set_thread_scope_stack(create_scope_stack()); + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "implicit-root-freshness", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let request = multi_turn_request(); + for (scope_name, scope_type, event_name) in [ + ("request-a", ScopeType::Custom, "root-fresh"), + ("request-b", ScopeType::Function, "root-stale"), + ] { + let scope = push_scope( + PushScopeParams::builder() + .name(scope_name) + .scope_type(scope_type) + .build(), + ) + .unwrap(); + llm_call( + LlmCallParams::builder() + .name(event_name) + .request(&request) + .annotated_request(multi_turn_annotation()) + .build(), + ) + .unwrap(); + pop_scope(PopScopeParams::builder().handle_uuid(&scope.uuid).build()).unwrap(); + } + + flush_subscribers().unwrap(); + assert!(deregister_subscriber("implicit-root-freshness").unwrap()); + let events = events.lock().unwrap(); + for (name, expected_messages) in [("root-fresh", 4), ("root-stale", 2)] { + let start = events + .iter() + .find(|event| { + event.name() == name && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap_or_else(|| panic!("missing LLM start event {name}")); + assert_eq!( + start.annotated_request().unwrap().messages.len(), + expected_messages, + "non-agent scopes should inherit the implicit root freshness for {name}" + ); + } +} + +#[test] +fn concurrent_starts_consume_freshness_exactly_once_without_ordering() { + let _guard = lock_global_runtime(); + reset_global(); + + let shared_stack = create_scope_stack(); + set_thread_scope_stack(shared_stack.clone()); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "concurrent-freshness", + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let barrier = Arc::new(Barrier::new(3)); + let mut workers = Vec::new(); + for name in ["concurrent-a", "concurrent-b"] { + let stack = shared_stack.clone(); + let barrier = barrier.clone(); + workers.push(std::thread::spawn(move || { + set_thread_scope_stack(stack); + barrier.wait(); + let request = multi_turn_request(); + llm_call( + LlmCallParams::builder() + .name(name) + .request(&request) + .annotated_request(multi_turn_annotation()) + .build(), + ) + .unwrap(); + })); + } + barrier.wait(); + for worker in workers { + worker.join().unwrap(); + } + + flush_subscribers().unwrap(); + assert!(deregister_subscriber("concurrent-freshness").unwrap()); + let events = events.lock().unwrap(); + let mut history_lengths = events + .iter() + .filter(|event| { + event.name().starts_with("concurrent-") + && event.scope_category() == Some(ScopeCategory::Start) + }) + .map(|event| event.annotated_request().unwrap().messages.len()) + .collect::>(); + history_lengths.sort_unstable(); + assert_eq!(history_lengths, vec![2, 4]); +} + #[test] fn rejected_optimization_mark_queue_keeps_cursor_and_summary_evidence() { let _guard = lock_global_runtime();