fix(cache): stable layout under text dialects; RunPolicy::host_renders_tool_catalogue

crates/tinyagents-harness/src/agent_loop/dialect.rs

@@ -12,8 +12,10 @@
 
 use std::sync::Arc;
 
-use tinyinference_llm::message::ContentBlock;
-use tinyinference_llm::model::{ModelRequest, ModelResponse, ToolChoice};
+use tinyinference_llm::message::{ContentBlock, Message};
+use tinyinference_llm::model::{
+    ModelRequest, ModelResponse, PromptSegment, SegmentRole, ToolChoice,
+};
 use tinyinference_llm::tool::{ToolCall, ToolSchema};
 use tinytools_agent::dialect::{CodeDialect, CodeStyle, PFormatDialect};
 use tinytools_agent::types::{ParseOptions, ParsedToolCall};
@@ -101,7 +103,24 @@ impl RunDialect {
     /// is folded into forms a prompt-guided model can read, the protocol block
     /// and catalogue go into the system prompt, and no schema goes on the
     /// wire. A no-op for [`Self::Native`] or when no tools are offered.
-    pub(super) fn apply_to_request(&self, request: &mut ModelRequest) {
+    ///
+    /// With `host_renders_catalogue` the schemas still leave the wire (the
+    /// registry built from them before this call is what parses the answer),
+    /// and nothing from the run's ordinary catalogue is appended: the host's
+    /// own prompt already carries the protocol block and the catalogue for
+    /// this dialect. `synthesized` is the exception — tool schemas minted
+    /// *this turn*, after the host's static prompt was already composed (the
+    /// structured-output fallback tool `StructuredStrategy::ToolCall` /
+    /// `ToolCallUnion` push onto `request.tools`). The host cannot have
+    /// rendered a schema it did not know about yet, so their catalogue
+    /// entries are appended here even in the host-rendered case, or the
+    /// model never learns the shape it is being forced to call.
+    pub(super) fn apply_to_request(
+        &self,
+        request: &mut ModelRequest,
+        host_renders_catalogue: bool,
+        synthesized: &[ToolSchema],
+    ) {
         if !self.is_text() || request.tools.is_empty() || request.tool_choice == ToolChoice::None {
             return;
         }
@@ -110,6 +129,34 @@ impl RunDialect {
         let tools = std::mem::take(&mut request.tools);
         let messages = prompt_tools::coalesce_tool_results(&request.messages);
         let messages = prompt_tools::ensure_resolvable_user_turn(&messages);
+        let had_leading_system = matches!(messages.first(), Some(Message::System(_)));
+        if host_renders_catalogue {
+            let mut block = String::new();
+            if !synthesized.is_empty() {
+                block.push_str(&self.render_catalogue(synthesized));
+            }
+            // Only a forced choice still has to be said, since the host's
+            // prompt was composed before the choice was known.
+            match &request.tool_choice {
+                ToolChoice::Required => block.push_str("You must emit at least one tool call.\n"),
+                ToolChoice::Tool(name) => {
+                    block.push_str(&format!("You must call the `{name}` tool.\n"));
+                }
+                ToolChoice::Auto | ToolChoice::None => {}
+            }
+            request.messages = if block.is_empty() {
+                // Nothing to say: no synthesized tool to advertise and no
+                // forced choice, so this rewrite leaves `messages` — and in
+                // particular whether a leading system message exists —
+                // exactly as it already was.
+                messages
+            } else {
+                prompt_tools::append_system_block(&messages, &block)
+            };
+            request.tool_choice = ToolChoice::Auto;
+            sync_stripped_tools_cache_segment(request, had_leading_system);
+            return;
+        }
         request.messages = match self {
             Self::Xml | Self::Native => {
                 prompt_tools::with_tool_instructions(&messages, &tools, &request.tool_choice)
@@ -158,7 +205,118 @@ impl RunDialect {
             }
         };
         request.tool_choice = ToolChoice::Auto;
+        sync_stripped_tools_cache_segment(request, had_leading_system);
+    }
+
+    /// Renders `tools` into this dialect's catalogue shape alone (no
+    /// protocol instructions): the `Self::Xml` full-schema form, the
+    /// `Self::PFormat` positional-signature form, or the `Self::Code`
+    /// function-signature form. Used to advertise a schema the host's own
+    /// static catalogue could not have carried — see
+    /// [`Self::apply_to_request`]'s `synthesized` parameter.
+    fn render_catalogue(&self, tools: &[ToolSchema]) -> String {
+        let specs: Vec<tinytools_agent::tinytools::ToolSpec> = tools
+            .iter()
+            .map(|schema| tinytools_agent::tinytools::ToolSpec {
+                name: schema.name.clone(),
+                description: schema.description.clone(),
+                parameters: schema.parameters.clone(),
+            })
+            .collect();
+        match self {
+            Self::Xml | Self::Native => tinytools_agent::render::render_json_catalogue(&specs),
+            Self::PFormat(_) => tinytools_agent::render::render_pformat_catalogue(&specs),
+            Self::Code(style, _) => tinytools_agent::render::render_code_catalogue(&specs, *style),
+        }
+    }
+}
+
+/// Keeps a harness-declared `cache_segments` layout in sync with a text
+/// dialect's rewrite, using the one thing only this call site still knows
+/// for certain: whether `pre_rewrite_messages` already had a leading system
+/// message *before* the protocol block gets folded in below.
+///
+/// `request.cache_segments` may declare a trailing canonical tools segment
+/// (`{id: "tools", role: Tools, cacheable: true}`) that is about to
+/// disappear once `request.tools` is cleared. When a leading system message
+/// already existed, dropping that trailing segment is all that is needed —
+/// the declared head still names the same messages it always did, and later
+/// fingerprinting (`refresh_prompt_cache_fingerprint`) can verify that by
+/// simple equality. But when none existed yet,
+/// `tinyinference_llm::prompt_tools::append_system_block` (used by both the
+/// host-rendered and ordinary rewrite paths below) synthesizes exactly one
+/// new leading system message for the protocol block — a segment no
+/// declaration could have named in advance. That case is resolved *here*,
+/// with certain knowledge of the pre-rewrite shape, rather than left for
+/// `refresh_prompt_cache_fingerprint` to guess from the rewritten request
+/// alone: reconstructing it after the fact from the post-rewrite shape alone
+/// cannot tell an actually-synthesized segment apart from a custom
+/// declaration that deliberately left an already-present system message out
+/// of the cache key, and conflating the two would silently widen what a
+/// middleware asked to keep out of the stable prefix.
+///
+/// A declaration that is not exactly `[.., tools_segment]` — anything with a
+/// head that does not otherwise account for the messages, or no declaration
+/// at all — is left untouched, so `refresh_prompt_cache_fingerprint` keeps
+/// taking the conservative whole-request digest for it.
+fn sync_stripped_tools_cache_segment(request: &mut ModelRequest, had_leading_system: bool) {
+    let canonical_tools_segment = PromptSegment {
+        id: "tools".to_string(),
+        role: SegmentRole::Tools,
+        cacheable: true,
+    };
+    let Some((last, head)) = request.cache_segments.split_last() else {
+        return;
+    };
+    if *last != canonical_tools_segment {
+        return;
+    }
+    // The real, post-rewrite leading-system-message count — the exact same
+    // thing `refresh_prompt_cache_fingerprint` will independently derive
+    // from `request.messages` a moment later. Deriving the comparison
+    // against this, rather than against `head`'s own declared shape, is
+    // what lets every case below be a plain equality check instead of a
+    // guess: whatever the rewrite actually did to the messages is the one
+    // fact this function can trust.
+    let final_system_end = request
+        .messages
+        .iter()
+        .take_while(|message| matches!(message, Message::System(_)))
+        .count();
+    let canonical_head: Vec<PromptSegment> = (0..final_system_end)
+        .map(|index| PromptSegment {
+            id: crate::prompt::system_segment_id(index),
+            role: SegmentRole::System,
+            cacheable: true,
+        })
+        .collect();
+    if head == canonical_head {
+        // The declared head already names exactly the messages that are
+        // really there (including the trivial `final_system_end == 0`
+        // case, where both sides are empty because this rewrite turned out
+        // to touch nothing — e.g. a host-rendered request with no
+        // synthesized tool and an `Auto` choice leaves `messages`
+        // untouched); only the now-gone trailing tools segment is stale.
+        request.cache_segments = canonical_head;
+    } else if head.is_empty() && !had_leading_system && final_system_end == 1 {
+        // No declared head and no existing leading system message before
+        // this call: the rewrite is the sole source of the new leading
+        // segment (`prompt_tools::append_system_block` inserts exactly one
+        // when none exists), so this is unambiguously the harness's own
+        // synthesis rather than something a declaration could have named in
+        // advance.
+        request.cache_segments = canonical_head;
     }
+    // Every other combination is left completely untouched, including the
+    // trailing tools segment: `head.is_empty() && had_leading_system` is a
+    // declaration that deliberately named nothing ahead of the tools
+    // segment even though a system message already existed (dropping to
+    // `head` would leave an *empty* `cache_segments`, which
+    // `refresh_prompt_cache_fingerprint` reads as "nothing declared yet"
+    // and promotes just the same); a non-empty `head` that does not match
+    // `canonical_head` is a custom declaration (a middleware-owned id, or a
+    // stale count) that must not be silently rewritten out from under it,
+    // partially or otherwise.
 }
 
 /// Builds the positional layout registry the P-Format and code dialects

crates/tinyagents-harness/src/agent_loop/dialect/test.rs

@@ -53,3 +53,167 @@ fn code_dialects_are_opt_in_and_share_the_positional_registry() {
         RunDialect::Code(CodeStyle::TypeScript, _)
     ));
 }
+
+#[test]
+fn a_host_that_renders_the_catalogue_gets_the_schemas_stripped_but_nothing_appended() {
+    use tinyinference_llm::message::Message;
+    use tinyinference_llm::model::{ModelRequest, ToolChoice};
+    use tinyinference_llm::tool::ToolSchema;
+
+    let tools = vec![ToolSchema::new(
+        "lookup",
+        "Looks something up.",
+        serde_json::json!({
+            "type": "object",
+            "properties": {"q": {"type": "string"}},
+            "required": ["q"]
+        }),
+    )];
+    let dialect = RunDialect::resolve(ToolDispatcher::Python, &tools, Some(true));
+    let messages = vec![
+        Message::system("host prompt with its own ## Tools block"),
+        Message::user("hi"),
+    ];
+
+    // Default: the loop appends the protocol block and the catalogue.
+    let mut appended = ModelRequest::new(messages.clone()).with_tools(tools.clone());
+    dialect.apply_to_request(&mut appended, false, &[]);
+    assert!(appended.tools.is_empty());
+    let system = appended.messages[0].text();
+    assert!(system.contains("def lookup("), "{system}");
+
+    // Host-rendered: schemas still leave the wire, the prompt is untouched.
+    let mut host = ModelRequest::new(messages.clone()).with_tools(tools.clone());
+    dialect.apply_to_request(&mut host, true, &[]);
+    assert!(host.tools.is_empty());
+    assert_eq!(host.messages, messages);
+    assert_eq!(host.tool_choice, ToolChoice::Auto);
+
+    // A forced choice is the one thing the host could not have said.
+    let mut forced = ModelRequest::new(messages).with_tools(tools);
+    forced.tool_choice = ToolChoice::Tool("lookup".into());
+    dialect.apply_to_request(&mut forced, true, &[]);
+    let system = forced.messages[0].text();
+    assert!(
+        system.contains("You must call the `lookup` tool."),
+        "{system}"
+    );
+    assert!(!system.contains("def lookup("));
+    assert_eq!(forced.tool_choice, ToolChoice::Auto);
+}
+
+/// Every case above starts from a transcript that already has a leading
+/// system message, so `prompt_tools::append_system_block`'s *other* branch —
+/// inserting a brand-new leading message when none exists yet — is never
+/// exercised. A regression there (failing to insert, inserting more than
+/// one, or inserting it somewhere other than the front) would pass every
+/// other test in this file undetected.
+#[test]
+fn a_run_with_no_leading_system_message_gets_exactly_one_synthesized_by_the_rewrite() {
+    use tinyinference_llm::message::Message;
+    use tinyinference_llm::model::{ModelRequest, ToolChoice};
+    use tinyinference_llm::tool::ToolSchema;
+
+    let tools = vec![ToolSchema::new(
+        "lookup",
+        "Looks something up.",
+        serde_json::json!({
+            "type": "object",
+            "properties": {"q": {"type": "string"}},
+            "required": ["q"]
+        }),
+    )];
+    let dialect = RunDialect::resolve(ToolDispatcher::Python, &tools, Some(true));
+    let no_leading_system = vec![Message::user("hi")];
+
+    // Ordinary rewrite: the loop's own protocol block and catalogue become
+    // the sole, newly-inserted leading system message.
+    let mut appended = ModelRequest::new(no_leading_system.clone()).with_tools(tools.clone());
+    dialect.apply_to_request(&mut appended, false, &[]);
+    assert_eq!(
+        appended.messages.len(),
+        2,
+        "exactly one system message must be inserted, not folded into an \
+         existing one or duplicated: {:?}",
+        appended.messages
+    );
+    assert!(
+        matches!(appended.messages[0], Message::System(_)),
+        "the synthesized message must be the new leading one: {:?}",
+        appended.messages
+    );
+    let system = appended.messages[0].text();
+    assert!(system.contains("def lookup("), "{system}");
+    assert_eq!(appended.messages[1], Message::user("hi"));
+
+    // Host-rendered with a forced choice: the host had no prompt at all to
+    // predate the synthesis, so the forced-choice sentence is what lands in
+    // the newly-inserted message; the catalogue stays the host's job.
+    let mut forced = ModelRequest::new(no_leading_system).with_tools(tools);
+    forced.tool_choice = ToolChoice::Tool("lookup".into());
+    dialect.apply_to_request(&mut forced, true, &[]);
+    assert_eq!(forced.messages.len(), 2, "{:?}", forced.messages);
+    assert!(matches!(forced.messages[0], Message::System(_)));
+    let system = forced.messages[0].text();
+    assert!(
+        system.contains("You must call the `lookup` tool."),
+        "{system}"
+    );
+    assert!(!system.contains("def lookup("));
+    assert_eq!(forced.tool_choice, ToolChoice::Auto);
+}
+
+#[test]
+fn a_host_that_renders_the_catalogue_still_learns_a_turn_synthesized_tool() {
+    use tinyinference_llm::message::Message;
+    use tinyinference_llm::model::{ModelRequest, ToolChoice};
+    use tinyinference_llm::tool::ToolSchema;
+
+    // The base tool the host's own static prompt already advertises.
+    let base = ToolSchema::new(
+        "lookup",
+        "Looks something up.",
+        serde_json::json!({
+            "type": "object",
+            "properties": {"q": {"type": "string"}},
+            "required": ["q"]
+        }),
+    );
+    // The structured-output fallback tool minted for this turn only, after
+    // the host's prompt was already composed — see
+    // `RunPolicy::host_renders_tool_catalogue` and
+    // `StructuredStrategy::ToolCall`.
+    let synthesized = ToolSchema::new(
+        "emit_result",
+        "Return the result as `emit_result`.",
+        serde_json::json!({
+            "type": "object",
+            "properties": {"total": {"type": "number"}},
+            "required": ["total"]
+        }),
+    );
+    let dialect = RunDialect::resolve(
+        ToolDispatcher::Python,
+        std::slice::from_ref(&base),
+        Some(true),
+    );
+    let messages = vec![
+        Message::system("host prompt with its own ## Tools block"),
+        Message::user("hi"),
+    ];
+
+    let mut request =
+        ModelRequest::new(messages).with_tools(vec![base.clone(), synthesized.clone()]);
+    request.tool_choice = ToolChoice::Tool("emit_result".into());
+    dialect.apply_to_request(&mut request, true, std::slice::from_ref(&synthesized));
+
+    assert!(request.tools.is_empty());
+    let system = request.messages[0].text();
+    // The base tool is left to the host's own (untouched) catalogue...
+    assert!(!system.contains("def lookup("), "{system}");
+    // ...but the synthesized one, the host could never have known about, is
+    // appended so the forced call has a schema to answer against.
+    assert!(system.contains("def emit_result("), "{system}");
+    assert!(system.contains("You must call the `emit_result` tool."));
+    assert_eq!(request.tool_choice, ToolChoice::Auto);
+}