diff --git a/Cargo.lock b/Cargo.lock index 9e3b3e1..7223e79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2517,7 +2517,7 @@ dependencies = [ [[package]] name = "kit" -version = "0.1.93" +version = "0.1.94" dependencies = [ "a2a-protocol-client", "a2a-protocol-server", diff --git a/Cargo.toml b/Cargo.toml index ce52dcf..18749dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "kit" -version = "0.1.93" +version = "0.1.94" edition = "2024" rust-version = "1.94.0" publish = false diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 2d60038..54129c0 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -16,7 +16,7 @@ use agentkit_acp::{ }, }; use agentkit_core::{ - CancellationController, FinishReason, Item, ItemKind, Part, SessionId, ToolOutput, + CancellationController, FinishReason, Item, ItemKind, Part, SessionId, ToolOutput, Usage, }; use agentkit_loop::{ AgentEvent, LoopDriver, LoopError, LoopInterrupt, LoopObserver, LoopStep, ModelSession, @@ -257,6 +257,21 @@ struct ResponseReplacementObserver { session_id: wire::SessionId, } +fn usage_update(usage: &Usage) -> Option { + let tokens = usage.tokens.as_ref()?; + let used = tokens.input_tokens.checked_add(tokens.output_tokens)?; + let size = [ + "context_window", + "context_window_tokens", + "model.context_window", + "model.context_length", + "openrouter.context_length", + ] + .iter() + .find_map(|key| usage.metadata.get(*key).and_then(|value| value.as_u64()))?; + Some(wire::UsageUpdate::new(used, size)) +} + impl ResponseReplacementObserver { fn new( inner: AcpIntegration, @@ -286,6 +301,19 @@ where S: AcpSessionUpdateSink + Clone, { fn handle_event(&self, event: ObservedEvent) { + if let AgentEvent::UsageUpdated(usage) = &event.event { + let Some(update) = usage_update(usage) else { + return; + }; + let notification = wire::UpdateSessionNotification::new( + self.session_id.clone(), + wire::SessionUpdate::UsageUpdate(update), + ); + if let Err(error) = self.sink.update(notification) { + tracing::debug!(%error, "failed to queue ACP v2 usage update"); + } + return; + } if matches!( &event.event, AgentEvent::ContentDelta(delta) if crate::response_attempt::is_marker(delta) @@ -1762,6 +1790,42 @@ mod tests { } } + #[test] + fn observer_reports_usage_with_a_known_context_window() { + let recording = RecordingSink::default(); + let sink = ResponseReplacementSink::new(recording.clone()); + let observer = ResponseReplacementObserver::new( + AcpIntegration::default(), + sink, + wire::SessionId::new("usage-session"), + ); + let loop_session_id = SessionId::new("usage-loop"); + let emit = |usage| { + observer.handle_event(ObservedEvent { + session_id: Arc::new(loop_session_id.clone()), + event: AgentEvent::UsageUpdated(usage), + }); + }; + + emit(agentkit_core::Usage::new(agentkit_core::TokenUsage::new( + 10, 2, + ))); + emit( + agentkit_core::Usage::new(agentkit_core::TokenUsage::new(50_000, 3_000)).with_metadata( + MetadataMap::from([("context_window".into(), json!(272_000))]), + ), + ); + + let updates = recording.updates.lock().unwrap(); + assert_eq!(updates.len(), 1); + let wire::SessionUpdate::UsageUpdate(usage) = &updates[0].update else { + panic!("expected usage update, got {:?}", updates[0].update); + }; + assert_eq!(usage.used, 53_000); + assert_eq!(usage.size, 272_000); + assert!(usage.cost.is_none()); + } + #[test] fn response_replacement_clears_and_remaps_message_ids_in_new_chunk_order() { let integration = AcpIntegration::default();