Skip to content

fix(acp): keep sessions usable after provider failures #72

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yan-ad wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(acp): keep sessions usable after provider failures #72

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
172 changes: 168 additions & 4 deletions src/acp/service.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1715,7 +1715,10 @@ impl AcpService {
cancelled = true;
break;
}
crate::llm::ChunkMessage::Failed(error) => failed = Some(error),
crate::llm::ChunkMessage::Failed(error) => {
failed = Some(error);
break;
}
crate::llm::ChunkMessage::TurnStopReason(reason) => {
turn_stop_reason = Some(reason)
}
Expand Down Expand Up @@ -1768,6 +1771,42 @@ impl AcpService {
}
}

if !cancelled && !cancellation.is_cancelled() {
if let Some(error) = failed.as_deref() {
let pending_ids = assistant
.parts
.iter()
.filter(|part| part.part_type == "tool_call")
.filter_map(|part| part.tool_id().map(str::to_string))
.filter(|id| {
assistant.tool_result_part_data(id).is_none_or(|data| {
matches!(data.get("status").and_then(|s| s.as_str()),
None | Some("running" | "pending"))
})
})
.collect::<Vec<_>>();
assistant.mark_running_tool_parts_failed(error);
for id in pending_ids {
let payload = assistant.tool_result_part_data(&id).unwrap();
send_tool_result(
connection,
session_id,
crate::llm::ToolCallResult {
tool_call_id: id,
role: "tool".to_string(),
name: payload["name"].as_str().unwrap_or("tool").to_string(),
content: payload.to_string(),
},
&session.cwd,
)?;
}
let text = format!(
"\n\nThe AI request failed: {error}\nYou can retry or send another message in this session."
);
assistant.append(&text);
send_text(connection, session_id, &message_id, text, false)?;
}
}
Ok(())
}
.await;
Expand All @@ -1776,12 +1815,19 @@ impl AcpService {
}
assistant.is_complete = true;
assistant.was_interrupted = cancelled || cancellation.is_cancelled();
if assistant.was_interrupted || failed.is_some() {
assistant.mark_running_tool_parts_failed(
failed
.as_deref()
.unwrap_or("Tool execution cancelled by user"),
);
}
{
let mut manager = self.session_manager.lock().map_err(|_| internal_error())?;
let persist_result = manager.add_message_to_session(session_id, &assistant);
let status = if assistant.was_interrupted {
crate::session::types::SessionStatus::Interrupted
} else if failed.is_some() {
} else if stream_result.is_err() {
crate::session::types::SessionStatus::Failed
} else {
crate::session::types::SessionStatus::Idle
Expand All @@ -1795,8 +1841,8 @@ impl AcpService {
return Ok(PromptResponse::new(StopReason::Cancelled));
}
stream_result?;
if let Some(error) = failed {
return Err(internal_error_with(&error));
if failed.is_some() {
return Ok(PromptResponse::new(StopReason::EndTurn));
}
Ok(PromptResponse::new(acp_stop_reason(turn_stop_reason)))
}
Expand Down Expand Up @@ -4224,6 +4270,124 @@ mod tests {
}
}

#[tokio::test(flavor = "current_thread")]
async fn provider_failure_without_tools_keeps_acp_session_usable() {
provider_failure_keeps_session_usable(false).await;
}

#[tokio::test(flavor = "current_thread")]
async fn provider_failure_marks_pending_tools_failed_and_keeps_acp_session_usable() {
provider_failure_keeps_session_usable(true).await;
}

async fn provider_failure_keeps_session_usable(with_tools: bool) {
use agent_client_protocol::Agent;
use std::time::Duration;

tokio::task::LocalSet::new().run_until(async {
let workspace = tempfile::tempdir().unwrap();
let mut manager = SessionManager::new()
.with_history_for_workspace(workspace.path()).unwrap();
let session_id = manager.create_session(Some("ACP provider failure regression".into()));
let db_id = manager.get_db_id(&session_id).unwrap();
let service = AcpService {
sessions: Arc::new(AsyncMutex::new(AcpSessions::default())),
session_manager: Arc::new(Mutex::new(manager)),
client_capabilities: Arc::new(Mutex::new(Default::default())),
};
let session = session_with_config(empty_config());
service.sessions.lock().await.attachments.insert(session_id.clone(), session.clone());
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
let client = Client.builder().on_receive_notification(
async move |notification: SessionNotification, _connection| {
let _ = observed_tx.send(serde_json::to_value(notification).unwrap());
Ok(())
},
agent_client_protocol::on_receive_notification!(),
);
Agent.builder().connect_with(client, async move |connection| {
let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap();
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
let error = "Streaming failed: API error 400: No tool output found for function call call_119.";
if with_tools {
sender.send(crate::llm::ChunkMessage::Text("Inspecting files.".into())).unwrap();
let calls = ["completed-tool", "call_119"].into_iter().map(|id| crate::llm::ToolCall {
id: id.into(),
call_type: "function".into(),
function: crate::llm::tool_calls::FunctionCall {
name: "read".into(), arguments: "{\"file_path\":\"src/lib.rs\"}".into(),
},
}).collect();
sender.send(crate::llm::ChunkMessage::ToolCalls(calls)).unwrap();
sender.send(crate::llm::ChunkMessage::ToolResult(crate::llm::ToolCallResult {
tool_call_id: "completed-tool".into(), role: "tool".into(), name: "read".into(),
content: serde_json::json!({"status":"ok", "output_preview":"file contents"}).to_string(),
})).unwrap();
}
sender.send(crate::llm::ChunkMessage::Failed(error.into())).unwrap();
sender.send(crate::llm::ChunkMessage::End).unwrap();
let response = service.receive_prompt_stream(
&session_id, &session, &connection, &cancellation, receiver, 0, 0.0,
).await.expect("provider errors must not become ACP RPC errors");
assert_eq!(response.stop_reason, StopReason::EndTurn);
service.sessions.lock().await.finish_prompt(&session_id);

let mut saw_failed_tool = false;
let mut saw_error_message = false;
tokio::time::timeout(Duration::from_secs(2), async {
while let Some(notification) = observed_rx.recv().await {
let update = &notification["update"];
if update["sessionUpdate"] == "tool_call_update" && update["toolCallId"] == "call_119" {
assert_eq!(update["status"], "failed");
saw_failed_tool = true;
}
if update["sessionUpdate"] == "agent_message_chunk"
&& update["content"]["text"].as_str().is_some_and(|text| text.contains(error)) {
saw_error_message = true;
break;
}
}
}).await.unwrap();
assert!(saw_error_message, "client must receive the provider error in-band");
assert_eq!(saw_failed_tool, with_tools);
{
let manager = service.session_manager.lock().unwrap();
let stored = manager.get_session_ref(&session_id).unwrap();
assert_eq!(stored.status, crate::session::types::SessionStatus::Idle);
let assistant = stored.messages.last().unwrap();
assert!(assistant.is_complete);
assert!(!assistant.was_interrupted);
assert!(!assistant.has_running_tool_parts());
assert!(assistant.content.contains(error));
if with_tools {
assert_eq!(assistant.tool_result_part_data("completed-tool").unwrap()["status"], "ok");
assert_eq!(assistant.tool_result_part_data("call_119").unwrap()["status"], "error");
assert_eq!(assistant.tool_result_part_data("call_119").unwrap()["output_preview"], error);
}
}
let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap();
assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "idle");
let messages = dao.get_messages(db_id).unwrap();
assert_eq!(messages.len(), 1);
assert!(messages[0].parts.iter().any(|part| part.part_type == "text"
&& part.data["text"].as_str().is_some_and(|text| text.contains(error))));

let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap();
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
sender.send(crate::llm::ChunkMessage::Text("The next turn works.".into())).unwrap();
sender.send(crate::llm::ChunkMessage::End).unwrap();
let response = service.receive_prompt_stream(
&session_id, &session, &connection, &cancellation, receiver, 0, 0.0,
).await.unwrap();
assert_eq!(response.stop_reason, StopReason::EndTurn);
service.sessions.lock().await.finish_prompt(&session_id);
assert_eq!(service.session_manager.lock().unwrap().get_session_ref(&session_id).unwrap().messages.len(), 2);
service.session_manager.lock().unwrap().try_delete_session(&session_id).unwrap();
Ok(())
}).await.unwrap();
}).await;
}

