magi-code 0.77.1

use super::protocol::{
    CapabilitiesResult, EmptyParams, InitializeParams, InitializeResult, MAX_ID_BYTES,
    MAX_NAME_BYTES, MAX_NEGOTIATION_ITEMS, MessageKind, PROTOCOL_VERSION, ProtocolLimits,
    RequestIdentity, ResponsePayload, ServiceErrorCode, ServiceMessage, ServiceRequest,
    ServiceResponse, StatusResult, payload_is_bounded,
};
use std::{
    collections::HashSet,
    fmt,
    sync::{Arc, Mutex},
};

const INITIALIZE_METHOD: &str = "initialize";
const STATUS_METHOD: &str = "status";
const CAPABILITIES_METHOD: &str = "capabilities";
const OPERATIONS: [&str; 20] = [
    INITIALIZE_METHOD,
    STATUS_METHOD,
    CAPABILITIES_METHOD,
    super::protocol::TURN_START_METHOD,
    super::protocol::TURN_CANCEL_METHOD,
    "auth.status",
    "auth.login.start",
    "auth.login.callback",
    "auth.login.cancel",
    "auth.logout",
    "session.list",
    "session.create",
    "session.open",
    "session.replay",
    "session.close",
    "catalog.providers",
    "catalog.models",
    "catalog.refresh",
    "config.get",
    "config.set",
];
const EVENTS: [&str; 6] = [
    super::protocol::TURN_STARTED_EVENT,
    super::protocol::ASSISTANT_DELTA_EVENT,
    super::protocol::TURN_TERMINAL_EVENT,
    super::activity::ACTIVITY_EVENT,
    "auth.login.progress",
    "auth.login.terminal",
];

#[derive(Debug, Clone)]
pub(crate) struct ServiceTransportCapabilities {
    names: Vec<String>,
}

impl ServiceTransportCapabilities {
    pub(crate) fn new(names: Vec<String>) -> Self {
        Self { names }
    }
}

#[derive(Clone)]
pub(crate) struct ServiceSnapshot {
    provider_auth_readiness: Arc<dyn Fn() -> anyhow::Result<bool> + Send + Sync>,
}

impl fmt::Debug for ServiceSnapshot {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("ServiceSnapshot")
            .finish_non_exhaustive()
    }
}

impl ServiceSnapshot {
    pub(crate) fn new(
        provider_auth_readiness: impl Fn() -> anyhow::Result<bool> + Send + Sync + 'static,
    ) -> Self {
        Self {
            provider_auth_readiness: Arc::new(provider_auth_readiness),
        }
    }

