pi_agent_rust 0.8.0

use super::super::test_server::Fixture;
use super::*;
use serde_json::json;

async fn observe_dispatched_request(peer: &Fixture, method: &str) {
    // The application lane belongs to the request under test. Observe the
    // actual peer through a separate transport request; a first Pending poll
    // alone proves queue admission, not that the writer sent anything.
    let (id, response) = peer.client.rpc.request("test/frames", json!({})).unwrap();
    let frames = crate::lsp::jsonrpc::await_completion(response, Duration::from_secs(5), || {
        peer.client.rpc.cancel_request(id);
    })
    .await
    .expect("peer observation completed")
    .expect("peer observation succeeded");
    assert!(
        frames
            .as_array()
            .expect("peer frames")
            .iter()
            .any(|frame| frame["method"] == method),
        "the peer must observe the request before its owner abandons it"
    );
}

fn saturate_request_admission(peer: &Fixture) -> Vec<u64> {
    let mut pending = Vec::new();
    // Reach the production limit with unanswered requests. The independent
    // bound makes a regression to unbounded admission fail promptly.
    for _ in 0..4096 {
        match peer.client.rpc.request("test/hang", Value::Null) {
            Ok((id, _response)) => pending.push(id),
            Err(TransportError::Backpressure(_)) => {
                assert!(!pending.is_empty());
                return pending;
            }
            Err(error) => panic!("unexpected admission error: {error:?}"),
        }
    }
    panic!("unanswered requests must have a bounded admission limit");
}

#[test]
fn admission_backpressure_waits_and_posts_a_command_only_once() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.configure(json!({"workspace/executeCommand":[{"result":{"executed":true}}]}));
    let pending = saturate_request_admission(&peer);
    peer.runtime.block_on(async {
        let mut command = Box::pin(peer.client.call(
            "workspace/executeCommand",
            json!({"command":"run-once"}),
            Duration::from_secs(5),
        ));
        assert!(
            futures::poll!(command.as_mut()).is_pending(),
            "saturation must suspend within the request budget"
        );
        assert!(peer.client.is_alive());
        peer.client.rpc.cancel_request(pending[0]);
        assert_eq!(command.await.unwrap(), json!({"executed":true}));
    });
    let frames = peer.frames();
    let commands: Vec<_> = frames
        .iter()
        .filter(|frame| frame["method"] == "workspace/executeCommand")
        .collect();
    assert_eq!(
        commands.len(),
        1,
        "queue admission must not replay a command"
    );
    assert_eq!(commands[0]["params"]["command"], "run-once");
    assert!(peer.client.is_alive());
}

#[test]
fn admission_backpressure_spends_the_original_deadline_without_dispatch() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let pending = saturate_request_admission(&peer);
    let started = std::time::Instant::now();
    let error = peer
        .runtime
        .block_on(
            peer.client
                .call("test/must-not-run", Value::Null, Duration::from_millis(30)),
        )
        .unwrap_err();
    assert!(matches!(error, LspCallError::Timeout { timeout_ms: 30 }));
    assert!(started.elapsed() < Duration::from_secs(2));
    assert!(peer.client.is_alive());
    assert!(matches!(
        peer.client.rpc.request("test/still-full", Value::Null),
        Err(TransportError::Backpressure(_))
    ));
    peer.client.rpc.cancel_request(pending[0]);
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "test/must-not-run")
    );
}

#[test]
fn owner_cancellation_while_waiting_for_capacity_prevents_dispatch() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let pending = saturate_request_admission(&peer);
    let owner = AgentCx::for_request();
    let budget = RequestBudget {
        start: now(&owner),
        owner,
        timeout: Duration::from_secs(5),
    };
    peer.runtime.block_on(async {
        let mut request = Box::pin(peer.client.call_with_budget(
            "test/must-not-run",
            Value::Null,
            &budget,
        ));
        assert!(futures::poll!(request.as_mut()).is_pending());
        budget.owner.cancel_with(
            asupersync::types::CancelKind::User,
            Some("cancel admission wait"),
        );
        assert!(matches!(request.await, Err(LspCallError::Cancelled)));
    });
    assert!(peer.client.is_alive());
    peer.client.rpc.cancel_request(pending[0]);
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "test/must-not-run")
    );
}

