all tools

cursor-server/tests/connect_wire.rs

@@ -0,0 +1,136 @@
+#[path = "support/fake_cursor.rs"]
+mod fake_cursor;
+#[path = "support/fake_provider.rs"]
+mod fake_provider;
+#[path = "support/fixtures.rs"]
+mod fixtures;
+
+use std::{io::Write, sync::Arc};
+
+use axum::{
+    body::{to_bytes, Body},
+    http::{header, Request, StatusCode},
+};
+use base64::{engine::general_purpose::STANDARD_NO_PAD, Engine};
+use cursor_server::{
+    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;
+use tower::ServiceExt;
+
+#[test]
+fn connect_envelope_is_flag_plus_big_endian_length_plus_protobuf() {
+    let message = pb::BidiRequestId {
+        request_id: "abc".into(),
+    };
+    let frame = connect::encode_message(&message).unwrap();
+    assert_eq!(&frame[..5], &[0, 0, 0, 0, 5]);
+    let decoded: pb::BidiRequestId = fake_cursor::decode_single(&frame).unwrap();
+    assert_eq!(decoded.request_id, "abc");
+}
+
+#[test]
+fn end_stream_matches_captured_connect_shape() {
+    assert_eq!(
+        connect::encode_end_stream().as_ref(),
+        &[2, 0, 0, 0, 2, b'{', b'}']
+    );
+}
+
+#[test]
+fn error_end_stream_is_flagged_json_not_protobuf() {
+    let frame = connect::encode_error_end_stream(&connect::ConnectStreamError {
+        code: connect::ConnectCode::Unavailable,
+        message: "overloaded".into(),
+        details: vec![connect::ConnectErrorDetail {
+            type_name: "aiserver.v1.ErrorDetails".into(),
+            value: "AQ".into(),
+        }],
+    })
+    .unwrap();
+    let (flags, payload) = connect::decode_frames(&frame).unwrap().pop().unwrap();
+    assert_eq!(flags, connect::END_STREAM_FLAG);
+    let json: serde_json::Value = serde_json::from_slice(&payload).unwrap();
+    assert_eq!(json["error"]["code"], "unavailable");
+    assert_eq!(json["error"]["message"], "overloaded");
+    assert_eq!(
+        json["error"]["details"][0]["type"],
+        "aiserver.v1.ErrorDetails"
+    );
+}
+
+#[test]
+fn cursor_error_details_subset_decodes_captured_wire_value() {
+    let captured = "CAISVQoUQXV0aGVudGljYXRpb24gZXJyb3ISMklmIHlvdSBhcmUgbG9nZ2VkIGluLCB0cnkgbG9nZ2luZyBvdXQgYW5kIGJhY2sgaW4uIABSBwoFbG9naW4YAQ";
+    let bytes = STANDARD_NO_PAD.decode(captured).unwrap();
+    let details = ai::ErrorDetails::decode(bytes.as_slice()).unwrap();
+    assert_eq!(details.error, 2, "ERROR_NOT_LOGGED_IN");
+    assert_eq!(details.is_expected, Some(true));
+    let custom = details.details.unwrap();
+    assert_eq!(custom.title, "Authentication error");
+    assert_eq!(custom.is_retryable, Some(false));
+}
+
+#[test]
+fn captured_kv_ack_hex_decodes_as_agent_client_message() {
+    let bytes = hex::decode("1a0408011a00").unwrap();
+    let message = pb::AgentClientMessage::decode(bytes.as_slice()).unwrap();
+    let Some(pb::agent_client_message::Message::KvClientMessage(kv)) = message.message else {
+        panic!("expected KV client message")
+    };
+    assert_eq!(kv.id, 1);
+    assert!(matches!(
+        kv.message,
+        Some(pb::kv_client_message::Message::SetBlobResult(_))
+    ));
+}
+
+#[tokio::test]
+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")
+            .as_path(),
+    )
+    .unwrap();
+    let registry = RunRegistry::new(
+        store,
+        Arc::new(fake_provider::FakeProvider::default()),
+        PromptCompiler::new(assets),
+        "test-model".into(),
+    );
+    let wire = ai::BidiAppendRequest {
+        request_id: Some(ai::BidiRequestId {
+            request_id: "gzip-request".into(),
+        }),
+        ..Default::default()
+    }
+    .encode_to_vec();
+    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
+    encoder.write_all(&wire).unwrap();
+    let compressed = encoder.finish().unwrap();
+
+    let response = handlers::router(registry)
+        .oneshot(
+            Request::post("/aiserver.v1.BidiService/BidiAppend")
+                .header(header::CONTENT_TYPE, "application/proto")
+                .header(header::CONTENT_ENCODING, "gzip")
+                .body(Body::from(compressed))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+    let body = to_bytes(response.into_body(), 4096).await.unwrap();
+    let text = std::str::from_utf8(&body).unwrap();
+    assert!(text.contains("BidiAppend contains no AgentClientMessage"));
+    assert!(!text.contains("protobuf decode error"));
+}