diff --git a/crates/tinyagents-harness/src/agent_loop/model_call.rs b/crates/tinyagents-harness/src/agent_loop/model_call.rs index 6b796940..c37e6e53 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_call.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_call.rs @@ -137,13 +137,42 @@ impl AgentHarness { } // Provider prompt-cache breakpoints are injected *after* the key is - // derived (they mutate `provider_options`, which the key covers) and - // only when the policy asks for prefix protection, so the common path - // never pays for a request clone. + // derived (they mutate `provider_options`, which the key covers), so + // the common path — no protection, no declared prefix — never pays for + // a request clone. + // + // The *effective* policy is stamped onto the clone. A request that + // carries no `cache_policy` of its own inherits the harness-level + // `RunPolicy::cache`, but that inheritance used to stop here: both + // `apply_prompt_cache_breakpoints` and the provider adapters read + // `request.cache_policy`, so a host that set `protect_prompt_prefix` on + // its run policy — the documented way — got no `prompt_cache_key` and + // no `cache_control` markers on the wire, while the layout guard kept + // reporting the prefix as protected. Stamping also runs when the run + // policy does *not* protect but a middleware declared cacheable + // segments: an adapter treats declared segments alone as the opt-in, + // and the run policy must be able to veto that. + let declares_prefix = request + .cache_segments + .iter() + .any(|segment| segment.cacheable); + let needs_stamp = + request.cache_policy.is_none() && (policy.protect_prompt_prefix || declares_prefix); let mut breakpointed; - let effective_request = if policy.protect_prompt_prefix { + let effective_request = if policy.protect_prompt_prefix || needs_stamp { breakpointed = request.clone(); - apply_prompt_cache_breakpoints(&mut breakpointed); + if needs_stamp { + breakpointed.cache_policy = Some(policy.clone()); + } + let injected = + policy.protect_prompt_prefix && apply_prompt_cache_breakpoints(&mut breakpointed); + tinyagents_tracing::debug!( + call_id = %call_id.as_str(), + protect_prompt_prefix = policy.protect_prompt_prefix, + prompt_cache_key_injected = injected, + cacheable_segments = breakpointed.cacheable_prefix_ids().len(), + "[cache] effective cache policy applied to the outgoing request" + ); &breakpointed } else { request diff --git a/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs b/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs index aa85316c..70a04650 100644 --- a/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs +++ b/crates/tinyagents-integration-tests/tests/wave2_cache_layout.rs @@ -307,3 +307,139 @@ async fn the_guard_still_reports_an_invalidation_inside_one_run() { "rewriting a stable segment's text inside one run is still an invalidation" ); } + +/// The run-policy path: a host that sets `protect_prompt_prefix` on +/// [`RunPolicy::cache`] — rather than stamping every request — must still get +/// a `prompt_cache_key` on the wire, and the provider adapter must be able to +/// see the policy it is acting under. +/// +/// Before the fix both readers consulted `request.cache_policy` alone, which is +/// `None` on every request the loop builds itself, so the harness-level flag +/// produced no breakpoint anywhere: the layout guard reported the prefix as +/// protected while the provider was told nothing. +mod run_policy_breakpoints { + use std::sync::{Arc, Mutex}; + + use async_trait::async_trait; + use tinyagents_harness::cache::PROMPT_CACHE_KEY_OPTION; + use tinyagents_harness::context::RunContext; + use tinyagents_harness::middleware::Middleware; + use tinyagents_harness::runtime::{AgentHarness, RunPolicy}; + use tinyinference::cache::CachePolicy; + use tinyinference::message::Message; + use tinyinference::model::{ + ChatModel, ModelRequest, ModelResponse, PromptSegment, SegmentRole, + }; + + struct RecordingModel { + seen: Mutex>, + } + + #[async_trait] + impl ChatModel<()> for RecordingModel { + async fn invoke( + &self, + _state: &(), + request: ModelRequest, + ) -> tinyinference::Result { + self.seen.lock().expect("poisoned").push(request); + Ok(ModelResponse::assistant("ok")) + } + } + + /// Declares the system prompt as the cacheable prefix, the way a host + /// that assembles its own messages (instead of using `PromptBuilder`) does. + struct DeclareSystemPrefix; + + #[async_trait] + impl Middleware<()> for DeclareSystemPrefix { + fn name(&self) -> &str { + "declare_system_prefix" + } + + async fn before_model( + &self, + _ctx: &mut RunContext<()>, + _state: &(), + request: &mut ModelRequest, + ) -> tinyagents_harness::Result<()> { + request.cache_segments = vec![PromptSegment { + id: "system".into(), + role: SegmentRole::System, + cacheable: true, + }]; + request.prompt_fingerprint = Some("fp".into()); + Ok(()) + } + } + + async fn run_with(policy: CachePolicy) -> ModelRequest { + let model = Arc::new(RecordingModel { + seen: Mutex::new(Vec::new()), + }); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("rec", model.clone()); + harness.push_middleware(Arc::new(DeclareSystemPrefix)); + harness.with_policy(RunPolicy { + cache: policy, + ..RunPolicy::default() + }); + harness + .invoke_default( + &(), + vec![Message::system("stable rules"), Message::user("go")], + ) + .await + .expect("run succeeds"); + let seen = model.seen.lock().expect("poisoned"); + assert_eq!(seen.len(), 1); + seen[0].clone() + } + + #[tokio::test] + async fn a_protecting_run_policy_reaches_the_provider_as_a_breakpoint() { + let request = run_with(CachePolicy { + protect_prompt_prefix: true, + ..CachePolicy::default() + }) + .await; + assert!( + request.wants_prompt_cache_breakpoints(), + "the provider adapter must see the effective policy, not None" + ); + assert!( + request + .cache_policy + .as_ref() + .is_some_and(|p| p.protect_prompt_prefix), + "the effective policy is stamped onto the outgoing request" + ); + let key = request.provider_options[PROMPT_CACHE_KEY_OPTION] + .as_str() + .expect("a prompt_cache_key was injected"); + assert!(key.starts_with("tap-"), "unexpected key shape: {key}"); + } + + /// The run policy is authoritative in both directions: with protection + /// off, a middleware-declared prefix must not turn into breakpoints on + /// the wire, because an adapter treats declared segments alone as the + /// opt-in. The stamped `protect_prompt_prefix: false` is what vetoes it. + #[tokio::test] + async fn an_unprotecting_run_policy_vetoes_declared_segments() { + let request = run_with(CachePolicy::default()).await; + assert!( + request + .cache_policy + .as_ref() + .is_some_and(|p| !p.protect_prompt_prefix), + "the unprotecting policy is stamped so the adapter can see the veto" + ); + assert!( + request + .provider_options + .get(PROMPT_CACHE_KEY_OPTION) + .is_none() + ); + assert!(!request.wants_prompt_cache_breakpoints()); + } +} diff --git a/docs/modules/harness/cache.md b/docs/modules/harness/cache.md index 06956d23..81aa8e19 100644 --- a/docs/modules/harness/cache.md +++ b/docs/modules/harness/cache.md @@ -125,6 +125,28 @@ the key and there is no minimum-prefix threshold. Unsafe or side-effecting tool calls should not be cached by default. +### Where `protect_prompt_prefix` is read + +The effective policy is the request's own `cache_policy` when present, +otherwise the harness-level `RunPolicy::cache`. The agent loop resolves that +once per model call and **stamps it onto the outgoing request** (a clone — +the original is what the response-cache key was derived from) before the +provider adapter sees it, whenever protection is on or a middleware declared +cacheable segments. Two readers depend on that stamp: + +- `apply_prompt_cache_breakpoints` injects the `prompt_cache_key` routing + hint into `provider_options`; +- provider adapters decide whether to emit explicit `cache_control` markers + via `ModelRequest::wants_prompt_cache_breakpoints`, where declared + cacheable segments are the opt-in and a stamped `protect_prompt_prefix: + false` is the veto. + +Both used to read `request.cache_policy` alone, which the loop never set, so a +host protecting the prefix on its run policy — the documented way — produced +no breakpoint anywhere while the layout guard reported the prefix as +protected. `wave2_cache_layout::run_policy_breakpoints` pins the fix in both +directions. + ### Key composition The key is a two-part composition, never the prompt alone: diff --git a/vendor/tinyinference b/vendor/tinyinference index d2e377ae..8895ece9 160000 --- a/vendor/tinyinference +++ b/vendor/tinyinference @@ -1 +1 @@ -Subproject commit d2e377ae19a6d27785927f4564526653ae8a8906 +Subproject commit 8895ece978b42bf77895cb4ad649877cf2e6e10d