zerogravity/src/proto.rs

//! Protobuf wire-format encoder — byte-exact match to the real Antigravity webview.
//!
//! This is a minimal, hand-rolled encoder. We do NOT use prost or any codegen
//! because we need precise control over field ordering and encoding to produce
//! byte-identical output to the captured webview traffic.
//!
//! The LS communicates with Google via gRPC over H2, using proto messages from
//! the `exa.*_pb` package family. Init metadata (field 34 of the init request)
//! carries the `detect_and_use_proxy` enum, model selection, and version info.
//! See `docs/ls-binary-analysis.md` for the full proto schema reverse engineering.

use crate::constants::{client_version, CLIENT_NAME};

// ─── Wire primitives ────────────────────────────────────────────────────────

/// Encode a varint (base-128, little-endian, MSB continuation).
pub fn varint(mut val: u64) -> Vec<u8> {
    if val == 0 {
        return vec![0x00];
    }
    let mut out = Vec::with_capacity(10);
    while val > 0x7F {
        out.push(((val & 0x7F) | 0x80) as u8);
        val >>= 7;
    }
    out.push((val & 0x7F) as u8);
    out
}

/// Encode a field tag (field_number << 3 | wire_type).
pub fn tag(field: u32, wire: u8) -> Vec<u8> {
    varint(((field as u64) << 3) | (wire as u64))
}

/// Wire type 2: length-delimited string/bytes field.
pub fn proto_string(field: u32, val: &[u8]) -> Vec<u8> {
    let mut out = tag(field, 2);
    out.extend(varint(val.len() as u64));
    out.extend_from_slice(val);
    out
}

/// Wire type 2: length-delimited sub-message field.
pub fn proto_message(field: u32, inner: &[u8]) -> Vec<u8> {
    let mut out = tag(field, 2);
    out.extend(varint(inner.len() as u64));
    out.extend_from_slice(inner);
    out
}

/// Wire type 0: boolean field (varint 0 or 1).
pub fn bool_field(field: u32, val: bool) -> Vec<u8> {
    let mut out = tag(field, 0);
    out.extend(varint(if val { 1 } else { 0 }));
    out
}

/// Wire type 0: varint field.
pub fn varint_field(field: u32, val: u64) -> Vec<u8> {
    let mut out = tag(field, 0);
    out.extend(varint(val));
    out
}

// ─── Init metadata builder (for standalone LS stdin) ─────────────────────────

/// Build the init metadata protobuf that the LS expects on stdin at startup.
///
/// This replaces the Python snippet in `standalone-ls.sh` with proper Rust encoding.
/// Fields match what the real Antigravity extension sends to the LS.
///
/// Field layout (from binary analysis):
///   1: api_key (string) — unique session key
///   3: ide_name (string) — "antigravity"
///   4: antigravity_version (string) — e.g. "1.107.0"
///   5: ide_version (string) — e.g. "1.16.5"
///   6: locale (string) — "en_US"
///  10: session_id (string) — unique session identifier
///  11: editor_name (string) — "antigravity"
///  34: detect_and_use_proxy (varint enum) — 1 = ENABLED
pub fn build_init_metadata(
    api_key: &str,
    antigravity_version: &str,
    ide_version: &str,
    session_id: &str,
    detect_and_use_proxy: u64,
) -> Vec<u8> {
    let mut buf = Vec::with_capacity(128);

    // Field 1: api_key
    buf.extend(proto_string(1, api_key.as_bytes()));
    // Field 3: ide_name
    buf.extend(proto_string(3, CLIENT_NAME.as_bytes()));
    // Field 4: antigravity version
    buf.extend(proto_string(4, antigravity_version.as_bytes()));
    // Field 5: IDE/client version
    buf.extend(proto_string(5, ide_version.as_bytes()));
    // Field 6: locale
    buf.extend(proto_string(6, b"en_US"));
    // Field 10: session_id
    buf.extend(proto_string(10, session_id.as_bytes()));
    // Field 11: editor_name
    buf.extend(proto_string(11, CLIENT_NAME.as_bytes()));
    // Field 34: detect_and_use_proxy enum (1 = ENABLED)
    buf.extend(varint_field(34, detect_and_use_proxy));

    buf
}

// ─── SendUserCascadeMessageRequest builder ───────────────────────────────────

