1use std::io::{Read, Write};
18use std::path::PathBuf;
19use std::time::{Duration, Instant};
20
21use serde_json::{json, Value};
22use tungstenite::{Message, WebSocket};
23
24use super::{
25 wake_prompt, ProviderKind, ProviderSession, SessionKind, WakeAdapter, WakeError, WakeLetter,
26 WakeOutcome,
27};
28
29pub const CODEX_APP_SERVER_ENV: &str = "M4A_CODEX_APP_SERVER";
31
32pub mod desktop;
33
34const RPC_TIMEOUT: Duration = Duration::from_secs(20);
35
36#[derive(Debug, Clone, PartialEq, Eq)]
38pub enum CodexEndpoint {
39 Desktop(PathBuf),
41 Unix(PathBuf),
43 Ws(String),
45}
46
47impl CodexEndpoint {
48 pub fn parse(value: &str) -> Option<Self> {
51 let value = value.trim();
52 if value == "desktop" {
53 return desktop::default_endpoint().map(CodexEndpoint::Desktop);
54 }
55 if let Some(path) = value.strip_prefix("desktop://") {
56 return (!path.is_empty()).then(|| CodexEndpoint::Desktop(path.into()));
57 }
58 if let Some(path) = value.strip_prefix("unix://") {
59 return (!path.is_empty()).then(|| CodexEndpoint::Unix(PathBuf::from(path)));
60 }
61 let rest = value.strip_prefix("ws://")?;
62 let host = rest.split(['/', '?']).next().unwrap_or("");
63 let host = host.rsplit_once(':').map(|(h, _)| h).unwrap_or(host);
64 matches!(host, "127.0.0.1" | "localhost" | "[::1]")
65 .then(|| CodexEndpoint::Ws(value.to_string()))
66 }
67
68 pub fn from_env() -> Option<Self> {
70 std::env::var(CODEX_APP_SERVER_ENV)
71 .ok()
72 .and_then(|value| Self::parse(&value))
73 }
74}
75
76pub struct CodexAppServerAdapter {
78 endpoint: Option<CodexEndpoint>,
79}
80
81impl CodexAppServerAdapter {
82 pub fn new(endpoint: Option<CodexEndpoint>) -> Self {
84 Self { endpoint }
85 }
86}
87
88impl WakeAdapter for CodexAppServerAdapter {
89 fn kind(&self) -> SessionKind {
90 SessionKind::local(ProviderKind::Codex)
91 }
92
93 fn probe(&self, session: &ProviderSession) -> Result<(), WakeError> {
94 match &self.endpoint {
95 None => Err(WakeError::Unavailable(
96 "codex app-server endpoint unset".into(),
97 )),
98 Some(CodexEndpoint::Desktop(path)) => desktop::probe(path, &session.session_id),
99 Some(CodexEndpoint::Unix(path)) if !path.exists() => Err(WakeError::Unavailable(
100 "codex app-server socket missing".into(),
101 )),
102 Some(_) => Ok(()),
103 }
104 }
105
106 fn wake(
107 &mut self,
108 session: &ProviderSession,
109 letter: &WakeLetter<'_>,
110 ) -> Result<WakeOutcome, WakeError> {
111 self.probe(session)?;
112 let text = wake_prompt(session, letter);
113 match self.endpoint.clone() {
114 Some(CodexEndpoint::Desktop(path)) => desktop::wake(&path, &session.session_id, &text),
115 Some(CodexEndpoint::Ws(url)) => {
116 let (mut ws, _) = tungstenite::connect(url.as_str())
117 .map_err(|err| WakeError::Unavailable(format!("codex ws: {err}")))?;
118 if let tungstenite::stream::MaybeTlsStream::Plain(stream) = ws.get_mut() {
119 stream.set_read_timeout(Some(RPC_TIMEOUT))
120 .map_err(|_| WakeError::Transport("codex read timeout setup failed".into()))?;
121 stream.set_write_timeout(Some(RPC_TIMEOUT))
122 .map_err(|_| WakeError::Transport("codex write timeout setup failed".into()))?;
123 }
124 start_turn(&mut ws, &session.session_id, &text)
125 }
126 #[cfg(unix)]
127 Some(CodexEndpoint::Unix(path)) => {
128 let stream = std::os::unix::net::UnixStream::connect(&path)
129 .map_err(|err| WakeError::Unavailable(format!("codex unix: {}", err.kind())))?;
130 stream
131 .set_read_timeout(Some(RPC_TIMEOUT))
132 .map_err(|err| WakeError::Transport(err.to_string()))?;
133 let (mut ws, _) = tungstenite::client("ws://localhost/", stream)
134 .map_err(|err| WakeError::Unavailable(format!("codex handshake: {err}")))?;
135 start_turn(&mut ws, &session.session_id, &text)
136 }
137 #[cfg(not(unix))]
138 Some(CodexEndpoint::Unix(_)) => Err(WakeError::Unavailable(
139 "unix:// app-server is unix-only here; use ws://127.0.0.1:PORT".into(),
140 )),
141 None => Err(WakeError::Unavailable(
142 "codex app-server endpoint unset".into(),
143 )),
144 }
145 }
146}
147
148pub fn initialize_request(id: i64) -> Value {
150 json!({
151 "id": id,
152 "method": "initialize",
153 "params": {"clientInfo": {"name": "mail4agent", "version": env!("CARGO_PKG_VERSION")}},
154 })
155}
156
157pub fn turn_start_request(id: i64, thread_id: &str, text: &str) -> Value {
159 json!({
160 "id": id,
161 "method": "turn/start",
162 "params": {"threadId": thread_id, "input": [{"type": "text", "text": text}]},
163 })
164}
165
166pub fn thread_resume_request(id: i64, thread_id: &str) -> Value {
168 json!({"id": id, "method": "thread/resume", "params": {"threadId": thread_id}})
169}
170
171pub fn start_turn<S: Read + Write>(
174 ws: &mut WebSocket<S>,
175 thread_id: &str,
176 text: &str,
177) -> Result<WakeOutcome, WakeError> {
178 call(ws, initialize_request(0))?
179 .map_err(|msg| WakeError::Transport(format!("initialize: {msg}")))?;
180 send(ws, json!({"method": "initialized"}))?;
181 match call(ws, turn_start_request(1, thread_id, text))? {
182 Ok(_) => Ok(WakeOutcome::Delivered),
183 Err(msg) if msg.contains("not found") || msg.contains("not loaded") => {
184 call(ws, thread_resume_request(2, thread_id))?
185 .map_err(|msg| WakeError::Transport(format!("thread/resume: {msg}")))?;
186 call(ws, turn_start_request(3, thread_id, text))?
187 .map(|_| WakeOutcome::Delivered)
188 .map_err(|msg| WakeError::Transport(format!("turn/start: {msg}")))
189 }
190 Err(msg) => Err(WakeError::Transport(format!("turn/start: {msg}"))),
191 }
192}
193
194fn send<S: Read + Write>(ws: &mut WebSocket<S>, value: Value) -> Result<(), WakeError> {
195 ws.send(Message::text(value.to_string()))
196 .map_err(|err| WakeError::Transport(format!("codex send: {err}")))
197}
198
199fn call<S: Read + Write>(
203 ws: &mut WebSocket<S>,
204 request: Value,
205) -> Result<Result<Value, String>, WakeError> {
206 let id = request["id"].clone();
207 send(ws, request)?;
208 let deadline = Instant::now() + RPC_TIMEOUT;
209 while Instant::now() < deadline {
210 let message = ws
211 .read()
212 .map_err(|err| WakeError::Transport(format!("codex read: {err}")))?;
213 let Message::Text(text) = message else {
214 continue;
215 };
216 let Ok(value) = serde_json::from_str::<Value>(text.as_str()) else {
217 continue;
218 };
219 if value.get("method").is_some() || value.get("id") != Some(&id) {
220 continue;
221 }
222 if let Some(error) = value.get("error") {
223 let msg = error["message"].as_str().unwrap_or("error").to_string();
224 return Ok(Err(msg));
225 }
226 return Ok(Ok(value.get("result").cloned().unwrap_or(Value::Null)));
227 }
228 Err(WakeError::Transport("codex rpc timed out".into()))
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234 use crate::tests::{letter, session};
235 use std::net::TcpListener;
236
237 fn fake_server(loaded: bool) -> (String, std::thread::JoinHandle<Vec<String>>) {
238 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
239 let endpoint = format!("ws://{}", listener.local_addr().unwrap());
240 let handle = std::thread::spawn(move || {
241 let (stream, _) = listener.accept().unwrap();
242 let mut ws = tungstenite::accept(stream).unwrap();
243 let mut seen = Vec::new();
244 let mut loaded = loaded;
245 loop {
246 let Ok(Message::Text(text)) = ws.read() else {
247 break;
248 };
249 let v: Value = serde_json::from_str(text.as_str()).unwrap();
250 let method = v["method"].as_str().unwrap_or("").to_string();
251 seen.push(method.clone());
252 assert!(v.get("jsonrpc").is_none());
253 let reply = match method.as_str() {
254 "initialize" => json!({"id": v["id"], "result": {"userAgent": "fake"}}),
255 "initialized" => continue,
256 "thread/resume" => {
257 loaded = true;
258 json!({"id": v["id"], "result": {"thread": {"id": "t"}}})
259 }
260 "turn/start" if loaded => {
261 assert_eq!(v["params"]["input"][0]["type"], "text");
262 ws.send(Message::text(
263 json!({"method": "turn/started", "params": {}}).to_string(),
264 ))
265 .unwrap();
266 json!({"id": v["id"], "result": {"turn": {"id": "u", "status": "inProgress"}}})
267 }
268 "turn/start" => {
269 json!({"id": v["id"], "error": {"code": -32600, "message": "thread not found: t"}})
270 }
271 _ => json!({"id": v["id"], "error": {"code": -32601, "message": "nope"}}),
272 };
273 ws.send(Message::text(reply.to_string())).unwrap();
274 if seen.iter().filter(|m| *m == "turn/start").count() >= 1
275 && loaded
276 && method == "turn/start"
277 {
278 break;
279 }
280 }
281 seen
282 });
283 (endpoint, handle)
284 }
285
286 #[test]
287 fn loaded_thread_gets_one_turn_start() {
288 let (endpoint, server) = fake_server(true);
289 let mut adapter = CodexAppServerAdapter::new(CodexEndpoint::parse(&endpoint));
290 let s = session(SessionKind::local(ProviderKind::Codex));
291 assert_eq!(
292 adapter.wake(&s, &letter("hi")).unwrap(),
293 WakeOutcome::Delivered
294 );
295 assert_eq!(
296 server.join().unwrap(),
297 ["initialize", "initialized", "turn/start"]
298 );
299 }
300
301 #[test]
302 fn unloaded_thread_is_resumed_then_started() {
303 let (endpoint, server) = fake_server(false);
304 let mut adapter = CodexAppServerAdapter::new(CodexEndpoint::parse(&endpoint));
305 let s = session(SessionKind::local(ProviderKind::Codex));
306 assert_eq!(
307 adapter.wake(&s, &letter("hi")).unwrap(),
308 WakeOutcome::Delivered
309 );
310 assert_eq!(
311 server.join().unwrap(),
312 [
313 "initialize",
314 "initialized",
315 "turn/start",
316 "thread/resume",
317 "turn/start"
318 ]
319 );
320 }
321
322 #[test]
323 fn endpoint_parse_is_loopback_only() {
324 assert_eq!(
325 CodexEndpoint::parse("unix:///tmp/a.sock"),
326 Some(CodexEndpoint::Unix("/tmp/a.sock".into()))
327 );
328 assert!(CodexEndpoint::parse("ws://127.0.0.1:4500").is_some());
329 assert!(CodexEndpoint::parse("ws://10.0.0.5:4500").is_none());
330 assert!(CodexEndpoint::parse("unix://").is_none());
331 }
332
333 #[test]
335 #[ignore]
336 fn live_turn_start() {
337 let endpoint = std::env::var("M4A_LIVE_CODEX_APP_SERVER").unwrap();
338 let thread = std::env::var("M4A_LIVE_CODEX_THREAD").unwrap();
339 let mut adapter = CodexAppServerAdapter::new(CodexEndpoint::parse(&endpoint));
340 let mut s = session(SessionKind::local(ProviderKind::Codex));
341 s.session_id = thread;
342 assert_eq!(
343 adapter.wake(&s, &letter("live probe")).unwrap(),
344 WakeOutcome::Delivered
345 );
346 }
347}