fix(harness): screen hosted extension content

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

Author: senamakel

crates/tinyagents-harness/src/context/stats.rs

@@ -69,21 +69,25 @@ pub fn estimate_context_tokens(messages: &[Message], tokenize: impl Fn(&str) ->
     messages
         .iter()
         .map(|message| {
-            let visible = match message {
-                Message::User(user) => user
-                    .content
-                    .iter()
-                    .filter_map(|block| match block {
-                        ContentBlock::Text(text) => Some(text.clone()),
-                        ContentBlock::Json(value) | ContentBlock::ProviderExtension(value) => {
-                            Some(value.to_string())
-                        }
-                        _ => None,
-                    })
-                    .collect::<Vec<_>>()
-                    .join("\n"),
-                _ => message.text(),
+            let content = match message {
+                Message::System(message) => &message.content,
+                Message::User(message) => &message.content,
+                Message::Assistant(message) => &message.content,
+                Message::Tool(message) => &message.content,
             };
+            let visible = content
+                .iter()
+                .filter_map(|block| match block {
+                    ContentBlock::Text(text) | ContentBlock::Thinking { text, .. } => {
+                        Some(text.clone())
+                    }
+                    ContentBlock::Json(value) | ContentBlock::ProviderExtension(value) => {
+                        Some(value.to_string())
+                    }
+                    _ => None,
+                })
+                .collect::<Vec<_>>()
+                .join("\n");
             tokenize(&visible)
         })
         .sum()

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

@@ -229,6 +229,37 @@ impl<State: Send + Sync, Ctx: Send + Sync> AgentHarness<State, Ctx> {
         }
     }
 
+    /// Collects a hosted turn through the streaming driver while preserving the
+    /// parent's exact capability bundle. Recursive streaming delegation uses
+    /// this rather than the unary entry point so model deltas and delta
+    /// middleware remain part of the shared parent event stream.
+    pub(crate) async fn invoke_agent_streaming_with_host_capabilities(
+        &self,
+        host: crate::host::HostCapabilities<State>,
+        request: AgentTurnRequest,
+        context: RunContext<Ctx>,
+        state: &State,
+    ) -> Result<AgentRun>
+    where
+        Ctx: 'static,
+        State: 'static,
+    {
+        let stream = self
+            .invoke_agent_stream_with_host_capabilities(host, request, context, state)
+            .await?;
+        futures::pin_mut!(stream);
+        while let Some(item) = stream.next().await {
+            match item {
+                AgentStreamItem::Completed(run) => return Ok(*run),
+                AgentStreamItem::Failed { error, .. } => return Err(TinyAgentsError::Model(error)),
+                AgentStreamItem::Event(_) => {}
+            }
+        }
+        Err(TinyAgentsError::Model(
+            "hosted stream ended without a terminal result".to_string(),
+        ))
+    }
+
     /// Starts a hosted streaming turn.
     ///
     /// The returned stream is the existing event projection, so host-driven
@@ -598,7 +629,10 @@ async fn screen_user_messages<State: Send + Sync>(
                             return Err(TinyAgentsError::Validation(reason));
                         }
                     }
-                } else if let tinyinference_llm::message::ContentBlock::Json(value) = block {
+                } else if let tinyinference_llm::message::ContentBlock::Json(value)
+                | tinyinference_llm::message::ContentBlock::ProviderExtension(value) =
+                    block
+                {
                     let rendered = value.to_string();
                     match host
                         .security

crates/tinyagents-harness/src/subagent/mod.rs

@@ -88,7 +88,7 @@ use crate::runtime::AgentHarness;
 use crate::tool::ToolDispatch;
 use tinyinference_llm::message::Message;
 
-impl<State: Send + Sync + 'static, Ctx: Send + Sync> SubAgent<State, Ctx> {
+impl<State: Send + Sync + 'static, Ctx: Send + Sync + 'static> SubAgent<State, Ctx> {
     /// Creates a sub-agent wrapping `harness` with a stable `name` and
     /// `description`.
     pub fn new(
@@ -306,9 +306,25 @@ impl<State: Send + Sync + 'static, Ctx: Send + Sync> SubAgent<State, Ctx> {
             // capability bundle. The child harness's installed host (including
             // no host at all) is intentionally irrelevant here: allowing it
             // to decide policy would make delegation authorization bypassable.
-            self.harness
-                .invoke_agent_with_host_capabilities(authority.host.clone(), request, ctx, state)
-                .await?
+            if streaming {
+                self.harness
+                    .invoke_agent_streaming_with_host_capabilities(
+                        authority.host.clone(),
+                        request,
+                        ctx,
+                        state,
+                    )
+                    .await?
+            } else {
+                self.harness
+                    .invoke_agent_with_host_capabilities(
+                        authority.host.clone(),
+                        request,
+                        ctx,
+                        state,
+                    )
+                    .await?
+            }
         } else if streaming {
             self.harness
                 .invoke_streaming_in_context(state, ctx, messages)
@@ -486,7 +502,7 @@ impl<State: Send + Sync, Ctx: Send + Sync> SubAgentSession<State, Ctx> {
     }
 }
 
-impl<State: Send + Sync + 'static, Ctx: Send + Sync> SubAgentTool<State, Ctx> {
+impl<State: Send + Sync + 'static, Ctx: Send + Sync + 'static> SubAgentTool<State, Ctx> {
     /// Default JSON Schema for a sub-agent tool: an object with one required
     /// string field named [`SUBAGENT_INPUT_FIELD`].
     fn default_parameters() -> Value {
@@ -620,7 +636,7 @@ impl<State: Send + Sync + 'static, Ctx: Send + Sync> SubAgentTool<State, Ctx> {
 impl<State, Ctx> ToolDispatch<State, Ctx> for SubAgentTool<State, Ctx>
 where
     State: Send + Sync + 'static,
-    Ctx: Send + Sync,
+    Ctx: Send + Sync + 'static,
 {
     fn tool(&self) -> Arc<dyn tinytools::Tool> {
         Arc::new(SubAgentToolDeclaration {