Skip to main content

construct/gateway/
terminal.rs

1//! WebSocket PTY terminal handler.
2//!
3//! Connect: `ws://host:port/ws/terminal?session_id=ID&token=TOKEN`
4//!
5//! The handler spawns a PTY shell and bridges I/O bidirectionally over
6//! the WebSocket. The frontend (xterm.js) sends raw keystrokes as text
7//! frames and receives terminal output as text or binary frames.
8//!
9//! ## Protocol
10//!
11//! ```text
12//! Client -> Server: raw keystroke data (text frame)
13//! Client -> Server: {"type":"resize","cols":120,"rows":40}
14//! Server -> Client: terminal output (text frame)
15//! ```
16//!
17//! ## Tool-aware Code tab (M2)
18//!
19//! When `tool=<claude|codex|opencode|gemini>` is passed in the query, this
20//! handler spawns that CLI inside the PTY instead of the user's shell. When
21//! `mcp_session` and `mcp_token` are also present, a per-session MCP config
22//! file is written to a temp dir and exported via env vars so the CLI can
23//! auto-register the in-process MCP server (run as a tokio task inside the
24//! main daemon; discovered via `~/.construct/mcp.json`). The temp dir is
25//! removed when the socket closes.
26
27use super::AppState;
28use super::mcp_discovery::read_construct_mcp;
29use axum::{
30    extract::{
31        Query, State, WebSocketUpgrade,
32        ws::{Message, WebSocket},
33    },
34    http::{HeaderMap, StatusCode, header},
35    response::IntoResponse,
36};
37use futures_util::{SinkExt, StreamExt};
38use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
39use serde::Deserialize;
40use serde_json::json;
41use std::io::{Read, Write};
42use std::path::PathBuf;
43use tracing::{debug, error, warn};
44use uuid::Uuid;
45
46/// The sub-protocol we support for terminal WebSocket.
47const WS_PROTOCOL: &str = "construct.v1";
48
49/// Prefix used in `Sec-WebSocket-Protocol` to carry a bearer token.
50const BEARER_SUBPROTO_PREFIX: &str = "bearer.";
51
52#[derive(Deserialize, Default)]
53pub struct TerminalQuery {
54    pub token: Option<String>,
55    pub session_id: Option<String>,
56    /// Optional CLI tool to launch instead of the default shell.
57    /// Known values: `claude` | `codex` | `opencode` | `gemini`.
58    /// Anything else (or `None`) falls back to `$SHELL -l`.
59    pub tool: Option<String>,
60    /// Optional explicit working directory for the spawned process. Tilde
61    /// expansion is applied. Rejected (error frame) if it does not resolve
62    /// to a directory.
63    pub cwd: Option<String>,
64    /// Session id issued by the in-process MCP server (see
65    /// `mcp_server` — runs as a tokio task inside this daemon).
66    pub mcp_session: Option<String>,
67    /// Bearer token issued by the in-process MCP server.
68    pub mcp_token: Option<String>,
69    /// Initial terminal column count (defaults to 80 if absent/zero). Supplied
70    /// by the frontend after xterm's FitAddon measures the container, so the
71    /// child process's first layout matches what the user actually sees —
72    /// avoiding a 80×24 → resize repaint race that garbles TUIs.
73    pub cols: Option<u16>,
74    /// Initial terminal row count (defaults to 24 if absent/zero).
75    pub rows: Option<u16>,
76}
77
78/// Resize message sent from xterm.js frontend.
79#[derive(Deserialize)]
80struct ResizeMsg {
81    #[serde(rename = "type")]
82    msg_type: String,
83    cols: u16,
84    rows: u16,
85}
86
87/// Known CLI tools the Code tab can spawn.
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum CodeTool {
90    Claude,
91    Codex,
92    OpenCode,
93    Gemini,
94}
95
96impl CodeTool {
97    /// Parse the `tool` query parameter. Unknown strings return `None`, in
98    /// which case the handler falls back to spawning the user's shell.
99    pub fn from_query(s: &str) -> Option<Self> {
100        match s {
101            "claude" => Some(Self::Claude),
102            "codex" => Some(Self::Codex),
103            "opencode" => Some(Self::OpenCode),
104            "gemini" => Some(Self::Gemini),
105            _ => None,
106        }
107    }
108
109    /// Binary name to look up in `$PATH`.
110    pub fn binary(self) -> &'static str {
111        match self {
112            Self::Claude => "claude",
113            Self::Codex => "codex",
114            Self::OpenCode => "opencode",
115            Self::Gemini => "gemini",
116        }
117    }
118
119    /// Env variable each CLI (allegedly) looks at to find an MCP config file.
120    ///
121    /// Kept for backwards-compat with any older CLI version that honored a
122    /// generic env var. The authoritative mechanism for M3 is the per-CLI
123    /// adapter in `write_cli_config` — these env names are best-effort and
124    /// not relied upon by callers.
125    pub fn config_env(self) -> &'static str {
126        match self {
127            Self::Claude => "CLAUDE_MCP_CONFIG",
128            Self::Codex => "CODEX_MCP_CONFIG",
129            Self::OpenCode => "OPENCODE_MCP_CONFIG",
130            Self::Gemini => "GEMINI_MCP_CONFIG",
131        }
132    }
133}
134
135/// Adapter result: concrete files written under the spawn's temp HOME, plus
136/// any CLI-specific args to pass on the command line.
137#[derive(Debug)]
138pub struct CliInjection {
139    /// Extra argv to append after the binary.
140    pub args: Vec<String>,
141    /// Files written, for test assertion & debugging: (relative_path, content_sample).
142    pub files_written: Vec<PathBuf>,
143}
144
145/// Write the per-CLI MCP config under `temp_home`, returning any args that
146/// need to be appended to the command line.
147///
148/// This is the M3 per-CLI adapter. Each branch cites the source that
149/// documents the config location / mechanism so future maintainers can
150/// re-verify when the upstream CLI changes.
151pub fn write_cli_config(
152    tool: CodeTool,
153    temp_home: &std::path::Path,
154    mcp_url: &str,
155    session_id: &str,
156    token: &str,
157) -> Result<CliInjection, String> {
158    match tool {
159        // ── Claude Code ────────────────────────────────────────────────
160        // Source: https://docs.claude.com/en/docs/claude-code/mcp
161        //   > claude --mcp-config <path-to-json>
162        // The flag accepts the same `{ "mcpServers": { ... } }` shape
163        // used by `.mcp.json`. We prefer the flag over dropping `.mcp.json`
164        // into the cwd because it avoids mutating the user's repo.
165        CodeTool::Claude => {
166            let cfg_path = temp_home.join(".mcp.json");
167            let cfg = build_mcp_config_json(mcp_url, session_id, token);
168            std::fs::write(
169                &cfg_path,
170                serde_json::to_vec_pretty(&cfg).expect("serialize claude mcp config"),
171            )
172            .map_err(|e| format!("writing claude mcp config: {e}"))?;
173            Ok(CliInjection {
174                args: vec![
175                    "--mcp-config".into(),
176                    cfg_path.to_string_lossy().into_owned(),
177                ],
178                files_written: vec![cfg_path],
179            })
180        }
181
182        // ── Codex (OpenAI) ─────────────────────────────────────────────
183        // Source: https://github.com/openai/codex — `~/.codex/config.toml`
184        // with `[mcp_servers.<name>]` blocks. Recent versions support a
185        // `url` + `transport = "http"` entry for Streamable HTTP; we rely
186        // on that here. Codex has no per-invocation config flag, so we
187        // must redirect HOME to the temp dir.
188        CodeTool::Codex => {
189            let dir = temp_home.join(".codex");
190            std::fs::create_dir_all(&dir).map_err(|e| format!("creating ~/.codex: {e}"))?;
191            let cfg_path = dir.join("config.toml");
192            let mut toml = String::new();
193            toml.push_str("[mcp_servers.construct]\n");
194            toml.push_str(&format!("url = {}\n", toml_string(mcp_url)));
195            toml.push_str("transport = \"http\"\n");
196            toml.push_str("[mcp_servers.construct.headers]\n");
197            toml.push_str(&format!(
198                "Authorization = {}\n",
199                toml_string(&format!("Bearer {token}"))
200            ));
201            toml.push_str(&format!(
202                "X-Construct-Session = {}\n",
203                toml_string(session_id)
204            ));
205            std::fs::write(&cfg_path, toml.as_bytes())
206                .map_err(|e| format!("writing codex config: {e}"))?;
207            Ok(CliInjection {
208                args: vec![],
209                files_written: vec![cfg_path],
210            })
211        }
212
213        // ── OpenCode ───────────────────────────────────────────────────
214        // Source: https://opencode.ai/docs — config at
215        // `~/.config/opencode/config.json` (XDG_CONFIG_HOME respected)
216        // with top-level `mcp` map keyed by server name: each value is
217        // `{ type: "remote", url, headers }` for HTTP servers.
218        CodeTool::OpenCode => {
219            let dir = temp_home.join(".config").join("opencode");
220            std::fs::create_dir_all(&dir)
221                .map_err(|e| format!("creating opencode config dir: {e}"))?;
222            let cfg_path = dir.join("config.json");
223            let cfg = json!({
224                "$schema": "https://opencode.ai/config.json",
225                "mcp": {
226                    "construct": {
227                        "type": "remote",
228                        "url": mcp_url,
229                        "enabled": true,
230                        "headers": {
231                            "Authorization": format!("Bearer {token}"),
232                            "X-Construct-Session": session_id,
233                        }
234                    }
235                }
236            });
237            std::fs::write(
238                &cfg_path,
239                serde_json::to_vec_pretty(&cfg).expect("serialize opencode config"),
240            )
241            .map_err(|e| format!("writing opencode config: {e}"))?;
242            Ok(CliInjection {
243                args: vec![],
244                files_written: vec![cfg_path],
245            })
246        }
247
248        // ── Gemini CLI ─────────────────────────────────────────────────
249        // Source: https://github.com/google-gemini/gemini-cli — settings
250        // file at `~/.gemini/settings.json`, with an `mcpServers` map
251        // matching the shape used by Claude/Codex. HTTP servers use the
252        // `httpUrl` key (not `url`), per the documented schema.
253        CodeTool::Gemini => {
254            let dir = temp_home.join(".gemini");
255            std::fs::create_dir_all(&dir).map_err(|e| format!("creating ~/.gemini: {e}"))?;
256            let cfg_path = dir.join("settings.json");
257            let cfg = json!({
258                "mcpServers": {
259                    "construct": {
260                        "httpUrl": mcp_url,
261                        "headers": {
262                            "Authorization": format!("Bearer {token}"),
263                            "X-Construct-Session": session_id,
264                        }
265                    }
266                }
267            });
268            std::fs::write(
269                &cfg_path,
270                serde_json::to_vec_pretty(&cfg).expect("serialize gemini config"),
271            )
272            .map_err(|e| format!("writing gemini config: {e}"))?;
273            Ok(CliInjection {
274                args: vec![],
275                files_written: vec![cfg_path],
276            })
277        }
278    }
279}
280
281/// Minimal TOML string escaper (double-quoted basic string form).
282/// Sufficient for the URL + bearer strings we write; not a general-purpose
283/// escaper. Escapes backslash and double-quote.
284fn toml_string(s: &str) -> String {
285    let mut out = String::with_capacity(s.len() + 2);
286    out.push('"');
287    for c in s.chars() {
288        match c {
289            '\\' => out.push_str("\\\\"),
290            '"' => out.push_str("\\\""),
291            '\n' => out.push_str("\\n"),
292            '\r' => out.push_str("\\r"),
293            '\t' => out.push_str("\\t"),
294            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04X}", c as u32)),
295            c => out.push(c),
296        }
297    }
298    out.push('"');
299    out
300}
301
302/// Extract bearer token from WS-compatible sources (header > subprotocol > query param).
303fn extract_ws_token<'a>(headers: &'a HeaderMap, query_token: Option<&'a str>) -> Option<&'a str> {
304    // 1. Authorization header
305    if let Some(t) = headers
306        .get(header::AUTHORIZATION)
307        .and_then(|v| v.to_str().ok())
308        .and_then(|auth| auth.strip_prefix("Bearer "))
309    {
310        if !t.is_empty() {
311            return Some(t);
312        }
313    }
314
315    // 2. Sec-WebSocket-Protocol: bearer.<token>
316    if let Some(t) = headers
317        .get("sec-websocket-protocol")
318        .and_then(|v| v.to_str().ok())
319        .and_then(|protos| {
320            protos
321                .split(',')
322                .map(|p| p.trim())
323                .find_map(|p| p.strip_prefix(BEARER_SUBPROTO_PREFIX))
324        })
325    {
326        if !t.is_empty() {
327            return Some(t);
328        }
329    }
330
331    // 3. ?token= query parameter
332    if let Some(t) = query_token {
333        if !t.is_empty() {
334            return Some(t);
335        }
336    }
337
338    None
339}
340
341/// Build the generic `mcpServers` config document pointed at the local daemon.
342///
343/// Format chosen to match the `mcp.json` convention adopted by most
344/// MCP-aware CLIs:
345///
346/// ```json
347/// { "mcpServers": { "construct": { "url": "...", "headers": { ... } } } }
348/// ```
349pub fn build_mcp_config_json(mcp_url: &str, session_id: &str, token: &str) -> serde_json::Value {
350    json!({
351        "mcpServers": {
352            "construct": {
353                "type": "http",
354                "url": mcp_url,
355                "headers": {
356                    "Authorization": format!("Bearer {token}"),
357                    "X-Construct-Session": session_id,
358                }
359            }
360        }
361    })
362}
363
364/// Holds a temp dir that is removed on drop. Used to clean up per-session
365/// MCP config files when the WS closes.
366struct TempSpawnDir(PathBuf);
367
368impl Drop for TempSpawnDir {
369    fn drop(&mut self) {
370        let _ = std::fs::remove_dir_all(&self.0);
371    }
372}
373
374/// Resolve an optional cwd string: tilde-expand, canonicalize, verify it
375/// is a directory. Returns `None` if unspecified, `Err` if invalid.
376fn resolve_cwd(raw: Option<&str>) -> Result<Option<PathBuf>, String> {
377    let Some(s) = raw.filter(|s| !s.is_empty()) else {
378        return Ok(None);
379    };
380    let expanded = shellexpand::tilde(s).into_owned();
381    let p = PathBuf::from(&expanded);
382    let canon = p.canonicalize().map_err(|e| format!("{s}: {e}"))?;
383    if !canon.is_dir() {
384        return Err(format!("{} is not a directory", canon.display()));
385    }
386    Ok(Some(canon))
387}
388
389/// GET /ws/terminal — WebSocket upgrade for PTY terminal
390pub async fn handle_ws_terminal(
391    State(state): State<AppState>,
392    Query(params): Query<TerminalQuery>,
393    headers: HeaderMap,
394    ws: WebSocketUpgrade,
395) -> impl IntoResponse {
396    // Auth check
397    if state.pairing.require_pairing() {
398        let token = extract_ws_token(&headers, params.token.as_deref()).unwrap_or("");
399        if !state.pairing.is_authenticated(token) {
400            return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
401        }
402    }
403
404    // Echo sub-protocol if client requests it
405    let ws = if headers
406        .get("sec-websocket-protocol")
407        .and_then(|v| v.to_str().ok())
408        .map_or(false, |protos| {
409            protos.split(',').any(|p| p.trim() == WS_PROTOCOL)
410        }) {
411        ws.protocols([WS_PROTOCOL])
412    } else {
413        ws
414    };
415
416    if let Some(ref logger) = state.audit_logger {
417        let _ = logger.log_security_event("dashboard", "WebSocket terminal session connected");
418    }
419
420    ws.on_upgrade(move |socket| handle_terminal_socket(socket, params))
421        .into_response()
422}
423
424/// Helper: send a red error frame over the WS. Silently swallows errors.
425async fn send_err(ws_sender: &mut futures_util::stream::SplitSink<WebSocket, Message>, msg: &str) {
426    let _ = ws_sender
427        .send(Message::Text(format!("\x1b[31m{msg}\x1b[0m\r\n").into()))
428        .await;
429}
430
431/// Spawn configuration assembled for the child process.
432struct SpawnPlan {
433    cmd: CommandBuilder,
434    /// Kept alive for the duration of the PTY session — Drop removes the temp dir.
435    _temp: Option<TempSpawnDir>,
436}
437
438/// Assemble the `CommandBuilder` for the spawned process. Split out from
439/// `build_command` (which got tangled) for readability.
440fn plan_spawn(
441    tool: Option<CodeTool>,
442    cwd: Option<PathBuf>,
443    mcp_session: Option<&str>,
444    mcp_token: Option<&str>,
445) -> Result<SpawnPlan, String> {
446    // `plan_spawn_with_discovery` is the real implementation; the non-suffixed
447    // variant calls `read_construct_mcp()` for the prod path. Split out so the
448    // per-CLI adapter tests can pass in a fake discovery URL without needing
449    // `~/.construct/mcp.json` on disk.
450    let discovery_url = if tool.is_some() && mcp_session.is_some() && mcp_token.is_some() {
451        Some(
452            read_construct_mcp()
453                .map_err(|e| format!("in-process MCP server not available: {e}"))?
454                .url,
455        )
456    } else {
457        None
458    };
459    plan_spawn_with_discovery(tool, cwd, mcp_session, mcp_token, discovery_url.as_deref())
460}
461
462fn plan_spawn_with_discovery(
463    tool: Option<CodeTool>,
464    cwd: Option<PathBuf>,
465    mcp_session: Option<&str>,
466    mcp_token: Option<&str>,
467    mcp_url: Option<&str>,
468) -> Result<SpawnPlan, String> {
469    let (mut cmd, temp) = match tool {
470        Some(t) => {
471            let bin = which::which(t.binary())
472                .map_err(|_| format!("{} not found in PATH", t.binary()))?;
473            let mut cmd = CommandBuilder::new(bin);
474
475            if let (Some(sess), Some(tok), Some(url)) = (mcp_session, mcp_token, mcp_url) {
476                let dir = std::env::temp_dir().join(format!("construct-code-{}", Uuid::new_v4()));
477                std::fs::create_dir_all(&dir).map_err(|e| format!("creating temp dir: {e}"))?;
478
479                // Per-CLI adapter: writes the right config file(s) under
480                // `dir` (which we then expose to the child as HOME) and
481                // returns any argv the CLI needs.
482                let injection = write_cli_config(t, &dir, url, sess, tok)?;
483                for a in &injection.args {
484                    cmd.arg(a);
485                }
486
487                // Redirect HOME + XDG_CONFIG_HOME so CLIs that read from
488                // `~/.codex/...`, `~/.gemini/...`, `~/.config/opencode/...`
489                // pick up our freshly written config instead of the user's
490                // real dotfiles.
491                cmd.env("HOME", &dir);
492                cmd.env("XDG_CONFIG_HOME", dir.join(".config"));
493
494                // Stable fallback env — harmless if unused.
495                cmd.env("CONSTRUCT_MCP_URL", url);
496                cmd.env("CONSTRUCT_MCP_SESSION", sess);
497                cmd.env("CONSTRUCT_MCP_TOKEN", tok);
498
499                (cmd, Some(TempSpawnDir(dir)))
500            } else {
501                (cmd, None)
502            }
503        }
504        None => {
505            // SHELL isn't set on Windows by default and `/bin/sh` doesn't
506            // exist there. Pick the platform-native default so spawn
507            // doesn't fail with "specified path cannot be found".
508            let shell = std::env::var("SHELL").unwrap_or_else(|_| {
509                if cfg!(windows) {
510                    std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string())
511                } else {
512                    "/bin/sh".to_string()
513                }
514            });
515            let mut cmd = CommandBuilder::new(&shell);
516            // `-l` is a Unix login-shell flag; cmd.exe / powershell.exe
517            // reject it.
518            if !cfg!(windows) {
519                cmd.arg("-l");
520            }
521            (cmd, None)
522        }
523    };
524
525    if let Some(c) = cwd {
526        cmd.cwd(c);
527    }
528
529    // portable_pty::CommandBuilder starts with an empty env — the child would
530    // otherwise have no PATH, TERM, locale, etc. Readline / zle need TERM to
531    // map ^? (DEL) to backward-delete-char; without it, backspace silently
532    // fails. We pin TERM/COLORTERM to values xterm.js speaks, then carry a
533    // small allow-list of user env vars forward.
534    cmd.env("TERM", "xterm-256color");
535    cmd.env("COLORTERM", "truecolor");
536    for key in [
537        "PATH", "LANG", "LC_ALL", "LC_CTYPE", "USER", "LOGNAME", "SHELL", "TZ",
538    ] {
539        if let Ok(val) = std::env::var(key) {
540            cmd.env(key, val);
541        }
542    }
543    // Tool-branch has already redirected HOME to a per-session temp dir (so the
544    // CLI picks up our freshly-written MCP config) — don't clobber that. The
545    // shell-fallback branch (`temp.is_none()`) needs the real HOME so rc files
546    // load normally.
547    if temp.is_none() {
548        if let Ok(home) = std::env::var("HOME") {
549            cmd.env("HOME", home);
550        }
551    }
552
553    Ok(SpawnPlan { cmd, _temp: temp })
554}
555
556async fn handle_terminal_socket(socket: WebSocket, params: TerminalQuery) {
557    let (mut ws_sender, mut ws_receiver) = socket.split();
558
559    // Validate the `tool` name: known -> CodeTool, unknown -> fallback to shell
560    // (matches docs: "anything else, keep current behavior").
561    let tool = params.tool.as_deref().and_then(CodeTool::from_query);
562
563    // Resolve cwd (Err = user-visible).
564    let cwd = match resolve_cwd(params.cwd.as_deref()) {
565        Ok(c) => c,
566        Err(msg) => {
567            send_err(&mut ws_sender, &format!("Invalid cwd: {msg}")).await;
568            return;
569        }
570    };
571
572    // Assemble the spawn plan (binary lookup, MCP config, env vars).
573    let plan = match plan_spawn(
574        tool,
575        cwd,
576        params.mcp_session.as_deref(),
577        params.mcp_token.as_deref(),
578    ) {
579        Ok(p) => p,
580        Err(msg) => {
581            send_err(&mut ws_sender, &msg).await;
582            let _ = ws_sender.send(Message::Close(None)).await;
583            return;
584        }
585    };
586
587    let initial_size = PtySize {
588        rows: params.rows.filter(|r| *r > 0).unwrap_or(24),
589        cols: params.cols.filter(|c| *c > 0).unwrap_or(80),
590        pixel_width: 0,
591        pixel_height: 0,
592    };
593
594    // Spawn PTY
595    let pty_system = NativePtySystem::default();
596    let pair = match pty_system.openpty(initial_size) {
597        Ok(pair) => pair,
598        Err(e) => {
599            error!(error = %e, "Failed to open PTY");
600            send_err(&mut ws_sender, &format!("Failed to open PTY: {e}")).await;
601            return;
602        }
603    };
604
605    let SpawnPlan { cmd, _temp } = plan;
606    let _child = match pair.slave.spawn_command(cmd) {
607        Ok(child) => child,
608        Err(e) => {
609            error!(error = %e, "Failed to spawn child");
610            send_err(&mut ws_sender, &format!("Failed to spawn child: {e}")).await;
611            return;
612        }
613    };
614    // Drop slave — master owns the PTY fd now
615    drop(pair.slave);
616
617    let master = pair.master;
618
619    let mut pty_reader = match master.try_clone_reader() {
620        Ok(r) => r,
621        Err(e) => {
622            error!(error = %e, "Failed to clone PTY reader");
623            return;
624        }
625    };
626
627    let mut pty_writer: Box<dyn Write + Send> = match master.take_writer() {
628        Ok(w) => w,
629        Err(e) => {
630            error!(error = %e, "Failed to take PTY writer");
631            return;
632        }
633    };
634
635    // Channels to bridge blocking PTY I/O with async WebSocket
636    let (pty_out_tx, mut pty_out_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(64);
637    let (resize_tx, mut resize_rx) = tokio::sync::mpsc::channel::<(u16, u16)>(4);
638
639    // Blocking task: PTY stdout -> mpsc channel
640    tokio::task::spawn_blocking(move || {
641        let mut buf = [0u8; 4096];
642        loop {
643            match pty_reader.read(&mut buf) {
644                Ok(0) => break,
645                Ok(n) => {
646                    if pty_out_tx.blocking_send(buf[..n].to_vec()).is_err() {
647                        break;
648                    }
649                }
650                Err(_) => break,
651            }
652        }
653    });
654
655    // Async task: handle resize requests
656    tokio::spawn(async move {
657        while let Some((cols, rows)) = resize_rx.recv().await {
658            let _ = master.resize(PtySize {
659                rows,
660                cols,
661                pixel_width: 0,
662                pixel_height: 0,
663            });
664        }
665    });
666
667    // Main loop: bridge WebSocket <-> PTY
668    loop {
669        tokio::select! {
670            // PTY output -> WebSocket
671            Some(data) = pty_out_rx.recv() => {
672                let text = String::from_utf8_lossy(&data).into_owned();
673                if ws_sender.send(Message::Text(text.into())).await.is_err() {
674                    break;
675                }
676            }
677            // WebSocket input -> PTY
678            msg = ws_receiver.next() => {
679                match msg {
680                    Some(Ok(Message::Text(text))) => {
681                        // Check if it's a resize message
682                        if let Ok(resize) = serde_json::from_str::<ResizeMsg>(&text) {
683                            if resize.msg_type == "resize" {
684                                let _ = resize_tx.send((resize.cols, resize.rows)).await;
685                                continue;
686                            }
687                        }
688                        // Raw keystroke input
689                        if pty_writer.write_all(text.as_bytes()).is_err() {
690                            break;
691                        }
692                    }
693                    Some(Ok(Message::Binary(data))) => {
694                        if pty_writer.write_all(&data).is_err() {
695                            break;
696                        }
697                    }
698                    Some(Ok(Message::Close(_))) | None => {
699                        debug!("Terminal WebSocket closed");
700                        break;
701                    }
702                    Some(Ok(_)) => {} // Ping/Pong handled by axum
703                    Some(Err(e)) => {
704                        warn!(error = %e, "Terminal WebSocket error");
705                        break;
706                    }
707                }
708            }
709        }
710    }
711
712    // `_temp` drops here, removing the per-session config dir.
713    drop(_temp);
714
715    debug!("Terminal session ended");
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721
722    #[test]
723    fn tool_mapping_known() {
724        assert_eq!(CodeTool::from_query("claude"), Some(CodeTool::Claude));
725        assert_eq!(CodeTool::from_query("codex"), Some(CodeTool::Codex));
726        assert_eq!(CodeTool::from_query("opencode"), Some(CodeTool::OpenCode));
727        assert_eq!(CodeTool::from_query("gemini"), Some(CodeTool::Gemini));
728    }
729
730    #[test]
731    fn tool_mapping_unknown_falls_back() {
732        assert_eq!(CodeTool::from_query(""), None);
733        assert_eq!(CodeTool::from_query("bash"), None);
734        assert_eq!(CodeTool::from_query("Claude"), None); // case-sensitive
735        assert_eq!(CodeTool::from_query("nonsense"), None);
736    }
737
738    #[test]
739    fn tool_binaries_match_docs() {
740        assert_eq!(CodeTool::Claude.binary(), "claude");
741        assert_eq!(CodeTool::Codex.binary(), "codex");
742        assert_eq!(CodeTool::OpenCode.binary(), "opencode");
743        assert_eq!(CodeTool::Gemini.binary(), "gemini");
744    }
745
746    #[test]
747    fn tool_config_env_vars() {
748        assert_eq!(CodeTool::Claude.config_env(), "CLAUDE_MCP_CONFIG");
749        assert_eq!(CodeTool::Codex.config_env(), "CODEX_MCP_CONFIG");
750        assert_eq!(CodeTool::OpenCode.config_env(), "OPENCODE_MCP_CONFIG");
751        assert_eq!(CodeTool::Gemini.config_env(), "GEMINI_MCP_CONFIG");
752    }
753
754    #[test]
755    fn mcp_config_json_has_expected_shape() {
756        let v = build_mcp_config_json("http://127.0.0.1:54500/mcp", "sess-abc", "tok-xyz");
757        let srv = &v["mcpServers"]["construct"];
758        assert_eq!(srv["url"], "http://127.0.0.1:54500/mcp");
759        assert_eq!(srv["type"], "http");
760        assert_eq!(srv["headers"]["Authorization"], "Bearer tok-xyz");
761        assert_eq!(srv["headers"]["X-Construct-Session"], "sess-abc");
762        // JSON round-trips cleanly.
763        let s = serde_json::to_string(&v).unwrap();
764        let back: serde_json::Value = serde_json::from_str(&s).unwrap();
765        assert_eq!(back, v);
766    }
767
768    #[test]
769    fn resolve_cwd_none_when_unset() {
770        assert!(matches!(resolve_cwd(None), Ok(None)));
771        assert!(matches!(resolve_cwd(Some("")), Ok(None)));
772    }
773
774    #[test]
775    fn resolve_cwd_rejects_missing_path() {
776        assert!(resolve_cwd(Some("/this/should/not/exist/construct-xyz")).is_err());
777    }
778
779    #[test]
780    fn resolve_cwd_accepts_tmp() {
781        let tmp = std::env::temp_dir();
782        let got = resolve_cwd(Some(tmp.to_str().unwrap())).unwrap().unwrap();
783        assert!(got.is_dir());
784    }
785
786    #[test]
787    fn plan_spawn_shell_fallback_no_tool() {
788        // No tool -> shell fallback; must not touch $PATH or MCP discovery.
789        let plan = plan_spawn(None, None, None, None).expect("shell fallback works");
790        // We can't inspect CommandBuilder's inner bin easily without unstable API,
791        // but we can at least confirm no temp dir was created.
792        assert!(plan._temp.is_none());
793    }
794
795    #[test]
796    fn plan_spawn_missing_binary_errors() {
797        // Unless the user happens to have `gemini` installed during `cargo test`,
798        // this should return a "not found in PATH" error. If it *is* installed,
799        // the call succeeds and we just accept that.
800        match plan_spawn(Some(CodeTool::Gemini), None, None, None) {
801            Ok(_) => {} // gemini installed on this machine
802            Err(msg) => assert!(msg.contains("not found in PATH"), "got: {msg}"),
803        }
804    }
805
806    // Backwards-compat: the default `TerminalQuery` (what deserialization of
807    // an empty query string produces) must have no tool selected, so
808    // `handle_terminal_socket` takes the shell fallback path.
809    #[test]
810    fn terminal_query_default_falls_back_to_shell() {
811        let q = TerminalQuery::default();
812        assert!(q.tool.is_none());
813        assert!(q.cwd.is_none());
814        assert!(q.mcp_session.is_none());
815        assert!(q.mcp_token.is_none());
816        assert!(q.tool.as_deref().and_then(CodeTool::from_query).is_none());
817    }
818
819    // ── Per-CLI adapter tests ───────────────────────────────────────
820
821    fn tempdir() -> PathBuf {
822        let p = std::env::temp_dir().join(format!("construct-test-{}", Uuid::new_v4()));
823        std::fs::create_dir_all(&p).unwrap();
824        p
825    }
826
827    const URL: &str = "http://127.0.0.1:54500/mcp";
828    const SESS: &str = "sess-abc";
829    const TOK: &str = "tok-xyz";
830
831    #[test]
832    fn claude_adapter_writes_mcp_json_and_passes_flag() {
833        let home = tempdir();
834        let inj = write_cli_config(CodeTool::Claude, &home, URL, SESS, TOK).unwrap();
835
836        // Flag layout: --mcp-config <path>
837        assert_eq!(inj.args.len(), 2);
838        assert_eq!(inj.args[0], "--mcp-config");
839        let cfg_path = PathBuf::from(&inj.args[1]);
840        assert!(cfg_path.starts_with(&home));
841        assert!(cfg_path.ends_with(".mcp.json"));
842        assert!(cfg_path.exists());
843
844        let content: serde_json::Value =
845            serde_json::from_slice(&std::fs::read(&cfg_path).unwrap()).unwrap();
846        assert_eq!(content["mcpServers"]["construct"]["url"], URL);
847        assert_eq!(
848            content["mcpServers"]["construct"]["headers"]["Authorization"],
849            format!("Bearer {TOK}")
850        );
851        assert_eq!(
852            content["mcpServers"]["construct"]["headers"]["X-Construct-Session"],
853            SESS
854        );
855        let _ = std::fs::remove_dir_all(&home);
856    }
857
858    #[test]
859    fn codex_adapter_writes_toml_at_home_dot_codex() {
860        let home = tempdir();
861        let inj = write_cli_config(CodeTool::Codex, &home, URL, SESS, TOK).unwrap();
862        assert!(inj.args.is_empty(), "codex has no flag mechanism");
863        let cfg = home.join(".codex").join("config.toml");
864        assert!(cfg.exists(), "{} should exist", cfg.display());
865        let body = std::fs::read_to_string(&cfg).unwrap();
866        assert!(body.contains("[mcp_servers.construct]"));
867        assert!(body.contains(&format!("url = \"{URL}\"")), "body: {body}");
868        assert!(body.contains("transport = \"http\""));
869        assert!(body.contains("[mcp_servers.construct.headers]"));
870        assert!(
871            body.contains(&format!("Authorization = \"Bearer {TOK}\"")),
872            "body: {body}"
873        );
874        assert!(
875            body.contains(&format!("X-Construct-Session = \"{SESS}\"")),
876            "body: {body}"
877        );
878        // Sanity: parses as TOML round-trip.
879        let _: toml::Value = toml::from_str(&body).expect("codex config should be valid TOML");
880        let _ = std::fs::remove_dir_all(&home);
881    }
882
883    #[test]
884    fn opencode_adapter_writes_xdg_json() {
885        let home = tempdir();
886        let inj = write_cli_config(CodeTool::OpenCode, &home, URL, SESS, TOK).unwrap();
887        assert!(inj.args.is_empty());
888        let cfg = home.join(".config").join("opencode").join("config.json");
889        assert!(cfg.exists());
890        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&cfg).unwrap()).unwrap();
891        let srv = &v["mcp"]["construct"];
892        assert_eq!(srv["type"], "remote");
893        assert_eq!(srv["url"], URL);
894        assert_eq!(srv["enabled"], true);
895        assert_eq!(srv["headers"]["Authorization"], format!("Bearer {TOK}"));
896        assert_eq!(srv["headers"]["X-Construct-Session"], SESS);
897        let _ = std::fs::remove_dir_all(&home);
898    }
899
900    #[test]
901    fn gemini_adapter_writes_settings_json() {
902        let home = tempdir();
903        let inj = write_cli_config(CodeTool::Gemini, &home, URL, SESS, TOK).unwrap();
904        assert!(inj.args.is_empty());
905        let cfg = home.join(".gemini").join("settings.json");
906        assert!(cfg.exists());
907        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&cfg).unwrap()).unwrap();
908        let srv = &v["mcpServers"]["construct"];
909        assert_eq!(srv["httpUrl"], URL);
910        assert_eq!(srv["headers"]["Authorization"], format!("Bearer {TOK}"));
911        assert_eq!(srv["headers"]["X-Construct-Session"], SESS);
912        let _ = std::fs::remove_dir_all(&home);
913    }
914
915    #[test]
916    fn plan_spawn_with_discovery_no_creds_no_tempdir() {
917        // Even with a discovery URL, without session+token we don't write
918        // anything — matches the gating in the handler.
919        let plan = plan_spawn_with_discovery(None, None, None, None, Some(URL))
920            .expect("shell fallback works");
921        assert!(plan._temp.is_none());
922    }
923}