salvor_server/state.rs
1//! [`AppState`]: the shared handle every request works through, plus the
2//! [`AgentFactory`] seam that turns a submitted definition into a live agent.
3//!
4//! The state owns exactly one thing that matters for durability: an
5//! `Arc<dyn EventStore>`. Everything a request needs, it builds fresh from
6//! that handle. A [`Runtime`] is cheap (a store handle plus two function
7//! pointers), so the server constructs one per request rather than sharing
8//! mutable run state; there is no per-run state living in the process that a
9//! restart would lose. That is the whole kill-safety story restated: the
10//! process holds handles, the store holds truth.
11//!
12//! # Why agent building is a seam, not baked in
13//!
14//! Turning a definition into an [`Agent`] means parsing the agent-definition
15//! format and spawning its MCP servers. That logic already exists in the CLI
16//! (`salvor-cli` owns the TOML schema), and putting a copy here would give the
17//! definition format two homes. Instead the server takes an [`AgentFactory`]:
18//! a caller-supplied function from a submitted [`AgentDefinition`] to a
19//! [`BuiltAgent`]. The `salvor serve` command passes the CLI's own builder, so
20//! there is one definition parser in the workspace; tests pass a factory that
21//! builds an agent with an in-process tool and a mock model, which is how the
22//! control plane is exercised over real HTTP with nothing on the network.
23
24use std::collections::{HashMap, HashSet};
25use std::future::Future;
26use std::pin::Pin;
27use std::sync::{Arc, Mutex};
28use std::time::Duration;
29
30use salvor_graph::Graph;
31use salvor_runtime::{Agent, ClockFn, RandomFn, RunCtx, Runtime, RuntimeError};
32use salvor_store::EventStore;
33use salvor_tools::mcp::McpServer;
34use time::OffsetDateTime;
35use tokio::task::JoinHandle;
36
37use salvor_core::{EventEnvelope, RunId};
38
39use crate::client_tools::ClientToolRegistry;
40use crate::executor::ModelExecutor;
41use crate::tool_registry::ToolRegistry;
42
43/// The format a submitted agent definition is written in.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum DefFormat {
46 /// The agent TOML the CLI reads from a file.
47 Toml,
48 /// The same definition as a JSON document (what a thin SDK sends).
49 Json,
50}
51
52/// A submitted agent definition: the raw bytes plus the format they are in.
53///
54/// The server never interprets the bytes itself; it hands them to the
55/// [`AgentFactory`]. Keeping the raw body (rather than a parsed structure)
56/// means the definition is rebuilt from exactly what was submitted on every
57/// start, resume, and recover, the same way the CLI rebuilds from the TOML
58/// file each time.
59#[derive(Debug, Clone)]
60pub struct AgentDefinition {
61 /// The format of `body`.
62 pub format: DefFormat,
63 /// The raw definition bytes.
64 pub body: Vec<u8>,
65}
66
67/// A live agent plus the MCP server sessions its tools hold.
68///
69/// The sessions must outlive the run: each MCP tool keeps a client-peer clone
70/// into its server's session, so dropping the sessions stops the tools. The
71/// run driver keeps them for the run's life and closes them when it ends.
72pub struct BuiltAgent {
73 /// The built agent the runtime drives.
74 pub agent: Agent,
75 /// The MCP sessions to keep alive for the run, then close.
76 pub servers: Vec<McpServer>,
77}
78
79/// The future an [`AgentFactory`] returns.
80pub type BuildFuture = Pin<Box<dyn Future<Output = Result<BuiltAgent, String>> + Send>>;
81
82/// Builds a live agent from a submitted definition.
83///
84/// The `Err` is a human message; the register and start handlers turn it into
85/// a `400`, because a definition that will not build is a client mistake.
86pub type AgentFactory = Arc<dyn Fn(AgentDefinition) -> BuildFuture + Send + Sync>;
87
88/// One registered agent definition.
89#[derive(Debug, Clone)]
90pub struct RegisteredAgent {
91 /// The submitted definition, kept for every rebuild.
92 pub definition: AgentDefinition,
93 /// The agent's content hash (`agent_def_hash`), the id clients reference.
94 pub agent_hash: String,
95 /// The agent's display name, when the definition declared one
96 /// (`Agent::name`, read off the built agent at registration time).
97 /// `None` when the definition carried no name: genuinely absent, not a
98 /// default to fall back on: [`agents::get`](crate::agents::get) and
99 /// [`agents::list`](crate::agents::list) omit the field entirely for
100 /// such an agent rather than emit `"name": null`.
101 pub name: Option<String>,
102}
103
104/// How often the wake sweeper looks for due timers, unless
105/// [`AppState::with_wake_interval`] says otherwise.
106///
107/// A minute. The unit a durable timer is written in is hours or days, so the
108/// resolution that matters is "within a minute of the deadline", and sweeping
109/// costs a fold of every run's log (status is not a stored column); a shorter
110/// interval would pay that repeatedly to sharpen a number nobody measures.
111pub const DEFAULT_WAKE_INTERVAL: Duration = Duration::from_secs(60);
112
113/// The shared, cheaply cloned handle every route works through.
114#[derive(Clone)]
115pub struct AppState {
116 inner: Arc<Inner>,
117}
118
119struct Inner {
120 store: Arc<dyn EventStore>,
121 factory: AgentFactory,
122 // The general model-executor seam the server performs a client-driven run's
123 // model step through. `None` until a host injects one (the `AgentFactory`
124 // pattern): the model-step endpoint then answers with a clear error rather
125 // than performing a call. `salvor serve` wires a default from its own
126 // client-construction path, so the feature works out of the box.
127 model_executor: Option<Arc<dyn ModelExecutor>>,
128 // The general tool-registry seam the server performs a client-driven run's
129 // tool step through. `None` until a host injects one (the same pattern as
130 // `model_executor`): the tool-step endpoint then answers with a clear error
131 // rather than dispatching. `salvor serve` wires an EMPTY registry, so any
132 // tool-step there is a clean `unknown_tool` until a tool is registered.
133 tool_registry: Option<Arc<ToolRegistry>>,
134 // The client-performed tool DECLARATIONS this server was started with, the
135 // declarative sibling of `tool_registry` above. Not an `Option`: an empty
136 // set is a complete, honest state (every client-tool intent is a clean
137 // `unknown_tool`), because nothing here is ever dispatched, so there is no
138 // "the host wired no mechanism" case to tell apart from "the operator
139 // declared nothing". Loaded by the operator, never over HTTP; see
140 // `crate::client_tools` for why that rule is load-bearing.
141 client_tools: Arc<ClientToolRegistry>,
142 hooks: Option<(ClockFn, RandomFn)>,
143 auth_token: Option<String>,
144 poll_interval: Duration,
145 // How often the wake sweeper looks for runs whose durable timer has come
146 // due. A `Duration` with a real default rather than an `Option`, exactly
147 // like `poll_interval` above: a sleeping run nobody re-drives never wakes,
148 // so the sweep is part of what serving a store means, not an opt-in. Zero
149 // is the off switch (`salvor serve --wake-interval 0`), for an operator who
150 // wakes runs from cron with `salvor wake` instead.
151 wake_interval: Duration,
152 agents: Mutex<HashMap<String, RegisteredAgent>>,
153 // The graph documents this process has accepted, keyed by their reproducible
154 // content hash (`salvor_engine::graph_hash`). In-memory, exactly like the
155 // agent registry above and for the same reason: a graph is pure data with a
156 // content hash, so re-submitting the identical document is idempotent and a
157 // restart re-accepts it under the same hash. Storing it here rather than in
158 // the event store keeps the store schema untouched (additive-migration
159 // discipline) and mirrors how a registered agent lives only in this map.
160 graphs: Mutex<HashMap<String, Graph>>,
161 // Which runs a driver task is still working on, and the handles to those
162 // tasks. The `active` set is membership only, inserted synchronously
163 // before a task is spawned so a concurrent stream can never miss a run
164 // that has just started; `handles` is populated after the spawn and used
165 // only to abort tasks at shutdown, where a stale finished handle is
166 // harmless.
167 active: Mutex<HashSet<RunId>>,
168 handles: Mutex<HashMap<RunId, JoinHandle<()>>>,
169 // The client-driven runs this process has opened, each with its current
170 // drive-token lease. This registry is what keeps the client-driven and
171 // server-driven modes from colliding over one store: the driving
172 // client-run endpoints operate only on runs recorded here, so a
173 // server-driven run is never reachable through them. It is in-memory
174 // because the drive token is a single-writer lease with a process
175 // lifetime: a token nobody is holding any more means nothing.
176 //
177 // Which is why membership here is not the whole answer to "is this run
178 // client-driven". This map knows only what this process opened; a run
179 // opened before a restart is absent from it while its client is still
180 // driving it. The durable half of the answer is the run's own
181 // `RunStarted`, which records `driven_by: client` (see
182 // `client_runs::log_is_client_driven`). The open endpoint adopts such a
183 // run back into this registry, and the surfaces that must not become a
184 // second writer (resume, the wake sweeper) consult the log directly.
185 client_runs: Mutex<HashMap<RunId, ClientRunLease>>,
186 // How long a client-driven run's lease stays "current" without the driver
187 // presenting its token again. Past this, the run reports no attached driver
188 // on GET /v1/runs (the client-driven half of the liveness evidence): the tab
189 // closed, the SDK exited, the driver crashed. It is also how long the run is
190 // another driver's to take: a re-open while the lease is current is refused
191 // (`409 lease_held`), and a lapsed one is not. Generous by default so a
192 // single long model call between drive operations never reads as a false
193 // stall, which would also be a window for a second driver to take a run its
194 // first driver is still working on; a test shortens it (see
195 // `with_client_lease_ttl`) to prove the lapse.
196 client_lease_ttl: Duration,
197 // Runs the wake sweeper has already warned about being unwakeable here (its
198 // agent or graph is not registered in this process). Only the first sighting
199 // per run logs at WARN; every later pass while a run's id stays in this set
200 // logs the same fields at DEBUG, so an operator who has not fixed the
201 // registration gap yet is not paged again every sweep interval, but the
202 // fields to find and fix it are still there for anyone who turns on
203 // debug-level logging. Cleared when the run wakes or drops out of the due
204 // set, so a run that becomes unwakeable again later (a fresh nap, a
205 // different recorded agent hash) warns again.
206 unwakeable_warned: Mutex<HashSet<RunId>>,
207}
208
209/// The per-run lease state for a client-driven run.
210#[derive(Debug, Clone)]
211pub struct ClientRunLease {
212 /// The opaque drive token the single writer must present on every append.
213 pub drive_token: String,
214 /// Whether the opener asked for model request bodies to be recorded. Stored
215 /// at open time; it governs the server-performed model step, not yet
216 /// implemented, and carries no effect on the generic append this surface
217 /// serves.
218 pub record_prompts: bool,
219 /// When the driver last proved it was alive: stamped at open and refreshed
220 /// on every guarded operation (append, model-step, tool-step, resolve), each
221 /// of which presents the drive token. That token is the driver's own proof
222 /// of life, so its arrival IS the heartbeat: there is no separate mechanism.
223 /// Read against the lease TTL to decide whether a driver is still attached to
224 /// a client-driven run (see
225 /// [`client_run_driver_live`](AppState::client_run_driver_live)).
226 pub last_seen: OffsetDateTime,
227}
228
229/// What asking to drop a client-driven run's lease came to (see
230/// [`release_client_run`](AppState::release_client_run)).
231///
232/// The three outcomes are kept apart because the release endpoint answers each
233/// one differently: a release that dropped a lease and a release that found
234/// none are both a job done (the run is unheld either way, which is all a
235/// caller wanted), while a caller presenting somebody else's token is refused
236/// rather than quietly obeyed.
237#[derive(Debug, Clone, Copy, PartialEq, Eq)]
238pub enum LeaseRelease {
239 /// The caller held the lease, and it is gone.
240 Released,
241 /// There was no lease on the run to drop: it lapsed, it was released
242 /// already, or this process never opened the run at all.
243 NoLease,
244 /// A lease stands on the run and the caller did not present its token.
245 /// Nothing was dropped.
246 NotTheHolder,
247}
248
249impl AppState {
250 /// Builds server state over `store`, using `factory` to turn submitted
251 /// definitions into live agents. Auth is off and the clock and random
252 /// source are the runtime defaults until set with the `with_*` methods.
253 #[must_use]
254 pub fn new(store: Arc<dyn EventStore>, factory: AgentFactory) -> Self {
255 Self {
256 inner: Arc::new(Inner {
257 store,
258 factory,
259 model_executor: None,
260 tool_registry: None,
261 client_tools: Arc::new(ClientToolRegistry::new()),
262 hooks: None,
263 auth_token: None,
264 poll_interval: Duration::from_millis(50),
265 wake_interval: DEFAULT_WAKE_INTERVAL,
266 agents: Mutex::new(HashMap::new()),
267 graphs: Mutex::new(HashMap::new()),
268 active: Mutex::new(HashSet::new()),
269 handles: Mutex::new(HashMap::new()),
270 client_runs: Mutex::new(HashMap::new()),
271 client_lease_ttl: Duration::from_secs(60),
272 unwakeable_warned: Mutex::new(HashSet::new()),
273 }),
274 }
275 }
276
277 /// Sets how long a client-driven run's lease stays current without the
278 /// driver presenting its token again (default 60s). Past this, the run
279 /// reports no attached driver and a re-open may take it from the driver
280 /// that has gone quiet. Additive and off-default; a test or a seed
281 /// shortens it to make a driverless client run observable quickly, exactly
282 /// as [`with_poll_interval`](Self::with_poll_interval) shortens the stream
283 /// poll. `salvor serve` reads it from `SALVOR_CLIENT_LEASE_TTL_SECS`.
284 #[must_use]
285 pub fn with_client_lease_ttl(mut self, ttl: Duration) -> Self {
286 Arc::get_mut(&mut self.inner)
287 .expect("with_client_lease_ttl is called before the state is shared")
288 .client_lease_ttl = ttl;
289 self
290 }
291
292 /// Requires `Authorization: Bearer <token>` on every request. Without this,
293 /// the server trusts its caller (the reverse-proxy posture).
294 #[must_use]
295 pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
296 Arc::get_mut(&mut self.inner)
297 .expect("with_auth_token is called before the state is shared")
298 .auth_token = Some(token.into());
299 self
300 }
301
302 /// Injects the general model executor the server performs a client-driven
303 /// run's model step through. Additive and off by default (the existing
304 /// [`new`](Self::new) leaves it unset), so no caller that predates it
305 /// changes behavior. `salvor serve` wires a default here; another host
306 /// injects its own, exactly as it supplies its own [`AgentFactory`].
307 #[must_use]
308 pub fn with_model_executor(mut self, executor: Arc<dyn ModelExecutor>) -> Self {
309 Arc::get_mut(&mut self.inner)
310 .expect("with_model_executor is called before the state is shared")
311 .model_executor = Some(executor);
312 self
313 }
314
315 /// Injects the general tool registry the server performs a client-driven
316 /// run's tool step through. Additive and off by default (the existing
317 /// [`new`](Self::new) leaves it unset), so no caller that predates it
318 /// changes behavior. `salvor serve` wires an empty registry here; another
319 /// host injects one holding its own tools, exactly as it supplies its own
320 /// [`AgentFactory`] and [`ModelExecutor`](crate::ModelExecutor).
321 #[must_use]
322 pub fn with_tool_registry(mut self, registry: Arc<ToolRegistry>) -> Self {
323 Arc::get_mut(&mut self.inner)
324 .expect("with_tool_registry is called before the state is shared")
325 .tool_registry = Some(registry);
326 self
327 }
328
329 /// Loads the client-performed tool declarations this server answers
330 /// client-tool intents against. Additive and empty by default, so no caller
331 /// that predates it changes behavior: without a declaration, every
332 /// client-tool intent is a clean `unknown_tool` and nothing is written.
333 ///
334 /// This is the ONLY way declarations enter the process. There is no
335 /// endpoint that accepts one, on purpose: a declaration fixes the effect
336 /// class, and a client that could declare its own would be choosing whether
337 /// its own write is subject to the write-ahead rule. See
338 /// [`crate::client_tools`] for the full argument.
339 #[must_use]
340 pub fn with_client_tools(mut self, decls: Arc<ClientToolRegistry>) -> Self {
341 Arc::get_mut(&mut self.inner)
342 .expect("with_client_tools is called before the state is shared")
343 .client_tools = decls;
344 self
345 }
346
347 /// Injects the clock and random source every [`Runtime`] this state builds
348 /// uses. Deterministic tests pass fixed functions so full logs compare
349 /// equal across a control run and a recovered one.
350 #[must_use]
351 pub fn with_hooks(mut self, clock: ClockFn, random: RandomFn) -> Self {
352 Arc::get_mut(&mut self.inner)
353 .expect("with_hooks is called before the state is shared")
354 .hooks = Some((clock, random));
355 self
356 }
357
358 /// Sets how often the event stream polls the store for new events (default
359 /// 50ms). Tests shorten it so a streamed run completes quickly.
360 #[must_use]
361 pub fn with_poll_interval(mut self, interval: Duration) -> Self {
362 Arc::get_mut(&mut self.inner)
363 .expect("with_poll_interval is called before the state is shared")
364 .poll_interval = interval;
365 self
366 }
367
368 /// Sets how often the wake sweeper looks for runs whose durable timer has
369 /// come due (default [`DEFAULT_WAKE_INTERVAL`]). `Duration::ZERO` turns the
370 /// sweeper off entirely, for a host that wakes runs some other way.
371 ///
372 /// Same shape as [`with_poll_interval`](Self::with_poll_interval), and a
373 /// test shortens it for the same reason: to make a sweep observable without
374 /// waiting on a wall clock.
375 #[must_use]
376 pub fn with_wake_interval(mut self, interval: Duration) -> Self {
377 Arc::get_mut(&mut self.inner)
378 .expect("with_wake_interval is called before the state is shared")
379 .wake_interval = interval;
380 self
381 }
382
383 /// The event store every request reads from and writes through.
384 #[must_use]
385 pub fn store(&self) -> Arc<dyn EventStore> {
386 self.inner.store.clone()
387 }
388
389 /// The expected bearer token, when auth is required.
390 #[must_use]
391 pub fn auth_token(&self) -> Option<&str> {
392 self.inner.auth_token.as_deref()
393 }
394
395 /// How often the event stream polls for new events.
396 #[must_use]
397 pub fn poll_interval(&self) -> Duration {
398 self.inner.poll_interval
399 }
400
401 /// How often the wake sweeper looks for due timers. `Duration::ZERO` means
402 /// no sweeper runs on this server.
403 #[must_use]
404 pub fn wake_interval(&self) -> Duration {
405 self.inner.wake_interval
406 }
407
408 /// The injected model executor, if a host wired one. `None` means the
409 /// server cannot perform a model step and the endpoint says so.
410 #[must_use]
411 pub fn model_executor(&self) -> Option<Arc<dyn ModelExecutor>> {
412 self.inner.model_executor.clone()
413 }
414
415 /// The injected tool registry, if a host wired one. `None` means the server
416 /// cannot perform a tool step and the endpoint says so; a wired-but-empty
417 /// registry instead reports each tool as `unknown_tool`.
418 #[must_use]
419 pub fn tool_registry(&self) -> Option<Arc<ToolRegistry>> {
420 self.inner.tool_registry.clone()
421 }
422
423 /// The client-performed tool declarations the operator loaded. Empty unless
424 /// [`with_client_tools`](Self::with_client_tools) was called, and an empty
425 /// set answers every client-tool intent with `unknown_tool`.
426 #[must_use]
427 pub fn client_tools(&self) -> Arc<ClientToolRegistry> {
428 self.inner.client_tools.clone()
429 }
430
431 /// Reads the current instant from this state's injected clock, or the real
432 /// UTC clock when none was injected. This stamps envelopes the server
433 /// records itself (the model-step intent and completion), the same clock
434 /// edge a [`Runtime`] would use, so deterministic tests still compare logs.
435 #[must_use]
436 pub fn now(&self) -> OffsetDateTime {
437 match &self.inner.hooks {
438 Some((clock, _)) => clock(),
439 None => OffsetDateTime::now_utc(),
440 }
441 }
442
443 /// A fresh runtime over the shared store, with this state's clock and
444 /// random source.
445 #[must_use]
446 pub fn runtime(&self) -> Runtime {
447 match &self.inner.hooks {
448 Some((clock, random)) => {
449 Runtime::with_hooks(self.inner.store.clone(), clock.clone(), random.clone())
450 }
451 None => Runtime::new(self.inner.store.clone()),
452 }
453 }
454
455 /// Builds a live agent from a submitted definition, through the factory.
456 ///
457 /// # Errors
458 ///
459 /// The factory's human message when the definition will not build.
460 pub async fn build_agent(&self, definition: AgentDefinition) -> Result<BuiltAgent, String> {
461 (self.inner.factory)(definition).await
462 }
463
464 /// Records a registered agent under its content hash, returning that hash.
465 pub fn register_agent(&self, registered: RegisteredAgent) -> String {
466 let hash = registered.agent_hash.clone();
467 self.inner
468 .agents
469 .lock()
470 .expect("agents registry lock")
471 .insert(hash.clone(), registered);
472 hash
473 }
474
475 /// The definition registered under `hash`, if any.
476 #[must_use]
477 pub fn agent(&self, hash: &str) -> Option<RegisteredAgent> {
478 self.inner
479 .agents
480 .lock()
481 .expect("agents registry lock")
482 .get(hash)
483 .cloned()
484 }
485
486 /// Every registered agent's hash, sorted for a stable listing.
487 #[must_use]
488 pub fn agent_hashes(&self) -> Vec<String> {
489 let mut hashes: Vec<String> = self
490 .inner
491 .agents
492 .lock()
493 .expect("agents registry lock")
494 .keys()
495 .cloned()
496 .collect();
497 hashes.sort();
498 hashes
499 }
500
501 /// Records a validated graph document under `hash`, returning whether it was
502 /// newly stored (`true`) or already present (`false`). Re-storing the
503 /// identical document is idempotent: the second call keeps the first and
504 /// reports `false`, the graph counterpart of an agent register's `created`.
505 pub fn store_graph(&self, hash: String, graph: Graph) -> bool {
506 let mut graphs = self.inner.graphs.lock().expect("graphs registry lock");
507 if graphs.contains_key(&hash) {
508 return false;
509 }
510 graphs.insert(hash, graph);
511 true
512 }
513
514 /// The graph document stored under `hash`, if any. `None` is the
515 /// `unknown_graph` case.
516 #[must_use]
517 pub fn graph(&self, hash: &str) -> Option<Graph> {
518 self.inner
519 .graphs
520 .lock()
521 .expect("graphs registry lock")
522 .get(hash)
523 .cloned()
524 }
525
526 /// Every stored graph's hash, sorted for a stable listing.
527 #[must_use]
528 pub fn graph_hashes(&self) -> Vec<String> {
529 let mut hashes: Vec<String> = self
530 .inner
531 .graphs
532 .lock()
533 .expect("graphs registry lock")
534 .keys()
535 .cloned()
536 .collect();
537 hashes.sort();
538 hashes
539 }
540
541 /// Builds a per-run [`RunCtx`] over `log`, with this state's clock and random
542 /// source, so the graph engine can drive a run through the same durability
543 /// substrate the built-in loop uses. This is the graph counterpart of
544 /// [`runtime`](Self::runtime): the built-in loop reaches the store through a
545 /// [`Runtime`]; the graph engine reaches it through a `RunCtx` it drives
546 /// directly, and both share the exact clock/random hooks so a deterministic
547 /// test's logs still compare equal.
548 ///
549 /// # Errors
550 ///
551 /// [`RuntimeError::Replay`] when `log` is not a well-formed run history.
552 pub fn run_ctx(&self, run_id: RunId, log: Vec<EventEnvelope>) -> Result<RunCtx, RuntimeError> {
553 match &self.inner.hooks {
554 Some((clock, random)) => RunCtx::with_hooks(
555 self.inner.store.clone(),
556 run_id,
557 log,
558 clock.clone(),
559 random.clone(),
560 ),
561 None => RunCtx::new(self.inner.store.clone(), run_id, log),
562 }
563 }
564
565 /// Marks a run as being driven. Call this synchronously before spawning
566 /// the driver task, so a stream opened at the same instant sees the run as
567 /// active rather than racing the task's first store write.
568 pub fn begin_run(&self, run_id: RunId) {
569 self.inner
570 .active
571 .lock()
572 .expect("active runs lock")
573 .insert(run_id);
574 }
575
576 /// Records the driver task's handle, for aborting at shutdown.
577 pub fn set_handle(&self, run_id: RunId, handle: JoinHandle<()>) {
578 self.inner
579 .handles
580 .lock()
581 .expect("handles lock")
582 .insert(run_id, handle);
583 }
584
585 /// Marks a run's drive as ended and drops its handle. The task calls this
586 /// as its last act, whether it completed, parked, or errored.
587 pub fn end_run(&self, run_id: RunId) {
588 self.inner
589 .active
590 .lock()
591 .expect("active runs lock")
592 .remove(&run_id);
593 self.inner
594 .handles
595 .lock()
596 .expect("handles lock")
597 .remove(&run_id);
598 }
599
600 /// Whether a run is still being driven by a task in this process.
601 #[must_use]
602 pub fn is_run_active(&self, run_id: RunId) -> bool {
603 self.inner
604 .active
605 .lock()
606 .expect("active runs lock")
607 .contains(&run_id)
608 }
609
610 /// Records (or re-leases) a client-driven run, returning a fresh drive
611 /// token. Called by the open endpoint for a new run, and for a re-open the
612 /// open endpoint has already decided is allowed: no lease stands, or the
613 /// one that does has lapsed, or the run is finished. Minting here is
614 /// unconditional, so it must never be reached while another driver's lease
615 /// is current; that is the whole point of the check in
616 /// [`client_runs::open`](crate::client_runs::open), which reads
617 /// [`current_client_lease`](Self::current_client_lease) first.
618 pub fn lease_client_run(&self, run_id: RunId, record_prompts: bool) -> String {
619 let drive_token = format!("dt_{}", uuid::Uuid::new_v4().simple());
620 self.inner
621 .client_runs
622 .lock()
623 .expect("client runs lock")
624 .insert(
625 run_id,
626 ClientRunLease {
627 drive_token: drive_token.clone(),
628 record_prompts,
629 last_seen: self.now(),
630 },
631 );
632 drive_token
633 }
634
635 /// Refreshes a client-driven run's `last_seen` to now, the driver's proof of
636 /// life. Called by the lease gate on every guarded operation (the driver
637 /// presented its token, so it is alive); a no-op for a run this process holds
638 /// no lease for.
639 pub fn touch_client_run(&self, run_id: RunId) {
640 let now = self.now();
641 if let Some(lease) = self
642 .inner
643 .client_runs
644 .lock()
645 .expect("client runs lock")
646 .get_mut(&run_id)
647 {
648 lease.last_seen = now;
649 }
650 }
651
652 /// A client-driven run's lease and how long is left before it lapses, when
653 /// this process holds one whose driver has proved it is alive within the
654 /// lease TTL. `None` covers both "no lease here" (a fresh process, a run
655 /// this server never opened) and "the lease lapsed" (the tab closed, the
656 /// SDK exited, the driver crashed), because to everything downstream those
657 /// are the same fact: nobody is driving this run.
658 ///
659 /// This is the one place the TTL comparison happens, so the liveness field
660 /// on `GET /v1/runs` and the re-open refusal in
661 /// [`client_runs::open`](crate::client_runs::open) can never disagree about
662 /// whether a driver is still attached. The remaining time comes back with
663 /// the lease because the refusal has to tell the second caller when to try
664 /// again, and computing it twice would invite the two answers to drift.
665 #[must_use]
666 pub fn current_client_lease(&self, run_id: RunId) -> Option<(ClientRunLease, Duration)> {
667 let now = self.now();
668 let leases = self.inner.client_runs.lock().expect("client runs lock");
669 let lease = leases.get(&run_id)?;
670 let quiet_for = (now - lease.last_seen).unsigned_abs();
671 // `checked_sub` is the lapse test: nothing left means the TTL has run
672 // out. A remaining of exactly zero is a lapse too, which keeps this
673 // strictly-less-than, the comparison this rule has always used.
674 let remaining = self.inner.client_lease_ttl.checked_sub(quiet_for)?;
675 (!remaining.is_zero()).then(|| (lease.clone(), remaining))
676 }
677
678 /// Whether a live driver is currently attached to a client-driven run: this
679 /// process holds a lease for it AND the driver presented its token within the
680 /// lease TTL. A lapsed lease (the tab closed, the SDK exited) reports `false`:
681 /// the client-driven half of the liveness evidence `GET /v1/runs` carries.
682 #[must_use]
683 pub fn client_run_driver_live(&self, run_id: RunId) -> bool {
684 self.current_client_lease(run_id).is_some()
685 }
686
687 /// The lease for a client-driven run, if this process opened one under
688 /// `run_id`, whether or not its driver has been heard from lately. The
689 /// drive-token gate wants exactly this: a driver that went quiet for longer
690 /// than the TTL and then presents its token again is still the run's writer,
691 /// as long as nobody took the run away from it in the meantime. Ask
692 /// [`current_client_lease`](Self::current_client_lease) instead when the
693 /// question is whether someone is driving right now.
694 #[must_use]
695 pub fn client_run(&self, run_id: RunId) -> Option<ClientRunLease> {
696 self.inner
697 .client_runs
698 .lock()
699 .expect("client runs lock")
700 .get(&run_id)
701 .cloned()
702 }
703
704 /// How long a client-driven run's lease stays current without the driver
705 /// presenting its token again, the TTL every freshness judgment here is
706 /// made against. Read by the heartbeat endpoint, which answers with it so a
707 /// driver about to be busy for a while knows how often it has to beat.
708 #[must_use]
709 pub fn client_lease_ttl(&self) -> Duration {
710 self.inner.client_lease_ttl
711 }
712
713 /// Hands a client-driven run's lease back when `presented` is the token
714 /// holding it, so the next open takes the run immediately instead of
715 /// waiting out the TTL.
716 ///
717 /// The token check and the removal happen under one lock, so a release can
718 /// only ever drop the lease the caller was actually holding, never one
719 /// minted in between by a driver that took the run over.
720 ///
721 /// Only the lease goes. Nothing about the run itself changes, its recorded
722 /// `driven_by: client` included, so a later open adopts it back exactly as
723 /// it would after a restart.
724 pub fn release_client_run(&self, run_id: RunId, presented: Option<&str>) -> LeaseRelease {
725 let mut leases = self.inner.client_runs.lock().expect("client runs lock");
726 let Some(lease) = leases.get(&run_id) else {
727 return LeaseRelease::NoLease;
728 };
729 if presented != Some(lease.drive_token.as_str()) {
730 return LeaseRelease::NotTheHolder;
731 }
732 leases.remove(&run_id);
733 LeaseRelease::Released
734 }
735
736 /// Drops a client-driven run's lease unless `keep` is the token that holds
737 /// it, reporting whether a lease went. What a resolve calls: recording a
738 /// dangling write by hand says the driver that opened that write never came
739 /// back, so the lease it left behind is holding the run for nobody and the
740 /// next open should not have to wait out the TTL for it.
741 ///
742 /// `keep` is how the client-driven resolve keeps its own lease. That caller
743 /// presented its current token to get in, which is the driver saying it is
744 /// right here, so the lease is not a dead one and taking it away would
745 /// strand a driver mid-run. Every other resolve passes `None`, because no
746 /// token was presented and nothing says a driver is still attached.
747 pub fn clear_client_lease(&self, run_id: RunId, keep: Option<&str>) -> bool {
748 let mut leases = self.inner.client_runs.lock().expect("client runs lock");
749 let Some(lease) = leases.get(&run_id) else {
750 return false;
751 };
752 if keep == Some(lease.drive_token.as_str()) {
753 return false;
754 }
755 leases.remove(&run_id);
756 true
757 }
758
759 /// Whether `run_id` names a client-driven run this process opened.
760 #[must_use]
761 pub fn is_client_run(&self, run_id: RunId) -> bool {
762 self.inner
763 .client_runs
764 .lock()
765 .expect("client runs lock")
766 .contains_key(&run_id)
767 }
768
769 /// Records that the wake sweeper has warned about this run being
770 /// unwakeable here. Returns `true` the first time (the caller logs at
771 /// WARN) and `false` on every later call for the same run while the
772 /// record stands (the caller logs the same fields at DEBUG instead).
773 pub fn mark_unwakeable_warned(&self, run_id: RunId) -> bool {
774 self.inner
775 .unwakeable_warned
776 .lock()
777 .expect("unwakeable warned lock")
778 .insert(run_id)
779 }
780
781 /// Whether the sweeper has already warned about this run. Read-only,
782 /// unlike [`mark_unwakeable_warned`](Self::mark_unwakeable_warned), which
783 /// always records a sighting; a test uses this to check the record
784 /// without flipping it.
785 #[must_use]
786 pub fn unwakeable_warned(&self, run_id: RunId) -> bool {
787 self.inner
788 .unwakeable_warned
789 .lock()
790 .expect("unwakeable warned lock")
791 .contains(&run_id)
792 }
793
794 /// Clears a run's unwakeable-warned record: it woke, so the next time it
795 /// naps and cannot be rebuilt here is a fresh first sighting.
796 pub fn clear_unwakeable_warned(&self, run_id: RunId) {
797 self.inner
798 .unwakeable_warned
799 .lock()
800 .expect("unwakeable warned lock")
801 .remove(&run_id);
802 }
803
804 /// Drops every unwakeable-warned record for a run not in `still_due`.
805 /// Called once per sweep pass before processing, so a run that leaves the
806 /// due set some other way than being driven (the only other way its
807 /// record could go stale) does not carry a warning into a future nap that
808 /// has nothing to do with this one.
809 pub fn prune_unwakeable_warned(&self, still_due: &HashSet<RunId>) {
810 self.inner
811 .unwakeable_warned
812 .lock()
813 .expect("unwakeable warned lock")
814 .retain(|run_id| still_due.contains(run_id));
815 }
816
817 /// Aborts every in-flight driver task. Durability is unaffected: each event
818 /// was persisted before the task moved on, so an aborted run is recoverable
819 /// exactly as after a `kill -9`.
820 pub fn abort_all(&self) {
821 let mut handles = self.inner.handles.lock().expect("handles lock");
822 for (_, handle) in handles.drain() {
823 handle.abort();
824 }
825 self.inner.active.lock().expect("active runs lock").clear();
826 }
827}