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
390pub(super) async fn drive_ping<F>(
396 call: F,
397 rx: &mut mpsc::UnboundedReceiver<StreamEvent>,
398 tap: &PingTap,
399 forward: Option<&mpsc::UnboundedSender<StreamEvent>>,
400) -> (Result<Value>, Vec<String>)
401where
402 F: std::future::Future<Output = Result<Value>>,
403{
404 tokio::pin!(call);
405 let mut notices = Vec::new();
406 let mut on_event = |event: StreamEvent| {
407 if let StreamEvent::Session(SessionEvent::Notice(n) | SessionEvent::Retry(n)) = &event {
408 notices.push(n.clone());
409 }
410 tap.usage
411 .lock()
412 .unwrap_or_else(std::sync::PoisonError::into_inner)
413 .add(&event);
414 if let Some(tx) = forward {
415 if matches!(
416 event,
417 StreamEvent::Session(SessionEvent::Usage { .. } | SessionEvent::Notice(_))
418 ) {
419 let _ = tx.send(event);
420 }
421 }
422 };
423 let result = loop {
424 tokio::select! {
425 r = &mut call => break r,
426 Some(event) = rx.recv() => on_event(event),
427 }
428 };
429 while let Ok(event) = rx.try_recv() {
430 on_event(event);
431 }
432 (result, notices)
433}
434
435impl Runtime {
436 pub fn cache_keepwarm_ineligibility(&self) -> Option<String> {
440 if !super::auth::model_is_anthropic(&self.model) {
441 return Some(format!(
442 "model {} is not on the native Anthropic transport",
443 self.model
444 ));
445 }
446 None
447 }
448
449 pub fn cache_anchor(&self) -> Option<CacheAnchor> {
451 self.cache_anchor.get()
452 }
453
454 pub fn cache_ttl_downgraded(&self) -> bool {
457 self.ttl_downgrade_notified.load(Ordering::Relaxed)
458 }
459
460 pub fn cache_1h_honoured(&self) -> bool {
462 self.saw_1h_honored.load(Ordering::Relaxed)
463 }
464
465 pub async fn cache_keepwarm_probe(&self, history: &[SharedMessage]) -> Result<KeepWarmProbe> {
471 let cancel = CancellationToken::new();
472 let tap = PingTap::default();
473 let (value, _) = self
474 .keepwarm_request(history, &cancel, KeepWarmMode::Probe, &tap)
475 .await;
476 let value = value?;
477 let fingerprint = value["keepwarm_prefix"].as_u64().ok_or_else(|| {
478 RuntimeError::Session("keep-warm probe produced no fingerprint".into())
479 })?;
480 Ok(KeepWarmProbe {
481 fingerprint,
482 max_tokens: value["keepwarm_max_tokens"].as_u64().unwrap_or(0),
483 context_blocks: value["keepwarm_context_blocks"]
484 .as_str()
485 .map(str::to_string),
486 })
487 }
488
489 pub async fn cache_keepwarm_ping(
495 &self,
496 history: &[SharedMessage],
497 expect: u64,
498 cancel: &CancellationToken,
499 tap: &PingTap,
500 ) -> PingReport {
501 let (value, notices) = self
502 .keepwarm_request(history, cancel, KeepWarmMode::Send { expect }, tap)
503 .await;
504 let result = match value {
505 Ok(value) => match value.get("keepwarm_prefix").and_then(Value::as_u64) {
506 Some(actual) => Ok(PingOutcome::PrefixChanged {
507 expected: expect,
508 actual,
509 }),
510 None => Ok(PingOutcome::Sent),
511 },
512 Err(e) => Err(classify_failure(&e, tap.status.load(Ordering::Relaxed))),
513 };
514 PingReport {
515 result,
516 usage: tap.usage(),
517 notices,
518 }
519 }
520
521 async fn keepwarm_request(
522 &self,
523 history: &[SharedMessage],
524 cancel: &CancellationToken,
525 mode: KeepWarmMode,
526 tap: &PingTap,
527 ) -> (Result<Value>, Vec<String>) {
528 match self
529 .keepwarm_request_inner(history, cancel, mode, tap)
530 .await
531 {
532 Ok(r) => r,
533 Err(e) => (Err(e), Vec::new()),
534 }
535 }
536
537 async fn keepwarm_request_inner(
540 &self,
541 history: &[SharedMessage],
542 cancel: &CancellationToken,
543 mode: KeepWarmMode,
544 tap: &PingTap,
545 ) -> Result<(Result<Value>, Vec<String>)> {
546 if let Some(reason) = self.cache_keepwarm_ineligibility() {
547 return Err(RuntimeError::Config(format!(
548 "cache keep-warm skipped: {reason}"
549 )));
550 }
551 if !history_ready(history) {
552 return Err(RuntimeError::Config(
553 "cache keep-warm skipped: history does not end with a finished assistant reply"
554 .into(),
555 ));
556 }
557 let (durability_blocked, context_enabled) = {
558 let state = self
559 .continuation
560 .lock()
561 .unwrap_or_else(std::sync::PoisonError::into_inner);
562 (state.durability_blocked, state.enabled())
563 };
564 if durability_blocked {
565 return Err(RuntimeError::Session(
566 "cache keep-warm skipped: context head save is unresolved".into(),
567 ));
568 }
569
570 self.validate_request_preflight().await?;
573 let anthropic_execution_plan = self.authorized_anthropic_plan().await?;
574 if matches!(mode, KeepWarmMode::Send { .. }) {
575 self.refresh_for_keepwarm().await?;
578 }
579
580 let system_prompt = self.effective_system_prompt().await;
582 let system_prompt = if context_enabled {
583 Some(format!(
584 "{}\n\n{}",
585 system_prompt.as_deref().unwrap_or_default(),
586 super::continuation::GUIDANCE
587 ))
588 } else {
589 system_prompt
590 };
591 let session_injection = match self.session_id.as_deref() {
592 Some(id) => self.hook_bus.session_injection_for(id).await,
593 None => None,
594 };
595 let system_prompt = match session_injection {
596 Some(content) => Some(super::stream::wrap_extension_context(
597 system_prompt.as_deref().unwrap_or_default(),
598 &content,
599 )),
600 None => system_prompt,
601 };
602
603 let mut request: Vec<SharedMessage> = history.to_vec();
608 request.push(Arc::new(
609 json!({"role": "user", "content": KEEPWARM_PING_TEXT}),
610 ));
611 let round_images =
612 super::image_retention::round_images(&self.model, self.codex_request_role(), &request)
613 .map_err(RuntimeError::Config)?;
614 let (request_messages, _) = round_images.views(&request);
615
616 let tools_snapshot = self.tools.read().await.clone();
621 let (request_tools_schema, anthropic_tool_split) = if self.progressive_tool_disclosure {
622 let fresh = super::continuation::context_tool_set(
623 self.host_tool_session.clone(),
624 tools_snapshot.catalog(),
625 true,
626 context_enabled,
627 );
628 let retained = self
631 .retained_tool_set
632 .lock()
633 .unwrap_or_else(std::sync::PoisonError::into_inner)
634 .clone();
635 let existing = retained.as_ref().map(|shared| {
636 shared
637 .read()
638 .unwrap_or_else(std::sync::PoisonError::into_inner)
639 });
640 let (mut set, _, _) = super::stream::select_turn_start_tool_set(
641 existing.as_deref(),
642 fresh,
643 tools_snapshot.catalog(),
644 true,
645 );
646 drop(existing);
647 let pending = self
652 .pending_tool_restore
653 .lock()
654 .unwrap_or_else(std::sync::PoisonError::into_inner)
655 .clone();
656 if let Some(carried) = pending {
657 if crate::tools::activation::carry_forward_enabled() {
658 let _ = super::stream::restore_pending_activations(
659 &mut set,
660 &carried,
661 tools_snapshot.catalog(),
662 self.activation_confirm,
663 );
664 }
665 }
666 let schema = tools_snapshot.session_tools_schema(&set).schema;
667 let split = self.deferred_tool_loading.then(|| {
668 let referenced = HelperMethods::referenced_tool_ids(request_messages);
669 Arc::new(tools_snapshot.deferred_tool_split(&set, &referenced))
670 });
671 (Some(Arc::new(schema)), split)
672 } else {
673 (None, None)
674 };
675
676 let options = ApiOptions {
677 memory_backend: None,
678 use_1m_context: self.context_window_override == Some(1_000_000),
679 cache_ttl: self.cache_ttl,
680 ttl_downgrade_notified: self.ttl_downgrade_notified.clone(),
684 saw_1h_honored: self.saw_1h_honored.clone(),
685 credential_source: self.credential_source.clone(),
686 token_cache: self.token_cache.clone(),
687 anthropic_base_url: None,
688 anthropic_stream_client: self.anthropic_stream_client.clone(),
689 request_compression: self.request_compression.clone(),
690 anthropic_execution_plan,
691 codex_request_role: self.codex_request_role(),
692 thinking_display: self.wire_thinking_display(),
694 retention_frontier: round_images.applied_frontier(),
695 trace: self.trace_ctx.clone(),
697 request_correlation: None,
698 suppress_stream_deltas: true,
699 telemetry: self.request_telemetry(),
700 tool_session_id: Some(self.host_tool_session.clone()),
701 session_tool_set: None,
702 request_tools_schema: request_tools_schema.clone(),
703 anthropic_tool_split,
704 usage_counters: None,
705 cache_anchor: Some(Arc::clone(&self.cache_anchor)),
708 keepwarm: Some(KeepWarmRequest {
709 mode,
710 status: Arc::clone(&tap.status),
711 sent: Arc::clone(&tap.sent),
712 connect_failed: Arc::clone(&tap.connect_failed),
713 }),
714 cache_diagnosis: self.cache_diagnosis(),
717 context_rollover: None,
718 };
719
720 let (tx, mut rx) = mpsc::unbounded_channel();
721 let call = ApiMethods::call_api_stream_inner(
722 &self.auth,
723 &self.client,
724 &self.model,
725 &tools_snapshot,
726 &system_prompt,
727 self.thinking_budget,
728 self.reasoning_level(),
729 request_messages,
730 tx,
731 cancel,
732 0,
733 0,
734 &options,
735 self.telemetry_level,
736 );
737 let (mut result, notices) = drive_ping(call, &mut rx, tap, None).await;
738 if let (Ok(value), KeepWarmMode::Probe) = (&mut result, mode) {
740 if context_enabled {
741 if let Some(why) = self.keepwarm_context_blocks(
742 history,
743 system_prompt.as_deref(),
744 request_tools_schema.as_deref().map(|s| s.as_slice()),
745 &tools_snapshot,
746 ) {
747 value["keepwarm_context_blocks"] = Value::String(why);
748 }
749 }
750 }
751 Ok((result, notices))
752 }
753
754 fn keepwarm_context_blocks(
762 &self,
763 history: &[SharedMessage],
764 system_prompt: Option<&str>,
765 projected_schema: Option<&[Value]>,
766 tools_snapshot: &crate::ToolRegistry,
767 ) -> Option<String> {
768 use agent_core::core::context_policy::{assess_context, ContextAction};
769 let mut next: Vec<SharedMessage> = history.to_vec();
770 next.push(Arc::new(json!({"role": "user", "content": ""})));
771 let round_images =
773 super::image_retention::round_images(&self.model, self.codex_request_role(), &next)
774 .ok();
775 let (_, request_messages) = match &round_images {
776 Some(round) => round.views(&next),
777 None => (&next[..], &next[..]),
778 };
779 let fallback_schema;
780 let schema = match projected_schema {
781 Some(s) => s,
782 None => {
783 fallback_schema = tools_snapshot.tools_schema();
784 fallback_schema.as_slice()
785 }
786 };
787 let anchor = self
788 .continuation
789 .lock()
790 .unwrap_or_else(std::sync::PoisonError::into_inner)
791 .usage_anchor_for(&self.model)
792 .cloned();
793 let (assessment, budget) = super::stream::pre_request_context_budget(
794 &self.model,
795 self.context_window(),
796 system_prompt,
797 schema,
798 request_messages,
799 self.thinking_budget as u64,
800 self.max_tool_output as u64,
801 anchor.as_ref(),
802 );
803 let state = self
804 .continuation
805 .lock()
806 .unwrap_or_else(std::sync::PoisonError::into_inner);
807 let decision = assess_context(&state.config, &state.policy, budget);
808 match decision.action {
809 ContextAction::Rollover => Some(format!(
810 "the next turn would roll the context over ({} tokens used)",
811 assessment.used_tokens()
812 )),
813 ContextAction::HardStop => Some(format!(
814 "the next turn would hit the context hard limit ({} tokens used)",
815 assessment.used_tokens()
816 )),
817 _ => None,
818 }
819 }
820}
821
822#[cfg(test)]
823mod tests {
824 use super::*;
825 use crate::core::config::CacheTtl;
826 use crate::extensions::hooks::events::{HookEvent, HookKind, HookResult};
827 use crate::extensions::permissions::{Permission, PermissionSet};
828 use agent_core::reasoning::ReasoningLevel;
829 use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::post, Router};
830 use futures::StreamExt;
831 use std::sync::atomic::AtomicUsize;
832 use std::sync::Mutex;
833
834 const SSE_HI: &str = concat!(
835 "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
836 "\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
837 "\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0,",
838 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}}\n\n",
839 "data: {\"type\":\"content_block_start\",\"index\":0,",
840 "\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
841 "data: {\"type\":\"content_block_delta\",\"index\":0,",
842 "\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n",
843 "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
844 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",",
845 "\"stop_sequence\":null},\"usage\":{\"input_tokens\":10,\"output_tokens\":1,",
846 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}\n\n",
847 "data: {\"type\":\"message_stop\"}\n\n",
848 );
849
850 const SSE_USAGE_THEN_ERROR: &str = concat!(
852 "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
853 "\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
854 "\"stop_sequence\":null,\"usage\":{\"input_tokens\":12,\"output_tokens\":0,",
855 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}}\n\n",
856 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":null,",
857 "\"stop_sequence\":null},\"usage\":{\"input_tokens\":12,\"output_tokens\":1,",
858 "\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}\n\n",
859 "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",",
860 "\"message\":\"synthetic\"}}\n\n",
861 );
862
863 #[derive(Clone)]
864 struct Mock {
865 bodies: Arc<Mutex<Vec<String>>>,
866 status: u16,
868 sse: &'static str,
869 }
870
871 async fn spawn_mock() -> (String, Mock) {
872 spawn_mock_full(200, SSE_HI).await
873 }
874
875 async fn spawn_mock_with(status: u16) -> (String, Mock) {
876 spawn_mock_full(status, SSE_HI).await
877 }
878
879 async fn spawn_mock_full(status: u16, sse: &'static str) -> (String, Mock) {
880 let mock = Mock {
881 bodies: Arc::default(),
882 status,
883 sse,
884 };
885 let app = Router::new()
886 .route(
887 "/v1/messages",
888 post(|State(m): State<Mock>, body: String| async move {
889 m.bodies.lock().unwrap().push(body);
890 if m.status == 200 {
891 (
892 StatusCode::OK,
893 [("content-type", "text/event-stream")],
894 m.sse.to_string(),
895 )
896 .into_response()
897 } else {
898 (
899 StatusCode::from_u16(m.status).unwrap(),
900 [("content-type", "application/json")],
901 r#"{"type":"error","error":{"type":"api_error","message":"synthetic"}}"#
902 .to_string(),
903 )
904 .into_response()
905 }
906 }),
907 )
908 .layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
909 .with_state(mock.clone());
910 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
911 let addr = listener.local_addr().unwrap();
912 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
913 (format!("http://{addr}"), mock)
914 }
915
916 struct CountingInject(Arc<AtomicUsize>);
918 #[async_trait::async_trait]
919 impl crate::extensions::runtime::ExtensionHandler for CountingInject {
920 fn id(&self) -> &str {
921 "keepwarm-counting-inject"
922 }
923 async fn handle(&self, _event: &HookEvent) -> HookResult {
924 self.0.fetch_add(1, Ordering::SeqCst);
925 HookResult::Inject {
926 content: "PER-TURN-INJECTED-CONTEXT".into(),
927 }
928 }
929 async fn shutdown(&self) {}
930 }
931
932 async fn real_turn(rt: &Runtime, messages: Vec<SharedMessage>) -> Vec<SharedMessage> {
934 let mut stream = rt
935 .run_stream_with_messages(messages, CancellationToken::new(), None, None, false)
936 .await;
937 let mut history = None;
938 while let Some(event) = stream.next().await {
939 match event {
940 StreamEvent::Session(SessionEvent::MessageHistory(h)) => history = Some(h),
941 StreamEvent::Session(SessionEvent::Error(e)) => panic!("turn failed: {e:?}"),
942 _ => {}
943 }
944 }
945 history.expect("turn published its history")
946 }
947
948 fn user(text: &str) -> SharedMessage {
949 Arc::new(json!({"role": "user", "content": text}))
950 }
951
952 fn canon(v: &Value) -> String {
953 serde_json::to_string(v).unwrap()
954 }
955
956 fn strip_cache_control(mut msg: Value) -> Value {
957 if let Some(blocks) = msg["content"].as_array_mut() {
958 for b in blocks {
959 if let Some(obj) = b.as_object_mut() {
960 obj.remove("cache_control");
961 }
962 }
963 }
964 msg
965 }
966
967 fn marker_count(v: &Value) -> usize {
968 match v {
969 Value::Object(map) => {
970 usize::from(map.contains_key("cache_control"))
971 + map.values().map(marker_count).sum::<usize>()
972 }
973 Value::Array(items) => items.iter().map(marker_count).sum(),
974 _ => 0,
975 }
976 }
977
978 struct Observed {
980 ping: Value,
981 real: Value,
982 raw_ping: String,
983 history_len: usize,
984 usage: KeepWarmUsage,
985 hook_calls: (usize, usize, usize),
987 }
988
989 async fn ping_then_next_turn(rt: &Runtime) -> Observed {
994 let (url, mock) = spawn_mock().await;
995 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
996 let calls = Arc::new(AtomicUsize::new(0));
997 let mut perms = PermissionSet::new();
998 perms.grant(Permission::LlmContent);
999 rt.hook_bus()
1000 .subscribe(
1001 HookKind::BeforeMessage,
1002 Arc::new(CountingInject(Arc::clone(&calls))),
1003 None,
1004 None,
1005 perms,
1006 )
1007 .await
1008 .unwrap();
1009
1010 assert_eq!(rt.cache_anchor(), None, "no request has run yet");
1011 let history = real_turn(rt, vec![user("hello")]).await;
1013 assert_eq!(history.last().unwrap()["role"], "assistant");
1014 let turn_anchor = rt
1015 .cache_anchor()
1016 .expect("an accepted turn request stamps the anchor");
1017 let before_ping = calls.load(Ordering::SeqCst);
1018 let snapshot = history.clone();
1019
1020 let fingerprint = rt
1021 .cache_keepwarm_probe(&history)
1022 .await
1023 .expect("probe")
1024 .fingerprint;
1025 assert_eq!(
1026 mock.bodies.lock().unwrap().len(),
1027 1,
1028 "a probe sends nothing"
1029 );
1030 assert_eq!(
1031 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1032 fingerprint,
1033 "the fingerprint is deterministic"
1034 );
1035
1036 let tap = PingTap::default();
1037 let report = rt
1038 .cache_keepwarm_ping(&history, fingerprint, &CancellationToken::new(), &tap)
1039 .await;
1040 assert_eq!(report.result, Ok(PingOutcome::Sent), "{report:?}");
1041 assert!(tap.was_sent(), "the tap saw the request leave");
1042 assert_eq!(tap.usage(), report.usage, "the tap saw the usage");
1043 let usage = report.usage;
1044 assert!(
1045 usage.output_tokens >= 1,
1046 "usage comes back from the private channel"
1047 );
1048 let ping_anchor = rt.cache_anchor().expect("anchor");
1049 assert!(
1050 ping_anchor.mono >= turn_anchor.mono,
1051 "the ping moved the anchor"
1052 );
1053 let after_ping = calls.load(Ordering::SeqCst);
1054 assert_eq!(history, snapshot, "the ping must not mutate history");
1055
1056 let mut next = history.clone();
1058 next.push(user("next question"));
1059 real_turn(rt, next).await;
1060 let after_next = calls.load(Ordering::SeqCst);
1061
1062 let bodies = mock.bodies.lock().unwrap().clone();
1063 assert_eq!(bodies.len(), 3, "previous turn, ping, next turn");
1064 Observed {
1065 ping: serde_json::from_str(&bodies[1]).unwrap(),
1066 real: serde_json::from_str(&bodies[2]).unwrap(),
1067 raw_ping: bodies[1].clone(),
1068 history_len: history.len(),
1069 usage,
1070 hook_calls: (before_ping, after_ping, after_next),
1071 }
1072 }
1073
1074 fn assert_prefix_identical(
1075 ping: &Value,
1076 real: &Value,
1077 n: usize,
1078 history_marker: &Value,
1079 ping_max_tokens: u64,
1080 ) {
1081 for key in [
1083 "model",
1084 "tools",
1085 "system",
1086 "thinking",
1087 "output_config",
1088 "stream",
1089 ] {
1090 assert_eq!(
1091 canon(&ping[key]),
1092 canon(&real[key]),
1093 "`{key}` differs between ping and the real next turn"
1094 );
1095 }
1096 let pm = ping["messages"].as_array().unwrap();
1097 let rm = real["messages"].as_array().unwrap();
1098 assert_eq!(pm.len(), n + 1, "history + one trailing ping message");
1099 assert!(rm.len() > n);
1100 for i in 0..n - 1 {
1102 assert_eq!(canon(&pm[i]), canon(&rm[i]), "message {i} differs");
1103 }
1104 let last = &pm[n - 1];
1107 let marker = &last["content"].as_array().unwrap().last().unwrap()["cache_control"];
1108 assert_eq!(marker, history_marker);
1109 assert_eq!(
1110 canon(&strip_cache_control(last.clone())),
1111 canon(&rm[n - 1]),
1112 "last history message differs beyond the cache marker"
1113 );
1114 assert_eq!(marker_count(&pm[n]), 0, "ping message must be unmarked");
1116 assert_eq!(pm[n]["role"], "user");
1117 assert_eq!(pm[n]["content"], KEEPWARM_PING_TEXT);
1118 assert_eq!(ping["max_tokens"], json!(ping_max_tokens));
1120 assert!(real["max_tokens"].as_u64().unwrap() > ping_max_tokens);
1121 assert!(
1122 marker_count(ping) <= 4,
1123 "Anthropic allows at most 4 breakpoints"
1124 );
1125 }
1126
1127 fn marker_1h() -> Value {
1128 json!({"type": "ephemeral", "ttl": "1h"})
1129 }
1130
1131 fn marker_5m() -> Value {
1132 json!({"type": "ephemeral"})
1133 }
1134
1135 fn opus(ttl: CacheTtl) -> Runtime {
1136 let mut rt = Runtime::new_headless();
1137 rt.set_model("claude-opus-4-7".to_string());
1138 assert_eq!(rt.model(), "claude-opus-4-7");
1139 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1140 rt.set_cache_ttl(ttl);
1141 rt.set_system_prompt("You are a keep-warm test.".into());
1142 assert_eq!(rt.cache_keepwarm_ineligibility(), None);
1143 rt
1144 }
1145
1146 #[tokio::test]
1150 #[serial_test::serial(synaps_base_dir)]
1151 async fn ping_prefix_matches_next_turn_with_adaptive_thinking() {
1152 let _base = crate::test_env::BaseDirGuard::new();
1153 let rt = opus(CacheTtl::OneHour);
1154
1155 let o = ping_then_next_turn(&rt).await;
1156 assert!(
1157 !o.ping["thinking"].is_null(),
1158 "thinking must stay on: {}",
1159 o.ping
1160 );
1161 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1162
1163 let (before, after_ping, after_next) = o.hook_calls;
1165 assert!(
1166 before >= 1,
1167 "control: the previous turn fired before_message"
1168 );
1169 assert_eq!(after_ping, before, "the ping fired a hook");
1170 assert!(
1171 after_next > after_ping,
1172 "control: the next turn fired before_message"
1173 );
1174 assert!(!o.raw_ping.contains("PER-TURN-INJECTED-CONTEXT"));
1175 assert!(!o.raw_ping.contains("Extension context"));
1176 assert!(canon(&o.real).contains("PER-TURN-INJECTED-CONTEXT"));
1178 }
1179
1180 #[tokio::test]
1182 #[serial_test::serial(synaps_base_dir)]
1183 async fn ping_prefix_matches_next_turn_with_5m_ttl() {
1184 let _base = crate::test_env::BaseDirGuard::new();
1185 let rt = opus(CacheTtl::FiveMinutes);
1186 let o = ping_then_next_turn(&rt).await;
1187 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
1188 }
1189
1190 #[tokio::test]
1193 #[serial_test::serial(synaps_base_dir)]
1194 async fn ping_prefix_matches_next_turn_with_hybrid_ttl() {
1195 let _base = crate::test_env::BaseDirGuard::new();
1196 let rt = opus(CacheTtl::Hybrid);
1197 let o = ping_then_next_turn(&rt).await;
1198 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
1199 assert!(
1200 canon(&o.ping["system"]).contains("\"ttl\":\"1h\""),
1201 "hybrid keeps the 1h stable-prefix marker: {}",
1202 o.ping["system"]
1203 );
1204 }
1205
1206 #[tokio::test]
1208 #[serial_test::serial(synaps_base_dir)]
1209 async fn ping_prefix_matches_next_turn_with_thinking_off() {
1210 let _base = crate::test_env::BaseDirGuard::new();
1211 let mut rt = Runtime::new_headless();
1212 rt.set_model("claude-sonnet-4-6".to_string());
1213 rt.set_reasoning_level_explicit(ReasoningLevel::Off);
1214 rt.set_cache_ttl(CacheTtl::OneHour);
1215 assert_eq!(rt.cache_keepwarm_ineligibility(), None);
1216
1217 let o = ping_then_next_turn(&rt).await;
1218 assert!(o.ping.get("thinking").is_none(), "thinking off is omitted");
1219 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1220 }
1221
1222 #[tokio::test]
1227 #[serial_test::serial(synaps_base_dir)]
1228 async fn budgeted_thinking_ping_keeps_budget_and_is_costed() {
1229 let _base = crate::test_env::BaseDirGuard::new();
1230 let mut rt = Runtime::new_headless();
1231 rt.set_model("claude-sonnet-4-6".to_string());
1232 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1233 rt.set_cache_ttl(CacheTtl::FiveMinutes);
1234 assert_eq!(
1235 rt.cache_keepwarm_ineligibility(),
1236 None,
1237 "budgeted thinking is eligible"
1238 );
1239
1240 let o = ping_then_next_turn(&rt).await;
1241 assert_eq!(o.ping["thinking"]["type"], "enabled", "{}", o.ping);
1242 let budget = o.ping["thinking"]["budget_tokens"]
1243 .as_u64()
1244 .expect("budget");
1245 assert_eq!(canon(&o.ping["thinking"]), canon(&o.real["thinking"]));
1246 assert_prefix_identical(
1247 &o.ping,
1248 &o.real,
1249 o.history_len,
1250 &marker_5m(),
1251 budget + KEEPWARM_OUTPUT_TOKENS,
1252 );
1253 assert!(
1254 o.usage.cost_usd("claude-sonnet-4-6") > 0.0,
1255 "the ping is costed"
1256 );
1257 }
1258
1259 #[tokio::test]
1263 #[serial_test::serial(synaps_base_dir)]
1264 async fn changed_prefix_is_refused_before_the_network() {
1265 let _base = crate::test_env::BaseDirGuard::new();
1266 let (url, mock) = spawn_mock().await;
1267 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1268 let mut rt = opus(CacheTtl::FiveMinutes);
1269 let history = real_turn(&rt, vec![user("hello")]).await;
1270 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1271 let sent = mock.bodies.lock().unwrap().len();
1272
1273 let mut longer = history.clone();
1275 longer.push(user("more"));
1276 longer.push(Arc::new(
1277 json!({"role": "assistant", "content": [{"type": "text", "text": "x"}]}),
1278 ));
1279 let changed = |r: PingReport| matches!(r.result, Ok(PingOutcome::PrefixChanged { .. }));
1280 let tap = PingTap::default();
1281 assert!(changed(
1282 rt.cache_keepwarm_ping(&longer, fp, &CancellationToken::new(), &tap)
1283 .await
1284 ));
1285
1286 rt.set_system_prompt("A different system prompt.".into());
1287 assert_ne!(
1288 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1289 fp
1290 );
1291 assert!(changed(
1292 rt.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
1293 .await
1294 ));
1295 rt.set_system_prompt("You are a keep-warm test.".into());
1296 assert_eq!(
1297 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1298 fp,
1299 "restored"
1300 );
1301
1302 rt.set_reasoning_level_explicit(ReasoningLevel::Low);
1303 assert_ne!(
1304 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1305 fp,
1306 "effort is in the key"
1307 );
1308 rt.set_reasoning_level_explicit(ReasoningLevel::High);
1309
1310 rt.set_model("claude-opus-4-6".to_string());
1311 assert_eq!(rt.model(), "claude-opus-4-6");
1312 assert_ne!(
1313 rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
1314 fp,
1315 "model is in the key"
1316 );
1317
1318 assert_eq!(
1319 mock.bodies.lock().unwrap().len(),
1320 sent,
1321 "no refused ping reached the provider"
1322 );
1323 assert!(!tap.was_sent(), "a refused ping never leaves");
1324 }
1325
1326 #[tokio::test]
1327 #[serial_test::serial(synaps_base_dir)]
1328 async fn http_4xx_is_fatal_and_5xx_is_transient() {
1329 let _base = crate::test_env::BaseDirGuard::new();
1330 let history = vec![
1331 user("hello"),
1332 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1333 ];
1334 for (status, fatal) in [(400u16, true), (404, true), (500, false), (429, false)] {
1335 let (url, mock) = spawn_mock_with(status).await;
1336 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1337 let rt = opus(CacheTtl::FiveMinutes);
1338 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1339 let err = rt
1340 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1341 .await
1342 .result
1343 .expect_err("non-2xx fails");
1344 assert_eq!(
1345 matches!(err, PingFailure::Fatal(_)),
1346 fatal,
1347 "HTTP {status}: {err:?}"
1348 );
1349 assert_eq!(
1350 mock.bodies.lock().unwrap().len(),
1351 1,
1352 "HTTP {status}: no retry"
1353 );
1354 assert_eq!(
1355 rt.cache_anchor(),
1356 None,
1357 "a rejected request refreshes nothing"
1358 );
1359 }
1360 }
1361
1362 async fn assert_identity_1h(rt: &Runtime) -> Observed {
1366 let o = ping_then_next_turn(rt).await;
1367 assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
1368 o
1369 }
1370
1371 #[tokio::test]
1375 #[serial_test::serial(synaps_base_dir)]
1376 async fn ping_prefix_matches_next_turn_with_progressive_disclosure() {
1377 let _base = crate::test_env::BaseDirGuard::new();
1378 let mut rt = opus(CacheTtl::OneHour);
1379 rt.progressive_tool_disclosure = true;
1380 assert_identity_1h(&rt).await;
1381 }
1382
1383 #[tokio::test]
1386 #[serial_test::serial(synaps_base_dir)]
1387 async fn ping_prefix_matches_next_turn_with_deferred_tools() {
1388 let _base = crate::test_env::BaseDirGuard::new();
1389 let mut rt = opus(CacheTtl::OneHour);
1390 rt.progressive_tool_disclosure = true;
1391 rt.deferred_tool_loading = true;
1392 assert_identity_1h(&rt).await;
1393 }
1394
1395 #[tokio::test]
1401 #[serial_test::serial(synaps_base_dir)]
1402 async fn ping_projects_carried_activations_the_turn_start_restores() {
1403 use crate::tools::activation::{activate_exact_for_user, SessionId, SessionToolSet};
1404 use crate::tools::catalog::ToolId;
1405 let _base = crate::test_env::BaseDirGuard::new();
1406 let (url, mock) = spawn_mock().await;
1407 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1408 let mut rt = opus(CacheTtl::OneHour);
1409 rt.progressive_tool_disclosure = true;
1410 rt.deferred_tool_loading = true;
1411 let carried = {
1412 let registry = rt.tools.read().await;
1413 let mut previous = SessionToolSet::progressive_core_for_catalog(
1414 SessionId::parse("previous-runtime").unwrap(),
1415 registry.catalog(),
1416 );
1417 activate_exact_for_user(
1418 &mut previous,
1419 registry.catalog(),
1420 &ToolId::builtin("subagent_status"),
1421 )
1422 .unwrap();
1423 previous.carried_activations()
1424 };
1425 rt.restore_tool_activations(carried);
1426
1427 let history = vec![
1428 user("hello"),
1429 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1430 ];
1431 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1432 let report = rt
1433 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1434 .await;
1435 assert_eq!(report.result, Ok(PingOutcome::Sent), "{report:?}");
1436 let mut next = history.clone();
1437 next.push(user("next"));
1438 real_turn(&rt, next).await;
1439
1440 let bodies: Vec<Value> = mock
1441 .bodies
1442 .lock()
1443 .unwrap()
1444 .iter()
1445 .map(|b| serde_json::from_str(b).unwrap())
1446 .collect();
1447 let (ping, real) = (&bodies[0], &bodies[1]);
1448 let loaded = |b: &Value| -> Vec<Value> {
1453 b["tools"]
1454 .as_array()
1455 .unwrap()
1456 .iter()
1457 .filter(|t| t.get("defer_loading").is_none())
1458 .cloned()
1459 .collect()
1460 };
1461 assert_eq!(
1462 serde_json::to_string(&loaded(ping)).unwrap(),
1463 serde_json::to_string(&loaded(real)).unwrap()
1464 );
1465 assert!(
1466 real["tools"].as_array().unwrap().iter().any(|t| {
1467 t["name"] == "subagent_status" && t.get("defer_loading").is_none()
1468 }),
1469 "the restored activation is a loaded tool: {}",
1470 real["tools"]
1471 );
1472 }
1473
1474 #[tokio::test]
1477 #[serial_test::serial(synaps_base_dir)]
1478 async fn ping_prefix_matches_next_turn_with_context_management_guidance() {
1479 let _base = crate::test_env::BaseDirGuard::new();
1480 let rt = opus(CacheTtl::OneHour);
1481 rt.continuation.lock().unwrap().config.mode =
1482 crate::core::config::ContextManagementMode::Auto;
1483 let o = assert_identity_1h(&rt).await;
1484 assert!(
1485 canon(&o.ping["system"]).contains("Work normally."),
1486 "guidance present: {}",
1487 o.ping["system"]
1488 );
1489 }
1490
1491 #[tokio::test]
1494 #[serial_test::serial(synaps_base_dir)]
1495 async fn ping_prefix_matches_next_turn_with_session_injection() {
1496 let _base = crate::test_env::BaseDirGuard::new();
1497 let mut rt = opus(CacheTtl::OneHour);
1498 rt.set_session_id(Some("kw-session".into()));
1499 rt.hook_bus()
1500 .set_session_injection_for("kw-session", "SESSION-SCOPED-INJECTION".into())
1501 .await;
1502 let o = assert_identity_1h(&rt).await;
1503 assert!(canon(&o.ping["system"]).contains("SESSION-SCOPED-INJECTION"));
1504 }
1505
1506 fn bind_oauth(rt: &mut Runtime) {
1510 rt.auth = Arc::new(tokio::sync::RwLock::new(super::super::types::AuthState {
1511 auth_token: "oauth-test-token".to_string(),
1512 auth_type: "oauth".to_string(),
1513 refresh_token: None,
1514 token_expires: Some(crate::epoch_millis() + 3_600_000),
1515 auth_file_mtime: None,
1516 bound_credential: Some(crate::runtime::auth::anthropic_binding(
1517 &crate::auth::CredentialSource::Local,
1518 &crate::auth::CredentialRef::default_for(crate::auth::OAuthProviderId::Anthropic),
1519 )),
1520 }));
1521 }
1522
1523 #[tokio::test]
1528 #[serial_test::serial(synaps_base_dir)]
1529 async fn an_idle_ping_on_a_cooling_account_sends_nothing() {
1530 let _base = crate::test_env::BaseDirGuard::new();
1531 let (url, mock) = spawn_mock().await;
1532 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1533 let mut rt = opus(CacheTtl::OneHour);
1534 bind_oauth(&mut rt);
1535 let history = real_turn(&rt, vec![user("hello")]).await;
1536 let fingerprint = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1537 let seat = crate::auth::CredentialRef::default_for(crate::auth::OAuthProviderId::Anthropic);
1538 crate::auth::broker_from_source(&rt.credential_source, &rt.token_cache, reqwest::Client::new())
1539 .report_cooldown(&seat, Some(crate::epoch_millis() + 3_600_000), "test")
1540 .await
1541 .unwrap();
1542 let bound = rt.auth.read().await.bound_credential.clone();
1543
1544 let tap = PingTap::default();
1545 let report = rt
1546 .cache_keepwarm_ping(&history, fingerprint, &CancellationToken::new(), &tap)
1547 .await;
1548 match &report.result {
1549 Err(PingFailure::Fatal(why)) => assert!(why.contains("account cooling"), "{why}"),
1550 other => panic!("the ping stops on the cooling account: {other:?}"),
1551 }
1552 assert!(!tap.was_sent());
1553 assert!(report.notices.is_empty(), "{:?}", report.notices);
1554 assert_eq!(mock.bodies.lock().unwrap().len(), 1, "only the turn was sent");
1555 assert_eq!(rt.auth.read().await.bound_credential, bound, "not re-bound");
1556 assert_eq!(rt.auth.read().await.auth_token, "oauth-test-token");
1557 }
1558
1559 fn markers(v: &Value) -> Vec<Value> {
1561 match v {
1562 Value::Object(map) => map
1563 .iter()
1564 .flat_map(|(k, v)| {
1565 if k == "cache_control" {
1566 vec![v.clone()]
1567 } else {
1568 markers(v)
1569 }
1570 })
1571 .collect(),
1572 Value::Array(items) => items.iter().flat_map(markers).collect(),
1573 _ => Vec::new(),
1574 }
1575 }
1576
1577 #[tokio::test]
1583 #[serial_test::serial(synaps_base_dir)]
1584 async fn auto_ttl_wire_markers_follow_the_credential_kind() {
1585 use crate::core::config::CacheTtlSetting;
1586 let cases = [
1587 (CacheTtlSetting::Auto, true, CacheTtl::OneHour),
1588 (CacheTtlSetting::Auto, false, CacheTtl::FiveMinutes),
1589 (CacheTtl::FiveMinutes.into(), true, CacheTtl::FiveMinutes),
1590 (CacheTtl::OneHour.into(), false, CacheTtl::OneHour),
1591 ];
1592 for (setting, oauth, expect) in cases {
1593 let ctx = format!("{setting} on {}", if oauth { "oauth" } else { "api_key" });
1594 let observe = |setting: CacheTtlSetting| async move {
1595 let _base = crate::test_env::BaseDirGuard::new();
1596 let mut rt = opus(CacheTtl::FiveMinutes);
1597 rt.set_cache_ttl(setting);
1598 if oauth {
1599 bind_oauth(&mut rt);
1600 }
1601 let resolved = rt.resolved_cache_ttl().await;
1602 (resolved, ping_then_next_turn(&rt).await)
1603 };
1604 let (resolved, o) = observe(setting).await;
1605 assert_eq!(resolved, expect, "{ctx}");
1606 let want = if expect == CacheTtl::OneHour {
1607 marker_1h()
1608 } else {
1609 marker_5m()
1610 };
1611 assert_prefix_identical(&o.ping, &o.real, o.history_len, &want, 1);
1612 for (name, body) in [("ping", &o.ping), ("next turn", &o.real)] {
1613 let all = markers(body);
1614 assert!(!all.is_empty(), "{ctx}: {name} has no markers");
1615 assert!(
1616 all.iter().all(|m| *m == want),
1617 "{ctx}: {name} markers {all:?}, want all {want}"
1618 );
1619 }
1620 if setting == CacheTtlSetting::Auto {
1621 let (_, fixed) = observe(expect.into()).await;
1622 assert_eq!(canon(&o.ping), canon(&fixed.ping), "{ctx}: ping bytes");
1623 assert_eq!(canon(&o.real), canon(&fixed.real), "{ctx}: turn bytes");
1624 }
1625 }
1626 }
1627
1628 #[tokio::test]
1631 #[serial_test::serial(synaps_base_dir)]
1632 async fn ping_prefix_matches_next_turn_with_oauth_system_blocks() {
1633 let _base = crate::test_env::BaseDirGuard::new();
1634 let mut rt = opus(CacheTtl::OneHour);
1635 bind_oauth(&mut rt);
1636 let o = assert_identity_1h(&rt).await;
1637 assert!(
1638 canon(&o.ping["system"]).contains("You are Claude Code"),
1639 "{}",
1640 o.ping["system"]
1641 );
1642 }
1643
1644 #[tokio::test]
1648 #[serial_test::serial(synaps_base_dir)]
1649 async fn probe_reports_a_context_rollover_ahead() {
1650 let _base = crate::test_env::BaseDirGuard::new();
1651 let (url, _mock) = spawn_mock().await;
1652 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1653 let rt = opus(CacheTtl::OneHour);
1654 rt.continuation.lock().unwrap().config.mode =
1655 crate::core::config::ContextManagementMode::Auto;
1656 let history = real_turn(&rt, vec![user("hello")]).await;
1657 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1658 assert_eq!(probe.context_blocks, None, "plenty of room");
1659 let before = rt.continuation.lock().unwrap().policy;
1660 rt.continuation.lock().unwrap().config.reserve_tokens = 10_000_000;
1662 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1663 assert!(probe.context_blocks.is_some(), "{probe:?}");
1664 assert_eq!(
1665 rt.continuation.lock().unwrap().policy,
1666 before,
1667 "the probe never commits policy state"
1668 );
1669 rt.continuation.lock().unwrap().config.mode =
1671 crate::core::config::ContextManagementMode::Off;
1672 assert_eq!(
1673 rt.cache_keepwarm_probe(&history)
1674 .await
1675 .unwrap()
1676 .context_blocks,
1677 None
1678 );
1679 }
1680
1681 #[tokio::test]
1686 #[serial_test::serial(synaps_base_dir)]
1687 async fn probe_warms_an_opus_200k_session_past_the_old_hard_stop() {
1688 let _base = crate::test_env::BaseDirGuard::new();
1689 let (url, _mock) = spawn_mock().await;
1690 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1691 let rt = opus(CacheTtl::OneHour);
1692 assert_eq!(rt.context_window(), 200_000);
1693 rt.continuation.lock().unwrap().config.mode =
1694 crate::core::config::ContextManagementMode::Auto;
1695 let history = real_turn(&rt, vec![user(&"x".repeat(150_000))]).await;
1696 let used = rt.assess_context(&history).await.used_tokens();
1697 assert!((60_000..100_000).contains(&used), "{used}");
1698 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1699 assert_eq!(probe.context_blocks, None, "{probe:?}");
1700 }
1701
1702 #[tokio::test]
1704 #[serial_test::serial(synaps_base_dir)]
1705 async fn probe_reports_max_tokens() {
1706 let _base = crate::test_env::BaseDirGuard::new();
1707 let history = vec![
1708 user("hello"),
1709 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1710 ];
1711 let rt = opus(CacheTtl::OneHour);
1712 assert_eq!(
1713 rt.cache_keepwarm_probe(&history).await.unwrap().max_tokens,
1714 1
1715 );
1716 let mut rt = Runtime::new_headless();
1717 rt.set_model("claude-sonnet-4-6".to_string());
1718 rt.set_reasoning_level_explicit(ReasoningLevel::XHigh);
1719 let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
1720 let budget = u64::from(rt.thinking_budget);
1721 assert!(budget >= 16_384, "budgeted thinking at xhigh: {budget}");
1722 assert_eq!(
1723 probe.max_tokens,
1724 budget + KEEPWARM_OUTPUT_TOKENS,
1725 "{probe:?}"
1726 );
1727 }
1728
1729 #[tokio::test]
1732 #[serial_test::serial(synaps_base_dir)]
1733 async fn ping_usage_is_returned_on_error() {
1734 let _base = crate::test_env::BaseDirGuard::new();
1735 let (url, mock) = spawn_mock_full(200, SSE_USAGE_THEN_ERROR).await;
1736 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1737 let history = vec![
1738 user("hello"),
1739 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1740 ];
1741 let rt = opus(CacheTtl::OneHour);
1742 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
1743 let tap = PingTap::default();
1744 let report = rt
1745 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
1746 .await;
1747 assert!(
1748 matches!(report.result, Err(PingFailure::Transient(_))),
1749 "{report:?}"
1750 );
1751 assert_eq!(report.usage.cache_read_tokens, 5_000, "{report:?}");
1752 assert!(report.usage.cost_usd("claude-opus-4-7") > 0.0);
1753 assert_eq!(tap.usage(), report.usage);
1754 assert!(tap.was_sent());
1755 assert_eq!(mock.bodies.lock().unwrap().len(), 1, "no transport retry");
1756 }
1757
1758 #[tokio::test]
1762 #[serial_test::serial(synaps_base_dir)]
1763 async fn ping_never_carries_or_moves_the_diagnostics_chain() {
1764 let _base = crate::test_env::BaseDirGuard::new();
1765 let (url, mock) = spawn_mock().await;
1766 let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
1767 let history = vec![
1768 user("hello"),
1769 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1770 ];
1771 let mut rt = opus(CacheTtl::OneHour);
1772 rt.set_cache_diagnostics(true);
1773 rt.cache_diagnosis().unwrap().record("msg_turn");
1774 let mut ping_rt = rt.clone();
1775 ping_rt.share_ttl_latches(&rt);
1776 let fp = ping_rt
1777 .cache_keepwarm_probe(&history)
1778 .await
1779 .unwrap()
1780 .fingerprint;
1781 let report = ping_rt
1782 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
1783 .await;
1784 assert!(report.result.is_ok(), "{report:?}");
1785 let bodies = mock.bodies.lock().unwrap().clone();
1786 assert_eq!(bodies.len(), 1);
1787 assert!(!bodies[0].contains("\"diagnostics\""), "{}", bodies[0]);
1788 assert_eq!(
1789 rt.cache_diagnosis()
1790 .unwrap()
1791 .previous_message_id()
1792 .as_deref(),
1793 Some("msg_turn"),
1794 "the ping's msg_kw never enters the chain"
1795 );
1796 }
1797
1798 #[test]
1802 fn diagnostics_chain_is_shared_like_the_ttl_latches() {
1803 let mut rt = opus(CacheTtl::FiveMinutes);
1804 assert!(rt.cache_diagnosis().is_none(), "off: wire unchanged");
1805 rt.set_cache_diagnostics(true);
1806 rt.cache_diagnosis().unwrap().record("msg_a");
1807 let fresh = rt.clone();
1808 assert_eq!(fresh.cache_diagnosis().unwrap().previous_message_id(), None);
1809 let mut shared = rt.clone();
1810 shared.share_ttl_latches(&rt);
1811 shared.cache_diagnosis().unwrap().record("msg_b");
1812 shared.cache_diagnosis().unwrap().reject(400);
1813 let chain = rt.cache_diagnosis().unwrap();
1814 assert_eq!(chain.previous_message_id().as_deref(), Some("msg_b"));
1815 assert!(!chain.active(), "the rejection latch is session-wide");
1816 let fresh_chain = fresh.cache_diagnosis().unwrap();
1819 assert!(!fresh_chain.active(), "latch shared across Clone");
1820 assert_eq!(fresh_chain.previous_message_id(), None);
1821 assert!(!rt.clone().cache_diagnosis().unwrap().active());
1822 }
1823
1824 #[test]
1827 fn ping_shares_the_session_ttl_latch() {
1828 let rt = opus(CacheTtl::OneHour);
1829 assert!(!rt.cache_ttl_downgraded());
1830 let mut ping_rt = rt.clone();
1831 ping_rt.share_ttl_latches(&rt);
1832 ping_rt
1833 .ttl_downgrade_notified
1834 .store(true, Ordering::Relaxed);
1835 assert!(rt.cache_ttl_downgraded(), "the latch is shared");
1836 }
1837
1838 #[test]
1839 fn failure_classification() {
1840 let api = || RuntimeError::ApiStatus("x".into());
1841 assert_eq!(
1842 classify_failure(&RuntimeError::Canceled, 0),
1843 PingFailure::Canceled
1844 );
1845 assert!(matches!(
1846 classify_failure(&api(), 400),
1847 PingFailure::Fatal(_)
1848 ));
1849 assert!(matches!(
1850 classify_failure(&api(), 401),
1851 PingFailure::Fatal(_)
1852 ));
1853 assert!(matches!(
1854 classify_failure(&api(), 429),
1855 PingFailure::Transient(_)
1856 ));
1857 assert!(matches!(
1858 classify_failure(&api(), 529),
1859 PingFailure::Transient(_)
1860 ));
1861 assert!(matches!(
1862 classify_failure(&api(), 0),
1863 PingFailure::Transient(_)
1864 ));
1865 assert!(matches!(
1867 classify_failure(&api(), 200),
1868 PingFailure::Transient(_)
1869 ));
1870 assert!(matches!(
1871 classify_failure(&RuntimeError::Config("x".into()), 0),
1872 PingFailure::Fatal(_)
1873 ));
1874 }
1875
1876 #[test]
1880 fn network_timeout_is_transient() {
1881 for msg in ["upload/header deadline", "stream idle"] {
1882 assert_eq!(
1883 classify_failure(&RuntimeError::NetworkTimeout(msg), 0),
1884 PingFailure::Transient("network".to_string()),
1885 "{msg}"
1886 );
1887 }
1888 }
1889
1890 #[test]
1891 fn every_ttl_and_thinking_mode_is_eligible() {
1892 for ttl in [CacheTtl::FiveMinutes, CacheTtl::OneHour, CacheTtl::Hybrid] {
1893 for (model, level) in [
1894 ("claude-opus-4-7", ReasoningLevel::High),
1895 ("claude-sonnet-4-6", ReasoningLevel::High),
1896 ("claude-sonnet-4-6", ReasoningLevel::Off),
1897 ] {
1898 let mut rt = Runtime::new_headless();
1899 rt.set_model(model.to_string());
1900 rt.set_reasoning_level_explicit(level);
1901 rt.set_cache_ttl(ttl);
1902 assert_eq!(
1903 rt.cache_keepwarm_ineligibility(),
1904 None,
1905 "{model} {level:?} {ttl:?}"
1906 );
1907 }
1908 }
1909 }
1910
1911 #[test]
1912 fn non_anthropic_models_are_ineligible() {
1913 let mut rt = Runtime::new_headless();
1914 rt.set_model("groq/llama-3.3-70b-versatile".to_string());
1915 if rt.model() == "groq/llama-3.3-70b-versatile" {
1916 assert!(rt.cache_keepwarm_ineligibility().is_some());
1917 }
1918 }
1919
1920 #[tokio::test]
1923 async fn ineligible_ping_sends_nothing() {
1924 let mut rt = Runtime::new_headless();
1925 rt.set_model("groq/llama-3.3-70b-versatile".to_string());
1926 if rt.model() != "groq/llama-3.3-70b-versatile" {
1927 return;
1928 }
1929 let history = vec![
1930 user("hello"),
1931 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
1932 ];
1933 let tap = PingTap::default();
1934 let err = rt
1935 .cache_keepwarm_ping(&history, 0, &CancellationToken::new(), &tap)
1936 .await
1937 .result
1938 .expect_err("non-Anthropic must refuse");
1939 assert!(!tap.was_sent());
1940 assert!(
1941 matches!(err, PingFailure::Fatal(ref m) if m.contains("keep-warm skipped")),
1942 "{err:?}"
1943 );
1944 }
1945
1946 #[test]
1947 fn anchor_cell_only_moves_forward() {
1948 let cell = CacheAnchorCell::default();
1949 assert_eq!(cell.get(), None);
1950 let a = CacheAnchor::now();
1951 let later = CacheAnchor {
1952 mono: a.mono + std::time::Duration::from_secs(5),
1953 wall: a.wall + std::time::Duration::from_secs(5),
1954 };
1955 cell.stamp(later);
1956 cell.stamp(a);
1957 assert_eq!(
1958 cell.get(),
1959 Some(later),
1960 "a late-arriving older stamp never rewinds"
1961 );
1962 }
1963
1964 #[test]
1965 fn only_a_finished_assistant_text_reply_is_markable() {
1966 assert!(tail_is_markable(
1967 &json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}, {"type": "text", "text": "ok"}]})
1968 ));
1969 assert!(!tail_is_markable(&json!({"role": "user", "content": "q"})));
1970 assert!(!tail_is_markable(
1971 &json!({"role": "assistant", "content": [{"type": "tool_use", "id": "t", "name": "x", "input": {}}]})
1972 ));
1973 assert!(!tail_is_markable(
1974 &json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}]})
1975 ));
1976 }
1977
1978 #[test]
1979 fn usage_is_summed_and_priced() {
1980 let mut u = KeepWarmUsage::default();
1981 u.add(&StreamEvent::Session(SessionEvent::Usage {
1982 input_tokens: 20,
1983 output_tokens: 16,
1984 cache_read_input_tokens: 500_000,
1985 cache_creation_input_tokens: 0,
1986 cache_creation_5m: None,
1987 cache_creation_1h: Some(0),
1988 model: Some("claude-opus-4-7".into()),
1989 }));
1990 assert_eq!(u.cache_read_tokens, 500_000);
1991 assert_eq!(u.output_tokens, 16);
1992 let cost = u.cost_usd("claude-opus-4-7");
1994 assert!((cost - 0.2505).abs() < 0.001, "{cost}");
1995 }
1996
1997 #[test]
1999 fn booking_rule_table() {
2000 use PingBooking::*;
2001 let table = [
2003 ((true, true, false, 200), Measured),
2004 ((true, false, false, 0), Measured),
2005 ((true, true, false, 529), Measured),
2006 ((false, true, false, 200), WorstCase), ((false, true, false, 0), WorstCase), ((false, true, false, 500), Nothing), ((false, true, false, 429), Nothing),
2010 ((false, true, false, 529), Nothing),
2011 ((false, true, true, 0), Nothing), ((false, true, true, 503), Nothing), ((false, false, false, 0), Nothing), ((false, false, true, 0), Nothing),
2015 ];
2016 for ((usage, sent, connect, status), want) in table {
2017 assert_eq!(
2018 ping_booking(usage, sent, connect, status),
2019 want,
2020 "usage={usage} sent={sent} connect_failed={connect} status={status}"
2021 );
2022 }
2023 let tap = PingTap::default();
2025 assert_eq!(tap.bill("claude-opus-4-7", 0.5).cost_usd, 0.0);
2026 tap.sent.store(true, Ordering::Relaxed);
2027 let bill = tap.bill("claude-opus-4-7", 0.5);
2028 assert_eq!((bill.booking, bill.cost_usd), (WorstCase, 0.5));
2029 tap.usage.lock().unwrap().cache_read_tokens = 1_000_000;
2030 let bill = tap.bill("claude-opus-4-7", 0.5);
2031 assert_eq!(bill.booking, Measured);
2032 assert!((bill.cost_usd - 0.5).abs() < 1e-9, "1M reads at $0.50/MTok");
2033 }
2034
2035 fn two_message_history() -> Vec<SharedMessage> {
2036 vec![
2037 user("hello"),
2038 Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
2039 ]
2040 }
2041
2042 #[tokio::test]
2044 #[serial_test::serial(synaps_base_dir)]
2045 async fn a_refused_connection_is_never_sent_and_books_nothing() {
2046 let _base = crate::test_env::BaseDirGuard::new();
2047 let addr = {
2048 let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
2049 l.local_addr().unwrap()
2050 }; let _url = crate::test_env::EnvVarGuard::set(
2052 "SYNAPS_ANTHROPIC_BASE_URL",
2053 &format!("http://{addr}"),
2054 );
2055 let history = two_message_history();
2056 let rt = opus(CacheTtl::OneHour);
2057 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
2058 let tap = PingTap::default();
2059 let report = rt
2060 .cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
2061 .await;
2062 assert!(report.result.is_err(), "{report:?}");
2063 assert!(!tap.was_sent(), "never dispatched");
2064 assert!(tap.connect_failed.load(Ordering::Relaxed));
2065 let bill = tap.bill("claude-opus-4-7", 0.75);
2066 assert_eq!((bill.booking, bill.cost_usd), (PingBooking::Nothing, 0.0));
2067 }
2068
2069 #[tokio::test]
2072 #[serial_test::serial(synaps_base_dir)]
2073 async fn a_dispatched_ping_without_usage_books_the_worst_case() {
2074 use tokio::io::AsyncReadExt;
2075 let _base = crate::test_env::BaseDirGuard::new();
2076 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2077 let addr = listener.local_addr().unwrap();
2078 tokio::spawn(async move {
2080 let (mut sock, _) = listener.accept().await.unwrap();
2081 let mut buf = vec![0u8; 64 * 1024];
2082 while sock.read(&mut buf).await.is_ok_and(|n| n > 0) {}
2083 });
2084 let _url = crate::test_env::EnvVarGuard::set(
2085 "SYNAPS_ANTHROPIC_BASE_URL",
2086 &format!("http://{addr}"),
2087 );
2088 let history = two_message_history();
2089 let rt = opus(CacheTtl::OneHour);
2090 let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
2091 let tap = PingTap::default();
2092 let cancel = CancellationToken::new();
2093 let ping = {
2094 let (rt, history, tap, cancel) =
2095 (rt.clone(), history.clone(), tap.clone(), cancel.clone());
2096 tokio::spawn(async move { rt.cache_keepwarm_ping(&history, fp, &cancel, &tap).await })
2097 };
2098 tokio::time::timeout(std::time::Duration::from_secs(10), async {
2099 while !tap.was_sent() {
2100 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
2101 }
2102 })
2103 .await
2104 .expect("the body was pulled over the established connection");
2105 cancel.cancel();
2106 let report = ping.await.unwrap();
2107 assert_eq!(report.result, Err(PingFailure::Canceled));
2108 assert!(!tap.connect_failed.load(Ordering::Relaxed));
2109 let bill = tap.bill("claude-opus-4-7", 0.75);
2110 assert_eq!(
2111 (bill.booking, bill.cost_usd),
2112 (PingBooking::WorstCase, 0.75)
2113 );
2114 }
2115
2116 #[test]
2117 fn reply_prefix_tokens_skip_thinking() {
2118 let text = "a".repeat(250); let with_thinking = vec![Arc::new(json!({"role": "assistant", "content": [
2120 {"type": "thinking", "thinking": "x".repeat(50_000), "signature": "s"},
2121 {"type": "text", "text": text},
2122 ]}))];
2123 let plain = vec![Arc::new(json!({"role": "assistant", "content": [
2124 {"type": "text", "text": text},
2125 ]}))];
2126 let t = reply_prefix_tokens(&with_thinking);
2127 assert_eq!(t, reply_prefix_tokens(&plain), "thinking is not counted");
2128 assert!((100..200).contains(&t), "{t}");
2129 assert_eq!(reply_prefix_tokens(&[]), 0);
2130 }
2131}