From 64ae8471326f4af231dc2f9e35fcc99c2d27fcd0 Mon Sep 17 00:00:00 2001
From: Ning Shi <zeegeek@gmail.com>
Date: Sun, 26 Apr 2026 14:36:09 -0400
Subject: [PATCH 1/2] fix: router delegates customize_prompt_builder with
 resolved model name

RouterProvider inherited the default no-op for customize_prompt_builder,
so provider-specific prompt customizations (e.g. Ollama injecting the
<|think|> token for Gemma models) were silently skipped when routing
through hints like "hint:gemma". Resolve the hint to the actual provider
and model name before delegating.
---
 src/agent/prompt.rs     |  1 +
 src/providers/router.rs | 97 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

diff --git a/src/agent/prompt.rs b/src/agent/prompt.rs
index 1cfa5642..75e998f7 100644
--- a/src/agent/prompt.rs
+++ b/src/agent/prompt.rs
@@ -197,6 +197,7 @@ pub fn timestamp_prefix(message: &str, context: Option<&str>) -> String {
     }
 }
 
+#[derive(Clone)]
 pub struct PromptContext<'a> {
     pub workspace_dir: &'a Path,
     pub model_name: &'a str,
diff --git a/src/providers/router.rs b/src/providers/router.rs
index 845330a2..83a20d6c 100644
--- a/src/providers/router.rs
+++ b/src/providers/router.rs
@@ -323,6 +323,18 @@ impl Provider for RouterProvider {
             .any(|(_, provider)| provider.supports_vision())
     }
 
+    fn customize_prompt_builder(
+        &self,
+        builder: &mut crate::agent::prompt::SystemPromptBuilder,
+        ctx: &crate::agent::prompt::PromptContext<'_>,
+    ) {
+        let (provider_idx, resolved_model) = self.resolve(ctx.model_name);
+        let mut resolved_ctx = ctx.clone();
+        resolved_ctx.model_name = &resolved_model;
+        let (_, provider) = &self.providers[provider_idx];
+        provider.customize_prompt_builder(builder, &resolved_ctx);
+    }
+
     async fn warmup(&self) -> anyhow::Result<()> {
         for (name, provider) in &self.providers {
             tracing::info!(provider = name, "Warming up routed provider");
@@ -1198,4 +1210,89 @@ mod tests {
         assert_eq!(streaming.tool_event_calls.load(Ordering::SeqCst), 1);
         assert_eq!(*streaming.last_stream_model.lock(), "claude-opus");
     }
+
+    /// Mock that records the model name seen by `customize_prompt_builder`.
+    struct PromptCaptureMock {
+        seen_model: parking_lot::Mutex<Option<String>>,
+    }
+
+    impl PromptCaptureMock {
+        fn new() -> Self {
+            Self {
+                seen_model: parking_lot::Mutex::new(None),
+            }
+        }
+    }
+
+    #[async_trait]
+    impl Provider for PromptCaptureMock {
+        async fn chat_with_system(
+            &self,
+            _: Option<&str>,
+            _: &str,
+            _: &str,
+            _: f64,
+        ) -> anyhow::Result<String> {
+            Ok(String::new())
+        }
+
+        fn customize_prompt_builder(
+            &self,
+            _builder: &mut crate::agent::prompt::SystemPromptBuilder,
+            ctx: &crate::agent::prompt::PromptContext<'_>,
+        ) {
+            *self.seen_model.lock() = Some(ctx.model_name.to_string());
+        }
+    }
+
+    #[test]
+    fn customize_prompt_builder_delegates_with_resolved_model() {
+        use crate::agent::prompt::{SystemPromptBuilder, test_helpers::make_test_ctx};
+
+        let mock = Arc::new(PromptCaptureMock::new());
+        let providers: Vec<(String, Box<dyn Provider>)> =
+            vec![("ollama".into(), Box::new(Arc::clone(&mock)) as Box<dyn Provider>)];
+        let routes = vec![(
+            "gemma".to_string(),
+            Route {
+                provider_name: "ollama".to_string(),
+                model: "gemma4".to_string(),
+            },
+        )];
+        let router = RouterProvider::new(providers, routes, "default-model".to_string());
+
+        let tools: Vec<Box<dyn crate::tools::Tool>> = vec![];
+        let mut ctx = make_test_ctx(&tools);
+        ctx.model_name = "hint:gemma";
+
+        let mut builder = SystemPromptBuilder::with_defaults();
+        router.customize_prompt_builder(&mut builder, &ctx);
+
+        assert_eq!(
+            mock.seen_model.lock().as_deref(),
+            Some("gemma4"),
+            "Router should resolve hint to actual model name before delegating"
+        );
+    }
+
+    #[async_trait]
+    impl Provider for Arc<PromptCaptureMock> {
+        async fn chat_with_system(
+            &self,
+            s: Option<&str>,
+            m: &str,
+            model: &str,
+            t: f64,
+        ) -> anyhow::Result<String> {
+            self.as_ref().chat_with_system(s, m, model, t).await
+        }
+
+        fn customize_prompt_builder(
+            &self,
+            builder: &mut crate::agent::prompt::SystemPromptBuilder,
+            ctx: &crate::agent::prompt::PromptContext<'_>,
+        ) {
+            self.as_ref().customize_prompt_builder(builder, ctx);
+        }
+    }
 }

From ee94ab4415faca7f00d57d12d4502af832187d5b Mon Sep 17 00:00:00 2001
From: Ning Shi <zeegeek@gmail.com>
Date: Sun, 26 Apr 2026 14:55:05 -0400
Subject: [PATCH 2/2] style: fix cargo fmt formatting in router test

---
 src/providers/router.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/providers/router.rs b/src/providers/router.rs
index 83a20d6c..b655c63a 100644
--- a/src/providers/router.rs
+++ b/src/providers/router.rs
@@ -1250,8 +1250,10 @@ mod tests {
         use crate::agent::prompt::{SystemPromptBuilder, test_helpers::make_test_ctx};
 
         let mock = Arc::new(PromptCaptureMock::new());
-        let providers: Vec<(String, Box<dyn Provider>)> =
-            vec![("ollama".into(), Box::new(Arc::clone(&mock)) as Box<dyn Provider>)];
+        let providers: Vec<(String, Box<dyn Provider>)> = vec![(
+            "ollama".into(),
+            Box::new(Arc::clone(&mock)) as Box<dyn Provider>,
+        )];
         let routes = vec![(
             "gemma".to_string(),
             Route {