feat(console): add initial console for Cursor BYOK with provider management and LLM call tracking

This commit is contained in:
leookun
2026-08-17 10:14:10 +08:00
parent 4db2061611
commit c96cb0a110
220 changed files with 29466 additions and 10769 deletions
@@ -0,0 +1,256 @@
#[path = "support/fake_provider.rs"]
mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::{collections::HashMap, sync::Arc};
use cursor_server::{
cursor::{
connect,
prompting::{PromptAssets, PromptCompiler},
proto::agent::v1 as pb,
CursorCommand, CursorSessionHandle, CursorSessionRegistry,
},
model::{ContentPart, MessageContent, ProjectedContent, Role},
provider::{FinishReason, ModelEvent},
};
use prost::Message;
const FOLLOW_UP: &str = "Perform any necessary follow-up actions in response to the subagent completion above. If no follow-up work is needed, no further action is required. If you mention an agent or subagent in your response, link it with the `[Name](id)` Don't use generic label such as `[agent]`, `[worker]`, or `[subagent]`.";
#[tokio::test]
async fn background_subagent_completion_starts_a_simulated_parent_turn() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(stop_response("model-call", "followed up"));
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry.get_or_create("completion-request").await.unwrap();
let (checkpoint, blobs) = drive_completion(
&handle,
completion_run(
"child-id",
"reusable-parent-run",
pb::ConversationStateStructure {
mode: Some(pb::AgentMode::Multitask as i32),
..Default::default()
},
),
)
.await;
let requests = provider.requests();
assert_eq!(requests.len(), 1);
let [runtime] = requests[0].history.as_slice() else {
panic!("completion Run must add exactly one runtime message")
};
assert_eq!(runtime.role, Role::User);
let ProjectedContent::Parts(parts) = &runtime.content else {
panic!("completion context must be text")
};
let [ContentPart::Text { text }] = parts.as_slice() else {
panic!("completion context must have one text part")
};
assert!(text.contains("Subagent ID: child-id"));
assert!(text.contains("child result"));
assert!(text.contains(FOLLOW_UP));
let messages = store
.load_current_messages(&cursor_server::model::ConversationId::new(
"parent-conversation",
))
.await
.unwrap();
assert!(messages.iter().any(|message| {
message.runtime_event_id.as_deref() == Some("run-request:completion-request")
&& matches!(&message.content, MessageContent::Parts { parts } if !parts.is_empty())
}));
let turn = pb::ConversationTurnStructure::decode(
blobs
.get(checkpoint.turns.last().expect("completion Turn"))
.expect("completion Turn Blob")
.as_slice(),
)
.unwrap();
let pb::conversation_turn_structure::Turn::AgentConversationTurn(turn) = turn.turn.unwrap()
else {
panic!("expected agent conversation Turn")
};
let user = pb::UserMessage::decode(
blobs
.get(&turn.user_message)
.expect("simulated UserMessage Blob")
.as_slice(),
)
.unwrap();
assert!(user.text.contains(FOLLOW_UP));
assert_eq!(user.is_simulated_msg, Some(true));
assert_eq!(
user.simulated_msg_reason,
Some(pb::SimulatedMsgReason::BackgroundTaskCompletion as i32)
);
assert_eq!(
user.simulated_message_metadata.unwrap().task_id.as_deref(),
Some("child-id")
);
provider.push(stop_response("model-call-2", "followed up again"));
let second = registry
.get_or_create("completion-request-2")
.await
.unwrap();
drive_completion(
&second,
completion_run("child-id-2", "reusable-parent-run-2", checkpoint),
)
.await;
let requests = provider.requests();
assert_eq!(requests.len(), 2);
let runtime_ids = requests[1]
.history
.iter()
.map(|message| message.message_id.as_str())
.filter(|id| id.starts_with("runtime:"))
.collect::<Vec<_>>();
assert_eq!(
runtime_ids,
[
"runtime:run-request:completion-request",
"runtime:run-request:completion-request-2"
]
);
}
async fn drive_completion(
handle: &CursorSessionHandle,
message: pb::AgentClientMessage,
) -> (pb::ConversationStateStructure, HashMap<Vec<u8>, Vec<u8>>) {
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(message),
})
.await
.unwrap();
let mut append_seqno = 1;
let mut blobs = HashMap::new();
let mut final_checkpoint = None;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
break;
}
let server = pb::AgentServerMessage::decode(payload).unwrap();
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
if let Some(pb::kv_server_message::Message::SetBlobArgs(set)) = &kv.message {
blobs.insert(set.blob_id.clone(), set.blob_data.clone());
}
handle
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(kv_ack(kv.id)),
})
.await
.unwrap();
append_seqno += 1;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state))
if state.pending_tool_calls.is_empty() =>
{
final_checkpoint = Some(state);
}
_ => {}
}
}
(
final_checkpoint.expect("settled completion checkpoint"),
blobs,
)
}
fn completion_run(
child_id: &str,
run_id: &str,
conversation_state: pb::ConversationStateStructure,
) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
action: Some(pb::ConversationAction {
action: Some(
pb::conversation_action::Action::BackgroundTaskCompletionAction(
pb::BackgroundTaskCompletionAction {
completions: vec![pb::BackgroundTaskCompletion {
task_id: child_id.into(),
kind: pb::BackgroundTaskKind::Subagent as i32,
status: pb::BackgroundTaskStatus::Success as i32,
title: "Inspect protocol".into(),
detail: Some("child result".into()),
output_path: Some("/tmp/child.jsonl".into()),
reason: pb::BackgroundTaskCompletionReason::TaskFinished as i32,
subagent_id: Some(child_id.into()),
tool_call_id: Some("task-call".into()),
..Default::default()
}],
},
),
),
..Default::default()
}),
conversation_id: Some("parent-conversation".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
conversation_state: Some(conversation_state),
run_id: Some(run_id.into()),
..Default::default()
},
)),
}
}
fn stop_response(model_call_id: &str, text: &str) -> Vec<ModelEvent> {
vec![
ModelEvent::Start {
model_call_id: model_call_id.into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta(text.into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]
}
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 },
)),
},
)),
}
}
+460 -50
View File
@@ -3,77 +3,487 @@ mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use std::{collections::HashSet, sync::Arc};
use cursor_server::{
cursor::{connect, proto::agent::v1 as pb},
prompting::{PromptAssets, PromptCompiler},
run::RunRegistry,
store::BlobId,
cursor::{
connect,
prompting::{PromptAssets, PromptCompiler},
proto::agent::v1 as pb,
CursorCommand, CursorSessionRegistry,
},
model::ToolRoundId,
provider::{FinishReason, ModelEvent},
store::{BlobEdge, BlobId},
};
use prost::Message;
#[tokio::test]
async fn checkpoint_dependency_is_not_ready_until_blob_ack() {
async fn checkpoint_dependencies_are_content_addressed_without_a_persistent_stream_outbox() {
let (_directory, store) = fixtures::temp_store().await;
store.create_pending_run("request").await.unwrap();
let id = BlobId::digest(b"tool result");
let key = format!("blob:{}", id.to_base64());
store
.enqueue_outbox("request", &key, "kv_set", b"tool result", &[])
.await
.unwrap();
assert!(!store
.dependencies_acked("request", std::slice::from_ref(&id))
.await
.unwrap());
assert!(store.ack_outbox("request", &key).await.unwrap());
assert!(store.dependencies_acked("request", &[id]).await.unwrap());
}
#[tokio::test]
async fn pending_blob_outbox_is_replayed_when_the_run_is_recreated() {
let (_directory, store) = fixtures::temp_store().await;
store.create_pending_run("recover-request").await.unwrap();
let data = b"durable tool result";
let blob_id = store.put_blob(data, &[]).await.unwrap();
store
.enqueue_outbox(
"recover-request",
&format!("blob:{}", blob_id.to_base64()),
"kv_set",
data,
&[],
let child = store.put_blob(b"message", &[]).await.unwrap();
let root = store
.put_blob(
b"checkpoint",
&[BlobEdge {
child: child.clone(),
field_name: "turns[0]".into(),
}],
)
.await
.unwrap();
assert_eq!(root, BlobId::digest(b"checkpoint"));
assert_eq!(store.get_blob(&child).await.unwrap().unwrap(), b"message");
let closure = store
.blob_closure(std::slice::from_ref(&root))
.await
.unwrap();
assert!(closure.contains(&root));
assert!(closure.contains(&child));
let outbox: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'outbox'",
)
.fetch_one(store.pool())
.await
.unwrap();
assert_eq!(outbox, 0);
}
#[tokio::test]
async fn eligible_pending_checkpoint_resumes_tools_before_the_next_model_call() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "model-1".into(),
},
ModelEvent::ToolCallStart {
index: 0,
call_id: "read-1".into(),
name: "Read".into(),
},
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{\"path\":\"/tmp/a\"}".into(),
},
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Done(FinishReason::ToolUse),
]);
provider.push(vec![
ModelEvent::Start {
model_call_id: "model-2".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("resumed".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
store,
Arc::new(fake_provider::FakeProvider::default()),
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider.clone()),
PromptCompiler::new(assets),
"test".into(),
Default::default(),
);
let handle = registry.get_or_create("recover-request").await.unwrap();
let mut output = handle.subscribe();
let frame = tokio::time::timeout(std::time::Duration::from_secs(2), output.recv())
let first = registry.get_or_create("first-run").await.unwrap();
let mut first_output = first.subscribe();
first
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(start_request()),
})
.await
.unwrap();
let mut first_seqno = 1;
let mut sent_blob_ids = HashSet::new();
let staged = loop {
let server = next_message(&mut first_output).await;
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
if let Some(pb::kv_server_message::Message::SetBlobArgs(args)) = &kv.message {
sent_blob_ids.insert(args.blob_id.clone());
}
acknowledge(&first, &mut first_seqno, kv.id).await;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state))
if state.pending_tool_calls.len() == 1 =>
{
break state;
}
_ => {}
}
};
assert!(
!sent_blob_ids.contains(
BlobId::digest(&staged.encode_to_vec())
.as_bytes()
.as_slice()
),
"ConversationStateStructure is inline and must not be sent as a Blob"
);
let staged_started_at_ms =
serde_json::from_str::<serde_json::Value>(staged.pending_tool_calls.first().unwrap())
.unwrap()["providerOptions"]["cursor"]["pendingToolCallStartedAtMs"]
.as_u64()
.unwrap();
assert_eq!(provider.requests().len(), 1);
first.cancel();
let resumed = registry.get_or_create("resumed-run").await.unwrap();
let mut resumed_output = resumed.subscribe();
let mut resumed_checkpoints = Vec::new();
let mut resumed_set_blob_ids = HashSet::new();
resumed
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(resume_request(staged.clone())),
})
.await
.unwrap();
let mut resumed_seqno = 1;
let exec_id = loop {
let server = next_message(&mut resumed_output).await;
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
if let Some(pb::kv_server_message::Message::SetBlobArgs(args)) = &kv.message {
resumed_set_blob_ids.insert(args.blob_id.clone());
}
acknowledge(&resumed, &mut resumed_seqno, kv.id).await;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state)) => {
resumed_checkpoints.push(state);
}
Some(pb::agent_server_message::Message::ExecServerMessage(exec)) => break exec.id,
_ => {}
}
};
assert_eq!(
provider.requests().len(),
1,
"resume must execute the pending batch before calling the model"
);
let resumed_round = store
.tool_round(&ToolRoundId::new("resumed-run:round:resume"))
.await
.unwrap()
.unwrap();
let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
let message = pb::AgentServerMessage::decode(payload).unwrap();
let Some(pb::agent_server_message::Message::KvServerMessage(kv)) = message.message else {
panic!("expected replayed KV SET")
assert_eq!(resumed_round.created_at_ms, staged_started_at_ms);
resumed
.command(CursorCommand::Append {
seqno: resumed_seqno,
message: Box::new(read_result(exec_id)),
})
.await
.unwrap();
resumed_seqno += 1;
let mut saw_settled_barrier_blob = false;
let mut saw_settled_checkpoint = false;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), resumed_output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
break;
}
let server = pb::AgentServerMessage::decode(payload).unwrap();
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
if let Some(pb::kv_server_message::Message::SetBlobArgs(args)) = &kv.message {
resumed_set_blob_ids.insert(args.blob_id.clone());
}
if !saw_settled_barrier_blob {
assert_eq!(
provider.requests().len(),
1,
"the next model call must wait for the settled Blob barrier"
);
saw_settled_barrier_blob = true;
}
acknowledge(&resumed, &mut resumed_seqno, kv.id).await;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state)) => {
if state.pending_tool_calls.is_empty() {
saw_settled_checkpoint = true;
}
resumed_checkpoints.push(state);
}
Some(pb::agent_server_message::Message::InteractionUpdate(update))
if matches!(
update.message,
Some(pb::interaction_update::Message::TextDelta(_))
) =>
{
assert!(
saw_settled_checkpoint,
"the next model round must not become visible before settled checkpoint"
);
}
_ => {}
}
}
assert!(saw_settled_barrier_blob);
assert!(saw_settled_checkpoint);
assert!(staged
.root_prompt_messages_json
.iter()
.all(|id| !resumed_set_blob_ids.contains(id)));
assert_eq!(provider.requests().len(), 2);
assert!(resumed_checkpoints
.last()
.unwrap()
.read_paths
.iter()
.any(|path| path == "/tmp/a"));
let mut previous_steps = Vec::new();
let mut saw_completed_read = false;
for state in resumed_checkpoints {
let Some(turn_id) = state.turns.last() else {
continue;
};
let turn = pb::ConversationTurnStructure::decode(
store
.get_blob(&BlobId::from_bytes(turn_id).unwrap())
.await
.unwrap()
.unwrap()
.as_slice(),
)
.unwrap();
let pb::conversation_turn_structure::Turn::AgentConversationTurn(turn) = turn.turn.unwrap()
else {
panic!("expected agent turn");
};
assert!(turn.steps.len() >= previous_steps.len());
assert_eq!(
previous_steps,
turn.steps[..previous_steps.len()],
"published Step BlobIDs must be an immutable prefix"
);
previous_steps = turn.steps.clone();
for step_id in &turn.steps {
let step = pb::ConversationStep::decode(
store
.get_blob(&BlobId::from_bytes(step_id).unwrap())
.await
.unwrap()
.unwrap()
.as_slice(),
)
.unwrap();
if let Some(pb::conversation_step::Message::ToolCall(call)) = step.message {
if call.tool_call_id.as_deref() == Some("read-1") {
assert!(call.started_at_ms.is_some());
assert!(call.completed_at_ms.is_some());
saw_completed_read = true;
}
}
}
}
assert!(
saw_completed_read,
"settled Turn must keep the typed result"
);
}
#[tokio::test]
async fn recovery_rejects_a_kv_get_payload_whose_hash_does_not_match_the_blob_id() {
let (_directory, store) = fixtures::temp_store().await;
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(fake_provider::FakeProvider::default()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry.get_or_create("bad-blob-run").await.unwrap();
let mut output = handle.subscribe();
let expected = BlobId::digest(b"expected");
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(resume_request(pb::ConversationStateStructure {
root_prompt_messages_json: vec![expected.as_bytes().to_vec()],
mode: Some(pb::AgentMode::Agent as i32),
..Default::default()
})),
})
.await
.unwrap();
let get_id = loop {
let server = next_message(&mut output).await;
if let Some(pb::agent_server_message::Message::KvServerMessage(kv)) = server.message {
if matches!(
kv.message,
Some(pb::kv_server_message::Message::GetBlobArgs(_))
) {
break kv.id;
}
}
};
let Some(pb::kv_server_message::Message::SetBlobArgs(set)) = kv.message else {
panic!("expected SetBlobArgs")
handle
.command(CursorCommand::Append {
seqno: 1,
message: Box::new(pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::KvClientMessage(
pb::KvClientMessage {
id: get_id,
message: Some(pb::kv_client_message::Message::GetBlobResult(
pb::GetBlobResult {
blob_data: Some(b"corrupt".to_vec()),
error: None,
},
)),
},
)),
}),
})
.await
.unwrap();
let error = loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
break serde_json::from_slice::<serde_json::Value>(&payload).unwrap();
}
};
assert_eq!(set.blob_id, blob_id.as_bytes());
assert_eq!(set.blob_data, data);
assert_eq!(error["error"]["code"], "invalid_argument");
assert!(error["error"]["message"]
.as_str()
.unwrap()
.contains("Blob hash mismatch"));
}
fn start_request() -> 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: "user-1".into(),
mode: pb::AgentMode::Agent as i32,
..Default::default()
}),
..Default::default()
},
)),
..Default::default()
}),
conversation_id: Some("conversation".into()),
run_id: Some("first-run".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
}
}
fn resume_request(state: pb::ConversationStateStructure) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
action: Some(pb::ConversationAction {
action: Some(pb::conversation_action::Action::ResumeAction(
pb::ResumeAction::default(),
)),
..Default::default()
}),
conversation_state: Some(state),
conversation_id: Some("conversation".into()),
run_id: Some("resumed-run".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
}
}
fn read_result(id: u32) -> pb::AgentClientMessage {
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ExecClientMessage(
pb::ExecClientMessage {
id,
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()
})),
},
)),
..Default::default()
},
)),
}
}
async fn next_message(
output: &mut tokio::sync::mpsc::UnboundedReceiver<bytes::Bytes>,
) -> pb::AgentServerMessage {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
assert_eq!(
flags & connect::END_STREAM_FLAG,
0,
"unexpected EndStream: {}",
String::from_utf8_lossy(&payload)
);
pb::AgentServerMessage::decode(payload).unwrap()
}
async fn acknowledge(
handle: &cursor_server::cursor::CursorSessionHandle,
seqno: &mut i64,
id: u32,
) {
handle
.command(CursorCommand::Append {
seqno: *seqno,
message: Box::new(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 },
)),
},
)),
}),
})
.await
.unwrap();
*seqno += 1;
}
+226
View File
@@ -0,0 +1,226 @@
#[path = "support/fake_provider.rs"]
mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use cursor_server::{
client::{session, ClientCommand, ClientEvent, CommitCause},
model::{
ConversationId, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId, RunKind,
ToolDefinition,
},
provider::{FinishReason, ModelEvent},
run::{RunEngine, RunOutcome},
};
use tokio_util::sync::CancellationToken;
#[tokio::test]
async fn a_client_without_checkpoint_protocol_runs_the_same_text_loop() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "call-1".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("hello".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let prepared = prepared(&store).await;
let (port, mut client) = session(32);
let engine = RunEngine::new(store.clone(), Arc::new(provider));
let run =
tokio::spawn(async move { engine.run(prepared, port, CancellationToken::new()).await });
let mut saw_final_commit = false;
while let Some(event) = client.events.recv().await {
match event {
ClientEvent::TextDelta(text) => assert_eq!(text, "hello"),
ClientEvent::StateCommitted(state) => {
saw_final_commit |= state.cause == CommitCause::FinalTurn;
state.barrier.complete(Ok(()));
}
ClientEvent::Ended(outcome) => {
assert_eq!(outcome, RunOutcome::Completed);
break;
}
_ => {}
}
}
assert!(saw_final_commit);
assert_eq!(run.await.unwrap(), RunOutcome::Completed);
}
#[tokio::test]
async fn a_failed_claim_cannot_overwrite_the_existing_run() {
let (_directory, store) = fixtures::temp_store().await;
let prepared = prepared(&store).await;
store.claim_run(&prepared).await.unwrap();
let (port, mut client) = session(8);
let outcome = RunEngine::new(
store.clone(),
Arc::new(fake_provider::FakeProvider::default()),
)
.run(prepared, port, CancellationToken::new())
.await;
assert!(matches!(outcome, RunOutcome::Failed(_)));
assert!(matches!(
client.events.recv().await,
Some(ClientEvent::Ended(RunOutcome::Failed(_)))
));
let status: String = sqlx::query_scalar("SELECT status FROM runs WHERE run_id = 'run'")
.fetch_one(store.pool())
.await
.unwrap();
assert_eq!(status, "running");
}
#[tokio::test]
async fn required_client_state_failure_prevents_a_completed_run() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "call-1".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("hello".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let prepared = prepared(&store).await;
let (port, mut client) = session(32);
let engine = RunEngine::new(store, Arc::new(provider));
let run =
tokio::spawn(async move { engine.run(prepared, port, CancellationToken::new()).await });
while let Some(event) = client.events.recv().await {
match event {
ClientEvent::StateCommitted(state) if state.cause == CommitCause::FinalTurn => {
state.barrier.complete(Err("snapshot failed".into()));
}
ClientEvent::StateCommitted(state) => state.barrier.complete(Ok(())),
ClientEvent::Ended(outcome) => {
assert_eq!(
outcome,
RunOutcome::Failed(cursor_server::run::RunFailure::Client(
"snapshot failed".into()
))
);
break;
}
_ => {}
}
}
assert_eq!(
run.await.unwrap(),
RunOutcome::Failed(cursor_server::run::RunFailure::Client(
"snapshot failed".into()
))
);
}
#[tokio::test]
async fn generic_engine_waits_for_every_tool_result_without_any_cursor_wire_id() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "call-1".into(),
},
tool_start(0, "A"),
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{}".into(),
},
ModelEvent::ToolCallEnd { index: 0 },
tool_start(1, "B"),
ModelEvent::ToolCallArgumentsDelta {
index: 1,
delta: "{}".into(),
},
ModelEvent::ToolCallEnd { index: 1 },
ModelEvent::Done(FinishReason::ToolUse),
]);
provider.push(vec![
ModelEvent::Start {
model_call_id: "call-2".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("done".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let prepared = prepared(&store).await;
let (port, mut client) = session(64);
let commands = client.commands.clone();
let engine = RunEngine::new(store, Arc::new(provider));
let run =
tokio::spawn(async move { engine.run(prepared, port, CancellationToken::new()).await });
while let Some(event) = client.events.recv().await {
match event {
ClientEvent::ExecuteToolRound { calls, .. } => {
assert_eq!(calls.len(), 2);
commands
.send(ClientCommand::ToolResult {
call_id: "B".into(),
content: "result-B".into(),
is_error: false,
})
.await
.unwrap();
commands
.send(ClientCommand::ToolResult {
call_id: "A".into(),
content: "result-A".into(),
is_error: false,
})
.await
.unwrap();
}
ClientEvent::StateCommitted(state) => state.barrier.complete(Ok(())),
ClientEvent::Ended(outcome) => {
assert_eq!(outcome, RunOutcome::Completed);
break;
}
_ => {}
}
}
assert_eq!(run.await.unwrap(), RunOutcome::Completed);
}
async fn prepared(store: &cursor_server::store::Store) -> PreparedRun {
let conversation_id = ConversationId::new("conversation");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
PreparedRun {
run_id: RunId::new("run"),
conversation_id,
kind: RunKind::Root,
model: ModelSpec::new("model"),
prompt: PromptSpec {
instructions: "system".into(),
tools: vec![ToolDefinition {
name: "Tool".into(),
description: "test".into(),
parameters: serde_json::json!({"type":"object"}),
}],
},
selected_subagent_models: Vec::new(),
subagent_model_overrides: Vec::new(),
initial_messages: vec![fixtures::user("user", "hello")],
action: RunAction::Start,
base_revision_id: root,
}
}
fn tool_start(index: usize, call_id: &str) -> ModelEvent {
ModelEvent::ToolCallStart {
index,
call_id: call_id.into(),
name: "Tool".into(),
}
}
+6 -5
View File
@@ -13,12 +13,12 @@ use axum::{
};
use base64::{engine::general_purpose::STANDARD_NO_PAD, Engine};
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::CursorSessionRegistry,
cursor::{
connect, handlers,
proto::{agent::v1 as pb, aiserver::v1 as ai},
},
prompting::{PromptAssets, PromptCompiler},
run::RunRegistry,
};
use flate2::{write::GzEncoder, Compression};
use prost::Message;
@@ -96,15 +96,15 @@ async fn bidi_append_gzip_body_is_decompressed_before_protobuf_decode() {
let (_directory, store) = fixtures::temp_store().await;
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store,
Arc::new(fake_provider::FakeProvider::default()),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let wire = ai::BidiAppendRequest {
request_id: Some(ai::BidiRequestId {
@@ -118,6 +118,7 @@ async fn bidi_append_gzip_body_is_decompressed_before_protobuf_decode() {
let compressed = encoder.finish().unwrap();
let response = handlers::router(registry)
.unwrap()
.oneshot(
Request::post("/aiserver.v1.BidiService/BidiAppend")
.header(header::CONTENT_TYPE, "application/proto")
+43 -25
View File
@@ -7,39 +7,39 @@ use std::sync::Arc;
use base64::{engine::general_purpose::STANDARD_NO_PAD, Engine};
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::{
connect,
proto::{agent::v1 as pb, aiserver::v1 as ai},
},
cursor::{CursorCommand, CursorSessionRegistry},
model::{MessageContent, Role},
prompting::{PromptAssets, PromptCompiler},
provider::{FinishReason, ResponseEvent},
run::{RunCommand, RunRegistry},
provider::{FinishReason, ModelEvent},
Error,
};
use prost::Message;
#[tokio::test]
async fn provider_failure_checkpoints_then_returns_structured_error_and_closes() {
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_error(Error::Provider("provider failed".into()));
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry.get_or_create("failed-request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run()),
})
@@ -47,7 +47,7 @@ async fn provider_failure_checkpoints_then_returns_structured_error_and_closes()
.unwrap();
let mut append_seqno = 1;
let mut saw_checkpoint = false;
let mut checkpoints = Vec::new();
let mut saw_turn_ended = false;
let error_json = loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
@@ -62,7 +62,7 @@ async fn provider_failure_checkpoints_then_returns_structured_error_and_closes()
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(kv_ack(kv.id)),
})
@@ -70,8 +70,8 @@ async fn provider_failure_checkpoints_then_returns_structured_error_and_closes()
.unwrap();
append_seqno += 1;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(_)) => {
saw_checkpoint = true;
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state)) => {
checkpoints.push(state);
}
Some(pb::agent_server_message::Message::InteractionUpdate(update)) => {
if matches!(
@@ -88,7 +88,12 @@ async fn provider_failure_checkpoints_then_returns_structured_error_and_closes()
}
};
assert!(saw_checkpoint);
assert_eq!(
checkpoints.len(),
1,
"the initial user state is checkpointed"
);
assert!(checkpoints[0].pending_tool_calls.is_empty());
assert!(!saw_turn_ended);
assert_eq!(error_json["error"]["code"], "unavailable");
let detail = &error_json["error"]["details"][0];
@@ -112,7 +117,12 @@ async fn provider_failure_checkpoints_then_returns_structured_error_and_closes()
None
);
let messages = store.load_messages("failed-conversation").await.unwrap();
let messages = store
.load_current_messages(&cursor_server::model::ConversationId::new(
"failed-conversation",
))
.await
.unwrap();
assert!(messages.iter().any(|message| message.role == Role::User));
assert!(!messages.iter().any(|message| {
matches!(
@@ -127,32 +137,32 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ResponseEvent::Start {
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ResponseEvent::ToolCallStart {
ModelEvent::ToolCallStart {
index: 0,
call_id: "call-1".into(),
name: "Read".into(),
},
ResponseEvent::ToolCallArgumentsDelta {
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{\"path\":\"/tmp/a\"}".into(),
},
ResponseEvent::ToolCallEnd { index: 0 },
ResponseEvent::Done(FinishReason::ToolUse),
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Done(FinishReason::ToolUse),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry
.get_or_create("protocol-failed-request")
@@ -160,7 +170,7 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
.unwrap();
let mut output = handle.subscribe();
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(protocol_client_run()),
})
@@ -182,7 +192,7 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(kv_ack(kv.id)),
})
@@ -193,7 +203,7 @@ async fn runtime_protocol_failure_returns_connect_error_end_stream_and_closes()
Some(pb::agent_server_message::Message::ExecServerMessage(exec)) => {
// An unknown numeric bridge id is a runtime protocol error.
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ExecClientMessage(
@@ -260,6 +270,10 @@ fn client_run() -> pb::AgentClientMessage {
}),
conversation_id: Some("failed-conversation".into()),
run_id: Some("failed-request".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
@@ -286,6 +300,10 @@ fn protocol_client_run() -> pb::AgentClientMessage {
}),
conversation_id: Some("protocol-failed-conversation".into()),
run_id: Some("protocol-failed-request".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
+101 -24
View File
@@ -6,27 +6,91 @@ mod fixtures;
use std::sync::Arc;
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::{connect, proto::agent::v1 as pb},
prompting::{PromptAssets, PromptCompiler},
provider::{FinishReason, ResponseEvent},
run::{RunCommand, RunRegistry},
cursor::{CursorCommand, CursorSessionRegistry},
model::{ConversationId, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId, RunKind},
provider::{FinishReason, ModelEvent},
run::RunRegistry,
store::RunStatus,
};
use prost::Message;
use tokio_util::sync::CancellationToken;
#[tokio::test]
async fn a_new_revision_invalidates_late_events_from_the_old_run() {
async fn generic_run_registry_cancels_the_previous_client_for_a_conversation() {
let registry = RunRegistry::default();
let conversation = cursor_server::model::ConversationId::new("conversation");
let first = CancellationToken::new();
let second = CancellationToken::new();
registry
.activate(
conversation.clone(),
cursor_server::model::RunId::new("first"),
first.clone(),
)
.await;
registry
.activate(
conversation.clone(),
cursor_server::model::RunId::new("second"),
second.clone(),
)
.await;
assert!(first.is_cancelled());
assert!(!second.is_cancelled());
registry
.release(&conversation, &cursor_server::model::RunId::new("first"))
.await;
registry.shutdown().await;
assert!(second.is_cancelled());
}
#[tokio::test]
async fn a_replaced_run_cannot_overwrite_its_cancelled_status() {
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);
let conversation_id = ConversationId::new("conversation");
let base_revision_id = store.ensure_conversation(&conversation_id).await.unwrap();
let prepared = |run_id: &str| PreparedRun {
run_id: RunId::new(run_id),
conversation_id: conversation_id.clone(),
kind: RunKind::Root,
model: ModelSpec::new("model"),
prompt: PromptSpec {
instructions: String::new(),
tools: Vec::new(),
},
selected_subagent_models: Vec::new(),
subagent_model_overrides: Vec::new(),
initial_messages: Vec::new(),
action: RunAction::Resume {
pending_tool_round: None,
},
base_revision_id,
};
let first = prepared("first");
let second = prepared("second");
store.claim_run(&first).await.unwrap();
store.claim_run(&second).await.unwrap();
assert!(!store
.revision_is_current("conversation", first)
.finish_run(&first.run_id, RunStatus::Completed, None, None,)
.await
.unwrap());
assert!(store
.revision_is_current("conversation", second)
let status: String = sqlx::query_scalar("SELECT status FROM runs WHERE run_id = 'first'")
.fetch_one(store.pool())
.await
.unwrap());
.unwrap();
let active: Option<String> = sqlx::query_scalar(
"SELECT active_run_id FROM conversations WHERE conversation_id = 'conversation'",
)
.fetch_one(store.pool())
.await
.unwrap();
assert_eq!(status, "cancelled");
assert_eq!(active.as_deref(), Some("second"));
}
#[tokio::test]
@@ -34,15 +98,15 @@ 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")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store,
Arc::new(fake_provider::FakeProvider::default()),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry.get_or_create("active-run").await.unwrap();
let mut output = handle.subscribe();
@@ -63,34 +127,37 @@ 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 {
ModelEvent::Start {
model_call_id: "ignored".into(),
},
ModelEvent::ToolCallStart {
index: 0,
call_id: "call-1".into(),
name: "Read".into(),
},
ResponseEvent::ToolCallArgumentsDelta {
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{\"path\":\"/tmp/a\"}".into(),
},
ResponseEvent::ToolCallEnd { index: 0 },
ResponseEvent::Done(FinishReason::ToolUse),
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Done(FinishReason::ToolUse),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry.get_or_create("cancel-request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run()),
})
@@ -103,12 +170,18 @@ async fn cancel_aborts_active_exec_before_canceled_end_stream() {
.await
.unwrap()
.unwrap();
let (_, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
assert_eq!(
flags & connect::END_STREAM_FLAG,
0,
"Run ended before Exec: {}",
String::from_utf8_lossy(&payload)
);
let server = pb::AgentServerMessage::decode(payload).unwrap();
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(kv_ack(kv.id)),
})
@@ -170,6 +243,10 @@ fn client_run() -> pb::AgentClientMessage {
}),
conversation_id: Some("cancel-conversation".into()),
run_id: Some("cancel-request".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
+107
View File
@@ -0,0 +1,107 @@
use std::time::Duration;
use axum::{http::header, response::IntoResponse, routing::post, Router};
use cursor_server::{
model::{
ModelInvocation, ModelRequest, ModelSpec, PromptSpec, ProviderEndpointInput,
ProviderModelInput, ProviderType,
},
provider::{ModelEvent, Provider, ProviderRouter},
store::Store,
};
use futures_util::StreamExt;
use tokio_util::sync::CancellationToken;
#[tokio::test]
async fn records_one_summary_and_raw_payloads_for_one_provider_request() {
let app = Router::new().route(
"/v1/chat/completions",
post(|| async {
(
[(header::CONTENT_TYPE, "text/event-stream")],
concat!(
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2,\"total_tokens\":12}}\n\n",
"data: [DONE]\n\n"
),
)
.into_response()
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let directory = tempfile::tempdir().unwrap();
let store = Store::connect(&format!(
"sqlite://{}",
directory.path().join("observability.db").display()
))
.await
.unwrap();
store.set_detailed_logging(true).await.unwrap();
let endpoint = store
.create_provider(&ProviderEndpointInput {
name: "test".into(),
provider_type: ProviderType::OpenAiChat,
base_url: format!("http://{address}/v1"),
api_key: Some("not-recorded".into()),
custom_headers: serde_json::json!({"x-safe":"visible","authorization":"hidden"}),
})
.await
.unwrap();
let model = store
.save_provider_model(
endpoint.provider_id,
&ProviderModelInput {
model_id: "actual-model".into(),
display_name: "Display Model".into(),
enabled: true,
sort_order: 0,
context_window_tokens: None,
max_output_tokens: None,
reasoning_enabled: false,
reasoning_effort: None,
extra_params: serde_json::json!({}),
},
)
.await
.unwrap();
let provider = ProviderRouter::new(store.clone(), Duration::from_secs(5));
let events = provider
.stream(
ModelInvocation {
call_id: "call-1".into(),
run_id: "run-1".into(),
conversation_id: "conversation-1".into(),
provider_call_index: 0,
request: ModelRequest {
prompt: PromptSpec {
instructions: "system".into(),
tools: Vec::new(),
},
model: ModelSpec::new(model.model_hash),
history: Vec::new(),
},
},
CancellationToken::new(),
)
.collect::<Vec<_>>()
.await;
assert!(events.iter().all(Result::is_ok));
assert!(events
.iter()
.any(|event| matches!(event, Ok(ModelEvent::Done(_)))));
let call = store.llm_call("call-1").await.unwrap().unwrap();
assert_eq!(call.status, "completed");
assert_eq!(call.total_tokens, Some(12));
assert!(call.ttfb_ms.is_some());
assert!(call.ttft_ms.is_some());
let request = store.llm_call_request("call-1").await.unwrap().unwrap();
assert_eq!(request.body["model"], "actual-model");
assert_eq!(request.headers["x-safe"], "visible");
assert!(request.headers.get("authorization").is_none());
assert!(!store.llm_call_chunks("call-1").await.unwrap().is_empty());
server.abort();
}
+400 -33
View File
@@ -1,10 +1,17 @@
#[path = "support/fixtures.rs"]
mod fixtures;
use std::collections::BTreeMap;
use cursor_server::{
model::{CanonicalMessage, MessageContent, Origin, Role, ToolCallContent, ToolResultContent},
prompting::{project_messages, Mode, PromptAssets, PromptCompiler, ToolDefinition},
cursor::prompting::{Mode, PromptAssets, PromptCompiler},
model::{project_messages, ProjectedContent},
model::{
CanonicalMessage, MessageContent, Origin, Role, ToolCallContent, ToolDefinition,
ToolResultContent,
},
};
use sha2::{Digest, Sha256};
#[test]
fn projecting_an_append_only_context_preserves_the_complete_prefix() {
@@ -25,15 +32,18 @@ fn every_tool_result_is_projected_as_string_content() {
];
let projected = project_messages(&messages).unwrap();
let object_text = projected[0]
.content
.as_str()
.expect("object ToolResult must be JSON-encoded into a string");
let ProjectedContent::ToolResult(object_result) = &projected[0].content else {
panic!("expected tool result")
};
let object_text = &object_result.content;
assert_eq!(
serde_json::from_str::<serde_json::Value>(object_text).unwrap(),
object
);
assert_eq!(projected[1].content.as_str(), Some("plain text"));
let ProjectedContent::ToolResult(string_result) = &projected[1].content else {
panic!("expected tool result")
};
assert_eq!(string_result.content, "plain text");
}
#[test]
@@ -45,15 +55,19 @@ fn assistant_text_and_thinking_remain_separate_during_projection() {
content: MessageContent::Assistant {
text: "visible answer".into(),
thinking: "private reasoning".into(),
model_call_id: Some("model-call".into()),
tool_round_id: Some("round".into()),
replay_state: None,
tool_calls: Vec::new(),
},
runtime_event_id: None,
}];
let projected = project_messages(&messages).unwrap();
assert_eq!(projected[0].content.as_str(), Some("visible answer"));
assert_eq!(projected[0].thinking.as_deref(), Some("private reasoning"));
let ProjectedContent::Assistant { text, thinking, .. } = &projected[0].content else {
panic!("expected assistant")
};
assert_eq!(text, "visible answer");
assert_eq!(thinking, "private reasoning");
}
#[test]
@@ -75,61 +89,409 @@ fn split_tool_pairs_reconstruct_the_original_provider_assistant_message() {
let projected = project_messages(&messages).unwrap();
assert_eq!(projected.len(), 3);
assert_eq!(projected[0].role, "assistant");
assert_eq!(projected[0].thinking.as_deref(), Some("complete reasoning"));
let calls = projected[0].tool_calls.as_ref().unwrap();
assert_eq!(calls[0]["id"], "call-first");
assert_eq!(calls[1]["id"], "call-second");
assert_eq!(projected[1].tool_call_id.as_deref(), Some("call-first"));
assert_eq!(projected[2].tool_call_id.as_deref(), Some("call-second"));
assert_eq!(projected[0].role, Role::Assistant);
let ProjectedContent::Assistant {
thinking, calls, ..
} = &projected[0].content
else {
panic!("expected assistant")
};
assert_eq!(thinking, "complete reasoning");
assert_eq!(calls[0].call_id, "call-first");
assert_eq!(calls[1].call_id, "call-second");
let ProjectedContent::ToolResult(second) = &projected[1].content else {
panic!("expected tool result")
};
let ProjectedContent::ToolResult(first) = &projected[2].content else {
panic!("expected tool result")
};
assert_eq!(second.call_id, "call-second");
assert_eq!(first.call_id, "call-first");
}
#[test]
fn every_prompt_mode_loads_the_captured_tool_set() {
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
assert_eq!(assets.mode(Mode::Agent).tools.len(), 20);
assert_eq!(assets.mode(Mode::Ask).tools.len(), 18);
assert_eq!(assets.mode(Mode::Plan).tools.len(), 16);
assert_eq!(assets.mode(Mode::Debug).tools.len(), 18);
assert_eq!(assets.mode(Mode::Multitask).tools.len(), 20);
assert_eq!(assets.mode(Mode::Subagent).tools.len(), 4);
assert_eq!(
assets
.mode(Mode::Agent)
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<Vec<_>>(),
vec![
"Shell",
"Grep",
"Delete",
"WebSearch",
"WebFetch",
"GenerateImage",
"EditNotebook",
"TodoWrite",
"StrReplace",
"Write",
"Read",
"ReadLints",
"Glob",
"AskQuestion",
"Task",
"AwaitShell",
"GetMcpTools",
"FetchMcpResource",
"SwitchMode",
"CallMcpTool",
]
);
assert_mode(
&assets,
Mode::Ask,
&[
"AskQuestion",
"CallMcpTool",
"Delete",
"FetchMcpResource",
"Glob",
"Grep",
"Read",
"ReadLints",
"Shell",
"StrReplace",
"Task",
"TodoWrite",
"WebFetch",
"WebSearch",
"Write",
],
"1ca80a9b5547546217ac4999cd5069225c39cbd2df48e8468bb0105c18215236",
);
assert_mode(
&assets,
Mode::Plan,
&[
"Shell",
"Glob",
"Grep",
"Read",
"TodoWrite",
"ReadLints",
"WebSearch",
"WebFetch",
"AskQuestion",
"CreatePlan",
"Task",
"FetchMcpResource",
"CallMcpTool",
],
"06f9137d74d06edd03d90cb4124533d127be439eb2a480be4760944cd7c36d85",
);
assert_mode(
&assets,
Mode::Debug,
&[
"AskQuestion",
"CallMcpTool",
"Delete",
"FetchMcpResource",
"Glob",
"Grep",
"Read",
"ReadLints",
"Shell",
"StrReplace",
"Task",
"TodoWrite",
"WebFetch",
"WebSearch",
"Write",
],
"1ca80a9b5547546217ac4999cd5069225c39cbd2df48e8468bb0105c18215236",
);
assert_mode(
&assets,
Mode::Multitask,
&[
"AskQuestion",
"CallMcpTool",
"Delete",
"FetchMcpResource",
"Glob",
"Grep",
"Read",
"ReadLints",
"Shell",
"StrReplace",
"SwitchMode",
"Task",
"TodoWrite",
"WebFetch",
"WebSearch",
"Write",
"GenerateImage",
],
"b5b1310e3012b93f3b32e1cd834dad4865db1744200f362922c28fb0735aa86f",
);
assert_mode(
&assets,
Mode::Subagent,
&[
"Shell",
"Grep",
"Delete",
"WebSearch",
"WebFetch",
"GenerateImage",
"ReadLints",
"EditNotebook",
"TodoWrite",
"StrReplace",
"Write",
"Read",
"Glob",
"Task",
"AwaitShell",
"GetMcpTools",
"FetchMcpResource",
"SwitchMode",
"UpdateCurrentStep",
"CallMcpTool",
],
"142a76a2d0764585ebfd4c1c9b335a22d859e5a1d2ef515cc513923d8c7138b6",
);
assert_mode(
&assets,
Mode::Compaction,
&[],
"4f53cda18c2baa0c0354bb5f9a3ecbe5ed12ab4d8e11ba873c2f11161202b945",
);
assert_eq!(
schema_digest(&assets.mode(Mode::Agent).tools),
"7e9055bf2ae45384762a417b11a82db7a224d71c3fe716b7af6f06d589319335"
);
let shell = assets
.mode(Mode::Agent)
.tools
.iter()
.find(|tool| tool.name == "Shell")
.unwrap();
assert!(
shell.parameters["properties"]["block_until_ms"]["description"]
.as_str()
.unwrap()
.contains("do not combine it with `nohup`, `&`, `disown`")
);
for mode in [
Mode::Agent,
Mode::Ask,
Mode::Debug,
Mode::Multitask,
Mode::Subagent,
Mode::Compaction,
] {
assert!(!assets
.mode(mode)
.tools
.iter()
.any(|tool| tool.name == "CreatePlan" || tool.name == "PatchEdit"));
}
}
#[test]
fn every_captured_mode_owns_and_renders_its_runtime_template() {
let compiler = PromptCompiler::new(
PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap(),
);
let values = BTreeMap::from([
("REQUEST_CONTEXT", String::new()),
("OPEN_FILES", String::new()),
("SELECTED_CONTEXT", String::new()),
("ACTION_CONTEXT", String::new()),
("TIMESTAMP", "Sunday, Aug 16, 2026, 11:31 PM (UTC+8)".into()),
("USER_QUERY", "question".into()),
("DEBUG_SERVER_ENDPOINT", "http://debug".into()),
("DEBUG_LOG_PATH", "/tmp/debug.log".into()),
("DEBUG_SESSION_ID", "session".into()),
]);
for (mode, marker) in [
(Mode::Agent, "You are still in **Agent Mode**"),
(Mode::Ask, "Ask mode is active."),
(Mode::Plan, "Plan mode is active."),
(Mode::Debug, "You are now in **DEBUG MODE**"),
(Mode::Multitask, "The user has engaged **Multitask Mode**"),
] {
let rendered = compiler.runtime_message(mode, &values).unwrap();
assert!(rendered.contains(marker), "missing {mode:?} marker");
assert!(rendered.contains("<user_query>\nquestion\n</user_query>"));
assert_eq!(rendered.matches("<user_query>").count(), 1);
}
}
fn assert_mode(assets: &PromptAssets, mode: Mode, expected: &[&str], digest: &str) {
assert_eq!(
assets
.mode(mode)
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<Vec<_>>(),
expected
);
assert_eq!(schema_digest(&assets.mode(mode).tools), digest);
}
fn schema_digest(tools: &[ToolDefinition]) -> String {
hex::encode(Sha256::digest(serde_json::to_vec(tools).unwrap()))
}
#[test]
fn dynamic_mcp_tools_are_appended_after_the_stable_mode_tool_prefix() {
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let compiler = PromptCompiler::new(assets);
let messages = vec![fixtures::user("u1", "one")];
let base = compiler
.compile(Mode::Agent, "model", "call", &messages)
.prompt_spec(Mode::Agent, "model", &[], false)
.unwrap();
let dynamic = compiler
.compile_with_dynamic_tools(
.prompt_spec(
Mode::Agent,
"model",
"call",
&messages,
&[ToolDefinition {
name: "mcp_repo_lookup".into(),
description: "lookup".into(),
input_schema: serde_json::json!({"type": "object"}),
parameters: serde_json::json!({"type": "object"}),
}],
false,
)
.unwrap();
assert_eq!(base.tools, dynamic.tools[..base.tools.len()]);
assert_eq!(dynamic.tools.last().unwrap().name, "mcp_repo_lookup");
}
#[test]
fn dynamic_mcp_tool_cannot_replace_a_mode_tool() {
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let compiler = PromptCompiler::new(assets);
let error = compiler
.prompt_spec(
Mode::Agent,
"model",
&[ToolDefinition {
name: "Read".into(),
description: "replacement".into(),
parameters: serde_json::json!({"type": "object"}),
}],
false,
)
.unwrap_err();
assert!(error
.to_string()
.contains("dynamic MCP tool conflicts with a mode tool: Read"));
}
#[test]
fn agent_system_prompt_is_static_and_substitutes_the_model_name() {
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let compiler = PromptCompiler::new(assets);
let request = compiler
.prompt_spec(Mode::Agent, "test-model", &[], false)
.unwrap();
let prompt = &request.instructions;
assert!(prompt.contains("powered by test-model"));
assert!(!prompt.contains("{{FAKE_MODEL_NAME}}"));
assert!(!prompt.contains("<user_info>"));
}
#[test]
fn subagent_uses_the_agent_prompt_and_only_the_captured_tool_delta() {
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let compiler = PromptCompiler::new(assets);
let agent_prompt = compiler
.prompt_spec(Mode::Agent, "model", &[], false)
.unwrap();
let subagent_prompt = compiler
.prompt_spec(Mode::Subagent, "model", &[], false)
.unwrap();
assert_eq!(agent_prompt.instructions, subagent_prompt.instructions);
let request = compiler
.prompt_spec(Mode::Subagent, "model", &[], false)
.unwrap();
assert_eq!(
request
.tools
.iter()
.map(|tool| tool.name.as_str())
.collect::<Vec<_>>(),
vec![
"Shell",
"Grep",
"Delete",
"WebSearch",
"WebFetch",
"GenerateImage",
"ReadLints",
"EditNotebook",
"TodoWrite",
"StrReplace",
"Write",
"Read",
"Glob",
"Task",
"AwaitShell",
"GetMcpTools",
"FetchMcpResource",
"SwitchMode",
"UpdateCurrentStep",
"CallMcpTool",
]
);
let task = request
.tools
.iter()
.find(|tool| tool.name == "Task")
.unwrap();
let properties = task.parameters["properties"].as_object().unwrap();
assert!(!properties.contains_key("environment"));
assert!(!properties.contains_key("cloud_base_branch"));
let suppressed = compiler
.prompt_spec(Mode::Subagent, "model", &[], true)
.unwrap();
assert!(!suppressed
.tools
.iter()
.any(|tool| tool.name == "UpdateCurrentStep"));
}
fn tool_result(id: &str, output: serde_json::Value) -> CanonicalMessage {
tool_result_with_call(id, &format!("call-{id}"), output)
}
@@ -139,6 +501,7 @@ fn tool_result_with_call(
call_id: &str,
output: impl Into<serde_json::Value>,
) -> CanonicalMessage {
let output = output.into();
CanonicalMessage {
message_id: id.into(),
role: Role::Tool,
@@ -146,7 +509,10 @@ fn tool_result_with_call(
content: MessageContent::ToolResult(ToolResultContent {
call_id: call_id.into(),
name: "Tool".into(),
output: output.into(),
content: output
.as_str()
.map(str::to_string)
.unwrap_or_else(|| output.to_string()),
is_error: false,
}),
runtime_event_id: None,
@@ -155,7 +521,7 @@ fn tool_result_with_call(
fn assistant_tool_pair(
id: &str,
model_call_id: &str,
tool_round_id: &str,
index: usize,
call_id: &str,
text: &str,
@@ -168,7 +534,8 @@ fn assistant_tool_pair(
content: MessageContent::Assistant {
text: text.into(),
thinking: thinking.into(),
model_call_id: Some(model_call_id.into()),
tool_round_id: Some(tool_round_id.into()),
replay_state: None,
tool_calls: vec![ToolCallContent {
index,
call_id: call_id.into(),
+157
View File
@@ -0,0 +1,157 @@
use cursor_server::{
model::{NewLlmCall, ProviderEndpointInput, ProviderModelInput, ProviderType, Usage},
store::Store,
};
async fn store() -> (tempfile::TempDir, Store) {
let directory = tempfile::tempdir().unwrap();
let url = format!("sqlite://{}", directory.path().join("test.db").display());
let store = Store::connect(&url).await.unwrap();
(directory, store)
}
#[tokio::test]
async fn provider_secret_is_write_only_and_model_hash_is_stable() {
let (_directory, store) = store().await;
let provider = store
.create_provider(&ProviderEndpointInput {
name: "Local".into(),
provider_type: ProviderType::OpenAiChat,
base_url: "https://example.com/v1/".into(),
api_key: Some("secret".into()),
custom_headers: serde_json::json!({"x-route":"one", "authorization":"header-secret"}),
})
.await
.unwrap();
assert!(provider.has_api_key);
assert!(!serde_json::to_string(&provider).unwrap().contains("secret"));
assert_eq!(
provider.custom_headers["authorization"],
serde_json::Value::Null
);
let updated = store
.update_provider(
provider.provider_id,
&ProviderEndpointInput {
name: "Renamed".into(),
provider_type: provider.provider_type,
base_url: provider.base_url.clone(),
api_key: None,
custom_headers: provider.custom_headers.clone(),
},
)
.await
.unwrap();
assert_eq!(updated.name, "Renamed");
assert_eq!(
store
.provider(provider.provider_id)
.await
.unwrap()
.unwrap()
.custom_headers["authorization"],
"header-secret"
);
let model = store
.save_provider_model(
provider.provider_id,
&ProviderModelInput {
model_id: "model-a".into(),
display_name: "Model A".into(),
enabled: true,
sort_order: 0,
context_window_tokens: Some(128_000),
max_output_tokens: Some(8_192),
reasoning_enabled: false,
reasoning_effort: None,
extra_params: serde_json::json!({"temperature":0}),
},
)
.await
.unwrap();
assert_eq!(model.model_hash, "f246010a");
}
#[tokio::test]
async fn call_summary_is_always_stored_and_payloads_follow_detailed_setting() {
let (_directory, store) = store().await;
let provider = store
.create_provider(&ProviderEndpointInput {
name: "Local".into(),
provider_type: ProviderType::OpenAiChat,
base_url: "https://example.com/v1".into(),
api_key: None,
custom_headers: serde_json::json!({}),
})
.await
.unwrap();
let model = store
.save_provider_model(
provider.provider_id,
&ProviderModelInput {
model_id: "model-a".into(),
display_name: "Model A".into(),
enabled: true,
sort_order: 0,
context_window_tokens: None,
max_output_tokens: None,
reasoning_enabled: false,
reasoning_effort: None,
extra_params: serde_json::json!({}),
},
)
.await
.unwrap();
let call = NewLlmCall {
call_id: "call-1".into(),
run_id: "run-1".into(),
conversation_id: "conversation-1".into(),
provider_call_index: 0,
model_hash: model.model_hash,
provider_type: ProviderType::OpenAiChat,
provider_url: provider.base_url,
model_id: model.model_id,
display_name: model.display_name,
message_count: 2,
tool_count: 3,
detailed: false,
};
store.start_llm_call(&call).await.unwrap();
store
.record_llm_request(
"call-1",
&serde_json::json!({}),
&serde_json::json!({"model":"model-a"}),
false,
)
.await
.unwrap();
store
.record_llm_chunk("call-1", 0, 4, b"data", false)
.await
.unwrap();
store
.record_llm_usage(
"call-1",
Usage {
input_tokens: Some(10),
output_tokens: Some(5),
total_tokens: Some(15),
..Default::default()
},
)
.await
.unwrap();
store
.finish_llm_call("call-1", "completed", Some("stop"), 9, None, None)
.await
.unwrap();
let summary = store.llm_call("call-1").await.unwrap().unwrap();
assert_eq!(summary.total_tokens, Some(15));
assert_eq!(summary.request_bytes, Some(19));
assert_eq!(summary.response_bytes, 4);
assert!(store.llm_call_request("call-1").await.unwrap().is_none());
assert!(store.llm_call_chunks("call-1").await.unwrap().is_empty());
}
+499
View File
@@ -0,0 +1,499 @@
use cursor_server::{
client::ClientEvent,
config::{ProviderConfig, ProviderKind},
model::{
ContentPart, ModelInvocation, ModelRequest, ModelSpec, ProjectedContent, ProjectedMessage,
PromptSpec, Role, Usage,
},
provider::{
FinishReason, ModelEvent, OpenAiChatProvider, OpenAiResponsesProvider, Provider,
ProviderStream,
},
run::{consume_model_cycle, RunFailure},
};
use futures_util::{stream, StreamExt};
use serde_json::Value;
use std::{sync::Arc, time::Duration};
use tokio_util::sync::CancellationToken;
fn provider_stream(events: Vec<ModelEvent>) -> ProviderStream {
Box::pin(stream::iter(events.into_iter().map(Ok)))
}
#[tokio::test]
async fn complete_tool_stream_is_validated_and_projected() {
let (sender, mut receiver) = tokio::sync::mpsc::channel(32);
let result = consume_model_cycle(
provider_stream(vec![
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ModelEvent::ThinkingStart,
ModelEvent::ThinkingDelta("why".into()),
ModelEvent::ThinkingEnd,
ModelEvent::ToolCallStart {
index: 0,
call_id: "call".into(),
name: "Read".into(),
},
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: r#"{"path":"/tmp/a"}"#.into(),
},
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Usage(Usage {
input_tokens: Some(10),
output_tokens: Some(2),
..Usage::default()
}),
ModelEvent::Done(FinishReason::ToolUse),
]),
&sender,
&CancellationToken::new(),
)
.await
.unwrap();
drop(sender);
assert_eq!(result.reasoning, "why");
assert_eq!(result.calls[0].arguments["path"], "/tmp/a");
assert_eq!(result.usage.unwrap().input_tokens, Some(10));
let mut events = Vec::new();
while let Some(event) = receiver.recv().await {
events.push(event);
}
assert!(events
.iter()
.any(|event| matches!(event, ClientEvent::ThinkingEnd { .. })));
}
#[tokio::test]
async fn eof_and_half_a_tool_call_are_failures_and_keep_only_diagnostics() {
let (sender, _receiver) = tokio::sync::mpsc::channel(8);
let failure = consume_model_cycle(
provider_stream(vec![
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ModelEvent::ToolCallStart {
index: 0,
call_id: "call".into(),
name: "Read".into(),
},
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{".into(),
},
]),
&sender,
&CancellationToken::new(),
)
.await
.unwrap_err();
assert!(matches!(failure.failure, RunFailure::Provider(_)));
}
#[tokio::test]
async fn done_with_open_blocks_and_events_after_done_are_rejected() {
let (sender, _receiver) = tokio::sync::mpsc::channel(8);
let open = consume_model_cycle(
provider_stream(vec![
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ModelEvent::TextStart,
ModelEvent::Done(FinishReason::Stop),
]),
&sender,
&CancellationToken::new(),
)
.await
.unwrap_err();
assert!(matches!(open.failure, RunFailure::Protocol(_)));
let after = consume_model_cycle(
provider_stream(vec![
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ModelEvent::Done(FinishReason::Stop),
ModelEvent::Usage(Usage::default()),
]),
&sender,
&CancellationToken::new(),
)
.await
.unwrap_err();
assert!(matches!(after.failure, RunFailure::Protocol(_)));
}
#[tokio::test]
async fn duplicate_usage_is_rejected_instead_of_guessing_which_total_is_final() {
let (sender, _receiver) = tokio::sync::mpsc::channel(8);
let failure = consume_model_cycle(
provider_stream(vec![
ModelEvent::Start {
model_call_id: "model-call".into(),
},
ModelEvent::Usage(Usage::default()),
ModelEvent::Usage(Usage::default()),
ModelEvent::Done(FinishReason::Stop),
]),
&sender,
&CancellationToken::new(),
)
.await
.unwrap_err();
assert!(matches!(failure.failure, RunFailure::Protocol(_)));
}
#[tokio::test]
async fn openai_chat_raw_stream_and_request_projection_match_the_endpoint() {
let (base_url, mut requests, server) = fixture_server(
"/v1/chat/completions",
concat!(
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"wh\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"y\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":2}}\n\n",
"data: [DONE]\n\n",
),
)
.await;
let provider = OpenAiChatProvider::new(
reqwest::Client::new(),
config(ProviderKind::OpenAiChat, base_url, None),
);
let events = collect(provider.stream(invocation(), CancellationToken::new())).await;
let body = requests.recv().await.unwrap();
let continued = continued_invocation(&events);
let _ = collect(provider.stream(continued, CancellationToken::new())).await;
let second_body = requests.recv().await.unwrap();
server.abort();
assert_eq!(body["messages"][0]["content"], "system");
assert_eq!(body["messages"][1]["content"][0]["text"], "hello");
assert_eq!(body["messages"][1]["content"][1]["type"], "image_url");
assert_eq!(
body["messages"][1]["content"][1]["image_url"]["url"],
"data:image/png;base64,AQID"
);
assert!(body.get("max_completion_tokens").is_none());
assert!(body.get("reasoning_effort").is_none());
assert_eq!(
events
.iter()
.filter_map(|event| match event {
ModelEvent::ThinkingDelta(text) => Some(text.as_str()),
_ => None,
})
.collect::<String>(),
"why"
);
assert!(events
.iter()
.any(|event| matches!(event, ModelEvent::Usage(usage)
if usage.input_tokens == Some(7)
&& usage.output_tokens == Some(2)
&& usage.cache_read_tokens.is_none()
&& usage.cache_write_tokens.is_none()
&& usage.reasoning_tokens.is_none())));
assert_eq!(events.last(), Some(&ModelEvent::Done(FinishReason::Stop)));
assert_array_prefix(&body["messages"], &second_body["messages"]);
assert_eq!(
second_body["messages"].as_array().unwrap().last().unwrap()["reasoning_content"],
"why"
);
}
#[tokio::test]
async fn openai_responses_raw_stream_does_not_invent_reasoning_effort() {
let (base_url, mut requests, server) = fixture_server(
"/v1/responses",
concat!(
"data: {\"type\":\"response.reasoning_summary_text.delta\",\"delta\":\"why\"}\n\n",
"data: {\"type\":\"response.reasoning_summary_text.done\"}\n\n",
"data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"reasoning\",\"id\":\"r1\",\"encrypted_content\":\"opaque-1\"}}\n\n",
"data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"reasoning\",\"id\":\"r2\",\"encrypted_content\":\"opaque-2\"}}\n\n",
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n",
"data: {\"type\":\"response.output_text.done\"}\n\n",
"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":8,\"output_tokens\":3}}}\n\n",
),
)
.await;
let mut request = invocation();
request.request.model.reasoning.enabled = true;
let provider = OpenAiResponsesProvider::new(
reqwest::Client::new(),
config(ProviderKind::OpenAiResponses, base_url, Some(4096)),
);
let events = collect(provider.stream(request, CancellationToken::new())).await;
let body = requests.recv().await.unwrap();
let continued = continued_invocation(&events);
let _ = collect(provider.stream(continued, CancellationToken::new())).await;
let second_body = requests.recv().await.unwrap();
server.abort();
assert_eq!(body["reasoning"]["summary"], "auto");
assert!(body["reasoning"].get("effort").is_none());
assert_eq!(body["max_output_tokens"], 4096);
assert_eq!(body["input"][0]["content"][1]["type"], "input_image");
assert_eq!(
body["input"][0]["content"][1]["image_url"],
"data:image/png;base64,AQID"
);
assert!(events.iter().any(|event| matches!(event, ModelEvent::ProviderReplayState(state) if state.provider_kind == "openai_responses")));
assert!(events
.iter()
.any(|event| matches!(event, ModelEvent::Usage(usage)
if usage.input_tokens == Some(8)
&& usage.output_tokens == Some(3)
&& usage.cache_read_tokens.is_none()
&& usage.cache_write_tokens.is_none()
&& usage.reasoning_tokens.is_none())));
assert_eq!(events.last(), Some(&ModelEvent::Done(FinishReason::Stop)));
assert_array_prefix(&body["input"], &second_body["input"]);
let replayed = second_body["input"]
.as_array()
.unwrap()
.iter()
.filter_map(|item| item.get("encrypted_content").and_then(Value::as_str))
.collect::<Vec<_>>();
assert_eq!(replayed, ["opaque-1", "opaque-2"]);
}
#[tokio::test]
async fn anthropic_raw_stream_uses_only_the_explicit_route_token_limit() {
let (base_url, mut requests, server) = fixture_server(
"/v1/messages",
concat!(
"event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":9}}}\n\n",
"event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\",\"signature\":\"\"}}\n\n",
"event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"first\"}}\n\n",
"event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"signature-1\"}}\n\n",
"event: content_block_stop\ndata: {\"index\":0}\n\n",
"event: content_block_start\ndata: {\"index\":1,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\",\"signature\":\"\"}}\n\n",
"event: content_block_delta\ndata: {\"index\":1,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"second\"}}\n\n",
"event: content_block_delta\ndata: {\"index\":1,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"signature-2\"}}\n\n",
"event: content_block_stop\ndata: {\"index\":1}\n\n",
"event: content_block_start\ndata: {\"index\":2,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
"event: content_block_delta\ndata: {\"index\":2,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n",
"event: content_block_stop\ndata: {\"index\":2}\n\n",
"event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n",
"event: message_stop\ndata: {}\n\n",
),
)
.await;
let provider = cursor_server::provider::AnthropicProvider::new(
reqwest::Client::new(),
config(ProviderKind::Anthropic, base_url, Some(1234)),
);
let events = collect(provider.stream(invocation(), CancellationToken::new())).await;
let body = requests.recv().await.unwrap();
let continued = continued_invocation(&events);
let _ = collect(provider.stream(continued, CancellationToken::new())).await;
let second_body = requests.recv().await.unwrap();
server.abort();
assert_eq!(body["max_tokens"], 1234);
assert_eq!(body["messages"][0]["content"][1]["type"], "image");
assert_eq!(
body["messages"][0]["content"][1]["source"]["media_type"],
"image/png"
);
assert_eq!(body["messages"][0]["content"][1]["source"]["data"], "AQID");
assert!(events
.iter()
.any(|event| matches!(event, ModelEvent::Usage(usage)
if usage.input_tokens == Some(9)
&& usage.output_tokens == Some(2)
&& usage.cache_read_tokens.is_none()
&& usage.cache_write_tokens.is_none()
&& usage.reasoning_tokens.is_none())));
assert_eq!(events.last(), Some(&ModelEvent::Done(FinishReason::Stop)));
assert_array_prefix(&body["messages"], &second_body["messages"]);
let signatures = second_body["messages"]
.as_array()
.unwrap()
.iter()
.flat_map(|message| message["content"].as_array().into_iter().flatten())
.filter_map(|block| block.get("signature").and_then(Value::as_str))
.collect::<Vec<_>>();
assert_eq!(signatures, ["signature-1", "signature-2"]);
}
#[tokio::test]
async fn every_provider_can_be_cancelled_while_waiting_for_response_headers() {
for kind in [
ProviderKind::OpenAiChat,
ProviderKind::OpenAiResponses,
ProviderKind::Anthropic,
] {
let (base_url, accepted, server) = hanging_server().await;
let config = config(kind.clone(), base_url, Some(1024));
let provider: Arc<dyn Provider> = match kind {
ProviderKind::OpenAiChat => {
Arc::new(OpenAiChatProvider::new(reqwest::Client::new(), config))
}
ProviderKind::OpenAiResponses => {
Arc::new(OpenAiResponsesProvider::new(reqwest::Client::new(), config))
}
ProviderKind::Anthropic => Arc::new(cursor_server::provider::AnthropicProvider::new(
reqwest::Client::new(),
config,
)),
};
let cancellation = CancellationToken::new();
let stream = provider.stream(invocation(), cancellation.clone());
let collect = tokio::spawn(async move { collect(stream).await });
accepted.await.unwrap();
cancellation.cancel();
let events = tokio::time::timeout(Duration::from_secs(1), collect)
.await
.expect("provider did not cancel while waiting for headers")
.unwrap();
assert!(events.is_empty());
server.abort();
}
}
fn invocation() -> ModelInvocation {
ModelInvocation {
call_id: "call-1".into(),
run_id: "run-1".into(),
conversation_id: "conversation-1".into(),
provider_call_index: 0,
request: ModelRequest {
prompt: PromptSpec {
instructions: "system".into(),
tools: Vec::new(),
},
model: ModelSpec::new("model"),
history: vec![ProjectedMessage {
message_id: "user-1".into(),
role: Role::User,
content: ProjectedContent::Parts(vec![
ContentPart::Text {
text: "hello".into(),
},
ContentPart::Image {
mime_type: "image/png".into(),
data: vec![1, 2, 3],
},
]),
}],
},
}
}
fn continued_invocation(events: &[ModelEvent]) -> ModelInvocation {
let mut invocation = invocation();
let text = events
.iter()
.filter_map(|event| match event {
ModelEvent::TextDelta(text) => Some(text.as_str()),
_ => None,
})
.collect::<String>();
let thinking = events
.iter()
.filter_map(|event| match event {
ModelEvent::ThinkingDelta(text) => Some(text.as_str()),
_ => None,
})
.collect::<String>();
let replay_state = events.iter().find_map(|event| match event {
ModelEvent::ProviderReplayState(state) => Some(state.clone()),
_ => None,
});
invocation.request.history.push(ProjectedMessage {
message_id: "assistant-1".into(),
role: Role::Assistant,
content: ProjectedContent::Assistant {
text,
thinking,
replay_state,
calls: Vec::new(),
},
});
invocation.call_id = "call-2".into();
invocation
}
fn assert_array_prefix(first: &Value, second: &Value) {
let first = first.as_array().unwrap();
let second = second.as_array().unwrap();
assert_eq!(first.as_slice(), &second[..first.len()]);
}
fn config(kind: ProviderKind, base_url: String, max_output_tokens: Option<u64>) -> ProviderConfig {
ProviderConfig {
kind,
base_url,
api_key: "test".into(),
custom_headers: Default::default(),
max_output_tokens,
request_timeout: Duration::from_secs(5),
}
}
async fn collect(stream: ProviderStream) -> Vec<ModelEvent> {
stream.map(|event| event.unwrap()).collect::<Vec<_>>().await
}
#[derive(Clone)]
struct FixtureState {
response: Arc<str>,
requests: tokio::sync::mpsc::UnboundedSender<Value>,
}
async fn fixture_server(
path: &'static str,
response: &'static str,
) -> (
String,
tokio::sync::mpsc::UnboundedReceiver<Value>,
tokio::task::JoinHandle<()>,
) {
async fn endpoint(
axum::extract::State(state): axum::extract::State<FixtureState>,
axum::Json(body): axum::Json<Value>,
) -> impl axum::response::IntoResponse {
let _ = state.requests.send(body);
(
[(axum::http::header::CONTENT_TYPE, "text/event-stream")],
state.response.to_string(),
)
}
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
let app = axum::Router::new()
.route(path, axum::routing::post(endpoint))
.with_state(FixtureState {
response: response.into(),
requests: sender,
});
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
(format!("http://{address}/v1"), receiver, server)
}
async fn hanging_server() -> (
String,
tokio::sync::oneshot::Receiver<()>,
tokio::task::JoinHandle<()>,
) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let (accepted, receiver) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (_socket, _) = listener.accept().await.unwrap();
let _ = accepted.send(());
std::future::pending::<()>().await;
});
(format!("http://{address}/v1"), receiver, server)
}
+192
View File
@@ -0,0 +1,192 @@
#[path = "support/fixtures.rs"]
mod fixtures;
use cursor_server::model::{
ConversationId, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId, RunKind,
};
fn prepared(
run_id: &str,
conversation_id: &ConversationId,
base_revision_id: cursor_server::model::RevisionId,
) -> PreparedRun {
PreparedRun {
run_id: RunId::new(run_id),
conversation_id: conversation_id.clone(),
kind: RunKind::Root,
model: ModelSpec::new("test-model"),
prompt: PromptSpec {
instructions: "test".into(),
tools: Vec::new(),
},
selected_subagent_models: Vec::new(),
subagent_model_overrides: Vec::new(),
initial_messages: Vec::new(),
action: RunAction::Resume {
pending_tool_round: None,
},
base_revision_id,
}
}
#[tokio::test]
async fn selecting_an_old_revision_creates_a_branch_without_old_suffixes() {
let (_directory, store) = fixtures::temp_store().await;
let conversation_id = ConversationId::new("conversation");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
let first = prepared("run-1", &conversation_id, root);
store.claim_run(&first).await.unwrap();
let a = fixtures::user("a", "A");
let revision_a = store
.append_revision(
&conversation_id,
&first.run_id,
root,
std::slice::from_ref(&a),
)
.await
.unwrap();
let b = fixtures::user("b", "B");
let revision_b = store
.append_revision(
&conversation_id,
&first.run_id,
revision_a,
std::slice::from_ref(&b),
)
.await
.unwrap();
let second = prepared("run-2", &conversation_id, revision_a);
let claimed = store.claim_run(&second).await.unwrap();
assert_eq!(claimed.replaced_run_id.as_ref(), Some(&first.run_id));
let c = fixtures::user("c", "C");
let revision_c = store
.append_revision(
&conversation_id,
&second.run_id,
revision_a,
std::slice::from_ref(&c),
)
.await
.unwrap();
assert_eq!(
store.load_revision_messages(revision_b).await.unwrap(),
vec![a.clone(), b]
);
assert_eq!(
store.load_revision_messages(revision_c).await.unwrap(),
vec![a, c]
);
assert!(store
.append_revision(
&conversation_id,
&first.run_id,
revision_b,
&[fixtures::user("late", "late")],
)
.await
.is_err());
}
#[tokio::test]
async fn identical_runtime_event_is_exactly_once_and_conflicts_are_rejected() {
let (_directory, store) = fixtures::temp_store().await;
let conversation_id = ConversationId::new("runtime");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
let run = prepared("run", &conversation_id, root);
store.claim_run(&run).await.unwrap();
let event = cursor_server::model::RuntimeEvent {
event_id: "branch:changed:7".into(),
text: "runtime state changed".into(),
}
.into_message();
let (revision, inserted) = store
.append_message_once(&conversation_id, &run.run_id, root, &event)
.await
.unwrap();
assert!(inserted);
let (same, inserted) = store
.append_message_once(&conversation_id, &run.run_id, revision, &event)
.await
.unwrap();
assert_eq!(same, revision);
assert!(!inserted);
let conflict = cursor_server::model::RuntimeEvent {
event_id: "branch:changed:7".into(),
text: "different".into(),
}
.into_message();
assert!(store
.append_message_once(&conversation_id, &run.run_id, revision, &conflict)
.await
.is_err());
}
#[tokio::test]
async fn editing_a_logical_input_discards_its_active_suffix() {
let (_directory, store) = fixtures::temp_store().await;
let conversation_id = ConversationId::new("edited-conversation");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
let input_id = "cursor:user:stable-id";
assert_eq!(
store
.anchor_input(&conversation_id, input_id, root)
.await
.unwrap(),
root
);
let first = prepared("first-run", &conversation_id, root);
store.claim_run(&first).await.unwrap();
let original = fixtures::user("original", "original text");
let original_revision = store
.append_revision(
&conversation_id,
&first.run_id,
root,
std::slice::from_ref(&original),
)
.await
.unwrap();
let suffix = fixtures::user("suffix", "old suffix");
let old_head = store
.append_revision(
&conversation_id,
&first.run_id,
original_revision,
std::slice::from_ref(&suffix),
)
.await
.unwrap();
let edit_base = store
.anchor_input(&conversation_id, input_id, old_head)
.await
.unwrap();
assert_eq!(edit_base, root);
let second = prepared("second-run", &conversation_id, edit_base);
store.claim_run(&second).await.unwrap();
let edited = fixtures::user("edited", "edited text");
let edited_head = store
.append_revision(
&conversation_id,
&second.run_id,
edit_base,
std::slice::from_ref(&edited),
)
.await
.unwrap();
assert_eq!(
store.load_revision_messages(edited_head).await.unwrap(),
vec![edited]
);
assert_eq!(
store.load_revision_messages(old_head).await.unwrap(),
vec![original, suffix]
);
}
+276
View File
@@ -0,0 +1,276 @@
#[path = "support/fake_provider.rs"]
mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use cursor_server::{
cursor::{
connect,
prompting::{PromptAssets, PromptCompiler},
proto::agent::v1 as pb,
CursorCommand, CursorSessionRegistry,
},
model::{ContentPart, ProjectedContent},
provider::{FinishReason, ModelEvent},
store::{BlobId, Store},
};
use prost::Message;
#[tokio::test]
async fn current_mode_and_referenced_context_are_consumed_by_one_runtime_message() {
let (_directory, store) = fixtures::temp_store().await;
let references = references(&store).await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "model".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("answer".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store,
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry.get_or_create("ask-request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(run_request(references)),
})
.await
.unwrap();
let mut seqno = 1;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
break;
}
let message = pb::AgentServerMessage::decode(payload).unwrap();
if let Some(pb::agent_server_message::Message::KvServerMessage(kv)) = message.message {
handle
.command(CursorCommand::Append {
seqno,
message: Box::new(kv_ack(kv.id)),
})
.await
.unwrap();
seqno += 1;
}
}
let requests = provider.requests();
let request = &requests[0];
assert!(request
.prompt
.tools
.iter()
.any(|tool| tool.name == "AskQuestion"));
assert!(!request
.prompt
.tools
.iter()
.any(|tool| tool.name == "GenerateImage"));
assert_eq!(request.history.len(), 1);
assert_eq!(
request.history[0].message_id,
"runtime:run-request:ask-request"
);
let ProjectedContent::Parts(parts) = &request.history[0].content else {
panic!("runtime message must use typed parts")
};
let [ContentPart::Text { text }] = parts.as_slice() else {
panic!("this fixture has no images")
};
for expected in [
"<user_rule>workspace rule</user_rule>",
"<agent_skill fullPath=\"/skills/test/SKILL.md\">test skill</agent_skill>",
"<subagent name=\"reviewer\">review code</subagent>",
"<mcp_meta_tool_server name=\"mcp-test\" tools=\"lookup\" />",
"Ask mode is active.",
"<user_query>\nexplain this\n</user_query>",
] {
assert!(
text.contains(expected),
"missing runtime section: {expected}"
);
}
assert!(text.contains("/workspace/src/main.rs"));
}
struct References {
rules: (BlobId, u32),
skills: (BlobId, u32),
subagents: (BlobId, u32),
mcps: (BlobId, u32),
}
async fn references(store: &Store) -> References {
References {
rules: put(
store,
&pb::RequestContextRulesPart {
rules: vec![pb::CursorRule {
full_path: "/workspace/AGENTS.md".into(),
content: "workspace rule".into(),
..Default::default()
}],
..Default::default()
},
)
.await,
skills: put(
store,
&pb::RequestContextSkillsPart {
agent_skills: vec![pb::AgentSkill {
full_path: "/skills/test/SKILL.md".into(),
description: "test skill".into(),
..Default::default()
}],
..Default::default()
},
)
.await,
subagents: put(
store,
&pb::RequestContextSubagentsPart {
custom_subagents: vec![pb::CustomSubagent {
name: "reviewer".into(),
description: "review code".into(),
..Default::default()
}],
},
)
.await,
mcps: put(
store,
&pb::RequestContextMcpsPart {
mcp_meta_tool_options: Some(pb::McpMetaToolOptions {
enabled: true,
mcp_descriptors: vec![pb::McpDescriptor {
server_identifier: "mcp-test".into(),
tools: vec![pb::McpToolDescriptor {
tool_name: "lookup".into(),
..Default::default()
}],
..Default::default()
}],
}),
..Default::default()
},
)
.await,
}
}
async fn put<T: Message>(store: &Store, value: &T) -> (BlobId, u32) {
let data = value.encode_to_vec();
let length = data.len() as u32;
(store.put_blob(&data, &[]).await.unwrap(), length)
}
fn run_request(references: References) -> pb::AgentClientMessage {
let (rules, rules_byte_length) = references.rules;
let (skills, skills_byte_length) = references.skills;
let (subagents, subagents_byte_length) = references.subagents;
let (mcps, mcps_byte_length) = references.mcps;
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
conversation_state: Some(pb::ConversationStateStructure {
mode: Some(pb::AgentMode::Agent as i32),
..Default::default()
}),
action: Some(pb::ConversationAction {
request_context_parts: Some(pb::RequestContextPartReferences {
rules_blob_id: rules.as_bytes().to_vec(),
rules_byte_length,
skills_blob_id: skills.as_bytes().to_vec(),
skills_byte_length,
subagents_blob_id: subagents.as_bytes().to_vec(),
subagents_byte_length,
mcps_blob_id: mcps.as_bytes().to_vec(),
mcps_byte_length,
dynamic_context: Some(pb::RequestContext {
env: Some(pb::RequestContextEnv {
os_version: "darwin".into(),
workspace_paths: vec!["/workspace".into()],
shell: "zsh".into(),
time_zone: "UTC".into(),
..Default::default()
}),
..Default::default()
}),
}),
action: Some(pb::conversation_action::Action::UserMessageAction(
pb::UserMessageAction {
user_message: Some(pb::UserMessage {
text: "explain this".into(),
message_id: "wire-user".into(),
mode: pb::AgentMode::Ask as i32,
selected_context: Some(pb::SelectedContext {
invocation_context: Some(pb::InvocationContext {
data: Some(pb::invocation_context::Data::IdeState(
pb::invocation_context::IdeState {
visible_files: vec![
pb::invocation_context::ide_state::File {
path: "/workspace/src/main.rs".into(),
total_lines: 10,
..Default::default()
},
],
..Default::default()
},
)),
}),
..Default::default()
}),
..Default::default()
}),
..Default::default()
},
)),
..Default::default()
}),
conversation_id: Some("mode-conversation".into()),
run_id: Some("wire-run".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..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 },
)),
},
)),
}
}
+39 -11
View File
@@ -1,24 +1,52 @@
#[path = "support/fixtures.rs"]
mod fixtures;
use cursor_server::model::RuntimeEvent;
use cursor_server::model::{
ConversationId, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId, RunKind, RuntimeEvent,
};
#[tokio::test]
async fn runtime_event_is_appended_exactly_once() {
let (_directory, store) = fixtures::temp_store().await;
let event = RuntimeEvent {
let conversation_id = ConversationId::new("conversation");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
let run = PreparedRun {
run_id: RunId::new("run"),
conversation_id: conversation_id.clone(),
kind: RunKind::Root,
model: ModelSpec::new("model"),
prompt: PromptSpec {
instructions: String::new(),
tools: Vec::new(),
},
selected_subagent_models: Vec::new(),
subagent_model_overrides: Vec::new(),
initial_messages: Vec::new(),
action: RunAction::Resume {
pending_tool_round: None,
},
base_revision_id: root,
};
store.claim_run(&run).await.unwrap();
let message = RuntimeEvent {
event_id: "branch:changed:7".into(),
text: "runtime state changed".into(),
};
assert!(store
.append_runtime_event_once("conversation", event.clone())
}
.into_message();
let (revision, inserted) = store
.append_message_once(&conversation_id, &run.run_id, root, &message)
.await
.unwrap());
assert!(!store
.append_runtime_event_once("conversation", event)
.await
.unwrap());
let messages = store.load_messages("conversation").await.unwrap();
.unwrap();
assert!(inserted);
assert!(
!store
.append_message_once(&conversation_id, &run.run_id, revision, &message)
.await
.unwrap()
.1
);
let messages = store.load_revision_messages(revision).await.unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(
messages[0].runtime_event_id.as_deref(),
+229
View File
@@ -0,0 +1,229 @@
#[path = "support/fake_provider.rs"]
mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use cursor_server::{
cursor::{
connect,
prompting::{PromptAssets, PromptCompiler},
proto::agent::v1 as pb,
CursorCommand, CursorSessionRegistry,
},
model::{ContentPart, ProjectedContent},
provider::{FinishReason, ModelEvent},
store::BlobId,
};
use prost::Message;
#[tokio::test]
async fn selected_image_bytes_flow_from_run_request_to_history_providers_and_checkpoint() {
let (_directory, store) = fixtures::temp_store().await;
let stored_image = vec![4, 5];
let stored_image_id = store.put_blob(&stored_image, &[]).await.unwrap();
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ModelEvent::Start {
model_call_id: "model".into(),
},
ModelEvent::TextStart,
ModelEvent::TextDelta("seen".into()),
ModelEvent::TextEnd,
ModelEvent::Done(FinishReason::Stop),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider.clone()),
PromptCompiler::new(assets),
Default::default(),
);
let handle = registry.get_or_create("image-run").await.unwrap();
let mut output = handle.subscribe();
handle
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(request(&stored_image_id)),
})
.await
.unwrap();
let mut seqno = 1;
let mut checkpoints = Vec::new();
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
.unwrap()
.unwrap();
let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
if flags & connect::END_STREAM_FLAG != 0 {
break;
}
let message = pb::AgentServerMessage::decode(payload).unwrap();
match message.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
handle
.command(CursorCommand::Append {
seqno,
message: Box::new(kv_ack(kv.id)),
})
.await
.unwrap();
seqno += 1;
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state)) => {
checkpoints.push(state)
}
_ => {}
}
}
let requests = provider.requests();
let user = requests[0]
.history
.iter()
.find(|message| message.message_id == "runtime:run-request:image-run")
.unwrap();
let ProjectedContent::Parts(parts) = &user.content else {
panic!("runtime user message must retain typed parts")
};
assert!(matches!(
&parts[0],
ContentPart::Text { text } if text.contains("<user_query>\nwhat is this?\n</user_query>")
));
assert_eq!(
&parts[1..],
&[
ContentPart::Image {
mime_type: "image/png".into(),
data: vec![1, 2, 3],
},
ContentPart::Image {
mime_type: "image/webp".into(),
data: stored_image,
},
ContentPart::Image {
mime_type: "image/jpeg".into(),
data: vec![6, 7, 8],
},
]
);
let serialized_request = serde_json::to_string(&requests[0]).unwrap();
assert!(!serialized_request.contains("private-image-uuid"));
assert!(!serialized_request.contains("/private/image/path"));
let state = checkpoints
.iter()
.find(|state| state.root_prompt_messages_json.len() >= 2)
.unwrap();
let mut user_root = None;
for raw_id in &state.root_prompt_messages_json {
let id = BlobId::from_bytes(raw_id).unwrap();
let bytes = store.get_blob(&id).await.unwrap().unwrap();
let value: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
if value["id"] == "runtime:run-request:image-run" {
user_root = Some(value);
break;
}
}
let user_root = user_root.unwrap();
assert_eq!(user_root["content"][1]["type"], "image");
assert_eq!(user_root["content"][1]["mimeType"], "image/png");
assert_eq!(user_root["content"][1]["data"], "AQID");
assert_eq!(user_root["content"][2]["mimeType"], "image/webp");
assert_eq!(user_root["content"][2]["data"], "BAU=");
assert_eq!(user_root["content"][3]["mimeType"], "image/jpeg");
assert_eq!(user_root["content"][3]["data"], "BgcI");
}
fn request(stored_image_id: &BlobId) -> pb::AgentClientMessage {
let data = vec![1, 2, 3];
let inline_blob_data = vec![6, 7, 8];
let inline_blob_id = BlobId::digest(&inline_blob_data);
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
action: Some(pb::ConversationAction {
action: Some(pb::conversation_action::Action::UserMessageAction(
pb::UserMessageAction {
user_message: Some(pb::UserMessage {
text: "what is this?".into(),
message_id: "image-user".into(),
selected_context: Some(pb::SelectedContext {
selected_images: vec![
pb::SelectedImage {
uuid: "private-image-uuid".into(),
path: "/private/image/path".into(),
mime_type: "image/png".into(),
data_or_blob_id: Some(
pb::selected_image::DataOrBlobId::Data(data),
),
..Default::default()
},
pb::SelectedImage {
uuid: "stored-image".into(),
path: "/private/stored/path".into(),
mime_type: "image/webp".into(),
data_or_blob_id: Some(
pb::selected_image::DataOrBlobId::BlobId(
stored_image_id.as_bytes().to_vec(),
),
),
..Default::default()
},
pb::SelectedImage {
uuid: "inline-blob-image".into(),
path: "/private/inline/path".into(),
mime_type: "image/jpeg".into(),
data_or_blob_id: Some(
pb::selected_image::DataOrBlobId::BlobIdWithData(
pb::selected_image::BlobIdWithData {
blob_id: inline_blob_id.as_bytes().to_vec(),
data: inline_blob_data,
},
),
),
..Default::default()
},
],
..Default::default()
}),
mode: pb::AgentMode::Agent as i32,
..Default::default()
}),
..Default::default()
},
)),
..Default::default()
}),
conversation_id: Some("image-conversation".into()),
run_id: Some("image-run".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 },
)),
},
)),
}
}
+181
View File
@@ -0,0 +1,181 @@
use std::collections::{BTreeMap, HashMap, HashSet};
use cursor_server::{
cursor::{
interaction,
proto::agent::v1 as pb,
tools::{
codec,
runtime::{CursorToolRuntime, ExecContext, SubagentModel},
ToolBatchState, ToolDispatcher,
},
},
model::{CanonicalMessage, MessageContent, Origin, Role, ToolCall},
};
#[test]
fn task_keeps_wire_type_model_parent_and_background_fields() {
let mut context = context();
context.subagent_models.insert(
"cursor-guide".into(),
SubagentModel::Model("guide-model".into()),
);
let call = task_call(serde_json::json!({
"description": "guide",
"prompt": "inspect",
"subagent_type": "cursor-guide",
"run_in_background": true,
"interrupt": false,
}));
let message = codec::request(7, &call, &context).unwrap();
let pb::agent_server_message::Message::ExecServerMessage(exec) = message.message.unwrap()
else {
panic!("expected ExecServerMessage")
};
assert_eq!(exec.id, 7);
assert_eq!(exec.exec_id, "task-call");
assert_eq!(exec.accept_hook_additional_contexts, Some(false));
let pb::exec_server_message::Message::SubagentArgs(args) = exec.message.unwrap() else {
panic!("expected SubagentArgs")
};
assert_eq!(args.subagent_type, "cursor-guide");
assert_eq!(args.model_id, "guide-model");
assert_eq!(args.parent_conversation_id.as_deref(), Some("child"));
assert_eq!(args.root_parent_conversation_id.as_deref(), Some("root"));
assert_eq!(args.run_in_background, Some(true));
assert_eq!(args.interrupt, Some(false));
}
#[test]
fn task_renders_general_typed_and_custom_subagent_types_without_aliases() {
let cases = [
("generalPurpose", "unspecified"),
("cursor-guide", "cursor-guide"),
("MyReviewer", "MyReviewer"),
];
for (name, expected) in cases {
let call = task_call(serde_json::json!({
"description": "task",
"prompt": "inspect",
"subagent_type": name,
}));
let rendered = interaction::render_tool_call(&call, false).unwrap();
let Some(pb::tool_call::Tool::TaskToolCall(tool)) = rendered.tool else {
panic!("expected TaskToolCall")
};
let subagent = tool.args.unwrap().subagent_type.unwrap().r#type.unwrap();
match (expected, subagent) {
("unspecified", pb::subagent_type::Type::Unspecified(_)) => {}
("cursor-guide", pb::subagent_type::Type::CursorGuide(_)) => {}
(custom, pb::subagent_type::Type::Custom(value)) => {
assert_eq!(value.name, custom)
}
_ => panic!("wrong subagent oneof for {name}"),
}
}
}
#[test]
fn disabled_task_model_is_a_protocol_error() {
let mut context = context();
context
.subagent_models
.insert("security-review".into(), SubagentModel::Disabled);
let call = task_call(serde_json::json!({
"description": "review",
"prompt": "inspect",
"subagent_type": "security-review",
}));
assert!(codec::request(1, &call, &context)
.unwrap_err()
.to_string()
.contains("security-review is disabled"));
}
#[tokio::test]
async fn update_current_step_uses_a_one_based_turn_message_index() {
let dispatcher = ToolDispatcher::new(CursorToolRuntime::default());
let call = ToolCall {
index: 0,
call_id: "update-call".into(),
model_call_id: "model-call".into(),
name: "UpdateCurrentStep".into(),
arguments_text: r#"{"current_step":"testing"}"#.into(),
arguments: serde_json::json!({"current_step":"testing"}),
};
let completed = HashSet::new();
let started = HashSet::new();
let messages = vec![
CanonicalMessage::text("old-runtime", Role::User, Origin::Runtime, "old turn"),
CanonicalMessage {
message_id: "old-assistant".into(),
role: Role::Assistant,
origin: Origin::Assistant,
content: MessageContent::Assistant {
text: "old response".into(),
thinking: String::new(),
tool_round_id: None,
replay_state: None,
tool_calls: Vec::new(),
},
runtime_event_id: None,
},
CanonicalMessage::text(
"current-runtime",
Role::User,
Origin::Runtime,
"current turn",
),
];
let dispatched = dispatcher
.start_batch(
&[call],
ToolBatchState {
completed: &completed,
started: &started,
response_text: "",
response_thinking: "",
},
&messages,
&BTreeMap::new(),
&context(),
)
.await
.unwrap();
let completion = dispatched[0].completion.as_ref().unwrap();
let Some(pb::tool_call::Tool::CommunicateUpdateToolCall(tool)) =
completion.tool_call().tool.as_ref()
else {
panic!("expected CommunicateUpdateToolCall")
};
let pb::communicate_update_result::Result::Success(success) =
tool.result.as_ref().unwrap().result.as_ref().unwrap()
else {
panic!("expected communicate update success")
};
assert_eq!(success.message_index, 1);
assert_eq!(success.current_step, "testing");
}
fn context() -> ExecContext {
ExecContext {
conversation_id: "child".into(),
root_conversation_id: "root".into(),
model_id: "parent-model".into(),
subagent_models: HashMap::new(),
terminals_folder: "/tmp/terminals".into(),
admin_command_denylist: Vec::new(),
}
}
fn task_call(arguments: serde_json::Value) -> ToolCall {
ToolCall {
index: 0,
call_id: "task-call".into(),
model_call_id: "model-call".into(),
name: "Task".into(),
arguments_text: arguments.to_string(),
arguments,
}
}
+10 -6
View File
@@ -6,14 +6,14 @@ use std::{
};
use cursor_server::{
prompting::ModelRequest,
provider::{Provider, ProviderStream, ResponseEvent},
model::{ModelInvocation, ModelRequest},
provider::{ModelEvent, Provider, ProviderStream},
Error,
};
use futures_util::stream;
use tokio_util::sync::CancellationToken;
type FakeResponse = Vec<Result<ResponseEvent, Error>>;
type FakeResponse = Vec<Result<ModelEvent, Error>>;
#[derive(Clone, Default)]
pub struct FakeProvider {
@@ -22,7 +22,7 @@ pub struct FakeProvider {
}
impl FakeProvider {
pub fn push(&self, events: Vec<ResponseEvent>) {
pub fn push(&self, events: Vec<ModelEvent>) {
self.responses
.lock()
.unwrap()
@@ -37,8 +37,12 @@ impl FakeProvider {
}
impl Provider for FakeProvider {
fn stream(&self, request: ModelRequest, _cancellation: CancellationToken) -> ProviderStream {
self.requests.lock().unwrap().push(request);
fn stream(
&self,
invocation: ModelInvocation,
_cancellation: CancellationToken,
) -> ProviderStream {
self.requests.lock().unwrap().push(invocation.request);
let events = self
.responses
.lock()
+196 -29
View File
@@ -3,13 +3,17 @@ mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::{connect, proto::agent::v1 as pb},
prompting::{PromptAssets, PromptCompiler},
provider::{FinishReason, ResponseEvent},
run::{RunCommand, RunRegistry},
cursor::{CursorCommand, CursorSessionRegistry},
model::{ProjectedContent, Role, Usage},
provider::{FinishReason, ModelEvent},
};
use prost::Message;
@@ -18,35 +22,43 @@ async fn text_turn_runs_from_bidi_request_through_checkpoint_and_end_stream() {
let (_directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ResponseEvent::Start {
ModelEvent::Start {
model_call_id: "ignored".into(),
},
ResponseEvent::ThinkingStart,
ResponseEvent::ThinkingDelta("reason".into()),
ResponseEvent::ThinkingEnd,
ResponseEvent::TextStart,
ResponseEvent::TextDelta("hello".into()),
ResponseEvent::TextEnd,
ResponseEvent::Done(FinishReason::Stop),
ModelEvent::ThinkingStart,
ModelEvent::ThinkingDelta("reason".into()),
ModelEvent::ThinkingEnd,
ModelEvent::TextStart,
ModelEvent::TextDelta("hello".into()),
ModelEvent::TextEnd,
ModelEvent::Usage(Usage {
input_tokens: Some(20_000),
output_tokens: Some(8),
total_tokens: Some(20_008),
cache_read_tokens: Some(80),
cache_write_tokens: None,
reasoning_tokens: Some(3),
}),
ModelEvent::Done(FinishReason::Stop),
]);
let assets = PromptAssets::load(
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../prompt")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider.clone()),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry.get_or_create("request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run("request", "conversation", "hello")),
message: Box::new(client_run("conversation", "hello")),
})
.await
.unwrap();
@@ -56,7 +68,12 @@ async fn text_turn_runs_from_bidi_request_through_checkpoint_and_end_stream() {
let mut thinking = String::new();
let mut thinking_duration_ms = None;
let mut saw_turn_ended = false;
let mut token_deltas = Vec::new();
let mut checkpoints = 0;
let mut after_turn_checkpoints = Vec::new();
let mut blobs = HashMap::<Vec<u8>, Vec<u8>>::new();
let mut set_blob_ids = HashSet::new();
let mut withheld_final_ack = false;
loop {
let frame = tokio::time::timeout(std::time::Duration::from_secs(5), output.recv())
.await
@@ -69,8 +86,24 @@ async fn text_turn_runs_from_bidi_request_through_checkpoint_and_end_stream() {
let server = pb::AgentServerMessage::decode(payload).unwrap();
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
if let Some(pb::kv_server_message::Message::SetBlobArgs(set)) = &kv.message {
assert!(
set_blob_ids.insert(set.blob_id.clone()),
"an acknowledged content-addressed Blob must not be SET twice"
);
blobs.insert(set.blob_id.clone(), set.blob_data.clone());
}
if text == "hello" && !withheld_final_ack {
withheld_final_ack = true;
assert!(
tokio::time::timeout(std::time::Duration::from_millis(100), output.recv())
.await
.is_err(),
"TurnEnded/checkpoint must wait for the final Blob ACK"
);
}
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: append_seqno,
message: Box::new(kv_ack(kv.id)),
})
@@ -93,12 +126,24 @@ async fn text_turn_runs_from_bidi_request_through_checkpoint_and_end_stream() {
Some(pb::interaction_update::Message::ThinkingCompleted(completed)) => {
thinking_duration_ms = Some(completed.thinking_duration_ms)
}
Some(pb::interaction_update::Message::TurnEnded(_)) => saw_turn_ended = true,
Some(pb::interaction_update::Message::TurnEnded(usage)) => {
assert_eq!(usage.input_tokens, Some(20_000));
assert_eq!(usage.output_tokens, Some(8));
assert_eq!(usage.cache_read_tokens, Some(80));
assert_eq!(usage.reasoning_tokens, Some(3));
saw_turn_ended = true;
}
Some(pb::interaction_update::Message::TokenDelta(delta)) => {
token_deltas.push(delta.tokens)
}
_ => {}
}
}
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(_)) => {
checkpoints += 1
Some(pb::agent_server_message::Message::ConversationCheckpointUpdate(state)) => {
checkpoints += 1;
if saw_turn_ended {
after_turn_checkpoints.push(state);
}
}
_ => {}
}
@@ -107,17 +152,111 @@ async fn text_turn_runs_from_bidi_request_through_checkpoint_and_end_stream() {
assert_eq!(thinking, "reason");
assert!(thinking_duration_ms.is_some_and(|duration| duration >= 1));
assert!(saw_turn_ended);
assert_eq!(checkpoints, 2, "final checkpoint is intentionally repeated");
assert_eq!(provider.requests().len(), 1);
assert!(store
.load_messages("conversation")
.await
assert_eq!(token_deltas, vec![8]);
assert!(withheld_final_ack);
assert!(
checkpoints >= 2,
"final checkpoint is intentionally repeated"
);
assert_eq!(after_turn_checkpoints.len(), 3);
assert_eq!(after_turn_checkpoints[0].pending_tool_calls.len(), 1);
assert!(after_turn_checkpoints[1].pending_tool_calls.is_empty());
assert_eq!(after_turn_checkpoints[1], after_turn_checkpoints[2]);
let token_details = after_turn_checkpoints[1].token_details.as_ref().unwrap();
assert_eq!(token_details.used_tokens, 20_008);
assert_eq!(token_details.max_tokens, 256_000);
let breakdown = token_details.breakdown.as_ref().unwrap();
assert_eq!(breakdown.total_used_tokens, 20_008);
assert_eq!(breakdown.max_tokens, 256_000);
assert_eq!(breakdown.categories.len(), 8);
assert_eq!(
breakdown
.categories
.iter()
.map(|category| category.estimated_tokens)
.sum::<u32>(),
20_008
);
assert_eq!(
after_turn_checkpoints[0].turns, after_turn_checkpoints[1].turns,
"staged and settled checkpoints reuse the same frozen Turn"
);
assert_eq!(
after_turn_checkpoints[0].root_prompt_messages_json.len() + 1,
after_turn_checkpoints[1].root_prompt_messages_json.len()
);
let pending: serde_json::Value =
serde_json::from_str(&after_turn_checkpoints[0].pending_tool_calls[0]).unwrap();
assert_eq!(pending["role"], "assistant");
assert!(pending["content"]
.as_array()
.unwrap()
.iter()
.any(|message| message.message_id == "user"));
.any(|part| part["type"] == "text" && part["text"] == "hello"));
let final_root = after_turn_checkpoints[1]
.root_prompt_messages_json
.last()
.unwrap();
let stable: serde_json::Value = serde_json::from_slice(blobs.get(final_root).unwrap()).unwrap();
assert_eq!(stable["role"], "assistant");
assert!(
stable.get("origin").is_none(),
"wire root is not CanonicalMessage JSON"
);
let turn = pb::ConversationTurnStructure::decode(
blobs
.get(after_turn_checkpoints[0].turns.last().unwrap())
.unwrap()
.as_slice(),
)
.unwrap();
let pb::conversation_turn_structure::Turn::AgentConversationTurn(turn) = turn.turn.unwrap()
else {
panic!("expected agent turn")
};
assert_eq!(turn.steps.len(), 2, "thinking and text are frozen once");
let requests = provider.requests();
assert_eq!(requests.len(), 1);
assert!(requests[0]
.prompt
.instructions
.contains("powered by test-model"));
let projected = &requests[0].history;
assert_eq!(projected[0].role, Role::User);
let ProjectedContent::Parts(runtime) = &projected[0].content else {
panic!("runtime context must be text")
};
assert!(matches!(
runtime.as_slice(),
[cursor_server::model::ContentPart::Text { text }]
if text.contains("<user_info>")
&& text.contains("<user_query>\nhello\n</user_query>")
));
assert_eq!(
projected.len(),
1,
"the raw UserMessage is not projected twice"
);
let messages = store
.load_current_messages(&cursor_server::model::ConversationId::new("conversation"))
.await
.unwrap();
assert_eq!(messages[0].message_id, "runtime:run-request:request");
assert_eq!(messages[0].role, Role::User);
assert_eq!(messages.len(), 2, "runtime user plus final assistant");
let stored_runs: Vec<String> = sqlx::query_scalar("SELECT run_id FROM runs ORDER BY run_id")
.fetch_all(store.pool())
.await
.unwrap();
assert_eq!(
stored_runs,
vec!["request"],
"the concrete request_id, not Cursor's reusable wire run_id, owns the execution"
);
}
fn client_run(request_id: &str, conversation_id: &str, text: &str) -> pb::AgentClientMessage {
fn client_run(conversation_id: &str, text: &str) -> pb::AgentClientMessage {
let user = pb::UserMessage {
text: text.into(),
message_id: "user".into(),
@@ -127,17 +266,45 @@ fn client_run(request_id: &str, conversation_id: &str, text: &str) -> pb::AgentC
pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::RunRequest(
pb::AgentRunRequest {
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
action: Some(pb::ConversationAction {
action: Some(pb::conversation_action::Action::UserMessageAction(
pb::UserMessageAction {
user_message: Some(user),
request_context: Some(pb::RequestContext {
env: Some(pb::RequestContextEnv {
os_version: "darwin".into(),
workspace_paths: vec!["/workspace".into()],
shell: "zsh".into(),
terminals_folder: "/terminals".into(),
agent_transcripts_folder: "/transcripts".into(),
..Default::default()
}),
git_repos: vec![pb::GitRepoInfo {
path: "/workspace".into(),
status: "M src/main.rs".into(),
..Default::default()
}],
..Default::default()
}),
..Default::default()
},
)),
..Default::default()
}),
conversation_id: Some(conversation_id.into()),
run_id: Some(request_id.into()),
// Cursor provides the effective context limit even for a fresh conversation.
conversation_state: Some(pb::ConversationStateStructure {
token_details: Some(pb::ConversationTokenDetails {
max_tokens: 256_000,
..Default::default()
}),
..Default::default()
}),
run_id: Some("reusable-wire-run-id".into()),
..Default::default()
},
)),
+377 -54
View File
@@ -3,18 +3,25 @@ mod fake_provider;
#[path = "support/fixtures.rs"]
mod fixtures;
use std::sync::Arc;
use std::{
collections::{BTreeMap, HashSet},
sync::Arc,
};
use cursor_server::{
cursor::prompting::{PromptAssets, PromptCompiler},
cursor::{
connect, exec,
pending::{ExecContext, PendingExecRegistry},
connect,
proto::agent::v1 as pb,
tools::{
codec,
runtime::{CursorToolRuntime, ExecContext},
ToolBatchState, ToolDispatcher,
},
},
cursor::{CursorCommand, CursorSessionRegistry},
model::{MessageContent, ToolCall},
prompting::{PromptAssets, PromptCompiler},
provider::{FinishReason, ResponseEvent},
run::{RunCommand, RunRegistry},
provider::{FinishReason, ModelEvent},
};
use prost::Message;
use serde_json::json;
@@ -33,6 +40,9 @@ fn call(id: &str, name: &str) -> ToolCall {
fn exec_context() -> ExecContext {
ExecContext {
conversation_id: "conversation".into(),
root_conversation_id: "conversation".into(),
model_id: "model".into(),
subagent_models: std::collections::HashMap::new(),
terminals_folder: "/tmp/terminals".into(),
admin_command_denylist: Vec::new(),
}
@@ -56,7 +66,7 @@ fn dynamic_mcp_call_routes_to_the_captured_exec_message() {
input_schema: None,
input_schema_json: None,
};
let message = cursor_server::cursor::exec::mcp_request(7, &call, &definition).unwrap();
let message = codec::mcp_request(7, &call, &definition).unwrap();
let Some(pb::agent_server_message::Message::ExecServerMessage(exec)) = message.message else {
panic!("expected ExecServerMessage")
};
@@ -72,6 +82,128 @@ fn dynamic_mcp_call_routes_to_the_captured_exec_message() {
);
}
#[tokio::test]
async fn call_mcp_tool_reuses_the_exact_definition_returned_by_mcp_state() {
let runtime = CursorToolRuntime::default();
let dispatcher = ToolDispatcher::new(runtime.clone());
let completed = HashSet::new();
let started = HashSet::new();
let state = ToolBatchState {
completed: &completed,
started: &started,
response_text: "",
response_thinking: "",
};
let discovery = call("get-tools", "GetMcpTools");
let request = dispatcher
.start_batch(&[discovery], state, &[], &BTreeMap::new(), &exec_context())
.await
.unwrap();
let Some(pb::agent_server_message::Message::ExecServerMessage(request)) =
request[0].messages[1].message.as_ref()
else {
panic!("expected MCP state Exec")
};
let event = codec::client_event(
&pb::ExecClientMessage {
id: request.id,
message: Some(pb::exec_client_message::Message::McpStateExecResult(
pb::McpStateExecResult {
result: Some(pb::mcp_state_exec_result::Result::Success(
pb::McpStateSuccess {
servers: vec![pb::McpStateServer {
server_name: "browser-use".into(),
server_identifier: "plugin-browser-use-browser-use".into(),
tools: vec![pb::McpToolDefinition {
name: "plugin-browser-use-browser-use-browser_exec".into(),
provider_identifier: "browser-use".into(),
tool_name: "browser_exec".into(),
description: "execute browser code".into(),
..Default::default()
}],
..Default::default()
}],
},
)),
},
)),
..Default::default()
},
&runtime,
)
.await
.unwrap();
assert!(matches!(event, codec::ClientExecEvent::Completed(_)));
let mut invocation = call("call-mcp", "CallMcpTool");
invocation.arguments = json!({
"server": "plugin-browser-use-browser-use",
"toolName": "browser_exec",
"description": "run browser code",
"arguments": {"code": "print('ok')"}
});
let requests = dispatcher
.start_batch(
&[invocation],
ToolBatchState {
completed: &completed,
started: &started,
response_text: "",
response_thinking: "",
},
&[],
&BTreeMap::new(),
&exec_context(),
)
.await
.unwrap();
let Some(pb::agent_server_message::Message::ExecServerMessage(exec)) =
requests[0].messages[1].message.as_ref()
else {
panic!("expected MCP Exec")
};
let Some(pb::exec_server_message::Message::McpArgs(args)) = exec.message.as_ref() else {
panic!("expected McpArgs")
};
assert_eq!(args.name, "plugin-browser-use-browser-use-browser_exec");
assert_eq!(args.provider_identifier, "browser-use");
assert_eq!(args.tool_name, "browser_exec");
assert_eq!(args.server_identifier, "plugin-browser-use-browser-use");
assert_eq!(
args.args["code"].kind,
Some(prost_types::value::Kind::StringValue("print('ok')".into()))
);
let event = codec::client_event(
&pb::ExecClientMessage {
id: exec.id,
message: Some(pb::exec_client_message::Message::McpResult(pb::McpResult {
result: Some(pb::mcp_result::Result::Success(pb::McpSuccess {
content: vec![pb::McpToolResultContentItem {
content: Some(pb::mcp_tool_result_content_item::Content::Text(
pb::McpTextContent {
text: "browser result".into(),
output_location: None,
},
)),
}],
is_error: false,
structured_content: None,
})),
})),
..Default::default()
},
&runtime,
)
.await
.unwrap();
let codec::ClientExecEvent::Completed(completion) = event else {
panic!("expected MCP completion")
};
assert_eq!(completion.result().content, "browser result");
assert!(!completion.result().is_error);
}
#[tokio::test]
async fn shell_uses_background_timeout_and_preserves_stream_identity() {
let mut shell = call("call-shell", "Shell");
@@ -82,7 +214,7 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
"description": "Start HTTP server"
});
let context = exec_context();
let request = exec::request(7, &shell, &context).unwrap();
let request = codec::request(7, &shell, &context).unwrap();
let Some(pb::agent_server_message::Message::ExecServerMessage(request)) = request.message
else {
panic!("expected ExecServerMessage")
@@ -102,9 +234,9 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
assert_eq!(args.conversation_id.as_deref(), Some("conversation"));
assert_eq!(args.file_output_threshold_bytes, Some(40_000));
let pending = PendingExecRegistry::default();
let id = pending.reserve(&shell, &context).await.unwrap();
let delta = exec::client_event(
let pending = CursorToolRuntime::default();
let id = pending.reserve_exec(&shell, &context).await.unwrap();
let delta = codec::client_event(
&pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::ShellStream(
@@ -120,7 +252,7 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
)
.await
.unwrap();
let exec::ClientExecEvent::Delta(delta) = delta else {
let codec::ClientExecEvent::Delta(delta) = delta else {
panic!("expected Shell stdout delta")
};
let Some(pb::agent_server_message::Message::InteractionUpdate(delta)) = delta.message else {
@@ -141,7 +273,7 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
};
assert_eq!(stdout.content, "Serving HTTP on port 8000\n");
let completion = exec::client_event(
let completion = codec::client_event(
&pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::ShellStream(
@@ -164,12 +296,12 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
)
.await
.unwrap();
let exec::ClientExecEvent::Completed(completion) = completion else {
let codec::ClientExecEvent::Completed(completion) = completion else {
panic!("expected background completion")
};
assert_eq!(
completion.result().output.as_str(),
Some(
completion.result().content,
(
"shell running in background shell_id=42 pid=1234 terminals_folder=/tmp/terminals\nServing HTTP on port 8000\n"
)
);
@@ -180,38 +312,51 @@ async fn shell_uses_background_timeout_and_preserves_stream_identity() {
assert_eq!(result.is_background, Some(true));
assert_eq!(result.terminals_folder.as_deref(), Some("/tmp/terminals"));
assert_eq!(result.pid, Some(1234));
assert!(
pending.drain_running().await.is_empty(),
"a backgrounded Shell is no longer an abortable Run Exec"
);
}
#[tokio::test]
async fn exec_ids_are_monotonic_and_released_ids_are_not_reused() {
let pending = PendingExecRegistry::default();
let pending = CursorToolRuntime::default();
let first = pending
.reserve(&call("call-1", "Read"), &exec_context())
.reserve_exec(&call("call-1", "Read"), &exec_context())
.await
.unwrap();
assert_eq!(first, 1);
assert_eq!(
pending.call(first).await.map(|call| call.call_id),
pending.exec_call(first).await.map(|call| call.call_id),
Some("call-1".into())
);
pending.discard(first).await;
assert!(pending.call(first).await.is_none());
pending.discard_exec(first).await;
assert!(pending.exec_call(first).await.is_none());
let second = pending
.reserve(&call("call-2", "Read"), &exec_context())
.reserve_exec(&call("call-2", "Read"), &exec_context())
.await
.unwrap();
assert_eq!(second, 2, "released Exec ids must not be reused in one Run");
let interaction = pending
.reserve_interaction(&call("call-3", "AskQuestion"), &exec_context())
.await
.unwrap();
assert_eq!(
interaction, 3,
"Exec and Interaction share one wire-id space"
);
}
#[tokio::test]
async fn empty_exec_client_message_is_not_a_terminal_result() {
let pending = PendingExecRegistry::default();
let pending = CursorToolRuntime::default();
let id = pending
.reserve(&call("call-1", "Read"), &exec_context())
.reserve_exec(&call("call-1", "Read"), &exec_context())
.await
.unwrap();
let event = exec::client_event(
let event = codec::client_event(
&pb::ExecClientMessage {
id,
message: None,
@@ -221,20 +366,20 @@ async fn empty_exec_client_message_is_not_a_terminal_result() {
)
.await
.unwrap();
assert!(matches!(event, exec::ClientExecEvent::Pending));
assert!(matches!(event, codec::ClientExecEvent::Pending));
assert_eq!(
pending.call(id).await.map(|call| call.call_id),
pending.exec_call(id).await.map(|call| call.call_id),
Some("call-1".into())
);
}
#[tokio::test]
async fn tool_success_is_not_inferred_from_debug_text() {
let pending = PendingExecRegistry::default();
let pending = CursorToolRuntime::default();
let mut write = call("call-1", "Write");
write.arguments = json!({"path": "/tmp/a", "contents": "x"});
let id = pending.reserve(&write, &exec_context()).await.unwrap();
let event = exec::client_event(
let id = pending.reserve_exec(&write, &exec_context()).await.unwrap();
let event = codec::client_event(
&pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::WriteResult(
@@ -252,7 +397,7 @@ async fn tool_success_is_not_inferred_from_debug_text() {
)
.await
.unwrap();
let exec::ClientExecEvent::Completed(completion) = event else {
let codec::ClientExecEvent::Completed(completion) = event else {
panic!("expected terminal write result")
};
assert!(!completion.result().is_error);
@@ -262,14 +407,63 @@ async fn tool_success_is_not_inferred_from_debug_text() {
));
}
#[tokio::test]
async fn new_task_result_exposes_the_subagent_name_and_id_to_the_model() {
let pending = CursorToolRuntime::default();
let mut task = call("call-task", "Task");
task.arguments = json!({
"description": "Analyze game logic",
"prompt": "Inspect the game",
"run_in_background": true,
"subagent_type": "generalPurpose"
});
let id = pending.reserve_exec(&task, &exec_context()).await.unwrap();
let event = codec::client_event(
&pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::SubagentResult(
pb::SubagentResult {
result: Some(pb::subagent_result::Result::Success(pb::SubagentSuccess {
agent_id: "child-id".into(),
..Default::default()
})),
},
)),
..Default::default()
},
&pending,
)
.await
.unwrap();
let codec::ClientExecEvent::Completed(completion) = event else {
panic!("expected terminal Task result")
};
assert_eq!(
completion.result().content,
"Subagent name: Analyze game logic\nSubagent ID: child-id"
);
let Some(pb::tool_call::Tool::TaskToolCall(tool)) = &completion.tool_call().tool else {
panic!("expected TaskToolCall")
};
let Some(pb::task_result::Result::Success(success)) = tool
.result
.as_ref()
.and_then(|result| result.result.as_ref())
else {
panic!("expected typed Task success")
};
assert_eq!(success.agent_id.as_deref(), Some("child-id"));
}
#[tokio::test]
async fn an_exec_result_must_match_the_reserved_tool() {
let pending = PendingExecRegistry::default();
let pending = CursorToolRuntime::default();
let id = pending
.reserve(&call("call-1", "Read"), &exec_context())
.reserve_exec(&call("call-1", "Read"), &exec_context())
.await
.unwrap();
let result = exec::client_event(
let result = codec::client_event(
&pb::ExecClientMessage {
id,
message: Some(pb::exec_client_message::Message::WriteResult(
@@ -291,7 +485,127 @@ async fn an_exec_result_must_match_the_reserved_tool() {
assert!(error
.to_string()
.contains("unexpected Exec result for tool Read"));
assert!(pending.call(id).await.is_none());
assert!(pending.exec_call(id).await.is_none());
assert_eq!(pending.completed_call(id).await.as_deref(), Some("call-1"));
let duplicate = codec::client_event(
&pb::ExecClientMessage {
id,
message: None,
..Default::default()
},
&pending,
)
.await;
let Err(duplicate) = duplicate else {
panic!("duplicate terminal result must fail")
};
assert!(duplicate.to_string().contains("duplicate terminal"));
}
#[tokio::test]
async fn unknown_exec_id_is_a_protocol_error() {
let result = codec::client_event(
&pb::ExecClientMessage {
id: 999,
message: Some(pb::exec_client_message::Message::ReadResult(
pb::ReadResult::default(),
)),
..Default::default()
},
&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"
));
}
#[tokio::test]
async fn await_shell_consumes_the_background_output_file_terminal_state() {
let runtime = CursorToolRuntime::default();
let dispatcher = ToolDispatcher::new(runtime.clone());
let mut await_call = call("await-call", "AwaitShell");
await_call.arguments = json!({
"shell_id": "42",
"block_until_ms": 1000,
"pattern": "ready",
});
await_call.arguments_text = await_call.arguments.to_string();
let completed = HashSet::new();
let started = HashSet::new();
let dispatched = dispatcher
.start_batch(
&[await_call],
ToolBatchState {
completed: &completed,
started: &started,
response_text: "",
response_thinking: "",
},
&[],
&BTreeMap::new(),
&exec_context(),
)
.await
.unwrap();
let exec = dispatched[0]
.messages
.iter()
.find_map(|message| match message.message.as_ref() {
Some(pb::agent_server_message::Message::ExecServerMessage(exec)) => Some(exec),
_ => None,
})
.unwrap();
let Some(pb::exec_server_message::Message::ReadArgs(read)) = exec.message.as_ref() else {
panic!("expected AwaitShell ReadArgs")
};
assert_eq!(read.path, "/tmp/terminals/42.txt");
let event = codec::client_event(
&pb::ExecClientMessage {
id: exec.id,
message: Some(pb::exec_client_message::Message::ReadResult(
pb::ReadResult {
result: Some(pb::read_result::Result::Success(pb::ReadSuccess {
output: Some(pb::read_success::Output::Content(
"server ready\nexit_code: 0\n".into(),
)),
..Default::default()
})),
},
)),
..Default::default()
},
&runtime,
)
.await
.unwrap();
let codec::ClientExecEvent::Completed(completion) = event else {
panic!("expected completed AwaitShell")
};
assert_eq!(completion.result().call_id, "await-call");
assert!(!completion.result().is_error);
let Some(pb::tool_call::Tool::AwaitToolCall(tool)) = completion.tool_call().tool.as_ref()
else {
panic!("expected AwaitToolCall")
};
let pb::await_result::Result::Success(success) =
tool.result.as_ref().unwrap().result.as_ref().unwrap()
else {
panic!("expected Await success")
};
let pb::await_success::AwaitResult::Complete(complete) = success.await_result.as_ref().unwrap()
else {
panic!("expected completed background task")
};
assert_eq!(complete.task_id, "42");
assert_eq!(complete.exit_code, Some(0));
assert_eq!(complete.regex_match.as_deref(), Some("ready"));
}
#[tokio::test]
@@ -299,46 +613,46 @@ async fn provider_tool_use_waits_for_client_result_then_calls_provider_again() {
let (directory, store) = fixtures::temp_store().await;
let provider = fake_provider::FakeProvider::default();
provider.push(vec![
ResponseEvent::Start {
ModelEvent::Start {
model_call_id: "ignored".into(),
},
ResponseEvent::ToolCallStart {
ModelEvent::ToolCallStart {
index: 0,
call_id: "call-1".into(),
name: "Read".into(),
},
ResponseEvent::ToolCallArgumentsDelta {
ModelEvent::ToolCallArgumentsDelta {
index: 0,
delta: "{\"path\":\"/tmp/a\"}".into(),
},
ResponseEvent::ToolCallEnd { index: 0 },
ResponseEvent::Done(FinishReason::ToolUse),
ModelEvent::ToolCallEnd { index: 0 },
ModelEvent::Done(FinishReason::ToolUse),
]);
provider.push(vec![
ResponseEvent::Start {
ModelEvent::Start {
model_call_id: "ignored".into(),
},
ResponseEvent::TextStart,
ResponseEvent::TextDelta("done".into()),
ResponseEvent::TextEnd,
ResponseEvent::Done(FinishReason::Stop),
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")
.join("../prompt/cursor")
.as_path(),
)
.unwrap();
let registry = RunRegistry::new(
let registry = CursorSessionRegistry::new(
store.clone(),
Arc::new(provider.clone()),
PromptCompiler::new(assets),
"test-model".into(),
Default::default(),
);
let handle = registry.get_or_create("tool-request").await.unwrap();
let mut output = handle.subscribe();
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno: 0,
message: Box::new(client_run()),
})
@@ -364,7 +678,7 @@ async fn provider_tool_use_waits_for_client_result_then_calls_provider_again() {
match server.message {
Some(pb::agent_server_message::Message::KvServerMessage(kv)) => {
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno,
message: Box::new(kv_ack(kv.id)),
})
@@ -376,7 +690,7 @@ async fn provider_tool_use_waits_for_client_result_then_calls_provider_again() {
saw_exec = true;
let exec_id = exec.id;
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno,
message: Box::new(pb::AgentClientMessage {
message: Some(pb::agent_client_message::Message::ExecClientMessage(
@@ -409,7 +723,7 @@ async fn provider_tool_use_waits_for_client_result_then_calls_provider_again() {
.unwrap();
seqno += 1;
handle
.command(RunCommand::Append {
.command(CursorCommand::Append {
seqno,
message: Box::new(pb::AgentClientMessage {
message: Some(
@@ -461,13 +775,18 @@ async fn provider_tool_use_waits_for_client_result_then_calls_provider_again() {
.await
.unwrap();
let provider_call_index: i64 =
sqlx::query_scalar("SELECT provider_call_index FROM runs WHERE request_id = ?")
sqlx::query_scalar("SELECT provider_call_index FROM runs WHERE run_id = ?")
.bind("tool-request")
.fetch_one(&database)
.await
.unwrap();
assert_eq!(provider_call_index, 1);
let messages = store.load_messages("tool-conversation").await.unwrap();
let messages = store
.load_current_messages(&cursor_server::model::ConversationId::new(
"tool-conversation",
))
.await
.unwrap();
let result_position = messages
.iter()
.position(|message| matches!(message.content, MessageContent::ToolResult(_)))
@@ -501,6 +820,10 @@ fn client_run() -> pb::AgentClientMessage {
}),
conversation_id: Some("tool-conversation".into()),
run_id: Some("tool-request".into()),
requested_model: Some(pb::RequestedModel {
model_id: "test-model".into(),
..Default::default()
}),
..Default::default()
},
)),
+118
View File
@@ -0,0 +1,118 @@
#[path = "support/fixtures.rs"]
mod fixtures;
use cursor_server::{
model::{
ConversationId, MessageContent, ModelSpec, PreparedRun, PromptSpec, RunAction, RunId,
RunKind, ToolCall, ToolResult, ToolRoundAssistant, ToolRoundId,
},
store::ToolRoundStatus,
};
#[tokio::test]
async fn results_commit_adjacent_pairs_in_arrival_order() {
let (_directory, store) = fixtures::temp_store().await;
let conversation_id = ConversationId::new("conversation");
let root = store.ensure_conversation(&conversation_id).await.unwrap();
let run = PreparedRun {
run_id: RunId::new("run"),
conversation_id: conversation_id.clone(),
kind: RunKind::Root,
model: ModelSpec::new("model"),
prompt: PromptSpec {
instructions: String::new(),
tools: Vec::new(),
},
selected_subagent_models: Vec::new(),
subagent_model_overrides: Vec::new(),
initial_messages: Vec::new(),
action: RunAction::Resume {
pending_tool_round: None,
},
base_revision_id: root,
};
store.claim_run(&run).await.unwrap();
let round_id = ToolRoundId::new("round");
let calls = [call(0, "A"), call(1, "B"), call(2, "C")];
store
.create_tool_round(
&round_id,
&run.run_id,
root,
&ToolRoundAssistant {
text: "answer prefix".into(),
thinking: "reasoning".into(),
model_call_id: "model-call".into(),
replay_state: None,
},
&calls,
None,
)
.await
.unwrap();
let b = store
.commit_tool_result(&conversation_id, &run.run_id, &round_id, &result("B"))
.await
.unwrap();
assert_eq!(b.completion_seq, 0);
assert!(!b.settled);
let a = store
.commit_tool_result(&conversation_id, &run.run_id, &round_id, &result("A"))
.await
.unwrap();
assert_eq!(a.completion_seq, 1);
let c = store
.commit_tool_result(&conversation_id, &run.run_id, &round_id, &result("C"))
.await
.unwrap();
assert!(c.settled);
let messages = store.load_revision_messages(c.revision_id).await.unwrap();
assert_eq!(messages.len(), 6);
let ids = messages
.chunks_exact(2)
.map(|pair| match (&pair[0].content, &pair[1].content) {
(MessageContent::Assistant { tool_calls, .. }, MessageContent::ToolResult(result)) => {
assert_eq!(tool_calls[0].call_id, result.call_id);
result.call_id.clone()
}
_ => panic!("tool result must be adjacent to its assistant call"),
})
.collect::<Vec<_>>();
assert_eq!(ids, ["B", "A", "C"]);
let MessageContent::Assistant { text, thinking, .. } = &messages[0].content else {
unreachable!()
};
assert_eq!(text, "answer prefix");
assert_eq!(thinking, "reasoning");
for message in [&messages[2], &messages[4]] {
let MessageContent::Assistant { text, thinking, .. } = &message.content else {
unreachable!()
};
assert!(text.is_empty());
assert!(thinking.is_empty());
}
let snapshot = store.tool_round(&round_id).await.unwrap().unwrap();
assert_eq!(snapshot.status, ToolRoundStatus::Settled);
assert_eq!(snapshot.completed_call_ids.len(), 3);
}
fn call(index: usize, call_id: &str) -> ToolCall {
ToolCall {
index,
call_id: call_id.into(),
model_call_id: "model-call".into(),
name: "Read".into(),
arguments_text: r#"{"path":"/tmp/a"}"#.into(),
arguments: serde_json::json!({"path":"/tmp/a"}),
}
}
fn result(call_id: &str) -> ToolResult {
ToolResult {
call_id: call_id.into(),
content: format!("result-{call_id}"),
is_error: false,
}
}