use std::sync::atomic::{AtomicU16, Ordering};
use std::sync::Arc;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::api::{ApiMethods, ApiOptions, KeepWarmMode, KeepWarmRequest};
use super::helpers::HelperMethods;
use super::{Runtime, SessionEvent, StreamEvent};
use crate::{Result, RuntimeError, SharedMessage};
pub const KEEPWARM_OUTPUT_TOKENS: u64 = 1;
pub const KEEPWARM_PING_TEXT: &str = "[automated prompt-cache keep-warm ping, not from the user] \
This needs no reasoning. Reply with exactly: ok";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CacheAnchor {
pub mono: tokio::time::Instant,
pub wall: std::time::SystemTime,
}
impl CacheAnchor {
pub fn now() -> Self {
Self {
mono: tokio::time::Instant::now(),
wall: std::time::SystemTime::now(),
}
}
}
#[derive(Debug, Default)]
pub struct CacheAnchorCell(std::sync::Mutex<Option<CacheAnchor>>);
impl CacheAnchorCell {
pub fn stamp(&self, at: CacheAnchor) {
let mut slot = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if slot.map_or(true, |prev| prev.mono <= at.mono) {
*slot = Some(at);
}
}
pub fn get(&self) -> Option<CacheAnchor> {
*self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct KeepWarmUsage {
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_read_tokens: u64,
pub cache_write_tokens: u64,
pub cache_write_5m: Option<u64>,
pub cache_write_1h: Option<u64>,
pub model: Option<String>,
}
impl KeepWarmUsage {
pub(crate) fn add(&mut self, event: &StreamEvent) {
if let StreamEvent::Session(SessionEvent::Usage {
input_tokens,
output_tokens,
cache_read_input_tokens,
cache_creation_input_tokens,
cache_creation_5m,
cache_creation_1h,
model,
}) = event
{
self.input_tokens += input_tokens;
self.output_tokens += output_tokens;
self.cache_read_tokens += cache_read_input_tokens;
self.cache_write_tokens += cache_creation_input_tokens;
if let Some(n) = cache_creation_5m {
*self.cache_write_5m.get_or_insert(0) += n;
}
if let Some(n) = cache_creation_1h {
*self.cache_write_1h.get_or_insert(0) += n;
}
if model.is_some() {
self.model = model.clone();
}
}
}
pub fn cost_usd(&self, fallback_model: &str) -> f64 {
crate::pricing::calculate_cost_optional_split(
self.model.as_deref().unwrap_or(fallback_model),
self.input_tokens,
self.output_tokens,
self.cache_read_tokens,
self.cache_write_tokens,
self.cache_write_5m,
self.cache_write_1h,
)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum PingOutcome {
Sent,
PrefixChanged { expected: u64, actual: u64 },
}
#[derive(Debug, Clone, PartialEq)]
pub struct PingReport {
pub result: std::result::Result<PingOutcome, PingFailure>,
pub usage: KeepWarmUsage,
pub notices: Vec<String>,
}
#[derive(Debug, Clone, Default)]
pub struct PingTap {
pub sent: Arc<std::sync::atomic::AtomicBool>,
pub connect_failed: Arc<std::sync::atomic::AtomicBool>,
pub status: Arc<AtomicU16>,
pub usage: Arc<std::sync::Mutex<KeepWarmUsage>>,
}
impl PingTap {
pub fn was_sent(&self) -> bool {
self.sent.load(Ordering::Relaxed)
}
pub fn usage(&self) -> KeepWarmUsage {
self.usage
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
pub fn bill(&self, model: &str, worst_usd: f64) -> PingBill {
let usage = self.usage();
let booking = ping_booking(
usage != KeepWarmUsage::default(),
self.was_sent(),
self.connect_failed.load(Ordering::Relaxed),
self.status.load(Ordering::Relaxed),
);
let cost_usd = match booking {
PingBooking::Measured => usage.cost_usd(model),
PingBooking::WorstCase => worst_usd,
PingBooking::Nothing => 0.0,
};
PingBill {
booking,
usage,
cost_usd,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PingBooking {
Measured,
WorstCase,
Nothing,
}
pub fn ping_booking(
usage_arrived: bool,
sent: bool,
connect_failed: bool,
http_status: u16,
) -> PingBooking {
if usage_arrived {
PingBooking::Measured
} else if http_status >= 400 {
PingBooking::Nothing
} else if sent && !(connect_failed && http_status == 0) {
PingBooking::WorstCase
} else {
PingBooking::Nothing
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct PingBill {
pub booking: PingBooking,
pub usage: KeepWarmUsage,
pub cost_usd: f64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct KeepWarmProbe {
pub fingerprint: u64,
pub max_tokens: u64,
pub context_blocks: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PingFailure {
Transient(String),
Fatal(String),
Canceled,
}
impl std::fmt::Display for PingFailure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PingFailure::Transient(m) => write!(f, "transient: {m}"),
PingFailure::Fatal(m) => write!(f, "{m}"),
PingFailure::Canceled => f.write_str("canceled"),
}
}
}
pub(crate) fn classify_failure(err: &RuntimeError, status: u16) -> PingFailure {
if matches!(err, RuntimeError::Canceled) {
return PingFailure::Canceled;
}
if (400..500).contains(&status) && status != 429 {
return PingFailure::Fatal(format!("HTTP {status}"));
}
match err {
RuntimeError::Api(_) | RuntimeError::ApiStatus(_) | RuntimeError::Timeout => {
PingFailure::Transient(if status == 0 {
"network".to_string()
} else {
format!("HTTP {status}")
})
}
other => PingFailure::Fatal(other.to_string()),
}
}
fn tail_is_markable(last: &Value) -> bool {
if last["role"] != "assistant" {
return false;
}
match last["content"].as_array() {
Some(blocks) => blocks.last().is_some_and(|b| {
b["type"] == "text" && b["text"].as_str().is_some_and(|t| !t.is_empty())
}),
None => last["content"].as_str().is_some_and(|t| !t.is_empty()),
}
}
pub fn history_ready(history: &[SharedMessage]) -> bool {
history.last().is_some_and(|m| tail_is_markable(m))
}
pub fn reply_prefix_tokens(history: &[SharedMessage]) -> u64 {
let Some(last) = history.last() else {
return 0;
};
let kept = match last["content"].as_array() {
Some(blocks) => Value::Array(
blocks
.iter()
.filter(|b| !matches!(b["type"].as_str(), Some("thinking" | "redacted_thinking")))
.cloned()
.collect(),
),
None => last["content"].clone(),
};
super::context::conservative_token_estimate(&kept.to_string())
}
impl Runtime {
pub fn cache_keepwarm_ineligibility(&self) -> Option<String> {
if !super::auth::model_is_anthropic(&self.model) {
return Some(format!(
"model {} is not on the native Anthropic transport",
self.model
));
}
None
}
pub fn cache_anchor(&self) -> Option<CacheAnchor> {
self.cache_anchor.get()
}
pub fn cache_ttl_downgraded(&self) -> bool {
self.ttl_downgrade_notified.load(Ordering::Relaxed)
}
pub fn cache_1h_honoured(&self) -> bool {
self.saw_1h_honored.load(Ordering::Relaxed)
}
pub async fn cache_keepwarm_probe(&self, history: &[SharedMessage]) -> Result<KeepWarmProbe> {
let cancel = CancellationToken::new();
let tap = PingTap::default();
let (value, _) = self
.keepwarm_request(history, &cancel, KeepWarmMode::Probe, &tap)
.await;
let value = value?;
let fingerprint = value["keepwarm_prefix"].as_u64().ok_or_else(|| {
RuntimeError::Session("keep-warm probe produced no fingerprint".into())
})?;
Ok(KeepWarmProbe {
fingerprint,
max_tokens: value["keepwarm_max_tokens"].as_u64().unwrap_or(0),
context_blocks: value["keepwarm_context_blocks"]
.as_str()
.map(str::to_string),
})
}
pub async fn cache_keepwarm_ping(
&self,
history: &[SharedMessage],
expect: u64,
cancel: &CancellationToken,
tap: &PingTap,
) -> PingReport {
let (value, notices) = self
.keepwarm_request(history, cancel, KeepWarmMode::Send { expect }, tap)
.await;
let result = match value {
Ok(value) => match value.get("keepwarm_prefix").and_then(Value::as_u64) {
Some(actual) => Ok(PingOutcome::PrefixChanged {
expected: expect,
actual,
}),
None => Ok(PingOutcome::Sent),
},
Err(e) => Err(classify_failure(&e, tap.status.load(Ordering::Relaxed))),
};
PingReport {
result,
usage: tap.usage(),
notices,
}
}
async fn keepwarm_request(
&self,
history: &[SharedMessage],
cancel: &CancellationToken,
mode: KeepWarmMode,
tap: &PingTap,
) -> (Result<Value>, Vec<String>) {
match self
.keepwarm_request_inner(history, cancel, mode, tap)
.await
{
Ok(r) => r,
Err(e) => (Err(e), Vec::new()),
}
}
async fn keepwarm_request_inner(
&self,
history: &[SharedMessage],
cancel: &CancellationToken,
mode: KeepWarmMode,
tap: &PingTap,
) -> Result<(Result<Value>, Vec<String>)> {
if let Some(reason) = self.cache_keepwarm_ineligibility() {
return Err(RuntimeError::Config(format!(
"cache keep-warm skipped: {reason}"
)));
}
if !history_ready(history) {
return Err(RuntimeError::Config(
"cache keep-warm skipped: history does not end with a finished assistant reply"
.into(),
));
}
let (durability_blocked, context_enabled) = {
let state = self
.continuation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
(state.durability_blocked, state.enabled())
};
if durability_blocked {
return Err(RuntimeError::Session(
"cache keep-warm skipped: context head save is unresolved".into(),
));
}
self.validate_request_preflight().await?;
let anthropic_execution_plan = self.authorized_anthropic_plan().await?;
if matches!(mode, KeepWarmMode::Send { .. }) {
self.refresh_if_needed().await?;
}
let system_prompt = self.effective_system_prompt().await;
let system_prompt = if context_enabled {
Some(format!(
"{}\n\n{}",
system_prompt.as_deref().unwrap_or_default(),
super::continuation::GUIDANCE
))
} else {
system_prompt
};
let session_injection = match self.session_id.as_deref() {
Some(id) => self.hook_bus.session_injection_for(id).await,
None => None,
};
let system_prompt = match session_injection {
Some(content) => Some(super::stream::wrap_extension_context(
system_prompt.as_deref().unwrap_or_default(),
&content,
)),
None => system_prompt,
};
let mut request: Vec<SharedMessage> = history.to_vec();
request.push(Arc::new(
json!({"role": "user", "content": KEEPWARM_PING_TEXT}),
));
let capped;
let request_messages: &[SharedMessage] = match super::stream::cap_history_image_bytes(
&request,
super::stream::HISTORY_IMAGE_BYTE_CAP,
) {
Some(c) => {
capped = c;
&capped
}
None => &request,
};
let tools_snapshot = self.tools.read().await.clone();
let (request_tools_schema, anthropic_tool_split) = if self.progressive_tool_disclosure {
let fresh = super::continuation::context_tool_set(
self.host_tool_session.clone(),
tools_snapshot.catalog(),
true,
context_enabled,
);
let retained = self
.retained_tool_set
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
let existing = retained.as_ref().map(|shared| {
shared
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
});
let (set, _, _) = super::stream::select_turn_start_tool_set(
existing.as_deref(),
fresh,
tools_snapshot.catalog(),
true,
);
drop(existing);
let schema = tools_snapshot.session_tools_schema(&set).schema;
let split = self.deferred_tool_loading.then(|| {
let referenced = HelperMethods::referenced_tool_ids(request_messages);
Arc::new(tools_snapshot.deferred_tool_split(&set, &referenced))
});
(Some(Arc::new(schema)), split)
} else {
(None, None)
};
let options = ApiOptions {
memory_backend: None,
use_1m_context: self.context_window_override == Some(1_000_000),
cache_ttl: self.cache_ttl,
ttl_downgrade_notified: self.ttl_downgrade_notified.clone(),
saw_1h_honored: self.saw_1h_honored.clone(),
credential_source: self.credential_source.clone(),
token_cache: self.token_cache.clone(),
anthropic_base_url: None,
anthropic_execution_plan,
codex_request_role: self.codex_request_role(),
trace: self.trace_ctx.clone(),
request_correlation: None,
suppress_stream_deltas: true,
telemetry: self.telemetry_writer.clone(),
tool_session_id: Some(self.host_tool_session.clone()),
session_tool_set: None,
request_tools_schema: request_tools_schema.clone(),
anthropic_tool_split,
usage_counters: None,
cache_anchor: Some(Arc::clone(&self.cache_anchor)),
keepwarm: Some(KeepWarmRequest {
mode,
status: Arc::clone(&tap.status),
sent: Arc::clone(&tap.sent),
connect_failed: Arc::clone(&tap.connect_failed),
}),
};
let (tx, mut rx) = mpsc::unbounded_channel();
let call = ApiMethods::call_api_stream_inner(
&self.auth,
&self.client,
&self.model,
&tools_snapshot,
&system_prompt,
self.thinking_budget,
self.reasoning_level(),
request_messages,
tx,
cancel,
0,
0,
&options,
self.telemetry_level,
);
tokio::pin!(call);
let mut notices = Vec::new();
let mut on_event = |event: StreamEvent| {
if let StreamEvent::Session(SessionEvent::Notice(n)) = &event {
notices.push(n.clone());
}
tap.usage
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.add(&event);
};
let mut result = loop {
tokio::select! {
r = &mut call => break r,
Some(event) = rx.recv() => on_event(event),
}
};
while let Ok(event) = rx.try_recv() {
on_event(event);
}
if let (Ok(value), KeepWarmMode::Probe) = (&mut result, mode) {
if context_enabled {
if let Some(why) = self.keepwarm_context_blocks(
history,
system_prompt.as_deref(),
request_tools_schema.as_deref().map(|s| s.as_slice()),
&tools_snapshot,
) {
value["keepwarm_context_blocks"] = Value::String(why);
}
}
}
Ok((result, notices))
}
fn keepwarm_context_blocks(
&self,
history: &[SharedMessage],
system_prompt: Option<&str>,
projected_schema: Option<&[Value]>,
tools_snapshot: &crate::ToolRegistry,
) -> Option<String> {
use agent_core::core::context_policy::{assess_context, ContextAction};
let mut next: Vec<SharedMessage> = history.to_vec();
next.push(Arc::new(json!({"role": "user", "content": ""})));
let capped =
super::stream::cap_history_image_bytes(&next, super::stream::HISTORY_IMAGE_BYTE_CAP);
let request_messages: &[SharedMessage] = capped.as_deref().unwrap_or(&next[..]);
let fallback_schema;
let schema = match projected_schema {
Some(s) => s,
None => {
fallback_schema = tools_snapshot.tools_schema();
fallback_schema.as_slice()
}
};
let (assessment, budget) = super::stream::pre_request_context_budget(
&self.model,
self.context_window(),
system_prompt,
schema,
request_messages,
self.thinking_budget as u64,
self.max_tool_output as u64,
);
let state = self
.continuation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let decision = assess_context(&state.config, &state.policy, budget);
match decision.action {
ContextAction::Rollover => Some(format!(
"the next turn would roll the context over ({} tokens used)",
assessment.used_tokens()
)),
ContextAction::HardStop => Some(format!(
"the next turn would hit the context hard limit ({} tokens used)",
assessment.used_tokens()
)),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::config::CacheTtl;
use crate::extensions::hooks::events::{HookEvent, HookKind, HookResult};
use crate::extensions::permissions::{Permission, PermissionSet};
use agent_core::reasoning::ReasoningLevel;
use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::post, Router};
use futures::StreamExt;
use std::sync::atomic::AtomicUsize;
use std::sync::Mutex;
const SSE_HI: &str = concat!(
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
"\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
"\"stop_sequence\":null,\"usage\":{\"input_tokens\":10,\"output_tokens\":0,",
"\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}}\n\n",
"data: {\"type\":\"content_block_start\",\"index\":0,",
"\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,",
"\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n",
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",",
"\"stop_sequence\":null},\"usage\":{\"input_tokens\":10,\"output_tokens\":1,",
"\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0}}\n\n",
"data: {\"type\":\"message_stop\"}\n\n",
);
const SSE_USAGE_THEN_ERROR: &str = concat!(
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_kw\",\"type\":\"message\",",
"\"role\":\"assistant\",\"content\":[],\"model\":\"claude-opus-4-7\",\"stop_reason\":null,",
"\"stop_sequence\":null,\"usage\":{\"input_tokens\":12,\"output_tokens\":0,",
"\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}}\n\n",
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":null,",
"\"stop_sequence\":null},\"usage\":{\"input_tokens\":12,\"output_tokens\":1,",
"\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":5000}}\n\n",
"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",",
"\"message\":\"synthetic\"}}\n\n",
);
#[derive(Clone)]
struct Mock {
bodies: Arc<Mutex<Vec<String>>>,
status: u16,
sse: &'static str,
}
async fn spawn_mock() -> (String, Mock) {
spawn_mock_full(200, SSE_HI).await
}
async fn spawn_mock_with(status: u16) -> (String, Mock) {
spawn_mock_full(status, SSE_HI).await
}
async fn spawn_mock_full(status: u16, sse: &'static str) -> (String, Mock) {
let mock = Mock {
bodies: Arc::default(),
status,
sse,
};
let app = Router::new()
.route(
"/v1/messages",
post(|State(m): State<Mock>, body: String| async move {
m.bodies.lock().unwrap().push(body);
if m.status == 200 {
(
StatusCode::OK,
[("content-type", "text/event-stream")],
m.sse.to_string(),
)
.into_response()
} else {
(
StatusCode::from_u16(m.status).unwrap(),
[("content-type", "application/json")],
r#"{"type":"error","error":{"type":"api_error","message":"synthetic"}}"#
.to_string(),
)
.into_response()
}
}),
)
.layer(axum::extract::DefaultBodyLimit::max(64 * 1024 * 1024))
.with_state(mock.clone());
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(format!("http://{addr}"), mock)
}
struct CountingInject(Arc<AtomicUsize>);
#[async_trait::async_trait]
impl crate::extensions::runtime::ExtensionHandler for CountingInject {
fn id(&self) -> &str {
"keepwarm-counting-inject"
}
async fn handle(&self, _event: &HookEvent) -> HookResult {
self.0.fetch_add(1, Ordering::SeqCst);
HookResult::Inject {
content: "PER-TURN-INJECTED-CONTEXT".into(),
}
}
async fn shutdown(&self) {}
}
async fn real_turn(rt: &Runtime, messages: Vec<SharedMessage>) -> Vec<SharedMessage> {
let mut stream = rt
.run_stream_with_messages(messages, CancellationToken::new(), None, None, false)
.await;
let mut history = None;
while let Some(event) = stream.next().await {
match event {
StreamEvent::Session(SessionEvent::MessageHistory(h)) => history = Some(h),
StreamEvent::Session(SessionEvent::Error(e)) => panic!("turn failed: {e:?}"),
_ => {}
}
}
history.expect("turn published its history")
}
fn user(text: &str) -> SharedMessage {
Arc::new(json!({"role": "user", "content": text}))
}
fn canon(v: &Value) -> String {
serde_json::to_string(v).unwrap()
}
fn strip_cache_control(mut msg: Value) -> Value {
if let Some(blocks) = msg["content"].as_array_mut() {
for b in blocks {
if let Some(obj) = b.as_object_mut() {
obj.remove("cache_control");
}
}
}
msg
}
fn marker_count(v: &Value) -> usize {
match v {
Value::Object(map) => {
usize::from(map.contains_key("cache_control"))
+ map.values().map(marker_count).sum::<usize>()
}
Value::Array(items) => items.iter().map(marker_count).sum(),
_ => 0,
}
}
struct Observed {
ping: Value,
real: Value,
raw_ping: String,
history_len: usize,
usage: KeepWarmUsage,
hook_calls: (usize, usize, usize),
}
async fn ping_then_next_turn(rt: &Runtime) -> Observed {
let (url, mock) = spawn_mock().await;
let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
let calls = Arc::new(AtomicUsize::new(0));
let mut perms = PermissionSet::new();
perms.grant(Permission::LlmContent);
rt.hook_bus()
.subscribe(
HookKind::BeforeMessage,
Arc::new(CountingInject(Arc::clone(&calls))),
None,
None,
perms,
)
.await
.unwrap();
assert_eq!(rt.cache_anchor(), None, "no request has run yet");
let history = real_turn(rt, vec![user("hello")]).await;
assert_eq!(history.last().unwrap()["role"], "assistant");
let turn_anchor = rt
.cache_anchor()
.expect("an accepted turn request stamps the anchor");
let before_ping = calls.load(Ordering::SeqCst);
let snapshot = history.clone();
let fingerprint = rt
.cache_keepwarm_probe(&history)
.await
.expect("probe")
.fingerprint;
assert_eq!(
mock.bodies.lock().unwrap().len(),
1,
"a probe sends nothing"
);
assert_eq!(
rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
fingerprint,
"the fingerprint is deterministic"
);
let tap = PingTap::default();
let report = rt
.cache_keepwarm_ping(&history, fingerprint, &CancellationToken::new(), &tap)
.await;
assert_eq!(report.result, Ok(PingOutcome::Sent), "{report:?}");
assert!(tap.was_sent(), "the tap saw the request leave");
assert_eq!(tap.usage(), report.usage, "the tap saw the usage");
let usage = report.usage;
assert!(
usage.output_tokens >= 1,
"usage comes back from the private channel"
);
let ping_anchor = rt.cache_anchor().expect("anchor");
assert!(
ping_anchor.mono >= turn_anchor.mono,
"the ping moved the anchor"
);
let after_ping = calls.load(Ordering::SeqCst);
assert_eq!(history, snapshot, "the ping must not mutate history");
let mut next = history.clone();
next.push(user("next question"));
real_turn(rt, next).await;
let after_next = calls.load(Ordering::SeqCst);
let bodies = mock.bodies.lock().unwrap().clone();
assert_eq!(bodies.len(), 3, "previous turn, ping, next turn");
Observed {
ping: serde_json::from_str(&bodies[1]).unwrap(),
real: serde_json::from_str(&bodies[2]).unwrap(),
raw_ping: bodies[1].clone(),
history_len: history.len(),
usage,
hook_calls: (before_ping, after_ping, after_next),
}
}
fn assert_prefix_identical(
ping: &Value,
real: &Value,
n: usize,
history_marker: &Value,
ping_max_tokens: u64,
) {
for key in [
"model",
"tools",
"system",
"thinking",
"output_config",
"stream",
] {
assert_eq!(
canon(&ping[key]),
canon(&real[key]),
"`{key}` differs between ping and the real next turn"
);
}
let pm = ping["messages"].as_array().unwrap();
let rm = real["messages"].as_array().unwrap();
assert_eq!(pm.len(), n + 1, "history + one trailing ping message");
assert!(rm.len() > n);
for i in 0..n - 1 {
assert_eq!(canon(&pm[i]), canon(&rm[i]), "message {i} differs");
}
let last = &pm[n - 1];
let marker = &last["content"].as_array().unwrap().last().unwrap()["cache_control"];
assert_eq!(marker, history_marker);
assert_eq!(
canon(&strip_cache_control(last.clone())),
canon(&rm[n - 1]),
"last history message differs beyond the cache marker"
);
assert_eq!(marker_count(&pm[n]), 0, "ping message must be unmarked");
assert_eq!(pm[n]["role"], "user");
assert_eq!(pm[n]["content"], KEEPWARM_PING_TEXT);
assert_eq!(ping["max_tokens"], json!(ping_max_tokens));
assert!(real["max_tokens"].as_u64().unwrap() > ping_max_tokens);
assert!(
marker_count(ping) <= 4,
"Anthropic allows at most 4 breakpoints"
);
}
fn marker_1h() -> Value {
json!({"type": "ephemeral", "ttl": "1h"})
}
fn marker_5m() -> Value {
json!({"type": "ephemeral"})
}
fn opus(ttl: CacheTtl) -> Runtime {
let mut rt = Runtime::new_headless();
rt.set_model("claude-opus-4-7".to_string());
assert_eq!(rt.model(), "claude-opus-4-7");
rt.set_reasoning_level_explicit(ReasoningLevel::High);
rt.set_cache_ttl(ttl);
rt.set_system_prompt("You are a keep-warm test.".into());
assert_eq!(rt.cache_keepwarm_ineligibility(), None);
rt
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_adaptive_thinking() {
let _base = crate::test_env::BaseDirGuard::new();
let rt = opus(CacheTtl::OneHour);
let o = ping_then_next_turn(&rt).await;
assert!(
!o.ping["thinking"].is_null(),
"thinking must stay on: {}",
o.ping
);
assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
let (before, after_ping, after_next) = o.hook_calls;
assert!(
before >= 1,
"control: the previous turn fired before_message"
);
assert_eq!(after_ping, before, "the ping fired a hook");
assert!(
after_next > after_ping,
"control: the next turn fired before_message"
);
assert!(!o.raw_ping.contains("PER-TURN-INJECTED-CONTEXT"));
assert!(!o.raw_ping.contains("Extension context"));
assert!(canon(&o.real).contains("PER-TURN-INJECTED-CONTEXT"));
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_5m_ttl() {
let _base = crate::test_env::BaseDirGuard::new();
let rt = opus(CacheTtl::FiveMinutes);
let o = ping_then_next_turn(&rt).await;
assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_hybrid_ttl() {
let _base = crate::test_env::BaseDirGuard::new();
let rt = opus(CacheTtl::Hybrid);
let o = ping_then_next_turn(&rt).await;
assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_5m(), 1);
assert!(
canon(&o.ping["system"]).contains("\"ttl\":\"1h\""),
"hybrid keeps the 1h stable-prefix marker: {}",
o.ping["system"]
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_thinking_off() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = Runtime::new_headless();
rt.set_model("claude-sonnet-4-6".to_string());
rt.set_reasoning_level_explicit(ReasoningLevel::Off);
rt.set_cache_ttl(CacheTtl::OneHour);
assert_eq!(rt.cache_keepwarm_ineligibility(), None);
let o = ping_then_next_turn(&rt).await;
assert!(o.ping.get("thinking").is_none(), "thinking off is omitted");
assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn budgeted_thinking_ping_keeps_budget_and_is_costed() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = Runtime::new_headless();
rt.set_model("claude-sonnet-4-6".to_string());
rt.set_reasoning_level_explicit(ReasoningLevel::High);
rt.set_cache_ttl(CacheTtl::FiveMinutes);
assert_eq!(
rt.cache_keepwarm_ineligibility(),
None,
"budgeted thinking is eligible"
);
let o = ping_then_next_turn(&rt).await;
assert_eq!(o.ping["thinking"]["type"], "enabled", "{}", o.ping);
let budget = o.ping["thinking"]["budget_tokens"]
.as_u64()
.expect("budget");
assert_eq!(canon(&o.ping["thinking"]), canon(&o.real["thinking"]));
assert_prefix_identical(
&o.ping,
&o.real,
o.history_len,
&marker_5m(),
budget + KEEPWARM_OUTPUT_TOKENS,
);
assert!(
o.usage.cost_usd("claude-sonnet-4-6") > 0.0,
"the ping is costed"
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn changed_prefix_is_refused_before_the_network() {
let _base = crate::test_env::BaseDirGuard::new();
let (url, mock) = spawn_mock().await;
let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
let mut rt = opus(CacheTtl::FiveMinutes);
let history = real_turn(&rt, vec![user("hello")]).await;
let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
let sent = mock.bodies.lock().unwrap().len();
let mut longer = history.clone();
longer.push(user("more"));
longer.push(Arc::new(
json!({"role": "assistant", "content": [{"type": "text", "text": "x"}]}),
));
let changed = |r: PingReport| matches!(r.result, Ok(PingOutcome::PrefixChanged { .. }));
let tap = PingTap::default();
assert!(changed(
rt.cache_keepwarm_ping(&longer, fp, &CancellationToken::new(), &tap)
.await
));
rt.set_system_prompt("A different system prompt.".into());
assert_ne!(
rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
fp
);
assert!(changed(
rt.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
.await
));
rt.set_system_prompt("You are a keep-warm test.".into());
assert_eq!(
rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
fp,
"restored"
);
rt.set_reasoning_level_explicit(ReasoningLevel::Low);
assert_ne!(
rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
fp,
"effort is in the key"
);
rt.set_reasoning_level_explicit(ReasoningLevel::High);
rt.set_model("claude-opus-4-6".to_string());
assert_eq!(rt.model(), "claude-opus-4-6");
assert_ne!(
rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint,
fp,
"model is in the key"
);
assert_eq!(
mock.bodies.lock().unwrap().len(),
sent,
"no refused ping reached the provider"
);
assert!(!tap.was_sent(), "a refused ping never leaves");
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn http_4xx_is_fatal_and_5xx_is_transient() {
let _base = crate::test_env::BaseDirGuard::new();
let history = vec![
user("hello"),
Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
];
for (status, fatal) in [(400u16, true), (404, true), (500, false), (429, false)] {
let (url, mock) = spawn_mock_with(status).await;
let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
let rt = opus(CacheTtl::FiveMinutes);
let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
let err = rt
.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &PingTap::default())
.await
.result
.expect_err("non-2xx fails");
assert_eq!(
matches!(err, PingFailure::Fatal(_)),
fatal,
"HTTP {status}: {err:?}"
);
assert_eq!(
mock.bodies.lock().unwrap().len(),
1,
"HTTP {status}: no retry"
);
assert_eq!(
rt.cache_anchor(),
None,
"a rejected request refreshes nothing"
);
}
}
async fn assert_identity_1h(rt: &Runtime) -> Observed {
let o = ping_then_next_turn(rt).await;
assert_prefix_identical(&o.ping, &o.real, o.history_len, &marker_1h(), 1);
o
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_progressive_disclosure() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = opus(CacheTtl::OneHour);
rt.progressive_tool_disclosure = true;
assert_identity_1h(&rt).await;
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_deferred_tools() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = opus(CacheTtl::OneHour);
rt.progressive_tool_disclosure = true;
rt.deferred_tool_loading = true;
assert_identity_1h(&rt).await;
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_context_management_guidance() {
let _base = crate::test_env::BaseDirGuard::new();
let rt = opus(CacheTtl::OneHour);
rt.continuation.lock().unwrap().config.mode =
crate::core::config::ContextManagementMode::Auto;
let o = assert_identity_1h(&rt).await;
assert!(
canon(&o.ping["system"]).contains("Work normally."),
"guidance present: {}",
o.ping["system"]
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_session_injection() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = opus(CacheTtl::OneHour);
rt.set_session_id(Some("kw-session".into()));
rt.hook_bus()
.set_session_injection_for("kw-session", "SESSION-SCOPED-INJECTION".into())
.await;
let o = assert_identity_1h(&rt).await;
assert!(canon(&o.ping["system"]).contains("SESSION-SCOPED-INJECTION"));
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_prefix_matches_next_turn_with_oauth_system_blocks() {
let _base = crate::test_env::BaseDirGuard::new();
let mut rt = opus(CacheTtl::OneHour);
rt.auth = Arc::new(tokio::sync::RwLock::new(super::super::types::AuthState {
auth_token: "oauth-test-token".to_string(),
auth_type: "oauth".to_string(),
refresh_token: None,
token_expires: Some(crate::epoch_millis() + 3_600_000),
}));
let o = assert_identity_1h(&rt).await;
assert!(
canon(&o.ping["system"]).contains("You are Claude Code"),
"{}",
o.ping["system"]
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn probe_reports_a_context_rollover_ahead() {
let _base = crate::test_env::BaseDirGuard::new();
let (url, _mock) = spawn_mock().await;
let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
let rt = opus(CacheTtl::OneHour);
rt.continuation.lock().unwrap().config.mode =
crate::core::config::ContextManagementMode::Auto;
let history = real_turn(&rt, vec![user("hello")]).await;
let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
assert_eq!(probe.context_blocks, None, "plenty of room");
let before = rt.continuation.lock().unwrap().policy;
rt.continuation.lock().unwrap().config.reserve_tokens = 10_000_000;
let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
assert!(probe.context_blocks.is_some(), "{probe:?}");
assert_eq!(
rt.continuation.lock().unwrap().policy,
before,
"the probe never commits policy state"
);
rt.continuation.lock().unwrap().config.mode =
crate::core::config::ContextManagementMode::Off;
assert_eq!(
rt.cache_keepwarm_probe(&history)
.await
.unwrap()
.context_blocks,
None
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn probe_reports_max_tokens() {
let _base = crate::test_env::BaseDirGuard::new();
let history = vec![
user("hello"),
Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
];
let rt = opus(CacheTtl::OneHour);
assert_eq!(
rt.cache_keepwarm_probe(&history).await.unwrap().max_tokens,
1
);
let mut rt = Runtime::new_headless();
rt.set_model("claude-sonnet-4-6".to_string());
rt.set_reasoning_level_explicit(ReasoningLevel::XHigh);
let probe = rt.cache_keepwarm_probe(&history).await.unwrap();
let budget = u64::from(rt.thinking_budget);
assert!(budget >= 16_384, "budgeted thinking at xhigh: {budget}");
assert_eq!(
probe.max_tokens,
budget + KEEPWARM_OUTPUT_TOKENS,
"{probe:?}"
);
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn ping_usage_is_returned_on_error() {
let _base = crate::test_env::BaseDirGuard::new();
let (url, mock) = spawn_mock_full(200, SSE_USAGE_THEN_ERROR).await;
let _url = crate::test_env::EnvVarGuard::set("SYNAPS_ANTHROPIC_BASE_URL", &url);
let history = vec![
user("hello"),
Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
];
let rt = opus(CacheTtl::OneHour);
let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
let tap = PingTap::default();
let report = rt
.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
.await;
assert!(
matches!(report.result, Err(PingFailure::Transient(_))),
"{report:?}"
);
assert_eq!(report.usage.cache_read_tokens, 5_000, "{report:?}");
assert!(report.usage.cost_usd("claude-opus-4-7") > 0.0);
assert_eq!(tap.usage(), report.usage);
assert!(tap.was_sent());
assert_eq!(mock.bodies.lock().unwrap().len(), 1, "no transport retry");
}
#[test]
fn ping_shares_the_session_ttl_latch() {
let rt = opus(CacheTtl::OneHour);
assert!(!rt.cache_ttl_downgraded());
let mut ping_rt = rt.clone();
ping_rt.share_ttl_latches(&rt);
ping_rt
.ttl_downgrade_notified
.store(true, Ordering::Relaxed);
assert!(rt.cache_ttl_downgraded(), "the latch is shared");
}
#[test]
fn failure_classification() {
let api = || RuntimeError::ApiStatus("x".into());
assert_eq!(
classify_failure(&RuntimeError::Canceled, 0),
PingFailure::Canceled
);
assert!(matches!(
classify_failure(&api(), 400),
PingFailure::Fatal(_)
));
assert!(matches!(
classify_failure(&api(), 401),
PingFailure::Fatal(_)
));
assert!(matches!(
classify_failure(&api(), 429),
PingFailure::Transient(_)
));
assert!(matches!(
classify_failure(&api(), 529),
PingFailure::Transient(_)
));
assert!(matches!(
classify_failure(&api(), 0),
PingFailure::Transient(_)
));
assert!(matches!(
classify_failure(&api(), 200),
PingFailure::Transient(_)
));
assert!(matches!(
classify_failure(&RuntimeError::Config("x".into()), 0),
PingFailure::Fatal(_)
));
}
#[test]
fn every_ttl_and_thinking_mode_is_eligible() {
for ttl in [CacheTtl::FiveMinutes, CacheTtl::OneHour, CacheTtl::Hybrid] {
for (model, level) in [
("claude-opus-4-7", ReasoningLevel::High),
("claude-sonnet-4-6", ReasoningLevel::High),
("claude-sonnet-4-6", ReasoningLevel::Off),
] {
let mut rt = Runtime::new_headless();
rt.set_model(model.to_string());
rt.set_reasoning_level_explicit(level);
rt.set_cache_ttl(ttl);
assert_eq!(
rt.cache_keepwarm_ineligibility(),
None,
"{model} {level:?} {ttl:?}"
);
}
}
}
#[test]
fn non_anthropic_models_are_ineligible() {
let mut rt = Runtime::new_headless();
rt.set_model("groq/llama-3.3-70b-versatile".to_string());
if rt.model() == "groq/llama-3.3-70b-versatile" {
assert!(rt.cache_keepwarm_ineligibility().is_some());
}
}
#[tokio::test]
async fn ineligible_ping_sends_nothing() {
let mut rt = Runtime::new_headless();
rt.set_model("groq/llama-3.3-70b-versatile".to_string());
if rt.model() != "groq/llama-3.3-70b-versatile" {
return;
}
let history = vec![
user("hello"),
Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
];
let tap = PingTap::default();
let err = rt
.cache_keepwarm_ping(&history, 0, &CancellationToken::new(), &tap)
.await
.result
.expect_err("non-Anthropic must refuse");
assert!(!tap.was_sent());
assert!(
matches!(err, PingFailure::Fatal(ref m) if m.contains("keep-warm skipped")),
"{err:?}"
);
}
#[test]
fn anchor_cell_only_moves_forward() {
let cell = CacheAnchorCell::default();
assert_eq!(cell.get(), None);
let a = CacheAnchor::now();
let later = CacheAnchor {
mono: a.mono + std::time::Duration::from_secs(5),
wall: a.wall + std::time::Duration::from_secs(5),
};
cell.stamp(later);
cell.stamp(a);
assert_eq!(
cell.get(),
Some(later),
"a late-arriving older stamp never rewinds"
);
}
#[test]
fn only_a_finished_assistant_text_reply_is_markable() {
assert!(tail_is_markable(
&json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}, {"type": "text", "text": "ok"}]})
));
assert!(!tail_is_markable(&json!({"role": "user", "content": "q"})));
assert!(!tail_is_markable(
&json!({"role": "assistant", "content": [{"type": "tool_use", "id": "t", "name": "x", "input": {}}]})
));
assert!(!tail_is_markable(
&json!({"role": "assistant", "content": [{"type": "thinking", "thinking": "t"}]})
));
}
#[test]
fn usage_is_summed_and_priced() {
let mut u = KeepWarmUsage::default();
u.add(&StreamEvent::Session(SessionEvent::Usage {
input_tokens: 20,
output_tokens: 16,
cache_read_input_tokens: 500_000,
cache_creation_input_tokens: 0,
cache_creation_5m: None,
cache_creation_1h: Some(0),
model: Some("claude-opus-4-7".into()),
}));
assert_eq!(u.cache_read_tokens, 500_000);
assert_eq!(u.output_tokens, 16);
let cost = u.cost_usd("claude-opus-4-7");
assert!((cost - 0.2505).abs() < 0.001, "{cost}");
}
#[test]
fn booking_rule_table() {
use PingBooking::*;
let table = [
((true, true, false, 200), Measured),
((true, false, false, 0), Measured),
((true, true, false, 529), Measured),
((false, true, false, 200), WorstCase), ((false, true, false, 0), WorstCase), ((false, true, false, 500), Nothing), ((false, true, false, 429), Nothing),
((false, true, false, 529), Nothing),
((false, true, true, 0), Nothing), ((false, true, true, 503), Nothing), ((false, false, false, 0), Nothing), ((false, false, true, 0), Nothing),
];
for ((usage, sent, connect, status), want) in table {
assert_eq!(
ping_booking(usage, sent, connect, status),
want,
"usage={usage} sent={sent} connect_failed={connect} status={status}"
);
}
let tap = PingTap::default();
assert_eq!(tap.bill("claude-opus-4-7", 0.5).cost_usd, 0.0);
tap.sent.store(true, Ordering::Relaxed);
let bill = tap.bill("claude-opus-4-7", 0.5);
assert_eq!((bill.booking, bill.cost_usd), (WorstCase, 0.5));
tap.usage.lock().unwrap().cache_read_tokens = 1_000_000;
let bill = tap.bill("claude-opus-4-7", 0.5);
assert_eq!(bill.booking, Measured);
assert!((bill.cost_usd - 0.5).abs() < 1e-9, "1M reads at $0.50/MTok");
}
fn two_message_history() -> Vec<SharedMessage> {
vec![
user("hello"),
Arc::new(json!({"role": "assistant", "content": [{"type": "text", "text": "hi"}]})),
]
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn a_refused_connection_is_never_sent_and_books_nothing() {
let _base = crate::test_env::BaseDirGuard::new();
let addr = {
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
l.local_addr().unwrap()
}; let _url = crate::test_env::EnvVarGuard::set(
"SYNAPS_ANTHROPIC_BASE_URL",
&format!("http://{addr}"),
);
let history = two_message_history();
let rt = opus(CacheTtl::OneHour);
let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
let tap = PingTap::default();
let report = rt
.cache_keepwarm_ping(&history, fp, &CancellationToken::new(), &tap)
.await;
assert!(report.result.is_err(), "{report:?}");
assert!(!tap.was_sent(), "never dispatched");
assert!(tap.connect_failed.load(Ordering::Relaxed));
let bill = tap.bill("claude-opus-4-7", 0.75);
assert_eq!((bill.booking, bill.cost_usd), (PingBooking::Nothing, 0.0));
}
#[tokio::test]
#[serial_test::serial(synaps_base_dir)]
async fn a_dispatched_ping_without_usage_books_the_worst_case() {
use tokio::io::AsyncReadExt;
let _base = crate::test_env::BaseDirGuard::new();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut sock, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 64 * 1024];
while sock.read(&mut buf).await.is_ok_and(|n| n > 0) {}
});
let _url = crate::test_env::EnvVarGuard::set(
"SYNAPS_ANTHROPIC_BASE_URL",
&format!("http://{addr}"),
);
let history = two_message_history();
let rt = opus(CacheTtl::OneHour);
let fp = rt.cache_keepwarm_probe(&history).await.unwrap().fingerprint;
let tap = PingTap::default();
let cancel = CancellationToken::new();
let ping = {
let (rt, history, tap, cancel) =
(rt.clone(), history.clone(), tap.clone(), cancel.clone());
tokio::spawn(async move { rt.cache_keepwarm_ping(&history, fp, &cancel, &tap).await })
};
tokio::time::timeout(std::time::Duration::from_secs(10), async {
while !tap.was_sent() {
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
})
.await
.expect("the body was pulled over the established connection");
cancel.cancel();
let report = ping.await.unwrap();
assert_eq!(report.result, Err(PingFailure::Canceled));
assert!(!tap.connect_failed.load(Ordering::Relaxed));
let bill = tap.bill("claude-opus-4-7", 0.75);
assert_eq!(
(bill.booking, bill.cost_usd),
(PingBooking::WorstCase, 0.75)
);
}
#[test]
fn reply_prefix_tokens_skip_thinking() {
let text = "a".repeat(250); let with_thinking = vec![Arc::new(json!({"role": "assistant", "content": [
{"type": "thinking", "thinking": "x".repeat(50_000), "signature": "s"},
{"type": "text", "text": text},
]}))];
let plain = vec![Arc::new(json!({"role": "assistant", "content": [
{"type": "text", "text": text},
]}))];
let t = reply_prefix_tokens(&with_thinking);
assert_eq!(t, reply_prefix_tokens(&plain), "thinking is not counted");
assert!((100..200).contains(&t), "{t}");
assert_eq!(reply_prefix_tokens(&[]), 0);
}
}