Skip to main content

mcpr_core/proxy/pipeline/
context.rs

1//! Per-request contexts threaded through the pipeline.
2//!
3//! [`RequestContext`] is a one-pass parse of the incoming HTTP request that
4//! every downstream stage reads from instead of re-parsing. Handlers update
5//! a small set of mutable fields (session id after upstream assigns one,
6//! client name/version after session lookup) before handing off to
7//! [`super::emit::emit_request_event`].
8//!
9//! [`ResponseContext`] carries the accumulating response through the
10//! response middleware chain: the raw body, an optional parsed JSON view,
11//! SSE-wrapping state, and JSON-RPC error info. Middleware mutate it; the
12//! final handler builds the `axum::Response` from `resp.body` + `resp.headers`.
13
14use std::time::Instant;
15
16use crate::protocol::session::ClientInfo;
17use crate::protocol::{McpMethod, ParsedBody};
18use axum::http::{HeaderMap, Method};
19use serde_json::Value;
20
21pub struct RequestContext {
22    pub start: Instant,
23
24    // ── HTTP ──
25    pub http_method: Method,
26    pub path: String,
27    pub request_size: usize,
28    pub wants_sse: bool,
29
30    // ── Session (set from header; overwritten when upstream assigns one) ──
31    pub session_id: Option<String>,
32
33    // ── JSON-RPC / MCP (None when the body is not JSON-RPC) ──
34    pub jsonrpc: Option<ParsedBody>,
35    pub mcp_method: Option<McpMethod>,
36    /// String form for event output. Set to the protocol method for MCP POSTs
37    /// and overwritten by specific handlers where appropriate (e.g. "SSE").
38    pub mcp_method_str: Option<String>,
39    /// `tools/call` tool name; `None` for other methods.
40    pub tool: Option<String>,
41    pub is_batch: bool,
42
43    // ── Client info ──
44    /// Parsed from `initialize` params. The Initialize success path stores
45    /// this into the session store.
46    pub client_info_from_init: Option<ClientInfo>,
47    /// Resolved from the session store by the handler before emit.
48    pub client_name: Option<String>,
49    pub client_version: Option<String>,
50
51    /// Transform tags pushed by handlers / middleware. The emit stage joins
52    /// them with `+` to build the `RequestEvent.note` field
53    /// (e.g. `["rewritten", "sse"]` → `"rewritten+sse"`).
54    pub tags: Vec<&'static str>,
55}
56
57/// Response-side state threaded through the response middleware chain.
58///
59/// Middleware mutate `body` and `json` in place. Handlers instantiate after
60/// reading the upstream body and finalize by building an `axum::Response` from
61/// `(status, headers, body)`.
62pub struct ResponseContext {
63    pub status: u16,
64    pub headers: HeaderMap,
65    /// Serialized response body — what gets returned to the client. Held
66    /// verbatim from the upstream until `EncodeResponseJson` overwrites it.
67    /// When no middleware mutates `json`, this retains the original bytes
68    /// byte-for-byte (preserving SSE framing, key order, etc.).
69    pub body: Vec<u8>,
70    /// True when the upstream sent SSE-wrapped JSON. `DecodeResponseJson`
71    /// sets it; `EncodeResponseJson` reads it to decide whether to re-wrap.
72    pub was_sse: bool,
73    /// Parsed JSON view of the body. Populated by `DecodeResponseJson`;
74    /// mutated by later middleware; serialized back into `body` by
75    /// `EncodeResponseJson` only when `json_mutated` is set.
76    pub json: Option<Value>,
77    /// Signals that some middleware mutated `json`. Set by any middleware
78    /// that takes `json.as_mut()` or reassigns `json`. `EncodeResponseJson`
79    /// skips re-serialization when false, leaving `body` untouched — the
80    /// byte-pass fast path.
81    pub json_mutated: bool,
82    /// JSON-RPC error extracted from `json` (when present).
83    pub rpc_error: Option<(i64, String)>,
84    pub upstream_us: Option<u64>,
85}
86
87impl ResponseContext {
88    pub fn new(status: u16, headers: HeaderMap, body: Vec<u8>, upstream_us: Option<u64>) -> Self {
89        Self {
90            status,
91            headers,
92            body,
93            was_sse: false,
94            json: None,
95            json_mutated: false,
96            rpc_error: None,
97            upstream_us,
98        }
99    }
100
101    /// Mutable access to the parsed JSON value, marking it as mutated so
102    /// `EncodeResponseJson` will re-serialize. Prefer this over direct
103    /// `json.as_mut()` — forgetting to set the flag causes silent staleness.
104    pub fn json_mut(&mut self) -> Option<&mut Value> {
105        if self.json.is_some() {
106            self.json_mutated = true;
107        }
108        self.json.as_mut()
109    }
110}