#[test]
fn dropping_an_admission_wait_releases_the_lane_without_cancelling_other_work() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let pending = saturate_request_admission(&peer);
    peer.runtime.block_on(async {
        let mut request = Box::pin(peer.client.call(
            "test/must-not-run",
            Value::Null,
            Duration::from_secs(30),
        ));
        assert!(futures::poll!(request.as_mut()).is_pending());
        drop(request);
        assert!(matches!(
            peer.client.rpc.request("test/still-full", Value::Null),
            Err(TransportError::Backpressure(_))
        ));
        peer.client.rpc.cancel_request(pending[0]);
        assert_eq!(
            peer.client
                .call("test/success", Value::Null, Duration::from_secs(5))
                .await
                .unwrap(),
            json!({"ok":true})
        );
    });
    let frames = peer.frames();
    assert!(
        !frames
            .iter()
            .any(|frame| frame["method"] == "test/must-not-run")
    );
    assert!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "$/cancelRequest")
            .all(|frame| frame["params"]["id"] == json!(pending[0]))
    );
    assert!(peer.client.is_alive());
}

#[test]
fn zero_budget_never_dispatches_a_request() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let error = peer
        .runtime
        .block_on(peer.client.call("test/success", json!({}), Duration::ZERO))
        .unwrap_err();
    assert!(matches!(error, LspCallError::Timeout { timeout_ms: 0 }));
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "test/success")
    );
}

#[test]
fn serialized_lane_admission_spends_the_request_timeout() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.runtime.block_on(async {
        let owner = AgentCx::for_current_or_request();
        let guard = OwnedMutexGuard::lock(Arc::clone(&peer.client.request_lane), owner.cx())
            .await
            .unwrap();
        let started = std::time::Instant::now();
        let error = peer
            .client
            .call("test/success", json!({}), Duration::from_millis(30))
            .await
            .unwrap_err();
        assert!(matches!(error, LspCallError::Timeout { timeout_ms: 30 }));
        assert!(started.elapsed() < Duration::from_secs(2));
        drop(guard);
    });
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "test/success")
    );
}

#[test]
fn dropping_a_written_request_cancels_it_before_the_next_dispatch() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.runtime.block_on(async {
        let mut request = Box::pin(peer.client.call(
            "test/hang",
            json!({}),
            Duration::from_secs(30),
        ));
        assert!(futures::poll!(request.as_mut()).is_pending());
        observe_dispatched_request(&peer, "test/hang").await;
        drop(request);
        assert_eq!(
            peer.client
                .call("test/success", json!({}), Duration::from_secs(5))
                .await
                .unwrap(),
            json!({"ok":true})
        );
    });
    let frames = peer.frames();
    let posted = frames
        .iter()
        .position(|frame| frame["method"] == "test/hang")
        .unwrap();
    let cancelled = frames
        .iter()
        .position(|frame| frame["method"] == "$/cancelRequest")
        .unwrap();
    let successor = frames
        .iter()
        .position(|frame| frame["method"] == "test/success")
        .unwrap();
    assert!(posted < cancelled && cancelled < successor);
    assert_eq!(frames[cancelled]["params"]["id"], frames[posted]["id"]);
    assert_eq!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "$/cancelRequest")
            .count(),
        1
    );
}

#[test]
fn timeout_cancels_once_and_releases_the_lane() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let error = peer.runtime.block_on(async {
        let mut request = Box::pin(peer.client.call(
            "test/hang",
            json!({}),
            Duration::from_millis(30),
        ));
        assert!(futures::poll!(request.as_mut()).is_pending());
        observe_dispatched_request(&peer, "test/hang").await;
        request.await.unwrap_err()
    });
    assert!(matches!(error, LspCallError::Timeout { timeout_ms: 30 }));
    let frames = peer.frames();
    let request = frames
        .iter()
        .find(|frame| frame["method"] == "test/hang")
        .unwrap();
    let cancellations: Vec<_> = frames
        .iter()
        .filter(|frame| frame["method"] == "$/cancelRequest")
        .collect();
    assert_eq!(cancellations.len(), 1);
    assert_eq!(cancellations[0]["params"]["id"], request["id"]);
    assert!(peer.client.is_alive());
}

