1use std::cell::RefCell;
7use std::collections::HashMap;
8use std::rc::Rc;
9
10use trace_stream::render::RenderOptions;
11use turbo_vision::core::event::Event;
12use turbo_vision::core::geometry::Rect;
13use turbo_vision::terminal::Terminal;
14use turbo_vision::views::view::View;
15
16use crate::pipeline::Pipeline;
17use crate::proto::StreamKind;
18use crate::registry::SessionId;
19use crate::streamview::StreamView;
20use crate::tracefmt::TraceRenderer;
21
22pub type SharedView = Rc<RefCell<StreamView>>;
24
25#[derive(Debug)]
27pub struct SharedStreamView(pub SharedView);
28
29impl View for SharedStreamView {
30 fn bounds(&self) -> Rect {
31 self.0.borrow().bounds()
32 }
33 fn set_bounds(&mut self, bounds: Rect) {
34 self.0.borrow_mut().set_bounds(bounds);
35 }
36 fn draw(&mut self, terminal: &mut Terminal) {
37 self.0.borrow_mut().draw(terminal);
38 }
39 fn handle_event(&mut self, event: &mut Event) {
40 self.0.borrow_mut().handle_event(event);
41 }
42 fn can_focus(&self) -> bool {
43 true
44 }
45 fn grow_mode(&self) -> turbo_vision::core::state::GrowFlags {
50 self.0.borrow().grow_mode()
51 }
52 fn set_grow_mode(&mut self, grow_mode: turbo_vision::core::state::GrowFlags) {
53 self.0.borrow_mut().set_grow_mode(grow_mode);
54 }
55 fn get_palette(&self) -> Option<turbo_vision::core::palette::Palette> {
56 None
57 }
58}
59
60#[derive(Debug)]
65enum Renderer {
66 Tokens(Box<Pipeline>),
67 Trace(TraceRenderer),
68}
69
70impl Renderer {
71 fn feed(&mut self, bytes: &[u8], view: &mut StreamView) {
72 match self {
73 Self::Tokens(p) => p.feed(bytes, view),
74 Self::Trace(t) => t.feed(bytes, view),
75 }
76 }
77
78 fn finish(&mut self, view: &mut StreamView) {
79 match self {
80 Self::Tokens(p) => p.finish(view),
81 Self::Trace(t) => t.finish(view),
82 }
83 }
84}
85
86#[derive(Debug)]
88pub struct SessionState {
89 pub name: String,
90 pub port: u16,
91 pub view: SharedView,
92 pub kind: StreamKind,
93 renderer: Renderer,
94 pub connected: bool,
95}
96
97impl SessionState {
98 #[must_use]
104 pub fn window_title(&self) -> String {
105 let base = format_title(&self.name, self.port);
106 let base = match self.kind {
107 StreamKind::Tokens => base,
108 StreamKind::Trace => format!("[trace] {base}"),
109 };
110 if self.connected {
111 base
112 } else {
113 format!("{base} [disconnected]")
114 }
115 }
116
117 pub fn feed(&mut self, bytes: &[u8]) {
119 let mut view = self.view.borrow_mut();
120 self.renderer.feed(bytes, &mut view);
121 }
122
123 pub fn finish(&mut self) {
125 let mut view = self.view.borrow_mut();
126 self.renderer.finish(&mut view);
127 }
128}
129
130#[must_use]
136pub fn format_title(name: &str, port: u16) -> String {
137 if port == 0 {
138 name.to_string()
139 } else {
140 format!("{name} :{port}")
141 }
142}
143
144#[derive(Debug, Default)]
146pub struct Sessions {
147 inner: HashMap<SessionId, SessionState>,
148}
149
150impl Sessions {
151 pub fn insert(
152 &mut self,
153 id: SessionId,
154 name: String,
155 port: u16,
156 kind: StreamKind,
157 view: SharedView,
158 opts: RenderOptions,
159 ) {
160 let renderer = match kind {
161 StreamKind::Tokens => Renderer::Tokens(Box::new(Pipeline::new(opts))),
162 StreamKind::Trace => Renderer::Trace(TraceRenderer::new()),
163 };
164 self.inner.insert(
165 id,
166 SessionState {
167 name,
168 port,
169 view,
170 kind,
171 renderer,
172 connected: false,
173 },
174 );
175 }
176
177 #[must_use]
181 pub fn is_connected(&self, id: SessionId) -> bool {
182 self.inner.get(&id).is_some_and(|s| s.connected)
183 }
184
185 pub fn get_mut(&mut self, id: SessionId) -> Option<&mut SessionState> {
186 self.inner.get_mut(&id)
187 }
188
189 pub fn remove(&mut self, id: SessionId) -> Option<SessionState> {
190 self.inner.remove(&id)
191 }
192
193 pub fn feed(&mut self, id: SessionId, data: &[u8]) {
195 if let Some(s) = self.inner.get_mut(&id) {
196 s.feed(data);
197 }
198 }
199
200 pub fn mark_reconnected(&mut self, id: SessionId) {
202 if let Some(s) = self.inner.get_mut(&id) {
203 s.connected = true;
204 s.feed(b"\n-- reconnected --\n");
205 }
206 }
207
208 pub fn mark_attached(&mut self, id: SessionId, reattached: bool) {
214 if reattached {
215 self.mark_reconnected(id);
216 } else if let Some(s) = self.inner.get_mut(&id) {
217 s.connected = true;
218 }
219 }
220
221 pub fn mark_disconnected(&mut self, id: SessionId) {
222 if let Some(s) = self.inner.get_mut(&id) {
223 s.connected = false;
224 s.finish();
225 }
226 }
227
228 pub fn clear(&mut self, id: SessionId) {
230 if let Some(s) = self.inner.get_mut(&id) {
231 s.view.borrow_mut().clear();
232 }
233 }
234
235 #[must_use]
237 pub fn plain_text(&self, id: SessionId) -> Option<String> {
238 self.inner.get(&id).map(|s| s.view.borrow().plain_text())
239 }
240
241 pub fn select_all(&mut self, id: SessionId) {
243 if let Some(s) = self.inner.get_mut(&id) {
244 s.view.borrow_mut().select_all();
245 }
246 }
247
248 #[must_use]
250 pub fn selected_text(&self, id: SessionId) -> Option<String> {
251 self.inner
252 .get(&id)
253 .and_then(|s| s.view.borrow().selected_text())
254 }
255
256 #[must_use]
260 pub fn window_title(&self, id: SessionId) -> Option<String> {
261 self.inner.get(&id).map(SessionState::window_title)
262 }
263}
264
265#[cfg(test)]
266mod tests {
267 use super::*;
268 use trace_stream::render::RenderOptions;
269
270 fn opts() -> RenderOptions {
271 RenderOptions {
272 use_color: true,
273 format_thinking: true,
274 format_markdown: true,
275 }
276 }
277
278 fn view() -> SharedView {
279 Rc::new(RefCell::new(StreamView::new(Rect::new(0, 0, 80, 24))))
280 }
281
282 #[test]
283 fn feed_reaches_the_session_view() {
284 let mut sessions = Sessions::default();
285 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
286 sessions.feed(1, b"hello\n");
287 assert!(sessions.plain_text(1).unwrap().contains("hello"));
288 }
289
290 #[test]
291 fn window_title_reflects_connection_state() {
292 let mut sessions = Sessions::default();
293 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
294 assert_eq!(
295 sessions.window_title(1).unwrap(),
296 "demo :4242 [disconnected]"
297 );
298 sessions.mark_reconnected(1);
299 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
300 sessions.mark_disconnected(1);
301 assert_eq!(
302 sessions.window_title(1).unwrap(),
303 "demo :4242 [disconnected]"
304 );
305 }
306
307 #[test]
308 fn window_title_omits_a_zero_port() {
309 let mut sessions = Sessions::default();
310 sessions.insert(1, "anon-1".into(), 0, StreamKind::Tokens, view(), opts());
311 assert_eq!(sessions.window_title(1).unwrap(), "anon-1 [disconnected]");
312 sessions.mark_reconnected(1);
313 assert_eq!(sessions.window_title(1).unwrap(), "anon-1");
314 }
315
316 #[test]
322 fn mark_attached_first_attach_connects_without_a_rule() {
323 let mut sessions = Sessions::default();
324 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
325 sessions.mark_attached(1, false);
326 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
327 assert!(!sessions.plain_text(1).unwrap().contains("reconnected"));
328 }
329
330 #[test]
333 fn mark_attached_reattach_connects_and_draws_a_rule() {
334 let mut sessions = Sessions::default();
335 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
336 sessions.mark_attached(1, true);
337 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
338 assert!(sessions.plain_text(1).unwrap().contains("reconnected"));
339 }
340
341 #[test]
342 fn mark_reconnected_draws_a_horizontal_rule() {
343 let mut sessions = Sessions::default();
344 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
345 sessions.mark_reconnected(1);
346 assert!(sessions.plain_text(1).unwrap().contains("reconnected"));
347 }
348
349 #[test]
350 fn clear_empties_the_scrollback() {
351 let mut sessions = Sessions::default();
352 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
353 sessions.feed(1, b"hello\n");
354 sessions.clear(1);
355 assert_eq!(sessions.plain_text(1).unwrap(), "");
356 }
357
358 #[test]
359 fn remove_drops_the_session() {
360 let mut sessions = Sessions::default();
361 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
362 assert!(sessions.remove(1).is_some());
363 assert!(sessions.plain_text(1).is_none());
364 }
365
366 #[test]
367 fn unknown_id_returns_none_everywhere() {
368 let sessions = Sessions::default();
369 assert!(sessions.plain_text(99).is_none());
370 assert!(sessions.window_title(99).is_none());
371 }
372
373 #[test]
374 fn a_trace_session_renders_through_tracefmt_not_the_pipeline() {
375 let mut sessions = Sessions::default();
376 sessions.insert(1, "myapp".into(), 4242, StreamKind::Trace, view(), opts());
377 sessions.feed(1, b"{\"level\":\"INFO\",\"fields\":{\"message\":\"hi\"}}\n");
378 assert_eq!(sessions.plain_text(1).unwrap(), "INFO hi");
379 }
380
381 #[test]
382 fn a_trace_session_window_title_is_tagged() {
383 let mut sessions = Sessions::default();
384 sessions.insert(1, "myapp".into(), 4242, StreamKind::Trace, view(), opts());
385 assert_eq!(
386 sessions.window_title(1).unwrap(),
387 "[trace] myapp :4242 [disconnected]"
388 );
389 sessions.mark_reconnected(1);
390 assert_eq!(sessions.window_title(1).unwrap(), "[trace] myapp :4242");
391 }
392}