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

cursor-server/tests/checkpoint_recovery.rs

@@ -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;
 }