use std::convert::Infallible;
use std::net::SocketAddr;
use axum::Json;
use axum::extract::{ConnectInfo, Extension, State};
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use futures_util::stream::{self, StreamExt};
use log::{debug, info, warn};
use tokio::sync::mpsc;
use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
use tokio_stream::wrappers::{BroadcastStream, ReceiverStream};
use crate::chat_job_queue;
use crate::types::CommandApprovalDecision;
use crate::user_message_file_refs::{
expand_at_file_refs_in_user_message, user_message_has_workspace_file_ref_syntax,
};
use crate::web::http_types::chat::{ApiError, ChatRequestBody};
use crate::clarification_questionnaire::merge_user_text_with_clarification_answers;
use crate::redact;
use crate::web::app_state::{
APPROVAL_SESSION_TTL, ApprovalSessionSlot, ConversationTurnSeed,
purge_expired_approval_sessions,
};
use crate::web::app_state_facets::WebChatTurnAppFacet;
use super::super::parse::{ensure_bearer_api_key_for_chat, normalize_approval_session_id};
use super::builtin_skills::run_web_builtin_command;
use super::turn_build::{
ChatStreamRequestParsed, build_messages_for_turn, parse_chat_stream_request,
parse_last_event_id,
};
pub(super) fn sse_event_with_id(seq: u64, data: String) -> Result<Event, Infallible> {
Ok(Event::default().id(seq.to_string()).data(data))
}
async fn chat_stream_resume_response(
state: &WebChatTurnAppFacet,
headers: &HeaderMap,
p: &ChatStreamRequestParsed,
) -> Result<Option<Response>, (StatusCode, Json<ApiError>)> {
let Some(sr) = p.resume.as_ref() else {
return Ok(None);
};
let job_id = sr.job_id;
if !state.sse_stream_hub.has_job(job_id) {
return Err((
StatusCode::GONE,
Json(ApiError::new(
"STREAM_JOB_GONE",
"流式任务已结束或不在本进程内存中，无法重连".to_string(),
)),
));
}
let after_header = parse_last_event_id(headers).unwrap_or(0);
let after_body = sr.after_seq.unwrap_or(0);
let after_seq = after_header.max(after_body);
let Some(sub) = state.sse_stream_hub.subscribe(job_id) else {
return Err((
StatusCode::GONE,
Json(ApiError::new(
"STREAM_JOB_GONE",
"流式任务已结束或不在本进程内存中，无法重连".to_string(),
)),
));
};
let replay = state
.sse_stream_hub
.replay_after(job_id, after_seq)
.unwrap_or_default();
let max_replayed = replay.last().map(|(s, _)| *s).unwrap_or(after_seq);
info!(
target: "crabmate",
"chat stream 断线重连 job_id={} after_seq={} replayed={}",
job_id,
after_seq,
replay.len()
);
let replay_st = stream::iter(replay).map(|(seq, data)| sse_event_with_id(seq, data));
let live_st = BroadcastStream::new(sub).filter_map(move |item| {
std::future::ready(match item {
Ok((seq, data)) if seq > max_replayed => Some(sse_event_with_id(seq, data)),
Ok(_) => None,
Err(BroadcastStreamRecvError::Lagged(n)) => {
warn!(
target: "crabmate",
"chat stream 重连 broadcast lag job_id={} skipped={}",
job_id,
n
);
None
}
})
});
let merged = replay_st.chain(live_st);
let mut resp = Sse::new(merged)
.keep_alive(KeepAlive::default())
.into_response();
if let Ok(v) = HeaderValue::from_str(&job_id.to_string()) {
resp.headers_mut().insert("x-stream-job-id", v);
}
if let Ok(v) = HeaderValue::from_str(&p.conversation_id) {
resp.headers_mut().insert("x-conversation-id", v);
}
Ok(Some(resp))
}
type ChatStreamHttpError = (StatusCode, Json<ApiError>);
async fn chat_stream_expand_at_files_and_clarify(
state: &WebChatTurnAppFacet,
p: &ChatStreamRequestParsed,
) -> Result<String, ChatStreamHttpError> {
let eff_ws_raw = state.effective_workspace_path().await;
let eff_ws = eff_ws_raw.trim().to_string();
if eff_ws.is_empty() && user_message_has_workspace_file_ref_syntax(&p.user_trim) {
return Err((
StatusCode::BAD_REQUEST,
Json(ApiError::new(
"WORKSPACE_NOT_SET",
"未设置工作区：无法在消息中使用 `file:///` / `@` 引用工作区内文件。请先在侧栏工作区面板选择或提交目录。",
)),
));
}
let work_dir_for_expand = std::path::PathBuf::from(eff_ws_raw);
let msg = {
let cfg = state.cfg.read().await;
expand_at_file_refs_in_user_message(&p.user_trim, work_dir_for_expand.as_path(), &cfg)
.map_err(|e| {
(
StatusCode::BAD_REQUEST,
Json(ApiError::new("INVALID_AT_FILE_REF", e)),
)
})?
};
Ok(merge_user_text_with_clarification_answers(
msg,
p.clarify.clone(),
))
}
async fn chat_stream_build_turn_seed(
state: &WebChatTurnAppFacet,
p: &ChatStreamRequestParsed,
msg: &str,
) -> Result<ConversationTurnSeed, ChatStreamHttpError> {
build_messages_for_turn(
state,
&p.conversation_id,
msg,
&p.image_urls,
p.agent_role.as_deref(),
p.session_mode.as_deref(),
)
.await
.map_err(|e| {
let (code, message) =
if let Some(msg) = crate::config::skills_slash::SkillSlashError::strip_turn_err(&e) {
("SKILL_INVOKE_FAILED", msg.to_string())
} else {
("INVALID_AGENT_ROLE", e)
};
(StatusCode::BAD_REQUEST, Json(ApiError::new(code, message)))
})
}
async fn chat_stream_open_approval_session_if_requested(
state: &WebChatTurnAppFacet,
body: &ChatRequestBody,
) -> Result<(Option<String>, Option<chat_job_queue::WebApprovalSession>), ChatStreamHttpError> {
let approval_session_id = match body.approval_session_id.as_deref() {
Some(v) => Some(normalize_approval_session_id(v).ok_or((
StatusCode::BAD_REQUEST,
Json(ApiError::new(
"INVALID_APPROVAL_SESSION_ID",
"approval_session_id 非法或为空".to_string(),
)),
))?),
None => None,
};
let mut web_approval_session = None;
if let Some(session_id) = approval_session_id.as_ref() {
let (approval_tx, approval_rx) = mpsc::channel::<CommandApprovalDecision>(8);
let mut guard = state.approval_sessions.write().await;
purge_expired_approval_sessions(&mut guard, APPROVAL_SESSION_TTL);
guard.insert(
session_id.clone(),
ApprovalSessionSlot {
tx: approval_tx,
created_at: std::time::Instant::now(),
},
);
web_approval_session = Some(chat_job_queue::WebApprovalSession {
session_id: session_id.clone(),
approval_rx,
});
}
Ok((approval_session_id, web_approval_session))
}
struct ChatStreamEnqueueCtx<'a> {
state: &'a WebChatTurnAppFacet,
headers: &'a HeaderMap,
peer: SocketAddr,
p: &'a ChatStreamRequestParsed,
msg: &'a str,
turn_seed: ConversationTurnSeed,
eff_ws: &'a str,
approval_session_id: Option<String>,
web_approval_session: Option<chat_job_queue::WebApprovalSession>,
request_id: Option<String>,
}
async fn chat_stream_try_enqueue_job(
ctx: ChatStreamEnqueueCtx<'_>,
) -> Result<(u64, mpsc::Receiver<(u64, String)>), ChatStreamHttpError> {
let ChatStreamEnqueueCtx {
state,
headers,
peer,
p,
msg,
turn_seed,
eff_ws,
approval_session_id,
web_approval_session,
request_id,
} = ctx;
let workspace_is_set = state.workspace_is_set().await;
let work_dir_for_job = if eff_ws.is_empty() {
let cfg = state.cfg.read().await;
std::path::PathBuf::from(cfg.command_exec.run_command_working_dir.clone())
} else {
std::path::PathBuf::from(eff_ws.to_string())
};
let job_id = state.chat.chat_queue.next_job_id();
let (tx, rx) = mpsc::channel::<(u64, String)>(1024);
debug!(
target: "crabmate",
"chat stream 请求摘要 job_id={} user_len={} user_preview={}",
job_id,
msg.len(),
redact::preview_chars(msg, redact::MESSAGE_LOG_PREVIEW_CHARS)
);
info!(target: "crabmate", "chat stream 任务入队 job_id={}", job_id);
let request_audit = super::enqueue::web_request_audit_for_turn(state, headers, peer).await;
if let Err(e) = state
.chat
.chat_queue
.try_submit_stream(chat_job_queue::StreamSubmitParams {
envelope: chat_job_queue::WebChatJobEnvelope {
job_id,
queue_deps: state.chat.chat_queue_job_deps.clone(),
app: state.chat_job_app_facet(),
conversation_id: p.conversation_id.clone(),
messages: turn_seed.messages,
expected_revision: turn_seed.expected_revision,
request_agent_role: p.agent_role.clone(),
persisted_active_agent_role: turn_seed.persisted_active_agent_role.clone(),
request_session_mode: p.session_mode.clone(),
persisted_active_session_mode: turn_seed.persisted_active_session_mode.clone(),
work_dir: work_dir_for_job,
workspace_is_set,
temperature_override: p.temperature_override,
seed_override: p.seed_override,
client_sse_protocol: p.client_sse_protocol,
llm_override: p.llm_override.clone(),
executor_llm_override: p.executor_llm_override.clone(),
readonly_tool_ttl_cache_secs: p.readonly_tool_ttl_cache_secs,
request_audit,
request_id,
github_token: crate::web::github_token_request::extract_github_token_from_headers(
headers,
),
},
stream_event_tx: tx,
web_approval_session,
})
{
if let Some(session_id) = approval_session_id {
state.approval_sessions.write().await.remove(&session_id);
}
return Err((
StatusCode::SERVICE_UNAVAILABLE,
Json(ApiError::new(
"QUEUE_FULL",
format!(
"对话任务队列已满（最多等待 {} 个），请稍后重试",
e.max_pending
),
)),
));
}
Ok((job_id, rx))
}
fn chat_stream_sse_response_with_meta(
rx: mpsc::Receiver<(u64, String)>,
p: &ChatStreamRequestParsed,
job_id: u64,
) -> Response {
let stream = ReceiverStream::new(rx).map(|(seq, data)| sse_event_with_id(seq, data));
let mut resp = Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response();
if let Ok(v) = HeaderValue::from_str(&p.conversation_id) {
resp.headers_mut().insert("x-conversation-id", v);
}
if let Ok(v) = HeaderValue::from_str(&job_id.to_string()) {
resp.headers_mut().insert("x-stream-job-id", v);
}
resp
}
fn chat_stream_builtin_sse_response(reply: String, conversation_id: &str) -> Response {
let stream = stream::iter(vec![(1_u64, reply)]).map(|(seq, data)| sse_event_with_id(seq, data));
let mut resp = Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_response();
if let Ok(v) = HeaderValue::from_str(conversation_id) {
resp.headers_mut().insert("x-conversation-id", v);
}
resp
}
pub(crate) async fn chat_stream_handler(
State(state): State<WebChatTurnAppFacet>,
Extension(request_id): Extension<crate::web::request_id::RequestId>,
headers: HeaderMap,
ConnectInfo(peer): ConnectInfo<SocketAddr>,
Json(body): Json<ChatRequestBody>,
) -> Result<Response, (StatusCode, Json<ApiError>)> {
let p = parse_chat_stream_request(&state, &body)?;
ensure_bearer_api_key_for_chat(&state, &p.llm_override).await?;
if let Some(reply) = run_web_builtin_command(&state, p.user_trim.as_str()).await {
return Ok(chat_stream_builtin_sse_response(reply, &p.conversation_id));
}
if let Some(resp) = chat_stream_resume_response(&state, &headers, &p).await? {
return Ok(resp);
}
let msg = chat_stream_expand_at_files_and_clarify(&state, &p).await?;
let turn_seed = chat_stream_build_turn_seed(&state, &p, &msg).await?;
let eff_ws = state.effective_workspace_path().await;
let eff_ws = eff_ws.trim();
let (approval_session_id, web_approval_session) =
chat_stream_open_approval_session_if_requested(&state, &body).await?;
let (job_id, rx) = chat_stream_try_enqueue_job(ChatStreamEnqueueCtx {
state: &state,
headers: &headers,
peer,
p: &p,
msg: &msg,
turn_seed,
eff_ws,
approval_session_id,
web_approval_session,
request_id: Some(request_id.0),
})
.await?;
Ok(chat_stream_sse_response_with_meta(rx, &p, job_id))
}