Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566//! Private capability API and background scheduler. The host owns all state.use crate::{ brain::{Brain, Generated, Generation}, routing::{self, Costs, Kind, Queue, Request as RouteRequest, Work}, simulation::{self, Command, HostedPet, Task}, store::{AccessDenied, Conflict, Store}, types::{DiaryPage, PetState}, world,};use anyhow::{Result, ensure};use axum::{ Json, Router, extract::{DefaultBodyLimit, Path, State}, http::{HeaderMap, StatusCode, header}, response::{IntoResponse, Response}, routing::{get, post},};use serde::{Deserialize, Serialize};use serde_json::{Value, json};use std::time::Instant;use std::{ collections::HashMap, sync::{ Arc, Mutex, atomic::{AtomicBool, Ordering}, },};use tokio::sync::{Notify, Semaphore};
pub trait Generator: Send { fn generate(&mut self, g: &Generation, cancel: &AtomicBool) -> Result<Generated>;}impl Generator for Brain { fn generate(&mut self, g: &Generation, cancel: &AtomicBool) -> Result<Generated> { self.generate_cancellable(g, || cancel.load(Ordering::Relaxed)) }}pub struct Worker { pub name: &'static str, pub engine: Arc<Mutex<Box<dyn Generator>>>, pub slot: Arc<Semaphore>, pub available: Arc<AtomicBool>, pub costs: Arc<Mutex<Costs>>,}impl Worker { pub fn new(name: &'static str, engine: impl Generator + 'static) -> Self { Self { name, engine: Arc::new(Mutex::new(Box::new(engine))), slot: Arc::new(Semaphore::new(1)), available: Arc::new(AtomicBool::new(true)), costs: Arc::new(Mutex::new(Costs::default())), } }}pub struct Lease { pub epoch: u64, pub context: u64, pub cancel: Arc<AtomicBool>, pub lane: usize, pub started: Instant, pub estimated_ms: Option<f64>,}pub struct Host { pub db: Mutex<Store>, pub model_id: String, pub features: Vec<String>, pub max_pets: usize, pub weather_seed: String, pub cpu: Worker, pub gpu: Option<Worker>, pub jobs: Mutex<HashMap<String, Lease>>, pub notify: Notify, pub queue: Mutex<Queue>,}impl Host { pub fn snapshot(&self, id: &str, row: HostedPet) -> Value { let thinking = self.jobs.lock().unwrap().contains_key(id) || row.pending.is_some(); let t = simulation::now(); let weather = simulation::weather_at(&row.state, t, &self.weather_seed); json!({"state":row.state,"epoch":row.epoch,"revision":row.revision,"authority":row.authority,"localDevice":row.local_device,"effects":row.effects,"game":row.game,"thinking":thinking,"lastError":row.last_error,"modelId":self.model_id,"features":self.features,"serverNow":t,"weather":weather}) } fn cancel(&self, id: &str) { if let Some(job) = self.jobs.lock().unwrap().get(id) { job.cancel.store(true, Ordering::Relaxed); } } fn cancel_stale(&self, id: &str, row: &HostedPet) { if let Some(job) = self.jobs.lock().unwrap().get(id) { if row.authority != "remote" || job.epoch != row.epoch || job.context != row.context_revision { job.cancel.store(true, Ordering::Relaxed); } } } pub fn router(self: Arc<Self>) -> Router { Router::new() .route("/api/health", get(health)) .route("/api/pets", post(create)) .route("/api/pets/{id}", get(snapshot).delete(delete)) .route("/api/pets/{id}/commands", post(command)) .route("/api/pets/{id}/transfer", post(transfer)) .layer(DefaultBodyLimit::max(8 * 1024 * 1024)) .with_state(self) } pub async fn run(self: Arc<Self>) { loop { if let Err(error) = self.dispatch() { eprintln!("pet scheduler failed: {error}"); } tokio::select! {_ = tokio::time::sleep(std::time::Duration::from_millis(250)) => {}, _ = self.notify.notified() => {}} } } fn dispatch(self: &Arc<Self>) -> Result<()> { let t = simulation::now(); let rows = self.db.lock().unwrap().advance(t, &self.weather_seed)?; let now = Instant::now(); let mut ready = Vec::new(); let mut remaining = [0.0; 2]; let jobs = self.jobs.lock().unwrap(); for job in jobs.values() { remaining[job.lane] = job.estimated_ms.map_or(f64::INFINITY, |ms| { (ms - job.started.elapsed().as_secs_f64() * 1000.0).max(ms * 0.1) }); } let mut queue = self.queue.lock().unwrap(); for (id, row) in rows { if row.authority != "remote" || row.state.pet.is_none() || row.game.is_some() { continue; } if let Some(job) = jobs.get(&id) { if row.epoch != job.epoch || row.context_revision != job.context { job.cancel.store(true, Ordering::Relaxed); } continue; } let Some(task) = row.choose_task(t) else { continue; }; let work = work(&task, &row, t, &self.weather_seed); let age = queue.age_ms(&id, row.epoch, row.context_revision, work, now); ready.push((id, row, task, RouteRequest { work, age_ms: age })); } queue .waiting .retain(|id, _| ready.iter().any(|r| &r.0 == id)); drop(queue); drop(jobs); ready.sort_by(|a, b| { routing::priority(a.3) .cmp(&routing::priority(b.3)) .then_with(|| b.3.age_ms.total_cmp(&a.3.age_ms)) .then_with(|| a.1.last_turn.total_cmp(&b.1.last_turn)) .then_with(|| a.0.cmp(&b.0)) }); let workers = [Some(&self.cpu), self.gpu.as_ref()]; let available = workers.map(|w| { w.is_some_and(|w| w.available.load(Ordering::Relaxed) && !w.engine.is_poisoned()) }); let free = workers.map(|w| w.is_some_and(|w| w.slot.available_permits() > 0)); let cpu_costs = self.cpu.costs.lock().unwrap().clone(); let gpu_costs = self .gpu .as_ref() .map(|w| w.costs.lock().unwrap().clone()) .unwrap_or_default(); let requests: Vec<_> = ready.iter().map(|r| r.3).collect(); let lanes = routing::plan( [&cpu_costs, &gpu_costs], remaining, available, free, &requests, ); for ((id, row, task, request), lane) in ready.into_iter().zip(lanes) { let Some(worker) = workers[lane].filter(|_| available[lane]) else { continue; }; let Ok(permit) = worker.slot.clone().try_acquire_owned() else { continue; }; let work = request.work; let estimate = routing::estimates([&cpu_costs, &gpu_costs], work)[lane]; let costs = worker.costs.clone(); let started = Instant::now(); let engine = worker.engine.clone(); let availability = worker.available.clone(); let backend = worker.name; let retry_task = task.clone(); let cancel = Arc::new(AtomicBool::new(false)); let captured = self.db .lock() .unwrap() .commit(&id, row.epoch, row.context_revision, |r| { r.engine.begin_turn(t, &mut simulation::dice()); }); if !captured? { continue; } self.jobs.lock().unwrap().insert( id.clone(), Lease { epoch: row.epoch, context: row.context_revision, cancel: cancel.clone(), lane, started, estimated_ms: estimate, }, ); self.queue.lock().unwrap().waiting.remove(&id); let host = self.clone(); let seed = self.weather_seed.clone(); tokio::spawn(async move { let captured = row.clone(); let cancelled = cancel.clone(); let result = tokio::task::spawn_blocking(move || { let mut engine = engine .lock() .map_err(|_| anyhow::anyhow!("brain worker is unavailable"))?; execute(&captured, task, engine.as_mut(), &cancelled, t, &seed) }) .await; let outcome = result.unwrap_or_else(|_| Err(anyhow::anyhow!("brain worker stopped"))); if outcome.is_ok() && !cancel.load(Ordering::Relaxed) { costs .lock() .unwrap() .observe(work, started.elapsed().as_secs_f64() * 1000.0); } let gpu_failed = backend == "gpu" && outcome .as_ref() .err() .is_some_and(|e| !e.is::<FilteredDiary>()) && !cancel.load(Ordering::Relaxed); if gpu_failed { availability.store(false, Ordering::Relaxed); } let at = simulation::now(); let committed = host.db.lock().unwrap().commit( &id, row.epoch, row.context_revision, |current| { if cancel.load(Ordering::Relaxed) { return; } match outcome { Ok(Output::Chat(body)) => { current.engine.commit_line( &body, at, &host.weather_seed, &mut simulation::dice(), ); current.last_error = None; } Ok(Output::Diary(page)) => { current.engine.commit_page(page, false); current.last_error = None; } Err(_) => { current.last_error = Some( "The brain could not finish this turn. It will retry.".into(), ); current.next_idle = at + if gpu_failed { 100.0 } else { 5000.0 }; if gpu_failed { if let Task::Chat { kind } = &retry_task { current.pending = Some(kind.clone()); } } current.revision += 1; } } }, ); if let Err(error) = committed { eprintln!("pet turn commit failed: {error}"); } host.jobs.lock().unwrap().remove(&id); drop(permit); host.notify.notify_one(); }); } Ok(()) }}
/// Approximate prompt size for routing: bounded UTF-8 bytes / 4, not tokens.fn work(task: &Task, row: &HostedPet, t: f64, seed: &str) -> Work { let (kind, bytes, cap) = match task { Task::Chat { kind } => ( if kind == "react" { Kind::Reactive } else { Kind::Idle }, world::header("chat", &row.world(t, seed)).len() + row.state.chat.iter().map(|s| s.len() + 1).sum::<usize>(), if kind == "react" { 956 } else { 468 }, ), Task::Diary { date } => { let log = if date == &row.state.today.date { Some(&row.state.today.log) } else { row.state.history.get(date) }; ( Kind::Diary, world::header("diary", &row.world(t, seed)).len() + log.map_or(0, |d| d.log.iter().map(|s| s.len() + 1).sum::<usize>()), 800, ) } }; Work::new(kind, bytes.div_ceil(4).min(cap))}enum Output { Chat(String), Diary(DiaryPage),}#[derive(Debug)]struct FilteredDiary;impl std::fmt::Display for FilteredDiary { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "diary text was filtered") }}impl std::error::Error for FilteredDiary {}/// Runs one task on a brain. The shared turn logic builds every request,/// appends lookup results and filters the text; only generation happens here.fn execute( row: &HostedPet, task: Task, engine: &mut dyn Generator, cancel: &AtomicBool, t: f64, seed: &str,) -> Result<Output> { match task { Task::Chat { kind } => { let mut turn = row.chat_turn(&kind, t, seed); loop { let reply = engine.generate(&turn.request(), cancel)?; if !turn.feed(&reply.text, &reply.reason) { break; } } Ok(Output::Chat(turn.body().unwrap_or_default().to_owned())) } Task::Diary { date } => { let turn = row .diary_turn(&date, t, seed, false) .context("diary day disappeared")?; let reply = engine.generate(&turn.request(), cancel)?; // Unsafe text drops the whole page; the day stays pending. Ok(Output::Diary( turn.finish(&reply.text).ok_or(FilteredDiary)?, )) } }}
use anyhow::Context;struct ApiError(anyhow::Error);impl From<anyhow::Error> for ApiError { fn from(e: anyhow::Error) -> Self { Self(e) }}impl IntoResponse for ApiError { fn into_response(self) -> Response { let status = if self.0.is::<AccessDenied>() { StatusCode::UNAUTHORIZED } else if self.0.is::<Conflict>() { StatusCode::CONFLICT } else { StatusCode::BAD_REQUEST }; (status, Json(json!({"error":self.0.to_string()}))).into_response() }}type ApiResult = std::result::Result<Json<Value>, ApiError>;fn key(headers: &HeaderMap) -> Result<String> { let key = headers .get(header::AUTHORIZATION) .and_then(|v| v.to_str().ok()) .and_then(|v| v.strip_prefix("Bearer ")) .unwrap_or(""); if key.len() != 64 || !key.bytes().all(|b| b.is_ascii_hexdigit()) { return Err(AccessDenied.into()); } Ok(key.into())}async fn health(State(host): State<Arc<Host>>) -> Json<Value> { Json( json!({"protocol":1,"modelId":host.model_id,"features":host.features,"cpu":!host.cpu.engine.is_poisoned(),"gpu":host.gpu.as_ref().is_some_and(|w|w.available.load(Ordering::Relaxed))}), )}#[derive(Deserialize)]#[serde(rename_all = "camelCase")]struct Create { id: String, state: PetState, model_id: String,}async fn create( State(host): State<Arc<Host>>, headers: HeaderMap, Json(input): Json<Create>,) -> ApiResult { ensure_model(&host, &input.model_id)?; let key = key(&headers)?; let row = host.db.lock().unwrap().create( &input.id, &key, HostedPet::new(input.state, simulation::now()), host.max_pets, )?; Ok(Json(host.snapshot(&input.id, row)))}fn ensure_model(host: &Host, id: &str) -> Result<()> { ensure!( host.model_id == id, "this host does not serve the pet's selected brain" ); Ok(())}async fn snapshot( State(host): State<Arc<Host>>, Path(id): Path<String>, headers: HeaderMap,) -> ApiResult { let row = host.db.lock().unwrap().get(&id, &key(&headers)?)?; Ok(Json(host.snapshot(&id, row)))}async fn delete( State(host): State<Arc<Host>>, Path(id): Path<String>, headers: HeaderMap,) -> ApiResult { host.db.lock().unwrap().delete(&id, &key(&headers)?)?; host.cancel(&id); Ok(Json(json!({"deleted":true})))}#[derive(Deserialize, Serialize)]#[serde(rename_all = "camelCase")]struct CommandRequest { epoch: u64, operation_id: String, command: Command,}async fn command( State(host): State<Arc<Host>>, Path(id): Path<String>, headers: HeaderMap, Json(input): Json<CommandRequest>,) -> ApiResult { let body = serde_json::to_string(&input).unwrap(); let command = input.command; let row = host.db.lock().unwrap().operation( &id, &key(&headers)?, input.epoch, &input.operation_id, &body, |r| r.command(command, simulation::now(), &host.weather_seed), )?; host.cancel_stale(&id, &row); host.notify.notify_one(); Ok(Json(host.snapshot(&id, row)))}#[derive(Deserialize, Serialize)]#[serde(rename_all = "camelCase")]struct Transfer { epoch: u64, operation_id: String, mode: String, device_id: String, state: Option<PetState>,}async fn transfer( State(host): State<Arc<Host>>, Path(id): Path<String>, headers: HeaderMap, Json(input): Json<Transfer>,) -> ApiResult { let body = serde_json::to_string(&input).unwrap(); if uuid::Uuid::parse_str(&input.device_id).is_err() { return Err(anyhow::anyhow!("invalid device identifier").into()); } let row = host.db.lock().unwrap().operation( &id, &key(&headers)?, input.epoch, &input.operation_id, &body, |r| { let t = simulation::now(); match input.mode.as_str() { "remote" => { if r.authority != "local" { return Err(Conflict("this pet already has a remote executor").into()); } if r.local_device .as_ref() .is_some_and(|d| d != &input.device_id) { return Err(Conflict("this pet has moved to another device").into()); } let state = input .state .context("local save is required for remote activation")?; state.validate()?; ensure_model( &host, state.hosting.brain_id.as_deref().unwrap_or(&host.model_id), )?; r.state = state; r.state.hosting.mode = "remote".into(); r.state.hosting.brain_id = Some(host.model_id.clone()); r.authority = "remote".into(); r.local_device = None; r.pending = r.state.pet.as_ref().map(|_| "react".into()); r.next_idle = t + 4000.0; r.effects.clear(); r.last_error = None; } "local" => { if r.authority == "local" && r.local_device .as_ref() .is_some_and(|d| d != &input.device_id) { return Err(Conflict("this pet has moved to another device").into()); } ensure!( r.game.is_none(), "finish the current game before moving the pet" ); if r.authority == "remote" { r.advance(t, &host.weather_seed); } r.authority = "local".into(); r.local_device = Some(input.device_id); r.state.hosting.mode = "local".into(); r.pending = None; } _ => anyhow::bail!("unknown execution mode"), } r.epoch += 1; r.context_revision += 1; r.revision += 1; Ok(()) }, )?; host.cancel_stale(&id, &row); host.notify.notify_one(); Ok(Json(host.snapshot(&id, row)))}