/// Build the `SendUserCascadeMessageRequest` protobuf binary.
///
/// Produces a byte-exact match to real Antigravity webview traffic.
/// Verified against Chrome DevTools network capture 2026-02-12.
///
/// Field layout:
///   1: cascade_id (string)
///   2: { 1: text } (message)
///   3: metadata { 1: client_name, 3: oauth_token, 4: "en", 7: version, 12: client_name }
///   5: PlannerConfig { 1: inner_config, 7: { 1: 1 } }
///      inner_config contains: f2 (conv mode), f13 (tool config), f15 (model), f21 (ephemeral), f32 (knowledge)
///  11: conversation_history = true
pub fn build_request(cascade_id: &str, text: &str, oauth_token: &str, model_enum: u32) -> Vec<u8> {
    let mut msg = Vec::with_capacity(256);

    // Field 1: cascade_id
    msg.extend(proto_string(1, cascade_id.as_bytes()));

    // Field 2: { field 1: text }
    msg.extend(proto_message(2, &proto_string(1, text.as_bytes())));

    // Field 3: Metadata (Auth + Client ID)
    let mut meta = Vec::new();
    meta.extend(proto_string(1, CLIENT_NAME.as_bytes()));
    meta.extend(proto_string(3, oauth_token.as_bytes()));
    meta.extend(proto_string(4, b"en"));
    meta.extend(proto_string(7, client_version().as_bytes()));
    meta.extend(proto_string(12, CLIENT_NAME.as_bytes()));
    msg.extend(proto_message(3, &meta));

    // Field 5: PlannerConfig
    let mut inner = Vec::new();

    // field 2: conversational mode { f4: 1, f14: 0 }
    let conv_mode = [varint_field(4, 1), varint_field(14, 0)].concat();
    inner.extend(proto_message(2, &conv_mode));

    // field 13: toolConfig
    //   field 8 (runCommand): field 3 (autoCommandConfig) -> field 6 (policy) = 3 (EAGER)
    //   field 33 (artifactReviewPolicy): field 1 = 2 (TURBO)
    let run_cmd = proto_message(3, &varint_field(6, 3));
    let tool_config = [
        proto_message(8, &run_cmd),
        proto_message(33, &varint_field(1, 2)),
    ]
    .concat();
    inner.extend(proto_message(13, &tool_config));

    // field 15: requested model { f1: model_enum }
    inner.extend(proto_message(15, &varint_field(1, model_enum as u64)));

    // field 21: ephemeral messages config { f1: 1 }
    inner.extend(proto_message(21, &varint_field(1, 1)));

    // field 32: knowledge config { f1: true }
    inner.extend(proto_message(32, &bool_field(1, true)));

    // Field 5 wraps: field 1 (inner config) + field 7 { f1: 1 }
    let f5_payload = [
        proto_message(1, &inner),
        proto_message(7, &varint_field(1, 1)),
    ]
    .concat();
    msg.extend(proto_message(5, &f5_payload));

    // Field 11: conversation history flag
    msg.extend(bool_field(11, true));

    msg
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_varint_zero() {
        assert_eq!(varint(0), vec![0x00]);
    }

    #[test]
    fn test_varint_small() {
        assert_eq!(varint(1), vec![0x01]);
        assert_eq!(varint(127), vec![0x7F]);
    }

    #[test]
    fn test_varint_multibyte() {
        assert_eq!(varint(128), vec![0x80, 0x01]);
        assert_eq!(varint(300), vec![0xAC, 0x02]);
    }

    #[test]
    fn test_varint_1026() {
        // model_enum 1026 = 0x402 → varint [0x82, 0x08]
        assert_eq!(varint(1026), vec![0x82, 0x08]);
    }

    #[test]
    fn test_tag() {
        // field 1, wire type 2 (LEN) = (1 << 3) | 2 = 0x0A
        assert_eq!(tag(1, 2), vec![0x0A]);
        // field 3, wire type 0 (VARINT) = (3 << 3) | 0 = 0x18
        assert_eq!(tag(3, 0), vec![0x18]);
        // field 33, wire type 2 = (33 << 3) | 2 = 266 → varint [0x8A, 0x02]
        assert_eq!(tag(33, 2), vec![0x8A, 0x02]);
    }

    #[test]
    fn test_proto_string() {
        let result = proto_string(1, b"hi");
        // tag(1,2) = 0x0A, len=2, 'h'=0x68, 'i'=0x69
        assert_eq!(result, vec![0x0A, 0x02, 0x68, 0x69]);
    }

    #[test]
    fn test_build_request_deterministic() {
        let a = build_request("cid", "hello", "ya29.tok", 1026);
        let b = build_request("cid", "hello", "ya29.tok", 1026);
        assert_eq!(a, b, "build_request must be deterministic");
    }

    #[test]
    fn test_build_request_structure() {
        let msg = build_request("test-cascade-id", "hello", "ya29.test-token", 1026);

        assert_eq!(msg[0], 0x0A, "first byte must be field 1 tag");

        let cascade_bytes = b"test-cascade-id";
        assert!(
            msg.windows(cascade_bytes.len())
                .any(|w| w == cascade_bytes),
            "cascade_id must appear in output"
        );

        assert!(
            msg.windows(5).any(|w| w == b"hello"),
            "text must appear in output"
        );

        let token_bytes = b"ya29.test-token";
        assert!(
            msg.windows(token_bytes.len()).any(|w| w == token_bytes),
            "oauth token must appear in output"
        );

        // model enum 1026 varint [0x82, 0x08]
        assert!(
            msg.windows(2).any(|w| w == [0x82, 0x08]),
            "model enum 1026 varint must appear in output"
        );

        // field 11 bool true at end: tag(11,0)=0x58, varint(1)=0x01
        let len = msg.len();
        assert_eq!(msg[len - 2], 0x58);
        assert_eq!(msg[len - 1], 0x01);
    }

    /// Cross-verified against Python output: 127/127 bytes identical.
    #[test]
    fn test_byte_exact_match_with_python() {
        let msg = build_request("test-cascade-id", "hello", "ya29.test-token", 1026);
        let hex: String = msg.iter().map(|b| format!("{:02x}", b)).collect();
        let expected = "0a0f746573742d636173636164652d696412070a0568656c6c6f\
            1a370a0b616e7469677261766974791a0f796132392e746573742d746f6b656e\
            2202656e3a06312e31362e35620b616e7469677261766974792a280a22120420\
            0170006a0b42041a0230038a020208027a03088208aa010208018202020801\
            3a0208015801";
        assert_eq!(hex, expected, "must be byte-exact match with Python");
        assert_eq!(msg.len(), 127);
    }
}