wip(agent): P0进度快照(会话历史/AskUser/取消)
This commit is contained in:
+179
-3
@@ -24,6 +24,7 @@
|
||||
use crate::auth::jwt::validate_jwt;
|
||||
use crate::services::agent_workbench::{
|
||||
approval::{self, ApprovalOutcome, SettledState},
|
||||
ask::{self, AskOutcome},
|
||||
bash_sandbox::SandboxedBashTool,
|
||||
permission::{classify_tool_risk, decide_permission, PermissionAction, PermissionMode},
|
||||
sandbox,
|
||||
@@ -39,12 +40,14 @@ use axum::{
|
||||
};
|
||||
use futures::stream::Stream;
|
||||
use open_agent_sdk::api::openai::OpenAIProvider;
|
||||
use open_agent_sdk::types::{ContentBlock, PermissionDecision};
|
||||
use open_agent_sdk::{Agent, AgentOptions, LLMProvider, SDKMessage};
|
||||
use open_agent_sdk::types::{ContentBlock, Message, PermissionDecision};
|
||||
use open_agent_sdk::{Agent, AgentOptions, CancellationToken, LLMProvider, SDKMessage};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use std::convert::Infallible;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
|
||||
use tokio_stream::wrappers::UnboundedReceiverStream;
|
||||
@@ -54,6 +57,8 @@ use uuid::Uuid;
|
||||
const MAX_TURNS: u32 = 20;
|
||||
/// How long a tool approval waits for the user before auto-denying.
|
||||
const APPROVAL_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
/// How long an AskUserQuestion waits for the user before giving up.
|
||||
const ASK_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
/// Body limit for a chat request (images/attachments can be large).
|
||||
const AGENT_CHAT_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
@@ -67,6 +72,7 @@ pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/agent/chat", post(chat_handler))
|
||||
.route("/api/agent/approve", post(approve_handler))
|
||||
.route("/api/agent/answer", post(answer_handler))
|
||||
.route("/api/agent/status", get(status_handler))
|
||||
.layer(DefaultBodyLimit::max(AGENT_CHAT_BODY_LIMIT_BYTES))
|
||||
}
|
||||
@@ -92,6 +98,13 @@ struct ApproveRequest {
|
||||
approved: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AnswerRequest {
|
||||
request_id: String,
|
||||
/// The user's free-text reply to an `AskUserQuestion` prompt.
|
||||
answer: String,
|
||||
}
|
||||
|
||||
/// Resolved caller identity for an agent request.
|
||||
#[derive(Debug, Clone)]
|
||||
struct AgentAuth {
|
||||
@@ -237,6 +250,53 @@ fn build_can_use_tool(
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the SDK `ask_fn` for the `AskUserQuestion` tool. When the agent needs
|
||||
/// the user to clarify something, this emits an `ask` SSE event and blocks the
|
||||
/// tool call until the user answers via `POST /api/agent/answer` (or the wait
|
||||
/// times out, which surfaces a cancellation back to the agent). Reuses the same
|
||||
/// SSE round-trip pattern as the approval flow.
|
||||
fn build_ask_fn(
|
||||
sse_tx: UnboundedSender<Result<Event, Infallible>>,
|
||||
) -> open_agent_sdk::tools::askuser::AskUserFn {
|
||||
use open_agent_sdk::tools::askuser::AskUserRequest;
|
||||
Arc::new(move |req: AskUserRequest| {
|
||||
let sse_tx = sse_tx.clone();
|
||||
Box::pin(async move {
|
||||
let request_id = Uuid::new_v4().to_string();
|
||||
// `_guard` removes the registry entry when this future ends (normal
|
||||
// completion or drop on session abort) — no leak.
|
||||
let (entry, _guard) = ask::register(&request_id);
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "ask",
|
||||
"data": {
|
||||
"request_id": request_id,
|
||||
"question": req.question,
|
||||
"options": req.options,
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let outcome = ask::wait_for_answer(&entry, ASK_TIMEOUT).await;
|
||||
// Tell the client the prompt is resolved so it can dismiss the modal.
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "ask_resolved",
|
||||
"data": { "request_id": request_id }
|
||||
}),
|
||||
);
|
||||
match outcome {
|
||||
AskOutcome::Answered(text) => Ok(text),
|
||||
AskOutcome::Cancelled => {
|
||||
Err("User did not answer the question in time".to_string())
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
async fn chat_handler(
|
||||
State(state): State<AppState>,
|
||||
@@ -259,6 +319,18 @@ async fn chat_handler(
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||||
// Refuse to resume a session owned by another tenant (which would leak/
|
||||
// overwrite their transcript); fall back to a fresh id instead of erroring.
|
||||
let session_id = match session_store::session_is_claimable(&state.db, &session_id, auth.user_id)
|
||||
.await
|
||||
{
|
||||
Ok(true) => session_id,
|
||||
Ok(false) => Uuid::new_v4().to_string(),
|
||||
Err(e) => {
|
||||
tracing::warn!(%e, "failed to check agent session ownership");
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
};
|
||||
let mode = PermissionMode::from_wire(body.permission_mode.as_deref().unwrap_or("default"));
|
||||
let model = body
|
||||
.model
|
||||
@@ -292,6 +364,26 @@ async fn chat_handler(
|
||||
tracing::warn!(%e, "failed to ensure agent session row");
|
||||
}
|
||||
|
||||
// Load any prior conversation for this session so the agent sees the full
|
||||
// multi-turn history (scoped to the caller — never another tenant's).
|
||||
let history: Vec<Message> = match session_store::load_transcript(
|
||||
&state.db,
|
||||
&session_id,
|
||||
auth.user_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(raw)) => serde_json::from_str(&raw).unwrap_or_else(|e| {
|
||||
tracing::warn!(%e, "failed to parse agent transcript; starting fresh");
|
||||
Vec::new()
|
||||
}),
|
||||
Ok(None) => Vec::new(),
|
||||
Err(e) => {
|
||||
tracing::warn!(%e, "failed to load agent transcript");
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
|
||||
// Route the agent's LLM calls back through BoosAPI's own aggregation using
|
||||
// the built-in downstream key.
|
||||
let base_url = boosapi_base(&state);
|
||||
@@ -299,9 +391,16 @@ async fn chat_handler(
|
||||
|
||||
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
||||
|
||||
// Cancellation: the token is handed to the agent loop (so a long LLM call or
|
||||
// tool can be interrupted mid-flight); a guard wrapping the SSE stream fires
|
||||
// it when the client disconnects (stream dropped). See `CancelGuard`.
|
||||
let cancel = CancellationToken::new();
|
||||
|
||||
let db = state.db.clone();
|
||||
let user_id = auth.user_id;
|
||||
let prompt = body.message;
|
||||
let session_for_event = session_id.clone();
|
||||
let agent_cancel = cancel.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// First event: hand the (possibly newly minted) session id to the client.
|
||||
@@ -318,6 +417,7 @@ async fn chat_handler(
|
||||
));
|
||||
|
||||
let can_use_tool = build_can_use_tool(tx.clone(), cwd.clone(), mode);
|
||||
let ask_fn = build_ask_fn(tx.clone());
|
||||
|
||||
// Replace the SDK's built-in Bash (a bare subprocess) with an
|
||||
// OS-sandboxed one confined to this session's directory. Registered as
|
||||
@@ -333,6 +433,8 @@ async fn chat_handler(
|
||||
system_prompt: Some(SYSTEM_PROMPT.to_string()),
|
||||
max_turns: Some(MAX_TURNS),
|
||||
can_use_tool: Some(can_use_tool),
|
||||
ask_fn: Some(ask_fn),
|
||||
abort_signal: Some(agent_cancel),
|
||||
custom_tools: vec![sandboxed_bash],
|
||||
session_id: Some(session_id.clone()),
|
||||
..Default::default()
|
||||
@@ -348,22 +450,84 @@ async fn chat_handler(
|
||||
}
|
||||
};
|
||||
|
||||
// Seed the agent with the prior conversation so this turn continues it.
|
||||
agent.messages = history;
|
||||
|
||||
let (mut rx_msg, handle) = agent.query(&prompt).await;
|
||||
|
||||
// The SDK's Result event carries the full updated message history; we
|
||||
// capture it here and persist it after the stream drains so the next
|
||||
// turn (same session id) resumes from it.
|
||||
let mut latest_transcript: Option<Vec<Message>> = None;
|
||||
|
||||
while let Some(event) = rx_msg.recv().await {
|
||||
// Stop early if the client disconnected.
|
||||
if tx.is_closed() {
|
||||
break;
|
||||
}
|
||||
if let SDKMessage::Result { messages, .. } = &event {
|
||||
latest_transcript = Some(messages.clone());
|
||||
}
|
||||
forward_sdk_event(&tx, &db, &session_id, event).await;
|
||||
}
|
||||
let _ = handle.await;
|
||||
agent.close().await;
|
||||
|
||||
// Persist the updated transcript (best-effort). Skipped if the turn was
|
||||
// cancelled before producing a Result, leaving the prior history intact.
|
||||
if let Some(messages) = latest_transcript {
|
||||
match serde_json::to_string(&messages) {
|
||||
Ok(raw) => {
|
||||
if let Err(e) =
|
||||
session_store::save_transcript(&db, &session_id, user_id, &raw).await
|
||||
{
|
||||
tracing::warn!(%e, "failed to save agent transcript");
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::warn!(%e, "failed to serialize agent transcript"),
|
||||
}
|
||||
}
|
||||
|
||||
send_sse(&tx, json!({ "type": "done" }));
|
||||
});
|
||||
|
||||
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
||||
// Wrap the stream so dropping it (client disconnect) cancels the agent loop.
|
||||
Ok(Sse::new(CancelGuard::new(
|
||||
UnboundedReceiverStream::new(rx),
|
||||
cancel,
|
||||
)))
|
||||
}
|
||||
|
||||
/// Wraps the SSE body stream and cancels a [`CancellationToken`] when the stream
|
||||
/// is dropped — i.e. when the client disconnects and Axum tears down the
|
||||
/// response. This propagates cancellation into the agent loop's `abort_signal`,
|
||||
/// so a long-running LLM call or tool execution is interrupted promptly instead
|
||||
/// of running to completion on a server nobody is listening to.
|
||||
struct CancelGuard<S> {
|
||||
inner: S,
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
|
||||
impl<S> CancelGuard<S> {
|
||||
fn new(inner: S, cancel: CancellationToken) -> Self {
|
||||
Self { inner, cancel }
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Stream + Unpin> Stream for CancelGuard<S> {
|
||||
type Item = S::Item;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
Pin::new(&mut self.inner).poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> Drop for CancelGuard<S> {
|
||||
fn drop(&mut self) {
|
||||
// Idempotent: a normal end-of-stream also drops this, harmlessly
|
||||
// cancelling an already-finished loop.
|
||||
self.cancel.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
/// Translate one SDK message into SSE event(s), and persist usage on Result.
|
||||
@@ -448,6 +612,18 @@ async fn approve_handler(
|
||||
Ok(Json(json!({ "ok": delivered })))
|
||||
}
|
||||
|
||||
async fn answer_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<AnswerRequest>,
|
||||
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
|
||||
// Auth required; the pending question is keyed by an unguessable uuid, so any
|
||||
// authenticated caller may answer their own running session's prompt.
|
||||
resolve_auth(&headers, &state).await?;
|
||||
let delivered = ask::answer(&body.request_id, body.answer);
|
||||
Ok(Json(json!({ "ok": delivered })))
|
||||
}
|
||||
|
||||
async fn status_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
|
||||
@@ -1047,6 +1047,7 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
payment: Default::default(),
|
||||
jimeng: Default::default(),
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user