    fn provider_auth_ready(&self) -> anyhow::Result<bool> {
        (self.provider_auth_readiness)()
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct RequestIdTracker {
    in_flight: Arc<Mutex<HashSet<String>>>,
}

impl RequestIdTracker {
    fn reserve(&self, request_id: &str) -> Result<RequestIdGuard, ServiceErrorCode> {
        if request_id.is_empty()
            || request_id.len() > MAX_ID_BYTES
            || request_id.chars().any(char::is_control)
        {
            return Err(ServiceErrorCode::InvalidRequest);
        }
        let mut in_flight = self
            .in_flight
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if !in_flight.insert(request_id.to_string()) {
            return Err(ServiceErrorCode::DuplicateRequestId);
        }
        Ok(RequestIdGuard {
            tracker: self.clone(),
            request_id: request_id.to_string(),
        })
    }
}

/// Holds a request ID for the duration of one in-flight operation.
///
/// IDs are removed when this guard is dropped. An asynchronous adapter can keep the guard with
/// its work item and pass it to `dispatch_with_request_id` when the response is ready.
#[derive(Debug)]
pub(crate) struct RequestIdGuard {
    tracker: RequestIdTracker,
    request_id: String,
}

impl RequestIdGuard {
    pub(crate) fn request_id(&self) -> &str {
        &self.request_id
    }
}

impl Drop for RequestIdGuard {
    fn drop(&mut self) {
        let mut in_flight = self
            .tracker
            .in_flight
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        in_flight.remove(&self.request_id);
    }
}

/// Owns one outbound message batch and, when present, the request ID reservation for it.
///
/// A transport must keep this value alive until every message in the batch has been encoded,
/// written, and flushed. This lets a future streaming adapter keep the reservation through its
/// terminal event instead of releasing it when dispatch merely constructs the batch.
#[derive(Debug)]
pub(crate) struct ServiceOutbound {
    messages: Vec<ServiceMessage>,
    _request_id_guard: Option<RequestIdGuard>,
}

impl ServiceOutbound {
    pub(crate) fn unguarded(messages: Vec<ServiceMessage>) -> Self {
        Self {
            messages,
            _request_id_guard: None,
        }
    }

    pub(crate) fn guarded(messages: Vec<ServiceMessage>, guard: RequestIdGuard) -> Self {
        Self {
            messages,
            _request_id_guard: Some(guard),
        }
    }

    pub(crate) fn messages(&self) -> &[ServiceMessage] {
        &self.messages
    }
}

#[derive(Debug)]
pub(crate) struct ServiceDispatcher {
    snapshot: ServiceSnapshot,
    initialized: bool,
    transport_capabilities: ServiceTransportCapabilities,
    request_ids: RequestIdTracker,
}

impl ServiceDispatcher {
    pub(crate) fn new(
        snapshot: ServiceSnapshot,
        transport_capabilities: ServiceTransportCapabilities,
    ) -> Self {
        Self {
            snapshot,
            initialized: false,
            transport_capabilities,
            request_ids: RequestIdTracker::default(),
        }
    }

    /// Reserve an ID before queuing asynchronous work.
    ///
    /// The reservation lasts until the returned guard is dropped. Completed synchronous requests
    /// do not remain in a history set, so a client may reuse their IDs after completion.
    pub(crate) fn begin_request(
        &self,
        request_id: &str,
    ) -> Result<RequestIdGuard, ServiceErrorCode> {
        self.request_ids.reserve(request_id)
    }

    pub(crate) fn is_initialized(&self) -> bool {
        self.initialized
    }

    pub(crate) fn request_error_code(&self, request: &ServiceRequest) -> Option<ServiceErrorCode> {
        request_error_code(request)
    }

    /// Build an error response while retaining a safe request ID reservation through output.
    /// Malformed records without a safe ID remain unguarded. For a safe active ID, duplicate
    /// rejection takes precedence over the original error and does not acquire a second guard.
    pub(crate) fn error_outbound(
        &mut self,
        identity: RequestIdentity,
        code: ServiceErrorCode,
    ) -> ServiceOutbound {
        let (code, guard) = match identity.request_id.as_deref() {
            Some(request_id) => match self.begin_request(request_id) {
                Ok(guard) => (code, Some(guard)),
                Err(ServiceErrorCode::DuplicateRequestId) => {
                    (ServiceErrorCode::DuplicateRequestId, None)
                }
                Err(_) => (code, None),
            },
            None => (code, None),
        };
        match guard {
            Some(guard) => ServiceOutbound::guarded(one_error(identity, code), guard),
            None => ServiceOutbound::unguarded(one_error(identity, code)),
        }
    }

    /// Dispatch one request. A valid request's ID remains reserved in the returned outbound batch
    /// until the transport drops that batch after successful output.
    pub(crate) fn dispatch(&mut self, request: ServiceRequest) -> ServiceOutbound {
        let identity = RequestIdentity::from_request(&request);
        if let Some(code) = request_error_code(&request) {
            return self.error_outbound(identity, code);
        }
        let guard = match self.begin_request(&request.request_id) {
            Ok(guard) => guard,
            Err(code) => return ServiceOutbound::unguarded(one_error(identity, code)),
        };
        self.dispatch_with_request_id(request, guard)
    }

    /// Dispatch a request whose ID was reserved before asynchronous work began.
    ///
    /// The returned outbound batch owns the guard. Keeping the batch until all of its messages are
    /// written prevents a second request with the same ID from being accepted while output is
    /// still in flight. A future adapter can put intermediate events and the terminal event in one
    /// batch.
    pub(crate) fn dispatch_with_request_id(
        &mut self,
        request: ServiceRequest,
        guard: RequestIdGuard,
    ) -> ServiceOutbound {
        let identity = RequestIdentity::from_request(&request);
        if guard.request_id() != request.request_id {
            return ServiceOutbound::guarded(
                one_error(identity, ServiceErrorCode::InvalidRequest),
                guard,
            );
        }
        if let Some(code) = request_error_code(&request) {
            return ServiceOutbound::guarded(one_error(identity, code), guard);
        }
        self.dispatch_reserved(request, guard)
    }

    fn dispatch_reserved(
        &mut self,
        request: ServiceRequest,
        guard: RequestIdGuard,
    ) -> ServiceOutbound {
        let identity = RequestIdentity::from_request(&request);
        let response = match request.method.as_str() {
            INITIALIZE_METHOD => self.initialize(&request, identity),
            STATUS_METHOD => self.status(&request, identity),
            CAPABILITIES_METHOD => self.capabilities(&request, identity),
            _ => ServiceResponse::error(identity, ServiceErrorCode::UnsupportedOperation),
        };
        ServiceOutbound::guarded(vec![ServiceMessage::response(response)], guard)
    }

    fn initialize(
        &mut self,
        request: &ServiceRequest,
        identity: RequestIdentity,
    ) -> ServiceResponse {
        let params = match serde_json::from_value::<InitializeParams>(request.payload.clone()) {
            Ok(params) => params,
            Err(_) => return ServiceResponse::error(identity, ServiceErrorCode::InvalidPayload),
        };
        if params.supported_protocol_versions.is_empty()
            || params.supported_protocol_versions.len() > MAX_NEGOTIATION_ITEMS
        {
            return ServiceResponse::error(identity, ServiceErrorCode::LimitExceeded);
        }
        if !params
            .supported_protocol_versions
            .contains(&PROTOCOL_VERSION)
        {
            return ServiceResponse::error(identity, ServiceErrorCode::UnsupportedVersion);
        }
        if params.requested_capabilities.len() > MAX_NEGOTIATION_ITEMS {
            return ServiceResponse::error(identity, ServiceErrorCode::LimitExceeded);
        }
        for capability in &params.requested_capabilities {
            if capability.len() > MAX_NAME_BYTES {
                return ServiceResponse::error(identity, ServiceErrorCode::LimitExceeded);
            }
            if capability.is_empty()
                || capability.chars().any(char::is_control)
                || !OPERATIONS.contains(&capability.as_str())
            {
                return ServiceResponse::error(identity, ServiceErrorCode::UnsupportedCapability);
            }
        }

        self.initialized = true;
        ServiceResponse::success(
            request,
            ResponsePayload::Initialize(InitializeResult {
                protocol_version: PROTOCOL_VERSION,
                server_name: "magi-code".to_string(),
                server_version: env!("CARGO_PKG_VERSION").to_string(),
                limits: ProtocolLimits::current(),
                capabilities: self.capabilities_result(),
            }),
        )
    }

    fn status(&self, request: &ServiceRequest, identity: RequestIdentity) -> ServiceResponse {
        if !self.initialized {
            return ServiceResponse::error(identity, ServiceErrorCode::NotInitialized);
        }
        if serde_json::from_value::<EmptyParams>(request.payload.clone()).is_err() {
            return ServiceResponse::error(identity, ServiceErrorCode::InvalidPayload);
        }
        let provider_auth_ready = match self.snapshot.provider_auth_ready() {
            Ok(provider_auth_ready) => provider_auth_ready,
            Err(_) => return ServiceResponse::error(identity, ServiceErrorCode::InternalError),
        };
        ServiceResponse::success(
            request,
            ResponsePayload::Status(StatusResult {
                service: "ready".to_string(),
                provider_auth_ready,
            }),
        )
    }

    fn capabilities(&self, request: &ServiceRequest, identity: RequestIdentity) -> ServiceResponse {
        if !self.initialized {
            return ServiceResponse::error(identity, ServiceErrorCode::NotInitialized);
        }
        if serde_json::from_value::<EmptyParams>(request.payload.clone()).is_err() {
            return ServiceResponse::error(identity, ServiceErrorCode::InvalidPayload);
        }
        ServiceResponse::success(
            request,
            ResponsePayload::Capabilities(self.capabilities_result()),
        )
    }

    fn capabilities_result(&self) -> CapabilitiesResult {
        CapabilitiesResult {
            protocol_versions: vec![PROTOCOL_VERSION],
            operations: OPERATIONS
                .iter()
                .map(|value| (*value).to_string())
                .collect(),
            events: EVENTS.iter().map(|value| (*value).to_string()).collect(),
            transports: self.transport_capabilities.names.clone(),
            limits: ProtocolLimits::current(),
            activity: super::activity::ActivityCapabilities::current(),
        }
    }
}

fn request_error_code(request: &ServiceRequest) -> Option<ServiceErrorCode> {
    let identity = RequestIdentity::from_request(request);
    if identity.request_id.as_deref() != Some(request.request_id.as_str())
        || identity.method.as_deref() != Some(request.method.as_str())
        || request.session_id.as_deref() != identity.session_id.as_deref()
    {
        return Some(ServiceErrorCode::InvalidRequest);
    }
    if request.kind != MessageKind::Request {
        return Some(ServiceErrorCode::InvalidRequest);
    }
    if let Err(code) = payload_is_bounded(&request.payload) {
        return Some(code);
    }
    if request.protocol_version != PROTOCOL_VERSION {
        return Some(ServiceErrorCode::UnsupportedVersion);
    }
    None
}

fn one_error(identity: RequestIdentity, code: ServiceErrorCode) -> Vec<ServiceMessage> {
    vec![ServiceMessage::response(ServiceResponse::error(
        identity, code,
    ))]
}

#[cfg(test)]
mod tests {
    use super::super::protocol::MAX_RECORD_BYTES;
    use super::*;
    use serde_json::{Value, json};

    use std::sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    };
    const TEST_TRANSPORT: &str = "test";

    fn dispatcher(provider_auth_ready: bool) -> ServiceDispatcher {
        ServiceDispatcher::new(
            ServiceSnapshot::new(move || Ok(provider_auth_ready)),
            ServiceTransportCapabilities::new(vec![TEST_TRANSPORT.to_string()]),
        )
    }

    fn request(id: &str, method: &str, payload: Value) -> ServiceRequest {
        ServiceRequest::new(id, method, payload)
    }

    fn response_value(outbound: ServiceOutbound) -> Value {
        let Some(ServiceMessage::Response(response)) = outbound.messages().first() else {
            panic!("dispatcher did not return one response");
        };
        serde_json::to_value(response).unwrap()
    }

    fn response_value_ref(outbound: &ServiceOutbound) -> Value {
        let Some(ServiceMessage::Response(response)) = outbound.messages().first() else {
            panic!("dispatcher did not return one response");
        };
        serde_json::to_value(response).unwrap()
    }

    fn decoded_error_outbound(dispatcher: &mut ServiceDispatcher, bytes: &[u8]) -> ServiceOutbound {
        let error = super::super::protocol::decode_request(bytes).unwrap_err();
        dispatcher.error_outbound(error.identity, error.code)
    }

    #[test]
    fn initialize_negotiates_version_and_returns_limits_capabilities_and_injected_transport() {
        let mut dispatcher = dispatcher(false);
        let outbound = dispatcher.dispatch(request(
            "init-1",
            INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1]}),
        ));
        assert!(matches!(
            dispatcher.begin_request("init-1"),
            Err(ServiceErrorCode::DuplicateRequestId)
        ));
        let value = response_value(outbound);

        assert_eq!(value["kind"], "response");
        assert_eq!(value["request_id"], "init-1");
        assert_eq!(value["method"], INITIALIZE_METHOD);
        assert_eq!(value["payload"]["protocol_version"], PROTOCOL_VERSION);
        assert_eq!(
            value["payload"]["limits"]["max_record_bytes"],
            MAX_RECORD_BYTES
        );
        assert_eq!(
            value["payload"]["capabilities"]["operations"],
            json!(OPERATIONS)
        );
        assert_eq!(
            value["payload"]["capabilities"]["transports"],
            json!([TEST_TRANSPORT])
        );
    }

    #[test]
    fn status_requires_initialization_and_only_exposes_readiness() {
        let mut dispatcher = dispatcher(true);
        let before_init =
            response_value(dispatcher.dispatch(request("status-0", STATUS_METHOD, json!({}))));
        assert_eq!(before_init["error"]["code"], "not_initialized");

        dispatcher.dispatch(request(
            "init-1",
            INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1]}),
        ));
        let status =
            response_value(dispatcher.dispatch(request("status-1", STATUS_METHOD, json!({}))));

        assert_eq!(status["request_id"], "status-1");
        assert_eq!(status["payload"]["service"], "ready");
        assert_eq!(status["payload"]["provider_auth_ready"], true);
        assert!(status["payload"].get("provider").is_none());
        assert!(status["payload"].get("model").is_none());
    }

    #[test]
    fn status_validates_before_reading_auth_and_sanitizes_read_failures() {
        let read_attempts = Arc::new(AtomicUsize::new(0));
        let attempts = Arc::clone(&read_attempts);
        let mut dispatcher = ServiceDispatcher::new(
            ServiceSnapshot::new(move || {
                attempts.fetch_add(1, Ordering::SeqCst);
                Err(anyhow::anyhow!(
                    "auth read failed at /private/auth.json containing raw-secret"
                ))
            }),
            ServiceTransportCapabilities::new(vec![TEST_TRANSPORT.to_string()]),
        );

        let before_init = response_value(dispatcher.dispatch(request(
            "status-before-init",
            STATUS_METHOD,
            json!({"unexpected": "raw-secret"}),
        )));
        assert_eq!(before_init["error"]["code"], "not_initialized");
        assert_eq!(read_attempts.load(Ordering::SeqCst), 0);

        dispatcher.dispatch(request(
            "init-1",
            INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1]}),
        ));
        let invalid_payload = response_value(dispatcher.dispatch(request(
            "status-invalid-payload",
            STATUS_METHOD,
            json!({"unexpected": "raw-secret"}),
        )));
        assert_eq!(invalid_payload["error"]["code"], "invalid_payload");
        assert_eq!(read_attempts.load(Ordering::SeqCst), 0);

        let mut status_request = request("status-read-failure", STATUS_METHOD, json!({}));
        status_request.session_id = Some("session-1".to_string());
        let read_failure = response_value(dispatcher.dispatch(status_request));
        assert_eq!(read_failure["request_id"], "status-read-failure");
        assert_eq!(read_failure["session_id"], "session-1");
        assert_eq!(read_failure["method"], STATUS_METHOD);
        assert!(read_failure["payload"].is_null());
        assert_eq!(read_failure["error"]["code"], "internal_error");
        assert_eq!(
            read_failure["error"]["message"],
            ServiceErrorCode::InternalError.message()
        );
        assert!(!read_failure.to_string().contains("/private/auth.json"));
        assert!(!read_failure.to_string().contains("raw-secret"));
        assert_eq!(read_attempts.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn repeated_successful_initialization_keeps_service_ready() {
        let mut dispatcher = dispatcher(false);
        for id in ["init-1", "init-2"] {
            let response = response_value(dispatcher.dispatch(request(
                id,
                INITIALIZE_METHOD,
                json!({"supported_protocol_versions": [1]}),
            )));
            assert_eq!(response["request_id"], id);
            assert!(response["error"].is_null());
        }

        let status =
            response_value(dispatcher.dispatch(request("status-1", STATUS_METHOD, json!({}))));
        assert_eq!(status["payload"]["service"], "ready");
    }

    #[test]
    fn failed_initialization_preserves_state_and_strict_payload_validation() {
        let cases = [
            (
                json!({"supported_protocol_versions": [99]}),
                "unsupported_version",
            ),
            (
                json!({"supported_protocol_versions": [1], "requested_capabilities": ["turn.activity"]}),
                "unsupported_capability",
            ),
            (
                json!({"supported_protocol_versions": [1], "requested_capabilities": ["approvals"]}),
                "unsupported_capability",
            ),
            (
                json!({"supported_protocol_versions": [1], "future": true}),
                "invalid_payload",
            ),
            (
                json!({"supported_protocol_versions": ["1"]}),
                "invalid_payload",
            ),
            (
                json!({"supported_protocol_versions": [1], "requested_capabilities": null}),
                "invalid_payload",
            ),
            (json!({"supported_protocol_versions": []}), "limit_exceeded"),
            (
                json!({"supported_protocol_versions": vec![1; MAX_NEGOTIATION_ITEMS + 1]}),
                "limit_exceeded",
            ),
        ];
        for initialized in [false, true] {
            let mut dispatcher = dispatcher(false);
            if initialized {
                dispatcher.dispatch(request(
                    "init",
                    INITIALIZE_METHOD,
                    json!({"supported_protocol_versions": [1]}),
                ));
            }
            for (payload, code) in &cases {
                let rejected =
                    dispatcher.dispatch(request("retry", INITIALIZE_METHOD, payload.clone()));
                let value = response_value_ref(&rejected);
                assert_eq!(value["error"]["code"], *code);
                assert_eq!(value["request_id"], "retry");
                assert!(value["payload"].is_null());
                assert!(matches!(
                    dispatcher.begin_request("retry"),
                    Err(ServiceErrorCode::DuplicateRequestId)
                ));
                drop(rejected);
                let status =
                    response_value(dispatcher.dispatch(request("retry", STATUS_METHOD, json!({}))));
                if initialized {
                    assert_eq!(status["payload"]["service"], "ready");
                } else {
                    assert_eq!(status["error"]["code"], "not_initialized");
                }
            }
            let recovered = response_value(dispatcher.dispatch(request(
                "retry",
                INITIALIZE_METHOD,
                json!({"supported_protocol_versions": [1]}),
            )));
            assert!(recovered["error"].is_null());
        }
    }

    #[test]
    fn negotiation_list_bound_does_not_truncate_discovery_or_subscribe_to_events() {
        let mut dispatcher = dispatcher(false);
        let rejected = response_value(dispatcher.dispatch(request(
            "too-many", INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1], "requested_capabilities": &OPERATIONS[..MAX_NEGOTIATION_ITEMS + 1]}),
        )));
        assert_eq!(rejected["error"]["code"], "limit_exceeded");
        assert!(!dispatcher.is_initialized());

        // Every advertised operation is a valid requirement, but not all fit in one request.
        for operations in OPERATIONS.chunks(MAX_NEGOTIATION_ITEMS) {
            let accepted = dispatcher.dispatch(request(
                "init",
                INITIALIZE_METHOD,
                json!({"supported_protocol_versions": [1], "requested_capabilities": operations}),
            ));
            for message in accepted.messages() {
                super::super::protocol::encode_message(message).unwrap();
            }
            let value = response_value(accepted);
            assert!(value["error"].is_null());
            let capabilities = &value["payload"]["capabilities"];
            assert_eq!(capabilities["operations"], json!(OPERATIONS));
            assert_eq!(capabilities["events"], json!(EVENTS));
            assert!(capabilities["operations"].as_array().unwrap().len() > MAX_NEGOTIATION_ITEMS);
            assert_eq!(capabilities["limits"], value["payload"]["limits"]);
            let discovered = response_value(dispatcher.dispatch(request(
                "discover",
                CAPABILITIES_METHOD,
                json!({}),
            )));
            assert_eq!(&discovered["payload"], capabilities);
        }
    }

    #[test]
    fn initialize_rejects_unknown_capability_without_leaking_request_data() {
        let mut dispatcher = dispatcher(false);
        let value = response_value(dispatcher.dispatch(request(
            "init-1",
            INITIALIZE_METHOD,
            json!({
                "supported_protocol_versions": [1],
                "requested_capabilities": ["future.secret-token"]
            }),
        )));

        assert_eq!(value["error"]["code"], "unsupported_capability");
        assert_eq!(
            value["error"]["message"],
            ServiceErrorCode::UnsupportedCapability.message()
        );
        assert!(!value.to_string().contains("future.secret-token"));
    }

    #[test]
    fn dispatcher_rejects_wrong_version_and_unknown_operation_with_correlation() {
        let mut dispatcher = dispatcher(false);
        let mut wrong_version = request("version-1", STATUS_METHOD, json!({}));
        wrong_version.protocol_version = 99;
        let version_outbound = dispatcher.dispatch(wrong_version);
        assert!(matches!(
            dispatcher.begin_request("version-1"),
            Err(ServiceErrorCode::DuplicateRequestId)
        ));
        let version = response_value(version_outbound);
        assert_eq!(version["error"]["code"], "unsupported_version");
        assert_eq!(version["request_id"], "version-1");

        let unknown = response_value(dispatcher.dispatch(request(
            "unknown-1",
            "not-an-operation",
            json!({}),
        )));
        assert_eq!(unknown["error"]["code"], "unsupported_operation");
        assert_eq!(unknown["request_id"], "unknown-1");
    }

    #[test]
    fn duplicate_request_id_is_rejected_only_while_the_first_request_is_in_flight() {
        let mut dispatcher = dispatcher(false);
        let first = dispatcher.begin_request("same-id").unwrap();
        let duplicate = response_value(dispatcher.dispatch(request(
            "same-id",
            INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1]}),
        )));
        assert_eq!(duplicate["error"]["code"], "duplicate_request_id");
        assert_eq!(
            duplicate["error"]["message"],
            ServiceErrorCode::DuplicateRequestId.message()
        );

        drop(first);
        let accepted = response_value(dispatcher.dispatch(request(
            "same-id",
            INITIALIZE_METHOD,
            json!({"supported_protocol_versions": [1]}),
        )));
        assert!(accepted["error"].is_null());
    }

    #[test]
    fn fresh_malformed_request_ids_keep_decoder_errors_and_safe_correlation_until_output() {
        let text = "x".repeat(super::super::protocol::MAX_STRING_BYTES);
        let cases = vec![
            (
                json!({
                    "protocol_version": PROTOCOL_VERSION,
                    "kind": "request",
                    "request_id": "invalid-id",
                    "session_id": "session-invalid",
                    "method": "status",
                    "payload": {},
                    "future": "malformed-secret",
                }),
                "invalid_request",
            ),
            (
                json!({
                    "protocol_version": PROTOCOL_VERSION,
                    "kind": "request",
                    "request_id": "limit-id",
                    "session_id": "session-limit",
                    "method": "status",
                    "payload": {"x".repeat(super::super::protocol::MAX_STRING_BYTES + 1): true},
                }),
                "limit_exceeded",
            ),
            (
                json!({
                    "protocol_version": PROTOCOL_VERSION,
                    "kind": "request",
                    "request_id": "payload-id",
                    "session_id": "session-payload",
                    "method": "status",
                    "payload": {"values": [text.clone(), text.clone(), text]},
                }),
                "payload_too_large",
            ),
        ];
        let mut dispatcher = dispatcher(false);

        for (request, expected_code) in cases {
            let request_id = request["request_id"].as_str().unwrap();
            let bytes = serde_json::to_vec(&request).unwrap();
            let outbound = decoded_error_outbound(&mut dispatcher, &bytes);
            let response = response_value_ref(&outbound);
            assert_eq!(response["error"]["code"], expected_code);
            assert_eq!(response["request_id"], request["request_id"]);
            assert_eq!(response["session_id"], request["session_id"]);
            assert_eq!(response["method"], request["method"]);
            let encoded =
                super::super::protocol::encode_message(outbound.messages().first().unwrap())
                    .unwrap();
            assert!(
                !String::from_utf8(encoded)
                    .unwrap()
                    .contains("malformed-secret")
            );
            assert!(matches!(
                dispatcher.begin_request(request_id),
                Err(ServiceErrorCode::DuplicateRequestId)
            ));

            drop(outbound);
            assert!(dispatcher.begin_request(request_id).is_ok());
        }
    }

    #[test]
    fn uncorrelatable_decode_errors_do_not_echo_or_reserve_request_ids() {
        const UNSAFE_REQUEST_ID: &str = "unsafe\nid";
        let unsafe_request = json!({
            "protocol_version": PROTOCOL_VERSION,
            "kind": "request",
            "request_id": UNSAFE_REQUEST_ID,
            "session_id": "safe-session",
            "method": "status",
            "payload": {},
            "future": "malformed-secret",
        });
        let missing_request_id = json!({
            "protocol_version": PROTOCOL_VERSION,
            "kind": "request",
            "session_id": "missing-id-session",
            "method": "status",
            "payload": {},
            "future": "malformed-secret",
        });
        let cases = vec![
            (b"not-json".to_vec(), "invalid_json", None, None),
            (
                vec![b'x'; super::super::protocol::MAX_RECORD_BYTES + 1],
                "record_too_large",
                None,
                None,
            ),
            (
                serde_json::to_vec(&unsafe_request).unwrap(),
                "invalid_request",
                Some("safe-session"),
                Some("status"),
            ),
            (
                serde_json::to_vec(&missing_request_id).unwrap(),
                "invalid_request",
                Some("missing-id-session"),
                Some("status"),
            ),
        ];
        let mut dispatcher = dispatcher(false);

        for (bytes, expected_code, expected_session, expected_method) in cases {
            let outbound = decoded_error_outbound(&mut dispatcher, &bytes);
            let response = response_value_ref(&outbound);
            assert_eq!(response["error"]["code"], expected_code);
            assert!(response["request_id"].is_null());
            assert_eq!(response["session_id"].as_str(), expected_session);
            assert_eq!(response["method"].as_str(), expected_method);
            let encoded =
                super::super::protocol::encode_message(outbound.messages().first().unwrap())
                    .unwrap();
            assert!(
                !String::from_utf8(encoded)
                    .unwrap()
                    .contains("malformed-secret")
            );
            drop(outbound);
        }

        assert!(matches!(
            dispatcher.begin_request(UNSAFE_REQUEST_ID),
            Err(ServiceErrorCode::InvalidRequest)
        ));
        assert!(dispatcher.begin_request("after-unsafe-id").is_ok());
    }

    #[test]
    fn reserved_request_id_can_be_consumed_after_async_work() {
        let mut dispatcher = dispatcher(false);
        let guard = dispatcher.begin_request("async-id").unwrap();
        let outbound = dispatcher.dispatch_with_request_id(
            request(
                "async-id",
                INITIALIZE_METHOD,
                json!({"supported_protocol_versions": [1]}),
            ),
            guard,
        );
        assert_eq!(outbound.messages().len(), 1);
        assert!(matches!(
            dispatcher.begin_request("async-id"),
            Err(ServiceErrorCode::DuplicateRequestId)
        ));
        assert_eq!(response_value(outbound)["request_id"], "async-id");
        assert!(dispatcher.begin_request("async-id").is_ok());
    }
}