magi-code 0.80.0

use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use uuid::Uuid;

pub(crate) const PROTOCOL_VERSION: u16 = 2;
pub(crate) const MAX_RECORD_BYTES: usize = 64 * 1024;
pub(crate) const MAX_PAYLOAD_BYTES: usize = 32 * 1024;
pub(crate) const MAX_ERROR_MESSAGE_BYTES: usize = 512;
pub(crate) const MAX_ID_BYTES: usize = 128;
pub(crate) const MAX_NAME_BYTES: usize = 64;
pub(crate) const MAX_STRING_BYTES: usize = 16 * 1024;
pub(crate) const MAX_JSON_DEPTH: usize = 32;
pub(crate) const MAX_ARRAY_ITEMS: usize = 64;
pub(crate) const MAX_OBJECT_FIELDS: usize = 64;
pub(crate) const MAX_NEGOTIATION_ITEMS: usize = 8;

/// Shared admission/worker budget: retain the existing 4 KiB snapshot envelope reserve.
/// Charge JSON encoding, including the field name, not raw manifest/string bytes.
pub(super) fn persistent_assistant_budget(
    manifest: Option<&Value>,
) -> Result<usize, ServiceErrorCode> {
    let reserved = serde_json::to_vec(&json!({"application_manifest":manifest}))
        .map_err(|_| ServiceErrorCode::InternalError)?
        .len();
    if reserved > 18_000 {
        return Err(ServiceErrorCode::LimitExceeded);
    }
    Ok(20_000 - reserved)
}

pub(crate) const TURN_START_METHOD: &str = "turn.start";
pub(crate) const TURN_CANCEL_METHOD: &str = "turn.cancel";
pub(crate) const TURN_STARTED_EVENT: &str = "turn.started";
pub(crate) const ASSISTANT_DELTA_EVENT: &str = "turn.assistant_delta";
pub(crate) const TURN_TERMINAL_EVENT: &str = "turn.terminal";

#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum MessageKind {
    Request,
    Response,
    Event,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ServiceRequest {
    pub(crate) protocol_version: u16,
    pub(crate) kind: MessageKind,
    pub(crate) request_id: String,
    #[serde(default)]
    pub(crate) session_id: Option<String>,
    pub(crate) method: String,
    pub(crate) payload: Value,
}

impl ServiceRequest {
    #[cfg(test)]
    pub(crate) fn new(request_id: &str, method: &str, payload: Value) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Request,
            request_id: request_id.to_string(),
            session_id: None,
            method: method.to_string(),
            payload,
        }
    }
}

#[derive(Debug, Serialize)]
pub(crate) struct ServiceResponse {
    pub(crate) protocol_version: u16,
    pub(crate) kind: MessageKind,
    pub(crate) request_id: Option<String>,
    pub(crate) session_id: Option<String>,
    pub(crate) method: Option<String>,
    pub(crate) payload: Option<ResponsePayload>,
    pub(crate) error: Option<ServiceErrorDto>,
}

#[derive(Debug, Serialize)]
#[serde(untagged)]
pub(crate) enum ResponsePayload {
    Initialize(InitializeResult),
    Status(StatusResult),
    Capabilities(CapabilitiesResult),
    TurnStart(TurnStartResult),
    TurnCancel(TurnCancelResult),
    Auth(serde_json::Value),
    Session(serde_json::Value),
    Configuration(serde_json::Value),
}

#[derive(Debug, Serialize)]
pub(crate) struct InitializeResult {
    pub(crate) protocol_version: u16,
    pub(crate) server_name: String,
    pub(crate) server_version: String,
    pub(crate) limits: ProtocolLimits,
    pub(crate) capabilities: CapabilitiesResult,
}

#[derive(Debug, Serialize)]
pub(crate) struct StatusResult {
    pub(crate) service: String,
    pub(crate) provider_auth_ready: bool,
}

#[derive(Debug, Serialize)]
pub(crate) struct CapabilitiesResult {
    pub(crate) protocol_versions: Vec<u16>,
    pub(crate) operations: Vec<String>,
    pub(crate) events: Vec<String>,
    pub(crate) transports: Vec<String>,
    pub(crate) limits: ProtocolLimits,
    pub(crate) activity: super::activity::ActivityCapabilities,
}

#[derive(Debug, Serialize)]
pub(crate) struct TurnStartResult {
    pub(crate) turn_id: String,
    pub(crate) session_id: String,
    pub(crate) status: String,
}

#[derive(Debug, Serialize)]
pub(crate) struct TurnCancelResult {
    pub(crate) turn_id: String,
    pub(crate) status: String,
}

