Skip to main content

wesichain_server/
agent_endpoint.rs

1//! Agent streaming endpoint: `POST /agent/chat` → SSE stream.
2
3use std::sync::Arc;
4
5use axum::{
6    extract::State,
7    response::IntoResponse,
8    routing::post,
9    Json, Router,
10};
11use futures::{stream::BoxStream, StreamExt};
12use serde::Deserialize;
13use wesichain_core::{StreamEvent, WesichainError};
14
15use crate::sse::stream_to_sse;
16
17#[derive(Debug, Deserialize)]
18pub struct AgentChatRequest {
19    pub message: String,
20}
21
22/// A factory function that turns a message into a streaming agent response.
23pub type AgentHandler =
24    Arc<dyn Fn(String) -> BoxStream<'static, Result<StreamEvent, WesichainError>> + Send + Sync>;
25
26/// Build a router with `POST /agent/chat` that streams agent events as SSE.
27///
28/// # Example
29/// ```ignore
30/// let handler: AgentHandler = Arc::new(move |msg| {
31///     Box::pin(my_agent.stream(msg))
32/// });
33/// let router = agent_router(handler);
34/// ```
35pub fn agent_router(handler: AgentHandler) -> Router {
36    Router::new()
37        .route("/agent/chat", post(agent_handler))
38        .with_state(handler)
39}
40
41async fn agent_handler(
42    State(handler): State<AgentHandler>,
43    Json(req): Json<AgentChatRequest>,
44) -> impl IntoResponse {
45    let stream = handler(req.message).boxed();
46    stream_to_sse(stream)
47}