fix(harness): reserve hosted terminal progress

Co-authored-by: Medulla <medulla@tinyhumans.ai>

Author: senamakel

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

@@ -826,20 +826,51 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> {
     ///   fail-closed behaviour `run_loop` already has for a pre-wrap override.
     ///   Silently substituting a different model is the one outcome that is
     ///   never acceptable.
-    fn rebind(
+    async fn rebind(
         &self,
         ctx: &mut RunContext<Ctx>,
         request: &ModelRequest,
-    ) -> ResolvedModelBinding<State> {
+    ) -> Result<ResolvedModelBinding<State>> {
         let captured = || ResolvedModelBinding {
             resolved: self.resolved.clone(),
             model: Arc::clone(&self.model),
         };
         let Some(requested) = request.model.as_deref() else {
-            return captured();
+            return Ok(captured());
         };
+        if let Some(host_run) = self.harness.host_run_binding(ctx.instance_id())? {
+            let mut resolve = crate::host::ModelResolveRequest::new(host_run.agent_id.clone());
+            if ctx.depth() == 0 {
+                resolve = resolve.as_team_lead();
+            }
+            if let Some(role) = host_run.role.clone() {
+                resolve = resolve.with_role(role);
+            }
+            if let Some(pin) = request.model.clone().or(host_run.model_pin.clone()) {
+                resolve = resolve.with_model_pin(pin);
+            }
+            if let Some(capabilities) = request.required_capabilities.clone() {
+                resolve = resolve.with_required_capabilities(capabilities);
+            }
+            let model = host_run.host.models.resolve(&resolve).await.map_err(|error| {
+                tinyagents_tracing::warn!(%error, agent_id = %host_run.agent_id, "[host] wrap model resolution failed");
+                TinyAgentsError::Model("host model resolution failed".to_string())
+            })?;
+            let name = model
+                .profile()
+                .and_then(|profile| profile.model.clone())
+                .unwrap_or_else(|| format!("host:{}", host_run.agent_id));
+            return Ok(ResolvedModelBinding {
+                resolved: ResolvedModel {
+                    name,
+                    requested: resolve.model_pin,
+                    source: ModelResolutionSource::RequestOverride,
+                },
+                model,
+            });
+        }
         if requested == self.resolved.name {
-            return captured();
+            return Ok(captured());
         }
         match self.harness.models.resolve_request(request, None, None) {
             Some(binding)
@@ -852,7 +883,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> {
                     to = %binding.resolved.name,
                     "[model] wrap layer overrode the model; re-resolved the binding"
                 );
-                binding
+                Ok(binding)
             }
             _ => {
                 tinyagents_tracing::warn!(
@@ -865,7 +896,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelCallBase<'_, State, Ctx> {
                     requested: requested.to_string(),
                     resolved: self.resolved.name.clone(),
                 });
-                captured()
+                Ok(captured())
             }
         }
     }
