heartbit_core/browser/
distill_tool.rs

1//! `DistillingTool` — an `Arc<dyn Tool>` decorator that distills snapshot output
2//! before it re-enters the agent's context (token-control wiring).
3//!
4//! [`super::distill::distill_snapshot`] is a pure transform, but nothing applied
5//! it to live tool output: every `take_snapshot` (and every mutating tool with
6//! `includeSnapshot:true`) returned the FULL accessibility tree, which the agent
7//! loop then re-sends on every subsequent turn. On a busy page that is the single
8//! largest input-token cost. This decorator runs the distiller on any output that
9//! looks like a snapshot (contains a `uid=` line), dropping the redundant echoed
10//! `StaticText` while preserving every interactive `uid` — so the agent sees a
11//! smaller, clearer tree and the conversation grows far more slowly.
12//!
13//! It composes with [`super::harness::ReliableInteractionTool`]: distillation
14//! wraps the OUTSIDE (`DistillingTool::wrap(reliable)`), so a mutating tool's
15//! forced fresh snapshot is also distilled. Distillation is idempotent and only
16//! touches `uid`-bearing content, so applying it uniformly is safe — non-snapshot
17//! outputs (network lists, console logs, plain text) pass through untouched.
18
19use std::sync::Arc;
20
21use serde_json::Value;
22
23use crate::error::Error;
24use crate::execution_context::ExecutionContext;
25use crate::llm::types::ToolDefinition;
26use crate::tool::{Tool, ToolOutput};
27
28use super::distill::{DistillConfig, distill_snapshot};
29
30/// Decorator that distills snapshot-shaped output of the wrapped tool.
31pub struct DistillingTool {
32    inner: Arc<dyn Tool>,
33    cfg: DistillConfig,
34}
35
36impl DistillingTool {
37    /// Wrap `inner`, distilling its snapshot output with `cfg`.
38    pub fn wrap(inner: Arc<dyn Tool>, cfg: DistillConfig) -> Self {
39        Self { inner, cfg }
40    }
41
42    /// Wrap every tool in a set; a drop-in replacement for an agent's tools.
43    pub fn wrap_all(tools: Vec<Arc<dyn Tool>>, cfg: DistillConfig) -> Vec<Arc<dyn Tool>> {
44        tools
45            .into_iter()
46            .map(|t| Arc::new(Self::wrap(t, cfg.clone())) as Arc<dyn Tool>)
47            .collect()
48    }
49}
50
51impl Tool for DistillingTool {
52    fn definition(&self) -> ToolDefinition {
53        self.inner.definition()
54    }
55
56    fn execute(
57        &self,
58        ctx: &ExecutionContext,
59        input: Value,
60    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<ToolOutput, Error>> + Send + '_>>
61    {
62        // Own everything the future needs so it borrows only `&self`'s elided
63        // lifetime (mirrors ReliableInteractionTool's ctx-clone fix).
64        let inner = Arc::clone(&self.inner);
65        let cfg = self.cfg.clone();
66        let ctx = ctx.clone();
67        Box::pin(async move {
68            let out = inner.execute(&ctx, input).await?;
69            // Only distill snapshot-shaped output; pass everything else through.
70            if out.content.contains("uid=") {
71                Ok(ToolOutput {
72                    content: distill_snapshot(&out.content, &cfg),
73                    is_error: out.is_error,
74                })
75            } else {
76                Ok(out)
77            }
78        })
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    /// A tool returning a fixed output, for decorator tests.
87    struct FixedTool {
88        name: String,
89        output: ToolOutput,
90    }
91    impl Tool for FixedTool {
92        fn definition(&self) -> ToolDefinition {
93            ToolDefinition {
94                name: self.name.clone(),
95                description: "fixed".into(),
96                input_schema: serde_json::json!({ "type": "object" }),
97            }
98        }
99        fn execute(
100            &self,
101            _ctx: &ExecutionContext,
102            _input: Value,
103        ) -> std::pin::Pin<
104            Box<dyn std::future::Future<Output = Result<ToolOutput, Error>> + Send + '_>,
105        > {
106            let out = self.output.clone();
107            Box::pin(async move { Ok(out) })
108        }
109    }
110
111    fn fixed(name: &str, content: &str, is_error: bool) -> Arc<dyn Tool> {
112        Arc::new(FixedTool {
113            name: name.to_string(),
114            output: ToolOutput {
115                content: content.to_string(),
116                is_error,
117            },
118        })
119    }
120
121    // A snapshot with the canonical redundancy: each link followed by a child
122    // StaticText echoing its name (which the distiller drops).
123    const SNAP: &str = r#"## Latest page snapshot
124uid=1_0 RootWebArea "Hacker News" url="https://news.ycombinator.com/"
125  uid=1_2 link "Hacker News" url="https://news.ycombinator.com/news"
126    uid=1_3 StaticText "Hacker News"
127  uid=1_4 link "new" url="https://news.ycombinator.com/newest"
128    uid=1_5 StaticText "new""#;
129
130    #[tokio::test]
131    async fn distills_snapshot_output_and_shrinks_it() {
132        let tool = DistillingTool::wrap(
133            fixed("take_snapshot", SNAP, false),
134            DistillConfig::default(),
135        );
136        let out = tool
137            .execute(&ExecutionContext::default(), serde_json::json!({}))
138            .await
139            .expect("ok");
140        assert!(
141            out.content.len() < SNAP.len(),
142            "distilled output must be smaller"
143        );
144        // Interactive uids survive...
145        assert!(out.content.contains("uid=1_2 link \"Hacker News\""));
146        assert!(out.content.contains("uid=1_4 link \"new\""));
147        // ...redundant echoed StaticText is gone.
148        assert!(!out.content.contains("uid=1_3 StaticText \"Hacker News\""));
149        // The header (no uid) is preserved.
150        assert!(out.content.contains("## Latest page snapshot"));
151    }
152
153    #[tokio::test]
154    async fn preserves_every_interactive_uid() {
155        let tool = DistillingTool::wrap(
156            fixed("take_snapshot", SNAP, false),
157            DistillConfig::default(),
158        );
159        let out = tool
160            .execute(&ExecutionContext::default(), serde_json::json!({}))
161            .await
162            .expect("ok");
163        for uid in ["uid=1_2", "uid=1_4"] {
164            assert!(out.content.contains(uid), "interactive {uid} must survive");
165        }
166    }
167
168    #[tokio::test]
169    async fn non_snapshot_output_passes_through_untouched() {
170        // A network-list / console output with no uid lines must be unchanged.
171        let raw = "reqid=3 GET https://example.com/ [304]\nreqid=4 GET /style.css [200]";
172        let tool = DistillingTool::wrap(
173            fixed("list_network_requests", raw, false),
174            DistillConfig::default(),
175        );
176        let out = tool
177            .execute(&ExecutionContext::default(), serde_json::json!({}))
178            .await
179            .expect("ok");
180        assert_eq!(out.content, raw, "non-snapshot output must be untouched");
181    }
182
183    #[tokio::test]
184    async fn preserves_is_error_flag() {
185        let tool = DistillingTool::wrap(
186            fixed("click", "Error: no element with uid=9_9", true),
187            DistillConfig::default(),
188        );
189        let out = tool
190            .execute(&ExecutionContext::default(), serde_json::json!({}))
191            .await
192            .expect("ok");
193        assert!(
194            out.is_error,
195            "error flag must be preserved through distillation"
196        );
197    }
198
199    #[tokio::test]
200    async fn definition_is_passthrough() {
201        let tool = DistillingTool::wrap(
202            fixed("take_snapshot", SNAP, false),
203            DistillConfig::default(),
204        );
205        assert_eq!(tool.definition().name, "take_snapshot");
206    }
207
208    #[test]
209    fn wrap_all_preserves_count_and_names() {
210        let tools = vec![
211            fixed("take_snapshot", SNAP, false),
212            fixed("click", "ok", false),
213        ];
214        let wrapped = DistillingTool::wrap_all(tools, DistillConfig::default());
215        assert_eq!(wrapped.len(), 2);
216        assert_eq!(wrapped[0].definition().name, "take_snapshot");
217        assert_eq!(wrapped[1].definition().name, "click");
218    }
219}
heartbit_core/browser/distill_tool.rs

heartbit_core/browser/
distill_tool.rs
