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 get_palette(&self) -> Option<turbo_vision::core::palette::Palette> {
46 None
47 }
48}
49
50#[derive(Debug)]
55enum Renderer {
56 Tokens(Box<Pipeline>),
57 Trace(TraceRenderer),
58}
59
60impl Renderer {
61 fn feed(&mut self, bytes: &[u8], view: &mut StreamView) {
62 match self {
63 Self::Tokens(p) => p.feed(bytes, view),
64 Self::Trace(t) => t.feed(bytes, view),
65 }
66 }
67
68 fn finish(&mut self, view: &mut StreamView) {
69 match self {
70 Self::Tokens(p) => p.finish(view),
71 Self::Trace(t) => t.finish(view),
72 }
73 }
74}
75
76#[derive(Debug)]
78pub struct SessionState {
79 pub name: String,
80 pub port: u16,
81 pub view: SharedView,
82 pub kind: StreamKind,
83 renderer: Renderer,
84 pub connected: bool,
85}
86
87impl SessionState {
88 #[must_use]
94 pub fn window_title(&self) -> String {
95 let base = format_title(&self.name, self.port);
96 let base = match self.kind {
97 StreamKind::Tokens => base,
98 StreamKind::Trace => format!("[trace] {base}"),
99 };
100 if self.connected {
101 base
102 } else {
103 format!("{base} [disconnected]")
104 }
105 }
106
107 pub fn feed(&mut self, bytes: &[u8]) {
109 let mut view = self.view.borrow_mut();
110 self.renderer.feed(bytes, &mut view);
111 }
112
113 pub fn finish(&mut self) {
115 let mut view = self.view.borrow_mut();
116 self.renderer.finish(&mut view);
117 }
118}
119
120#[must_use]
126pub fn format_title(name: &str, port: u16) -> String {
127 if port == 0 {
128 name.to_string()
129 } else {
130 format!("{name} :{port}")
131 }
132}
133
134#[derive(Debug, Default)]
136pub struct Sessions {
137 inner: HashMap<SessionId, SessionState>,
138}
139
140impl Sessions {
141 pub fn insert(
142 &mut self,
143 id: SessionId,
144 name: String,
145 port: u16,
146 kind: StreamKind,
147 view: SharedView,
148 opts: RenderOptions,
149 ) {
150 let renderer = match kind {
151 StreamKind::Tokens => Renderer::Tokens(Box::new(Pipeline::new(opts))),
152 StreamKind::Trace => Renderer::Trace(TraceRenderer::new()),
153 };
154 self.inner.insert(
155 id,
156 SessionState {
157 name,
158 port,
159 view,
160 kind,
161 renderer,
162 connected: false,
163 },
164 );
165 }
166
167 pub fn get_mut(&mut self, id: SessionId) -> Option<&mut SessionState> {
168 self.inner.get_mut(&id)
169 }
170
171 pub fn remove(&mut self, id: SessionId) -> Option<SessionState> {
172 self.inner.remove(&id)
173 }
174
175 pub fn feed(&mut self, id: SessionId, data: &[u8]) {
177 if let Some(s) = self.inner.get_mut(&id) {
178 s.feed(data);
179 }
180 }
181
182 pub fn mark_reconnected(&mut self, id: SessionId) {
184 if let Some(s) = self.inner.get_mut(&id) {
185 s.connected = true;
186 s.feed(b"\n-- reconnected --\n");
187 }
188 }
189
190 pub fn mark_attached(&mut self, id: SessionId, reattached: bool) {
196 if reattached {
197 self.mark_reconnected(id);
198 } else if let Some(s) = self.inner.get_mut(&id) {
199 s.connected = true;
200 }
201 }
202
203 pub fn mark_disconnected(&mut self, id: SessionId) {
204 if let Some(s) = self.inner.get_mut(&id) {
205 s.connected = false;
206 s.finish();
207 }
208 }
209
210 pub fn clear(&mut self, id: SessionId) {
212 if let Some(s) = self.inner.get_mut(&id) {
213 s.view.borrow_mut().clear();
214 }
215 }
216
217 #[must_use]
219 pub fn plain_text(&self, id: SessionId) -> Option<String> {
220 self.inner.get(&id).map(|s| s.view.borrow().plain_text())
221 }
222
223 #[must_use]
227 pub fn window_title(&self, id: SessionId) -> Option<String> {
228 self.inner.get(&id).map(SessionState::window_title)
229 }
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235 use trace_stream::render::RenderOptions;
236
237 fn opts() -> RenderOptions {
238 RenderOptions {
239 use_color: true,
240 format_thinking: true,
241 format_markdown: true,
242 }
243 }
244
245 fn view() -> SharedView {
246 Rc::new(RefCell::new(StreamView::new(Rect::new(0, 0, 80, 24))))
247 }
248
249 #[test]
250 fn feed_reaches_the_session_view() {
251 let mut sessions = Sessions::default();
252 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
253 sessions.feed(1, b"hello\n");
254 assert!(sessions.plain_text(1).unwrap().contains("hello"));
255 }
256
257 #[test]
258 fn window_title_reflects_connection_state() {
259 let mut sessions = Sessions::default();
260 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
261 assert_eq!(
262 sessions.window_title(1).unwrap(),
263 "demo :4242 [disconnected]"
264 );
265 sessions.mark_reconnected(1);
266 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
267 sessions.mark_disconnected(1);
268 assert_eq!(
269 sessions.window_title(1).unwrap(),
270 "demo :4242 [disconnected]"
271 );
272 }
273
274 #[test]
275 fn window_title_omits_a_zero_port() {
276 let mut sessions = Sessions::default();
277 sessions.insert(1, "anon-1".into(), 0, StreamKind::Tokens, view(), opts());
278 assert_eq!(sessions.window_title(1).unwrap(), "anon-1 [disconnected]");
279 sessions.mark_reconnected(1);
280 assert_eq!(sessions.window_title(1).unwrap(), "anon-1");
281 }
282
283 #[test]
289 fn mark_attached_first_attach_connects_without_a_rule() {
290 let mut sessions = Sessions::default();
291 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
292 sessions.mark_attached(1, false);
293 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
294 assert!(!sessions.plain_text(1).unwrap().contains("reconnected"));
295 }
296
297 #[test]
300 fn mark_attached_reattach_connects_and_draws_a_rule() {
301 let mut sessions = Sessions::default();
302 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
303 sessions.mark_attached(1, true);
304 assert_eq!(sessions.window_title(1).unwrap(), "demo :4242");
305 assert!(sessions.plain_text(1).unwrap().contains("reconnected"));
306 }
307
308 #[test]
309 fn mark_reconnected_draws_a_horizontal_rule() {
310 let mut sessions = Sessions::default();
311 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
312 sessions.mark_reconnected(1);
313 assert!(sessions.plain_text(1).unwrap().contains("reconnected"));
314 }
315
316 #[test]
317 fn clear_empties_the_scrollback() {
318 let mut sessions = Sessions::default();
319 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
320 sessions.feed(1, b"hello\n");
321 sessions.clear(1);
322 assert_eq!(sessions.plain_text(1).unwrap(), "");
323 }
324
325 #[test]
326 fn remove_drops_the_session() {
327 let mut sessions = Sessions::default();
328 sessions.insert(1, "demo".into(), 4242, StreamKind::Tokens, view(), opts());
329 assert!(sessions.remove(1).is_some());
330 assert!(sessions.plain_text(1).is_none());
331 }
332
333 #[test]
334 fn unknown_id_returns_none_everywhere() {
335 let sessions = Sessions::default();
336 assert!(sessions.plain_text(99).is_none());
337 assert!(sessions.window_title(99).is_none());
338 }
339
340 #[test]
341 fn a_trace_session_renders_through_tracefmt_not_the_pipeline() {
342 let mut sessions = Sessions::default();
343 sessions.insert(1, "myapp".into(), 4242, StreamKind::Trace, view(), opts());
344 sessions.feed(1, b"{\"level\":\"INFO\",\"fields\":{\"message\":\"hi\"}}\n");
345 assert_eq!(sessions.plain_text(1).unwrap(), "INFO hi");
346 }
347
348 #[test]
349 fn a_trace_session_window_title_is_tagged() {
350 let mut sessions = Sessions::default();
351 sessions.insert(1, "myapp".into(), 4242, StreamKind::Trace, view(), opts());
352 assert_eq!(
353 sessions.window_title(1).unwrap(),
354 "[trace] myapp :4242 [disconnected]"
355 );
356 sessions.mark_reconnected(1);
357 assert_eq!(sessions.window_title(1).unwrap(), "[trace] myapp :4242");
358 }
359}