From 1cd728c11692d64daa951eadaba499cbfa60f446 Mon Sep 17 00:00:00 2001 From: iret77 <63622643+iret77@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:33:48 +0200 Subject: [PATCH] feat(upload): add file upload from Mac into agent session (#146) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce the `upload` MCP tool and `/aiui:upload` slash-command — the first aiui data flow that runs Mac → agent-host instead of the reverse. Calling `upload` opens a native file picker on the Mac (Tauri dialog plugin); the chosen file streams back over the existing authenticated :7777 channel (the reverse direction of POST /media) and is written to `target_dir/` on the host the agent runs on. Replaces the "please scp me that file" round-trip. - Companion: new authenticated `POST /upload` endpoint (http.rs) that runs the picker and returns the file bytes + percent-encoded filename header, with 204 (cancel), 413 (>512 MB cap), and 500 (read error) paths. - Rust bridge (mcp.rs): `upload` tool + `do_upload` — resolves/validates target_dir, decodes the filename, sanitises to a base name, and does a no-clobber atomic write. Deterministic destination, never overwrites. - Python bridge (server.py): mirror implementation for remote SSH hosts via uvx, incl. a progress heartbeat so the held picker call keeps the client alive. - Prompt `upload` + INSTRUCTIONS trigger on both servers; docs/skill.md, python skill.md, and CHANGELOG updated. - Tests: filename percent-encode round-trip (Rust) and base-name / target-dir / no-clobber write (Python). --- CHANGELOG.md | 25 ++++ companion/src-tauri/src/http.rs | 216 ++++++++++++++++++++++++++++++++ companion/src-tauri/src/mcp.rs | 198 +++++++++++++++++++++++++++++ docs/skill.md | 36 ++++++ python/src/aiui_mcp/server.py | 192 ++++++++++++++++++++++++++++ python/src/aiui_mcp/skill.md | 20 +++ python/tests/test_upload.py | 75 +++++++++++ 7 files changed, 762 insertions(+) create mode 100644 python/tests/test_upload.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 37865cc..eca6481 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,31 @@ All notable changes to this project are documented here. +## [Unreleased] + +### Added + +- **File upload from the Mac into the agent session (`upload` tool + + `/aiui:upload` slash-command, #146).** The first aiui data flow that + runs *Mac → agent-host* instead of the other way around. Calling + `upload` opens a **native file picker** on the user's Mac (via the + Tauri dialog plugin); the chosen file streams back over the existing + authenticated `:7777` channel — the reverse direction of `POST /media` + — and is written to `target_dir/` on the host the agent runs + on (the remote for an SSH session). It replaces the "please `scp` me + that file" round-trip. `target_dir` is optional (defaults to the + bridge's cwd); the filename comes from the selection, so the write is + deterministic with no temp/staging path. Existing files are never + overwritten — a name clash returns an error rather than clobbering. + Robust error paths (`{status:"error", error}`) cover a cancelled + picker, an unreadable file, the 512 MB size cap (413), and a + missing/unwritable target directory. The tool blocks until the user + picks or dismisses the picker, with the usual `notifications/progress` + keepalive. Implemented consistently in the native Rust MCP server + (`companion/src-tauri/src/{mcp,http}.rs`) and the Python bridge + (`aiui-mcp`, used by remote SSH hosts via `uvx`), plus a new + `/aiui:upload` prompt and an `INSTRUCTIONS` trigger on both. + ## [0.8.3] — 2026-07-29 ### Added diff --git a/companion/src-tauri/src/http.rs b/companion/src-tauri/src/http.rs index 4b10bb5..9da3ea7 100644 --- a/companion/src-tauri/src/http.rs +++ b/companion/src-tauri/src/http.rs @@ -16,6 +16,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use crate::logging::trace; use tauri::{AppHandle, Emitter, Manager}; +use tauri_plugin_dialog::DialogExt; use tauri_plugin_updater::UpdaterExt; /// How long `/health` waits for a `ui:ping` round-trip from the frontend @@ -208,6 +209,12 @@ pub async fn serve( post(media_upload) .layer(DefaultBodyLimit::max(crate::media::MEDIA_FILE_CAP as usize)), ) + // Inbound file transfer (#146): the bridge asks the Mac to open a + // native file picker; on selection the picked file's bytes stream + // back over the same :7777 channel (the reverse direction of + // `POST /media`). This is the "get a Mac file into the agent + // session" path. + .route("/upload", post(upload_pick)) // Capability-URL playback: unauthenticated (filename is a UUID), // range-capable for video seeking via tower-http's ServeDir. .nest_service( @@ -370,6 +377,161 @@ async fn media_upload( .into_response() } +/// Largest single inbound file accepted through `POST /upload` (#146). The +/// picked file is buffered in memory once before it streams back to the +/// bridge, so this caps a runaway pick (a multi-GB file the user selected by +/// mistake) rather than letting it exhaust RAM. Mirrors the outbound +/// `media::MEDIA_FILE_CAP` so both directions share one ceiling. +const UPLOAD_FILE_CAP: u64 = 512 * 1024 * 1024; + +/// HTTP header carrying the picked file's base name (percent-encoded, RFC +/// 3986) on a successful `POST /upload`. The bridge decodes it, sanitises it +/// to a base name, and writes `target_dir/`. +const UPLOAD_FILENAME_HEADER: &str = "x-aiui-filename"; + +/// Percent-encode a filename for transport in an ASCII HTTP header. Encodes +/// every byte that isn't an RFC-3986 unreserved char, so UTF-8 names, spaces, +/// and control bytes all survive the round-trip and the header stays valid. +fn pct_encode_filename(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for b in s.bytes() { + match b { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => { + out.push(b as char) + } + _ => { + out.push('%'); + out.push_str(&format!("{b:02X}")); + } + } + } + out +} + +/// `POST /upload` — open a native file picker on the Mac and stream the picked +/// file's bytes back to the caller (#146). This is the reverse of +/// `POST /media`: bytes flow Mac → agent-host, over the same authenticated +/// :7777 channel (loopback locally, the SSH reverse-tunnel remotely). +/// +/// Responses: +/// - `200 OK` — body is the raw file bytes; `x-aiui-filename` header carries +/// the percent-encoded base name. `Content-Length` gives the byte count. +/// A legitimately empty file is still a 200 with a `0`-length body and the +/// filename header present — distinct from the cancel case below. +/// - `204 No Content` — the user dismissed the picker without choosing a file. +/// No body, no filename header. +/// - `413 Payload Too Large` — the picked file exceeds `UPLOAD_FILE_CAP`. +/// - `500` — the file could not be read, or the picker failed unexpectedly. +/// +/// The handler blocks until the user picks or cancels; the caller's MCP +/// progress notifications keep the client alive meanwhile (same keepalive the +/// dialog tools use). Authenticated like every mutating endpoint. +async fn upload_pick( + State(state): State, + headers: HeaderMap, +) -> impl IntoResponse { + if !auth_ok(&headers, &state.cfg.token) { + return (StatusCode::UNAUTHORIZED, "unauthorized").into_response(); + } + + // The native picker is callback-based; bridge it to async via a oneshot. + // `pick_file` dispatches to the main thread internally (rfd requirement on + // macOS), so calling it from this tokio task is safe. + let (tx, rx) = tokio::sync::oneshot::channel(); + state.app.dialog().file().pick_file(move |picked| { + let _ = tx.send(picked); + }); + + let picked = match rx.await { + Ok(p) => p, + Err(_) => { + trace("upload_pick: picker channel dropped"); + return (StatusCode::INTERNAL_SERVER_ERROR, "picker closed unexpectedly") + .into_response(); + } + }; + + let Some(file_path) = picked else { + trace("upload_pick: user cancelled the picker"); + return StatusCode::NO_CONTENT.into_response(); + }; + + let path = match file_path.into_path() { + Ok(p) => p, + Err(e) => { + trace(&format!("upload_pick: non-filesystem selection: {e}")); + return (StatusCode::INTERNAL_SERVER_ERROR, "selection is not a local file") + .into_response(); + } + }; + + let meta = match tokio::fs::metadata(&path).await { + Ok(m) => m, + Err(e) => { + trace(&format!("upload_pick: stat failed: {e}")); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("cannot read selected file: {e}"), + ) + .into_response(); + } + }; + if !meta.is_file() { + return (StatusCode::INTERNAL_SERVER_ERROR, "selection is not a regular file") + .into_response(); + } + if meta.len() > UPLOAD_FILE_CAP { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + format!( + "selected file is {} bytes (max {})", + meta.len(), + UPLOAD_FILE_CAP + ), + ) + .into_response(); + } + + let filename = path + .file_name() + .and_then(|n| n.to_str()) + .map(String::from) + .unwrap_or_else(|| "upload.bin".to_string()); + + let bytes = match tokio::fs::read(&path).await { + Ok(b) => b, + Err(e) => { + trace(&format!("upload_pick: read failed: {e}")); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("cannot read selected file: {e}"), + ) + .into_response(); + } + }; + + trace(&format!( + "upload_pick: delivering '{}' ({} bytes)", + filename, + bytes.len() + )); + ( + StatusCode::OK, + [ + ( + axum::http::header::CONTENT_TYPE, + "application/octet-stream".to_string(), + ), + ( + axum::http::HeaderName::from_static(UPLOAD_FILENAME_HEADER), + pct_encode_filename(&filename), + ), + ], + bytes, + ) + .into_response() +} + /// Composite health check. Probes the WebView event loop with a `ui:ping` /// round-trip, reads live counters from the dialog registry and lifetime /// tracker, and reports `ready` only when all three are healthy. Computed @@ -1237,3 +1399,57 @@ mod async_render_tests { assert!(slots.is_empty()); } } + +#[cfg(test)] +mod upload_tests { + use super::pct_encode_filename; + + #[test] + fn plain_ascii_name_is_unchanged() { + assert_eq!(pct_encode_filename("report.pdf"), "report.pdf"); + assert_eq!(pct_encode_filename("a-b_c.1~2"), "a-b_c.1~2"); + } + + #[test] + fn spaces_and_specials_are_encoded() { + assert_eq!(pct_encode_filename("my file.txt"), "my%20file.txt"); + assert_eq!(pct_encode_filename("a/b"), "a%2Fb"); + assert_eq!(pct_encode_filename("a+b&c"), "a%2Bb%26c"); + } + + #[test] + fn utf8_survives_roundtrip() { + // Umlaut + emoji: every non-unreserved byte becomes %XX, so the + // header stays pure ASCII and the bridge can reconstruct the name. + let encoded = pct_encode_filename("Prüfung.md"); + assert!(encoded.is_ascii()); + assert!(encoded.starts_with("Pr%")); + assert!(encoded.ends_with("fung.md")); + // Decoding the percent-escapes yields the original UTF-8 bytes. + let decoded = pct_decode(&encoded); + assert_eq!(decoded, "Prüfung.md".as_bytes()); + } + + /// Reference percent-decoder mirroring what the bridges do, used to prove + /// the encoder round-trips. Not used in production Rust (the encoder lives + /// on the companion; the bridges own decoding). + fn pct_decode(s: &str) -> Vec { + let mut out = Vec::with_capacity(s.len()); + let bytes = s.as_bytes(); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'%' && i + 2 < bytes.len() { + let hi = (bytes[i + 1] as char).to_digit(16); + let lo = (bytes[i + 2] as char).to_digit(16); + if let (Some(h), Some(l)) = (hi, lo) { + out.push((h * 16 + l) as u8); + i += 3; + continue; + } + } + out.push(bytes[i]); + i += 1; + } + out + } +} diff --git a/companion/src-tauri/src/mcp.rs b/companion/src-tauri/src/mcp.rs index 7f1df10..4ec32f8 100644 --- a/companion/src-tauri/src/mcp.rs +++ b/companion/src-tauri/src/mcp.rs @@ -44,6 +44,9 @@ instead of asking via chat. Default behaviour for this session: - Pick-one-of-N options where context per option matters → call `ask`. - Multiple related inputs, secret, date, slider, sortable order, \ table-row triage, image confirm/grid → call `form`. +- User wants to hand you a file from their Mac (`/aiui:upload`, \ + \"take this file\", \"upload …\") → call `upload` with the target \ + directory on your host; don't ask them to `scp` it. - Pure information the user only reads → keep it in chat. Type `/aiui:teach` for the full widget catalog when composing a \ @@ -96,6 +99,11 @@ Report the outcome in one line: \"aiui ok — you clicked '{label}'\" if the \ window opened and returned, or the underlying error if it didn't. "; +const UPLOAD_PROMPT: &str = "\ +Call the `upload` tool to let me hand you a file from my Mac. \ +Use my current working directory as the target unless I say otherwise. +"; + const REMOTES_PROMPT: &str = "\ Show the user a quick rundown of their registered aiui remotes — same as \ the Settings window's \"Eingerichtete Remote-Hosts\" section, but in chat. \ @@ -408,6 +416,17 @@ fn tools_list() -> Value { } } }, + { + "name": "upload", + "description": "Pull a file FROM the user's Mac INTO this agent session. Calling this opens a native file picker on the user's Mac; the file they choose is streamed back over aiui's channel and written to `target_dir` on YOUR host (the machine you run on — the remote for an SSH session). This is the counterpart to the user having to `scp` a file over: reach for it whenever the user says \"take this file\", \"upload …\", \"here's the file/screenshot/PDF\", or triggers `/aiui:upload`. **`target_dir` is optional and you should almost always pass it:** set it to the directory the file belongs in given the conversation — usually your current working directory or the active project dir. Do NOT ask the user where to put it or which file to pick; just call the tool and let them choose the file in the native dialog. If you have no context at all, omit `target_dir` (defaults to your process's cwd) or ask in one short sentence. The filename comes from the user's selection — the file lands at `target_dir/`, a deterministic path, no temp/staging dir. **Existing files are never overwritten:** if `target_dir/` already exists the call returns an error rather than clobbering — pick a different `target_dir` or move the old file first. Returns `{status: \"ok\", path, filename, bytes}` on success, or `{status: \"error\", error}` on any failure — user cancelled the picker, file unreadable, file too large (512 MB cap), or target directory missing/not writable. Report the result briefly; on `ok` mention the path the file landed at. **This tool blocks until the user picks a file or dismisses the picker. Response can take a while — do not assume aiui is broken; the user is choosing a file. Progress notifications fire every ~10 s while waiting.**", + "inputSchema": { + "type": "object", + "properties": { + "target_dir": { "type": "string", "description": "Absolute or `~/`-rooted directory ON YOUR HOST where the picked file is written as `/`. Optional; defaults to your process's current working directory. Relative paths are rejected (no stable cwd contract). The directory must already exist and be writable." }, + "session": { "type": "string", "description": "Optional short human label for the session this upload belongs to (project/task name), shown in aiui's window chrome." } + } + } + }, { "name": "aiui_health", "description": "Reachability check against the local aiui companion. Returns version + ready flag if the companion is running and responding.", @@ -441,6 +460,14 @@ fn tools_list() -> Value { /// enough to surface the diagnostic message. const COLDSTART_WAIT: std::time::Duration = std::time::Duration::from_secs(30); +/// How long the `upload` tool waits on `POST /upload` before giving up. Unlike +/// a dialog render — which the async poll loop keeps off a single held +/// connection — the picker + byte transfer runs on one request, and the user +/// may browse their filesystem for a while before choosing. A generous ceiling +/// (well above any think-time a file picker realistically takes) keeps the call +/// alive; the MCP progress notifications reassure the client meanwhile. #146. +const UPLOAD_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(900); + /// Poll `/ping` until the HTTP server answers, or `COLDSTART_WAIT` elapses. /// `/ping` is unauthenticated and cheap, returning `pong` in plain text — /// any 2xx means aiui is bound and serving. Returns `true` once reachable, @@ -667,6 +694,8 @@ async fn tools_call( format_dialog_result, ), + "upload" => Ok(do_upload(&args, cfg, http).await), + "aiui_health" => get_json(http, cfg, "/health").await.map(value_to_tool_text), "version" => get_json(http, cfg, "/version").await.map(value_to_tool_text), "update" => post_empty(http, cfg, "/update") @@ -757,6 +786,169 @@ async fn upload_media( .ok_or_else(|| "/media response missing url".to_string()) } +/// Percent-decode a filename the companion sent in the `x-aiui-filename` +/// header (RFC-3986, produced by `http::pct_encode_filename`). Returns the raw +/// bytes; invalid `%` sequences are passed through literally rather than +/// erroring, so a mangled header degrades to a slightly-odd name, never a lost +/// upload. +fn pct_decode_bytes(s: &str) -> Vec { + let bytes = s.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'%' && i + 2 < bytes.len() { + let hi = (bytes[i + 1] as char).to_digit(16); + let lo = (bytes[i + 2] as char).to_digit(16); + if let (Some(h), Some(l)) = (hi, lo) { + out.push((h * 16 + l) as u8); + i += 3; + continue; + } + } + out.push(bytes[i]); + i += 1; + } + out +} + +/// Reduce a filename to a safe base name: strip any directory components (the +/// companion sends a base name already, but a hostile/odd selection must never +/// escape `target_dir`), reject `.`/`..`/empty. Returns `None` if nothing safe +/// remains. +fn safe_base_name(raw: &str) -> Option { + let base = std::path::Path::new(raw) + .file_name() + .and_then(|n| n.to_str())? + .trim(); + if base.is_empty() || base == "." || base == ".." { + return None; + } + Some(base.to_string()) +} + +/// Expand a `~/`-rooted or absolute directory path. Relative paths return +/// `None` — there is no stable cwd contract to resolve them against, so we +/// treat "relative" as a caller error rather than guessing (the cwd default is +/// applied by the caller *before* this, only when `target_dir` is absent). +fn expand_dir(raw: &str) -> Option { + if let Some(rest) = raw.strip_prefix("~/") { + return dirs::home_dir().map(|h| h.join(rest)); + } + if raw == "~" { + return dirs::home_dir(); + } + let p = std::path::PathBuf::from(raw); + if p.is_absolute() { + Some(p) + } else { + None + } +} + +fn upload_error(msg: impl Into) -> Value { + value_to_tool_text(json!({ "status": "error", "error": msg.into() })) +} + +/// Implements the `upload` tool (#146): ask the companion to open a native file +/// picker on the Mac, receive the picked file's bytes over the :7777 channel, +/// and write them to `target_dir/` on THIS host. Returns the MCP +/// tool-result shape wrapping `{status, path, filename, bytes}` (ok) or +/// `{status, error}` (any failure, including a user-cancelled picker). +async fn do_upload(args: &Value, cfg: &AppConfig, http: &reqwest::Client) -> Value { + // Resolve the destination directory up front so a bad `target_dir` fails + // before we even open the picker (nothing more annoying than picking a file + // only to be told the target was invalid). + let target_dir = match args.get("target_dir").and_then(|v| v.as_str()) { + Some(s) if !s.trim().is_empty() => match expand_dir(s.trim()) { + Some(p) => p, + None => { + return upload_error(format!( + "target_dir must be an absolute or ~/-rooted path, got '{s}'" + )) + } + }, + // No target_dir → default to this process's cwd. + _ => match std::env::current_dir() { + Ok(p) => p, + Err(e) => return upload_error(format!("no target_dir given and cwd unavailable: {e}")), + }, + }; + if !target_dir.is_dir() { + return upload_error(format!( + "target directory does not exist: {}", + target_dir.display() + )); + } + + let token = match load_token(cfg) { + Ok(t) => t, + Err(e) => return upload_error(e), + }; + let url = format!("{}/upload", base_url(cfg)); + let resp = match http + .post(&url) + .bearer_auth(&token) + .timeout(UPLOAD_TIMEOUT) + .send() + .await + { + Ok(r) => r, + Err(e) => return upload_error(format!("POST /upload: {e}")), + }; + + match resp.status() { + reqwest::StatusCode::NO_CONTENT => { + return upload_error("upload cancelled — no file was selected"); + } + reqwest::StatusCode::PAYLOAD_TOO_LARGE => { + let detail = resp.text().await.unwrap_or_default(); + return upload_error(format!("selected file too large: {detail}")); + } + s if !s.is_success() => { + let detail = resp.text().await.unwrap_or_default(); + return upload_error(format!("companion /upload failed ({s}): {detail}")); + } + _ => {} + } + + // Filename travels in a header (percent-encoded); the body is the raw bytes. + let filename_raw = resp + .headers() + .get("x-aiui-filename") + .and_then(|v| v.to_str().ok()) + .map(|s| String::from_utf8_lossy(&pct_decode_bytes(s)).into_owned()) + .unwrap_or_default(); + let filename = match safe_base_name(&filename_raw) { + Some(f) => f, + None => "upload.bin".to_string(), + }; + + let bytes = match resp.bytes().await { + Ok(b) => b, + Err(e) => return upload_error(format!("reading uploaded bytes: {e}")), + }; + + let dest = target_dir.join(&filename); + // Never clobber: a deterministic path is the point, but silently + // overwriting the user's existing file is not. Fail loudly instead. + if dest.exists() { + return upload_error(format!( + "target already exists, not overwriting: {}", + dest.display() + )); + } + if let Err(e) = crate::fsutil::atomic_write(&dest, &bytes) { + return upload_error(format!("writing {}: {e}", dest.display())); + } + + value_to_tool_text(json!({ + "status": "ok", + "path": dest.display().to_string(), + "filename": filename, + "bytes": bytes.len(), + })) +} + /// Per-call dialog rendering can fail in two structurally different /// ways. v0.4.36 splits them so the tool dispatcher can convert /// `Busy` into a structured tool result (with retry-vs-tell-user @@ -1099,6 +1291,11 @@ fn prompts_list() -> Value { "name": "remotes", "description": "List the user's registered aiui remotes in chat (same set the Settings window shows).", "arguments": [] + }, + { + "name": "upload", + "description": "Hand a file from your Mac to the agent session — opens a native file picker and writes the chosen file to the agent host.", + "arguments": [] } ]) } @@ -1116,6 +1313,7 @@ fn prompts_get(params: Value) -> Result { "health" => HEALTH_PROMPT, "test-dialog" => TEST_DIALOG_PROMPT, "remotes" => REMOTES_PROMPT, + "upload" => UPLOAD_PROMPT, _ => { return Err(RpcError { code: -32602, diff --git a/docs/skill.md b/docs/skill.md index 16f994d..85c44bb 100644 --- a/docs/skill.md +++ b/docs/skill.md @@ -269,6 +269,42 @@ batch". Use `confirm`+`image` for a single yes/no sign-off, and `ask`+`thumbnail` / `image_grid` when the task is *picking* among candidates rather than judging each one. +## File upload: `upload` (Mac → your host) + +Every other aiui data flow goes *your host → Mac* (dialog specs down, +image bytes inlined). `upload` is the one that reverses it: it pulls a +file **from the user's Mac into your session**, over the same +authenticated `:7777` channel the dialogs use (loopback locally, the SSH +reverse-tunnel remotely). It's the native replacement for "please `scp` +me that file". + +Call `upload` whenever the user wants to hand you a local file — "take +this file", "here's the screenshot/PDF/CSV", or the `/aiui:upload` +slash-command. It opens a **native file picker on the Mac**; the file the +user chooses is streamed back and written to `target_dir/` on +**your** host (the remote for an SSH session). + +- **Pass `target_dir`** — an absolute or `~/`-rooted directory on your + host, chosen from context (usually your cwd or the active project dir). + Omit it only when you genuinely have no context (defaults to your + process's cwd). Relative paths are rejected; the directory must already + exist and be writable. +- **Don't ask which file or where** — the user picks the file in the + native dialog, and you infer the target. One tool call, no back-and-forth. +- **Deterministic destination:** the filename comes from the selection, + so the file lands at exactly `target_dir/` — no temp/staging + path. Existing files are **never overwritten**; a name clash returns an + error instead of clobbering (pick another `target_dir` or move the old + file first). +- **Result:** `{status: "ok", path, filename, bytes}` on success, or + `{status: "error", error}` for a cancelled picker, an unreadable file, a + file over the 512 MB cap, or a missing/unwritable target dir. Report it + briefly; on `ok`, name the path the file landed at. + +Blocks until the user picks or dismisses the picker, exactly like the +dialog tools — progress notifications fire every ~10 s while you wait, so +a slow response just means the user is browsing, not that aiui broke. + ## Starting window size: `size` / `width` / `height` `form` and `gallery` accept an optional **`size`** hint — `"s"`, `"m"`, or diff --git a/python/src/aiui_mcp/server.py b/python/src/aiui_mcp/server.py index f61985d..c7488d7 100644 --- a/python/src/aiui_mcp/server.py +++ b/python/src/aiui_mcp/server.py @@ -27,6 +27,7 @@ import sys import tempfile import time +import urllib.parse from datetime import datetime, timezone from pathlib import Path from typing import Any @@ -106,6 +107,14 @@ def _default_token_path() -> str: # before we time out, letting us re-poll cleanly. ASYNC_POLL_TIMEOUT_S = 40.0 +# Timeout for the `upload` tool's held `POST /upload` (#146). The picker + byte +# transfer runs on one request and the user may browse their filesystem before +# choosing, so this is deliberately generous — far beyond any realistic +# file-picker think-time. A periodic progress notification keeps the MCP client +# reassured while the call is held. +UPLOAD_TIMEOUT_S = float(os.environ.get("AIUI_UPLOAD_TIMEOUT_S", "900")) +UPLOAD_FILE_CAP = 512 * 1024 * 1024 # mirrors the companion's cap + _INSTRUCTIONS = """\ aiui is connected — you can render native dialogs on the user's Mac \ instead of asking via chat. Default behaviour for this session: @@ -115,6 +124,9 @@ def _default_token_path() -> str: - Pick-one-of-N options where context per option matters → call `ask`. - Multiple related inputs, secret, date, slider, sortable order, \ table-row triage, image confirm/grid → call `form`. +- User wants to hand you a file from their Mac (`/aiui:upload`, \ + "take this file", "upload …") → call `upload` with the target \ + directory on your host; don't ask them to `scp` it. - Pure information the user only reads → keep it in chat. Type `/aiui:teach` for the full widget catalog when composing a \ @@ -536,6 +548,73 @@ def _apply_target_writes(spec: dict[str, Any], data: dict[str, Any]) -> None: values[name] = {"value": v, **outcome} +def _upload_safe_base_name(raw: str) -> str | None: + """Reduce a filename to a safe base name (#146): strip any directory + components so a selection can never escape the target dir, reject + empty / `.` / `..`. Mirrors `safe_base_name` in the Rust bridge. + """ + base = os.path.basename(raw).strip() + if not base or base in (".", ".."): + return None + return base + + +def _upload_expand_dir(raw: str) -> Path | None: + """Expand a `~/`-rooted or absolute target directory. Returns None for a + relative path — there is no stable cwd contract to resolve it against, so + it's treated as a caller error (the cwd default is applied by the caller + only when `target_dir` is absent). Mirrors `expand_dir` in the Rust bridge. + """ + if raw.startswith("~"): + return Path(raw).expanduser() + p = Path(raw) + return p if p.is_absolute() else None + + +def _upload_write(dest_dir: Path, filename: str, data: bytes) -> dict[str, Any]: + """Atomically write the uploaded bytes to `dest_dir/` on THIS host, + never overwriting an existing file. Mirrors `do_upload`'s write half in the + Rust bridge. Returns the `{status, …}` payload. + """ + dest = dest_dir / filename + if dest.exists(): + return {"status": "error", "error": f"target already exists, not overwriting: {dest}"} + try: + fd, tmp = tempfile.mkstemp(prefix=".aiui-upload-", dir=str(dest_dir)) + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + os.replace(tmp, dest) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + except OSError as e: + return {"status": "error", "error": f"writing {dest}: {e}"} + return {"status": "ok", "path": str(dest), "filename": filename, "bytes": len(data)} + + +async def _upload_heartbeat(ctx: Context) -> None: + """Emit a progress notification every ~10 s while `POST /upload` is held, so + the MCP client knows the tool is alive while the user browses the picker. + Cancelled by the caller once the POST returns. Best-effort — a missing + progressToken or any reporting hiccup must never break the upload. + """ + iteration = 0 + try: + while True: + await asyncio.sleep(10) + iteration += 1 + try: + await ctx.report_progress(progress=float(iteration), total=None) + except Exception as e: # noqa: BLE001 + log.debug("upload progress skipped: %s", _explain_exc(e)) + except asyncio.CancelledError: + pass + + async def _wait_for_aiui() -> None: """Poll the unauthenticated `/ping` until the companion answers or `COLDSTART_WAIT_S` elapses (Step 3, parity with the Rust bridge). @@ -986,6 +1065,105 @@ async def gallery( return _format_result(await _post_render(spec, ctx, session)) +@mcp.tool() +async def upload( + target_dir: str | None = None, + session: str | None = None, + ctx: Context | None = None, +) -> dict[str, Any]: + """Pull a file FROM the user's Mac INTO this agent session. + + Calling this opens a native file picker on the user's Mac; the file they + choose is streamed back over aiui's channel and written to `target_dir` on + YOUR host (the machine you run on — the remote for an SSH session). This is + the counterpart to the user having to `scp` a file over: reach for it + whenever the user says "take this file", "upload …", "here's the + file/screenshot/PDF", or triggers `/aiui:upload`. + + WHEN TO USE: the user wants to give you a local Mac file. Do NOT ask them + which file — they pick it in the native dialog. Do NOT ask where to put it; + infer `target_dir` from the conversation (usually your cwd or the active + project dir) and pass it. + + BEHAVIOUR: + - `target_dir` is optional but you should almost always pass it — an + absolute or `~/`-rooted directory ON YOUR HOST. Omit it only when you + have no context (defaults to your process's cwd). Relative paths are + rejected. The directory must already exist and be writable. + - The filename comes from the user's selection; the file lands at + `target_dir/` — a deterministic path, no temp/staging dir. + - Existing files are never overwritten: if `target_dir/` already + exists the call errors instead of clobbering. Pick a different + `target_dir` or move the old file first. + - Blocks until the user picks a file or dismisses the picker. Progress + notifications fire every ~10 s meanwhile — a slow response just means the + user is choosing a file, not that aiui is broken. + + Returns `{status: "ok", path, filename, bytes}` on success, or + `{status: "error", error}` on any failure (user cancelled, file unreadable, + file too large — 512 MB cap, target directory missing/not writable). Report + briefly; on `ok`, mention the path the file landed at. + + Args: + target_dir: Absolute or `~/`-rooted directory on your host where the + picked file is written as `/`. Defaults to + your process's cwd. + session: Short human label for this session, shown in aiui's window + chrome so parallel dialogs stay distinguishable. + """ + # Resolve the destination up front so a bad target_dir fails before the + # picker even opens (nothing worse than picking a file only to be rejected). + if target_dir is not None and target_dir.strip(): + dest_dir = _upload_expand_dir(target_dir.strip()) + if dest_dir is None: + return { + "status": "error", + "error": f"target_dir must be an absolute or ~/-rooted path, got '{target_dir}'", + } + else: + dest_dir = Path.cwd() + if not dest_dir.is_dir(): + return {"status": "error", "error": f"target directory does not exist: {dest_dir}"} + + await _wait_for_aiui() + await _preflight() + + heartbeat = asyncio.create_task(_upload_heartbeat(ctx)) if ctx is not None else None + try: + async with httpx.AsyncClient(timeout=UPLOAD_TIMEOUT_S) as client: + r = await client.post( + f"{ENDPOINT}/upload", + headers={"Authorization": f"Bearer {_token()}"}, + ) + except httpx.HTTPError as e: + return {"status": "error", "error": f"POST /upload failed: {_explain_exc(e)}"} + finally: + if heartbeat is not None: + heartbeat.cancel() + + if r.status_code == 204: + return {"status": "error", "error": "upload cancelled — no file was selected"} + if r.status_code == 413: + return {"status": "error", "error": f"selected file too large: {r.text[:200]}"} + if r.status_code != 200: + return { + "status": "error", + "error": f"companion /upload failed ({r.status_code}): {r.text[:200]}", + } + + # Filename travels in a percent-encoded header; the body is the raw bytes. + raw_header = r.headers.get("x-aiui-filename", "") + decoded = urllib.parse.unquote_to_bytes(raw_header).decode("utf-8", "replace") + filename = _upload_safe_base_name(decoded) or "upload.bin" + data = r.content + if len(data) > UPLOAD_FILE_CAP: + return {"status": "error", "error": f"uploaded file exceeds cap: {len(data)} bytes"} + + result = _upload_write(dest_dir, filename, data) + log.info("upload ← filename=%s bytes=%d status=%s", filename, len(data), result.get("status")) + return result + + @mcp.prompt(name="teach") def teach_prompt() -> str: """Brief the agent on aiui. Loads the full widget catalog, design @@ -1102,6 +1280,20 @@ def remotes_prompt() -> str: return _REMOTES_PROMPT +_UPLOAD_PROMPT = """\ +Call the `upload` tool to let me hand you a file from my Mac. \ +Use my current working directory as the target unless I say otherwise. +""" + + +@mcp.prompt(name="upload") +def upload_prompt() -> str: + """Hand a file from the Mac to the agent session — opens a native file + picker and writes the chosen file to the agent host. Surfaces as + `/aiui:upload` in Claude Code.""" + return _UPLOAD_PROMPT + + @mcp.tool() async def aiui_health() -> dict[str, Any]: """Reachability + token check against the aiui companion. diff --git a/python/src/aiui_mcp/skill.md b/python/src/aiui_mcp/skill.md index 5483d54..cc422d9 100644 --- a/python/src/aiui_mcp/skill.md +++ b/python/src/aiui_mcp/skill.md @@ -198,6 +198,26 @@ hosting it anywhere. Result: `{cancelled, decisions: {"": {decision, comment?}}}`. Only touched items appear — an untouched item means "no verdict", not a default. +## File upload: `upload` (Mac → your host) + +The one reversed flow: `upload` pulls a file **from the user's Mac into +your session** over the same `:7777` channel (loopback locally, the SSH +reverse-tunnel remotely) — the native replacement for "please `scp` me +that file". Call it whenever the user wants to hand you a local file +("take this file", "here's the screenshot/PDF", `/aiui:upload`). It opens +a **native file picker on the Mac**; the chosen file is written to +`target_dir/` on **your** host. + +- Pass `target_dir` (absolute or `~/`-rooted, on your host) inferred from + context — usually your cwd/project dir. Omit only with no context + (defaults to cwd). Relative paths rejected; dir must exist and be writable. +- Don't ask which file or where — the user picks in the dialog, you infer + the target. Deterministic: lands at exactly `target_dir/`, no + staging path. Existing files are never overwritten (a clash errors). +- Returns `{status:"ok", path, filename, bytes}` or `{status:"error", error}` + (cancelled picker, unreadable file, >512 MB cap, missing/unwritable dir). + Blocks until pick/cancel; progress fires every ~10 s meanwhile. + ## Starting window size: `size` `form` and `gallery` take an optional `size` hint — `"s"`, `"m"`, `"l"` — diff --git a/python/tests/test_upload.py b/python/tests/test_upload.py new file mode 100644 index 0000000..d1bec00 --- /dev/null +++ b/python/tests/test_upload.py @@ -0,0 +1,75 @@ +"""Bridge-side upload tests (#146). + +Cover the pure helpers behind the `upload` tool — filename sanitisation, +target-dir expansion, and the no-clobber atomic write — without needing a +running companion. Mirrors the Rust bridge's `do_upload` helpers. +""" +from __future__ import annotations + +import urllib.parse +from pathlib import Path + +from aiui_mcp.server import ( + _upload_expand_dir, + _upload_safe_base_name, + _upload_write, +) + + +def test_safe_base_name_strips_directories() -> None: + assert _upload_safe_base_name("report.pdf") == "report.pdf" + assert _upload_safe_base_name("/Users/me/Downloads/report.pdf") == "report.pdf" + assert _upload_safe_base_name("../../etc/passwd") == "passwd" + assert _upload_safe_base_name(" spaced.txt ") == "spaced.txt" + + +def test_safe_base_name_rejects_empty_and_dots() -> None: + assert _upload_safe_base_name("") is None + assert _upload_safe_base_name(".") is None + assert _upload_safe_base_name("..") is None + assert _upload_safe_base_name("/") is None + + +def test_filename_header_roundtrip() -> None: + # The companion percent-encodes UTF-8 filenames into an ASCII header; + # the bridge decodes with unquote_to_bytes → utf-8. Prove a name with an + # umlaut and a space survives. + original = "Prüfung final.md" + encoded = urllib.parse.quote(original, safe="") + assert encoded.isascii() + decoded = urllib.parse.unquote_to_bytes(encoded).decode("utf-8", "replace") + assert _upload_safe_base_name(decoded) == original + + +def test_expand_dir_absolute_and_tilde() -> None: + assert _upload_expand_dir("/tmp/x") == Path("/tmp/x") + assert _upload_expand_dir("~/Downloads") == Path.home() / "Downloads" + # Relative paths are rejected — no stable cwd contract. + assert _upload_expand_dir("relative/dir") is None + assert _upload_expand_dir("./here") is None + + +def test_write_creates_file(tmp_path: Path) -> None: + out = _upload_write(tmp_path, "hello.txt", b"hi there") + assert out["status"] == "ok" + assert out["filename"] == "hello.txt" + assert out["bytes"] == 8 + dest = tmp_path / "hello.txt" + assert dest.read_bytes() == b"hi there" + assert out["path"] == str(dest) + + +def test_write_refuses_to_clobber(tmp_path: Path) -> None: + dest = tmp_path / "existing.txt" + dest.write_text("original") + out = _upload_write(tmp_path, "existing.txt", b"new content") + assert out["status"] == "error" + assert "already exists" in out["error"] + # The existing file is untouched. + assert dest.read_text() == "original" + + +def test_write_leaves_no_temp_files(tmp_path: Path) -> None: + _upload_write(tmp_path, "a.bin", b"\x00\x01\x02") + names = sorted(p.name for p in tmp_path.iterdir()) + assert names == ["a.bin"], f"stray temp files: {names}"