all tools

2026-08-18 12:07:02 +08:00 · 2026-08-16 17:29:29 +08:00
parent eafede22e3
commit 4db2061611
95 changed files with 19023 additions and 5637 deletions
@@ -0,0 +1,190 @@
+#[path = "support/fake_provider.rs"]
+mod fake_provider;
+#[path = "support/fixtures.rs"]
+mod fixtures;
+
+use std::sync::Arc;
+
+use cursor_server::{
+    cursor::{connect, proto::agent::v1 as pb},
+    prompting::{PromptAssets, PromptCompiler},
+    provider::{FinishReason, ResponseEvent},
+    run::{RunCommand, RunRegistry},
+};
+use prost::Message;
+
+#[tokio::test]
+async fn a_new_revision_invalidates_late_events_from_the_old_run() {
+    let (_directory, store) = fixtures::temp_store().await;
+    let first = store.begin_revision("conversation").await.unwrap();
+    let second = store.begin_revision("conversation").await.unwrap();
+    assert!(second > first);
+    assert!(!store
+        .revision_is_current("conversation", first)
+        .await
+        .unwrap());
+    assert!(store
+        .revision_is_current("conversation", second)
+        .await
+        .unwrap());
+}
+
+#[tokio::test]
+async fn registry_shutdown_cancels_runs_and_closes_run_sse_outputs() {
+    let (_directory, store) = fixtures::temp_store().await;
+    let assets = PromptAssets::load(
+        std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("../prompt")
+            .as_path(),
+    )
+    .unwrap();
+    let registry = RunRegistry::new(
+        store,
+        Arc::new(fake_provider::FakeProvider::default()),
+        PromptCompiler::new(assets),
+        "test-model".into(),
+    );
+    let handle = registry.get_or_create("active-run").await.unwrap();
+    let mut output = handle.subscribe();
+
+    registry.shutdown().await;
+
+    assert!(handle.cancellation().is_cancelled());
+    let terminal = output.recv().await.expect("canceled EndStream");
+    let (flags, payload) = connect::decode_frames(&terminal).unwrap().pop().unwrap();
+    assert_eq!(flags, connect::END_STREAM_FLAG);
+    let payload: serde_json::Value = serde_json::from_slice(&payload).unwrap();
+    assert_eq!(payload["error"]["code"], "canceled");
+    assert_eq!(output.recv().await, None);
+}
+
+#[tokio::test]
+async fn cancel_aborts_active_exec_before_canceled_end_stream() {
+    let (_directory, store) = fixtures::temp_store().await;
+    let provider = fake_provider::FakeProvider::default();
+    provider.push(vec![
+        ResponseEvent::ToolCallStart {
+            index: 0,
+            call_id: "call-1".into(),
+            name: "Read".into(),
+        },
+        ResponseEvent::ToolCallArgumentsDelta {
+            index: 0,
+            delta: "{\"path\":\"/tmp/a\"}".into(),
+        },
+        ResponseEvent::ToolCallEnd { index: 0 },
+        ResponseEvent::Done(FinishReason::ToolUse),
+    ]);
+    let assets = PromptAssets::load(
+        std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("../prompt")
+            .as_path(),
+    )
+    .unwrap();
+    let registry = RunRegistry::new(
+        store,
+        Arc::new(provider),
+        PromptCompiler::new(assets),
+        "test-model".into(),
+    );
+    let handle = registry.get_or_create("cancel-request").await.unwrap();
+    let mut output = handle.subscribe();
+    handle
+        .command(RunCommand::Append {
+            seqno: 0,
+            message: Box::new(client_run()),
+        })
+        .await
+        .unwrap();
+
+    let mut append_seqno = 1;
+    let exec_id = loop {
+        let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
+        let server = pb::AgentServerMessage::decode(payload).unwrap();
+        match server.message {
+            Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
+                handle
+                    .command(RunCommand::Append {
+                        seqno: append_seqno,
+                        message: Box::new(kv_ack(kv.id)),
+                    })
+                    .await
+                    .unwrap();
+                append_seqno += 1;
+            }
+            Some(pb::agent_server_message::Message::ExecServerMessage(exec)) => break exec.id,
+            _ => {}
+        }
+    };
+
+    handle.cancel();
+    let mut saw_abort = false;
+    loop {
+        let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
+            .await
+            .unwrap()
+            .expect("RunSSE closed before canceled EndStream");
+        let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
+        if flags & connect::END_STREAM_FLAG != 0 {
+            let json: serde_json::Value = serde_json::from_slice(&payload).unwrap();
+            assert_eq!(json["error"]["code"], "canceled");
+            assert!(saw_abort, "ExecServerAbort must precede canceled EndStream");
+            break;
+        }
+        let server = pb::AgentServerMessage::decode(payload).unwrap();
+        if let Some(pb::agent_server_message::Message::ExecServerControlMessage(control)) =
+            server.message
+        {
+            let Some(pb::exec_server_control_message::Message::Abort(abort)) = control.message
+            else {
+                panic!("expected ExecServerAbort")
+            };
+            assert_eq!(abort.id, exec_id);
+            saw_abort = true;
+        }
+    }
+    assert_eq!(output.recv().await, None);
+}
+
+fn client_run() -> pb::AgentClientMessage {
+    pb::AgentClientMessage {
+        message: Some(pb::agent_client_message::Message::RunRequest(
+            pb::AgentRunRequest {
+                action: Some(pb::ConversationAction {
+                    action: Some(pb::conversation_action::Action::UserMessageAction(
+                        pb::UserMessageAction {
+                            user_message: Some(pb::UserMessage {
+                                text: "read".into(),
+                                message_id: "cancel-user".into(),
+                                mode: pb::AgentMode::Agent as i32,
+                                ..Default::default()
+                            }),
+                            ..Default::default()
+                        },
+                    )),
+                    ..Default::default()
+                }),
+                conversation_id: Some("cancel-conversation".into()),
+                run_id: Some("cancel-request".into()),
+                ..Default::default()
+            },
+        )),
+    }
+}
+
+fn kv_ack(id: u32) -> pb::AgentClientMessage {
+    pb::AgentClientMessage {
+        message: Some(pb::agent_client_message::Message::KvClientMessage(
+            pb::KvClientMessage {
+                id,
+                message: Some(pb::kv_client_message::Message::SetBlobResult(
+                    pb::SetBlobResult { error: None },
+                )),
+            },
+        )),
+    }
+}