#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum TurnTerminalStatus {
    Completed,
    Cancelled,
    Failed,
}

#[derive(Debug, Serialize, Clone, Copy)]
pub(crate) struct ProtocolLimits {
    pub(crate) max_record_bytes: usize,
    pub(crate) max_payload_bytes: usize,
    pub(crate) max_error_message_bytes: usize,
    pub(crate) max_id_bytes: usize,
    pub(crate) max_name_bytes: usize,
    pub(crate) max_string_bytes: usize,
    pub(crate) max_json_depth: usize,
    pub(crate) max_array_items: usize,
    pub(crate) max_object_fields: usize,
    pub(crate) max_negotiation_items: usize,
    pub(crate) max_session_attachments: usize,
    pub(crate) max_session_page_items: usize,
    pub(crate) max_session_replay_page_bytes: usize,
    pub(crate) max_session_replay_window_bytes: usize,
    pub(crate) max_session_list_scan_entries: usize,
    pub(crate) max_catalog_page_items: usize,
}

impl ProtocolLimits {
    pub(crate) const fn current() -> Self {
        Self {
            max_record_bytes: MAX_RECORD_BYTES,
            max_payload_bytes: MAX_PAYLOAD_BYTES,
            max_error_message_bytes: MAX_ERROR_MESSAGE_BYTES,
            max_id_bytes: MAX_ID_BYTES,
            max_name_bytes: MAX_NAME_BYTES,
            max_string_bytes: MAX_STRING_BYTES,
            max_json_depth: MAX_JSON_DEPTH,
            max_array_items: MAX_ARRAY_ITEMS,
            max_object_fields: MAX_OBJECT_FIELDS,
            max_negotiation_items: MAX_NEGOTIATION_ITEMS,
            max_session_attachments: super::sessions::MAX_ATTACHMENTS,
            max_session_page_items: crate::sessions::FRONTEND_PAGE_LIMIT,
            max_session_replay_page_bytes: crate::sessions::FRONTEND_REPLAY_PAGE_BYTES,
            max_session_replay_window_bytes: crate::sessions::FRONTEND_REPLAY_WINDOW_BYTES,
            max_session_list_scan_entries: crate::sessions::FRONTEND_LIST_SCAN_ENTRIES,
            max_catalog_page_items: super::configuration::MAX_CATALOG_PAGE_ITEMS,
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub(crate) struct ServiceEvent {
    pub(crate) protocol_version: u16,
    pub(crate) kind: MessageKind,
    pub(crate) event_id: String,
    pub(crate) request_id: String,
    pub(crate) session_id: Option<String>,
    pub(crate) event: String,
    pub(crate) payload: Value,
}

impl ServiceEvent {
    pub(crate) fn new(
        request_id: String,
        session_id: Option<String>,
        event: &'static str,
        payload: Value,
    ) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Event,
            event_id: Uuid::new_v4().to_string(),
            request_id,
            session_id,
            event: event.to_string(),
            payload,
        }
    }

    pub(crate) fn turn_started(request_id: String, session_id: String, turn_id: String) -> Self {
        Self::new(
            request_id,
            Some(session_id.clone()),
            TURN_STARTED_EVENT,
            json!({
                "turn_id": turn_id,
                "session_id": session_id,
                "status": "started",
                "sequence": 0,
            }),
        )
    }

    pub(crate) fn assistant_delta(
        request_id: String,
        session_id: String,
        turn_id: String,
        text: String,
    ) -> Self {
        Self::new(
            request_id,
            Some(session_id),
            ASSISTANT_DELTA_EVENT,
            json!({"turn_id": turn_id, "text": text}),
        )
    }

    pub(crate) fn turn_terminal(
        request_id: String,
        session_id: String,
        turn_id: String,
        status: TurnTerminalStatus,
        assistant_text: String,
    ) -> Self {
        Self::new(
            request_id,
            Some(session_id.clone()),
            TURN_TERMINAL_EVENT,
            json!({
                "turn_id": turn_id,
                "session_id": session_id,
                "status": status,
                "assistant_text": assistant_text,
                "sequence": 0,
                "activity_dropped": 0,
                "activity_replay_available": false,
                "replay_required": true,
            }),
        )
    }
}

#[derive(Debug, Serialize)]
#[serde(untagged)]
pub(crate) enum ServiceMessage {
    Response(Box<ServiceResponse>),
    Event(ServiceEvent),
}

impl ServiceMessage {
    pub(crate) fn response(response: ServiceResponse) -> Self {
        Self::Response(Box::new(response))
    }
}

#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ServiceErrorCode {
    InvalidJson,
    RecordTooLarge,
    InvalidRequest,
    PayloadTooLarge,
    LimitExceeded,
    UnsupportedVersion,
    UnsupportedCapability,
    UnsupportedOperation,
    NotInitialized,
    InvalidPayload,
    Conflict,
    ResourceNotFound,
    DuplicateRequestId,
    TooManyActiveTurns,
    SessionBusy,
    SessionUnavailable,
    InvalidSessionId,
    UnknownTurn,
    AuthBusy,
    UnknownLogin,
    ProviderUnavailable,
    ConfirmationRequired,
    ConfigurationUnavailable,
    ConfigurationBusy,
    CatalogRefreshFailed,
    SettingsWriteFailed,
    InvalidScope,
    InvalidSelection,
    UnsupportedReasoning,
    CatalogUnavailable,
    CatalogStale,
    InternalError,
    StaleInstance,
    StaleConnection,
    StaleGrant,
    SessionControlled,
    OperationIdentityMismatch,
    OperationAlreadyKnown,
    OperationExpired,
    OperationCapacity,
    AlreadyInitialized,
    RequestTimeout,
    PersistenceFailed,
    OutputLimitExceeded,
}

