Separate SSE streaming from POST work submission (#7834)

This commit is contained in:
jh-block
2026-03-18 18:00:06 +01:00
committed by GitHub
parent c2d10cea9b
commit 8d72344602
15 changed files with 1830 additions and 204 deletions
+1
View File
@@ -3,6 +3,7 @@ pub mod configuration;
pub mod error;
pub mod openapi;
pub mod routes;
pub mod session_event_bus;
pub mod state;
pub mod tls;
pub mod tunnel;
+1
View File
@@ -4,6 +4,7 @@ mod error;
mod logging;
mod openapi;
mod routes;
mod session_event_bus;
mod state;
mod tunnel;
+6
View File
@@ -435,6 +435,9 @@ derive_utoipa!(Icon as IconSchema);
super::routes::agent::update_session,
super::routes::action_required::confirm_tool_action,
super::routes::reply::reply,
super::routes::session_events::session_events,
super::routes::session_events::session_reply,
super::routes::session_events::session_cancel,
super::routes::session::list_sessions,
super::routes::session::search_sessions,
super::routes::session::get_session,
@@ -523,6 +526,9 @@ derive_utoipa!(Icon as IconSchema);
goose::prompt_template::Template,
super::routes::action_required::ConfirmToolActionRequest,
super::routes::reply::ChatRequest,
super::routes::session_events::SessionReplyRequest,
super::routes::session_events::SessionReplyResponse,
super::routes::session_events::CancelRequest,
super::routes::session::ImportSessionRequest,
super::routes::session::SessionListResponse,
super::routes::session::UpdateSessionNameRequest,
+2
View File
@@ -14,6 +14,7 @@ pub mod reply;
pub mod sampling;
pub mod schedule;
pub mod session;
pub mod session_events;
pub mod setup;
pub mod status;
pub mod telemetry;
@@ -44,5 +45,6 @@ pub fn configure(state: Arc<crate::state::AppState>, secret_key: String) -> Rout
.merge(gateway::routes(state.clone()))
.merge(mcp_ui_proxy::routes(secret_key.clone()))
.merge(mcp_app_proxy::routes(secret_key))
.merge(session_events::routes(state.clone()))
.merge(sampling::routes(state))
}
+8 -3
View File
@@ -29,7 +29,7 @@ use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) {
pub fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) {
match content {
MessageContent::ToolRequest(tool_request) => {
if let Ok(tool_call) = &tool_request.tool_call {
@@ -123,7 +123,7 @@ impl IntoResponse for SseResponse {
}
}
#[derive(Debug, Serialize, utoipa::ToSchema)]
#[derive(Debug, Clone, Serialize, utoipa::ToSchema)]
#[serde(tag = "type")]
pub enum MessageEvent {
Message {
@@ -149,10 +149,15 @@ pub enum MessageEvent {
UpdateConversation {
conversation: Conversation,
},
/// Sent at the start of an SSE stream to inform the client about
/// in-flight requests it can reattach to.
ActiveRequests {
request_ids: Vec<String>,
},
Ping,
}
async fn get_token_state(session_manager: &SessionManager, session_id: &str) -> TokenState {
pub async fn get_token_state(session_manager: &SessionManager, session_id: &str) -> TokenState {
session_manager
.get_session(session_id, false)
.await
@@ -296,6 +296,13 @@ async fn delete_session(
}
})?;
// Cancel any in-flight replies before dropping the bus, so spawned
// agent tasks stop consuming tokens for a deleted session.
if let Some(bus) = state.get_event_bus(&session_id).await {
bus.cancel_all_requests().await;
}
state.remove_event_bus(&session_id).await;
Ok(StatusCode::OK)
}
@@ -0,0 +1,615 @@
use crate::routes::errors::ErrorResponse;
use crate::routes::reply::{get_token_state, track_tool_telemetry, MessageEvent};
use crate::state::AppState;
use axum::{
extract::{DefaultBodyLimit, Path, State},
http::{self, HeaderMap},
response::IntoResponse,
routing::{get, post},
Json, Router,
};
use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::agents::{AgentEvent, SessionConfig};
use goose::conversation::message::Message;
use goose::conversation::Conversation;
use serde::{Deserialize, Serialize};
use std::{
convert::Infallible,
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
// ── Request / Response types ────────────────────────────────────────────
#[derive(Debug, Deserialize, Serialize, utoipa::ToSchema)]
pub struct SessionReplyRequest {
/// Client-generated UUIDv7 identifying this request.
pub request_id: String,
pub user_message: Message,
#[serde(default)]
pub override_conversation: Option<Vec<Message>>,
pub recipe_name: Option<String>,
pub recipe_version: Option<String>,
}
#[derive(Debug, Serialize, utoipa::ToSchema)]
pub struct SessionReplyResponse {
pub request_id: String,
}
#[derive(Debug, Deserialize, Serialize, utoipa::ToSchema)]
pub struct CancelRequest {
pub request_id: String,
}
// ── SSE Event Stream Response ───────────────────────────────────────────
/// An SSE response that includes `id:` lines for Last-Event-ID reconnection.
pub struct SseEventStream {
rx: ReceiverStream<String>,
}
impl SseEventStream {
fn new(rx: ReceiverStream<String>) -> Self {
Self { rx }
}
}
impl Stream for SseEventStream {
type Item = Result<Bytes, Infallible>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.rx)
.poll_next(cx)
.map(|opt| opt.map(|s| Ok(Bytes::from(s))))
}
}
impl IntoResponse for SseEventStream {
fn into_response(self) -> axum::response::Response {
let body = axum::body::Body::from_stream(self);
http::Response::builder()
.header("Content-Type", "text/event-stream")
.header("Cache-Control", "no-cache")
.header("Connection", "keep-alive")
.body(body)
.unwrap()
}
}
// ── Helpers ─────────────────────────────────────────────────────────────
fn format_sse_event(seq: u64, json: &str) -> String {
format!("id: {}\ndata: {}\n\n", seq, json)
}
fn serialize_session_event(seq: u64, request_id: Option<&str>, event: &MessageEvent) -> String {
// Build JSON payload: { request_id?: string, ...event_fields }
// We flatten request_id into the event JSON.
let mut event_json = serde_json::to_value(event).unwrap_or_else(
|e| serde_json::json!({"type": "Error", "error": format!("Serialization error: {}", e)}),
);
if let Some(rid) = request_id {
if let serde_json::Value::Object(ref mut map) = event_json {
// Always insert chat_request_id for routing (the chat UUID that
// the frontend registered its listener under).
map.insert(
"chat_request_id".to_string(),
serde_json::Value::String(rid.to_string()),
);
// Also set request_id if the event doesn't already carry one
// (e.g. Notification events have their own request_id for tool-call matching)
map.entry("request_id")
.or_insert_with(|| serde_json::Value::String(rid.to_string()));
}
}
let json_str = serde_json::to_string(&event_json).unwrap_or_default();
format_sse_event(seq, &json_str)
}
// ── GET /sessions/{id}/events ───────────────────────────────────────────
#[utoipa::path(
get,
path = "/sessions/{id}/events",
params(
("id" = String, Path, description = "Session ID"),
),
responses(
(status = 200, description = "SSE event stream",
body = MessageEvent,
content_type = "text/event-stream"),
(status = 404, description = "Session not found"),
)
)]
pub async fn session_events(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
headers: HeaderMap,
) -> Result<SseEventStream, axum::http::StatusCode> {
// Validate the session exists before creating an event bus.
state
.session_manager()
.get_session(&session_id, false)
.await
.map_err(|_| axum::http::StatusCode::NOT_FOUND)?;
let last_event_id: Option<u64> = headers
.get("Last-Event-ID")
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse().ok());
let bus = state.get_or_create_event_bus(&session_id).await;
let (replay, replay_max_seq, mut live_rx) = match bus.subscribe(last_event_id).await {
Ok(result) => result,
Err(_) => {
// Client's Last-Event-ID has been evicted from the replay buffer.
// Send a single error event so the client knows to reload.
let (tx, rx) = mpsc::channel::<String>(1);
let stream = ReceiverStream::new(rx);
let seq = 0;
let error_event = MessageEvent::Error {
error: "Client too far behind — reload conversation".to_string(),
};
let frame = serialize_session_event(seq, None, &error_event);
tokio::spawn(async move {
let _ = tx.send(frame).await;
});
return Ok(SseEventStream::new(stream));
}
};
let (tx, rx) = mpsc::channel::<String>(256);
let stream = ReceiverStream::new(rx);
let task_bus = bus.clone();
tokio::spawn(async move {
let bus = task_bus;
// Notify the client about any in-flight requests BEFORE replay
// so it can register event handlers before replayed events arrive.
// Emitted without an SSE `id:` field so it doesn't regress the
// client's Last-Event-ID cursor.
let active_ids = bus.active_request_ids().await;
if !active_ids.is_empty() {
let event = MessageEvent::ActiveRequests {
request_ids: active_ids,
};
let json_str = serde_json::to_string(&serde_json::to_value(&event).unwrap_or_default())
.unwrap_or_default();
let frame = format!("data: {}\n\n", json_str);
if tx.send(frame).await.is_err() {
return;
}
}
// Send replayed events
for event in &replay {
let frame =
serialize_session_event(event.seq, event.request_id.as_deref(), &event.event);
if tx.send(frame).await.is_err() {
return;
}
}
// Send live events + heartbeat pings
let mut heartbeat_interval = tokio::time::interval(Duration::from_millis(500));
// Heartbeat uses a local counter — not stored in the replay buffer
let mut heartbeat_seq = 0u64;
loop {
tokio::select! {
_ = heartbeat_interval.tick() => {
// Send heartbeat directly without publishing to the bus,
// so pings don't evict real events from the replay buffer.
// Use a comment-style SSE id so it won't interfere with Last-Event-ID.
let frame = format!(": ping {}\n\n", heartbeat_seq);
heartbeat_seq += 1;
if tx.send(frame).await.is_err() {
return;
}
}
result = live_rx.recv() => {
match result {
Ok(event) => {
// Skip events already covered by replay to avoid duplicates
// at the replay/live handoff boundary.
if event.seq <= replay_max_seq {
continue;
}
let frame = serialize_session_event(
event.seq,
event.request_id.as_deref(),
&event.event,
);
if tx.send(frame).await.is_err() {
return;
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!("SSE subscriber lagged by {} events, closing stream so client reconnects with Last-Event-ID", n);
// Close the stream so the client reconnects and
// replays missed events from the buffer.
return;
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
return;
}
}
}
}
}
});
Ok(SseEventStream::new(stream))
}
// ── POST /sessions/{id}/reply ───────────────────────────────────────────
#[utoipa::path(
post,
path = "/sessions/{id}/reply",
params(
("id" = String, Path, description = "Session ID"),
),
request_body = SessionReplyRequest,
responses(
(status = 200, description = "Request accepted",
body = SessionReplyResponse),
(status = 400, description = "Invalid request"),
(status = 404, description = "Session not found"),
(status = 424, description = "Agent not initialized"),
(status = 500, description = "Internal server error"),
)
)]
pub async fn session_reply(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(request): Json<SessionReplyRequest>,
) -> Result<Json<SessionReplyResponse>, ErrorResponse> {
let request_id = request.request_id.clone();
// Validate request_id is a valid UUID
if uuid::Uuid::parse_str(&request_id).is_err() {
return Err(ErrorResponse::bad_request(
"request_id must be a valid UUID",
));
}
// Validate session exists before allocating a bus/registering work
state
.session_manager()
.get_session(&session_id, false)
.await
.map_err(|_| ErrorResponse::not_found(format!("Session {} not found", session_id)))?;
let session_start = std::time::Instant::now();
tracing::info!(
monotonic_counter.goose.session_starts = 1,
session_type = "app",
interface = "ui",
"Session started"
);
if let Some(recipe_name) = request.recipe_name.clone() {
if state.mark_recipe_run_if_absent(&session_id).await {
let recipe_version = request
.recipe_version
.clone()
.unwrap_or_else(|| "unknown".to_string());
tracing::info!(
monotonic_counter.goose.recipe_runs = 1,
recipe_name = %recipe_name,
recipe_version = %recipe_version,
session_type = "app",
interface = "ui",
"Recipe execution started"
);
}
}
let bus = state.get_or_create_event_bus(&session_id).await;
let cancel_token = bus.register_request(request_id.clone()).await;
let user_message = request.user_message;
let override_conversation = request.override_conversation;
let task_state = state.clone();
let task_session_id = session_id.clone();
let task_request_id = request_id.clone();
let task_cancel = cancel_token.clone();
let task_bus = bus.clone();
drop(tokio::spawn(async move {
let publish = |rid: Option<String>, event: MessageEvent| {
let bus = task_bus.clone();
async move {
bus.publish(rid, event).await;
}
};
let agent = match task_state.get_agent(task_session_id.clone()).await {
Ok(agent) => agent,
Err(e) => {
tracing::error!("Failed to get session agent: {}", e);
publish(
Some(task_request_id.clone()),
MessageEvent::Error {
error: format!("Failed to get session agent: {}", e),
},
)
.await;
task_bus.cleanup_request(&task_request_id).await;
return;
}
};
let session = match task_state
.session_manager()
.get_session(&task_session_id, true)
.await
{
Ok(metadata) => metadata,
Err(e) => {
tracing::error!("Failed to read session for {}: {}", task_session_id, e);
publish(
Some(task_request_id.clone()),
MessageEvent::Error {
error: format!("Failed to read session: {}", e),
},
)
.await;
task_bus.cleanup_request(&task_request_id).await;
return;
}
};
let session_config = SessionConfig {
id: task_session_id.clone(),
schedule_id: session.schedule_id.clone(),
max_turns: None,
retry_config: None,
};
let mut all_messages = match override_conversation {
Some(history) => {
let conv = Conversation::new_unvalidated(history);
if let Err(e) = task_state
.session_manager()
.replace_conversation(&task_session_id, &conv)
.await
{
tracing::warn!(
"Failed to replace session conversation for {}: {}",
task_session_id,
e
);
}
conv
}
None => session.conversation.unwrap_or_default(),
};
all_messages.push(user_message.clone());
let mut stream = match agent
.reply(
user_message.clone(),
session_config,
Some(task_cancel.clone()),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
publish(
Some(task_request_id.clone()),
MessageEvent::Error {
error: e.to_string(),
},
)
.await;
task_bus.cleanup_request(&task_request_id).await;
return;
}
};
loop {
tokio::select! {
_ = task_cancel.cancelled() => {
tracing::info!("Agent task cancelled for request {}", task_request_id);
break;
}
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(AgentEvent::Message(message)))) => {
for content in &message.content {
track_tool_telemetry(content, all_messages.messages());
}
all_messages.push(message.clone());
let token_state = get_token_state(
task_state.session_manager(),
&task_session_id,
)
.await;
publish(
Some(task_request_id.clone()),
MessageEvent::Message {
message,
token_state,
},
)
.await;
}
Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => {
all_messages = new_messages.clone();
publish(
Some(task_request_id.clone()),
MessageEvent::UpdateConversation {
conversation: new_messages,
},
)
.await;
}
Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
publish(
Some(task_request_id.clone()),
MessageEvent::ModelChange { model, mode },
)
.await;
}
Ok(Some(Ok(AgentEvent::McpNotification((notification_request_id, n))))) => {
publish(
Some(task_request_id.clone()),
MessageEvent::Notification {
request_id: notification_request_id,
message: n,
},
)
.await;
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
publish(
Some(task_request_id.clone()),
MessageEvent::Error {
error: e.to_string(),
},
)
.await;
break;
}
Ok(None) => {
break;
}
Err(_) => {
// Timeout — check if the bus still has subscribers
continue;
}
}
}
}
}
// Telemetry
let session_duration = session_start.elapsed();
if let Ok(session) = task_state
.session_manager()
.get_session(&task_session_id, true)
.await
{
let total_tokens = session.total_tokens.unwrap_or(0);
tracing::info!(
monotonic_counter.goose.session_completions = 1,
session_type = "app",
interface = "ui",
exit_type = "normal",
duration_ms = session_duration.as_millis() as u64,
total_tokens = total_tokens,
message_count = session.message_count,
"Session completed"
);
tracing::info!(
monotonic_counter.goose.session_duration_ms = session_duration.as_millis() as u64,
session_type = "app",
interface = "ui",
"Session duration"
);
if total_tokens > 0 {
tracing::info!(
monotonic_counter.goose.session_tokens = total_tokens,
session_type = "app",
interface = "ui",
"Session tokens"
);
}
} else {
tracing::info!(
monotonic_counter.goose.session_completions = 1,
session_type = "app",
interface = "ui",
exit_type = "normal",
duration_ms = session_duration.as_millis() as u64,
total_tokens = 0u64,
message_count = all_messages.len(),
"Session completed"
);
tracing::info!(
monotonic_counter.goose.session_duration_ms = session_duration.as_millis() as u64,
session_type = "app",
interface = "ui",
"Session duration"
);
}
let final_token_state =
get_token_state(task_state.session_manager(), &task_session_id).await;
publish(
Some(task_request_id.clone()),
MessageEvent::Finish {
reason: "stop".to_string(),
token_state: final_token_state,
},
)
.await;
task_bus.cleanup_request(&task_request_id).await;
}));
Ok(Json(SessionReplyResponse { request_id }))
}
// ── POST /sessions/{id}/cancel ──────────────────────────────────────────
#[utoipa::path(
post,
path = "/sessions/{id}/cancel",
params(
("id" = String, Path, description = "Session ID"),
),
request_body = CancelRequest,
responses(
(status = 200, description = "Cancellation accepted"),
)
)]
pub async fn session_cancel(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(request): Json<CancelRequest>,
) -> axum::http::StatusCode {
let bus = match state.get_event_bus(&session_id).await {
Some(bus) => bus,
None => return axum::http::StatusCode::NOT_FOUND,
};
bus.cancel_request(&request.request_id).await;
axum::http::StatusCode::OK
}
// ── Route registration ──────────────────────────────────────────────────
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/sessions/{id}/events", get(session_events))
.route(
"/sessions/{id}/reply",
post(session_reply).layer(DefaultBodyLimit::max(50 * 1024 * 1024)),
)
.route("/sessions/{id}/cancel", post(session_cancel))
.with_state(state)
}
@@ -0,0 +1,259 @@
use crate::routes::reply::MessageEvent;
use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{broadcast, Mutex};
use tokio_util::sync::CancellationToken;
const BROADCAST_CAPACITY: usize = 256;
const REPLAY_BUFFER_CAPACITY: usize = 512;
/// Error returned by [`SessionEventBus::subscribe`].
#[derive(Debug)]
pub enum SubscribeError {
/// The client's `Last-Event-ID` has been evicted from the replay buffer,
/// so events have been irrecoverably lost.
ClientTooFarBehind,
}
#[derive(Clone, Debug)]
pub struct SessionEvent {
/// Monotonic sequence number, written as SSE `id:` frame (not in JSON payload).
pub seq: u64,
/// None for Ping events, Some for events associated with a specific request.
pub request_id: Option<String>,
/// The event payload.
pub event: MessageEvent,
}
pub struct SessionEventBus {
tx: broadcast::Sender<SessionEvent>,
buffer: Mutex<VecDeque<SessionEvent>>,
next_seq: AtomicU64,
active_requests: Mutex<HashMap<String, CancellationToken>>,
}
impl SessionEventBus {
pub fn new() -> Self {
let (tx, _) = broadcast::channel(BROADCAST_CAPACITY);
Self {
tx,
buffer: Mutex::new(VecDeque::with_capacity(REPLAY_BUFFER_CAPACITY)),
next_seq: AtomicU64::new(1),
active_requests: Mutex::new(HashMap::new()),
}
}
/// Publish an event to the bus. Assigns a monotonic sequence number.
///
/// The sequence ID is assigned under the buffer lock so that concurrent
/// callers cannot reorder events (i.e. seq=2 published before seq=1).
pub async fn publish(&self, request_id: Option<String>, event: MessageEvent) -> u64 {
let session_event = {
let mut buf = self.buffer.lock().await;
let seq = self.next_seq.fetch_add(1, Ordering::Relaxed);
let session_event = SessionEvent {
seq,
request_id,
event,
};
buf.push_back(session_event.clone());
while buf.len() > REPLAY_BUFFER_CAPACITY {
buf.pop_front();
}
session_event
};
// Send on broadcast channel (ignore error if no subscribers)
let _ = self.tx.send(session_event.clone());
session_event.seq
}
/// Subscribe to live events. If `last_event_id` is provided, replay buffered
/// events with seq > last_event_id. Returns (replay_events, replay_max_seq, live_receiver).
///
/// Returns `Err(SubscribeError::ClientTooFarBehind)` when `last_event_id`
/// refers to an event that has already been evicted from the replay buffer,
/// meaning the client has irrecoverably missed events.
///
/// The live receiver is created *before* snapshotting the buffer so that
/// no event can fall into the gap between the two steps. The caller must
/// skip live events with `seq <= replay_max_seq` to deduplicate.
pub async fn subscribe(
&self,
last_event_id: Option<u64>,
) -> Result<(Vec<SessionEvent>, u64, broadcast::Receiver<SessionEvent>), SubscribeError> {
// Subscribe first so that any event published while we hold the
// buffer lock is guaranteed to appear in `rx` (possibly duplicating
// a replay entry). The caller deduplicates via replay_max_seq.
let rx = self.tx.subscribe();
let (replay, replay_max_seq) = {
let buf = self.buffer.lock().await;
let buf_max = buf.back().map(|e| e.seq).unwrap_or(0);
let buf_min = buf.front().map(|e| e.seq).unwrap_or(0);
let last_id = last_event_id.unwrap_or(0);
// If the client sent a Last-Event-ID that has been evicted from
// the buffer, they have irrecoverably missed events.
if last_id > 0 && buf_min > 0 && last_id < buf_min {
return Err(SubscribeError::ClientTooFarBehind);
}
// Clamp to the actual buffer max so a stale Last-Event-ID
// (e.g. from before a server restart) doesn't suppress live events.
let events: Vec<_> = buf.iter().filter(|e| e.seq > last_id).cloned().collect();
let max_seq = events.last().map(|e| e.seq).unwrap_or(last_id.min(buf_max));
(events, max_seq)
};
Ok((replay, replay_max_seq, rx))
}
/// Return the IDs of all currently active (in-flight) requests.
pub async fn active_request_ids(&self) -> Vec<String> {
let requests = self.active_requests.lock().await;
requests.keys().cloned().collect()
}
/// Register a new request and return its cancellation token.
pub async fn register_request(&self, request_id: String) -> CancellationToken {
let token = CancellationToken::new();
let mut requests = self.active_requests.lock().await;
requests.insert(request_id, token.clone());
token
}
/// Cancel a specific request by request_id.
pub async fn cancel_request(&self, request_id: &str) -> bool {
let requests = self.active_requests.lock().await;
if let Some(token) = requests.get(request_id) {
token.cancel();
true
} else {
false
}
}
/// Cancel all active requests (e.g. when deleting a session).
pub async fn cancel_all_requests(&self) {
let requests = self.active_requests.lock().await;
for token in requests.values() {
token.cancel();
}
}
/// Remove the cancellation token for a completed request.
pub async fn cleanup_request(&self, request_id: &str) {
let mut requests = self.active_requests.lock().await;
requests.remove(request_id);
}
}
impl Default for SessionEventBus {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use goose::conversation::message::TokenState;
#[tokio::test]
async fn test_publish_and_subscribe() {
let bus = SessionEventBus::new();
// Publish some events
bus.publish(Some("req-1".to_string()), MessageEvent::Ping)
.await;
bus.publish(
Some("req-1".to_string()),
MessageEvent::Finish {
reason: "stop".to_string(),
token_state: TokenState::default(),
},
)
.await;
// Subscribe with replay
let (replay, replay_max_seq, _rx) = bus.subscribe(Some(0)).await.unwrap();
assert_eq!(replay.len(), 2);
assert_eq!(replay[0].seq, 1);
assert_eq!(replay[1].seq, 2);
assert_eq!(replay_max_seq, 2);
}
#[tokio::test]
async fn test_subscribe_with_last_event_id() {
let bus = SessionEventBus::new();
bus.publish(None, MessageEvent::Ping).await;
bus.publish(None, MessageEvent::Ping).await;
bus.publish(None, MessageEvent::Ping).await;
// Only get events after seq 2
let (replay, replay_max_seq, _rx) = bus.subscribe(Some(2)).await.unwrap();
assert_eq!(replay.len(), 1);
assert_eq!(replay[0].seq, 3);
assert_eq!(replay_max_seq, 3);
}
#[tokio::test]
async fn test_subscribe_without_last_event_id_replays_all() {
let bus = SessionEventBus::new();
bus.publish(None, MessageEvent::Ping).await;
bus.publish(None, MessageEvent::Ping).await;
// First connect (no Last-Event-ID) should replay all buffered events
let (replay, replay_max_seq, _rx) = bus.subscribe(None).await.unwrap();
assert_eq!(replay.len(), 2);
assert_eq!(replay_max_seq, 2);
}
#[tokio::test]
async fn test_subscribe_with_stale_last_event_id() {
let bus = SessionEventBus::new();
// Buffer has seq 1..3, but client sends Last-Event-ID: 9999
bus.publish(None, MessageEvent::Ping).await;
bus.publish(None, MessageEvent::Ping).await;
bus.publish(None, MessageEvent::Ping).await;
let (replay, replay_max_seq, _rx) = bus.subscribe(Some(9999)).await.unwrap();
// No replay events (all are below 9999)
assert_eq!(replay.len(), 0);
// replay_max_seq should be clamped to buf_max (3), not 9999
assert_eq!(replay_max_seq, 3);
}
#[tokio::test]
async fn test_cancel_request() {
let bus = SessionEventBus::new();
let token = bus.register_request("req-1".to_string()).await;
assert!(!token.is_cancelled());
let cancelled = bus.cancel_request("req-1").await;
assert!(cancelled);
assert!(token.is_cancelled());
// Non-existent request
let cancelled = bus.cancel_request("req-999").await;
assert!(!cancelled);
}
#[tokio::test]
async fn test_cleanup_request() {
let bus = SessionEventBus::new();
bus.register_request("req-1".to_string()).await;
bus.cleanup_request("req-1").await;
// Should return false since it was cleaned up
let cancelled = bus.cancel_request("req-1").await;
assert!(!cancelled);
}
}
+23
View File
@@ -9,6 +9,7 @@ use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use crate::session_event_bus::SessionEventBus;
use crate::tunnel::TunnelManager;
use goose::agents::ExtensionLoadResult;
use goose::gateway::manager::GatewayManager;
@@ -26,6 +27,7 @@ pub struct AppState {
pub gateway_manager: Arc<GatewayManager>,
pub extension_loading_tasks: ExtensionLoadingTasks,
pub inference_runtime: Arc<InferenceRuntime>,
session_buses: Arc<Mutex<HashMap<String, Arc<SessionEventBus>>>>,
}
impl AppState {
@@ -44,6 +46,7 @@ impl AppState {
gateway_manager,
extension_loading_tasks: Arc::new(Mutex::new(HashMap::new())),
inference_runtime: InferenceRuntime::get_or_init(),
session_buses: Arc::new(Mutex::new(HashMap::new())),
}))
}
@@ -107,6 +110,26 @@ impl AppState {
}
}
pub async fn get_or_create_event_bus(&self, session_id: &str) -> Arc<SessionEventBus> {
let mut buses = self.session_buses.lock().await;
buses
.entry(session_id.to_string())
.or_insert_with(|| Arc::new(SessionEventBus::new()))
.clone()
}
/// Get an existing event bus for a session without creating one.
pub async fn get_event_bus(&self, session_id: &str) -> Option<Arc<SessionEventBus>> {
let buses = self.session_buses.lock().await;
buses.get(session_id).cloned()
}
/// Remove the event bus for a session, freeing its replay buffer.
pub async fn remove_event_bus(&self, session_id: &str) {
let mut buses = self.session_buses.lock().await;
buses.remove(session_id);
}
pub async fn get_agent(&self, session_id: String) -> anyhow::Result<Arc<goose::agents::Agent>> {
self.agent_manager.get_or_create_agent(session_id).await
}