diff --git a/src/commands/mod.rs b/src/commands/mod.rs index e49b868..80ac6b5 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -10,6 +10,7 @@ mod lab; mod login; mod model; mod new; +mod pair; mod python; mod remove; mod shell; @@ -33,6 +34,7 @@ use lab::{lab, Lab}; use login::{login, Login}; use model::{model, Model}; use new::{new, New}; +use pair::{pair, Pair}; use python::{python, Python}; use remove::{remove, Remove}; use shell::{shell, Shell}; @@ -77,6 +79,7 @@ pub enum Commands { #[command(subcommand)] args: Job, }, + Pair(Pair), Login(Login), Auth { #[command(subcommand)] @@ -110,6 +113,7 @@ impl Cli { Commands::Dataset { args } => dataset(args, global).await, Commands::Model { args } => model(args, global).await, Commands::Job { args } => job(args, global).await, + Commands::Pair(args) => pair(args, global).await, Commands::Login(args) => login(args, global).await, Commands::Auth { args } => auth(args, global).await, Commands::Python(args) => python(args, global).await, diff --git a/src/commands/pair/agent.rs b/src/commands/pair/agent.rs new file mode 100644 index 0000000..c132368 --- /dev/null +++ b/src/commands/pair/agent.rs @@ -0,0 +1,249 @@ +use std::path::PathBuf; + +use crate::error::{self, Result}; + +const SKILL_NAME: &str = "marimo-pair"; +const SKILL_FILE: &str = "SKILL.md"; + +pub const INSTALL_HINT: &str = "npx skills add marimo-team/marimo-pair"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Agent { + Claude, + Codex, + Opencode, +} + +/// What an agent is missing, if anything. Both halves are needed to pair: the +/// CLI to launch, and the skill the prompt tells it to use. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Availability { + pub binary: bool, + pub skill: bool, +} + +impl Availability { + fn is_ready(&self) -> bool { + self.binary && self.skill + } +} + +impl Agent { + /// The order auto-detection tries. + pub const ALL: [Agent; 3] = [Agent::Claude, Agent::Codex, Agent::Opencode]; + + pub fn display_name(&self) -> &'static str { + match self { + Agent::Claude => "Claude Code", + Agent::Codex => "Codex", + Agent::Opencode => "opencode", + } + } + + pub fn binary(&self) -> &'static str { + match self { + Agent::Claude => "claude", + Agent::Codex => "codex", + Agent::Opencode => "opencode", + } + } + + /// The prompt goes in as a single argument. It carries no token — only the + /// path to one — so it is safe in `ps` and in shell history. + pub fn command(&self, prompt: &str) -> tokio::process::Command { + let mut command = tokio::process::Command::new(self.binary()); + match self { + Agent::Claude | Agent::Codex => command.arg(prompt), + Agent::Opencode => command.args(["--prompt", prompt]), + }; + command + } + + pub fn availability(&self) -> Availability { + Availability { + binary: which::which(self.binary()).is_ok(), + skill: self.has_skill(), + } + } + + /// Directories that may hold `/marimo-pair/SKILL.md`, mirroring the + /// layouts `marimo pair prompt` looks in. + fn skill_dirs(&self) -> Vec { + let home = dirs::home_dir(); + let cwd = std::env::current_dir().ok(); + let roots = |sub: &[&str]| -> Vec { + let mut dirs = Vec::new(); + for root in [home.as_ref(), cwd.as_ref()].into_iter().flatten() { + dirs.push(sub.iter().fold(root.clone(), |path, part| path.join(part))); + } + dirs + }; + match self { + Agent::Claude => [ + roots(&[".claude", "skills"]), + roots(&[".claude", "plugins"]), + roots(&[".claude", "plugins", "marketplaces"]), + ] + .concat(), + Agent::Codex => roots(&[".codex", "skills"]), + Agent::Opencode => [ + roots(&[".opencode", "skills"]), + roots(&[".config", "opencode", "skills"]), + roots(&[".claude", "skills"]), + roots(&[".agents", "skills"]), + ] + .concat(), + } + } + + /// `Path::exists` follows symlinks, which matters — skills are commonly + /// installed once and symlinked into each agent's directory. + pub fn has_skill(&self) -> bool { + self.skill_dirs() + .into_iter() + .any(|dir| dir.join(SKILL_NAME).join(SKILL_FILE).exists()) + } +} + +fn skill_missing(agent: Agent) -> error::Error { + error::user( + &format!( + "The {SKILL_NAME} skill for {} could not be found", + agent.display_name() + ), + &format!("Install it with:\n\n {INSTALL_HINT}"), + ) +} + +/// Decide which agent to launch. Selection is local, so it runs before anything +/// is resolved or opened and a misconfigured machine fails without side effects. +pub fn select( + explicit: Option, + availability: impl Fn(Agent) -> Availability, +) -> Result { + if let Some(agent) = explicit { + let available = availability(agent); + if !available.binary { + return Err(error::user( + &format!("{} is not installed", agent.display_name()), + &format!( + "Install it and make sure '{}' is on your PATH, or pair with another agent.", + agent.binary() + ), + )); + } + if !available.skill { + return Err(skill_missing(agent)); + } + return Ok(agent); + } + + Agent::ALL + .into_iter() + .find(|agent| availability(*agent).is_ready()) + .ok_or_else(|| { + error::user( + "No agent is ready to pair", + &format!( + "Pairing needs claude, codex or opencode on your PATH with the \ + {SKILL_NAME} skill installed.\n\nInstall the skill with:\n\n \ + {INSTALL_HINT}\n\nOr print the prompt for another agent with --prompt-only.", + ), + ) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const READY: Availability = Availability { + binary: true, + skill: true, + }; + const NO_BINARY: Availability = Availability { + binary: false, + skill: true, + }; + const NO_SKILL: Availability = Availability { + binary: true, + skill: false, + }; + const MISSING: Availability = Availability { + binary: false, + skill: false, + }; + + /// Availability by agent, in `Agent::ALL` order. + fn given(all: [Availability; 3]) -> impl Fn(Agent) -> Availability { + move |agent| all[Agent::ALL.iter().position(|a| *a == agent).unwrap()] + } + + #[test] + fn auto_picks_the_first_fully_available_agent() { + let selected = select(None, given([READY, READY, READY])).unwrap(); + assert_eq!(selected, Agent::Claude); + } + + #[test] + fn auto_skips_an_agent_that_is_missing_its_skill() { + let selected = select(None, given([NO_SKILL, READY, READY])).unwrap(); + assert_eq!(selected, Agent::Codex); + } + + #[test] + fn auto_skips_an_agent_that_is_not_installed() { + let selected = select(None, given([NO_BINARY, NO_BINARY, READY])).unwrap(); + assert_eq!(selected, Agent::Opencode); + } + + #[test] + fn auto_errors_when_no_agent_is_ready() { + let err = select(None, given([NO_SKILL, NO_BINARY, MISSING])).unwrap_err(); + assert!(err.is_user()); + assert!(err.to_string().contains(INSTALL_HINT), "{err}"); + } + + /// The prompt is one argv element, never split and never a shell string. + fn argv(agent: Agent) -> Vec { + let command = agent.command("pair with me"); + std::iter::once(command.as_std().get_program()) + .chain(command.as_std().get_args()) + .map(|arg| arg.to_string_lossy().into_owned()) + .collect() + } + + #[test] + fn claude_and_codex_take_the_prompt_as_their_first_argument() { + assert_eq!(argv(Agent::Claude), ["claude", "pair with me"]); + assert_eq!(argv(Agent::Codex), ["codex", "pair with me"]); + } + + #[test] + fn opencode_takes_the_prompt_behind_its_prompt_flag() { + assert_eq!( + argv(Agent::Opencode), + ["opencode", "--prompt", "pair with me"] + ); + } + + #[test] + fn an_explicit_agent_is_used_when_it_is_ready() { + let selected = select(Some(Agent::Opencode), given([READY, READY, READY])).unwrap(); + assert_eq!(selected, Agent::Opencode); + } + + #[test] + fn an_explicit_agent_errors_when_its_binary_is_missing() { + let err = select(Some(Agent::Codex), given([READY, NO_BINARY, READY])).unwrap_err(); + assert!(err.is_user()); + assert!(err.to_string().contains("codex"), "{err}"); + } + + #[test] + fn an_explicit_agent_errors_when_its_skill_is_missing() { + let err = select(Some(Agent::Claude), given([NO_SKILL, READY, READY])).unwrap_err(); + assert!(err.is_user()); + assert!(err.to_string().contains(INSTALL_HINT), "{err}"); + } +} diff --git a/src/commands/pair/mod.rs b/src/commands/pair/mod.rs new file mode 100644 index 0000000..b468056 --- /dev/null +++ b/src/commands/pair/mod.rs @@ -0,0 +1,177 @@ +mod agent; +mod prompt; +mod session; +mod target; + +use clap::Args; +use serde::Serialize; + +use crate::commands::GlobalArgs; +use crate::error::Result; + +use agent::{select, Agent}; +use prompt::{build_prompt, write_token}; +use session::Sessions; +use target::{resolve_dataset_editor, resolve_editor, PairTarget}; + +#[derive(Args, Debug, Serialize)] +#[command(about = "Pair an agent CLI with a workspace's marimo notebook")] +pub struct Pair { + /// The workspace, or dataset with --dataset, to pair on, as + /// "{owner}/{slug}" with an optional "@{version}". Defaults to the newest + /// draft version. + target: String, + /// Pair on a dataset's notebook instead of a workspace's + #[arg(long)] + dataset: bool, + /// The notebook to open, defaulting to the workspace's overview notebook + #[arg(long)] + notebook: Option, + /// Pair with Claude Code instead of the first agent found + #[arg(long, group = "agent")] + claude: bool, + /// Pair with Codex instead of the first agent found + #[arg(long, group = "agent")] + codex: bool, + /// Pair with opencode instead of the first agent found + #[arg(long, group = "agent")] + opencode: bool, + /// Do not open the notebook in a browser + #[arg(long)] + no_open: bool, + /// Print the prompt instead of launching an agent + #[arg(long, conflicts_with = "agent")] + prompt_only: bool, +} + +impl Pair { + fn agent(&self) -> Option { + match (self.claude, self.codex, self.opencode) { + (true, _, _) => Some(Agent::Claude), + (_, true, _) => Some(Agent::Codex), + (_, _, true) => Some(Agent::Opencode), + _ => None, + } + } +} + +pub async fn pair(args: Pair, global: GlobalArgs) -> Result<()> { + let target: PairTarget = args.target.parse()?; + + // Selection only reads the local machine, so an agent that cannot pair + // fails here, before anything is resolved or opened. + let agent = if args.prompt_only { + None + } else { + Some(select(args.agent(), |agent| agent.availability())?) + }; + + let pb = global + .spinner() + .with_message(format!("Resolving the editor for {target}")); + let client = global.graphql_client().await?; + let editor = if args.dataset { + resolve_dataset_editor(&client, &target, args.notebook).await? + } else { + resolve_editor(&client, &target, args.notebook).await? + }; + pb.set_message(format!("Editor for {target} is {}", editor.phase)); + + let (token_dir, token_path) = write_token(&editor.token)?; + let editor_page = global + .aqora_url()? + .join(&format!("workspaces/{}/edit", editor.editor_page_id))?; + + // A session only exists while the notebook is open in a browser, so open it + // — but not if the user already has it open. + let sessions = Sessions::new(&editor, global.allow_insecure_host)?; + if !sessions.is_ready().await { + if args.no_open { + pb.println(format!("Please open {editor_page} to start the notebook")); + } else { + pb.set_message(format!("Opening {editor_page}")); + if open::that(editor_page.as_str()).is_err() { + pb.println(format!( + "Could not open a browser. Please open {editor_page}" + )); + } + } + pb.set_message("Waiting for the notebook to connect"); + sessions.wait(&pb, &editor_page).await?; + } + + let prompt = build_prompt(&editor, &token_path, &editor_page); + match agent { + Some(agent) => { + pb.finish_with_message(format!("Launching {}", agent.display_name())); + agent.command(&prompt).spawn()?.wait().await?; + // The agent is done with the token now. + drop(token_dir); + } + None => { + // The agent this prompt is for starts after we exit, so the token + // has to outlive us. + let _ = token_dir.into_path(); + pb.finish_and_clear(); + println!("{prompt}"); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::commands::{Cli, Commands}; + use clap::Parser; + + fn parse(args: &[&str]) -> std::result::Result { + // `Cli`'s version string reports the embedded interpreter's version. + pyo3::Python::initialize(); + let argv = [&["aqora", "pair"], args].concat(); + match Cli::try_parse_from(argv)?.commands { + Commands::Pair(pair) => Ok(pair), + other => panic!("parsed as {other:?}"), + } + } + + #[test] + fn parses_a_bare_target() { + let args = parse(&["alice/ws"]).unwrap(); + assert_eq!(args.target, "alice/ws"); + assert_eq!(args.agent(), None); + assert!(!args.no_open); + assert!(!args.prompt_only); + } + + #[test] + fn an_agent_flag_selects_that_agent() { + assert_eq!( + parse(&["alice/ws", "--codex"]).unwrap().agent(), + Some(Agent::Codex) + ); + } + + #[test] + fn parses_the_dataset_flag() { + assert!(parse(&["alice/ds", "--dataset"]).unwrap().dataset); + } + + #[test] + fn rejects_two_agent_flags() { + assert!(parse(&["alice/ws", "--claude", "--codex"]).is_err()); + } + + #[test] + fn rejects_an_agent_flag_with_prompt_only() { + assert!(parse(&["alice/ws", "--claude", "--prompt-only"]).is_err()); + } + + #[test] + fn prompt_only_and_no_open_are_independent() { + let args = parse(&["alice/ws", "--prompt-only", "--no-open"]).unwrap(); + assert!(args.prompt_only); + assert!(args.no_open); + } +} diff --git a/src/commands/pair/prompt.rs b/src/commands/pair/prompt.rs new file mode 100644 index 0000000..3229b7d --- /dev/null +++ b/src/commands/pair/prompt.rs @@ -0,0 +1,140 @@ +use std::io::Write; +use std::path::{Path, PathBuf}; + +use tempfile::TempDir; +use url::Url; + +use crate::error::Result; + +use super::target::PairEditor; + +/// POSIX single-quoting, so a quote or a space in a URL or a path cannot change +/// how the agent's shell parses the command it is handed. +fn quote(value: impl std::fmt::Display) -> String { + format!("'{}'", value.to_string().replace('\'', r"'\''")) +} + +pub fn build_prompt(editor: &PairEditor, token_path: &Path, editor_page: &Url) -> String { + format!( + "Use the /marimo-pair skill to pair-program on a running marimo notebook. + +Connect to the notebook at: {base_url} + +Use `execute-code.sh --url {quoted_url}` from the marimo-pair skill to execute code in the \ +notebook. + +An auth token is stored at {token_path}. Pass it via `execute-code.sh --url {quoted_url} \ +--token \"$(cat {quoted_token_path})\"`. + +The notebook must be open in a browser for a session to exist. If the server reports no \ +active sessions, ask the user to open {editor_page} and then try again. + +Once you are connected, send a fun toast (mo.status.toast(...)) to the user inside marimo \ +letting them know you're ready to pair.", + base_url = editor.base_url, + quoted_url = quote(&editor.base_url), + token_path = token_path.display(), + quoted_token_path = quote(token_path.display()), + editor_page = editor_page, + ) +} + +/// Keep the token out of the prompt text (and so out of shell history and the +/// agent's transcript) by handing the agent a path instead, the way +/// `marimo pair prompt --with-token` does. +/// +/// The token lives as long as the returned directory: dropping it takes the +/// token with it once the agent has exited. A caller that prints the prompt +/// rather than launching an agent must leak the directory with `into_path`, +/// since the agent it is printed for outlives this process. +pub fn write_token(token: &str) -> Result<(TempDir, PathBuf)> { + let dir = tempfile::Builder::new().prefix("aqora-pair-").tempdir()?; + make_private_dir(dir.path())?; + let path = dir.path().join("token.txt"); + let mut file = open_private(&path)?; + file.write_all(token.as_bytes())?; + file.sync_all()?; + Ok((dir, path)) +} + +/// The umask decides the temp dir's mode, so narrow it explicitly rather than +/// relying on the token file's own 0600 alone. +#[cfg(unix)] +fn make_private_dir(dir: &Path) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + Ok(std::fs::set_permissions( + dir, + std::fs::Permissions::from_mode(0o700), + )?) +} + +#[cfg(not(unix))] +fn make_private_dir(_dir: &Path) -> Result<()> { + Ok(()) +} + +#[cfg(unix)] +fn open_private(path: &Path) -> Result { + use std::os::unix::fs::OpenOptionsExt; + Ok(std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .mode(0o600) + .open(path)?) +} + +#[cfg(not(unix))] +fn open_private(path: &Path) -> Result { + // The containing directory is already private to the user on Windows. + Ok(std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(path)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn commands_carry_a_quoted_url_and_token_path() { + let editor = PairEditor { + base_url: Url::parse("http://host/runner/it's/").unwrap(), + token: "unused".into(), + phase: "READY".into(), + editor_page_id: "workspace-id".into(), + }; + let editor_page = Url::parse("https://aqora.io/workspaces/workspace-id/edit").unwrap(); + + let prompt = build_prompt(&editor, Path::new("/tmp/it's dir/token.txt"), &editor_page); + + assert!( + prompt.contains(r#"--url 'http://host/runner/it'\''s/'"#), + "{prompt}" + ); + assert!( + prompt.contains(r#"cat '/tmp/it'\''s dir/token.txt'"#), + "{prompt}" + ); + } + + #[test] + fn the_token_file_is_removed_when_its_handle_drops() { + let (dir, path) = write_token("s3cret").unwrap(); + assert_eq!(std::fs::read_to_string(&path).unwrap(), "s3cret"); + + drop(dir); + + assert!(!path.exists(), "{} outlived its handle", path.display()); + } + + #[test] + fn a_leaked_token_file_outlives_its_handle() { + let (dir, path) = write_token("s3cret").unwrap(); + + let _ = dir.into_path(); + + assert!(path.exists(), "{} was removed", path.display()); + std::fs::remove_dir_all(path.parent().unwrap()).unwrap(); + } +} diff --git a/src/commands/pair/session.rs b/src/commands/pair/session.rs new file mode 100644 index 0000000..fe80dd0 --- /dev/null +++ b/src/commands/pair/session.rs @@ -0,0 +1,211 @@ +use std::time::Duration; + +use indicatif::ProgressBar; +use url::Url; + +use crate::error::{self, Result}; + +use super::target::PairEditor; + +const POLL_INTERVAL: Duration = Duration::from_secs(1); +const TIMEOUT: Duration = Duration::from_secs(60); +/// Bounds one probe. A runner that accepts a connection and then goes quiet +/// would otherwise hang the poll loop straight through its deadline. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); + +/// The runner's `/api/sessions`, which is what the marimo-pair scripts use to +/// find a session. A session only exists while the notebook is open in a +/// browser, and the prompt is useless without one. +pub struct Sessions { + client: reqwest::Client, + url: Url, + token: String, +} + +impl Sessions { + pub fn new(editor: &PairEditor, allow_insecure_host: bool) -> Result { + // The runner URL may arrive without a trailing slash, which would make + // `join` replace its last path segment instead of appending. + let mut base = editor.base_url.clone(); + if !base.path().ends_with('/') { + base.set_path(&format!("{}/", base.path())); + } + let client = reqwest::Client::builder() + .danger_accept_invalid_certs(allow_insecure_host) + .timeout(REQUEST_TIMEOUT) + .build()?; + Ok(Self { + client, + url: base.join("api/sessions")?, + token: editor.token.clone(), + }) + } + + /// Whether a session is live right now. A runner that is still starting + /// refuses connections and answers errors, so anything that is not a clear + /// "yes" counts as "not yet" — the caller decides how long to keep asking. + pub async fn is_ready(&self) -> bool { + match self.query().await { + Ok(ready) => ready, + Err(err) => { + tracing::debug!("Could not read sessions from {}: {err}", self.url); + false + } + } + } + + async fn query(&self) -> Result { + let response = self + .client + .get(self.url.clone()) + .bearer_auth(&self.token) + .send() + .await? + .error_for_status()?; + has_session(&response.bytes().await?) + } + + /// Poll until the notebook connects, or give up. + pub async fn wait(&self, pb: &ProgressBar, editor_page: &Url) -> Result<()> { + let deadline = tokio::time::Instant::now() + TIMEOUT; + loop { + if self.is_ready().await { + return Ok(()); + } + // Give up rather than sleeping through the deadline first. + if tokio::time::Instant::now() + POLL_INTERVAL >= deadline { + break; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + pb.finish_and_clear(); + Err(error::user( + "The notebook never connected", + &format!( + "Pairing needs the notebook open in a browser. Open {editor_page} and \ + try again." + ), + )) + } +} + +/// `/api/sessions` answers an object keyed by session id, so an empty object +/// means no notebook is open. +fn has_session(body: &[u8]) -> Result { + let sessions: serde_json::Map = serde_json::from_slice(body)?; + Ok(!sessions.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A one-shot HTTP server standing in for the runner. Returns the URL it is + /// serving and the request it received. + async fn serve_once(body: &'static str) -> (Url, tokio::task::JoinHandle) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = Url::parse(&format!( + "http://{}/runner/abc/", + listener.local_addr().unwrap() + )) + .unwrap(); + let handle = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buf = [0u8; 1024]; + while !request.windows(4).any(|w| w == b"\r\n\r\n") { + let read = stream.read(&mut buf).await.unwrap(); + if read == 0 { + break; + } + request.extend_from_slice(&buf[..read]); + } + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + stream.flush().await.unwrap(); + String::from_utf8_lossy(&request).into_owned() + }); + (url, handle) + } + + fn editor(base_url: Url) -> PairEditor { + PairEditor { + base_url, + token: "s3cret".into(), + phase: "Running".into(), + editor_page_id: "id".into(), + } + } + + #[tokio::test] + async fn asks_the_runner_for_its_sessions_with_the_token() { + let (url, served) = serve_once(r#"{"s_1": {"path": "overview.py"}}"#).await; + let sessions = Sessions::new(&editor(url), false).unwrap(); + + assert!(sessions.is_ready().await); + + let request = served.await.unwrap().to_lowercase(); + assert!( + request.contains("get /runner/abc/api/sessions "), + "{request}" + ); + assert!( + request.contains("authorization: bearer s3cret"), + "{request}" + ); + } + + /// A runner that accepts the connection and then goes quiet would otherwise + /// hang the poll loop past its deadline, forever. + #[tokio::test] + async fn gives_up_on_a_runner_that_never_answers() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = Url::parse(&format!( + "http://{}/runner/abc/", + listener.local_addr().unwrap() + )) + .unwrap(); + // Hold the connection open rather than dropping it, which would answer + // the request with a reset. + let _silent = tokio::spawn(async move { + let _connection = listener.accept().await.unwrap(); + std::future::pending::<()>().await; + }); + let sessions = Sessions::new(&editor(url), false).unwrap(); + + let ready = tokio::time::timeout(Duration::from_secs(15), sessions.is_ready()) + .await + .expect("is_ready never gave up"); + + assert!(!ready); + } + + #[tokio::test] + async fn is_not_ready_when_the_runner_reports_no_sessions() { + let (url, _served) = serve_once("{}").await; + let sessions = Sessions::new(&editor(url), false).unwrap(); + + assert!(!sessions.is_ready().await); + } + + #[test] + fn no_session_when_the_map_is_empty() { + assert!(!has_session(b"{}").unwrap()); + } + + #[test] + fn a_session_is_ready_when_the_map_has_an_entry() { + let body = br#"{"s_1234": {"path": "/notebooks/overview.py"}}"#; + assert!(has_session(body).unwrap()); + } + + #[test] + fn errors_on_a_body_that_is_not_json() { + assert!(has_session(b"not marimo").is_err()); + } +} diff --git a/src/commands/pair/target.rs b/src/commands/pair/target.rs new file mode 100644 index 0000000..1e48e2b --- /dev/null +++ b/src/commands/pair/target.rs @@ -0,0 +1,684 @@ +use std::str::FromStr; + +use graphql_client::GraphQLQuery; +use url::Url; + +use crate::{ + error::{self, Result}, + graphql_client::{custom_scalars::*, GraphQLClient}, +}; + +#[derive(GraphQLQuery)] +#[graphql( + query_path = "src/graphql/workspace_pair_editor.graphql", + schema_path = "schema.graphql", + response_derives = "Debug" +)] +pub struct WorkspacePairEditor; + +#[derive(GraphQLQuery)] +#[graphql( + query_path = "src/graphql/workspace_version_pair_editor.graphql", + schema_path = "schema.graphql", + response_derives = "Debug" +)] +pub struct WorkspaceVersionPairEditor; + +#[derive(GraphQLQuery)] +#[graphql( + query_path = "src/graphql/dataset_pair_editor.graphql", + schema_path = "schema.graphql", + response_derives = "Debug" +)] +pub struct DatasetPairEditor; + +#[derive(GraphQLQuery)] +#[graphql( + query_path = "src/graphql/dataset_version_pair_editor.graphql", + schema_path = "schema.graphql", + response_derives = "Debug" +)] +pub struct DatasetVersionPairEditor; + +/// A workspace to pair on, written `owner/slug` with an optional `@version`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PairTarget { + pub owner: String, + pub slug: String, + pub version: Option, +} + +const TARGET_ADVICE: &str = "Expected a workspace like: {owner}/{workspace}[@{version}]"; + +impl FromStr for PairTarget { + type Err = crate::error::Error; + + fn from_str(input: &str) -> Result { + let input = input.strip_prefix('@').unwrap_or(input); + let (owner, rest) = input + .split_once('/') + .ok_or_else(|| error::user("Malformed workspace", TARGET_ADVICE))?; + + // Only an `@` *after* the slash introduces a version, so a leading + // `@owner` stays part of the owner. + let (slug, version) = match rest.split_once('@') { + Some((slug, version)) => { + let version = semver::Version::parse(version).map_err(|err| { + error::user( + &format!("Malformed version '{version}': {err}"), + "Versions are semver, like: 1.2.3", + ) + })?; + (slug, Some(version)) + } + None => (rest, None), + }; + + if owner.is_empty() || slug.is_empty() { + return Err(error::user("Malformed workspace", TARGET_ADVICE)); + } + + Ok(PairTarget { + owner: owner.to_string(), + slug: slug.to_string(), + version, + }) + } +} + +impl std::fmt::Display for PairTarget { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}/{}", self.owner, self.slug)?; + if let Some(version) = &self.version { + write!(f, "@{version}")?; + } + Ok(()) + } +} + +/// Everything the emitted prompt needs about a workspace's editor runner. +#[derive(Debug)] +pub struct PairEditor { + /// The runner's base URL, with the `access_token` query stripped off. + pub base_url: Url, + pub token: String, + pub phase: String, + /// Node id of whatever owns the editor, for the `/workspaces/{id}/edit` page. + pub editor_page_id: String, +} + +/// The platform hands back `{host}/runner/{id}/?file=…&access_token=…`; the +/// marimo-pair scripts want the bare base URL plus the token as a header. +fn split_url_and_token(mut url: Url) -> Result<(Url, String)> { + let token = url + .query_pairs() + .find(|(key, _)| key == "access_token") + .map(|(_, value)| value.into_owned()) + .ok_or_else(|| { + error::system( + "Runner URL carried no access_token", + "The platform returned an unexpected runner URL. Please report this.", + ) + })?; + url.set_query(None); + Ok((url, token)) +} + +fn no_editor(target: &PairTarget, published: bool) -> crate::error::Error { + if published { + error::user( + &format!("{target} is published and has no editor"), + "Published versions are read-only. Pair on the workspace's draft version \ + by dropping the @version, or create a new draft.", + ) + } else { + error::user( + &format!("No editor available for {target}"), + "The version has no editor runner. Open the workspace on aqora.io to start \ + one, then try again.", + ) + } +} + +fn no_draft(target: &PairTarget) -> crate::error::Error { + error::user( + &format!("{target} has no draft version"), + "Pairing edits a workspace's draft version. Create a draft version on aqora.io, \ + then try again.", + ) +} + +fn cannot_edit(target: &PairTarget) -> crate::error::Error { + error::user( + &format!("You cannot edit {target}"), + "Pairing needs edit access to the version. Check that you are logged in as a user \ + who can edit this workspace.", + ) +} + +pub async fn resolve_editor( + client: &GraphQLClient, + target: &PairTarget, + notebook: Option, +) -> Result { + match &target.version { + Some(version) => resolve_pinned(client, target, version, notebook).await, + None => resolve_draft(client, target, notebook).await, + } +} + +async fn resolve_pinned( + client: &GraphQLClient, + target: &PairTarget, + version: &semver::Version, + notebook: Option, +) -> Result { + let workspace = client + .send::(workspace_version_pair_editor::Variables { + owner: target.owner.clone(), + slug: target.slug.clone(), + version: version.to_string(), + notebook, + }) + .await? + .workspace_by_slug + .ok_or_else(|| workspace_not_found(target))?; + + let version = workspace.version.ok_or_else(|| { + error::user( + &format!("{target} does not exist"), + &format!( + "Check the version with 'aqora pair {}/{}' or on aqora.io", + target.owner, target.slug + ), + ) + })?; + + // Published first: a read-only version is a clearer answer than telling + // someone they cannot edit it. + if version.published_at.is_some() { + return Err(no_editor(target, true)); + } + if !version.viewer_can_edit { + return Err(cannot_edit(target)); + } + + let editor = version.editor.ok_or_else(|| no_editor(target, false))?; + let (base_url, token) = split_url_and_token(editor.url)?; + + Ok(PairEditor { + base_url, + token, + phase: format!("{:?}", editor.phase), + editor_page_id: version.id, + }) +} + +async fn resolve_draft( + client: &GraphQLClient, + target: &PairTarget, + notebook: Option, +) -> Result { + let workspace = client + .send::(workspace_pair_editor::Variables { + owner: target.owner.clone(), + slug: target.slug.clone(), + notebook, + }) + .await? + .workspace_by_slug + .ok_or_else(|| workspace_not_found(target))?; + + // Newest draft first, so this is the version the workspace is edited + // through. A workspace that owns a runner directly is not one of these. + let draft = workspace + .versions + .nodes + .into_iter() + .next() + .ok_or_else(|| no_draft(target))?; + + if !draft.viewer_can_edit { + return Err(cannot_edit(target)); + } + + let editor = draft.editor.ok_or_else(|| no_editor(target, false))?; + let (base_url, token) = split_url_and_token(editor.url)?; + + Ok(PairEditor { + base_url, + token, + phase: format!("{:?}", editor.phase), + editor_page_id: draft.id, + }) +} + +fn workspace_not_found(target: &PairTarget) -> crate::error::Error { + error::user( + &format!("Workspace {}/{} not found", target.owner, target.slug), + "Please double check the workspace on aqora.io", + ) +} + +fn dataset_not_found(target: &PairTarget) -> crate::error::Error { + error::user( + &format!("Dataset {}/{} not found", target.owner, target.slug), + "Please double check the dataset on aqora.io", + ) +} + +/// A dataset is edited through the workspace its version owns. +pub async fn resolve_dataset_editor( + client: &GraphQLClient, + target: &PairTarget, + notebook: Option, +) -> Result { + match &target.version { + Some(version) => resolve_dataset_pinned(client, target, version, notebook).await, + None => resolve_dataset_draft(client, target, notebook).await, + } +} + +async fn resolve_dataset_pinned( + client: &GraphQLClient, + target: &PairTarget, + version: &semver::Version, + notebook: Option, +) -> Result { + // Datasets are pinned by their three numbers, with nowhere to put the rest + // of a semver — better to say so than to quietly resolve a different one. + if !version.pre.is_empty() || !version.build.is_empty() { + return Err(error::user( + &format!("Cannot pair on {target}"), + "Datasets are pinned by major.minor.patch, so a prerelease or build suffix \ + cannot be resolved. Drop the suffix, or drop the @version to use the newest \ + draft.", + )); + } + + let dataset = client + .send::(dataset_version_pair_editor::Variables { + owner: target.owner.clone(), + local_slug: target.slug.clone(), + major: version.major as i64, + minor: version.minor as i64, + patch: version.patch as i64, + notebook, + }) + .await? + .dataset_by_slug + .ok_or_else(|| dataset_not_found(target))?; + + let version = dataset.version.ok_or_else(|| { + error::user( + &format!("{target} does not exist"), + &format!( + "Check the version with 'aqora pair --dataset {}/{}' or on aqora.io", + target.owner, target.slug + ), + ) + })?; + + if version.published_at.is_some() { + return Err(no_editor(target, true)); + } + + let workspace = version.workspace.ok_or_else(|| no_editor(target, false))?; + if !workspace.viewer_can_edit { + return Err(cannot_edit(target)); + } + + let editor = workspace.editor.ok_or_else(|| no_editor(target, false))?; + let (base_url, token) = split_url_and_token(editor.url)?; + + Ok(PairEditor { + base_url, + token, + phase: format!("{:?}", editor.phase), + editor_page_id: workspace.id, + }) +} + +async fn resolve_dataset_draft( + client: &GraphQLClient, + target: &PairTarget, + notebook: Option, +) -> Result { + let dataset = client + .send::(dataset_pair_editor::Variables { + owner: target.owner.clone(), + local_slug: target.slug.clone(), + notebook, + }) + .await? + .dataset_by_slug + .ok_or_else(|| dataset_not_found(target))?; + + let draft = dataset + .versions + .nodes + .into_iter() + .next() + .ok_or_else(|| no_draft(target))?; + + let workspace = draft.workspace.ok_or_else(|| no_editor(target, false))?; + if !workspace.viewer_can_edit { + return Err(cannot_edit(target)); + } + + let editor = workspace.editor.ok_or_else(|| no_editor(target, false))?; + let (base_url, token) = split_url_and_token(editor.url)?; + + Ok(PairEditor { + base_url, + token, + phase: format!("{:?}", editor.phase), + editor_page_id: workspace.id, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + use aqora_client::ClientOptions; + + use crate::graphql_client::unauthenticated_client; + + fn parse(input: &str) -> PairTarget { + input.parse().unwrap() + } + + /// Whether a request has been read in full, so the canned answer is not + /// written before the query has arrived. + fn is_complete(request: &[u8]) -> bool { + let text = String::from_utf8_lossy(request); + let Some((headers, body)) = text.split_once("\r\n\r\n") else { + return false; + }; + let length = headers + .lines() + .find_map(|line| { + line.to_lowercase() + .strip_prefix("content-length:") + .and_then(|value| value.trim().parse::().ok()) + }) + .unwrap_or(0); + body.len() >= length + } + + /// A one-shot GraphQL server answering with a canned response. Hands back + /// the request it was sent, so a caller can check what was asked for. + async fn serve_graphql( + response: &'static str, + ) -> (GraphQLClient, tokio::task::JoinHandle) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = Url::parse(&format!("http://{}/", listener.local_addr().unwrap())).unwrap(); + let handle = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buf = [0u8; 1024]; + while !is_complete(&request) { + let read = stream.read(&mut buf).await.unwrap(); + if read == 0 { + break; + } + request.extend_from_slice(&buf[..read]); + } + let body = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{response}", + response.len() + ); + stream.write_all(body.as_bytes()).await.unwrap(); + stream.flush().await.unwrap(); + String::from_utf8_lossy(&request).into_owned() + }); + let client = unauthenticated_client(url, ClientOptions::default()).unwrap(); + (client, handle) + } + + /// A workspace that owns an editor directly *and* has a draft version, so + /// the two candidate runners can be told apart. + const WORKSPACE_WITH_DRAFT: &str = r#"{"data":{"workspaceBySlug":{ + "id": "workspace-id", + "editor": {"id":"r-workspace","phase":"READY", + "url":"http://localhost:8080/runner/workspace/?access_token=workspace-token"}, + "versions":{"nodes":[{"id":"version-id","version":"0.1.0","viewerCanEdit":true, + "editor":{"id":"r-draft","phase":"READY", + "url":"http://localhost:8080/runner/draft/?access_token=draft-token"}}]} + }}}"#; + + #[tokio::test] + async fn a_draft_target_uses_the_draft_versions_editor() { + let (client, _server) = serve_graphql(WORKSPACE_WITH_DRAFT).await; + + let editor = resolve_editor(&client, &parse("alice/ws"), None) + .await + .unwrap(); + + assert_eq!( + editor.base_url.as_str(), + "http://localhost:8080/runner/draft/" + ); + assert_eq!(editor.token, "draft-token"); + assert_eq!(editor.editor_page_id, "version-id"); + } + + #[tokio::test] + async fn a_draft_target_errors_when_the_workspace_has_no_draft_version() { + let (client, _server) = serve_graphql( + r#"{"data":{"workspaceBySlug":{"id":"workspace-id","versions":{"nodes":[]}}}}"#, + ) + .await; + + let err = resolve_editor(&client, &parse("alice/ws"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("has no draft version"), "{err}"); + assert!(err.to_string().contains("Create a draft version"), "{err}"); + } + + #[tokio::test] + async fn a_draft_target_errors_when_the_viewer_cannot_edit_the_draft() { + let (client, _server) = serve_graphql( + r#"{"data":{"workspaceBySlug":{"id":"workspace-id","versions":{"nodes":[ + {"id":"version-id","version":"0.1.0","viewerCanEdit":false, + "editor":{"id":"r-draft","phase":"READY", + "url":"http://localhost:8080/runner/draft/?access_token=draft-token"}} + ]}}}}"#, + ) + .await; + + let err = resolve_editor(&client, &parse("alice/ws"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("edit"), "{err}"); + } + + /// A dataset version is edited through the workspace it owns. + const DATASET_WITH_DRAFT: &str = r#"{"data":{"datasetBySlug":{ + "id": "dataset-id", + "versions":{"nodes":[{"id":"dataset-version-id","version":"0.1.0","workspace":{ + "id": "workspace-id", + "viewerCanEdit": true, + "editor":{"id":"r-dataset","phase":"READY", + "url":"http://localhost:8080/runner/dataset/?access_token=dataset-token"}}}]} + }}}"#; + + #[tokio::test] + async fn a_dataset_target_uses_the_draft_versions_workspace_editor() { + let (client, _server) = serve_graphql(DATASET_WITH_DRAFT).await; + + let editor = resolve_dataset_editor(&client, &parse("alice/ds"), None) + .await + .unwrap(); + + assert_eq!( + editor.base_url.as_str(), + "http://localhost:8080/runner/dataset/" + ); + assert_eq!(editor.token, "dataset-token"); + // The edit page resolves a workspace, not a dataset version. + assert_eq!(editor.editor_page_id, "workspace-id"); + } + + #[tokio::test] + async fn a_dataset_target_errors_when_the_dataset_has_no_draft_version() { + let (client, _server) = serve_graphql( + r#"{"data":{"datasetBySlug":{"id":"dataset-id","versions":{"nodes":[]}}}}"#, + ) + .await; + + let err = resolve_dataset_editor(&client, &parse("alice/ds"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("has no draft version"), "{err}"); + } + + #[tokio::test] + async fn a_dataset_target_errors_when_the_dataset_does_not_exist() { + let (client, _server) = serve_graphql(r#"{"data":{"datasetBySlug":null}}"#).await; + + let err = resolve_dataset_editor(&client, &parse("alice/ds"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!( + err.to_string().contains("Dataset alice/ds not found"), + "{err}" + ); + } + + #[tokio::test] + async fn a_dataset_target_errors_when_the_viewer_cannot_edit_the_workspace() { + let (client, _server) = serve_graphql( + r#"{"data":{"datasetBySlug":{"id":"dataset-id","versions":{"nodes":[ + {"id":"dataset-version-id","version":"0.1.0","workspace":{ + "id":"workspace-id","viewerCanEdit":false, + "editor":{"id":"r-dataset","phase":"READY", + "url":"http://localhost:8080/runner/dataset/?access_token=t"}}} + ]}}}}"#, + ) + .await; + + let err = resolve_dataset_editor(&client, &parse("alice/ds"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("edit"), "{err}"); + } + + #[tokio::test] + async fn a_pinned_dataset_target_rejects_a_prerelease_version() { + let (client, _server) = serve_graphql(DATASET_WITH_DRAFT).await; + + let err = resolve_dataset_editor(&client, &parse("alice/ds@1.2.3-beta.1"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("prerelease"), "{err}"); + } + + #[tokio::test] + async fn a_pinned_dataset_target_asks_for_that_version() { + let (client, served) = serve_graphql( + r#"{"data":{"datasetBySlug":{"id":"dataset-id","version":{ + "id":"dataset-version-id","version":"1.2.3","publishedAt":null,"workspace":{ + "id":"workspace-id","viewerCanEdit":true, + "editor":{"id":"r-dataset","phase":"READY", + "url":"http://localhost:8080/runner/dataset/?access_token=t"}}} + }}}"#, + ) + .await; + + let editor = resolve_dataset_editor(&client, &parse("alice/ds@1.2.3"), None) + .await + .unwrap(); + + assert_eq!(editor.editor_page_id, "workspace-id"); + let request = served.await.unwrap(); + assert!(request.contains(r#""major":1"#), "{request}"); + assert!(request.contains(r#""minor":2"#), "{request}"); + assert!(request.contains(r#""patch":3"#), "{request}"); + } + + #[tokio::test] + async fn a_pinned_target_errors_when_the_viewer_cannot_edit_the_version() { + let (client, _server) = serve_graphql( + r#"{"data":{"workspaceBySlug":{"id":"workspace-id","version": + {"id":"version-id","version":"1.2.3","publishedAt":null,"viewerCanEdit":false, + "editor":{"id":"r-draft","phase":"READY", + "url":"http://localhost:8080/runner/draft/?access_token=draft-token"}} + }}}"#, + ) + .await; + + let err = resolve_editor(&client, &parse("alice/ws@1.2.3"), None) + .await + .unwrap_err(); + + assert!(err.is_user()); + assert!(err.to_string().contains("edit"), "{err}"); + } + + #[test] + fn parses_owner_and_slug() { + let target = parse("alice/my-workspace"); + assert_eq!(target.owner, "alice"); + assert_eq!(target.slug, "my-workspace"); + assert_eq!(target.version, None); + } + + #[test] + fn strips_leading_at_from_owner() { + assert_eq!(parse("@alice/my-workspace"), parse("alice/my-workspace")); + } + + #[test] + fn parses_version() { + let target = parse("@alice/my-workspace@1.2.3"); + assert_eq!(target.owner, "alice"); + assert_eq!(target.slug, "my-workspace"); + assert_eq!(target.version, Some(semver::Version::new(1, 2, 3))); + } + + #[test] + fn round_trips_through_display() { + for input in ["alice/my-workspace", "alice/my-workspace@1.2.3"] { + assert_eq!(parse(input).to_string(), input); + } + } + + #[test] + fn rejects_malformed_targets() { + for input in ["my-workspace", "alice/", "/my-workspace", "alice/ws@nope"] { + assert!(input.parse::().is_err(), "accepted {input:?}"); + } + } + + #[test] + fn splits_token_off_the_runner_url() { + let url = Url::parse("http://localhost:8080/runner/abc/?file=readme.py&access_token=tok") + .unwrap(); + let (base, token) = split_url_and_token(url).unwrap(); + assert_eq!(base.as_str(), "http://localhost:8080/runner/abc/"); + assert_eq!(token, "tok"); + } + + #[test] + fn rejects_a_runner_url_without_a_token() { + let url = Url::parse("http://localhost:8080/runner/abc/?file=readme.py").unwrap(); + assert!(split_url_and_token(url).is_err()); + } +} diff --git a/src/graphql/dataset_pair_editor.graphql b/src/graphql/dataset_pair_editor.graphql new file mode 100644 index 0000000..fce757e --- /dev/null +++ b/src/graphql/dataset_pair_editor.graphql @@ -0,0 +1,20 @@ +query DatasetPairEditor($owner: String!, $localSlug: String!, $notebook: String) { + datasetBySlug(owner: $owner, localSlug: $localSlug) { + id + versions(first: 1, filters: { published: false }) { + nodes { + id + version + workspace { + id + viewerCanEdit: viewerCan(action: UPDATE_WORKSPACE) + editor { + id + phase + url(notebook: $notebook) + } + } + } + } + } +} diff --git a/src/graphql/dataset_version_pair_editor.graphql b/src/graphql/dataset_version_pair_editor.graphql new file mode 100644 index 0000000..c5b187e --- /dev/null +++ b/src/graphql/dataset_version_pair_editor.graphql @@ -0,0 +1,26 @@ +query DatasetVersionPairEditor( + $owner: String! + $localSlug: String! + $major: Int! + $minor: Int! + $patch: Int! + $notebook: String +) { + datasetBySlug(owner: $owner, localSlug: $localSlug) { + id + version(major: $major, minor: $minor, patch: $patch) { + id + version + publishedAt + workspace { + id + viewerCanEdit: viewerCan(action: UPDATE_WORKSPACE) + editor { + id + phase + url(notebook: $notebook) + } + } + } + } +} diff --git a/src/graphql/workspace_pair_editor.graphql b/src/graphql/workspace_pair_editor.graphql new file mode 100644 index 0000000..e2dca8f --- /dev/null +++ b/src/graphql/workspace_pair_editor.graphql @@ -0,0 +1,17 @@ +query WorkspacePairEditor($owner: String!, $slug: String!, $notebook: String) { + workspaceBySlug(owner: $owner, slug: $slug) { + id + versions(first: 1, filters: { published: false }) { + nodes { + id + version + viewerCanEdit: viewerCan(action: UPDATE_WORKSPACE_VERSION) + editor { + id + phase + url(notebook: $notebook) + } + } + } + } +} diff --git a/src/graphql/workspace_version_pair_editor.graphql b/src/graphql/workspace_version_pair_editor.graphql new file mode 100644 index 0000000..fee35de --- /dev/null +++ b/src/graphql/workspace_version_pair_editor.graphql @@ -0,0 +1,21 @@ +query WorkspaceVersionPairEditor( + $owner: String! + $slug: String! + $version: Semver! + $notebook: String +) { + workspaceBySlug(owner: $owner, slug: $slug) { + id + version(version: $version) { + id + version + publishedAt + viewerCanEdit: viewerCan(action: UPDATE_WORKSPACE_VERSION) + editor { + id + phase + url(notebook: $notebook) + } + } + } +}