impl ServiceErrorCode {
    pub(crate) const fn message(self) -> &'static str {
        match self {
            Self::InvalidJson => "request is not valid JSON",
            Self::RecordTooLarge => "request record exceeds the maximum size",
            Self::InvalidRequest => "request envelope is invalid",
            Self::PayloadTooLarge => "request payload exceeds the maximum size",
            Self::LimitExceeded => "request exceeds a protocol limit",
            Self::UnsupportedVersion => "protocol version is not supported",
            Self::UnsupportedCapability => "requested capability is not supported",
            Self::UnsupportedOperation => "operation is not supported",
            Self::NotInitialized => "initialize must complete before this operation",
            Self::InvalidPayload => "operation payload is invalid",
            Self::Conflict => "resource revision conflicts with current state",
            Self::ResourceNotFound => "resource is not registered",
            Self::DuplicateRequestId => "request_id is already in flight",
            Self::TooManyActiveTurns => "the service has reached its active turn limit",
            Self::SessionBusy => "session has a conflicting writer or is closing",
            Self::SessionUnavailable => "session is unavailable",
            Self::InvalidSessionId => "session ID is invalid",
            Self::UnknownTurn => "turn is unknown or already terminal",
            Self::AuthBusy => "authentication operation is already running",
            Self::UnknownLogin => "login is unknown or already terminal",
            Self::ProviderUnavailable => "provider does not support this operation",
            Self::ConfirmationRequired => "logout requires explicit confirmation",
            Self::ConfigurationUnavailable => "configuration could not be read; check settings",
            Self::ConfigurationBusy => {
                "a configuration operation is running; retry after its response"
            }
            Self::CatalogRefreshFailed => {
                "catalog refresh failed; check provider readiness and retry"
            }
            Self::SettingsWriteFailed => {
                "settings could not be persisted; check settings and retry"
            }
            Self::InvalidScope => "scope is invalid or the setting requires global scope",
            Self::InvalidSelection => "provider or model is unknown, disabled, or invalid",
            Self::UnsupportedReasoning => "reasoning choice is unsupported for this model",
            Self::CatalogUnavailable => "model catalog is unavailable; refresh it before selecting",
            Self::CatalogStale => "model catalog is stale; refresh it before selecting",
            Self::InternalError => "application service failed to handle the request",
            Self::StaleInstance => "Daemon instance changed.",
            Self::StaleConnection => "Connection is no longer valid.",
            Self::StaleGrant => "Control is no longer valid.",
            Self::SessionControlled => "Session already has a controller.",
            Self::OperationIdentityMismatch => "Operation identity has different intent.",
            Self::OperationAlreadyKnown => "Look up the existing operation.",
            Self::OperationExpired => "Operation outcome expired.",
            Self::OperationCapacity => "Operation registry is full.",
            Self::AlreadyInitialized => "Connection is already initialized.",
            Self::RequestTimeout => "Request was not admitted before timeout.",
            Self::PersistenceFailed => "Terminal state could not be saved.",
            Self::OutputLimitExceeded => "Assistant output exceeded the negotiated bound.",
        }
    }
}

#[derive(Debug, Serialize)]
pub(crate) struct ServiceErrorDto {
    pub(crate) code: ServiceErrorCode,
    pub(crate) message: String,
}

