use std::ffi::OsString;
use clap::{Args, Subcommand};
use crate::clients::ClientKind;
#[must_use]
pub fn protect_client_arguments(arguments: Vec<OsString>, nested: bool) -> Vec<OsString> {
let start = if nested {
arguments
.iter()
.position(|argument| argument == "with")
.map_or(arguments.len(), |position| position + 1)
} else {
1
};
let value_options = [
"--server",
"--token",
"--model",
"--run-ttl-hours",
"--run-max-requests",
];
let clients: Vec<&str> = crate::clients::ClientKind::ALL
.iter()
.flat_map(|kind| [kind.canonical_name(), kind.legacy_name()])
.collect();
let boolean_options = [
"--global",
"--undo",
"--non-interactive",
"--interactive",
"--token-stdin",
];
let mut position = start;
while position < arguments.len() {
let value = arguments[position].to_string_lossy();
if value_options.contains(&value.as_ref()) {
position += 2;
continue;
}
if value_options
.iter()
.any(|option| value.starts_with(&format!("{option}=")))
{
position += 1;
continue;
}
if clients.contains(&value.as_ref()) {
let client = arguments[position].clone();
let prefix = arguments[..position].to_vec();
let mut wrapper = Vec::new();
let mut forwarded = Vec::new();
let mut cursor = position + 1;
let mut explicit_boundary = false;
while cursor < arguments.len() {
let item = arguments[cursor].to_string_lossy();
if explicit_boundary {
forwarded.push(arguments[cursor].clone());
cursor += 1;
continue;
}
if item == "--" {
explicit_boundary = true;
cursor += 1;
continue;
}
if boolean_options.contains(&item.as_ref()) {
wrapper.push(arguments[cursor].clone());
cursor += 1;
continue;
}
if value_options.contains(&item.as_ref()) {
wrapper.push(arguments[cursor].clone());
if let Some(value) = arguments.get(cursor + 1) {
wrapper.push(value.clone());
cursor += 2;
} else {
cursor += 1;
}
continue;
}
if value_options
.iter()
.any(|option| item.starts_with(&format!("{option}=")))
{
wrapper.push(arguments[cursor].clone());
cursor += 1;
continue;
}
forwarded.push(arguments[cursor].clone());
cursor += 1;
}
let mut normalized = prefix;
normalized.extend(wrapper);
normalized.push(client);
if !forwarded.is_empty() {
normalized.push("--".into());
normalized.extend(forwarded);
}
return normalized;
}
position += 1;
}
arguments
}
#[derive(Clone, Debug, Args)]
#[command(trailing_var_arg = true)]
pub struct WithArgs {
#[arg(long)]
pub global: bool,
#[arg(long, requires = "global")]
pub undo: bool,
#[arg(long, conflicts_with = "interactive")]
pub non_interactive: bool,
#[arg(long, conflicts_with = "non_interactive")]
pub interactive: bool,
#[arg(long)]
pub server: Option<String>,
#[arg(long, hide_env_values = true, conflicts_with = "token_stdin")]
pub token: Option<String>,
#[arg(long, conflicts_with = "token")]
pub token_stdin: bool,
#[arg(long)]
pub model: Option<String>,
#[arg(long, default_value_t = 1)]
pub run_ttl_hours: i64,
#[arg(long)]
pub run_max_requests: Option<u64>,
#[arg(value_enum)]
pub client: ClientKind,
#[arg(value_name = "CLIENT_ARGS", allow_hyphen_values = true)]
pub client_args: Vec<OsString>,
}
#[derive(Debug, Subcommand)]
pub enum ServerOp {
Use {
server: Option<String>,
#[arg(long, hide_env_values = true, conflicts_with = "token_stdin")]
token: Option<String>,
#[arg(long, conflicts_with = "token")]
token_stdin: bool,
#[arg(long)]
clear: bool,
#[arg(long)]
run_max_requests: Option<u64>,
},
Status,
Start,
Claim,
Stop,
Remove {
#[arg(long)]
yes: bool,
},
#[command(hide = true)]
Reap { pid: u32 },
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn wrapper_flags_after_client_are_protected() {
for option in [
"--global",
"--undo",
"--non-interactive",
"--interactive",
"--token-stdin",
] {
let arguments = ["router", "with", "codex", option, "prompt"]
.into_iter()
.map(OsString::from)
.collect();
assert_eq!(
protect_client_arguments(arguments, true),
["router", "with", option, "codex", "--", "prompt"].map(OsString::from),
"{option} after the client must remain wrapper-owned"
);
}
}
#[test]
fn value_wrapper_flags_after_client_are_accepted() {
for (option, value) in [
("--server", "https://router.test"),
("--token", "test-token"),
("--model", "gpt-test"),
("--run-ttl-hours", "2"),
("--run-max-requests", "3"),
] {
let arguments = ["with-router", "codex", option, value, "hi"]
.into_iter()
.map(OsString::from)
.collect();
assert_eq!(
protect_client_arguments(arguments, false),
["with-router", option, value, "codex", "--", "hi"].map(OsString::from),
"{option} VALUE after the client must remain wrapper-owned"
);
let equals = format!("{option}={value}");
let arguments = ["with-router", "codex", &equals, "hi"]
.into_iter()
.map(OsString::from)
.collect();
assert_eq!(
protect_client_arguments(arguments, false),
[
OsString::from("with-router"),
OsString::from(&equals),
OsString::from("codex"),
OsString::from("--"),
OsString::from("hi"),
],
"{option}=VALUE after the client must remain wrapper-owned"
);
}
}
#[test]
fn explicit_boundary_forwards_every_colliding_wrapper_flag_verbatim() {
for option in [
"--global",
"--undo",
"--non-interactive",
"--interactive",
"--token-stdin",
"--server",
"--token",
"--model",
"--run-ttl-hours",
"--run-max-requests",
] {
let arguments = ["with-router", "codex", "--", option, "client-value"]
.into_iter()
.map(OsString::from)
.collect::<Vec<_>>();
assert_eq!(
protect_client_arguments(arguments.clone(), false),
arguments,
"{option} after -- must be forwarded to the client"
);
}
}
#[test]
fn option_values_that_match_clients_are_not_boundaries() {
let arguments = ["with-router", "--model", "codex", "qwen", "hello"]
.into_iter()
.map(OsString::from)
.collect();
assert_eq!(
protect_client_arguments(arguments, false),
["with-router", "--model", "codex", "qwen", "--", "hello"].map(OsString::from)
);
}
}