refactor: remove retry_count from ProviderConfig and enhance error handling in tool execution

- Removed the `retry_count` field from `ProviderConfig` as it is no longer needed.
- Introduced `argument_error` field in `ToolCall` to capture errors related to tool arguments.
- Updated various components to handle argument errors more gracefully, including in the `ToolDispatcher` and `ConversationOutput`.
- Enhanced tests to validate the new error handling and ensure proper functionality of tool calls.

server/tests/compaction.rs

@@ -62,6 +62,7 @@ async fn summarize_replaces_model_history_and_preserves_cursor_history() {
         ModelEvent::TextEnd,
         ModelEvent::Usage(Usage {
             input_tokens: Some(4_012),
+            context_input_tokens: Some(4_012),
             output_tokens: Some(9),
             total_tokens: Some(4_021),
             ..Default::default()
@@ -442,6 +443,7 @@ fn text_response(text: &str, input: u64, output: u64) -> Vec<ModelEvent> {
         ModelEvent::TextEnd,
         ModelEvent::Usage(Usage {
             input_tokens: Some(input),
+            context_input_tokens: Some(input),
             output_tokens: Some(output),
             total_tokens: Some(input + output),
             ..Default::default()

server/tests/error_lifecycle.rs

@@ -51,10 +51,35 @@ async fn abort_command_cancels_the_run_and_closes_output() {
 }
 
 #[tokio::test]
-async fn provider_failure_keeps_the_initial_checkpoint_then_returns_structured_error() {
+async fn provider_failure_retries_from_the_current_checkpoint_without_hiding_partial_output() {
     let (_directory, store) = fixtures::temp_store().await;
     let provider = fake_provider::FakeProvider::default();
-    provider.push_error(Error::Provider("provider failed".into()));
+    provider.push_results(vec![
+        Ok(ModelEvent::Start {
+            model_call_id: "attempt-0".into(),
+        }),
+        Ok(ModelEvent::TextStart),
+        Ok(ModelEvent::TextDelta("partial ".into())),
+        Ok(ModelEvent::ToolCallStart {
+            index: 0,
+            call_id: "failed-tool".into(),
+            name: "Read".into(),
+        }),
+        Ok(ModelEvent::ToolCallArgumentsDelta {
+            index: 0,
+            delta: "{\"path\":".into(),
+        }),
+        Err(Error::Provider("stream disconnected".into())),
+    ]);
+    provider.push(vec![
+        ModelEvent::Start {
+            model_call_id: "attempt-1".into(),
+        },
+        ModelEvent::TextStart,
+        ModelEvent::TextDelta("completed".into()),
+        ModelEvent::TextEnd,
+        ModelEvent::Done(FinishReason::Stop),
+    ]);
     let assets = PromptAssets::load(
         std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
             .join("prompt/cursor")
@@ -63,7 +88,106 @@ async fn provider_failure_keeps_the_initial_checkpoint_then_returns_structured_e
     .unwrap();
     let registry = TransportRegistry::new(
         store.clone(),
-        Arc::new(provider),
+        Arc::new(provider.clone()),
+        PromptCompiler::new(assets),
+    );
+    let handle = registry.get_or_create("retry-request").await.unwrap();
+    let mut output = handle.subscribe();
+    handle
+        .command(TransportCommand::Append {
+            seqno: 0,
+            message: Box::new(protocol_client_run("retry", "retry-user")),
+        })
+        .await
+        .unwrap();
+
+    let mut seqno = 1;
+    let mut text = String::new();
+    let mut failed_tool_completed = false;
+    loop {
+        let frame = tokio::time::timeout(std::time::Duration::from_secs(60), output.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
+        if flags & connect::END_STREAM_FLAG != 0 {
+            assert_eq!(
+                serde_json::from_slice::<serde_json::Value>(&payload).unwrap(),
+                serde_json::json!({})
+            );
+            break;
+        }
+        let server = pb::AgentServerMessage::decode(payload).unwrap();
+        match server.message {
+            Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
+                handle
+                    .command(TransportCommand::Append {
+                        seqno,
+                        message: Box::new(kv_ack(kv.id)),
+                    })
+                    .await
+                    .unwrap();
+                seqno += 1;
+            }
+            Some(pb::agent_server_message::Message::InteractionUpdate(update)) => {
+                match update.message {
+                    Some(pb::interaction_update::Message::TextDelta(delta)) => {
+                        text.push_str(&delta.text);
+                    }
+                    Some(pb::interaction_update::Message::ToolCallCompleted(completed))
+                        if completed.call_id == "failed-tool" =>
+                    {
+                        let tool = completed.tool_call.expect("failed tool completion");
+                        failed_tool_completed = tool.completed_at_ms.is_some();
+                    }
+                    _ => {}
+                }
+            }
+            _ => {}
+        }
+    }
+
+    assert_eq!(text, "partial completed");
+    assert!(
+        failed_tool_completed,
+        "failed attempt must terminate its partial tool card"
+    );
+    let requests = provider.requests();
+    assert_eq!(requests.len(), 2);
+    assert_eq!(requests[0], requests[1]);
+    let messages = store
+        .load_current_messages(&cursor_server::model::ConversationId::new(
+            "protocol-failed-conversation",
+        ))
+        .await
+        .unwrap();
+    assert!(messages.iter().any(|message| {
+        matches!(&message.content, MessageContent::Assistant { text, .. } if text == "completed")
+    }));
+    assert!(!messages.iter().any(|message| {
+        matches!(&message.content, MessageContent::Assistant { text, .. } if text.contains("partial"))
+    }));
+}
+
+#[tokio::test]
+async fn provider_failure_keeps_the_initial_checkpoint_then_returns_structured_error() {
+    let (_directory, store) = fixtures::temp_store().await;
+    let provider = fake_provider::FakeProvider::default();
+    provider.push(vec![
+        ModelEvent::Start {
+            model_call_id: "length-limited".into(),
+        },
+        ModelEvent::Done(FinishReason::Length),
+    ]);
+    let assets = PromptAssets::load(
+        std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("prompt/cursor")
+            .as_path(),
+    )
+    .unwrap();
+    let registry = TransportRegistry::new(
+        store.clone(),
+        Arc::new(provider.clone()),
         PromptCompiler::new(assets),
     );
     let handle = registry.get_or_create("failed-request").await.unwrap();
@@ -148,6 +272,7 @@ async fn provider_failure_keeps_the_initial_checkpoint_then_returns_structured_e
         None
     );
 
+    assert_eq!(provider.requests().len(), 1, "Length must not retry");
     let messages = store
         .load_current_messages(&cursor_server::model::ConversationId::new(
             "failed-conversation",
@@ -164,7 +289,7 @@ async fn provider_failure_keeps_the_initial_checkpoint_then_returns_structured_e
 }
 
 #[tokio::test]
-async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes() {
+async fn unknown_tool_response_id_is_ignored_and_the_run_continues() {
     let (_directory, store) = fixtures::temp_store().await;
     let provider = fake_provider::FakeProvider::default();
     provider.push(vec![
@@ -183,6 +308,15 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
         ModelEvent::ToolCallEnd { index: 0 },
         ModelEvent::Done(FinishReason::ToolUse),
     ]);
+    provider.push(vec![
+        ModelEvent::Start {
+            model_call_id: "model-call-2".into(),
+        },
+        ModelEvent::TextStart,
+        ModelEvent::TextDelta("done".into()),
+        ModelEvent::TextEnd,
+        ModelEvent::Done(FinishReason::Stop),
+    ]);
     let assets = PromptAssets::load(
         std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
             .join("prompt/cursor")
@@ -209,7 +343,7 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
 
     let mut append_seqno = 1;
     let mut saw_turn_ended = false;
-    let error_json = loop {
+    let end_stream = loop {
         let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
             .await
             .unwrap()
@@ -231,24 +365,42 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
                 append_seqno += 1;
             }
             Some(pb::agent_server_message::Message::ExecServerMessage(exec)) => {
-                // An unknown numeric bridge id is a runtime protocol error.
-                handle
-                    .command(TransportCommand::Append {
-                        seqno: append_seqno,
-                        message: Box::new(pb::AgentClientMessage {
-                            message: Some(pb::agent_client_message::Message::ExecClientMessage(
-                                pb::ExecClientMessage {
-                                    id: exec.id + 1_000,
-                                    exec_id: String::new(),
-                                    message: None,
+                // Unknown bridge ids are ignored; the valid response still completes the tool.
+                for message in [
+                    pb::ExecClientMessage {
+                        id: exec.id + 1_000,
+                        exec_id: String::new(),
+                        message: None,
+                        ..Default::default()
+                    },
+                    pb::ExecClientMessage {
+                        id: exec.id,
+                        exec_id: String::new(),
+                        message: Some(pb::exec_client_message::Message::ReadResult(
+                            pb::ReadResult {
+                                result: Some(pb::read_result::Result::Success(pb::ReadSuccess {
+                                    path: "/tmp/a".into(),
+                                    output: Some(pb::read_success::Output::Content("value".into())),
                                     ..Default::default()
-                                },
-                            )),
-                        }),
-                    })
-                    .await
-                    .unwrap();
-                append_seqno += 1;
+                                })),
+                            },
+                        )),
+                        ..Default::default()
+                    },
+                ] {
+                    handle
+                        .command(TransportCommand::Append {
+                            seqno: append_seqno,
+                            message: Box::new(pb::AgentClientMessage {
+                                message: Some(
+                                    pb::agent_client_message::Message::ExecClientMessage(message),
+                                ),
+                            }),
+                        })
+                        .await
+                        .unwrap();
+                    append_seqno += 1;
+                }
             }
             Some(pb::agent_server_message::Message::InteractionUpdate(update)) => {
                 if matches!(
@@ -266,12 +418,8 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
         }
     };
 
-    assert!(!saw_turn_ended);
-    assert_eq!(error_json["error"]["code"], "invalid_argument");
-    assert_eq!(
-        error_json["error"]["message"],
-        "unknown ExecClientMessage id: 1001"
-    );
+    assert!(saw_turn_ended);
+    assert_eq!(end_stream, serde_json::json!({}));
     assert_eq!(
         tokio::time::timeout(std::time::Duration::from_secs(1), output.recv())
             .await
@@ -279,28 +427,14 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
         None
     );
 
-    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
-    let (status, failure_summary) = loop {
-        let row: (String, Option<String>) =
-            sqlx::query_as("SELECT status, failure_summary FROM runs WHERE cursor_request_id = ?")
-                .bind("protocol-failed-request")
-                .fetch_one(store.pool())
-                .await
-                .unwrap();
-        if row.0 != "running" {
-            break row;
-        }
-        assert!(
-            tokio::time::Instant::now() < deadline,
-            "Run remained running after the Cursor session failed"
-        );
-        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
-    };
-    assert_eq!(status, "failed");
-    assert_eq!(
-        failure_summary.as_deref(),
-        Some("unknown ExecClientMessage id: 1001")
-    );
+    let (status, failure_summary): (String, Option<String>) =
+        sqlx::query_as("SELECT status, failure_summary FROM runs WHERE cursor_request_id = ?")
+            .bind("protocol-failed-request")
+            .fetch_one(store.pool())
+            .await
+            .unwrap();
+    assert_eq!(status, "completed");
+    assert_eq!(failure_summary, None);
 }
 
 #[tokio::test]

server/tests/interrupt.rs

@@ -1580,6 +1580,7 @@ fn text_response(text: &str) -> Vec<ModelEvent> {
         ModelEvent::TextEnd,
         ModelEvent::Usage(Usage {
             input_tokens: Some(1),
+            context_input_tokens: Some(1),
             output_tokens: Some(1),
             total_tokens: Some(2),
             ..Default::default()

server/tests/support/fake_provider.rs

@@ -31,10 +31,13 @@ pub struct FakeProvider {
 
 impl FakeProvider {
     pub fn push(&self, events: Vec<ModelEvent>) {
+        self.push_results(events.into_iter().map(Ok).collect());
+    }
+    pub fn push_results(&self, events: Vec<Result<ModelEvent, Error>>) {
         self.responses
             .lock()
             .unwrap()
-            .push_back(FakeResponse::Events(events.into_iter().map(Ok).collect()));
+            .push_back(FakeResponse::Events(events));
     }
     pub fn push_error(&self, error: Error) {
         self.responses

server/tests/tool_round.rs

@@ -25,9 +25,11 @@ use cursor_server::{
         OPENAI_CHAT_ENDPOINT,
     },
     provider::{FinishReason, ModelEvent},
+    run::consume_model_cycle,
 };
 use prost::Message;
 use serde_json::json;
+use tokio_util::sync::CancellationToken;
 
 fn call(id: &str, name: &str) -> ToolCall {
     ToolCall {
@@ -37,6 +39,7 @@ fn call(id: &str, name: &str) -> ToolCall {
         name: name.into(),
         arguments_text: "{}".into(),
         arguments: json!({}),
+        argument_error: None,
     }
 }
 
@@ -68,6 +71,37 @@ fn mcp_context(server: &str, provider: &str, tool: &str) -> ExecContext {
     context
 }
 
+#[tokio::test]
+async fn malformed_provider_tool_json_is_kept_as_a_tool_validation_error() {
+    let stream = Box::pin(futures_util::stream::iter(vec![
+        Ok(ModelEvent::Start {
+            model_call_id: "model-call".into(),
+        }),
+        Ok(ModelEvent::ToolCallStart {
+            index: 0,
+            call_id: "call-malformed".into(),
+            name: "Read".into(),
+        }),
+        Ok(ModelEvent::ToolCallArgumentsDelta {
+            index: 0,
+            delta: "{\"path\":".into(),
+        }),
+        Ok(ModelEvent::ToolCallEnd { index: 0 }),
+        Ok(ModelEvent::Done(FinishReason::ToolUse)),
+    ]));
+    let (events, _receiver) = tokio::sync::mpsc::channel(16);
+    let result = consume_model_cycle(stream, &events, &CancellationToken::new())
+        .await
+        .expect("malformed tool JSON must not fail the model cycle");
+
+    assert_eq!(result.calls.len(), 1);
+    assert!(result.calls[0]
+        .argument_error
+        .as_deref()
+        .is_some_and(|message| message.contains("not valid JSON")));
+    assert_eq!(result.calls[0].arguments, json!({}));
+}
+
 #[test]
 fn dynamic_mcp_call_routes_to_the_captured_exec_message() {
     let call = ToolCall {
@@ -77,6 +111,7 @@ fn dynamic_mcp_call_routes_to_the_captured_exec_message() {
         name: "mcp_repo_lookup".into(),
         arguments_text: "{\"query\":\"x\"}".into(),
         arguments: json!({"query": "x"}),
+        argument_error: None,
     };
     let definition = pb::McpToolDefinition {
         name: "mcp_repo_lookup".into(),
@@ -372,6 +407,52 @@ async fn unknown_mcp_descriptor_returns_a_tool_error_without_client_discovery()
     assert!(completion.result().content.contains("descriptor not found"));
 }
 
+#[tokio::test]
+async fn invalid_tool_arguments_complete_as_tool_errors() {
+    let dispatcher = ToolDispatcher::new(CursorToolRuntime::default());
+    let completed = HashSet::new();
+    let started = HashSet::new();
+    let state = || ToolBatchState {
+        completed: &completed,
+        started: &started,
+        response_text: "",
+        response_thinking: "",
+    };
+
+    let mut malformed = call("call-malformed", "Read");
+    malformed.argument_error = Some("Read arguments are not valid JSON".into());
+    let mut missing = call("call-missing", "Shell");
+    missing.arguments = json!({"description": "missing command"});
+    let mut wrong_type = call("call-type", "Shell");
+    wrong_type.arguments = json!({"command": 42});
+    let mut invalid_timeout = call("call-timeout", "Shell");
+    invalid_timeout.arguments = json!({"command": "pwd", "block_until_ms": -1});
+
+    for (invocation, expected) in [
+        (malformed, "not valid JSON"),
+        (missing, "missing command"),
+        (wrong_type, "missing command"),
+        (invalid_timeout, "out of range"),
+    ] {
+        let dispatched = dispatcher
+            .start_batch(
+                &[invocation],
+                state(),
+                &[],
+                &BTreeMap::new(),
+                &exec_context(),
+            )
+            .await
+            .unwrap();
+        let completion = dispatched[0]
+            .completion
+            .as_ref()
+            .expect("invalid arguments must complete as a tool error");
+        assert!(completion.result().is_error);
+        assert!(completion.result().content.contains(expected));
+    }
+}
+
 #[tokio::test]
 async fn shell_uses_background_timeout_and_preserves_stream_identity() {
     let mut shell = call("call-shell", "Shell");
@@ -701,12 +782,15 @@ async fn an_exec_result_must_match_the_reserved_tool() {
         },
         &pending,
     )
-    .await;
-    let Err(error) = result else {
-        panic!("mismatched result must fail")
+    .await
+    .unwrap();
+    let codec::ClientExecEvent::Completed(completion) = result else {
+        panic!("mismatched result must complete as a tool error")
     };
-    assert!(error
-        .to_string()
+    assert!(completion.result().is_error);
+    assert!(completion
+        .result()
+        .content
         .contains("unexpected Exec result for tool Read"));
     assert!(pending.exec_call(id).await.is_none());
     assert_eq!(pending.completed_call(id).await.as_deref(), Some("call-1"));
@@ -718,15 +802,13 @@ async fn an_exec_result_must_match_the_reserved_tool() {
         },
         &pending,
     )
-    .await;
-    let Err(duplicate) = duplicate else {
-        panic!("duplicate terminal result must fail")
-    };
-    assert!(duplicate.to_string().contains("duplicate terminal"));
+    .await
+    .unwrap();
+    assert!(matches!(duplicate, codec::ClientExecEvent::Pending));
 }
 
 #[tokio::test]
-async fn unknown_exec_id_is_a_protocol_error() {
+async fn unknown_exec_id_is_ignored() {
     let result = codec::client_event(
         &pb::ExecClientMessage {
             id: 999,
@@ -737,15 +819,9 @@ async fn unknown_exec_id_is_a_protocol_error() {
         },
         &CursorToolRuntime::default(),
     )
-    .await;
-    let Err(error) = result else {
-        panic!("unknown Exec id must fail")
-    };
-    assert!(matches!(
-        error,
-        cursor_server::Error::Protocol(message)
-            if message == "unknown ExecClientMessage id: 999"
-    ));
+    .await
+    .unwrap();
+    assert!(matches!(result, codec::ClientExecEvent::Pending));
 }
 
 #[tokio::test]