@@ -881,7 +912,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> ModelBaseCall<State, Ctx>
         request: ModelRequest,
     ) -> BoxModelFuture<'a> {
         Box::pin(async move {
-            let binding = self.rebind(ctx, &request);
+            let binding = self.rebind(ctx, &request).await?;
             self.harness
                 .invoke_model_with_retry(
                     state,

crates/tinyagents-harness/src/runtime/agent.rs

@@ -128,7 +128,28 @@ struct PreparedAgentTurn<State: Send + Sync> {
     run_id: crate::ids::RunId,
     input_text: String,
     messages: Vec<tinyinference_llm::message::Message>,
-    progress: Option<tokio::sync::mpsc::Sender<ProgressEvent>>,
+    progress: Option<ProgressSender>,
+}
+
+#[derive(Clone)]
+pub(crate) struct ProgressSender {
+    tx: tokio::sync::mpsc::Sender<ProgressEvent>,
+    nonterminal_slots: std::sync::Arc<tokio::sync::Semaphore>,
+}
+
+impl ProgressSender {
+    fn send_nonterminal(&self, event: ProgressEvent) {
+        let Ok(permit) = self.nonterminal_slots.clone().try_acquire_owned() else {
+            return;
+        };
+        if self.tx.try_send(event).is_ok() {
+            permit.forget();
+        }
+    }
+
+    fn send_terminal(&self, event: ProgressEvent) {
+        let _ = self.tx.try_send(event);
+    }
 }
 
 impl<State: Send + Sync> Clone for PreparedAgentTurn<State> {
@@ -450,11 +471,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> {
         let Some(progress) = binding.progress else {
             return;
         };
-        if progress.try_send(event).is_err() {
-            tinyagents_tracing::debug!(
-                "[host] dropping progress event because the bounded queue is full or closed"
-            );
-        }
+        progress.send_nonterminal(event);
     }
 }
 
@@ -484,12 +501,12 @@ async fn finish_host_turn<State: Send + Sync>(
     // consumer completely outside the agent's critical path.
     if let Some(progress) = &prepared.progress {
         if let Some(message) = error {
-            let _ = progress.try_send(ProgressEvent::Error {
+            progress.send_terminal(ProgressEvent::Error {
                 run: prepared.run_id.clone(),
                 message,
             });
         } else {
-            let _ = progress.try_send(ProgressEvent::Finished {
+            progress.send_terminal(ProgressEvent::Finished {
                 run: prepared.run_id.clone(),
                 usage: Some(run.usage.usage),
             });
@@ -533,21 +550,29 @@ async fn finish_host_turn<State: Send + Sync>(
 
 fn start_progress_dispatcher(
     sink: Option<std::sync::Arc<dyn crate::host::ProgressSink>>,
-) -> Option<tokio::sync::mpsc::Sender<ProgressEvent>> {
+) -> Option<ProgressSender> {
     let sink = sink?;
     let handle = tokio::runtime::Handle::try_current().ok()?;
     // Progress is observational. Bound it so a slow sink cannot retain every
     // streamed token; producers use `try_send` and drop overflowed updates.
     // One slot is reserved for the single terminal outcome. Producers use
     // `try_send` for ordinary progress, so at most 128 nonterminal events can
     // fill before finalization claims the remaining slot.
-    let (tx, mut rx) = tokio::sync::mpsc::channel(129);
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<ProgressEvent>(129);
+    let nonterminal_slots = std::sync::Arc::new(tokio::sync::Semaphore::new(128));
+    let released_slots = nonterminal_slots.clone();
     handle.spawn(async move {
         while let Some(event) = rx.recv().await {
+            if !event.is_terminal() {
+                released_slots.add_permits(1);
+            }
             sink.emit(event).await;
         }
     });
-    Some(tx)
+    Some(ProgressSender {
+        tx,
+        nonterminal_slots,
+    })
 }
 
 async fn screen_user_messages<State: Send + Sync>(

crates/tinyagents-harness/src/runtime/test.rs

@@ -49,6 +49,7 @@ struct RetryableClassifier;
 struct LeadRecordingResolver {
     model: Arc<dyn ChatModel<()>>,
     team_lead_flags: Mutex<Vec<bool>>,
+    model_pins: Mutex<Vec<Option<String>>>,
 }
 
 struct RecordingBudget {
@@ -146,6 +147,10 @@ impl crate::host::ModelResolver<()> for LeadRecordingResolver {
             .lock()
             .expect("resolver lock")
             .push(request.is_team_lead);
+        self.model_pins
+            .lock()
+            .expect("resolver lock")
+            .push(request.model_pin.clone());
         Ok(Arc::clone(&self.model))
     }
 }
@@ -524,6 +529,7 @@ async fn hosted_model_resolution_marks_only_root_contexts_as_team_leads() {
     let resolver = Arc::new(LeadRecordingResolver {
         model: model.clone(),
         team_lead_flags: Mutex::new(Vec::new()),
+        model_pins: Mutex::new(Vec::new()),
     });
     let host = crate::host::HostCapabilities::new(
         Arc::new(StaticContextComposer::empty()),

crates/tinyagents-harness/src/runtime/types.rs

@@ -50,7 +50,7 @@ pub(crate) struct HostRunBinding<State: Send + Sync> {
     /// list is a host boundary enforced for schemas and dispatch alike.
     pub(crate) allowed_tools: HashSet<String>,
     /// Per-turn ordered, nonblocking projection to the optional progress sink.
-    pub(crate) progress: Option<tokio::sync::mpsc::Sender<crate::host::ProgressEvent>>,
+    pub(crate) progress: Option<super::agent::ProgressSender>,
 }
 
 impl<State: Send + Sync> Clone for HostRunBinding<State> {