1use std::sync::atomic::{AtomicU16, Ordering};
42use std::sync::Arc;
43
44use serde_json::{json, Value};
45use tokio::sync::mpsc;
46use tokio_util::sync::CancellationToken;
47
48use super::api::{ApiMethods, ApiOptions, KeepWarmMode, KeepWarmRequest};
49use super::helpers::HelperMethods;
50use super::{Runtime, SessionEvent, StreamEvent};
51use crate::{Result, RuntimeError, SharedMessage};
52
53pub const KEEPWARM_OUTPUT_TOKENS: u64 = 1;
56
57pub const KEEPWARM_PING_TEXT: &str = "[automated prompt-cache keep-warm ping, not from the user] \
61 This needs no reasoning. Reply with exactly: ok";
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub struct CacheAnchor {
69 pub mono: tokio::time::Instant,
71 pub wall: std::time::SystemTime,
74}
75
76impl CacheAnchor {
77 pub fn now() -> Self {
78 Self {
79 mono: tokio::time::Instant::now(),
80 wall: std::time::SystemTime::now(),
81 }
82 }
83}
84
85#[derive(Debug, Default)]
89pub struct CacheAnchorCell(std::sync::Mutex<Option<CacheAnchor>>);
90
91impl CacheAnchorCell {
92 pub fn stamp(&self, at: CacheAnchor) {
93 let mut slot = self
94 .0
95 .lock()
96 .unwrap_or_else(std::sync::PoisonError::into_inner);
97 if slot.is_none_or(|prev| prev.mono <= at.mono) {
98 *slot = Some(at);
99 }
100 }
101
102 pub fn get(&self) -> Option<CacheAnchor> {
103 *self
104 .0
105 .lock()
106 .unwrap_or_else(std::sync::PoisonError::into_inner)
107 }
108}
109
110#[derive(Debug, Clone, Default, PartialEq)]
112pub struct KeepWarmUsage {
113 pub input_tokens: u64,
114 pub output_tokens: u64,
115 pub cache_read_tokens: u64,
116 pub cache_write_tokens: u64,
117 pub cache_write_5m: Option<u64>,
118 pub cache_write_1h: Option<u64>,
119 pub model: Option<String>,
121}
122
123impl KeepWarmUsage {
124 pub(crate) fn add(&mut self, event: &StreamEvent) {
125 if let StreamEvent::Session(SessionEvent::Usage {
126 input_tokens,
127 output_tokens,
128 cache_read_input_tokens,
129 cache_creation_input_tokens,
130 cache_creation_5m,
131 cache_creation_1h,
132 model,
133 }) = event
134 {
135 self.input_tokens += input_tokens;
136 self.output_tokens += output_tokens;
137 self.cache_read_tokens += cache_read_input_tokens;
138 self.cache_write_tokens += cache_creation_input_tokens;
139 if let Some(n) = cache_creation_5m {
140 *self.cache_write_5m.get_or_insert(0) += n;
141 }
142 if let Some(n) = cache_creation_1h {
143 *self.cache_write_1h.get_or_insert(0) += n;
144 }
145 if model.is_some() {
146 self.model = model.clone();
147 }
148 }
149 }
150
151 pub fn cost_usd(&self, fallback_model: &str) -> f64 {
154 crate::pricing::calculate_cost_optional_split(
155 self.model.as_deref().unwrap_or(fallback_model),
156 self.input_tokens,
157 self.output_tokens,
158 self.cache_read_tokens,
159 self.cache_write_tokens,
160 self.cache_write_5m,
161 self.cache_write_1h,
162 )
163 }
164}
165
166#[derive(Debug, Clone, PartialEq)]
168pub enum PingOutcome {
169 Sent,
171 PrefixChanged { expected: u64, actual: u64 },
174}
175
176#[derive(Debug, Clone, PartialEq)]
180pub struct PingReport {
181 pub result: std::result::Result<PingOutcome, PingFailure>,
182 pub usage: KeepWarmUsage,
183 pub notices: Vec<String>,
187}
188
189#[derive(Debug, Clone, Default)]
197pub struct PingTap {
198 pub sent: Arc<std::sync::atomic::AtomicBool>,
199 pub connect_failed: Arc<std::sync::atomic::AtomicBool>,
200 pub status: Arc<AtomicU16>,
201 pub usage: Arc<std::sync::Mutex<KeepWarmUsage>>,
202}
203
204impl PingTap {
205 pub fn was_sent(&self) -> bool {
206 self.sent.load(Ordering::Relaxed)
207 }
208
209 pub fn usage(&self) -> KeepWarmUsage {
210 self.usage
211 .lock()
212 .unwrap_or_else(std::sync::PoisonError::into_inner)
213 .clone()
214 }
215
216 pub fn bill(&self, model: &str, worst_usd: f64) -> PingBill {
221 let usage = self.usage();
222 let booking = ping_booking(
223 usage != KeepWarmUsage::default(),
224 self.was_sent(),
225 self.connect_failed.load(Ordering::Relaxed),
226 self.status.load(Ordering::Relaxed),
227 );
228 let cost_usd = match booking {
229 PingBooking::Measured => usage.cost_usd(model),
230 PingBooking::WorstCase => worst_usd,
231 PingBooking::Nothing => 0.0,
232 };
233 PingBill {
234 booking,
235 usage,
236 cost_usd,
237 }
238 }
239}
240
241#[derive(Debug, Clone, Copy, PartialEq, Eq)]
243pub enum PingBooking {
244 Measured,
246 WorstCase,
249 Nothing,
255}
256
257pub fn ping_booking(
262 usage_arrived: bool,
263 sent: bool,
264 connect_failed: bool,
265 http_status: u16,
266) -> PingBooking {
267 if usage_arrived {
268 PingBooking::Measured
269 } else if http_status >= 400 {
270 PingBooking::Nothing
271 } else if sent && !(connect_failed && http_status == 0) {
272 PingBooking::WorstCase
273 } else {
274 PingBooking::Nothing
275 }
276}
277
278#[derive(Debug, Clone, PartialEq)]
280pub struct PingBill {
281 pub booking: PingBooking,
282 pub usage: KeepWarmUsage,
283 pub cost_usd: f64,
284}
285
286#[derive(Debug, Clone, PartialEq)]
288pub struct KeepWarmProbe {
289 pub fingerprint: u64,
291 pub max_tokens: u64,
294 pub context_blocks: Option<String>,
299}
300
301#[derive(Debug, Clone, PartialEq, Eq)]
303pub enum PingFailure {
304 Transient(String),
306 Fatal(String),
308 Canceled,
310}
311
312impl std::fmt::Display for PingFailure {
313 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314 match self {
315 PingFailure::Transient(m) => write!(f, "transient: {m}"),
316 PingFailure::Fatal(m) => write!(f, "{m}"),
317 PingFailure::Canceled => f.write_str("canceled"),
318 }
319 }
320}
321
322pub(crate) fn classify_failure(err: &RuntimeError, status: u16) -> PingFailure {
325 if matches!(err, RuntimeError::Canceled) {
326 return PingFailure::Canceled;
327 }
328 if (400..500).contains(&status) && status != 429 {
329 return PingFailure::Fatal(format!("HTTP {status}"));
330 }
331 match err {
332 RuntimeError::Api(_)
335 | RuntimeError::ApiStatus(_)
336 | RuntimeError::Timeout
337 | RuntimeError::NetworkTimeout(_) => PingFailure::Transient(if status == 0 {
338 "network".to_string()
339 } else {
340 format!("HTTP {status}")
341 }),
342 other => PingFailure::Fatal(other.to_string()),
343 }
344}
345
346fn tail_is_markable(last: &Value) -> bool {
350 if last["role"] != "assistant" {
351 return false;
352 }
353 match last["content"].as_array() {
354 Some(blocks) => blocks.last().is_some_and(|b| {
355 b["type"] == "text" && b["text"].as_str().is_some_and(|t| !t.is_empty())
356 }),
357 None => last["content"].as_str().is_some_and(|t| !t.is_empty()),
358 }
359}
360
361pub fn history_ready(history: &[SharedMessage]) -> bool {
364 history.last().is_some_and(|m| tail_is_markable(m))
365}
366
367pub fn reply_prefix_tokens(history: &[SharedMessage]) -> u64 {
374 let Some(last) = history.last() else {
375 return 0;
376 };
377 let kept = match last["content"].as_array() {
378 Some(blocks) => Value::Array(
379 blocks
380 .iter()
381 .filter(|b| !matches!(b["type"].as_str(), Some("thinking" | "redacted_thinking")))
382 .cloned()
383 .collect(),
384 ),
385 None => last["content"].clone(),
386 };
387 super::context::conservative_token_estimate(&kept.to_string())
388}
389
390impl Runtime {
391 pub fn cache_keepwarm_ineligibility(&self) -> Option<String> {
395 if !super::auth::model_is_anthropic(&self.model) {
396 return Some(format!(
397 "model {} is not on the native Anthropic transport",
398 self.model
399 ));
400 }
401 None
402 }
403
404 pub fn cache_anchor(&self) -> Option<CacheAnchor> {
406 self.cache_anchor.get()
407 }
408
409 pub fn cache_ttl_downgraded(&self) -> bool {
412 self.ttl_downgrade_notified.load(Ordering::Relaxed)
413 }
414
415 pub fn cache_1h_honoured(&self) -> bool {
417 self.saw_1h_honored.load(Ordering::Relaxed)
418 }
419
420 pub async fn cache_keepwarm_probe(&self, history: &[SharedMessage]) -> Result<KeepWarmProbe> {
426 let cancel = CancellationToken::new();
427 let tap = PingTap::default();
428 let (value, _) = self
429 .keepwarm_request(history, &cancel, KeepWarmMode::Probe, &tap)
430 .await;
431 let value = value?;
432 let fingerprint = value["keepwarm_prefix"].as_u64().ok_or_else(|| {
433 RuntimeError::Session("keep-warm probe produced no fingerprint".into())
434 })?;
435 Ok(KeepWarmProbe {
436 fingerprint,
437 max_tokens: value["keepwarm_max_tokens"].as_u64().unwrap_or(0),
438 context_blocks: value["keepwarm_context_blocks"]
439 .as_str()
440 .map(str::to_string),
441 })
442 }
443
444 pub async fn cache_keepwarm_ping(
450 &self,
451 history: &[SharedMessage],
452 expect: u64,
453 cancel: &CancellationToken,
454 tap: &PingTap,
455 ) -> PingReport {
456 let (value, notices) = self
457 .keepwarm_request(history, cancel, KeepWarmMode::Send { expect }, tap)
458 .await;
459 let result = match value {
460 Ok(value) => match value.get("keepwarm_prefix").and_then(Value::as_u64) {
461 Some(actual) => Ok(PingOutcome::PrefixChanged {
462 expected: expect,
463 actual,
464 }),
465 None => Ok(PingOutcome::Sent),
466 },
467 Err(e) => Err(classify_failure(&e, tap.status.load(Ordering::Relaxed))),
468 };
469 PingReport {
470 result,
471 usage: tap.usage(),
472 notices,
473 }
474 }
475
476 async fn keepwarm_request(
477 &self,
478 history: &[SharedMessage],
479 cancel: &CancellationToken,
480 mode: KeepWarmMode,
481 tap: &PingTap,
482 ) -> (Result<Value>, Vec<String>) {
483 match self
484 .keepwarm_request_inner(history, cancel, mode, tap)
485 .await
486 {
487 Ok(r) => r,
488 Err(e) => (Err(e), Vec::new()),
489 }
490 }
491
492 async fn keepwarm_request_inner(
495 &self,
496 history: &[SharedMessage],
497 cancel: &CancellationToken,
498 mode: KeepWarmMode,
499 tap: &PingTap,
500 ) -> Result<(Result<Value>, Vec<String>)> {
501 if let Some(reason) = self.cache_keepwarm_ineligibility() {
502 return Err(RuntimeError::Config(format!(
503 "cache keep-warm skipped: {reason}"
504 )));
505 }
506 if !history_ready(history) {
507 return Err(RuntimeError::Config(
508 "cache keep-warm skipped: history does not end with a finished assistant reply"
509 .into(),
510 ));
511 }
512 let (durability_blocked, context_enabled) = {
513 let state = self
514 .continuation
515 .lock()
516 .unwrap_or_else(std::sync::PoisonError::into_inner);
517 (state.durability_blocked, state.enabled())
518 };
519 if durability_blocked {
520 return Err(RuntimeError::Session(
521 "cache keep-warm skipped: context head save is unresolved".into(),
522 ));
523 }
524
525 self.validate_request_preflight().await?;
528 let anthropic_execution_plan = self.authorized_anthropic_plan().await?;
529 if matches!(mode, KeepWarmMode::Send { .. }) {
530 self.refresh_if_needed().await?;
531 }
532
533 let system_prompt = self.effective_system_prompt().await;
535 let system_prompt = if context_enabled {
536 Some(format!(
537 "{}\n\n{}",
538 system_prompt.as_deref().unwrap_or_default(),
539 super::continuation::GUIDANCE
540 ))
541 } else {
542 system_prompt
543 };
544 let session_injection = match self.session_id.as_deref() {
545 Some(id) => self.hook_bus.session_injection_for(id).await,
546 None => None,
547 };
548 let system_prompt = match session_injection {
549 Some(content) => Some(super::stream::wrap_extension_context(
550 system_prompt.as_deref().unwrap_or_default(),
551 &content,
552 )),
553 None => system_prompt,
554 };
555
556 let mut request: Vec<SharedMessage> = history.to_vec();
559 request.push(Arc::new(
560 json!({"role": "user", "content": KEEPWARM_PING_TEXT}),
561 ));
562 let capped;
563 let request_messages: &[SharedMessage] = match super::stream::cap_history_image_bytes(
564 &request,
565 super::stream::HISTORY_IMAGE_BYTE_CAP,
566 ) {
567 Some(c) => {
568 capped = c;
569 &capped
570 }
571 None => &request,
572 };
573
574 let tools_snapshot = self.tools.read().await.clone();
579 let (request_tools_schema, anthropic_tool_split) = if self.progressive_tool_disclosure {
580 let fresh = super::continuation::context_tool_set(
581 self.host_tool_session.clone(),
582 tools_snapshot.catalog(),
583 true,
584 context_enabled,
585 );
586 let retained = self
589 .retained_tool_set
590 .lock()
591 .unwrap_or_else(std::sync::PoisonError::into_inner)
592 .clone();
593 let existing = retained.as_ref().map(|shared| {
594 shared
595 .read()
596 .unwrap_or_else(std::sync::PoisonError::into_inner)
597 });
598 let (mut set, _, _) = super::stream::select_turn_start_tool_set(
599 existing.as_deref(),
600 fresh,
601 tools_snapshot.catalog(),
602 true,
603 );
604 drop(existing);
605 let pending = self
610 .pending_tool_restore
611 .lock()
612 .unwrap_or_else(std::sync::PoisonError::into_inner)
613 .clone();
614 if let Some(carried) = pending {
615 if crate::tools::activation::carry_forward_enabled() {
616 let _ = super::stream::restore_pending_activations(
617 &mut set,
618 &carried,
619 tools_snapshot.catalog(),
620 self.activation_confirm,
621 );
622 }
623 }
624 let schema = tools_snapshot.session_tools_schema(&set).schema;
625 let split = self.deferred_tool_loading.then(|| {
626 let referenced = HelperMethods::referenced_tool_ids(request_messages);
627 Arc::new(tools_snapshot.deferred_tool_split(&set, &referenced))
628 });
629 (Some(Arc::new(schema)), split)
630 } else {
631 (None, None)
632 };
633
634 let options = ApiOptions {
635 memory_backend: None,
636 use_1m_context: self.context_window_override == Some(1_000_000),
637 cache_ttl: self.cache_ttl,
638 ttl_downgrade_notified: self.ttl_downgrade_notified.clone(),
642 saw_1h_honored: self.saw_1h_honored.clone(),
643 credential_source: self.credential_source.clone(),
644 token_cache: self.token_cache.clone(),
645 anthropic_base_url: None,
646 anthropic_stream_client: self.anthropic_stream_client.clone(),
647 anthropic_execution_plan,
648 codex_request_role: self.codex_request_role(),
649 trace: self.trace_ctx.clone(),
651 request_correlation: None,
652 suppress_stream_deltas: true,
653 telemetry: self.telemetry_writer.clone(),
654 tool_session_id: Some(self.host_tool_session.clone()),
655 session_tool_set: None,
656 request_tools_schema: request_tools_schema.clone(),
657 anthropic_tool_split,
658 usage_counters: None,
659 cache_anchor: Some(Arc::clone(&self.cache_anchor)),
662 keepwarm: Some(KeepWarmRequest {
663 mode,
664 status: Arc::clone(&tap.status),
665 sent: Arc::clone(&tap.sent),
666 connect_failed: Arc::clone(&tap.connect_failed),
667 }),
668 cache_diagnosis: self.cache_diagnosis(),
671 };
672
673 let (tx, mut rx) = mpsc::unbounded_channel();
674 let call = ApiMethods::call_api_stream_inner(
675 &self.auth,
676 &self.client,
677 &self.model,
678 &tools_snapshot,
679 &system_prompt,
680 self.thinking_budget,
681 self.reasoning_level(),
682 request_messages,
683 tx,
684 cancel,
685 0,
686 0,
687 &options,
688 self.telemetry_level,
689 );
690 tokio::pin!(call);
691 let mut notices = Vec::new();
694 let mut on_event = |event: StreamEvent| {
695 if let StreamEvent::Session(SessionEvent::Notice(n) | SessionEvent::Retry(n)) = &event {
696 notices.push(n.clone());
697 }
698 tap.usage
699 .lock()
700 .unwrap_or_else(std::sync::PoisonError::into_inner)
701 .add(&event);
702 };
703 let mut result = loop {
704 tokio::select! {
705 r = &mut call => break r,
706 Some(event) = rx.recv() => on_event(event),
707 }
708 };
709 while let Ok(event) = rx.try_recv() {
710 on_event(event);
711 }
712 if let (Ok(value), KeepWarmMode::Probe) = (&mut result, mode) {
714 if context_enabled {
715 if let Some(why) = self.keepwarm_context_blocks(
716 history,
717 system_prompt.as_deref(),
718 request_tools_schema.as_deref().map(|s| s.as_slice()),
719 &tools_snapshot,
720 ) {
721 value["keepwarm_context_blocks"] = Value::String(why);
722 }
723 }
724 }
725 Ok((result, notices))
726 }
727
728 fn keepwarm_context_blocks(
736 &self,
737 history: &[SharedMessage],
738 system_prompt: Option<&str>,
739 projected_schema: Option<&[Value]>,
740 tools_snapshot: &crate::ToolRegistry,
741 ) -> Option<String> {
742 use agent_core::core::context_policy::{assess_context, ContextAction};
743 let mut next: Vec<SharedMessage> = history.to_vec();
744 next.push(Arc::new(json!({"role": "user", "content": ""})));
745 let capped =
746 super::stream::cap_history_image_bytes(&next, super::stream::HISTORY_IMAGE_BYTE_CAP);
747 let request_messages: &[SharedMessage] = capped.as_deref().unwrap_or(&next[..]);
748 let fallback_schema;
749 let schema = match projected_schema {
750 Some(s) => s,
751 None => {
752 fallback_schema = tools_snapshot.tools_schema();
753 fallback_schema.as_slice()
754 }
755 };
756 let (assessment, budget) = super::stream::pre_request_context_budget(
757 &self.model,
758 self.context_window(),
759 system_prompt,
760 schema,
761 request_messages,
762 self.thinking_budget as u64,
763 self.max_tool_output as u64,
764 );
765 let state = self
766 .continuation
767 .lock()
768 .unwrap_or_else(std::sync::PoisonError::into_inner);
769 let decision = assess_context(&state.config, &state.policy, budget);
770 match decision.action {
771 ContextAction::Rollover => Some(format!(
772 "the next turn would roll the context over ({} tokens used)",
773 assessment.used_tokens()
774 )),
775 ContextAction::HardStop => Some(format!(
776 "the next turn would hit the context hard limit ({} tokens used)",
777 assessment.used_tokens()
778 )),
779 _ => None,
780 }
781 }
782}
783
784#[cfg(test)]
785mod tests {
786 use super::*;
787 use crate::core::config::CacheTtl;
788 use crate::extensions::hooks::events::{HookEvent, HookKind, HookResult};
789 use crate::extensions::permissions::{Permission, PermissionSet};
790 use agent_core::reasoning::ReasoningLevel;
791 use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::post, Router};
792 use futures::StreamExt;
793 use std::sync::atomic::AtomicUsize;
794 use std::sync::Mutex;
795
796 const SSE_HI: &str = concat!(
797 "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
798 "\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
799 "\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0,",
800 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}}\n\n",
801 "data: {\"type\":\"content_block_start\",\"index\":0,",
802 "\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
803 "data: {\"type\":\"content_block_delta\",\"index\":0,",
804 "\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n",
805 "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
806 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",",
807 "\"stop_sequence\":null},\"usage\":{\"input_tokens\":10,\"output_tokens\":1,",
808 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}\n\n",
809 "data: {\"type\":\"message_stop\"}\n\n",
810 );
811
812 const SSE_USAGE_THEN_ERROR: &str = concat!(
814 "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
815 "\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
816 "\"stop_sequence\":null,\"usage\":{\"input_tokens\":12,\"output_tokens\":0,",
817 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}}\n\n",
818 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":null,",
819 "\"stop_sequence\":null},\"usage\":{\"input_tokens\":12,\"output_tokens\":1,",
820 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}\n\n",
821 "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",",
822 "\"message\":\"synthetic\"}}\n\n",
823 );
824
825 #[derive(Clone)]
826 struct Mock {
827 bodies: Arc<Mutex<Vec<String>>>,
828 status: u16,
830 sse: &'static str,
831 }
832
833 async fn spawn_mock() -> (String, Mock) {
834 spawn_mock_full(200, SSE_HI).await
835 }
836
837 async fn spawn_mock_with(status: u16) -> (String, Mock) {
838 spawn_mock_full(status, SSE_HI).await
839 }
840
841 async fn spawn_mock_full(status: u16, sse: &'static str) -> (String, Mock) {
842 let mock = Mock {
843 bodies: Arc::default(),
844 status,
845 sse,
846 };
847 let app = Router::new()
848 .route(
849 "/v1/messages",
850 post(|State(m): State<Mock>, body: String| async move {
851 m.bodies.lock().unwrap().push(body);
852 if m.status == 200 {
853 (
854 StatusCode::OK,
855 [("content-type", "text/event-stream")],
856 m.sse.to_string(),
857 )
858 .into_response()
859 } else {
860 (
861 StatusCode::from_u16(m.status).unwrap(),
862 [("content-type", "application/json")],
863 r#"{"type":"error","error":{"type":"api_error","message":"synthetic"}}"#
864 .to_string(),
865 )
866 .into_response()
867 }
868 }),
869 )
870 .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
871 .with_state(mock.clone());
872 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
873 let addr = listener.local_addr().unwrap();
874 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
875 (format!("http://{addr}"), mock)
876 }
877
878 struct CountingInject(Arc<AtomicUsize>);
880 #[async_trait::async_trait]
881 impl crate::extensions::runtime::ExtensionHandler for CountingInject {
882 fn id(&self) -> &str {
883 "keepwarm-counting-inject"
884 }
885 async fn handle(&self, _event: &HookEvent) -> HookResult {
886 self.0.fetch_add(1, Ordering::SeqCst);
887 HookResult::Inject {
888 content: "PER-TURN-INJECTED-CONTEXT".into(),
889 }
890 }
891 async fn shutdown(&self) {}
892 }
893
894 async fn real_turn(rt: &Runtime, messages: Vec<SharedMessage>) -> Vec<SharedMessage> {
896 let mut stream = rt
897 .run_stream_with_messages(messages, CancellationToken::new(), None, None, false)
898 .await;
899 let mut history = None;
900 while let Some(event) = stream.next().await {
901 match event {
902 StreamEvent::Session(SessionEvent::MessageHistory(h)) => history = Some(h),
903 StreamEvent::Session(SessionEvent::Error(e)) => panic!("turn failed: {e:?}"),
904 _ => {}
905 }
906 }
907 history.expect("turn published its history")
908 }
909
910 fn user(text: &str) -> SharedMessage {
911 Arc::new(json!({"role": "user", "content": text}))
912 }
913
914 fn canon(v: &Value) -> String {
915 serde_json::to_string(v).unwrap()
916 }
917
918 fn strip_cache_control(mut msg: Value) -> Value {
919 if let Some(blocks) = msg["content"].as_array_mut() {
920 for b in blocks {
921 if let Some(obj) = b.as_object_mut() {
922 obj.remove("cache_control");
923 }
924 }
925 }
926 msg
927 }
928
929 fn marker_count(v: &Value) -> usize {
930 match v {
931 Value::Object(map) => {
932 usize::from(map.contains_key("cache_control"))
933 + map.values().map(marker_count).sum::<usize>()
934 }
935 Value::Array(items) => items.iter().map(marker_count).sum(),
936 _ => 0,
937 }
938 }
939
940 struct Observed {
942 ping: Value,
943 real: Value,
944 raw_ping: String,
945 history_len: usize,
946 usage: KeepWarmUsage,
947 hook_calls: (usize, usize, usize),
949 }
950
951 async fn ping_then_next_turn(rt: &Runtime) -> Observed {
956 let (url, mock) = spawn_mock().await;
957 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
958 let calls = Arc::new(AtomicUsize::new(0));
959 let mut perms = PermissionSet::new();
960 perms.grant(Permission::LlmContent);
961 rt.hook_bus()
962 .subscribe(
963 HookKind::BeforeMessage,
964 Arc::new(CountingInject(Arc::clone(&calls))),
965 None,
966 None,
967 perms,
968 )
969 .await
970 .unwrap();
971
972 assert_eq!(rt.cache_anchor(), None, "no request has run yet");
973 let history = real_turn(rt, vec![user("hello")]).await;
975 assert_eq!(history.last().unwrap()["role"], "assistant");
976 let turn_anchor = rt
977 .cache_anchor()
978 .expect("an accepted turn request stamps the anchor");
979 let before_ping = calls.load(Ordering::SeqCst);
980 let snapshot = history.clone();
981
982 let fingerprint = rt
983 .cache_keepwarm_probe(&history)
984 .await
985 .expect("probe")
986 .fingerprint;
987 assert_eq!(
988 mock.bodies.lock().unwrap().len(),
989 1,
990 "a probe sends nothing"
991 );
992 assert_eq!(
993 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
994 fingerprint,
995 "the fingerprint is deterministic"
996 );
997
998 let tap = PingTap::default();
999 let report = rt
1000 .cache_keepwarm_ping(&history, fingerprint, &CancellationToken::new(), &tap)
1001 .await;
1002 assert_eq!(report.result, Ok(PingOutcome::Sent), "{report:?}");
1003 assert!(tap.was_sent(), "the tap saw the request leave");
1004 assert_eq!(tap.usage(), report.usage, "the tap saw the usage");
1005 let usage = report.usage;
1006 assert!(
1007 usage.output_tokens >= 1,
1008 "usage comes back from the private channel"
1009 );
1010 let ping_anchor = rt.cache_anchor().expect("anchor");
1011 assert!(
1012 ping_anchor.mono >= turn_anchor.mono,
1013 "the ping moved the anchor"
1014 );
1015 let after_ping = calls.load(Ordering::SeqCst);
1016 assert_eq!(history, snapshot, "the ping must not mutate history");
1017
1018 let mut next = history.clone();
1020 next.push(user("next question"));
1021 real_turn(rt, next).await;
1022 let after_next = calls.load(Ordering::SeqCst);
1023
1024 let bodies = mock.bodies.lock().unwrap().clone();
1025 assert_eq!(bodies.len(), 3, "previous turn, ping, next turn");
1026 Observed {
1027 ping: serde_json::from_str(&bodies[1]).unwrap(),
1028 real: serde_json::from_str(&bodies[2]).unwrap(),
1029 raw_ping: bodies[1].clone(),
1030 history_len: history.len(),
1031 usage,
1032 hook_calls: (before_ping, after_ping, after_next),
1033 }
1034 }
1035
1036 fn assert_prefix_identical(
1037 ping: &Value,
1038 real: &Value,
1039 n: usize,
1040 history_marker: &Value,
1041 ping_max_tokens: u64,
1042 ) {
1043 for key in [
1045 "model",
1046 "tools",
1047 "system",
1048 "thinking",
1049 "output_config",
1050 "stream",
1051 ] {
1052 assert_eq!(
1053 canon(&ping[key]),
1054 canon(&real[key]),
1055 "`{key}` differs between ping and the real next turn"
1056 );
1057 }
1058 let pm = ping["messages"].as_array().unwrap();
1059 let rm = real["messages"].as_array().unwrap();
1060 assert_eq!(pm.len(), n + 1, "history + one trailing ping message");
1061 assert!(rm.len() > n);
1062 for i in 0..n - 1 {
1064 assert_eq!(canon(&pm[i]), canon(&rm[i]), "message {i} differs");
1065 }
1066 let last = &pm[n - 1];
1069 let marker = &last["content"].as_array().unwrap().last().unwrap()["cache_control"];
1070 assert_eq!(marker, history_marker);
1071 assert_eq!(
1072 canon(&strip_cache_control(last.clone())),
1073 canon(&rm[n - 1]),
1074 "last history message differs beyond the cache marker"
1075 );
1076 assert_eq!(marker_count(&pm[n]), 0, "ping message must be unmarked");
1078 assert_eq!(pm[n]["role"], "user");
1079 assert_eq!(pm[n]["content"], KEEPWARM_PING_TEXT);
1080 assert_eq!(ping["max_tokens"], json!(ping_max_tokens));
1082 assert!(real["max_tokens"].as_u64().unwrap() > ping_max_tokens);
1083 assert!(
1084 marker_count(ping) <= 4,
1085 "Anthropic allows at most 4 breakpoints"
1086 );
1087 }
1088
1089 fn marker_1h() -> Value {
1090 json!({"type": "ephemeral", "ttl": "1h"})
1091 }
1092
1093 fn marker_5m() -> Value {
1094 json!({"type": "ephemeral"})
1095 }
1096
1097 fn opus(ttl: CacheTtl) -> Runtime {
1098 let mut rt = Runtime::new_headless();
1099 rt.set_model("claude-opus-4-7".to_string());
1100 assert_eq!(rt.model(), "claude-opus-4-7");
1101 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1102 rt.set_cache_ttl(ttl);
1103 rt.set_system_prompt("You are a keep-warm test.".into());
1104 assert_eq!(rt.cache_keepwarm_ineligibility(), None);
1105 rt
1106 }
1107
1108 #[tokio::test]
1112 #[serial_test::serial(synaps_base_dir)]
1113 async fn ping_prefix_matches_next_turn_with_adaptive_thinking() {
1114 let _base = crate::test_env::BaseDirGuard::new();
1115 let rt = opus(CacheTtl::OneHour);
1116
1117 let o = ping_then_next_turn(&rt).await;
1118 assert!(
1119 !o.ping["thinking"].is_null(),
1120 "thinking must stay on: {}",
1121 o.ping
1122 );
1123 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1124
1125 let (before, after_ping, after_next) = o.hook_calls;
1127 assert!(
1128 before >= 1,
1129 "control: the previous turn fired before_message"
1130 );
1131 assert_eq!(after_ping, before, "the ping fired a hook");
1132 assert!(
1133 after_next > after_ping,
1134 "control: the next turn fired before_message"
1135 );
1136 assert!(!o.raw_ping.contains("PER-TURN-INJECTED-CONTEXT"));
1137 assert!(!o.raw_ping.contains("Extension context"));
1138 assert!(canon(&o.real).contains("PER-TURN-INJECTED-CONTEXT"));
1140 }
1141
1142 #[tokio::test]
1144 #[serial_test::serial(synaps_base_dir)]
1145 async fn ping_prefix_matches_next_turn_with_5m_ttl() {
1146 let _base = crate::test_env::BaseDirGuard::new();
1147 let rt = opus(CacheTtl::FiveMinutes);
1148 let o = ping_then_next_turn(&rt).await;
1149 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
1150 }
1151
1152 #[tokio::test]
1155 #[serial_test::serial(synaps_base_dir)]
1156 async fn ping_prefix_matches_next_turn_with_hybrid_ttl() {
1157 let _base = crate::test_env::BaseDirGuard::new();
1158 let rt = opus(CacheTtl::Hybrid);
1159 let o = ping_then_next_turn(&rt).await;
1160 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
1161 assert!(
1162 canon(&o.ping["system"]).contains("\"ttl\":\"1h\""),
1163 "hybrid keeps the 1h stable-prefix marker: {}",
1164 o.ping["system"]
1165 );
1166 }
1167
1168 #[tokio::test]
1170 #[serial_test::serial(synaps_base_dir)]
1171 async fn ping_prefix_matches_next_turn_with_thinking_off() {
1172 let _base = crate::test_env::BaseDirGuard::new();
1173 let mut rt = Runtime::new_headless();
1174 rt.set_model("claude-sonnet-4-6".to_string());
1175 rt.set_reasoning_level_explicit(ReasoningLevel::Off);
1176 rt.set_cache_ttl(CacheTtl::OneHour);
1177 assert_eq!(rt.cache_keepwarm_ineligibility(), None);
1178
1179 let o = ping_then_next_turn(&rt).await;
1180 assert!(o.ping.get("thinking").is_none(), "thinking off is omitted");
1181 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1182 }
1183
1184 #[tokio::test]
1189 #[serial_test::serial(synaps_base_dir)]
1190 async fn budgeted_thinking_ping_keeps_budget_and_is_costed() {
1191 let _base = crate::test_env::BaseDirGuard::new();
1192 let mut rt = Runtime::new_headless();
1193 rt.set_model("claude-sonnet-4-6".to_string());
1194 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1195 rt.set_cache_ttl(CacheTtl::FiveMinutes);
1196 assert_eq!(
1197 rt.cache_keepwarm_ineligibility(),
1198 None,
1199 "budgeted thinking is eligible"
1200 );
1201
1202 let o = ping_then_next_turn(&rt).await;
1203 assert_eq!(o.ping["thinking"]["type"], "enabled", "{}", o.ping);
1204 let budget = o.ping["thinking"]["budget_tokens"]
1205 .as_u64()
1206 .expect("budget");
1207 assert_eq!(canon(&o.ping["thinking"]), canon(&o.real["thinking"]));
1208 assert_prefix_identical(
1209 &o.ping,
1210 &o.real,
1211 o.history_len,
1212 &marker_5m(),
1213 budget + KEEPWARM_OUTPUT_TOKENS,
1214 );
1215 assert!(
1216 o.usage.cost_usd("claude-sonnet-4-6") > 0.0,
1217 "the ping is costed"
1218 );
1219 }
1220
1221 #[tokio::test]
1225 #[serial_test::serial(synaps_base_dir)]
1226 async fn changed_prefix_is_refused_before_the_network() {
1227 let _base = crate::test_env::BaseDirGuard::new();
1228 let (url, mock) = spawn_mock().await;
1229 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1230 let mut rt = opus(CacheTtl::FiveMinutes);
1231 let history = real_turn(&rt, vec![user("hello")]).await;
1232 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1233 let sent = mock.bodies.lock().unwrap().len();
1234
1235 let mut longer = history.clone();
1237 longer.push(user("more"));
1238 longer.push(Arc::new(
1239 json!({"role": "assistant", "content": [{"type": "text", "text": "x"}]}),
1240 ));
1241 let changed = |r: PingReport| matches!(r.result, Ok(PingOutcome::PrefixChanged { .. }));
1242 let tap = PingTap::default();
1243 assert!(changed(
1244 rt.cache_keepwarm_ping(&longer, fp, &CancellationToken::new(), &tap)
1245 .await
1246 ));
1247
1248 rt.set_system_prompt("A different system prompt.".into());
1249 assert_ne!(
1250 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1251 fp
1252 );
1253 assert!(changed(
1254 rt.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
1255 .await
1256 ));
1257 rt.set_system_prompt("You are a keep-warm test.".into());
1258 assert_eq!(
1259 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1260 fp,
1261 "restored"
1262 );
1263
1264 rt.set_reasoning_level_explicit(ReasoningLevel::Low);
1265 assert_ne!(
1266 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1267 fp,
1268 "effort is in the key"
1269 );
1270 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1271
1272 rt.set_model("claude-opus-4-6".to_string());
1273 assert_eq!(rt.model(), "claude-opus-4-6");
1274 assert_ne!(
1275 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1276 fp,
1277 "model is in the key"
1278 );
1279
1280 assert_eq!(
1281 mock.bodies.lock().unwrap().len(),
1282 sent,
1283 "no refused ping reached the provider"
1284 );
1285 assert!(!tap.was_sent(), "a refused ping never leaves");
1286 }
1287
1288 #[tokio::test]
1289 #[serial_test::serial(synaps_base_dir)]
1290 async fn http_4xx_is_fatal_and_5xx_is_transient() {
1291 let _base = crate::test_env::BaseDirGuard::new();
1292 let history = vec![
1293 user("hello"),
1294 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1295 ];
1296 for (status, fatal) in [(400u16, true), (404, true), (500, false), (429, false)] {
1297 let (url, mock) = spawn_mock_with(status).await;
1298 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1299 let rt = opus(CacheTtl::FiveMinutes);
1300 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1301 let err = rt
1302 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1303 .await
1304 .result
1305 .expect_err("non-2xx fails");
1306 assert_eq!(
1307 matches!(err, PingFailure::Fatal(_)),
1308 fatal,
1309 "HTTP {status}: {err:?}"
1310 );
1311 assert_eq!(
1312 mock.bodies.lock().unwrap().len(),
1313 1,
1314 "HTTP {status}: no retry"
1315 );
1316 assert_eq!(
1317 rt.cache_anchor(),
1318 None,
1319 "a rejected request refreshes nothing"
1320 );
1321 }
1322 }
1323
1324 async fn assert_identity_1h(rt: &Runtime) -> Observed {
1328 let o = ping_then_next_turn(rt).await;
1329 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1330 o
1331 }
1332
1333 #[tokio::test]
1337 #[serial_test::serial(synaps_base_dir)]
1338 async fn ping_prefix_matches_next_turn_with_progressive_disclosure() {
1339 let _base = crate::test_env::BaseDirGuard::new();
1340 let mut rt = opus(CacheTtl::OneHour);
1341 rt.progressive_tool_disclosure = true;
1342 assert_identity_1h(&rt).await;
1343 }
1344
1345 #[tokio::test]
1348 #[serial_test::serial(synaps_base_dir)]
1349 async fn ping_prefix_matches_next_turn_with_deferred_tools() {
1350 let _base = crate::test_env::BaseDirGuard::new();
1351 let mut rt = opus(CacheTtl::OneHour);
1352 rt.progressive_tool_disclosure = true;
1353 rt.deferred_tool_loading = true;
1354 assert_identity_1h(&rt).await;
1355 }
1356
1357 #[tokio::test]
1363 #[serial_test::serial(synaps_base_dir)]
1364 async fn ping_projects_carried_activations_the_turn_start_restores() {
1365 use crate::tools::activation::{activate_exact_for_user, SessionId, SessionToolSet};
1366 use crate::tools::catalog::ToolId;
1367 let _base = crate::test_env::BaseDirGuard::new();
1368 let (url, mock) = spawn_mock().await;
1369 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1370 let mut rt = opus(CacheTtl::OneHour);
1371 rt.progressive_tool_disclosure = true;
1372 rt.deferred_tool_loading = true;
1373 let carried = {
1374 let registry = rt.tools.read().await;
1375 let mut previous = SessionToolSet::progressive_core_for_catalog(
1376 SessionId::parse("previous-runtime").unwrap(),
1377 registry.catalog(),
1378 );
1379 activate_exact_for_user(
1380 &mut previous,
1381 registry.catalog(),
1382 &ToolId::builtin("subagent_status"),
1383 )
1384 .unwrap();
1385 previous.carried_activations()
1386 };
1387 rt.restore_tool_activations(carried);
1388
1389 let history = vec![
1390 user("hello"),
1391 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1392 ];
1393 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1394 let report = rt
1395 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1396 .await;
1397 assert_eq!(report.result, Ok(PingOutcome::Sent), "{report:?}");
1398 let mut next = history.clone();
1399 next.push(user("next"));
1400 real_turn(&rt, next).await;
1401
1402 let bodies: Vec<Value> = mock
1403 .bodies
1404 .lock()
1405 .unwrap()
1406 .iter()
1407 .map(|b| serde_json::from_str(b).unwrap())
1408 .collect();
1409 let (ping, real) = (&bodies[0], &bodies[1]);
1410 let loaded = |b: &Value| -> Vec<Value> {
1415 b["tools"]
1416 .as_array()
1417 .unwrap()
1418 .iter()
1419 .filter(|t| t.get("defer_loading").is_none())
1420 .cloned()
1421 .collect()
1422 };
1423 assert_eq!(
1424 serde_json::to_string(&loaded(ping)).unwrap(),
1425 serde_json::to_string(&loaded(real)).unwrap()
1426 );
1427 assert!(
1428 real["tools"].as_array().unwrap().iter().any(|t| {
1429 t["name"] == "subagent_status" && t.get("defer_loading").is_none()
1430 }),
1431 "the restored activation is a loaded tool: {}",
1432 real["tools"]
1433 );
1434 }
1435
1436 #[tokio::test]
1439 #[serial_test::serial(synaps_base_dir)]
1440 async fn ping_prefix_matches_next_turn_with_context_management_guidance() {
1441 let _base = crate::test_env::BaseDirGuard::new();
1442 let rt = opus(CacheTtl::OneHour);
1443 rt.continuation.lock().unwrap().config.mode =
1444 crate::core::config::ContextManagementMode::Auto;
1445 let o = assert_identity_1h(&rt).await;
1446 assert!(
1447 canon(&o.ping["system"]).contains("Work normally."),
1448 "guidance present: {}",
1449 o.ping["system"]
1450 );
1451 }
1452
1453 #[tokio::test]
1456 #[serial_test::serial(synaps_base_dir)]
1457 async fn ping_prefix_matches_next_turn_with_session_injection() {
1458 let _base = crate::test_env::BaseDirGuard::new();
1459 let mut rt = opus(CacheTtl::OneHour);
1460 rt.set_session_id(Some("kw-session".into()));
1461 rt.hook_bus()
1462 .set_session_injection_for("kw-session", "SESSION-SCOPED-INJECTION".into())
1463 .await;
1464 let o = assert_identity_1h(&rt).await;
1465 assert!(canon(&o.ping["system"]).contains("SESSION-SCOPED-INJECTION"));
1466 }
1467
1468 #[tokio::test]
1471 #[serial_test::serial(synaps_base_dir)]
1472 async fn ping_prefix_matches_next_turn_with_oauth_system_blocks() {
1473 let _base = crate::test_env::BaseDirGuard::new();
1474 let mut rt = opus(CacheTtl::OneHour);
1475 rt.auth = Arc::new(tokio::sync::RwLock::new(super::super::types::AuthState {
1476 auth_token: "oauth-test-token".to_string(),
1477 auth_type: "oauth".to_string(),
1478 refresh_token: None,
1479 token_expires: Some(crate::epoch_millis() + 3_600_000),
1480 auth_file_mtime: None,
1481 bound_credential: Some(crate::runtime::auth::anthropic_binding(
1484 &crate::auth::CredentialSource::Local,
1485 &crate::auth::CredentialRef::default_for(crate::auth::OAuthProviderId::Anthropic),
1486 )),
1487 }));
1488 let o = assert_identity_1h(&rt).await;
1489 assert!(
1490 canon(&o.ping["system"]).contains("You are Claude Code"),
1491 "{}",
1492 o.ping["system"]
1493 );
1494 }
1495
1496 #[tokio::test]
1500 #[serial_test::serial(synaps_base_dir)]
1501 async fn probe_reports_a_context_rollover_ahead() {
1502 let _base = crate::test_env::BaseDirGuard::new();
1503 let (url, _mock) = spawn_mock().await;
1504 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1505 let rt = opus(CacheTtl::OneHour);
1506 rt.continuation.lock().unwrap().config.mode =
1507 crate::core::config::ContextManagementMode::Auto;
1508 let history = real_turn(&rt, vec![user("hello")]).await;
1509 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1510 assert_eq!(probe.context_blocks, None, "plenty of room");
1511 let before = rt.continuation.lock().unwrap().policy;
1512 rt.continuation.lock().unwrap().config.reserve_tokens = 10_000_000;
1514 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1515 assert!(probe.context_blocks.is_some(), "{probe:?}");
1516 assert_eq!(
1517 rt.continuation.lock().unwrap().policy,
1518 before,
1519 "the probe never commits policy state"
1520 );
1521 rt.continuation.lock().unwrap().config.mode =
1523 crate::core::config::ContextManagementMode::Off;
1524 assert_eq!(
1525 rt.cache_keepwarm_probe(&history)
1526 .await
1527 .unwrap()
1528 .context_blocks,
1529 None
1530 );
1531 }
1532
1533 #[tokio::test]
1535 #[serial_test::serial(synaps_base_dir)]
1536 async fn probe_reports_max_tokens() {
1537 let _base = crate::test_env::BaseDirGuard::new();
1538 let history = vec![
1539 user("hello"),
1540 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1541 ];
1542 let rt = opus(CacheTtl::OneHour);
1543 assert_eq!(
1544 rt.cache_keepwarm_probe(&history).await.unwrap().max_tokens,
1545 1
1546 );
1547 let mut rt = Runtime::new_headless();
1548 rt.set_model("claude-sonnet-4-6".to_string());
1549 rt.set_reasoning_level_explicit(ReasoningLevel::XHigh);
1550 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1551 let budget = u64::from(rt.thinking_budget);
1552 assert!(budget >= 16_384, "budgeted thinking at xhigh: {budget}");
1553 assert_eq!(
1554 probe.max_tokens,
1555 budget + KEEPWARM_OUTPUT_TOKENS,
1556 "{probe:?}"
1557 );
1558 }
1559
1560 #[tokio::test]
1563 #[serial_test::serial(synaps_base_dir)]
1564 async fn ping_usage_is_returned_on_error() {
1565 let _base = crate::test_env::BaseDirGuard::new();
1566 let (url, mock) = spawn_mock_full(200, SSE_USAGE_THEN_ERROR).await;
1567 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1568 let history = vec![
1569 user("hello"),
1570 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1571 ];
1572 let rt = opus(CacheTtl::OneHour);
1573 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1574 let tap = PingTap::default();
1575 let report = rt
1576 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
1577 .await;
1578 assert!(
1579 matches!(report.result, Err(PingFailure::Transient(_))),
1580 "{report:?}"
1581 );
1582 assert_eq!(report.usage.cache_read_tokens, 5_000, "{report:?}");
1583 assert!(report.usage.cost_usd("claude-opus-4-7") > 0.0);
1584 assert_eq!(tap.usage(), report.usage);
1585 assert!(tap.was_sent());
1586 assert_eq!(mock.bodies.lock().unwrap().len(), 1, "no transport retry");
1587 }
1588
1589 #[tokio::test]
1593 #[serial_test::serial(synaps_base_dir)]
1594 async fn ping_never_carries_or_moves_the_diagnostics_chain() {
1595 let _base = crate::test_env::BaseDirGuard::new();
1596 let (url, mock) = spawn_mock().await;
1597 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1598 let history = vec![
1599 user("hello"),
1600 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1601 ];
1602 let mut rt = opus(CacheTtl::OneHour);
1603 rt.set_cache_diagnostics(true);
1604 rt.cache_diagnosis().unwrap().record("msg_turn");
1605 let mut ping_rt = rt.clone();
1606 ping_rt.share_ttl_latches(&rt);
1607 let fp = ping_rt
1608 .cache_keepwarm_probe(&history)
1609 .await
1610 .unwrap()
1611 .fingerprint;
1612 let report = ping_rt
1613 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1614 .await;
1615 assert!(report.result.is_ok(), "{report:?}");
1616 let bodies = mock.bodies.lock().unwrap().clone();
1617 assert_eq!(bodies.len(), 1);
1618 assert!(!bodies[0].contains("\"diagnostics\""), "{}", bodies[0]);
1619 assert_eq!(
1620 rt.cache_diagnosis()
1621 .unwrap()
1622 .previous_message_id()
1623 .as_deref(),
1624 Some("msg_turn"),
1625 "the ping's msg_kw never enters the chain"
1626 );
1627 }
1628
1629 #[test]
1633 fn diagnostics_chain_is_shared_like_the_ttl_latches() {
1634 let mut rt = opus(CacheTtl::FiveMinutes);
1635 assert!(rt.cache_diagnosis().is_none(), "off: wire unchanged");
1636 rt.set_cache_diagnostics(true);
1637 rt.cache_diagnosis().unwrap().record("msg_a");
1638 let fresh = rt.clone();
1639 assert_eq!(fresh.cache_diagnosis().unwrap().previous_message_id(), None);
1640 let mut shared = rt.clone();
1641 shared.share_ttl_latches(&rt);
1642 shared.cache_diagnosis().unwrap().record("msg_b");
1643 shared.cache_diagnosis().unwrap().reject(400);
1644 let chain = rt.cache_diagnosis().unwrap();
1645 assert_eq!(chain.previous_message_id().as_deref(), Some("msg_b"));
1646 assert!(!chain.active(), "the rejection latch is session-wide");
1647 let fresh_chain = fresh.cache_diagnosis().unwrap();
1650 assert!(!fresh_chain.active(), "latch shared across Clone");
1651 assert_eq!(fresh_chain.previous_message_id(), None);
1652 assert!(!rt.clone().cache_diagnosis().unwrap().active());
1653 }
1654
1655 #[test]
1658 fn ping_shares_the_session_ttl_latch() {
1659 let rt = opus(CacheTtl::OneHour);
1660 assert!(!rt.cache_ttl_downgraded());
1661 let mut ping_rt = rt.clone();
1662 ping_rt.share_ttl_latches(&rt);
1663 ping_rt
1664 .ttl_downgrade_notified
1665 .store(true, Ordering::Relaxed);
1666 assert!(rt.cache_ttl_downgraded(), "the latch is shared");
1667 }
1668
1669 #[test]
1670 fn failure_classification() {
1671 let api = || RuntimeError::ApiStatus("x".into());
1672 assert_eq!(
1673 classify_failure(&RuntimeError::Canceled, 0),
1674 PingFailure::Canceled
1675 );
1676 assert!(matches!(
1677 classify_failure(&api(), 400),
1678 PingFailure::Fatal(_)
1679 ));
1680 assert!(matches!(
1681 classify_failure(&api(), 401),
1682 PingFailure::Fatal(_)
1683 ));
1684 assert!(matches!(
1685 classify_failure(&api(), 429),
1686 PingFailure::Transient(_)
1687 ));
1688 assert!(matches!(
1689 classify_failure(&api(), 529),
1690 PingFailure::Transient(_)
1691 ));
1692 assert!(matches!(
1693 classify_failure(&api(), 0),
1694 PingFailure::Transient(_)
1695 ));
1696 assert!(matches!(
1698 classify_failure(&api(), 200),
1699 PingFailure::Transient(_)
1700 ));
1701 assert!(matches!(
1702 classify_failure(&RuntimeError::Config("x".into()), 0),
1703 PingFailure::Fatal(_)
1704 ));
1705 }
1706
1707 #[test]
1711 fn network_timeout_is_transient() {
1712 for msg in ["upload/header deadline", "stream idle"] {
1713 assert_eq!(
1714 classify_failure(&RuntimeError::NetworkTimeout(msg), 0),
1715 PingFailure::Transient("network".to_string()),
1716 "{msg}"
1717 );
1718 }
1719 }
1720
1721 #[test]
1722 fn every_ttl_and_thinking_mode_is_eligible() {
1723 for ttl in [CacheTtl::FiveMinutes, CacheTtl::OneHour, CacheTtl::Hybrid] {
1724 for (model, level) in [
1725 ("claude-opus-4-7", ReasoningLevel::High),
1726 ("claude-sonnet-4-6", ReasoningLevel::High),
1727 ("claude-sonnet-4-6", ReasoningLevel::Off),
1728 ] {
1729 let mut rt = Runtime::new_headless();
1730 rt.set_model(model.to_string());
1731 rt.set_reasoning_level_explicit(level);
1732 rt.set_cache_ttl(ttl);
1733 assert_eq!(
1734 rt.cache_keepwarm_ineligibility(),
1735 None,
1736 "{model} {level:?} {ttl:?}"
1737 );
1738 }
1739 }
1740 }
1741
1742 #[test]
1743 fn non_anthropic_models_are_ineligible() {
1744 let mut rt = Runtime::new_headless();
1745 rt.set_model("groq/llama-3.3-70b-versatile".to_string());
1746 if rt.model() == "groq/llama-3.3-70b-versatile" {
1747 assert!(rt.cache_keepwarm_ineligibility().is_some());
1748 }
1749 }
1750
1751 #[tokio::test]
1754 async fn ineligible_ping_sends_nothing() {
1755 let mut rt = Runtime::new_headless();
1756 rt.set_model("groq/llama-3.3-70b-versatile".to_string());
1757 if rt.model() != "groq/llama-3.3-70b-versatile" {
1758 return;
1759 }
1760 let history = vec![
1761 user("hello"),
1762 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1763 ];
1764 let tap = PingTap::default();
1765 let err = rt
1766 .cache_keepwarm_ping(&history, 0, &CancellationToken::new(), &tap)
1767 .await
1768 .result
1769 .expect_err("non-Anthropic must refuse");
1770 assert!(!tap.was_sent());
1771 assert!(
1772 matches!(err, PingFailure::Fatal(ref m) if m.contains("keep-warm skipped")),
1773 "{err:?}"
1774 );
1775 }
1776
1777 #[test]
1778 fn anchor_cell_only_moves_forward() {
1779 let cell = CacheAnchorCell::default();
1780 assert_eq!(cell.get(), None);
1781 let a = CacheAnchor::now();
1782 let later = CacheAnchor {
1783 mono: a.mono + std::time::Duration::from_secs(5),
1784 wall: a.wall + std::time::Duration::from_secs(5),
1785 };
1786 cell.stamp(later);
1787 cell.stamp(a);
1788 assert_eq!(
1789 cell.get(),
1790 Some(later),
1791 "a late-arriving older stamp never rewinds"
1792 );
1793 }
1794
1795 #[test]
1796 fn only_a_finished_assistant_text_reply_is_markable() {
1797 assert!(tail_is_markable(
1798 &json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}, {"type": "text", "text": "ok"}]})
1799 ));
1800 assert!(!tail_is_markable(&json!({"role": "user", "content": "q"})));
1801 assert!(!tail_is_markable(
1802 &json!({"role": "assistant", "content": [{"type": "tool_use", "id": "t", "name": "x", "input": {}}]})
1803 ));
1804 assert!(!tail_is_markable(
1805 &json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}]})
1806 ));
1807 }
1808
1809 #[test]
1810 fn usage_is_summed_and_priced() {
1811 let mut u = KeepWarmUsage::default();
1812 u.add(&StreamEvent::Session(SessionEvent::Usage {
1813 input_tokens: 20,
1814 output_tokens: 16,
1815 cache_read_input_tokens: 500_000,
1816 cache_creation_input_tokens: 0,
1817 cache_creation_5m: None,
1818 cache_creation_1h: Some(0),
1819 model: Some("claude-opus-4-7".into()),
1820 }));
1821 assert_eq!(u.cache_read_tokens, 500_000);
1822 assert_eq!(u.output_tokens, 16);
1823 let cost = u.cost_usd("claude-opus-4-7");
1825 assert!((cost - 0.2505).abs() < 0.001, "{cost}");
1826 }
1827
1828 #[test]
1830 fn booking_rule_table() {
1831 use PingBooking::*;
1832 let table = [
1834 ((true, true, false, 200), Measured),
1835 ((true, false, false, 0), Measured),
1836 ((true, true, false, 529), Measured),
1837 ((false, true, false, 200), WorstCase), ((false, true, false, 0), WorstCase), ((false, true, false, 500), Nothing), ((false, true, false, 429), Nothing),
1841 ((false, true, false, 529), Nothing),
1842 ((false, true, true, 0), Nothing), ((false, true, true, 503), Nothing), ((false, false, false, 0), Nothing), ((false, false, true, 0), Nothing),
1846 ];
1847 for ((usage, sent, connect, status), want) in table {
1848 assert_eq!(
1849 ping_booking(usage, sent, connect, status),
1850 want,
1851 "usage={usage} sent={sent} connect_failed={connect} status={status}"
1852 );
1853 }
1854 let tap = PingTap::default();
1856 assert_eq!(tap.bill("claude-opus-4-7", 0.5).cost_usd, 0.0);
1857 tap.sent.store(true, Ordering::Relaxed);
1858 let bill = tap.bill("claude-opus-4-7", 0.5);
1859 assert_eq!((bill.booking, bill.cost_usd), (WorstCase, 0.5));
1860 tap.usage.lock().unwrap().cache_read_tokens = 1_000_000;
1861 let bill = tap.bill("claude-opus-4-7", 0.5);
1862 assert_eq!(bill.booking, Measured);
1863 assert!((bill.cost_usd - 0.5).abs() < 1e-9, "1M reads at $0.50/MTok");
1864 }
1865
1866 fn two_message_history() -> Vec<SharedMessage> {
1867 vec![
1868 user("hello"),
1869 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1870 ]
1871 }
1872
1873 #[tokio::test]
1875 #[serial_test::serial(synaps_base_dir)]
1876 async fn a_refused_connection_is_never_sent_and_books_nothing() {
1877 let _base = crate::test_env::BaseDirGuard::new();
1878 let addr = {
1879 let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1880 l.local_addr().unwrap()
1881 }; let _url = crate::test_env::EnvVarGuard::set(
1883 "SYNAPS_ANTHROPIC_BASE_URL",
1884 &format!("http://{addr}"),
1885 );
1886 let history = two_message_history();
1887 let rt = opus(CacheTtl::OneHour);
1888 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1889 let tap = PingTap::default();
1890 let report = rt
1891 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
1892 .await;
1893 assert!(report.result.is_err(), "{report:?}");
1894 assert!(!tap.was_sent(), "never dispatched");
1895 assert!(tap.connect_failed.load(Ordering::Relaxed));
1896 let bill = tap.bill("claude-opus-4-7", 0.75);
1897 assert_eq!((bill.booking, bill.cost_usd), (PingBooking::Nothing, 0.0));
1898 }
1899
1900 #[tokio::test]
1903 #[serial_test::serial(synaps_base_dir)]
1904 async fn a_dispatched_ping_without_usage_books_the_worst_case() {
1905 use tokio::io::AsyncReadExt;
1906 let _base = crate::test_env::BaseDirGuard::new();
1907 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1908 let addr = listener.local_addr().unwrap();
1909 tokio::spawn(async move {
1911 let (mut sock, _) = listener.accept().await.unwrap();
1912 let mut buf = vec![0u8; 64 * 1024];
1913 while sock.read(&mut buf).await.is_ok_and(|n| n > 0) {}
1914 });
1915 let _url = crate::test_env::EnvVarGuard::set(
1916 "SYNAPS_ANTHROPIC_BASE_URL",
1917 &format!("http://{addr}"),
1918 );
1919 let history = two_message_history();
1920 let rt = opus(CacheTtl::OneHour);
1921 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1922 let tap = PingTap::default();
1923 let cancel = CancellationToken::new();
1924 let ping = {
1925 let (rt, history, tap, cancel) =
1926 (rt.clone(), history.clone(), tap.clone(), cancel.clone());
1927 tokio::spawn(async move { rt.cache_keepwarm_ping(&history, fp, &cancel, &tap).await })
1928 };
1929 tokio::time::timeout(std::time::Duration::from_secs(10), async {
1930 while !tap.was_sent() {
1931 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1932 }
1933 })
1934 .await
1935 .expect("the body was pulled over the established connection");
1936 cancel.cancel();
1937 let report = ping.await.unwrap();
1938 assert_eq!(report.result, Err(PingFailure::Canceled));
1939 assert!(!tap.connect_failed.load(Ordering::Relaxed));
1940 let bill = tap.bill("claude-opus-4-7", 0.75);
1941 assert_eq!(
1942 (bill.booking, bill.cost_usd),
1943 (PingBooking::WorstCase, 0.75)
1944 );
1945 }
1946
1947 #[test]
1948 fn reply_prefix_tokens_skip_thinking() {
1949 let text = "a".repeat(250); let with_thinking = vec![Arc::new(json!({"role": "assistant", "content": [
1951 {"type": "thinking", "thinking": "x".repeat(50_000), "signature": "s"},
1952 {"type": "text", "text": text},
1953 ]}))];
1954 let plain = vec![Arc::new(json!({"role": "assistant", "content": [
1955 {"type": "text", "text": text},
1956 ]}))];
1957 let t = reply_prefix_tokens(&with_thinking);
1958 assert_eq!(t, reply_prefix_tokens(&plain), "thinking is not counted");
1959 assert!((100..200).contains(&t), "{t}");
1960 assert_eq!(reply_prefix_tokens(&[]), 0);
1961 }
1962}