#[derive(Debug, Default, Clone)]
pub(crate) struct RequestIdentity {
    pub(crate) request_id: Option<String>,
    pub(crate) session_id: Option<String>,
    pub(crate) method: Option<String>,
}

impl RequestIdentity {
    pub(crate) fn from_request(request: &ServiceRequest) -> Self {
        Self {
            request_id: bounded_text(&request.request_id, MAX_ID_BYTES),
            session_id: request
                .session_id
                .as_deref()
                .and_then(|value| bounded_text(value, MAX_ID_BYTES)),
            method: bounded_text(&request.method, MAX_NAME_BYTES),
        }
    }
}

#[derive(Debug)]
pub(crate) struct EncodeError;

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct InitializeParams {
    pub(crate) supported_protocol_versions: Vec<u16>,
    #[serde(default)]
    pub(crate) requested_capabilities: Vec<String>,
}

#[derive(Debug, Default, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct EmptyParams {}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct TurnStartParams {
    pub(crate) prompt: String,
    #[serde(default, rename = "options")]
    pub(crate) _options: EmptyParams,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct TurnCancelParams {
    pub(crate) turn_id: String,
}

pub(crate) fn validate_service_id(value: &str) -> Result<(), ServiceErrorCode> {
    validate_text(value, MAX_ID_BYTES, ServiceErrorCode::InvalidPayload)
}

fn validate_text(
    value: &str,
    max_bytes: usize,
    invalid_code: ServiceErrorCode,
) -> Result<(), ServiceErrorCode> {
    if value.len() > max_bytes {
        return Err(ServiceErrorCode::LimitExceeded);
    }
    if value.is_empty() || value.chars().any(char::is_control) {
        return Err(invalid_code);
    }
    Ok(())
}

pub(crate) fn validate_value_limits(value: &Value) -> Result<(), ServiceErrorCode> {
    validate_value_limits_at(value, 0)
}

fn validate_value_limits_at(value: &Value, depth: usize) -> Result<(), ServiceErrorCode> {
    if depth > MAX_JSON_DEPTH {
        return Err(ServiceErrorCode::LimitExceeded);
    }
    match value {
        Value::Array(values) => {
            if values.len() > MAX_ARRAY_ITEMS {
                return Err(ServiceErrorCode::LimitExceeded);
            }
            for value in values {
                validate_value_limits_at(value, depth + 1)?;
            }
        }
        Value::Object(fields) => {
            if fields.len() > MAX_OBJECT_FIELDS {
                return Err(ServiceErrorCode::LimitExceeded);
            }
            for (key, value) in fields {
                if key.len() > MAX_STRING_BYTES {
                    return Err(ServiceErrorCode::LimitExceeded);
                }
                validate_value_limits_at(value, depth + 1)?;
            }
        }
        Value::String(text) if text.len() > MAX_STRING_BYTES => {
            return Err(ServiceErrorCode::LimitExceeded);
        }
        _ => {}
    }
    Ok(())
}

pub(crate) fn payload_is_bounded(payload: &Value) -> Result<(), ServiceErrorCode> {
    validate_value_limits(payload)?;
    let size = serde_json::to_vec(payload)
        .map_err(|_| ServiceErrorCode::InvalidPayload)?
        .len();
    if size > MAX_PAYLOAD_BYTES {
        return Err(ServiceErrorCode::PayloadTooLarge);
    }
    Ok(())
}

pub(crate) fn encode_message(message: &ServiceMessage) -> Result<Vec<u8>, EncodeError> {
    validate_message(message)?;
    let mut bytes = serde_json::to_vec(message).map_err(|_| EncodeError)?;
    if bytes.len().saturating_add(1) > MAX_RECORD_BYTES {
        return Err(EncodeError);
    }
    bytes.push(b'\n');
    Ok(bytes)
}

pub(crate) fn terminal_text_fits(
    request_id: &str,
    session_id: &str,
    turn_id: &str,
    text: &str,
) -> bool {
    if text.len() > MAX_STRING_BYTES {
        return false;
    }
    let mut event = ServiceEvent::turn_terminal(
        request_id.to_string(),
        session_id.to_string(),
        turn_id.to_string(),
        TurnTerminalStatus::Completed,
        text.to_string(),
    );
    // Reserve the largest live sequence/counter encoding before accepting text.
    event.payload["sequence"] = json!(u64::MAX);
    event.payload["activity_dropped"] = json!(u64::MAX);
    encode_message(&ServiceMessage::Event(event)).is_ok()
}

