use std::sync::Arc;
use agent_client_protocol::schema::v1::StopReason;
use axum::{
extract::{
Path, State, WebSocketUpgrade,
ws::{Message, WebSocket},
},
http::HeaderMap,
response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tracing::{debug, info};
use crate::AppState;
use crate::acp::chat_persistence;
use crate::acp::permission::PermissionRequestEvent;
use crate::acp::terminal::TerminalActivity;
use crate::acp::turn_accumulator::TurnTiming;
use crate::acp::{AcpClient, FileInput, ImageInput, ResourceInput, TurnEndEvent, client};
use crate::api::agents::load_agent;
const MAX_PROMPT_FRAME_BYTES: usize = 12 * 1024 * 1024;
fn mib(bytes: usize) -> f64 {
bytes as f64 / (1024.0 * 1024.0)
}
const MAX_AT_REFERENCES: usize = 8;
const MAX_AT_FILE_BYTES: usize = 64 * 1024;
fn extract_at_paths(text: &str) -> Vec<String> {
static RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
let re = RE.get_or_init(|| regex::Regex::new(r"(?:^|\s)@([^\s@]+)").unwrap());
let mut out: Vec<String> = Vec::new();
for cap in re.captures_iter(text) {
let p = &cap[1];
if !out.iter().any(|e| e == p) {
out.push(p.to_string());
if out.len() >= MAX_AT_REFERENCES {
break;
}
}
}
out
}
async fn resolve_at_references(
db: &sqlx::SqlitePool,
session_id: &str,
text: &str,
) -> Vec<ResourceInput> {
let paths = extract_at_paths(text);
if paths.is_empty() {
return Vec::new();
}
let row: Option<(String,)> = sqlx::query_as("SELECT workspace_path FROM sessions WHERE id = ?")
.bind(session_id)
.fetch_optional(db)
.await
.ok()
.flatten();
let Some((ws_path,)) = row else {
return Vec::new();
};
let base = std::path::PathBuf::from(ws_path);
let mut out = Vec::new();
for rel in paths {
let abs = match crate::fs::sanitize_path(&base, &rel) {
Ok(p) => p,
Err(e) => {
debug!("@ 引用跳过(路径无效): {}: {}", rel, e);
continue;
}
};
if abs.is_dir() {
debug!("@ 引用跳过(是目录): {}", rel);
continue;
}
let content = match tokio::fs::read_to_string(&abs).await {
Ok(c) => c,
Err(e) => {
debug!("@ 引用跳过(读取失败): {}: {}", rel, e);
continue;
}
};
let text = if content.len() > MAX_AT_FILE_BYTES {
let mut end = MAX_AT_FILE_BYTES;
while !content.is_char_boundary(end) {
end -= 1;
}
format!("{}\n… [content truncated at 64KB]", &content[..end])
} else {
content
};
out.push(ResourceInput { uri: format!("file://{}", abs.display()), label: rel, text });
}
out
}
pub async fn ws_acp_handler(
ws: WebSocketUpgrade,
Path(session_id): Path<String>,
State(state): State<AppState>,
headers: HeaderMap,
) -> impl IntoResponse {
if let Some(rejected) = crate::ws::enforce_ws_origin(&headers) {
return rejected;
}
info!("ACP WS upgrade request: session_id={}", session_id);
ws.on_upgrade(move |socket| handle_acp_ws(socket, session_id, state))
}
#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
enum AcpClientMessage {
#[serde(rename = "prompt")]
Prompt {
text: String,
#[serde(default)]
images: Vec<ImageInput>,
#[serde(default)]
files: Vec<FileInput>,
},
#[serde(rename = "cancel")]
Cancel,
#[serde(rename = "load_session")]
LoadSession,
#[serde(rename = "permission_response")]
PermissionResponse { id: String, option_id: String },
#[serde(rename = "set_config_option")]
SetConfigOption { config_id: String, value: String },
}
#[derive(Debug, Serialize)]
#[serde(tag = "type")]
enum AcpServerMessage<'a> {
#[serde(rename = "error")]
Error {
#[serde(skip_serializing_if = "Option::is_none")]
code: Option<&'a str>,
message: &'a str,
},
#[serde(rename = "session_update")]
SessionUpdate {
data: serde_json::Value,
#[serde(skip_serializing_if = "Option::is_none")]
seq: Option<u64>,
},
#[serde(rename = "prompt_done")]
PromptDone {
stop_reason: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
row_id: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
duration: Option<&'a TurnTiming>,
#[serde(skip_serializing_if = "std::ops::Not::not")]
abnormal: bool,
},
#[serde(rename = "prompt_error")]
PromptError { message: &'a str },
#[serde(rename = "system_message")]
SystemMessage {
label: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
detail: Option<&'a serde_json::Value>,
},
#[serde(rename = "terminal_activity")]
TerminalActivity {
id: String,
command: String,
args: Vec<String>,
status: String,
exit_code: Option<u32>,
},
#[serde(rename = "replay_start")]
ReplayStart,
#[serde(rename = "replay_end")]
ReplayEnd,
#[serde(rename = "process_alive")]
ProcessAlive { alive: bool },
#[serde(rename = "permission_request")]
PermissionRequest { id: &'a str, request: &'a serde_json::Value },
#[serde(rename = "permission_resolved")]
PermissionResolved { id: &'a str },
#[serde(rename = "permissions_synced")]
PermissionsSynced,
#[serde(rename = "capabilities")]
Capabilities { image: bool, embedded_context: bool, agent_name: String },
#[serde(rename = "turn_state")]
TurnState { active: bool },
#[serde(rename = "turn_snapshot")]
TurnSnapshot { row_id: String, text: String, blocks: String, seq: u64 },
}
async fn forward_session_update(
tx: &tokio::sync::mpsc::Sender<Message>,
notif: &crate::acp::handler::SeqNotification,
) -> bool {
let data = serde_json::to_value(¬if.notification).unwrap_or_default();
let frame = serde_json::to_string(&AcpServerMessage::SessionUpdate { data, seq: notif.seq })
.unwrap_or_default();
tx.send(Message::Text(frame.into())).await.is_ok()
}
async fn spawn_notify_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<crate::acp::handler::SeqNotification>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(seq_notif) => {
let data = serde_json::to_value(&seq_notif.notification).unwrap_or_default();
let msg = serde_json::to_string(&AcpServerMessage::SessionUpdate {
data,
seq: seq_notif.seq,
})
.unwrap_or_default();
tracing::debug!(
session_id = %session_id,
"ACP notify task: forwarding session_update seq={:?}",
seq_notif.seq
);
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
tracing::warn!(session_id = %session_id, "ACP notify task: notify_tx closed (WS gone), exiting");
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
tracing::warn!(
session_id = %session_id,
"ACP notify task: session_update broadcast closed (client dropped), exiting"
);
break;
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP WS subscriber lagged by {} messages; dropped stale updates",
n
);
}
}
}
});
}
async fn spawn_turn_end_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<TurnEndEvent>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(event) => {
let msg = match &event {
TurnEndEvent::Done { stop_reason, row_id, duration, abnormal } => {
serde_json::to_string(&AcpServerMessage::PromptDone {
stop_reason,
row_id: row_id.as_deref(),
duration: duration.as_ref(),
abnormal: *abnormal,
})
}
TurnEndEvent::Error { message } => {
serde_json::to_string(&AcpServerMessage::PromptError { message })
}
}
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP turn-end subscriber lagged by {} messages; dropped stale events",
n
);
}
}
}
});
}
async fn spawn_crash_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<String>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(reason) => {
let msg =
serde_json::to_string(&AcpServerMessage::PromptError { message: &reason })
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP crash-event subscriber lagged by {} messages; dropped stale events",
n
);
}
}
}
});
}
async fn spawn_system_notice_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<crate::acp::client::SystemNotice>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(notice) => {
let msg = serde_json::to_string(&AcpServerMessage::SystemMessage {
label: ¬ice.label,
detail: notice.detail.as_ref(),
})
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP system-notice subscriber lagged by {} messages; dropped stale notices",
n
);
}
}
}
});
}
async fn spawn_terminal_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<TerminalActivity>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(ev) => {
let (id, command, args, status, exit_code) = match ev {
TerminalActivity::Created { id, command, args } => {
(id, command, args, "created".to_string(), None)
}
TerminalActivity::Exited { id, exit_code } => {
(id, String::new(), Vec::new(), "exited".to_string(), exit_code)
}
};
let msg = serde_json::to_string(&AcpServerMessage::TerminalActivity {
id,
command,
args,
status,
exit_code,
})
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP terminal-event subscriber lagged by {} messages; dropped stale events",
n
);
}
}
}
});
}
async fn spawn_permission_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<PermissionRequestEvent>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(event) => {
let msg = serde_json::to_string(&AcpServerMessage::PermissionRequest {
id: &event.id,
request: &event.request,
})
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP permission-event subscriber lagged by {} messages; dropped stale events",
n
);
}
}
}
});
}
async fn spawn_permission_resolved_task(
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<String>,
notify_tx: tokio::sync::mpsc::Sender<Message>,
) {
let session_id = session_id.to_string();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(id) => {
let msg =
serde_json::to_string(&AcpServerMessage::PermissionResolved { id: &id })
.unwrap_or_default();
if notify_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
"ACP permission-resolved subscriber lagged by {} messages; dropped stale events",
n
);
}
}
}
});
}
async fn query_agent_name(db: &sqlx::SqlitePool, session_id: &str) -> String {
let row: Option<(String,)> =
sqlx::query_as("SELECT a.display_name FROM sessions s JOIN agents a ON a.id = s.agent_id WHERE s.id = ? AND s.runtime_kind = 'acp'")
.bind(session_id)
.fetch_optional(db)
.await
.ok()
.flatten();
row.map(|(name,)| name).unwrap_or_default()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StopEndClass {
Normal,
Cancelled,
Abnormal,
}
fn classify_stop_reason(r: &StopReason) -> StopEndClass {
match r {
StopReason::EndTurn | StopReason::MaxTokens | StopReason::MaxTurnRequests => {
StopEndClass::Normal
}
StopReason::Cancelled => StopEndClass::Cancelled,
StopReason::Refusal => StopEndClass::Abnormal,
_ => StopEndClass::Abnormal,
}
}
fn stop_reason_wire(r: &StopReason) -> String {
serde_json::to_string(r)
.ok()
.map(|s| s.trim_matches('"').to_string())
.unwrap_or_else(|| format!("{:?}", r))
}
const SYSTEM_LABEL_TURN_FAILED_REFUSAL: &str = "system.turnFailed.refusal";
const SYSTEM_LABEL_TURN_FAILED_CANCELLED: &str = "system.turnFailed.cancelled";
const SYSTEM_LABEL_TURN_FAILED_OTHER: &str = "system.turnFailed.other";
struct TurnEndNotice {
label: &'static str,
text: String,
detail: serde_json::Value,
}
fn build_turn_end_notice(r: &StopReason, class: StopEndClass) -> Option<TurnEndNotice> {
let wire = stop_reason_wire(r);
let (label, text) = match (class, r) {
(StopEndClass::Normal, _) => return None,
(StopEndClass::Cancelled, _) => {
(SYSTEM_LABEL_TURN_FAILED_CANCELLED, format!("这一轮已被取消(stopReason={wire})。"))
}
(StopEndClass::Abnormal, StopReason::Refusal) => (
SYSTEM_LABEL_TURN_FAILED_REFUSAL,
format!("这一轮未正常完成:agent 拒绝继续(stopReason={wire})。"),
),
(StopEndClass::Abnormal, _) => (
SYSTEM_LABEL_TURN_FAILED_OTHER,
format!("这一轮以非正常原因结束(stopReason={wire})。"),
),
};
Some(TurnEndNotice {
label,
text,
detail: serde_json::json!({ "stop_reason": wire }),
})
}
async fn persist_and_broadcast_turn_end_notice(
db: &sqlx::SqlitePool,
session_id: &str,
client: &AcpClient,
notice: &TurnEndNotice,
) {
let blocks = serde_json::json!([
{ "type": "system", "label": notice.label, "detail": notice.detail }
])
.to_string();
if let Err(e) =
chat_persistence::insert_message(db, session_id, "system", ¬ice.text, Some(&blocks))
.await
{
tracing::warn!(
session_id = %session_id,
error = %e,
"turn-end: failed to persist abnormal-end system message"
);
}
client.notify_system_message(client::SystemNotice {
label: notice.label.to_string(),
detail: Some(notice.detail.clone()),
});
}
async fn notice_turn_end_if_abnormal(
c: &AcpClient,
db: &sqlx::SqlitePool,
session_id: &str,
stop_reason: &StopReason,
) -> StopEndClass {
let class = classify_stop_reason(stop_reason);
let Some(notice) = build_turn_end_notice(stop_reason, class) else {
return class;
};
if !c.claim_turn_end_notice_for_current_turn() {
return class;
}
persist_and_broadcast_turn_end_notice(db, session_id, c, ¬ice).await;
class
}
async fn dispatch_prompt(
c: Arc<AcpClient>,
db: sqlx::SqlitePool,
session_id: String,
text: String,
images: Vec<ImageInput>,
resources: Vec<ResourceInput>,
files: Vec<FileInput>,
) {
c.mark_prompt_active();
match c.send_prompt(&text, images, resources, files).await {
Ok(resp) => {
c.mark_prompt_idle();
let class = notice_turn_end_if_abnormal(&c, &db, &session_id, &resp.stop_reason).await;
c.notify_turn_end(TurnEndEvent::Done {
stop_reason: format!("{:?}", resp.stop_reason),
row_id: c.turn_row_id(),
duration: c.turn_timing(),
abnormal: class == StopEndClass::Abnormal,
});
}
Err(e) => {
c.mark_prompt_idle();
let message = if c.is_alive() {
format!("{}", e)
} else {
"会话进程已释放,请重新发送以自动恢复连接".to_string()
};
c.notify_turn_end(TurnEndEvent::Error { message });
}
}
}
async fn restore_acp_session(
state: &AppState,
db: &sqlx::SqlitePool,
sid: &str,
client: &mut Option<Arc<AcpClient>>,
notify_tx: &tokio::sync::mpsc::Sender<Message>,
) -> Result<(Arc<AcpClient>, tokio::task::JoinHandle<Result<(), String>>), String> {
let row: Option<(String, String, String)> = sqlx::query_as(
"SELECT agent_id, acp_session_id, workspace_path FROM sessions WHERE id = ? AND runtime_kind = 'acp'",
)
.bind(sid)
.fetch_optional(db)
.await
.ok()
.flatten();
let Some((agent_id, acp_sid, ws_path)) = row else {
return Err("session row not found or not ACP".to_string());
};
let Some(agent) = load_agent(db, &agent_id).await else {
return Err("agent config not found".to_string());
};
let cwd = std::path::PathBuf::from(&ws_path);
let agent_display_name = agent.display_name.clone();
let new_client =
AcpClient::spawn_and_load(agent, cwd.clone(), acp_sid.clone(), &state.api_keys)
.await
.map_err(|e| format!("failed to spawn agent: {}", e))?;
let new_client = Arc::new(new_client);
if !new_client.supports_load_session() {
new_client.shutdown().await;
return Err("agent does not support session/load".to_string());
}
if let Some(old) = state.acp_supervisor.dispose(sid).await {
old.shutdown().await;
}
state.acp_supervisor.insert(sid.to_string(), new_client.clone()).await;
new_client.attach_persistence(db.clone(), sid.to_string());
new_client.attach_config_prefs(db.clone(), sid.to_string(), agent_id).await;
let perm_rx = new_client.permission_subscribe();
spawn_permission_task(sid, perm_rx, notify_tx.clone()).await;
spawn_permission_resolved_task(
sid,
new_client.permission_resolved_subscribe(),
notify_tx.clone(),
)
.await;
let synced = serde_json::to_string(&AcpServerMessage::PermissionsSynced).unwrap_or_default();
let _ = notify_tx.send(Message::Text(synced.into())).await;
let crash_rx = new_client.crash_subscribe();
spawn_crash_task(sid, crash_rx, notify_tx.clone()).await;
let system_notice_rx = new_client.system_notice_subscribe();
spawn_system_notice_task(sid, system_notice_rx, notify_tx.clone()).await;
let turn_end_rx = new_client.turn_end_subscribe();
spawn_turn_end_task(sid, turn_end_rx, notify_tx.clone()).await;
let term_rx = new_client.terminal_event_subscribe();
spawn_terminal_task(sid, term_rx, notify_tx.clone()).await;
*client = Some(new_client.clone());
let cap_msg = serde_json::to_string(&AcpServerMessage::Capabilities {
image: new_client.supports_image(),
embedded_context: new_client.supports_embedded_context(),
agent_name: agent_display_name,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(cap_msg.into())).await;
let replay_msg = serde_json::to_string(&AcpServerMessage::ReplayStart).unwrap_or_default();
let _ = notify_tx.send(Message::Text(replay_msg.into())).await;
let task_client = new_client.clone();
let tx = notify_tx.clone();
let replay_sid = acp_sid.clone();
let replay_cwd = cwd.clone();
let restore_state = state.clone();
let restore_sid = sid.to_string();
let handle = tokio::spawn(async move {
let mut replay_rx = task_client.session_update_subscribe();
let load_fut = task_client.load_session(&replay_sid, replay_cwd);
tokio::pin!(load_fut);
let result = loop {
tokio::select! {
r = &mut load_fut => break r,
recved = replay_rx.recv() => match recved {
Ok(notif) => {
let _ = forward_session_update(&tx, ¬if).await;
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(session_id = %restore_sid, "ACP replay subscriber lagged by {} messages; dropped frames", n);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
break (&mut load_fut).await;
}
},
}
};
let load_err: Option<String> =
result.as_ref().err().map(|e| format!("session/load failed: {}", e));
tracing::info!(session_id = %restore_sid, "ACP replay task: load_session done, ok={}", load_err.is_none());
if load_err.is_none() {
task_client.restore_config_prefs().await;
}
tracing::info!(session_id = %restore_sid, "ACP replay task: restore_config_prefs done");
loop {
match replay_rx.try_recv() {
Ok(notif) => {
if !forward_session_update(&tx, ¬if).await {
break;
}
}
Err(tokio::sync::broadcast::error::TryRecvError::Lagged(n)) => {
tracing::warn!(session_id = %restore_sid, "ACP replay drain lagged by {} messages; dropped frames", n);
}
Err(_) => break, }
}
let msg = match &load_err {
None => serde_json::to_string(&AcpServerMessage::ReplayEnd).unwrap_or_default(),
Some(e) => serde_json::to_string(&AcpServerMessage::Error {
code: Some("load_failed"),
message: e,
})
.unwrap_or_default(),
};
let _ = tx.send(Message::Text(msg.into())).await;
tracing::info!(session_id = %restore_sid, "ACP replay task: replay_end sent, moving replay_rx into notify task");
spawn_notify_task(&restore_sid, replay_rx, tx.clone()).await;
tracing::info!(session_id = %restore_sid, "ACP replay task: notify task spawned, replay task finishing");
match load_err {
None => Ok(()),
Some(e) => {
if let Some(cur) = restore_state.acp_supervisor.get(&restore_sid).await
&& Arc::ptr_eq(&cur, &task_client)
{
let _ = restore_state.acp_supervisor.dispose(&restore_sid).await;
task_client.shutdown().await;
}
Err(e)
}
}
});
Ok((new_client, handle))
}
async fn handle_acp_ws(socket: WebSocket, session_id: String, state: AppState) {
let (mut ws_tx, mut ws_rx) = socket.split();
let (notify_tx, mut notify_rx) = tokio::sync::mpsc::channel::<Message>(64);
let mut client: Option<Arc<AcpClient>> = match state.acp_supervisor.get(&session_id).await {
Some(c) => {
info!("ACP WS connected: session_id={} (supervisor hit)", session_id);
let rx = c.session_update_subscribe();
let turn_end_rx = c.turn_end_subscribe();
let snap = c.turn_snapshot();
let ts_msg =
serde_json::to_string(&AcpServerMessage::TurnState { active: snap.active })
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(ts_msg.into())).await;
if snap.active
&& let Some(row_id) = snap.row_id
{
let snap_msg = serde_json::to_string(&AcpServerMessage::TurnSnapshot {
row_id,
text: snap.text,
blocks: snap.blocks,
seq: snap.seq,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(snap_msg.into())).await;
}
spawn_notify_task(&session_id, rx, notify_tx.clone()).await;
let perm_rx = c.permission_subscribe();
spawn_permission_task(&session_id, perm_rx, notify_tx.clone()).await;
spawn_permission_resolved_task(
&session_id,
c.permission_resolved_subscribe(),
notify_tx.clone(),
)
.await;
for event in c.pending_permission_events().await {
let msg = serde_json::to_string(&AcpServerMessage::PermissionRequest {
id: &event.id,
request: &event.request,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
let synced =
serde_json::to_string(&AcpServerMessage::PermissionsSynced).unwrap_or_default();
let _ = notify_tx.send(Message::Text(synced.into())).await;
let crash_rx = c.crash_subscribe();
spawn_crash_task(&session_id, crash_rx, notify_tx.clone()).await;
let system_notice_rx = c.system_notice_subscribe();
spawn_system_notice_task(&session_id, system_notice_rx, notify_tx.clone()).await;
spawn_turn_end_task(&session_id, turn_end_rx, notify_tx.clone()).await;
let term_rx = c.terminal_event_subscribe();
spawn_terminal_task(&session_id, term_rx, notify_tx.clone()).await;
if let Some(notif) = c.initial_config_notification() {
let data = serde_json::to_value(¬if).unwrap_or_default();
let msg =
serde_json::to_string(&AcpServerMessage::SessionUpdate { data, seq: None })
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
if let Some(notif) = c.initial_commands_notification() {
let data = serde_json::to_value(¬if).unwrap_or_default();
let msg =
serde_json::to_string(&AcpServerMessage::SessionUpdate { data, seq: None })
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
let agent_name = query_agent_name(&state.db, &session_id).await;
let msg = serde_json::to_string(&AcpServerMessage::Capabilities {
image: c.supports_image(),
embedded_context: c.supports_embedded_context(),
agent_name,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
Some(c)
}
None => {
info!("ACP WS: session_id={} not in supervisor, keeping alive for restore", session_id);
let row_exists: bool =
sqlx::query_scalar("SELECT COUNT(*) > 0 FROM sessions WHERE id = ?")
.bind(&session_id)
.fetch_one(&state.db)
.await
.unwrap_or(true); if !row_exists {
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: Some("session_not_found"),
message: "ACP session not found",
})
.unwrap();
let _ = ws_tx.send(Message::Text(msg.into())).await;
}
None
}
};
let mut proc_rx = state.acp_supervisor.process_event_subscribe();
let _ = notify_tx
.send(Message::Text(
serde_json::to_string(&AcpServerMessage::ProcessAlive { alive: client.is_some() })
.unwrap_or_default()
.into(),
))
.await;
let db = state.db.clone();
let sid = session_id.clone();
loop {
tokio::select! {
msg = ws_rx.next() => {
match msg {
Some(Ok(Message::Text(text))) => {
if text.len() > MAX_PROMPT_FRAME_BYTES {
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: Some("message_too_large"),
message: &format!(
"message too large ({:.1} MiB); the WebSocket pipe holds {:.1} MiB",
mib(text.len()),
mib(MAX_PROMPT_FRAME_BYTES),
),
})
.unwrap_or_default();
let _ = ws_tx.send(Message::Text(msg.into())).await;
continue;
}
match serde_json::from_str::<AcpClientMessage>(&text) {
Ok(AcpClientMessage::Prompt { text: prompt_text, images, files }) => {
let blocks_json = if images.is_empty() && files.is_empty() {
None
} else {
let mut arr = Vec::new();
if !prompt_text.is_empty() {
arr.push(serde_json::json!({
"type": "text", "text": prompt_text,
}));
}
for img in &images {
let (data, mime_type) = match &img.thumb {
Some(t) => (t.data.as_str(), t.mime_type.as_str()),
None => (img.data.as_str(), img.mime_type.as_str()),
};
arr.push(serde_json::json!({
"type": "image",
"mimeType": mime_type,
"data": data,
}));
}
for f in &files {
arr.push(serde_json::json!({
"type": "file",
"name": f.name,
"mimeType": f.mime_type,
"size": f.size,
}));
}
serde_json::to_string(&arr).ok()
};
let _ = chat_persistence::insert_message(
&db, &sid, "user", &prompt_text, blocks_json.as_deref(),
).await;
let resources = resolve_at_references(&db, &sid, &prompt_text).await;
let (live, replay_handle) =
match client.as_ref().and_then(|c| c.is_alive().then(|| c.clone())) {
Some(c) => (c, None),
None => match restore_acp_session(
&state, &db, &sid, &mut client, ¬ify_tx,
)
.await
{
Ok((c, handle)) => (c, Some(handle)),
Err(e) => {
let msg = serde_json::to_string(
&AcpServerMessage::Error {
code: Some("session_released"),
message: &e,
},
)
.unwrap_or_default();
let _ = ws_tx.send(Message::Text(msg.into())).await;
continue;
}
},
};
if !images.is_empty() && !live.supports_image() {
let msg = serde_json::to_string(&AcpServerMessage::PromptError {
message: "agent does not support image input",
}).unwrap_or_default();
let _ = ws_tx.send(Message::Text(msg.into())).await;
continue;
}
if !files.is_empty() && !live.supports_embedded_context() {
let msg = serde_json::to_string(&AcpServerMessage::PromptError {
message: "agent does not support file attachments (embedded context)",
}).unwrap_or_default();
let _ = ws_tx.send(Message::Text(msg.into())).await;
continue;
}
if let Some(handle) = replay_handle {
let c = live;
let tx = notify_tx.clone();
let db = db.clone();
let sid = sid.clone();
tokio::spawn(async move {
match handle.await {
Ok(Ok(())) => {
dispatch_prompt(
c,
db,
sid,
prompt_text,
images,
resources,
files,
)
.await;
}
Ok(Err(_)) => {} Err(e) => {
let err_msg = format!("replay task failed: {}", e);
let msg = serde_json::to_string(
&AcpServerMessage::Error {
code: Some("load_failed"),
message: &err_msg,
},
)
.unwrap_or_default();
let _ = tx.send(Message::Text(msg.into())).await;
}
}
});
} else {
tokio::spawn(dispatch_prompt(
live,
db.clone(),
sid.clone(),
prompt_text,
images,
resources,
files,
));
}
}
Ok(AcpClientMessage::Cancel) => {
if let Some(ref c) = client {
c.spawn_cancel_turn_fallback();
if let Err(e) = c.cancel() {
let err_msg = format!("取消 agent 失败: {}", e);
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: Some("cancel_failed"),
message: &err_msg,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
}
}
Ok(AcpClientMessage::LoadSession) => {
if let Err(e) = restore_acp_session(
&state, &db, &sid, &mut client, ¬ify_tx,
)
.await
{
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: None,
message: &e,
})
.unwrap_or_default();
let _ = ws_tx.send(Message::Text(msg.into())).await;
}
}
Ok(AcpClientMessage::PermissionResponse { id, option_id }) => {
let resolved = match client {
Some(ref c) => c.resolve_permission(&id, &option_id).await,
None => false,
};
if !resolved {
let msg = serde_json::to_string(
&AcpServerMessage::PermissionResolved { id: &id },
)
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
}
Ok(AcpClientMessage::SetConfigOption { config_id, value }) => {
if let Some(ref c) = client
&& let Err(e) = c.set_config_option(&config_id, &value).await {
let err_msg = format!("配置项 {} 设置失败: {}", config_id, e);
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: Some("config_option_failed"),
message: &err_msg,
})
.unwrap_or_default();
let _ = notify_tx.send(Message::Text(msg.into())).await;
}
}
Err(e) => {
let err_msg = format!("invalid message: {}", e);
let msg = serde_json::to_string(&AcpServerMessage::Error {
code: None,
message: &err_msg,
})
.unwrap_or_default();
if ws_tx.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
}
}
Some(Ok(Message::Close(_))) | None => break,
other => {
tracing::warn!(
session_id = %session_id,
?other,
"received unsupported websocket frame (non-text/non-close); ignoring"
);
}
}
}
msg = notify_rx.recv() => {
match msg {
Some(ws_msg) => {
if ws_tx.send(ws_msg).await.is_err() {
break;
}
}
None => break,
}
}
msg = proc_rx.recv() => {
match msg {
Ok(evt) if evt.session_id == session_id => {
let frame = serde_json::to_string(&AcpServerMessage::ProcessAlive {
alive: evt.alive,
})
.unwrap_or_default();
if notify_tx.send(Message::Text(frame.into())).await.is_err() {
break;
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(
session_id = %session_id,
skipped = n,
"process-alive channel lagged; process events may be stale"
);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
tracing::debug!(session_id = %session_id, "process-alive channel closed");
}
Ok(_) => {}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::{
AcpClientMessage, AcpServerMessage, SYSTEM_LABEL_TURN_FAILED_CANCELLED,
SYSTEM_LABEL_TURN_FAILED_OTHER, SYSTEM_LABEL_TURN_FAILED_REFUSAL, StopEndClass,
build_turn_end_notice, classify_stop_reason, extract_at_paths, stop_reason_wire,
};
use agent_client_protocol::schema::v1::StopReason;
#[test]
fn prompt_frame_without_files_deserializes_empty_files() {
let msg: AcpClientMessage =
serde_json::from_str(r#"{"type":"prompt","text":"hi"}"#).unwrap();
let AcpClientMessage::Prompt { text, images, files } = msg else {
panic!("expected prompt");
};
assert_eq!(text, "hi");
assert!(images.is_empty());
assert!(files.is_empty());
}
#[test]
fn prompt_frame_with_files_parses_metadata() {
let raw = r#"{"type":"prompt","text":"see","files":[{"name":"a.pdf","mime_type":"application/pdf","size":12,"data":"AAA"}]}"#;
let msg: AcpClientMessage = serde_json::from_str(raw).unwrap();
let AcpClientMessage::Prompt { files, .. } = msg else {
panic!("expected prompt");
};
assert_eq!(files.len(), 1);
assert_eq!(files[0].name, "a.pdf");
assert_eq!(files[0].mime_type, "application/pdf");
assert_eq!(files[0].size, 12);
assert_eq!(files[0].data, "AAA");
}
#[test]
fn prompt_frame_file_size_and_data_default_gracefully() {
let raw = r#"{"type":"prompt","text":"","files":[{"name":"a.bin","mime_type":"application/octet-stream","data":""}]}"#;
let msg: AcpClientMessage = serde_json::from_str(raw).unwrap();
let AcpClientMessage::Prompt { files, .. } = msg else {
panic!("expected prompt");
};
assert_eq!(files[0].size, 0);
assert_eq!(files[0].data, "");
}
#[test]
fn extracts_basic_paths() {
assert_eq!(
extract_at_paths("看看 @src/main.rs 和 @README.md 的内容"),
vec!["src/main.rs", "README.md"]
);
}
#[test]
fn extracts_at_start_and_after_newline() {
assert_eq!(extract_at_paths("@a.txt first"), vec!["a.txt"]);
assert_eq!(extract_at_paths("line1\n@b.txt"), vec!["b.txt"]);
}
#[test]
fn dedupes_preserving_order() {
assert_eq!(extract_at_paths("@a.rs @b.rs @a.rs"), vec!["a.rs", "b.rs"]);
}
#[test]
fn caps_at_max_references() {
let text = (1..=10).map(|i| format!("@f{}.rs", i)).collect::<Vec<_>>().join(" ");
assert_eq!(extract_at_paths(&text).len(), super::MAX_AT_REFERENCES);
}
#[test]
fn ignores_email_like_tokens() {
assert_eq!(extract_at_paths("联系 user@example.com 谢谢"), Vec::<String>::new());
}
#[test]
fn ignores_bare_at() {
assert_eq!(extract_at_paths("@ 后面是空格"), Vec::<String>::new());
assert_eq!(extract_at_paths("no refs here"), Vec::<String>::new());
}
#[test]
fn stop_reason_whitelist_classification() {
for r in [StopReason::EndTurn, StopReason::MaxTokens, StopReason::MaxTurnRequests] {
assert_eq!(classify_stop_reason(&r), StopEndClass::Normal, "{r:?} 应判为正常");
}
assert_eq!(classify_stop_reason(&StopReason::Cancelled), StopEndClass::Cancelled);
assert_eq!(classify_stop_reason(&StopReason::Refusal), StopEndClass::Abnormal);
}
#[test]
fn unknown_variant_is_unreachable_today_documented() {
let wires: Vec<String> = [
StopReason::EndTurn,
StopReason::MaxTokens,
StopReason::MaxTurnRequests,
StopReason::Refusal,
StopReason::Cancelled,
]
.iter()
.map(stop_reason_wire)
.collect();
assert_eq!(wires, ["end_turn", "max_tokens", "max_turn_requests", "refusal", "cancelled"]);
}
#[test]
fn stop_reason_wire_uses_protocol_snake_case() {
assert_eq!(stop_reason_wire(&StopReason::EndTurn), "end_turn");
assert_eq!(stop_reason_wire(&StopReason::MaxTokens), "max_tokens");
assert_eq!(stop_reason_wire(&StopReason::Refusal), "refusal");
assert_eq!(stop_reason_wire(&StopReason::Cancelled), "cancelled");
assert_ne!(stop_reason_wire(&StopReason::EndTurn), "EndTurn");
}
#[test]
fn turn_end_notice_carries_i18n_key_text_and_wire_detail() {
let cases = [
(StopReason::Refusal, StopEndClass::Abnormal, SYSTEM_LABEL_TURN_FAILED_REFUSAL),
(StopReason::Cancelled, StopEndClass::Cancelled, SYSTEM_LABEL_TURN_FAILED_CANCELLED),
];
for (reason, class, label) in cases {
let notice = build_turn_end_notice(&reason, class).expect("非正常/取消必须留痕");
assert_eq!(notice.label, label);
assert!(!notice.text.is_empty(), "text 是中文兜底,不能空");
assert_eq!(notice.detail["stop_reason"], stop_reason_wire(&reason));
assert!(notice.text.contains(stop_reason_wire(&reason).as_str()), "{}", notice.text);
}
}
#[test]
fn normal_stop_reason_writes_no_notice() {
for r in [StopReason::EndTurn, StopReason::MaxTokens, StopReason::MaxTurnRequests] {
assert!(
build_turn_end_notice(&r, StopEndClass::Normal).is_none(),
"{r:?} 是白名单正常值,不得留痕"
);
}
}
#[test]
fn cancelled_notice_is_not_the_refusal_copy() {
let cancelled =
build_turn_end_notice(&StopReason::Cancelled, StopEndClass::Cancelled).unwrap();
let refusal = build_turn_end_notice(&StopReason::Refusal, StopEndClass::Abnormal).unwrap();
assert_ne!(cancelled.label, refusal.label);
assert_ne!(cancelled.text, refusal.text);
assert!(cancelled.text.contains("取消"), "{}", cancelled.text);
assert!(refusal.text.contains("拒绝"), "{}", refusal.text);
}
#[test]
fn turn_failed_i18n_keys_are_the_agreed_literals() {
assert_eq!(SYSTEM_LABEL_TURN_FAILED_REFUSAL, "system.turnFailed.refusal");
assert_eq!(SYSTEM_LABEL_TURN_FAILED_CANCELLED, "system.turnFailed.cancelled");
assert_eq!(SYSTEM_LABEL_TURN_FAILED_OTHER, "system.turnFailed.other");
}
#[test]
fn other_notice_branch_is_unreachable_today_documented() {
let refusal = build_turn_end_notice(&StopReason::Refusal, StopEndClass::Abnormal).unwrap();
assert_eq!(refusal.label, SYSTEM_LABEL_TURN_FAILED_REFUSAL);
let cancelled =
build_turn_end_notice(&StopReason::Cancelled, StopEndClass::Cancelled).unwrap();
assert_eq!(cancelled.label, SYSTEM_LABEL_TURN_FAILED_CANCELLED);
assert_ne!(SYSTEM_LABEL_TURN_FAILED_OTHER, SYSTEM_LABEL_TURN_FAILED_REFUSAL);
}
#[test]
fn prompt_done_frame_omits_abnormal_when_false() {
let ok = serde_json::to_string(&AcpServerMessage::PromptDone {
stop_reason: "EndTurn",
row_id: None,
duration: None,
abnormal: false,
})
.unwrap();
assert!(!ok.contains("abnormal"), "false 必须省略字段,得到: {ok}");
let bad = serde_json::to_string(&AcpServerMessage::PromptDone {
stop_reason: "Refusal",
row_id: None,
duration: None,
abnormal: true,
})
.unwrap();
assert!(bad.contains("\"abnormal\":true"), "true 必须显式下发,得到: {bad}");
assert!(bad.contains("\"stop_reason\":\"Refusal\""), "得到: {bad}");
}
#[test]
fn only_abnormal_class_maps_to_frame_abnormal() {
for (class, expected) in [
(StopEndClass::Normal, false),
(StopEndClass::Cancelled, false),
(StopEndClass::Abnormal, true),
] {
assert_eq!(class == StopEndClass::Abnormal, expected, "{class:?} 的映射不符");
}
}
}
#[cfg(all(test, target_os = "linux"))]
mod notice_tests {
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use agent_client_protocol::schema::v1::StopReason;
use sqlx::sqlite::SqlitePoolOptions;
use super::{
SYSTEM_LABEL_TURN_FAILED_REFUSAL, StopEndClass, TurnEndEvent, build_turn_end_notice,
dispatch_prompt, notice_turn_end_if_abnormal, stop_reason_wire,
};
use crate::acp::agent_proc::spawn_test_lock_async;
use crate::acp::chat_persistence::{insert_message, list_messages_page};
use crate::acp::client::AcpClient;
use crate::models::agent::{Agent, AgentEnvVar};
const SESSION: &str = "s-notice";
const AGENT_ID: &str = "a-notice";
const DISPATCH_EVENT_TIMEOUT: Duration = Duration::from_secs(10);
const FAKE_AGENT_SCRIPT: &str = r#"#!/bin/sh
# test-only fake ACP agent:只响应 initialize + session/new 后常驻(stdin 保持
# 打开,`while read` 阻塞),让测试拿到一个活的 AcpClient 即可终止。
while IFS= read -r line; do
case "$line" in
*'"method":"initialize"'*)
id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([0-9a-f-][0-9a-f-]*\)".*/\1/p')
printf '{"jsonrpc":"2.0","id":"%s","result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true,"promptCapabilities":{}}}}\n' "$id"
;;
*'"method":"session/new"'*)
id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([0-9a-f-][0-9a-f-]*\)".*/\1/p')
printf '{"jsonrpc":"2.0","id":"%s","result":{"sessionId":"notice-session"}}\n' "$id"
;;
esac
done
exit 0
"#;
const DISPATCH_FAKE_AGENT_SCRIPT: &str = r#"#!/bin/sh
# test-only fake ACP agent:握手 + 对 session/prompt 推一条正文 chunk 后,
# 回注 env 指定的 stopReason。
while IFS= read -r line; do
case "$line" in
*'"method":"initialize"'*)
id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([0-9a-f-][0-9a-f-]*\)".*/\1/p')
printf '{"jsonrpc":"2.0","id":"%s","result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true,"promptCapabilities":{}}}}\n' "$id"
;;
*'"method":"session/new"'*)
id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([0-9a-f-][0-9a-f-]*\)".*/\1/p')
printf '{"jsonrpc":"2.0","id":"%s","result":{"sessionId":"dispatch-session"}}\n' "$id"
;;
*'"method":"session/prompt"'*)
# 先推正文:让累积器在本 turn 内真的开一行(row_id 有值、定稿会写库),
# 这样才能测到「真有一次 agent 回复且它被定稿落库」,而不是只测一个空
# turn 的收尾。notification 没有 id 字段(JSON-RPC 通知),不回抄。
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"dispatch-session","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"正文一句话"}}}}\n'
id=$(printf '%s' "$line" | sed -n 's/.*"id":"\([0-9a-f-][0-9a-f-]*\)".*/\1/p')
printf '{"jsonrpc":"2.0","id":"%s","result":{"stopReason":"@STOP_REASON@"}}\n' "$id"
;;
esac
done
exit 0
"#;
fn unique_dir(tag: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!(
"omniterm-turn-end-notice-{tag}-{}-{}",
std::process::id(),
chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)
));
std::fs::create_dir_all(&dir).expect("create test dir");
dir
}
fn write_fake_agent(dir: &Path) -> PathBuf {
let script = dir.join("fake-agent.sh");
std::fs::write(&script, FAKE_AGENT_SCRIPT).expect("write fake agent script");
script
}
fn agent_for(script: &Path) -> Agent {
Agent {
id: AGENT_ID.into(),
display_name: "Fake Notice Agent".into(),
command: "/bin/sh".into(),
args: vec![script.to_string_lossy().to_string()],
env: vec![AgentEnvVar { key: "PATH".into(), value: "/usr/bin:/bin".into() }],
npm_package: None,
created_at: String::new(),
updated_at: String::new(),
}
}
async fn fresh_db() -> sqlx::SqlitePool {
let db = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory sqlite pool");
sqlx::migrate!("./migrations").run(&db).await.expect("run migrations");
let now = chrono::Utc::now().to_rfc3339();
sqlx::query(
"INSERT INTO agents (id, display_name, command, args, env, created_at, updated_at) \
VALUES (?, 'Fake', '/bin/sh', '[]', '[]', ?, ?)",
)
.bind(AGENT_ID)
.bind(&now)
.bind(&now)
.execute(&db)
.await
.expect("seed agent");
sqlx::query(
"INSERT INTO projects (id, name, path, created_at) VALUES ('p-n', 'p', '/tmp', ?)",
)
.bind(&now)
.execute(&db)
.await
.expect("seed project");
sqlx::query(
"INSERT INTO sessions (id, project_id, workspace_path, created_at, runtime_kind, agent_id) \
VALUES (?, 'p-n', '/tmp', ?, 'acp', ?)",
)
.bind(SESSION)
.bind(&now)
.bind(AGENT_ID)
.execute(&db)
.await
.expect("seed session");
db
}
async fn rows(db: &sqlx::SqlitePool) -> Vec<(String, String, Option<String>)> {
let page = list_messages_page(db, SESSION, None, 100, usize::MAX)
.await
.expect("read back messages");
page.rows.iter().map(|r| (r.role.clone(), r.text.clone(), r.blocks.clone())).collect()
}
async fn system_count(db: &sqlx::SqlitePool) -> usize {
rows(db).await.iter().filter(|(role, _, _)| role == "system").count()
}
async fn live_client(dir: &Path) -> AcpClient {
let script = write_fake_agent(dir);
let workspace = dir.join("ws");
std::fs::create_dir_all(&workspace).expect("create workspace");
tokio::time::timeout(
Duration::from_secs(15),
AcpClient::spawn_and_connect(agent_for(&script), workspace, &HashMap::new()),
)
.await
.expect("spawn 超时:fake agent 未响应握手")
.expect("spawn_and_connect 失败")
}
async fn dispatch_client(
dir: &Path,
stop_reason: &str,
db: sqlx::SqlitePool,
) -> Arc<AcpClient> {
let script = dir.join("dispatch-agent.sh");
std::fs::write(&script, DISPATCH_FAKE_AGENT_SCRIPT.replace("@STOP_REASON@", stop_reason))
.expect("write dispatch fake agent script");
let workspace = dir.join("ws");
std::fs::create_dir_all(&workspace).expect("create workspace");
let env = vec![
AgentEnvVar { key: "STOP_REASON".into(), value: stop_reason.to_string() },
AgentEnvVar { key: "PATH".into(), value: "/usr/bin:/bin".into() },
];
let agent = Agent {
id: AGENT_ID.into(),
display_name: "Fake Dispatch Agent".into(),
command: "/bin/sh".into(),
args: vec![script.to_string_lossy().to_string()],
env,
npm_package: None,
created_at: String::new(),
updated_at: String::new(),
};
let client = tokio::time::timeout(
Duration::from_secs(15),
AcpClient::spawn_and_connect(agent, workspace, &HashMap::new()),
)
.await
.expect("spawn 超时:fake agent 未响应握手")
.expect("spawn_and_connect 失败");
client.attach_persistence(db, SESSION.to_string());
Arc::new(client)
}
async fn drive_dispatch_prompt(
client: &Arc<AcpClient>,
db: &sqlx::SqlitePool,
text: &str,
) -> Result<TurnEndEvent, String> {
let mut rx = client.turn_end_subscribe();
dispatch_prompt(
Arc::clone(client),
db.clone(),
SESSION.to_string(),
text.to_string(),
vec![],
vec![],
vec![],
)
.await;
match tokio::time::timeout(DISPATCH_EVENT_TIMEOUT, rx.recv()).await {
Ok(Ok(event)) => Ok(event),
Ok(Err(e)) => Err(format!("turn_end 广播接收失败: {e}")),
Err(_) => {
Err(format!("dispatch_prompt 未在 {DISPATCH_EVENT_TIMEOUT:?} 内广播 turn 结束事件"))
}
}
}
#[tokio::test]
async fn normal_stop_reasons_write_no_system_row() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("normal");
let client = live_client(&dir).await;
for reason in [StopReason::EndTurn, StopReason::MaxTokens, StopReason::MaxTurnRequests] {
let class = notice_turn_end_if_abnormal(&client, &db, SESSION, &reason).await;
assert_eq!(class, StopEndClass::Normal, "{reason:?} 应判为 Normal");
}
assert_eq!(system_count(&db).await, 0, "白名单正常值不得留 system 行");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn refusal_writes_exactly_one_row_even_when_repeated_for_same_turn() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("refusal");
let client = live_client(&dir).await;
client.mark_prompt_active();
let generation = client.prompt_generation();
assert_eq!(generation, 1, "首个 prompt 后世代应为 1(构造时为 0)");
let class = notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
assert_eq!(class, StopEndClass::Abnormal, "refusal 必须是 error 语义");
assert_eq!(system_count(&db).await, 1, "refusal 必须留一条痕");
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
assert_eq!(system_count(&db).await, 1, "重复定稿不得重复写(只写一条)");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn cancelled_writes_one_row_with_the_cancelled_copy() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("cancelled");
let client = live_client(&dir).await;
client.mark_prompt_active();
let class =
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Cancelled).await;
assert_eq!(class, StopEndClass::Cancelled, "cancelled 是单独类目");
let rows = rows(&db).await;
assert_eq!(rows.len(), 1, "只写一条");
let (role, text, _blocks) = &rows[0];
assert_eq!(role, "system", "必须是 role='system' 行(D2 载体)");
assert!(text.contains("取消"), "取消文案,得到: {text}");
let notice =
build_turn_end_notice(&StopReason::Cancelled, StopEndClass::Cancelled).unwrap();
assert_eq!(text, ¬ice.text, "落库 text 与文案构造器一致");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn persisted_row_carries_role_label_and_stop_reason_detail() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("rowshape");
let client = live_client(&dir).await;
client.mark_prompt_active();
insert_message(&db, SESSION, "user", "帮我做件事", None).await.expect("seed user row");
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
let rows = rows(&db).await;
assert_eq!(rows.len(), 2, "正文行 + 一条 system 留痕");
assert_eq!(rows[0].0, "user", "正文行在前");
assert_eq!(rows[1].0, "system", "失败提示在后(不得渲染到正文之前)");
let (role, text, blocks) = &rows[1];
assert_eq!(role, "system");
assert!(text.contains("拒绝"), "refusal 文案,得到: {text}");
let blocks = blocks.as_deref().expect("system 行必须带 blocks");
let parsed: serde_json::Value = serde_json::from_str(blocks).expect("blocks 是 JSON 数组");
let block = &parsed[0];
assert_eq!(block["type"], "system", "前端按 type='system' 分发渲染");
assert_eq!(block["label"], "system.turnFailed.refusal", "i18n key 必须在 blocks 里");
assert_eq!(
block["detail"]["stop_reason"], "refusal",
"detail 带协议原文({{reason}} 插值)"
);
assert_eq!(
text,
&build_turn_end_notice(&StopReason::Refusal, StopEndClass::Abnormal).unwrap().text
);
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn next_generation_gets_its_own_notice() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("generation");
let client = live_client(&dir).await;
client.mark_prompt_active();
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
assert_eq!(system_count(&db).await, 1);
client.mark_prompt_active();
assert_eq!(client.prompt_generation(), 2, "世代应递增");
notice_turn_end_if_abnormal(&client, &db, SESSION, &StopReason::Refusal).await;
assert_eq!(system_count(&db).await, 2, "第二个失败 turn 必须也有提示");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
fn assert_refusal_notice_row(role: &str, text: &str, blocks: Option<&str>) {
assert_eq!(role, "system", "必须是 role='system' 行(D2 载体)");
assert!(text.contains("拒绝"), "refusal 文案,得到: {text}");
let blocks = blocks.expect("system 行必须带 blocks");
let parsed: serde_json::Value = serde_json::from_str(blocks).expect("blocks 是 JSON 数组");
let block = &parsed[0];
assert_eq!(block["type"], "system", "前端按 type='system' 分发渲染");
assert_eq!(
block["label"], SYSTEM_LABEL_TURN_FAILED_REFUSAL,
"i18n key 必须在 blocks 里(前端命中才翻译,未命中原样显示)"
);
assert_eq!(
block["detail"]["stop_reason"], "refusal",
"detail 带协议原文({{reason}} 插值)"
);
assert_eq!(
text,
&build_turn_end_notice(&StopReason::Refusal, StopEndClass::Abnormal)
.expect("refusal 必须留痕")
.text
);
}
#[tokio::test]
async fn dispatch_prompt_refusal_persists_notice_and_marks_abnormal() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("dispatch-refusal");
let client =
dispatch_client(&dir, &stop_reason_wire(&StopReason::Refusal), db.clone()).await;
let event = drive_dispatch_prompt(&client, &db, "帮我做件事")
.await
.expect("dispatch_prompt 必须广播 turn 结束事件");
match event {
TurnEndEvent::Done { stop_reason, abnormal, .. } => {
assert_eq!(stop_reason, "Refusal", "stop_reason 用 Debug 形态下发");
assert!(abnormal, "refusal 是 Abnormal,prompt_done.abnormal 必须为 true");
}
other => panic!("refusal 必须走 Done 帧,得到: {other:?}"),
}
let all = rows(&db).await;
let notice: Vec<_> = all.iter().filter(|(role, _, _)| role == "system").collect();
assert_eq!(notice.len(), 1, "refusal 必须只留一条痕,得到 {} 行: {all:?}", all.len());
assert_refusal_notice_row(¬ice[0].0, ¬ice[0].1, notice[0].2.as_deref());
assert_eq!(system_count(&db).await, 1, "system 行必须恰好一条");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn dispatch_prompt_end_turn_writes_no_notice_and_not_abnormal() {
let _guard = spawn_test_lock_async().await;
let db = fresh_db().await;
let dir = unique_dir("dispatch-end-turn");
let client =
dispatch_client(&dir, &stop_reason_wire(&StopReason::EndTurn), db.clone()).await;
let event = drive_dispatch_prompt(&client, &db, "帮我做件事")
.await
.expect("dispatch_prompt 必须广播 turn 结束事件");
match event {
TurnEndEvent::Done { stop_reason, abnormal, .. } => {
assert_eq!(stop_reason, "EndTurn");
assert!(!abnormal, "end_turn 是白名单正常值,不得标成错误");
}
other => panic!("end_turn 必须走 Done 帧,得到: {other:?}"),
}
assert_eq!(system_count(&db).await, 0, "正常结束不得留 system 行(噪音)");
client.shutdown().await;
let _ = std::fs::remove_dir_all(&dir);
}
}