#[test]
fn completed_responses_do_not_emit_spurious_cancellation() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let result = peer
        .runtime
        .block_on(
            peer.client
                .call("test/success", json!({}), Duration::from_secs(5)),
        )
        .unwrap();
    assert_eq!(result, json!({"ok":true}));
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "$/cancelRequest")
    );
}

#[test]
fn completed_command_errors_are_not_replayed_or_cancelled() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.configure(json!({"workspace/executeCommand":[{"error":{"code":-32801,"message":"content modified"}}]}));
    let error = peer
        .runtime
        .block_on(peer.client.call(
            "workspace/executeCommand",
            json!({"command":"fixture"}),
            Duration::from_secs(5),
        ))
        .unwrap_err();
    assert!(matches!(
        error,
        LspCallError::Transport(TransportError::Server(_))
    ));
    let frames = peer.frames();
    assert_eq!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "workspace/executeCommand")
            .count(),
        1
    );
    assert!(
        !frames
            .iter()
            .any(|frame| frame["method"] == "$/cancelRequest")
    );
}

#[test]
fn owner_cancellation_before_admission_prevents_dispatch() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    let owner = AgentCx::for_request();
    let budget = RequestBudget {
        start: now(&owner),
        owner,
        timeout: Duration::from_secs(5),
    };
    budget.owner.cancel_with(
        asupersync::types::CancelKind::User,
        Some("test owner cancelled"),
    );
    let error = peer
        .runtime
        .block_on(
            peer.client
                .call_with_budget("test/success", json!({}), &budget),
        )
        .unwrap_err();
    assert!(matches!(error, LspCallError::Cancelled));
    assert!(
        !peer
            .frames()
            .iter()
            .any(|frame| frame["method"] == "test/success")
    );
}

#[test]
fn prepare_rename_retries_during_warmup_on_no_references_found() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.configure(json!({
        "textDocument/prepareRename": [
            {"error": {"code": -32602, "message": "No references found at position"}},
            {"result": {"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 5}}}
        ]
    }));
    let result = peer
        .runtime
        .block_on(peer.client.call(
            "textDocument/prepareRename",
            json!({"textDocument": {"uri": "file:///test.rs"}, "position": {"line": 0, "character": 0}}),
            Duration::from_secs(5),
        ))
        .unwrap();
    assert_eq!(
        result,
        json!({"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 5}})
    );
    let frames = peer.frames();
    assert_eq!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "textDocument/prepareRename")
            .count(),
        2
    );
}

#[test]
fn prepare_rename_retries_during_warmup_on_null_until_target_ready() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.configure(json!({
        "textDocument/prepareRename": [
            {"result": null},
            {"result": {"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 5}}}
        ]
    }));
    let result = peer
        .runtime
        .block_on(peer.client.call(
            "textDocument/prepareRename",
            json!({"textDocument": {"uri": "file:///test.rs"}, "position": {"line": 0, "character": 0}}),
            Duration::from_secs(5),
        ))
        .unwrap();
    assert_eq!(
        result,
        json!({"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 5}})
    );
    let frames = peer.frames();
    assert_eq!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "textDocument/prepareRename")
            .count(),
        2
    );
}

#[test]
fn prepare_rename_does_not_retry_null_when_server_is_quiescent() {
    let temp = tempfile::tempdir().unwrap();
    let Some(peer) = Fixture::connect(temp.path(), json!({})) else {
        return;
    };
    peer.client.quiescent.store(true, Ordering::SeqCst);
    peer.configure(json!({
        "textDocument/prepareRename": [
            {"result": null}
        ]
    }));
    let result = peer
        .runtime
        .block_on(peer.client.call(
            "textDocument/prepareRename",
            json!({"textDocument": {"uri": "file:///test.rs"}, "position": {"line": 0, "character": 0}}),
            Duration::from_secs(5),
        ))
        .unwrap();
    assert_eq!(result, Value::Null);
    let frames = peer.frames();
    assert_eq!(
        frames
            .iter()
            .filter(|frame| frame["method"] == "textDocument/prepareRename")
            .count(),
        1
    );
}