cookie_agent_tui/
client.rs

1//! TUI adapter for the shared protocol client.
2
3pub use cookie_agent_server::{
4    ClientDelivery, ClientError, ClientProtocol, validate_websocket_url,
5};
6
7use std::{ops::Deref, sync::Arc};
8
9use cookie_agent_protocol::Transport;
10use cookie_agent_server::{Server, WebSocketTransport, in_process_pair};
11
12/// Display-only protocol client. Raw output is retained by the runtime, not the TUI.
13#[derive(Clone)]
14pub struct Client(cookie_agent_protocol::Client);
15
16impl Client {
17    pub fn connect_stream<T: Transport + 'static>(transport: T) -> Self {
18        Self(cookie_agent_protocol::Client::connect_display_stream(
19            transport,
20        ))
21    }
22
23    pub fn connect_in_process(server: Arc<Server>) -> Self {
24        let (client, service) = in_process_pair(128);
25        tokio::spawn(async move {
26            let _ = server.serve_stream(service).await;
27        });
28        Self::connect_stream(client)
29    }
30
31    pub async fn connect_websocket_with_token(url: &str, token: &str) -> Result<Self, ClientError> {
32        WebSocketTransport::connect_with_token(url, token)
33            .await
34            .map(Self::connect_stream)
35            .map_err(|error| {
36                ClientError::WebSocket(cookie_agent_protocol::diagnostics::sanitize(
37                    &cookie_agent_protocol::diagnostics::error_chain(&error),
38                    4096,
39                ))
40            })
41    }
42}
43
44impl Deref for Client {
45    type Target = cookie_agent_protocol::Client;
46
47    fn deref(&self) -> &Self::Target {
48        &self.0
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55    use crate::state::{DeliveryOutcome, StateStore, ToolStatus};
56    use base64::{Engine as _, engine::general_purpose::STANDARD};
57    use cookie_agent_protocol::{EventPayload, MessageFrame, SessionId, StoredEvent, ToolCallId};
58
59    #[tokio::test]
60    async fn sustained_raw_output_with_fast_tui_consumer() {
61        sustained_raw_output(false).await;
62    }
63
64    #[tokio::test]
65    async fn sustained_raw_output_with_stalled_tui_consumer() {
66        sustained_raw_output(true).await;
67    }
68
69    async fn sustained_raw_output(stalled: bool) {
70        let session_id = SessionId::new_v7();
71        let run_id = cookie_agent_protocol::RunId::new_v7();
72        let call_id = ToolCallId::new_v7();
73        let owner = cookie_agent_protocol::AssistantToolCallRef {
74            model_turn_seq: 1,
75            content_index: 0,
76            model_call_id: cookie_agent_protocol::ModelCallId::new("test").unwrap(),
77            provider_item_id: None,
78        };
79        let events = [
80            EventPayload::ToolCallStarted {
81                start: cookie_agent_protocol::ToolCallStart {
82                    tool_call_id: call_id,
83                    owner: owner.clone(),
84                    output: Default::default(),
85                    presentation: cookie_agent_protocol::ToolCallPresentation {
86                        title: cookie_agent_protocol::SafeDisplayText::new("Test").unwrap(),
87                        primary_argument: None,
88                    },
89                    operation_fingerprint: serde_json::from_value(
90                        serde_json::json!({"digest": "1".repeat(64)}),
91                    )
92                    .unwrap(),
93                },
94            },
95            EventPayload::ToolCallProgress {
96                tool_call_id: call_id,
97                message: cookie_agent_protocol::SafeDisplayText::new("working").unwrap(),
98                display: Some("live display".into()),
99            },
100            EventPayload::ToolCallTerminated {
101                termination: cookie_agent_protocol::ToolCallTermination {
102                    tool_call_id: call_id,
103                    owner,
104                    outcome: cookie_agent_protocol::ToolTerminationOutcome::Completed,
105                    error: None,
106                    result: Some(cookie_agent_protocol::PersistedToolResult {
107                        title: cookie_agent_protocol::SafeDisplayText::new("Test").unwrap(),
108                        output: "authoritative preview".into(),
109                        display: Some("final display".into()),
110                        retained_output: None,
111                        metadata: serde_json::json!({}),
112                        truncation: None,
113                        attachments: Vec::new(),
114                        additional_messages: Vec::new(),
115                    }),
116                },
117            },
118            EventPayload::RunCompleted { final_text: None },
119        ];
120        let (transport, mut server) = in_process_pair(1);
121        let client = Client::connect_stream(transport);
122        let mut receiver = client.subscribe_deliveries().unwrap();
123        let consume = move |mut receiver: tokio::sync::mpsc::UnboundedReceiver<ClientDelivery>| async move {
124            let mut store = StateStore::default();
125            for seq in 1..=4 {
126                let delivery = receiver.recv().await.unwrap();
127                assert!(matches!(delivery, ClientDelivery::Live { .. }));
128                assert_eq!(store.apply_delivery(delivery), DeliveryOutcome::Applied);
129                assert_eq!(store.sessions[&session_id].last_seq, seq);
130                if seq == 2 {
131                    assert_eq!(
132                        store.sessions[&session_id].tools[&call_id].detail,
133                        "live display"
134                    );
135                    assert_eq!(
136                        store.sessions[&session_id].tools[&call_id].status,
137                        ToolStatus::Running
138                    );
139                }
140            }
141            assert!(receiver.is_empty());
142            assert_eq!(
143                store.sessions[&session_id].tools[&call_id].detail,
144                "final display"
145            );
146            assert_eq!(
147                store.sessions[&session_id].tools[&call_id].status,
148                ToolStatus::Completed
149            );
150        };
151        let consumer = if stalled {
152            None
153        } else {
154            let (_, empty) = tokio::sync::mpsc::unbounded_channel();
155            Some(tokio::spawn(consume(std::mem::replace(
156                &mut receiver,
157                empty,
158            ))))
159        };
160        let producer = tokio::spawn(async move {
161            let data = STANDARD.encode(vec![b'x'; 64 * 1024]);
162            for (index, payload) in events.into_iter().enumerate() {
163                if index == 2 {
164                    for chunk in 0..1600 {
165                        let delta = cookie_agent_protocol::OutputDelta {
166                            call_id,
167                            stream: cookie_agent_protocol::OutputStream::Single,
168                            byte_offset: chunk * 64 * 1024,
169                            data: data.clone(),
170                        };
171                        server.send(MessageFrame::Value(serde_json::json!({
172                            "jsonrpc": "2.0", "method": "events.tool_output_delta", "params": delta
173                        }))).await.unwrap();
174                        if chunk % 400 == 0 {
175                            server.send(MessageFrame::Value(serde_json::json!({
176                                "jsonrpc": "2.0", "method": "events.tool_output_snapshot",
177                                "params": {"stream": delta.stream, "snapshot": {
178                                    "call_id": call_id, "start_offset": delta.byte_offset,
179                                    "end_offset": delta.byte_offset + 64 * 1024, "chunks": [delta]
180                                }}
181                            }))).await.unwrap();
182                            server.send(MessageFrame::Value(serde_json::json!({
183                                "jsonrpc": "2.0", "method": "events.tool_output_gap",
184                                "params": {"call_id": call_id, "stream": delta.stream, "next_offset": delta.byte_offset}
185                            }))).await.unwrap();
186                        }
187                    }
188                }
189                let event = StoredEvent {
190                    engine_version: None,
191                    origin: None,
192                    session_id,
193                    run_id: Some(run_id),
194                    seq: index as u64 + 1,
195                    timestamp: jiff::Timestamp::now(),
196                    payload,
197                };
198                server.send(MessageFrame::Value(serde_json::json!({
199                    "jsonrpc": "2.0", "method": "events.subscription",
200                    "params": cookie_agent_protocol::EventSubscriptionMessage::Event { event: Box::new(event) }
201                }))).await.unwrap();
202            }
203            let request = server.recv().await.unwrap().unwrap();
204            let MessageFrame::Value(request) = request else {
205                panic!("expected request")
206            };
207            server
208                .send(MessageFrame::Value(serde_json::json!({
209                    "jsonrpc": "2.0", "id": request["id"], "result": null
210                })))
211                .await
212                .unwrap();
213        });
214        // The response fences all prior deliveries without relying on sleeps.
215        client
216            .call::<_, serde_json::Value>("test.fence", &())
217            .await
218            .unwrap();
219        if let Some(consumer) = consumer {
220            consumer.await.unwrap();
221        } else {
222            assert_eq!(
223                receiver.len(),
224                4,
225                "100 MiB raw output must not enter the TUI queue"
226            );
227            consume(receiver).await;
228        }
229        producer.await.unwrap();
230        client.shutdown();
231    }
232}
cookie_agent_tui/client.rs

cookie_agent_tui/
client.rs
