imagegen_bridge_codex_app_server/
rpc.rs

1//! Bounded JSONL request/response correlation for Codex app-server.
2
3use std::{
4    collections::HashMap,
5    sync::{
6        Arc,
7        atomic::{AtomicU64, Ordering},
8    },
9    time::{Duration, Instant},
10};
11
12use futures_util::StreamExt as _;
13use imagegen_bridge_core::{BridgeError, ErrorCode};
14use parking_lot::Mutex as ParkingMutex;
15use serde_json::{Value, json};
16use tokio::{
17    io::{AsyncRead, AsyncWrite, AsyncWriteExt as _},
18    sync::{Mutex, broadcast, oneshot, watch},
19};
20use tokio_util::{
21    codec::{FramedRead, LinesCodec},
22    sync::CancellationToken,
23};
24
25type PendingSender = oneshot::Sender<Result<Value, BridgeError>>;
26const MAX_RPC_MESSAGE_BYTES: usize = 64 * 1024 * 1024;
27const MAX_NOTIFICATION_CAPACITY: usize = 64;
28const MAX_NOTIFICATION_BYTES: usize = 48 * 1024 * 1024;
29const MAX_NOTIFICATION_RING_BYTES: usize = 256 * 1024 * 1024;
30
31/// Bounded connection settings.
32#[derive(Debug, Clone, Copy)]
33pub struct RpcConfig {
34    /// Maximum incoming or outgoing JSONL message bytes.
35    pub max_message_bytes: usize,
36    /// Maximum bytes accepted for one notification line.
37    pub max_notification_bytes: usize,
38    /// Default request timeout.
39    pub request_timeout: Duration,
40    /// Notification ring capacity.
41    pub notification_capacity: usize,
42}
43
44impl Default for RpcConfig {
45    fn default() -> Self {
46        Self {
47            max_message_bytes: 64 * 1024 * 1024,
48            max_notification_bytes: 48 * 1024 * 1024,
49            request_timeout: Duration::from_secs(60),
50            notification_capacity: 4,
51        }
52    }
53}
54
55/// One app-server notification with its method and parsed params.
56#[derive(Debug, Clone)]
57pub struct RpcNotification {
58    /// Notification method.
59    pub method: String,
60    /// Parsed params, or null when absent.
61    pub params: Value,
62}
63
64/// Initialized app-server connection.
65pub struct AppServerRpc {
66    writer: Mutex<Box<dyn AsyncWrite + Send + Unpin>>,
67    pending: Arc<ParkingMutex<HashMap<u64, PendingSender>>>,
68    notifications: broadcast::Sender<RpcNotification>,
69    closed: watch::Receiver<Option<BridgeError>>,
70    closed_sender: watch::Sender<Option<BridgeError>>,
71    next_id: AtomicU64,
72    config: RpcConfig,
73}
74
75impl std::fmt::Debug for AppServerRpc {
76    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        formatter
78            .debug_struct("AppServerRpc")
79            .field("pending_count", &self.pending.lock().len())
80            .field("config", &self.config)
81            .finish_non_exhaustive()
82    }
83}
84
85impl AppServerRpc {
86    /// Connects to app-server and completes initialize/initialized exactly once.
87    pub async fn connect<R, W>(
88        reader: R,
89        writer: W,
90        config: RpcConfig,
91    ) -> Result<Arc<Self>, BridgeError>
92    where
93        R: AsyncRead + Send + Unpin + 'static,
94        W: AsyncWrite + Send + Unpin + 'static,
95    {
96        let notification_slots = config
97            .notification_capacity
98            .checked_next_power_of_two()
99            .ok_or_else(|| protocol_error("notification capacity is too large"))?;
100        let notification_budget = config
101            .max_notification_bytes
102            .checked_mul(notification_slots)
103            .ok_or_else(|| protocol_error("notification ring budget overflowed"))?;
104        if config.max_message_bytes == 0
105            || config.max_message_bytes > MAX_RPC_MESSAGE_BYTES
106            || config.max_notification_bytes == 0
107            || config.max_notification_bytes > config.max_message_bytes
108            || config.max_notification_bytes > MAX_NOTIFICATION_BYTES
109            || config.notification_capacity == 0
110            || config.notification_capacity > MAX_NOTIFICATION_CAPACITY
111            || notification_budget > MAX_NOTIFICATION_RING_BYTES
112        {
113            return Err(protocol_error("RPC limits must be greater than zero"));
114        }
115        let (notifications, _) = broadcast::channel(config.notification_capacity);
116        let (closed_sender, closed) = watch::channel(None);
117        let pending = Arc::new(ParkingMutex::new(HashMap::new()));
118        tokio::spawn(read_loop(
119            reader,
120            config.max_message_bytes,
121            config.max_notification_bytes,
122            Arc::clone(&pending),
123            notifications.clone(),
124            closed_sender.clone(),
125        ));
126        let rpc = Arc::new(Self {
127            writer: Mutex::new(Box::new(writer)),
128            pending,
129            notifications,
130            closed,
131            closed_sender,
132            next_id: AtomicU64::new(1),
133            config,
134        });
135        rpc.request(
136            "initialize",
137            json!({
138                "clientInfo": {
139                    "name": "imagegen-bridge",
140                    "title": "Imagegen Bridge",
141                    "version": env!("CARGO_PKG_VERSION")
142                },
143                "capabilities": {
144                    "experimentalApi": false,
145                    "requestAttestation": false,
146                    "mcpServerOpenaiFormElicitation": false
147                }
148            }),
149        )
150        .await?;
151        rpc.notify("initialized", None).await?;
152        Ok(rpc)
153    }
154
155    /// Subscribes before starting a turn so no notification is missed.
156    #[must_use]
157    pub fn subscribe(&self) -> broadcast::Receiver<RpcNotification> {
158        self.notifications.subscribe()
159    }
160
161    /// Returns whether the read or write side has failed permanently.
162    #[must_use]
163    pub fn is_closed(&self) -> bool {
164        self.closed.borrow().is_some()
165    }
166
167    /// Sends a request using the default timeout.
168    pub async fn request(&self, method: &str, params: Value) -> Result<Value, BridgeError> {
169        self.request_until(
170            method,
171            params,
172            Instant::now() + self.config.request_timeout,
173            CancellationToken::new(),
174        )
175        .await
176    }
177
178    /// Sends a request bounded by an absolute deadline and cancellation token.
179    pub async fn request_until(
180        &self,
181        method: &str,
182        params: Value,
183        deadline: Instant,
184        cancellation: CancellationToken,
185    ) -> Result<Value, BridgeError> {
186        self.request_with_delivery(method, params, deadline, cancellation, false)
187            .await
188    }
189
190    /// Sends a side-effecting request and marks post-flush ambiguity explicitly.
191    pub async fn request_side_effecting_until(
192        &self,
193        method: &str,
194        params: Value,
195        deadline: Instant,
196        cancellation: CancellationToken,
197    ) -> Result<Value, BridgeError> {
198        self.request_with_delivery(method, params, deadline, cancellation, true)
199            .await
200    }
201
202    async fn request_with_delivery(
203        &self,
204        method: &str,
205        params: Value,
206        deadline: Instant,
207        cancellation: CancellationToken,
208        side_effecting: bool,
209    ) -> Result<Value, BridgeError> {
210        if let Some(error) = self.closed.borrow().clone() {
211            return Err(error);
212        }
213        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
214        let message = json!({"method": method, "id": id, "params": params});
215        let rendered = render_message(&message, self.config.max_message_bytes)?;
216        let (sender, receiver) = oneshot::channel();
217        self.pending.lock().insert(id, sender);
218        if let Err(error) = self.write(&rendered).await {
219            self.pending.lock().remove(&id);
220            return Err(error);
221        }
222
223        let timeout = deadline.saturating_duration_since(Instant::now());
224        tokio::select! {
225            result = receiver => result
226                .map_err(|_| protocol_error("app-server response channel closed"))?,
227            () = cancellation.cancelled() => {
228                self.pending.lock().remove(&id);
229                Err(request_interrupted(
230                    ErrorCode::Cancelled,
231                    "app-server request was cancelled",
232                    side_effecting,
233                    id,
234                ))
235            }
236            () = tokio::time::sleep(timeout) => {
237                self.pending.lock().remove(&id);
238                Err(request_interrupted(
239                    ErrorCode::Timeout,
240                    "app-server request timed out",
241                    side_effecting,
242                    id,
243                ))
244            }
245        }
246    }
247
248    /// Sends a JSONL notification.
249    pub async fn notify(&self, method: &str, params: Option<Value>) -> Result<(), BridgeError> {
250        let mut message = json!({"method": method});
251        if let Some(params) = params {
252            message["params"] = params;
253        }
254        let rendered = render_message(&message, self.config.max_message_bytes)?;
255        self.write(&rendered).await
256    }
257
258    async fn write(&self, rendered: &[u8]) -> Result<(), BridgeError> {
259        let mut writer = self.writer.lock().await;
260        if let Err(_error) = writer.write_all(rendered).await {
261            let error = protocol_error("could not write to app-server");
262            fail_connection(error.clone(), &self.pending, &self.closed_sender);
263            return Err(error);
264        }
265        if let Err(_error) = writer.flush().await {
266            let error = protocol_error("could not flush app-server input");
267            fail_connection(error.clone(), &self.pending, &self.closed_sender);
268            return Err(error);
269        }
270        Ok(())
271    }
272}
273
274fn render_message(message: &Value, maximum: usize) -> Result<Vec<u8>, BridgeError> {
275    let mut rendered = serde_json::to_vec(message)
276        .map_err(|_| protocol_error("could not serialize app-server message"))?;
277    if rendered.len() > maximum {
278        return Err(protocol_error("outgoing app-server message exceeds limit"));
279    }
280    rendered.push(b'\n');
281    Ok(rendered)
282}
283
284async fn read_loop<R>(
285    reader: R,
286    maximum: usize,
287    maximum_notification: usize,
288    pending: Arc<ParkingMutex<HashMap<u64, PendingSender>>>,
289    notifications: broadcast::Sender<RpcNotification>,
290    closed: watch::Sender<Option<BridgeError>>,
291) where
292    R: AsyncRead + Send + Unpin + 'static,
293{
294    let mut lines = FramedRead::new(reader, LinesCodec::new_with_max_length(maximum));
295    while let Some(line) = lines.next().await {
296        let result = match line {
297            Ok(line) => dispatch_message(&line, maximum_notification, &pending, &notifications),
298            Err(_) => Err(protocol_error(
299                "app-server message exceeds limit or is not valid UTF-8",
300            )),
301        };
302        if let Err(error) = result {
303            fail_connection(error, &pending, &closed);
304            return;
305        }
306    }
307    fail_connection(
308        protocol_error("app-server connection closed"),
309        &pending,
310        &closed,
311    );
312}
313
314fn dispatch_message(
315    line: &str,
316    maximum_notification: usize,
317    pending: &ParkingMutex<HashMap<u64, PendingSender>>,
318    notifications: &broadcast::Sender<RpcNotification>,
319) -> Result<(), BridgeError> {
320    let message: Value =
321        serde_json::from_str(line).map_err(|_| protocol_error("app-server sent invalid JSON"))?;
322    let object = message
323        .as_object()
324        .ok_or_else(|| protocol_error("app-server message is not an object"))?;
325    if let Some(id) = object.get("id").and_then(Value::as_u64) {
326        if let Some(sender) = pending.lock().remove(&id) {
327            let result = if let Some(error) = object.get("error") {
328                Err(rpc_response_error(error))
329            } else {
330                Ok(object.get("result").cloned().unwrap_or(Value::Null))
331            };
332            let _ = sender.send(result);
333        }
334        return Ok(());
335    }
336    if line.len() > maximum_notification {
337        return Err(protocol_error("app-server notification exceeds limit"));
338    }
339    let method = object
340        .get("method")
341        .and_then(Value::as_str)
342        .ok_or_else(|| protocol_error("app-server notification has no method"))?;
343    if !is_forwarded_notification(method) {
344        return Ok(());
345    }
346    let notification = RpcNotification {
347        method: method.to_owned(),
348        params: object.get("params").cloned().unwrap_or(Value::Null),
349    };
350    let _ = notifications.send(notification);
351    Ok(())
352}
353
354fn is_forwarded_notification(method: &str) -> bool {
355    matches!(method, "item/completed" | "turn/completed")
356}
357
358fn rpc_response_error(value: &Value) -> BridgeError {
359    let safe_string_code = value.get("code").and_then(Value::as_str).filter(|code| {
360        !code.is_empty()
361            && code.len() <= 64
362            && code
363                .bytes()
364                .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
365    });
366    let mut error = if safe_string_code.is_some_and(is_safety_code) {
367        BridgeError::safety_rejected("Codex app-server rejected the image request")
368    } else {
369        BridgeError::new(ErrorCode::Upstream, "Codex app-server request failed")
370    }
371    .with_provider("codex-app-server");
372    if let Some(code) = value.get("code").and_then(Value::as_i64) {
373        error = error.with_detail("rpc_code", code);
374    } else if let Some(code) = safe_string_code {
375        error = error.with_detail("rpc_code", code);
376    }
377    error
378}
379
380fn is_safety_code(code: &str) -> bool {
381    let lower = code.to_ascii_lowercase();
382    lower.contains("safety")
383        || lower.contains("content_policy")
384        || lower.contains("moderation")
385        || lower.contains("refusal")
386}
387
388fn fail_connection(
389    error: BridgeError,
390    pending: &ParkingMutex<HashMap<u64, PendingSender>>,
391    closed: &watch::Sender<Option<BridgeError>>,
392) {
393    for (_, sender) in pending.lock().drain() {
394        let _ = sender.send(Err(error.clone()));
395    }
396    closed.send_replace(Some(error));
397}
398
399fn protocol_error(message: &str) -> BridgeError {
400    BridgeError::new(ErrorCode::Protocol, message).with_provider("codex-app-server")
401}
402
403fn request_interrupted(
404    code: ErrorCode,
405    message: &str,
406    side_effecting: bool,
407    request_id: u64,
408) -> BridgeError {
409    let error = BridgeError::new(code, message).with_provider("codex-app-server");
410    if side_effecting {
411        error
412            .retryable(false)
413            .with_detail("outcome", "unknown")
414            .with_detail("request_sent", true)
415            .with_detail("rpc_id", request_id)
416    } else {
417        error.retryable(code == ErrorCode::Timeout)
418    }
419}
420
421#[cfg(test)]
422mod tests {
423    #![allow(clippy::unwrap_used)]
424
425    use futures_util::{SinkExt as _, StreamExt as _};
426    use tokio::io::duplex;
427    use tokio_util::codec::{Framed, LinesCodec};
428
429    use super::*;
430
431    async fn initialized_connection() -> (
432        Arc<AppServerRpc>,
433        Framed<tokio::io::DuplexStream, LinesCodec>,
434    ) {
435        let (client, server) = duplex(64 * 1024);
436        let (client_read, client_write) = tokio::io::split(client);
437        let mut server = Framed::new(server, LinesCodec::new());
438        let connect = tokio::spawn(AppServerRpc::connect(
439            client_read,
440            client_write,
441            RpcConfig {
442                max_message_bytes: 4096,
443                max_notification_bytes: 4096,
444                request_timeout: Duration::from_secs(1),
445                notification_capacity: 8,
446            },
447        ));
448        let initialize: Value =
449            serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
450        assert_eq!(initialize["method"], "initialize");
451        server
452            .send(json!({"id": initialize["id"], "result": {"userAgent": "codex/0.0"}}).to_string())
453            .await
454            .unwrap();
455        let rpc = connect.await.unwrap().unwrap();
456        let initialized: Value =
457            serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
458        assert_eq!(initialized["method"], "initialized");
459        (rpc, server)
460    }
461
462    #[tokio::test]
463    async fn correlates_concurrent_responses_by_id() {
464        let (rpc, mut server) = initialized_connection().await;
465        let first = tokio::spawn({
466            let rpc = Arc::clone(&rpc);
467            async move { rpc.request("one", json!({})).await.unwrap() }
468        });
469        let second = tokio::spawn({
470            let rpc = Arc::clone(&rpc);
471            async move { rpc.request("two", json!({})).await.unwrap() }
472        });
473        let request_a: Value =
474            serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
475        let request_b: Value =
476            serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
477        server
478            .send(json!({"id": request_b["id"], "result": request_b["method"]}).to_string())
479            .await
480            .unwrap();
481        server
482            .send(json!({"id": request_a["id"], "result": request_a["method"]}).to_string())
483            .await
484            .unwrap();
485        let values = [first.await.unwrap(), second.await.unwrap()];
486        assert!(values.contains(&json!("one")));
487        assert!(values.contains(&json!("two")));
488    }
489
490    #[tokio::test]
491    async fn side_effecting_timeout_records_post_flush_unknown_outcome() {
492        let (rpc, mut server) = initialized_connection().await;
493        let request = tokio::spawn({
494            let rpc = Arc::clone(&rpc);
495            async move {
496                rpc.request_side_effecting_until(
497                    "turn/start",
498                    json!({"threadId": "thread-1"}),
499                    Instant::now() + Duration::from_millis(20),
500                    CancellationToken::new(),
501                )
502                .await
503            }
504        });
505        let sent: Value = serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
506        assert_eq!(sent["method"], "turn/start");
507        let error = request.await.unwrap().unwrap_err();
508        assert_eq!(error.code, ErrorCode::Timeout);
509        assert!(!error.retryable);
510        assert_eq!(error.details["outcome"], "unknown");
511        assert_eq!(error.details["request_sent"], true);
512        assert_eq!(error.details["rpc_id"], sent["id"]);
513    }
514
515    #[tokio::test]
516    async fn rejects_notification_ring_above_the_absolute_byte_budget() {
517        let (client, _server) = duplex(64);
518        let (reader, writer) = tokio::io::split(client);
519        let error = AppServerRpc::connect(
520            reader,
521            writer,
522            RpcConfig {
523                max_message_bytes: 64 * 1024 * 1024,
524                max_notification_bytes: 48 * 1024 * 1024,
525                request_timeout: Duration::from_secs(1),
526                notification_capacity: 8,
527            },
528        )
529        .await
530        .unwrap_err();
531        assert_eq!(error.code, ErrorCode::Protocol);
532        assert!(error.message.contains("RPC limits"));
533    }
534
535    #[tokio::test]
536    async fn forwards_parsed_notifications_without_raw_logging() {
537        let (rpc, mut server) = initialized_connection().await;
538        let mut notifications = rpc.subscribe();
539        server
540            .send(json!({"method": "turn/completed", "params": {"threadId": "t"}}).to_string())
541            .await
542            .unwrap();
543        let notification = notifications.recv().await.unwrap();
544        assert_eq!(notification.method, "turn/completed");
545        assert_eq!(notification.params["threadId"], "t");
546    }
547
548    #[tokio::test]
549    async fn ignores_high_volume_notifications_that_no_consumer_uses() {
550        let (rpc, mut server) = initialized_connection().await;
551        let mut notifications = rpc.subscribe();
552        for index in 0..128 {
553            server
554                .send(
555                    json!({"method": "item/agentMessage/delta", "params": {"index": index}})
556                        .to_string(),
557                )
558                .await
559                .unwrap();
560        }
561        server
562            .send(json!({"method": "turn/completed", "params": {"threadId": "t"}}).to_string())
563            .await
564            .unwrap();
565
566        let notification = notifications.recv().await.unwrap();
567        assert_eq!(notification.method, "turn/completed");
568        assert_eq!(notification.params["threadId"], "t");
569    }
570
571    #[tokio::test]
572    async fn rejects_outgoing_message_over_limit() {
573        let (rpc, _server) = initialized_connection().await;
574        let error = rpc
575            .request("large", json!({"data": "x".repeat(5000)}))
576            .await
577            .unwrap_err();
578        assert_eq!(error.code, ErrorCode::Protocol);
579    }
580
581    #[tokio::test]
582    async fn connection_close_fails_pending_callers_and_marks_rpc_closed() {
583        let (rpc, mut server) = initialized_connection().await;
584        let pending = tokio::spawn({
585            let rpc = Arc::clone(&rpc);
586            async move { rpc.request("pending", json!({})).await }
587        });
588        let request: Value = serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
589        assert_eq!(request["method"], "pending");
590        drop(server);
591        let error = pending.await.unwrap().unwrap_err();
592        assert_eq!(error.code, ErrorCode::Protocol);
593        assert!(rpc.is_closed());
594    }
595
596    #[tokio::test]
597    async fn upstream_error_messages_are_not_reflected_to_clients() {
598        let (rpc, mut server) = initialized_connection().await;
599        let request = tokio::spawn({
600            let rpc = Arc::clone(&rpc);
601            async move { rpc.request("fails", json!({})).await }
602        });
603        let incoming: Value = serde_json::from_str(&server.next().await.unwrap().unwrap()).unwrap();
604        server
605            .send(
606                json!({
607                    "id": incoming["id"],
608                    "error": {"code": "invalid_request", "message": "secret prompt and /private/path"}
609                })
610                .to_string(),
611            )
612            .await
613            .unwrap();
614        let error = request.await.unwrap().unwrap_err();
615        assert!(!error.message.contains("secret"));
616        assert!(!error.message.contains("/private"));
617        assert_eq!(error.details["rpc_code"], "invalid_request");
618    }
619
620    #[test]
621    fn safety_rpc_codes_return_structured_recovery_guidance() {
622        let error = rpc_response_error(&json!({
623            "code": "content_policy_violation",
624            "message": "untrusted upstream detail"
625        }));
626        assert_eq!(error.code, ErrorCode::SafetyRejected);
627        assert_eq!(error.details["recovery"], "revise_prompt_or_inputs");
628        assert_eq!(error.details["rpc_code"], "content_policy_violation");
629        assert!(!error.message.contains("untrusted"));
630    }
631}
imagegen_bridge_codex_app_server/rpc.rs

imagegen_bridge_codex_app_server/
rpc.rs
