1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
//! In-flight assistant turns that outlive a single WebSocket connection.
//!
//! When the client disconnects mid-turn, the Gemini/tool loop keeps running and
//! persists to the DB. A reconnecting socket can [`LiveTurns::subscribe`] to
//! receive subsequent [`StreamEvent`]s.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;
use super::genai::Content;
const EVENT_CAPACITY: usize = 256;
/// Events emitted during a streaming turn (WS layer builds OOB HTML).
#[derive(Debug, Clone)]
pub enum StreamEvent {
UserSaved {
session_id: i64,
user: Content,
/// Set when this prompt became the session title.
title: Option<String>,
},
/// Live stream chunks (UI no longer shows a stream panel).
Partial(Content),
/// Model turn that includes function calls (args shown in transcript).
ToolCall(Content),
Tool(Content),
Final(Content),
}
/// Process-local registry of active turns keyed by session id.
#[derive(Clone, Default)]
pub struct LiveTurns {
inner: Arc<Mutex<HashMap<i64, broadcast::Sender<StreamEvent>>>>,
}
impl LiveTurns {
pub fn new() -> Self {
Self::default()
}
/// Register a new turn publisher (call before spawning the turn task).
pub fn insert(&self, session_id: i64, tx: broadcast::Sender<StreamEvent>) {
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(session_id, tx);
}
pub fn remove(&self, session_id: i64) {
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.remove(&session_id);
}
pub fn contains(&self, session_id: i64) -> bool {
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.contains_key(&session_id)
}
/// Subscribe to live events for an in-flight turn, if any.
pub fn subscribe(&self, session_id: i64) -> Option<broadcast::Receiver<StreamEvent>> {
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&session_id)
.map(|tx| tx.subscribe())
}
}
/// Create a broadcast channel for a new turn.
pub fn new_turn_channel() -> (
broadcast::Sender<StreamEvent>,
broadcast::Receiver<StreamEvent>,
) {
broadcast::channel(EVENT_CAPACITY)
}
/// Emit a stream event; zero receivers (detached client) is normal.
pub fn emit(tx: &broadcast::Sender<StreamEvent>, event: StreamEvent) {
let _ = tx.send(event);
}