abyss_agent_hook/
hook.rs

1//! Single coordinator for Harness detection, protocol parsing, correlation, and delivery.
2
3use std::sync::Arc;
4
5use abyss_mitm::{FlowContext, HookError, HookFuture, HttpExchange, MitmHook, WebSocketMessage};
6use abyss_plugin_protocol::event::LlmProvider;
7use http::{HeaderMap, header::HOST};
8
9use crate::{
10    config::{HarnessUsageHookConfig, HooksRuntimeConfig},
11    correlation::CorrelationRegistry,
12    delivery::{AgentEventSink, EventDelivery},
13    event::{NormalizableExchange, NormalizedUsageEvent, normalize_exchange},
14    harness::{HarnessDetection, HarnessRegistry},
15    protocol::{
16        AnthropicProtocol, ClaudeWebProtocol, LlmProtocol, OpenAiProtocol, ProtocolDetector,
17        anthropic::{
18            ClaudeWebContext, parse_anthropic_messages_exchange, parse_claude_web_exchange,
19        },
20        openai::{
21            OpenAiWebSocketAccumulator, parse_openai_exchange, parse_openai_websocket_message,
22        },
23    },
24    provider::ProviderResolver,
25};
26
27/// Broker Hook that owns the complete Harness usage pipeline.
28pub struct HarnessUsageHook<S>
29where
30    S: AgentEventSink,
31{
32    config: Arc<HarnessUsageHookConfig>,
33    runtime_config: HooksRuntimeConfig,
34    delivery: EventDelivery<S>,
35    correlation: CorrelationRegistry,
36    claude_web_context: ClaudeWebContext,
37    openai_websocket: OpenAiWebSocketAccumulator,
38}
39
40impl<S> HarnessUsageHook<S>
41where
42    S: AgentEventSink,
43{
44    /// Creates the unified pipeline with shared policy and event delivery.
45    #[must_use]
46    pub fn with_runtime_config_and_event_sink(
47        config: HarnessUsageHookConfig,
48        runtime_config: HooksRuntimeConfig,
49        sink: S,
50    ) -> Self {
51        Self {
52            config: Arc::new(config),
53            runtime_config,
54            delivery: EventDelivery::new(sink),
55            correlation: CorrelationRegistry::default(),
56            claude_web_context: ClaudeWebContext::default(),
57            openai_websocket: OpenAiWebSocketAccumulator::default(),
58        }
59    }
60
61    fn http_events(&self, exchange: &HttpExchange) -> Vec<NormalizedUsageEvent> {
62        let policy = self.runtime_config.snapshot();
63        let Some(detection) = HarnessRegistry::detect_http(exchange, &policy.harness_usage.config)
64        else {
65            return Vec::new();
66        };
67        if self.claude_web_context.record_upload(exchange) {
68            return Vec::new();
69        }
70        let Some(protocol) = ProtocolDetector::detect_http(exchange) else {
71            return Vec::new();
72        };
73        let provider = ProviderResolver::resolve(&http_host(exchange), protocol);
74        let content = policy
75            .harness_usage
76            .config
77            .content_for_harness(detection.harness_id.as_str());
78
79        match protocol {
80            LlmProtocol::OpenAi(OpenAiProtocol::Responses | OpenAiProtocol::ChatCompletions) => {
81                let Some(parsed) = parse_openai_exchange(exchange) else {
82                    return Vec::new();
83                };
84                self.normalize(&parsed, &content, &detection, &provider)
85            }
86            LlmProtocol::Anthropic(AnthropicProtocol::Messages) => {
87                let Some(parsed) = parse_anthropic_messages_exchange(exchange) else {
88                    return Vec::new();
89                };
90                self.normalize(&parsed, &content, &detection, &provider)
91            }
92            LlmProtocol::Anthropic(AnthropicProtocol::ClaudeWeb(
93                ClaudeWebProtocol::ConversationCompletion,
94            )) => {
95                let Some(mut parsed) = parse_claude_web_exchange(exchange) else {
96                    return Vec::new();
97                };
98                self.claude_web_context
99                    .attach_referenced_images(&mut parsed);
100                self.normalize(&parsed, &content, &detection, &provider)
101            }
102        }
103    }
104
105    fn normalize<E>(
106        &self,
107        parsed: &E,
108        content: &crate::config::HarnessUsageContentConfig,
109        detection: &HarnessDetection,
110        provider: &LlmProvider,
111    ) -> Vec<NormalizedUsageEvent>
112    where
113        E: NormalizableExchange,
114    {
115        let correlation =
116            self.correlation
117                .assign(detection, parsed.session_id(), parsed.protocol_turn_id());
118        normalize_exchange(
119            &self.config,
120            parsed,
121            content,
122            detection,
123            provider,
124            &correlation,
125        )
126    }
127
128    fn websocket_events(&self, message: &WebSocketMessage) -> Vec<NormalizedUsageEvent> {
129        let policy = self.runtime_config.snapshot();
130        let Some(detection) =
131            HarnessRegistry::detect_websocket(message, &policy.harness_usage.config)
132        else {
133            return Vec::new();
134        };
135        let Some(protocol @ LlmProtocol::OpenAi(OpenAiProtocol::Responses)) =
136            ProtocolDetector::detect_websocket(message)
137        else {
138            return Vec::new();
139        };
140        let parsed = parse_openai_websocket_message(message);
141        let Some(parsed) = self.openai_websocket.push(&message.flow.flow_id, parsed) else {
142            return Vec::new();
143        };
144        let provider = ProviderResolver::resolve(&websocket_host(message), protocol);
145        let content = policy
146            .harness_usage
147            .config
148            .content_for_harness(detection.harness_id.as_str());
149        self.normalize(&parsed, &content, &detection, &provider)
150    }
151}
152
153impl<S> MitmHook for HarnessUsageHook<S>
154where
155    S: AgentEventSink,
156{
157    fn enabled(&self) -> bool {
158        self.runtime_config.snapshot().harness_usage.enabled
159    }
160
161    fn matches(&self, _flow: &FlowContext) -> bool {
162        self.enabled()
163    }
164
165    fn on_http_exchange<'a>(&'a self, exchange: &'a HttpExchange) -> HookFuture<'a> {
166        let events = self.http_events(exchange);
167        Box::pin(async move {
168            self.delivery
169                .deliver(events)
170                .await
171                .map_err(|error| HookError::failed(error.to_string()))
172        })
173    }
174
175    fn on_websocket_message<'a>(&'a self, message: &'a WebSocketMessage) -> HookFuture<'a> {
176        let events = self.websocket_events(message);
177        Box::pin(async move {
178            self.delivery
179                .deliver(events)
180                .await
181                .map_err(|error| HookError::failed(error.to_string()))
182        })
183    }
184}
185
186fn http_host(exchange: &HttpExchange) -> String {
187    request_host(
188        exchange.request.headers(),
189        exchange.request.uri(),
190        exchange.flow.destination_host.as_deref(),
191    )
192}
193
194fn websocket_host(message: &WebSocketMessage) -> String {
195    request_host(
196        message.upgrade_request.headers(),
197        message.upgrade_request.uri(),
198        message.flow.destination_host.as_deref(),
199    )
200}
201
202fn request_host(headers: &HeaderMap, uri: &http::Uri, destination: Option<&str>) -> String {
203    uri.host()
204        .or_else(|| headers.get(HOST).and_then(|value| value.to_str().ok()))
205        .or(destination)
206        .unwrap_or("unknown")
207        .to_owned()
208}
209
210#[cfg(test)]
211mod tests {
212    use std::{
213        convert::Infallible,
214        future::Future,
215        net::SocketAddr,
216        sync::{Arc, Mutex},
217    };
218
219    use abyss_mitm::{
220        CapturedBody, FlowContext, HttpExchange, MitmHook, OriginalDestination, SourceProcess,
221        TransparentProtocol,
222    };
223    use abyss_plugin_protocol::event::AgentEvent;
224    use http::{Request, Response};
225
226    use crate::{
227        AgentEventSink, DeviceIdentity, HarnessUsageHookConfig, HooksConfig, HooksRuntimeConfig,
228    };
229
230    use super::HarnessUsageHook;
231
232    #[derive(Debug, Default)]
233    struct RecordingSink {
234        events: Mutex<Vec<AgentEvent>>,
235    }
236
237    impl AgentEventSink for RecordingSink {
238        type Error = Infallible;
239
240        fn publish(
241            &self,
242            event: AgentEvent,
243        ) -> impl Future<Output = Result<(), Self::Error>> + Send {
244            self.events.lock().unwrap().push(event);
245            std::future::ready(Ok(()))
246        }
247    }
248
249    #[tokio::test]
250    async fn custom_harness_reuses_protocol_and_gateway_provider() {
251        let sink = Arc::new(RecordingSink::default());
252        let hook = hook(Arc::clone(&sink));
253
254        hook.on_http_exchange(&openai_exchange("acme-agent"))
255            .await
256            .unwrap();
257
258        let events = sink.events.lock().unwrap();
259        assert_eq!(events.len(), 2);
260        assert!(events.iter().all(|event| event.agent.name == "acme-agent"));
261        assert!(
262            events
263                .iter()
264                .all(|event| { event.llm.provider.wire_name() == "gateway.example" })
265        );
266        drop(events);
267    }
268
269    #[tokio::test]
270    async fn protocol_shape_without_harness_evidence_produces_no_event() {
271        let sink = Arc::new(RecordingSink::default());
272        let hook = hook(Arc::clone(&sink));
273
274        hook.on_http_exchange(&openai_exchange("ordinary-sdk"))
275            .await
276            .unwrap();
277
278        assert!(sink.events.lock().unwrap().is_empty());
279    }
280
281    #[tokio::test]
282    async fn custom_harness_reuses_anthropic_messages_protocol() {
283        let sink = Arc::new(RecordingSink::default());
284        let hook = hook(Arc::clone(&sink));
285
286        hook.on_http_exchange(&anthropic_exchange("acme-agent"))
287            .await
288            .unwrap();
289
290        let events = sink.events.lock().unwrap();
291        assert_eq!(events.len(), 2);
292        assert!(events.iter().all(|event| event.agent.name == "acme-agent"));
293        assert!(
294            events
295                .iter()
296                .all(|event| event.llm.provider.wire_name() == "anthropic")
297        );
298        drop(events);
299    }
300
301    fn hook(sink: Arc<RecordingSink>) -> HarnessUsageHook<Arc<RecordingSink>> {
302        let config = serde_json::from_value::<HooksConfig>(serde_json::json!({
303            "harness_usage": {"config": {"harnesses": {"acme-agent": {
304                "enabled": true,
305                "matchers": [{"process_names": ["acme-agent"]}]
306            }}}}
307        }))
308        .unwrap();
309        HarnessUsageHook::with_runtime_config_and_event_sink(
310            HarnessUsageHookConfig::new(DeviceIdentity::new()),
311            HooksRuntimeConfig::new(config),
312            sink,
313        )
314    }
315
316    fn openai_exchange(process_name: &str) -> HttpExchange {
317        let flow = FlowContext::from_optional_addrs(
318            None,
319            None,
320            OriginalDestination::from(SocketAddr::from(([127, 0, 0, 1], 443))),
321            TransparentProtocol::PlainHttp,
322            Some(SourceProcess::new(
323                Some(42),
324                Some(process_name.to_owned()),
325                None,
326            )),
327        );
328        let request = Request::builder()
329            .method("POST")
330            .uri("/v1/responses")
331            .header("host", "gateway.example")
332            .header("content-type", "application/json")
333            .body(CapturedBody::from_bytes(
334                serde_json::to_vec(&serde_json::json!({
335                    "model": "gpt-test",
336                    "input": [{"role": "user", "content": "hello"}]
337                }))
338                .unwrap()
339                .into(),
340            ))
341            .unwrap();
342        let response = Response::builder()
343            .header("content-type", "application/json")
344            .body(CapturedBody::from_bytes(
345                serde_json::to_vec(&serde_json::json!({
346                    "id": "resp-test",
347                    "model": "gpt-test",
348                    "output": [{"content": [{"type": "output_text", "text": "hi"}]}],
349                    "usage": {"input_tokens": 1_i32, "output_tokens": 1_i32, "total_tokens": 2_i32}
350                }))
351                .unwrap()
352                .into(),
353            ))
354            .unwrap();
355        HttpExchange::new(flow, request, response)
356    }
357
358    fn anthropic_exchange(process_name: &str) -> HttpExchange {
359        let flow = FlowContext::from_optional_addrs(
360            None,
361            None,
362            OriginalDestination::from(SocketAddr::from(([127, 0, 0, 1], 443))),
363            TransparentProtocol::PlainHttp,
364            Some(SourceProcess::new(
365                Some(42),
366                Some(process_name.to_owned()),
367                None,
368            )),
369        );
370        let request = Request::builder()
371            .method("POST")
372            .uri("/v1/messages")
373            .header("host", "api.anthropic.com")
374            .header("content-type", "application/json")
375            .body(CapturedBody::from_bytes(
376                serde_json::to_vec(&serde_json::json!({
377                    "model": "claude-test",
378                    "messages": [{"role": "user", "content": "hello"}]
379                }))
380                .unwrap()
381                .into(),
382            ))
383            .unwrap();
384        let response = Response::builder()
385            .header("content-type", "application/json")
386            .body(CapturedBody::from_bytes(
387                serde_json::to_vec(&serde_json::json!({
388                    "id": "msg_test",
389                    "model": "claude-test",
390                    "content": [{"type": "text", "text": "hi"}],
391                    "usage": {"input_tokens": 1_i32, "output_tokens": 1_i32}
392                }))
393                .unwrap()
394                .into(),
395            ))
396            .unwrap();
397        HttpExchange::new(flow, request, response)
398    }
399}
abyss_agent_hook/hook.rs

abyss_agent_hook/
hook.rs