fn validate_message(message: &ServiceMessage) -> Result<(), EncodeError> {
    match message {
        ServiceMessage::Response(response) => validate_response(response),
        ServiceMessage::Event(event) => {
            if event.protocol_version != PROTOCOL_VERSION
                || event.kind != MessageKind::Event
                || bounded_text(&event.event_id, MAX_ID_BYTES).is_none()
                || bounded_text(&event.request_id, MAX_ID_BYTES).is_none()
                || bounded_text(&event.event, MAX_NAME_BYTES).is_none()
                || event
                    .session_id
                    .as_deref()
                    .is_some_and(|value| bounded_text(value, MAX_ID_BYTES).is_none())
                || payload_is_bounded(&event.payload).is_err()
            {
                return Err(EncodeError);
            }
            Ok(())
        }
    }
}

fn validate_response(response: &ServiceResponse) -> Result<(), EncodeError> {
    if response.protocol_version != PROTOCOL_VERSION
        || response.kind != MessageKind::Response
        || response.payload.is_none() == response.error.is_none()
        || response
            .request_id
            .as_ref()
            .is_some_and(|value| bounded_text(value, MAX_ID_BYTES).is_none())
        || response
            .session_id
            .as_ref()
            .is_some_and(|value| bounded_text(value, MAX_ID_BYTES).is_none())
        || response
            .method
            .as_ref()
            .is_some_and(|value| bounded_text(value, MAX_NAME_BYTES).is_none())
    {
        return Err(EncodeError);
    }
    if let Some(error) = response.error.as_ref()
        && bounded_text(&error.message, MAX_ERROR_MESSAGE_BYTES).is_none()
    {
        return Err(EncodeError);
    }
    if let Some(payload) = response.payload.as_ref() {
        let value = serde_json::to_value(payload).map_err(|_| EncodeError)?;
        if payload_is_bounded(&value).is_err() {
            return Err(EncodeError);
        }
    }
    Ok(())
}

impl ServiceResponse {
    pub(crate) fn success(request: &ServiceRequest, payload: ResponsePayload) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Response,
            request_id: Some(request.request_id.clone()),
            session_id: request.session_id.clone(),
            method: Some(request.method.clone()),
            payload: Some(payload),
            error: None,
        }
    }

    pub(crate) fn success_with_identity(
        request_id: String,
        session_id: Option<String>,
        method: Option<String>,
        payload: ResponsePayload,
    ) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Response,
            request_id: Some(request_id),
            session_id,
            method,
            payload: Some(payload),
            error: None,
        }
    }

    pub(crate) fn error(identity: RequestIdentity, code: ServiceErrorCode) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Response,
            request_id: identity.request_id,
            session_id: identity.session_id,
            method: identity.method,
            payload: None,
            error: Some(ServiceErrorDto {
                code,
                message: code.message().to_string(),
            }),
        }
    }
}

fn bounded_text(value: &str, max_bytes: usize) -> Option<String> {
    (!value.is_empty() && value.len() <= max_bytes && !value.chars().any(char::is_control))
        .then(|| value.to_string())
}

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

    #[test]
    fn encode_event_preserves_explicit_correlation_and_jsonl_shape() {
        let event = ServiceEvent {
            protocol_version: PROTOCOL_VERSION,
            kind: MessageKind::Event,
            event_id: "event-1".to_string(),
            request_id: "request-1".to_string(),
            session_id: Some("session-1".to_string()),
            event: "activity".to_string(),
            payload: json!({"state": "running"}),
        };

        let bytes = encode_message(&ServiceMessage::Event(event)).unwrap();
        let value: Value = serde_json::from_slice(&bytes).unwrap();

        assert_eq!(value["kind"], "event");
        assert_eq!(value["event_id"], "event-1");
        assert_eq!(value["request_id"], "request-1");
        assert_eq!(value["session_id"], "session-1");
        assert!(bytes.ends_with(b"\n"));
    }

    #[test]
    fn encode_message_keeps_error_envelopes_bounded_and_explicit() {
        let response =
            ServiceResponse::error(RequestIdentity::default(), ServiceErrorCode::InvalidJson);
        let bytes = encode_message(&ServiceMessage::response(response)).unwrap();
        let value: Value = serde_json::from_slice(&bytes).unwrap();

        assert_eq!(value["kind"], "response");
        assert_eq!(value["error"]["code"], "invalid_json");
        assert_eq!(value["error"]["message"], "request is not valid JSON");
        assert!(bytes.len() <= MAX_RECORD_BYTES);
    }
}