#[tokio::test(flavor = "current_thread")]
async fn cancelled_stream_persists_partial_transcript_and_interrupted_status() {
cancelled_stream_finalizes(false, false).await;
Expand Down
127 changes: 90 additions & 37 deletions src/aisdk/providers/openai.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::tool::Tool;
use async_trait::async_trait;
use eventsource_stream::{EventStreamError, Eventsource};
use futures::{SinkExt, StreamExt};
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::error::Error as StdError;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
Expand Down Expand Up @@ -2359,50 +2359,59 @@ fn build_openai_messages(
/// reuses an id fails with "Duplicate function_call_output". Remap each
/// repeated call — and its matching output — to a unique id; ids are opaque.
fn dedupe_responses_call_ids(mut items: Vec<serde_json::Value>) -> Vec<serde_json::Value> {
let mut reserved_ids: HashSet<String> = items
.iter()
.filter_map(|item| item.get("call_id").and_then(|value| value.as_str()))
.map(str::to_string)
.collect();
let mut replacements: HashMap<(String, usize), String> = HashMap::new();
// Occurrence k of a call_id pairs with occurrence k of its output, so
// track counts per call_id instead of assuming strict call/output order.
let mut call_occurrences: HashMap<String, usize> = HashMap::new();
let mut output_occurrences: HashMap<String, usize> = HashMap::new();

for item in items.iter_mut() {
match item.get("type").and_then(|value| value.as_str()) {
Some("function_call") => {
let Some(call_id) = item
.get("call_id")
.and_then(|value| value.as_str())
.map(str::to_string)
else {
continue;
};
let occurrence = call_occurrences.entry(call_id.clone()).or_insert(0);
*occurrence += 1;
if *occurrence == 1 {
continue;
}
let new_id = format!("{call_id}_dup{occurrence}");
item["call_id"] = serde_json::Value::String(new_id);
// The item id belongs to the original response; drop it so the
// rewritten pair cannot collide on item ids either.
if let Some(obj) = item.as_object_mut() {
obj.remove("id");
}
}
Some("function_call_output") => {
let Some(call_id) = item
.get("call_id")
.and_then(|value| value.as_str())
.map(str::to_string)
else {
continue;
};
let occurrence = output_occurrences.entry(call_id.clone()).or_insert(0);
*occurrence += 1;
if *occurrence > 1 {
item["call_id"] =
serde_json::Value::String(format!("{call_id}_dup{occurrence}"));
let is_call = match item.get("type").and_then(|value| value.as_str()) {
Some("function_call") => true,
Some("function_call_output") => false,
_ => continue,
};
let Some(call_id) = item
.get("call_id")
.and_then(|value| value.as_str())
.map(str::to_string)
else {
continue;
};
let occurrences = if is_call {
&mut call_occurrences
} else {
&mut output_occurrences
};
let occurrence = occurrences.entry(call_id.clone()).or_insert(0);
*occurrence += 1;
if *occurrence == 1 {
continue;
}
let new_id = replacements
.entry((call_id.clone(), *occurrence))
.or_insert_with(|| {
let mut suffix = *occurrence;
loop {
let candidate = format!("{call_id}_dup{suffix}");
if reserved_ids.insert(candidate.clone()) {
break candidate;
}
suffix += 1;
}
});
item["call_id"] = serde_json::Value::String(new_id.clone());
// The item id belongs to the original response; drop it so the
// rewritten pair cannot collide on item ids either.
if is_call {
if let Some(obj) = item.as_object_mut() {
obj.remove("id");
}
_ => {}
}
}

Expand Down Expand Up @@ -3029,6 +3038,50 @@ mod tests {
assert!(input[3].get("id").is_none());
}

#[test]
fn dedupes_responses_call_ids_without_colliding_with_original_suffixes() {
let input = build_openai_messages(
&[
Message::tool_call("call_1", "read", "{}"),
Message::tool_output("call_1", "read", "first", false),
Message::tool_call_with_item_id("fc_repeated", "call_1", "read", "{}"),
Message::tool_call_with_item_id("fc_original", "call_1_dup2", "glob", "{}"),
Message::tool_call("call_2", "grep", "{}"),
Message::tool_output("call_2", "grep", "concurrent third", false),
Message::tool_output("call_1_dup2", "glob", "concurrent second", false),
Message::tool_output("call_1", "read", "concurrent first", false),
Message::tool_call("call_1", "read", "{}"),
Message::tool_output("call_1", "read", "last", false),
],
false,
false,
);

let call_ids: Vec<&str> = input
.iter()
.map(|item| item["call_id"].as_str().expect("call_id"))
.collect();
assert_eq!(
call_ids,
[
"call_1",
"call_1",
"call_1_dup3",
"call_1_dup2",
"call_2",
"call_2",
"call_1_dup2",
"call_1_dup3",
"call_1_dup4",
"call_1_dup4",
]
);
assert!(input[2].get("id").is_none());
assert_eq!(input[3]["id"], "fc_original");
assert_eq!(input[6]["output"], "concurrent second");
assert_eq!(input[7]["output"], "concurrent first");
}

#[test]
fn serializes_tool_image_output_for_responses_input() {
let input = build_openai_messages(
Expand Down
Loading