Skip to main content

kimi_wire/
message.rs

1//! Message parsing and dispatch for the Wire protocol.
2//!
3//! Converts raw [`crate::protocol::RawWireMessage`]s into typed dispatch enums
4//! ([`WireMessage`](crate::message::WireMessage)) so callers don't have to manually inspect
5//! `method`, `id`, `result`, and `error` fields.
6
7use serde_json::Value;
8
9use crate::error::WireError;
10use crate::protocol::{
11    Event, JsonRpcErrorResponse, JsonRpcNotification, JsonRpcRequest, JsonRpcSuccessResponse,
12    RawWireMessage, Request,
13};
14
15/// A union type for all incoming wire messages.
16#[derive(Debug, Clone)]
17pub enum WireMessage {
18    /// An incoming request from the agent (e.g. `ApprovalRequest`).
19    Request(JsonRpcRequest<Request>),
20    /// An incoming event from the agent (e.g. `TurnBegin`, `ContentPart`).
21    Event(JsonRpcNotification<Event>),
22    /// A successful JSON-RPC response to a previous request.
23    SuccessResponse(JsonRpcSuccessResponse<Value>),
24    /// A JSON-RPC error response.
25    ErrorResponse(JsonRpcErrorResponse),
26}
27
28/// Parse a raw wire message into a typed [`WireMessage`].
29///
30/// # Errors
31///
32/// Returns [`WireError::JsonParse`] if the payload cannot be deserialized,
33/// or [`WireError::UnknownMessageType`] if the `method` field is unrecognized.
34pub fn parse_wire_message(raw: RawWireMessage) -> Result<WireMessage, WireError> {
35    if let Some(method) = raw.method {
36        match method.as_str() {
37            "request" => {
38                let params: Value = raw.params.ok_or(WireError::InvalidPayloadType)?;
39                let req: JsonRpcRequest<Request> = JsonRpcRequest {
40                    jsonrpc: raw.jsonrpc,
41                    method,
42                    id: raw.id.unwrap_or_default(),
43                    params: serde_json::from_value(params).map_err(|e| {
44                        WireError::JsonParse(format!("failed to parse request params: {e}"))
45                    })?,
46                };
47                Ok(WireMessage::Request(req))
48            }
49            "event" => {
50                let params: Value = raw.params.ok_or(WireError::InvalidPayloadType)?;
51                let ev: Event = serde_json::from_value(params).map_err(|e| {
52                    WireError::JsonParse(format!("failed to parse event params: {e}"))
53                })?;
54                let notification = JsonRpcNotification {
55                    jsonrpc: raw.jsonrpc,
56                    method,
57                    params: ev,
58                };
59                Ok(WireMessage::Event(notification))
60            }
61            other => Err(WireError::UnknownMessageType(other.to_string())),
62        }
63    } else if raw.error.is_some() {
64        let id = raw.id.unwrap_or_default();
65        let error = raw
66            .error
67            .ok_or_else(|| WireError::Internal("error response missing error field".to_string()))?;
68        Ok(WireMessage::ErrorResponse(JsonRpcErrorResponse {
69            jsonrpc: raw.jsonrpc,
70            id,
71            error,
72        }))
73    } else if raw.result.is_some() {
74        let id = raw.id.unwrap_or_default();
75        let result = raw.result.ok_or_else(|| {
76            WireError::Internal("success response missing result field".to_string())
77        })?;
78        Ok(WireMessage::SuccessResponse(JsonRpcSuccessResponse {
79            jsonrpc: raw.jsonrpc,
80            id,
81            result,
82        }))
83    } else {
84        Err(WireError::UnknownMessageType(
85            "unrecognized wire message shape".to_string(),
86        ))
87    }
88}