iris-ai
8c323fc7a9
A llama.cpp session had its own `llama-server`: two sessions on one model held two copies of it in memory, a model change bought a load only that session benefited from, and the process was a session's to end. A machine's models are now served by one `llama-server` in **router mode** -- no `-m`, a preset file naming models and their flags, a child server per model asked for, and each request routed by its `model` field. So one server per model with that model's own settings is what a machine runs, while this backend has one process, one port and one record per machine to keep track of. The record is the mechanism every other driver already uses, so a restart adopts it; a session records the same pid in its own directory as `Detail::Shared`, and `process::signal` refuses to signal one of those -- which is what keeps stopping, deleting or cleaning up after one session from unloading a model every other session is using. Nothing stops a router on its own. That is deliberate (a loaded model is minutes of disk) and it is why the machines tab now has a card per provider that opens its own screen: how each model is loaded, how many stay in memory, Unload, and Stop. How a model is *loaded* therefore belongs to the model on its machine rather than to a session -- context size, GPU layers, threads, slots, speculative decoding -- written into the preset as llama-server's own argument names. Saving them re-reads that file, which unloads the model; that is the change taking effect, and the dialog says so before you save. What stays a session's is everything that rides on a request, including which tools it offers: the router hosts one set for the machine and the choice is a filter applied here, so it costs no reload (2,181 tokens of prompt with all seven, 698 with none). Verified end to end against the scratch backend and the emulator: two sessions sharing one loaded model with one child process, a second session joining it with a 26ms prefill, a backend restart adopting the router and answering with the prompt cache intact, the same over ssh to this VM, a model's settings reaching the running server, Unload, and Stop leaving every session `exited` with no error line.
471 lines
16 KiB
Rust
471 lines
16 KiB
Rust
//! Interactive provider login carried between a CLI on a configured machine
|
|
//! and the phone. The CLI remains the only credential writer: this layer keeps
|
|
//! its short-lived process and relays only the authorization URL and the code
|
|
//! a person copies back from the browser.
|
|
|
|
use std::collections::HashMap;
|
|
use std::io::{BufRead, BufReader, Read, Write};
|
|
use std::process::Stdio;
|
|
use std::sync::{Arc, Mutex, mpsc};
|
|
use std::time::{Duration, Instant};
|
|
|
|
use anyhow::{Context, Result};
|
|
use rand::Rng;
|
|
use serde::Serialize;
|
|
|
|
use crate::config::{MachineConfig, ProviderConfig};
|
|
use crate::session::transport::Transport;
|
|
use crate::usage::UsageMonitor;
|
|
|
|
const LOGIN_TIMEOUT: Duration = Duration::from_secs(10 * 60);
|
|
const AUTHORIZATION_URL_TIMEOUT: Duration = Duration::from_secs(15);
|
|
const OUTPUT_POLL: Duration = Duration::from_millis(50);
|
|
const START_WAIT: Duration = Duration::from_secs(16);
|
|
|
|
type Key = (String, String);
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(tag = "state", rename_all = "camelCase")]
|
|
pub enum LoginState {
|
|
Starting,
|
|
WaitingForCode {
|
|
#[serde(rename = "authorizationUrl")]
|
|
authorization_url: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
detail: Option<String>,
|
|
},
|
|
Submitting,
|
|
Succeeded,
|
|
Failed {
|
|
detail: String,
|
|
},
|
|
Cancelled,
|
|
}
|
|
|
|
impl LoginState {
|
|
fn terminal(&self) -> bool {
|
|
matches!(
|
|
self,
|
|
Self::Succeeded | Self::Failed { .. } | Self::Cancelled
|
|
)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct LoginInfo {
|
|
pub attempt: String,
|
|
#[serde(flatten)]
|
|
pub state: LoginState,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct Attempt {
|
|
id: String,
|
|
state: Arc<Mutex<LoginState>>,
|
|
input: mpsc::Sender<Input>,
|
|
}
|
|
|
|
impl Attempt {
|
|
fn info(&self) -> LoginInfo {
|
|
LoginInfo {
|
|
attempt: self.id.clone(),
|
|
state: self.state.lock().unwrap().clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
enum Input {
|
|
Code(String),
|
|
Cancel,
|
|
}
|
|
|
|
/// The active login per machine and provider. Completed attempts stay until a
|
|
/// new one replaces them, so a phone that briefly loses its connection can ask
|
|
/// how the operation ended rather than being handed an ambiguous 404.
|
|
pub struct LoginManager {
|
|
attempts: Mutex<HashMap<Key, Attempt>>,
|
|
usage: Arc<UsageMonitor>,
|
|
}
|
|
|
|
impl LoginManager {
|
|
pub fn new(usage: Arc<UsageMonitor>) -> Self {
|
|
Self {
|
|
attempts: Mutex::new(HashMap::new()),
|
|
usage,
|
|
}
|
|
}
|
|
|
|
pub fn start(&self, machine: MachineConfig, provider: ProviderConfig) -> LoginInfo {
|
|
let provider_key = provider
|
|
.kind
|
|
.usage_provider()
|
|
.expect("a login route only accepts a metered provider")
|
|
.to_string();
|
|
let key = (machine.id.clone(), provider_key);
|
|
let mut attempts = self.attempts.lock().unwrap();
|
|
if let Some(attempt) = attempts.get(&key)
|
|
&& !attempt.state.lock().unwrap().terminal()
|
|
{
|
|
return attempt.info();
|
|
}
|
|
|
|
let id = attempt_id();
|
|
let state = Arc::new(Mutex::new(LoginState::Starting));
|
|
let (input, commands) = mpsc::channel();
|
|
let attempt = Attempt {
|
|
id: id.clone(),
|
|
state: Arc::clone(&state),
|
|
input,
|
|
};
|
|
attempts.insert(key, attempt.clone());
|
|
drop(attempts);
|
|
|
|
// Seed the cache before the worker takes the gate. A usage request in
|
|
// that small handoff window sees a fresh, truthful state and cannot
|
|
// start a second CLI against the same credential.
|
|
let gate = self.usage.claude_authentication_started(&machine);
|
|
let usage = Arc::clone(&self.usage);
|
|
std::thread::spawn(move || {
|
|
let _guard = gate.lock().unwrap();
|
|
run_login(&machine, &provider, commands, &state);
|
|
usage.claude_authentication_finished(&machine.id);
|
|
});
|
|
|
|
attempt.info()
|
|
}
|
|
|
|
pub fn wait_until_ready(&self, machine: &str, provider: &str, attempt: &str) -> LoginInfo {
|
|
let started = Instant::now();
|
|
loop {
|
|
let info = self.read(machine, provider, attempt).unwrap_or(LoginInfo {
|
|
attempt: attempt.to_string(),
|
|
state: LoginState::Failed {
|
|
detail: "the sign-in attempt disappeared".to_string(),
|
|
},
|
|
});
|
|
if !matches!(info.state, LoginState::Starting) || started.elapsed() >= START_WAIT {
|
|
return info;
|
|
}
|
|
std::thread::sleep(OUTPUT_POLL);
|
|
}
|
|
}
|
|
|
|
pub fn read(&self, machine: &str, provider: &str, attempt: &str) -> Option<LoginInfo> {
|
|
let attempts = self.attempts.lock().unwrap();
|
|
let found = attempts.get(&(machine.to_string(), provider.to_string()))?;
|
|
(found.id == attempt).then(|| found.info())
|
|
}
|
|
|
|
pub fn submit(
|
|
&self,
|
|
machine: &str,
|
|
provider: &str,
|
|
attempt: &str,
|
|
code: &str,
|
|
) -> Result<LoginInfo> {
|
|
let code = valid_code(code)?;
|
|
let attempts = self.attempts.lock().unwrap();
|
|
let found = attempts
|
|
.get(&(machine.to_string(), provider.to_string()))
|
|
.filter(|found| found.id == attempt)
|
|
.context("no such sign-in attempt")?;
|
|
if found.state.lock().unwrap().terminal() {
|
|
return Ok(found.info());
|
|
}
|
|
*found.state.lock().unwrap() = LoginState::Submitting;
|
|
if found.input.send(Input::Code(code.to_string())).is_err() {
|
|
*found.state.lock().unwrap() = LoginState::Failed {
|
|
detail: "the sign-in process has stopped".to_string(),
|
|
};
|
|
anyhow::bail!("the sign-in process has stopped");
|
|
}
|
|
Ok(found.info())
|
|
}
|
|
|
|
pub fn cancel(&self, machine: &str, provider: &str, attempt: &str) -> Result<LoginInfo> {
|
|
let attempts = self.attempts.lock().unwrap();
|
|
let found = attempts
|
|
.get(&(machine.to_string(), provider.to_string()))
|
|
.filter(|found| found.id == attempt)
|
|
.context("no such sign-in attempt")?;
|
|
if !found.state.lock().unwrap().terminal() {
|
|
let _ = found.input.send(Input::Cancel);
|
|
}
|
|
Ok(found.info())
|
|
}
|
|
|
|
/// Interactive helpers are unlike sessions: nothing adopts them after a
|
|
/// server restart. End every one while the process is still here to reap
|
|
/// the child it launched.
|
|
pub fn cancel_all(&self) {
|
|
let attempts = self.attempts.lock().unwrap();
|
|
for attempt in attempts.values() {
|
|
if !attempt.state.lock().unwrap().terminal() {
|
|
let _ = attempt.input.send(Input::Cancel);
|
|
}
|
|
}
|
|
let pending: Vec<_> = attempts
|
|
.values()
|
|
.map(|attempt| Arc::clone(&attempt.state))
|
|
.collect();
|
|
drop(attempts);
|
|
let deadline = Instant::now() + Duration::from_secs(2);
|
|
while Instant::now() < deadline
|
|
&& pending
|
|
.iter()
|
|
.any(|state| !state.lock().unwrap().terminal())
|
|
{
|
|
std::thread::sleep(OUTPUT_POLL);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn run_login(
|
|
machine: &MachineConfig,
|
|
provider: &ProviderConfig,
|
|
commands: mpsc::Receiver<Input>,
|
|
state: &Arc<Mutex<LoginState>>,
|
|
) {
|
|
if let Err(err) = run_login_inner(machine, provider, commands, state) {
|
|
*state.lock().unwrap() = LoginState::Failed {
|
|
detail: format!("couldn't sign in to Claude on {}: {err:#}", machine.name),
|
|
};
|
|
}
|
|
}
|
|
|
|
fn run_login_inner(
|
|
machine: &MachineConfig,
|
|
provider: &ProviderConfig,
|
|
commands: mpsc::Receiver<Input>,
|
|
state: &Arc<Mutex<LoginState>>,
|
|
) -> Result<()> {
|
|
let transport = Transport::for_machine(machine);
|
|
let args = vec![
|
|
"BROWSER=/bin/false".to_string(),
|
|
provider.program().to_string(),
|
|
"auth".to_string(),
|
|
"login".to_string(),
|
|
"--claudeai".to_string(),
|
|
];
|
|
let host = match &transport {
|
|
Transport::Here => None,
|
|
Transport::Ssh { ssh, .. } => Some(ssh),
|
|
};
|
|
let mut command = crate::ssh::command(host, "env", &args, None, None);
|
|
command
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
let mut child = command
|
|
.spawn()
|
|
.with_context(|| format!("couldn't run {} auth login", provider.program()))?;
|
|
let mut stdin = child
|
|
.stdin
|
|
.take()
|
|
.context("the login process has no stdin")?;
|
|
let stdout = child
|
|
.stdout
|
|
.take()
|
|
.context("the login process has no stdout")?;
|
|
let stderr = child
|
|
.stderr
|
|
.take()
|
|
.context("the login process has no stderr")?;
|
|
let (output, lines) = mpsc::channel();
|
|
read_lines(stdout, output.clone());
|
|
read_lines(stderr, output);
|
|
|
|
let started = Instant::now();
|
|
let mut authorization_url = None;
|
|
let mut last_line = None;
|
|
loop {
|
|
while let Ok(line) = lines.try_recv() {
|
|
if let Some(url) = authorization_url_in(&line) {
|
|
authorization_url = Some(url.to_string());
|
|
*state.lock().unwrap() = LoginState::WaitingForCode {
|
|
authorization_url: url.to_string(),
|
|
detail: None,
|
|
};
|
|
} else if line.to_ascii_lowercase().contains("invalid code") {
|
|
if let Some(url) = &authorization_url {
|
|
*state.lock().unwrap() = LoginState::WaitingForCode {
|
|
authorization_url: url.clone(),
|
|
detail: Some(
|
|
"That code was not accepted. Copy the complete code and try again."
|
|
.to_string(),
|
|
),
|
|
};
|
|
}
|
|
} else if !line.trim().is_empty() {
|
|
last_line = Some(line.trim().chars().take(500).collect::<String>());
|
|
}
|
|
}
|
|
|
|
match commands.recv_timeout(OUTPUT_POLL) {
|
|
Ok(Input::Code(code)) => {
|
|
*state.lock().unwrap() = LoginState::Submitting;
|
|
writeln!(stdin, "{code}").context("couldn't send the login code")?;
|
|
stdin.flush().context("couldn't send the login code")?;
|
|
}
|
|
Ok(Input::Cancel) => {
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
*state.lock().unwrap() = LoginState::Cancelled;
|
|
return Ok(());
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) => {
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
anyhow::bail!("the phone disconnected from the sign-in attempt");
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Timeout) => {}
|
|
}
|
|
|
|
if let Some(status) = child
|
|
.try_wait()
|
|
.context("couldn't check the login process")?
|
|
{
|
|
*state.lock().unwrap() = if status.success() {
|
|
LoginState::Succeeded
|
|
} else {
|
|
LoginState::Failed {
|
|
detail: last_line
|
|
.unwrap_or_else(|| format!("Claude's login process exited with {status}")),
|
|
}
|
|
};
|
|
return Ok(());
|
|
}
|
|
if authorization_url.is_none() && started.elapsed() >= AUTHORIZATION_URL_TIMEOUT {
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
anyhow::bail!("the Claude CLI did not provide an authorization URL");
|
|
}
|
|
if started.elapsed() >= LOGIN_TIMEOUT {
|
|
let _ = child.kill();
|
|
let _ = child.wait();
|
|
anyhow::bail!("the sign-in attempt expired; start it again");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn read_lines(reader: impl Read + Send + 'static, output: mpsc::Sender<String>) {
|
|
std::thread::spawn(move || {
|
|
for line in BufReader::new(reader).lines().map_while(Result::ok) {
|
|
let _ = output.send(line);
|
|
}
|
|
});
|
|
}
|
|
|
|
fn authorization_url_in(line: &str) -> Option<&str> {
|
|
let start = line.find("https://")?;
|
|
let tail = &line[start..];
|
|
let end = tail
|
|
.find(|character: char| character.is_whitespace() || character == '\u{1b}')
|
|
.unwrap_or(tail.len());
|
|
let url = &tail[..end];
|
|
(url.starts_with("https://claude.com/") || url.starts_with("https://platform.claude.com/"))
|
|
.then_some(url)
|
|
}
|
|
|
|
fn valid_code(code: &str) -> Result<&str> {
|
|
let code = code.trim();
|
|
anyhow::ensure!(!code.is_empty(), "the login code is empty");
|
|
anyhow::ensure!(code.len() <= 4096, "the login code is too long");
|
|
anyhow::ensure!(
|
|
!code.chars().any(char::is_control),
|
|
"the login code contains a line break or control character"
|
|
);
|
|
Ok(code)
|
|
}
|
|
|
|
fn attempt_id() -> String {
|
|
let mut bytes = [0u8; 16];
|
|
rand::rng().fill_bytes(&mut bytes);
|
|
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::DriverKind;
|
|
|
|
#[test]
|
|
fn extracts_only_anthropics_https_login_url() {
|
|
assert_eq!(
|
|
authorization_url_in("visit: https://claude.com/cai/oauth/authorize?state=x"),
|
|
Some("https://claude.com/cai/oauth/authorize?state=x")
|
|
);
|
|
assert!(authorization_url_in("visit: http://claude.com/nope").is_none());
|
|
assert!(authorization_url_in("visit: https://example.com/nope").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn login_code_is_one_bounded_line() {
|
|
assert_eq!(valid_code(" abc#state ").unwrap(), "abc#state");
|
|
assert!(valid_code("\n").is_err());
|
|
assert!(valid_code("a\nb").is_err());
|
|
assert!(valid_code(&"x".repeat(4097)).is_err());
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
#[test]
|
|
fn relays_a_headless_cli_login_without_taking_over_its_credentials() {
|
|
use std::os::unix::fs::PermissionsExt;
|
|
|
|
let dir = tempfile::tempdir().expect("tempdir");
|
|
let cli = dir.path().join("fake-claude");
|
|
std::fs::write(
|
|
&cli,
|
|
"#!/bin/sh\necho 'https://claude.com/cai/oauth/authorize?state=test'\nIFS= read -r code\n[ \"$code\" = 'the-code' ]\n",
|
|
)
|
|
.expect("write fake CLI");
|
|
std::fs::set_permissions(&cli, std::fs::Permissions::from_mode(0o700))
|
|
.expect("make fake CLI executable");
|
|
|
|
let monitor = Arc::new(UsageMonitor::new(Default::default()));
|
|
let logins = LoginManager::new(monitor);
|
|
let machine = MachineConfig {
|
|
id: "vm".to_string(),
|
|
name: "test vm".to_string(),
|
|
ssh: None,
|
|
providers: Vec::new(),
|
|
};
|
|
let provider = ProviderConfig {
|
|
name: "claude-cli".to_string(),
|
|
kind: DriverKind::ClaudeCli,
|
|
command: Some(cli.display().to_string()),
|
|
models: Vec::new(),
|
|
mcp_servers: Vec::new(),
|
|
model_settings: Default::default(),
|
|
max_loaded: None,
|
|
};
|
|
|
|
let started = logins.start(machine, provider);
|
|
let ready = logins.wait_until_ready("vm", "claude", &started.attempt);
|
|
assert!(matches!(ready.state, LoginState::WaitingForCode { .. }));
|
|
let wire = serde_json::to_value(&ready).expect("serialize login state");
|
|
assert!(wire.get("authorizationUrl").is_some(), "{wire}");
|
|
assert!(wire.get("authorization_url").is_none(), "{wire}");
|
|
let submitted = logins
|
|
.submit("vm", "claude", &started.attempt, "the-code")
|
|
.expect("submit code");
|
|
assert!(matches!(submitted.state, LoginState::Submitting));
|
|
|
|
let deadline = Instant::now() + Duration::from_secs(2);
|
|
loop {
|
|
let finished = logins
|
|
.read("vm", "claude", &started.attempt)
|
|
.expect("attempt remains readable");
|
|
if matches!(finished.state, LoginState::Succeeded) {
|
|
break;
|
|
}
|
|
assert!(
|
|
Instant::now() < deadline,
|
|
"login did not finish: {finished:?}"
|
|
);
|
|
std::thread::sleep(OUTPUT_POLL);
|
|
}
|
|
}
|
|
}
|