tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Agent loop defensive (#3554)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-07-22 18:21:47 +02:00
committed by GitHub
parent b3cd03ef61
commit 9f356e7009
19 changed files with 349 additions and 773 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/reply.rs
+113 -274
View File
@@ -11,7 +11,7 @@ use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::{
agents::{AgentEvent, SessionConfig},
message::{push_message, Message, MessageContent},
message::{push_message, Message},
permission::permission_confirmation::PrincipalType,
};
use goose::{
@@ -19,7 +19,7 @@ use goose::{
session,
};
use mcp_core::{protocol::JsonRpcMessage, ToolResult};
use rmcp::model::{Content, Role};
use rmcp::model::Content;
use serde::{Deserialize, Serialize};
use serde_json::json;
use serde_json::Value;
@@ -34,9 +34,10 @@ use std::{
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use utoipa::ToSchema;
#[derive(Debug, Deserialize)]
#[derive(Debug, Deserialize, Serialize)]
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
@@ -113,7 +114,7 @@ async fn stream_event(
tx.send(format!("data: {}\n\n", json)).await
}
async fn handler(
async fn reply_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(request): Json<ChatRequest>,
@@ -122,6 +123,7 @@ async fn handler(
let (tx, rx) = mpsc::channel(100);
let stream = ReceiverStream::new(rx);
let cancel_token = CancellationToken::new();
let messages = request.messages;
let session_working_dir = request.session_working_dir.clone();
@@ -130,65 +132,35 @@ async fn handler(
.session_id
.unwrap_or_else(session::generate_session_id);
tokio::spawn(async move {
let agent = state.get_agent().await;
let agent = match agent {
Ok(agent) => {
let provider = agent.provider().await;
match provider {
Ok(_) => agent,
Err(_) => {
let _ = stream_event(
MessageEvent::Error {
error: "No provider configured".to_string(),
},
&tx,
)
.await;
let _ = stream_event(
MessageEvent::Finish {
reason: "error".to_string(),
},
&tx,
)
.await;
return;
}
}
}
let task_cancel = cancel_token.clone();
let task_tx = tx.clone();
std::mem::drop(tokio::spawn(async move {
let agent = match state.get_agent().await {
Ok(agent) => agent,
Err(_) => {
let _ = stream_event(
MessageEvent::Error {
error: "No agent configured".to_string(),
},
&tx,
)
.await;
let _ = stream_event(
MessageEvent::Finish {
reason: "error".to_string(),
},
&tx,
&task_tx,
)
.await;
return;
}
};
let provider = agent.provider().await;
let session_config = SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(&session_working_dir),
schedule_id: request.scheduled_job_id.clone(),
execution_mode: None,
max_turns: None,
retry_config: None,
};
let mut stream = match agent
.reply(
&messages,
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(&session_working_dir),
schedule_id: request.scheduled_job_id.clone(),
execution_mode: None,
max_turns: None,
retry_config: None,
}),
)
.reply(&messages, Some(session_config), Some(task_cancel.clone()))
.await
{
Ok(stream) => stream,
@@ -198,14 +170,7 @@ async fn handler(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
)
.await;
let _ = stream_event(
MessageEvent::Finish {
reason: "error".to_string(),
},
&tx,
&task_tx,
)
.await;
return;
@@ -221,7 +186,7 @@ async fn handler(
MessageEvent::Error {
error: format!("Failed to get session path: {}", e),
},
&tx,
&task_tx,
)
.await;
return;
@@ -231,222 +196,104 @@ async fn handler(
loop {
tokio::select! {
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(AgentEvent::Message(message)))) => {
push_message(&mut all_messages, message.clone());
if let Err(e) = stream_event(MessageEvent::Message { message }, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
_ = task_cancel.cancelled() => {
tracing::info!("Agent task cancelled");
break;
}
}
Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
if let Err(e) = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await {
tracing::error!("Error sending model change through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
}
}
Ok(Some(Ok(AgentEvent::McpNotification((request_id, n))))) => {
if let Err(e) = stream_event(MessageEvent::Notification{
request_id: request_id.clone(),
message: n,
}, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
}
}
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(AgentEvent::Message(message)))) => {
push_message(&mut all_messages, message.clone());
if let Err(e) = stream_event(MessageEvent::Message { message }, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
break;
}
}
Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
if let Err(e) = stream_event(MessageEvent::ModelChange { model, mode }, &tx).await {
tracing::error!("Error sending model change through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
}
}
Ok(Some(Ok(AgentEvent::McpNotification((request_id, n))))) => {
if let Err(e) = stream_event(MessageEvent::Notification{
request_id: request_id.clone(),
message: n,
}, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
}
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
break;
}
Ok(None) => {
break;
}
Err(_) => { // Heartbeat, used to detect disconnected clients
if tx.is_closed() {
break;
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: e.to_string(),
},
&tx,
).await;
break;
}
Ok(None) => {
break;
}
Err(_) => {
if tx.is_closed() {
break;
}
continue;
}
}
}
continue;
}
}
}
}
}
if all_messages.len() > saved_message_count {
let provider = Arc::clone(provider.as_ref().unwrap());
tokio::spawn(async move {
if let Err(e) = session::persist_messages(
&session_path,
&all_messages,
Some(provider),
Some(PathBuf::from(&session_working_dir)),
)
.await
{
tracing::error!("Failed to store session history: {:?}", e);
}
});
if let Ok(provider) = agent.provider().await {
let provider = Arc::clone(&provider);
tokio::spawn(async move {
if let Err(e) = session::persist_messages(
&session_path,
&all_messages,
Some(provider),
Some(PathBuf::from(&session_working_dir)),
)
.await
{
tracing::error!("Failed to store session history: {:?}", e);
}
});
}
}
let _ = stream_event(
MessageEvent::Finish {
reason: "stop".to_string(),
},
&tx,
&task_tx,
)
.await;
});
}));
Ok(SseResponse::new(stream))
}
#[derive(Debug, Deserialize, Serialize)]
struct AskRequest {
prompt: String,
session_id: Option<String>,
session_working_dir: String,
scheduled_job_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct AskResponse {
response: String,
}
async fn ask_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(request): Json<AskRequest>,
) -> Result<Json<AskResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let session_working_dir = request.session_working_dir.clone();
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
let agent = state
.get_agent()
.await
.map_err(|_| StatusCode::PRECONDITION_FAILED)?;
let provider = agent.provider().await;
let messages = vec![Message::user().with_text(request.prompt)];
let mut response_text = String::new();
let mut stream = match agent
.reply(
&messages,
Some(SessionConfig {
id: session::Identifier::Name(session_id.clone()),
working_dir: PathBuf::from(&session_working_dir),
schedule_id: request.scheduled_job_id.clone(),
execution_mode: None,
max_turns: None,
retry_config: None,
}),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let mut all_messages = messages.clone();
let mut response_message = Message::assistant();
while let Some(response) = stream.next().await {
match response {
Ok(AgentEvent::Message(message)) => {
if message.role == Role::Assistant {
for content in &message.content {
if let MessageContent::Text(text) = content {
response_text.push_str(&text.text);
response_text.push('\n');
}
response_message.content.push(content.clone());
}
}
}
Ok(AgentEvent::ModelChange { model, mode }) => {
// Log model change for non-streaming
tracing::info!("Model changed to {} in {} mode", model, mode);
}
Ok(AgentEvent::McpNotification(n)) => {
// Handle notifications if needed
tracing::info!("Received notification: {:?}", n);
}
Err(e) => {
tracing::error!("Error processing as_ai message: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
}
}
if !response_message.content.is_empty() {
push_message(&mut all_messages, response_message);
}
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
Ok(path) => path,
Err(e) => {
tracing::error!("Failed to get session path: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let session_path_clone = session_path.clone();
let messages = all_messages.clone();
let provider = Arc::clone(provider.as_ref().unwrap());
let session_working_dir_clone = session_working_dir.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(
&session_path_clone,
&messages,
Some(provider),
Some(PathBuf::from(session_working_dir_clone)),
)
.await
{
tracing::error!("Failed to store session history: {:?}", e);
}
});
Ok(Json(AskResponse {
response: response_text.trim().to_string(),
}))
}
#[derive(Debug, Deserialize, Serialize, ToSchema)]
pub struct PermissionConfirmationRequest {
id: String,
@@ -509,7 +356,7 @@ struct ToolResultRequest {
async fn submit_tool_result(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
raw: axum::extract::Json<serde_json::Value>,
raw: Json<Value>,
) -> Result<Json<Value>, StatusCode> {
verify_secret_key(&headers, &state)?;
@@ -540,8 +387,7 @@ async fn submit_tool_result(
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/reply", post(handler))
.route("/ask", post(ask_handler))
.route("/reply", post(reply_handler))
.route("/confirm", post(confirm_permission))
.route("/tool_result", post(submit_tool_result))
.with_state(state)
@@ -571,10 +417,6 @@ mod tests {
goose::providers::base::ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete(
&self,
_system: &str,
@@ -586,6 +428,10 @@ mod tests {
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
}
mod integration_tests {
@@ -595,7 +441,7 @@ mod tests {
use tower::ServiceExt;
#[tokio::test]
async fn test_ask_endpoint() {
async fn test_reply_endpoint() {
let mock_model_config = ModelConfig::new("test-model".to_string());
let mock_provider = Arc::new(MockProvider {
model_config: mock_model_config,
@@ -603,24 +449,17 @@ mod tests {
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let state = AppState::new(Arc::new(agent), "test-secret".to_string()).await;
let scheduler_path = goose::scheduler::get_default_scheduler_storage_path()
.expect("Failed to get default scheduler storage path");
let scheduler =
goose::scheduler_factory::SchedulerFactory::create_legacy(scheduler_path)
.await
.unwrap();
state.set_scheduler(scheduler).await;
let app = routes(state);
let request = Request::builder()
.uri("/ask")
.uri("/reply")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test-secret")
.body(Body::from(
serde_json::to_string(&AskRequest {
prompt: "test prompt".to_string(),
serde_json::to_string(&ChatRequest {
messages: vec![Message::user().with_text("test message")],
session_id: Some("test-session".to_string()),
session_working_dir: "test-working-dir".to_string(),
scheduled_job_id: None,