Skip to main content

harness_webfetch/
types.rs

1use harness_core::{PermissionPolicy, ToolError};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5
6use crate::engine::WebFetchEngine;
7
8#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "UPPERCASE")]
10pub enum WebFetchMethod {
11    Get,
12    Post,
13}
14
15impl WebFetchMethod {
16    pub fn as_str(&self) -> &'static str {
17        match self {
18            Self::Get => "GET",
19            Self::Post => "POST",
20        }
21    }
22}
23
24#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
25#[serde(rename_all = "snake_case")]
26pub enum WebFetchExtract {
27    Markdown,
28    Raw,
29    Both,
30}
31
32impl WebFetchExtract {
33    pub fn as_str(&self) -> &'static str {
34        match self {
35            Self::Markdown => "markdown",
36            Self::Raw => "raw",
37            Self::Both => "both",
38        }
39    }
40}
41
42/// Session permission policy plus the autonomous escape hatch for tests.
43#[derive(Clone)]
44pub struct WebFetchPermissionPolicy {
45    pub inner: PermissionPolicy,
46    pub unsafe_allow_fetch_without_hook: bool,
47}
48
49impl WebFetchPermissionPolicy {
50    pub fn new(inner: PermissionPolicy) -> Self {
51        Self {
52            inner,
53            unsafe_allow_fetch_without_hook: false,
54        }
55    }
56
57    pub fn with_unsafe_bypass(mut self, v: bool) -> Self {
58        self.unsafe_allow_fetch_without_hook = v;
59        self
60    }
61}
62
63impl std::fmt::Debug for WebFetchPermissionPolicy {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("WebFetchPermissionPolicy")
66            .field(
67                "unsafe_allow_fetch_without_hook",
68                &self.unsafe_allow_fetch_without_hook,
69            )
70            .field("inner", &self.inner)
71            .finish()
72    }
73}
74
75pub type WebFetchCache = Arc<Mutex<HashMap<String, CachedResponse>>>;
76
77#[derive(Clone)]
78pub struct WebFetchSessionConfig {
79    pub permissions: WebFetchPermissionPolicy,
80    pub engine: Arc<dyn WebFetchEngine>,
81    pub default_headers: Option<HashMap<String, String>>,
82    pub allow_loopback: bool,
83    pub allow_private_networks: bool,
84    pub allow_metadata: bool,
85    pub resolve_once: bool,
86    pub default_timeout_ms: Option<u64>,
87    pub session_backstop_ms: Option<u64>,
88    pub max_redirects: Option<u32>,
89    pub inline_markdown_cap: Option<usize>,
90    pub inline_raw_cap: Option<usize>,
91    pub spill_hard_cap: Option<usize>,
92    pub cache_ttl_ms: Option<u64>,
93    pub spill_dir: Option<String>,
94    pub session_id: Option<String>,
95    pub cache: Option<WebFetchCache>,
96}
97
98impl WebFetchSessionConfig {
99    pub fn new(permissions: WebFetchPermissionPolicy, engine: Arc<dyn WebFetchEngine>) -> Self {
100        Self {
101            permissions,
102            engine,
103            default_headers: None,
104            allow_loopback: false,
105            allow_private_networks: false,
106            allow_metadata: false,
107            resolve_once: true,
108            default_timeout_ms: None,
109            session_backstop_ms: None,
110            max_redirects: None,
111            inline_markdown_cap: None,
112            inline_raw_cap: None,
113            spill_hard_cap: None,
114            cache_ttl_ms: None,
115            spill_dir: None,
116            session_id: None,
117            cache: None,
118        }
119    }
120
121    pub fn with_cache(mut self) -> Self {
122        self.cache = Some(Arc::new(Mutex::new(HashMap::new())));
123        self
124    }
125}
126
127impl std::fmt::Debug for WebFetchSessionConfig {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        f.debug_struct("WebFetchSessionConfig")
130            .field("permissions", &self.permissions)
131            .field("allow_loopback", &self.allow_loopback)
132            .field("allow_private_networks", &self.allow_private_networks)
133            .field("allow_metadata", &self.allow_metadata)
134            .field("has_cache", &self.cache.is_some())
135            .finish()
136    }
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct CachedResponse {
141    pub at_ms: u64,
142    pub status: u16,
143    pub final_url: String,
144    pub redirect_chain: Vec<String>,
145    pub content_type: String,
146    pub body: Vec<u8>,
147    pub extract: WebFetchExtract,
148    #[serde(default, skip_serializing_if = "Option::is_none")]
149    pub extracted_markdown: Option<String>,
150}
151
152#[derive(Debug, Clone, Serialize, Deserialize)]
153pub struct FetchMetadata {
154    pub url: String,
155    pub final_url: String,
156    pub method: WebFetchMethod,
157    pub status: u16,
158    pub content_type: String,
159    pub redirect_chain: Vec<String>,
160    pub fetched_ms: u64,
161    pub from_cache: bool,
162    #[serde(default, skip_serializing_if = "Option::is_none")]
163    pub cache_age_sec: Option<u64>,
164}
165
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct WebFetchOk {
168    pub output: String,
169    pub meta: FetchMetadata,
170    #[serde(default, skip_serializing_if = "Option::is_none")]
171    pub body_markdown: Option<String>,
172    #[serde(default, skip_serializing_if = "Option::is_none")]
173    pub body_raw: Option<String>,
174    #[serde(default, skip_serializing_if = "Option::is_none")]
175    pub log_path: Option<String>,
176    pub byte_cap: bool,
177}
178
179#[derive(Debug, Clone, Serialize, Deserialize)]
180pub struct WebFetchRedirectLoop {
181    pub output: String,
182    pub meta: FetchMetadata,
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize)]
186pub struct WebFetchHttpError {
187    pub output: String,
188    pub meta: FetchMetadata,
189    pub body_raw: String,
190    #[serde(default, skip_serializing_if = "Option::is_none")]
191    pub log_path: Option<String>,
192    pub byte_cap: bool,
193}
194
195#[derive(Debug, Clone, Serialize, Deserialize)]
196pub struct WebFetchError {
197    pub error: ToolError,
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize)]
201#[serde(tag = "kind", rename_all = "snake_case")]
202pub enum WebFetchResult {
203    #[serde(rename = "ok")]
204    Ok(WebFetchOk),
205    #[serde(rename = "redirect_loop")]
206    RedirectLoop(WebFetchRedirectLoop),
207    #[serde(rename = "http_error")]
208    HttpError(WebFetchHttpError),
209    #[serde(rename = "error")]
210    Error(WebFetchError),
211}
212
213impl From<WebFetchError> for WebFetchResult {
214    fn from(e: WebFetchError) -> Self {
215        WebFetchResult::Error(e)
216    }
217}