diff --git a/crates/astrid-capsule/src/dispatcher.rs b/crates/astrid-capsule/src/dispatcher.rs index 2c2acaaf8..d6510c247 100644 --- a/crates/astrid-capsule/src/dispatcher.rs +++ b/crates/astrid-capsule/src/dispatcher.rs @@ -406,6 +406,7 @@ impl EventDispatcher { dispatch_to_capsule_queues( &capsule_queues, &self.chain_locks, + &self.event_bus, matches, topic, payload_bytes, @@ -417,6 +418,99 @@ impl EventDispatcher { } } +/// IPC topic carrying a structured "an interceptor dispatch failed" signal. +/// +/// Dispatch is fire-and-forget by design, which historically meant a failed +/// invocation was only a host-side `warn!` — the ORIGIN of the event (e.g. +/// react awaiting an `llm.v1.request.generate.*` it published) waited forever +/// with no signal, and the requester's session hung until an outer timeout. +/// The 2026-07-28 all-principals wedge rode exactly this: every LLM dispatch +/// failed instantly on a poisoned instance, and every turn hung silently. +/// +/// On a (non-`NotSupported`) invocation error, the dispatcher now publishes +/// this topic with `{ failed_topic, capsule_id, action, error, request_id?, +/// session_id? }` (the ids are best-effort correlation extracted from the +/// failed event's own payload), stamped with the failed event's principal so a +/// per-principal subscriber (react) handles it under the right KV scope and +/// can fail the owning turn loudly. Failures of THIS topic's own handling are +/// never re-signalled (loop guard in [`emit_dispatch_failed`]). +pub const DISPATCH_FAILED_TOPIC: &str = "astrid.v1.dispatch.failed"; + +/// Best-effort extraction of `request_id` / `session_id` from a failed +/// event's guest payload bytes, for the [`DISPATCH_FAILED_TOPIC`] signal. +/// +/// Looks at the top level, then one object level deep (tagged-enum payload +/// shapes nest their fields). Anything unparseable — or oversized, guarding +/// the failure path against burning CPU on a multi-megabyte prompt payload — +/// yields `(None, None)`: the signal is still published, just without +/// correlation ids. +fn extract_correlation(payload: &[u8]) -> (Option, Option) { + const MAX_PARSE_BYTES: usize = 512 * 1024; + if payload.len() > MAX_PARSE_BYTES { + return (None, None); + } + let Ok(value) = serde_json::from_slice::(payload) else { + return (None, None); + }; + fn find(value: &serde_json::Value, key: &str) -> Option { + let obj = value.as_object()?; + if let Some(s) = obj.get(key).and_then(serde_json::Value::as_str) { + return Some(s.to_string()); + } + obj.values().find_map(|nested| { + nested + .as_object() + .and_then(|o| o.get(key)) + .and_then(serde_json::Value::as_str) + .map(str::to_string) + }) + } + (find(&value, "request_id"), find(&value, "session_id")) +} + +/// Publish the [`DISPATCH_FAILED_TOPIC`] signal for a failed interceptor +/// invocation. See the topic doc for the payload contract; the principal is +/// copied from the failed event so the signal routes to the same +/// per-(capsule, principal) scope the origin published under. +fn emit_dispatch_failed( + event_bus: &EventBus, + capsule_id: &CapsuleId, + action: &str, + topic: &str, + payload: &[u8], + ipc_message: Option<&astrid_events::ipc::IpcMessage>, + error: &crate::error::CapsuleError, +) { + // Loop guard: a failure while HANDLING a dispatch-failed signal must not + // mint another signal — that would recurse for as long as the handler + // keeps failing. + if topic == DISPATCH_FAILED_TOPIC { + return; + } + let (request_id, session_id) = extract_correlation(payload); + let mut msg = astrid_events::ipc::IpcMessage::new( + astrid_events::ipc::Topic::from_raw(DISPATCH_FAILED_TOPIC), + astrid_events::ipc::IpcPayload::Custom { + data: serde_json::json!({ + "failed_topic": topic, + "capsule_id": capsule_id.as_str(), + "action": action, + "error": error.to_string(), + "request_id": request_id, + "session_id": session_id, + }), + }, + uuid::Uuid::new_v4(), + ); + if let Some(principal) = ipc_message.and_then(|m| m.principal.clone()) { + msg = msg.with_principal(principal); + } + event_bus.publish(AstridEvent::Ipc { + metadata: astrid_events::EventMetadata::new("dispatcher"), + message: msg, + }); +} + /// Dispatch matching interceptors for an event. /// /// Matches at DISTINCT priorities form an ordered middleware chain: called @@ -438,6 +532,7 @@ impl EventDispatcher { fn dispatch_to_capsule_queues( queues: &CapsuleQueues, chain_locks: &ChainLocks, + event_bus: &Arc, matches: Vec<(Arc, String, u32)>, topic: Arc, payload_bytes: Arc>, @@ -454,6 +549,7 @@ fn dispatch_to_capsule_queues( let (capsule, action, _priority) = matches.into_iter().next().unwrap(); dispatch_single( queues, + Arc::clone(event_bus), capsule, action, topic, @@ -485,6 +581,7 @@ fn dispatch_to_capsule_queues( for (capsule, action, _priority) in matches { dispatch_single( queues, + Arc::clone(event_bus), capsule, action, Arc::clone(&topic), @@ -503,6 +600,7 @@ fn dispatch_to_capsule_queues( let topic_clone = Arc::clone(&topic); let ipc_clone = ipc_message.clone(); let chain_locks_clone = Arc::clone(chain_locks); + let event_bus_clone = Arc::clone(event_bus); tokio::task::spawn(async move { let mut current_payload = (*payload_bytes).clone(); @@ -580,6 +678,18 @@ fn dispatch_to_capsule_queues( error = %e, "Interceptor invocation failed — continuing chain" ); + // Signal the origin so it can fail its request loudly + // instead of waiting out a timeout (fire-and-forget + // dispatch otherwise swallows the failure entirely). + emit_dispatch_failed( + &event_bus_clone, + capsule.id(), + action, + &topic_clone, + ¤t_payload, + ipc_clone.as_deref(), + &e, + ); // Continue chain on error — don't let a broken capsule // block the entire pipeline }, @@ -608,6 +718,7 @@ fn get_or_spawn_consumer( queues: &CapsuleQueues, capsule: &Arc, key: (CapsuleId, PrincipalKey), + event_bus: Arc, ) -> mpsc::Sender { let mut guard = queues.lock(); // Never hand back a CLOSED sender. The mapped entry can be stale: an @@ -665,7 +776,7 @@ fn get_or_spawn_consumer( let queues_arc = Arc::clone(queues); let cleanup_key = effective_key.clone(); tokio::task::spawn(async move { - run_consumer(rx, capsule_arc, queues_arc, cleanup_key).await; + run_consumer(rx, capsule_arc, queues_arc, cleanup_key, event_bus).await; }); tx } @@ -680,6 +791,7 @@ async fn run_consumer( capsule: Arc, queues: CapsuleQueues, key: (CapsuleId, PrincipalKey), + event_bus: Arc, ) { loop { match tokio::time::timeout(idle_consumer_grace(), rx.recv()).await { @@ -735,6 +847,20 @@ async fn run_consumer( error = %e, "Interceptor invocation failed" ); + // Signal the origin so it can fail its request loudly + // instead of waiting out a timeout — fire-and-forget + // dispatch otherwise swallows the failure entirely + // (the 2026-07-28 wedge: every LLM dispatch failed + // instantly, every turn hung silently). + emit_dispatch_failed( + &event_bus, + capsule.id(), + &work.action, + &work.topic, + &work.payload, + work.ipc_message.as_deref(), + &e, + ); }, } }, @@ -810,8 +936,13 @@ async fn run_consumer( /// Keying on the full `PrincipalKey` (Option) means alice's /// events don't head-of-line block bob's on the same capsule, even /// when both fall in the same `PrincipalClass` (#813 Layer 3). +// One over the 7-argument threshold: the `event_bus` rides along solely so a +// failed invocation can publish the dispatch-failed signal; bundling it into +// a struct would obscure an otherwise positional, single-call-site helper. +#[allow(clippy::too_many_arguments)] fn dispatch_single( queues: &CapsuleQueues, + event_bus: Arc, capsule: Arc, action: String, topic: Arc, @@ -820,7 +951,7 @@ fn dispatch_single( principal_key: PrincipalKey, ) { let key = (capsule.id().clone(), principal_key); - let sender = get_or_spawn_consumer(queues, &capsule, key.clone()); + let sender = get_or_spawn_consumer(queues, &capsule, key.clone(), Arc::clone(&event_bus)); let work = InterceptorWork { action, @@ -842,7 +973,7 @@ fn dispatch_single( // stall under a 100-wide prompt burst — the route's consumer closed // and every later prompt was dropped.) The re-spawn just spawned its // consumer, so the retry cannot hit the same race. - let sender = get_or_spawn_consumer(queues, &capsule, key); + let sender = get_or_spawn_consumer(queues, &capsule, key, event_bus); match sender.try_send(work) { Ok(()) => {}, // `Full` after a fresh re-spawn is the same intended shed-load diff --git a/crates/astrid-capsule/src/dispatcher_tests.rs b/crates/astrid-capsule/src/dispatcher_tests.rs index e7656ce44..0298fdedb 100644 --- a/crates/astrid-capsule/src/dispatcher_tests.rs +++ b/crates/astrid-capsule/src/dispatcher_tests.rs @@ -36,6 +36,9 @@ struct MockCapsule { principal_log: Option>>>, /// Optional shared counter incremented on every invoke. invoke_counter: Option>, + /// When set, `invoke_interceptor` returns `Err(WasmError(..))` with this + /// message — models a trapped/broken capsule for dispatch-failed tests. + error_override: Option, } impl MockCapsule { @@ -107,6 +110,7 @@ impl MockCapsule { result_override: None, principal_log: None, invoke_counter: None, + error_override: None, }; (capsule, invoked) } @@ -148,6 +152,9 @@ impl Capsule for MockCapsule { if let Some(ref c) = self.invoke_counter { c.fetch_add(1, Ordering::SeqCst); } + if let Some(ref e) = self.error_override { + return Err(crate::error::CapsuleError::WasmError(e.clone())); + } if let Some(ref result) = self.result_override { return Ok(result.clone()); } @@ -873,6 +880,7 @@ async fn dispatch_respawns_when_mapped_consumer_is_closed() { // deliver rather than hand back the dead sender and drop. dispatch_single( &queues, + Arc::new(EventBus::with_capacity(64)), Arc::clone(&capsule), "test_action".to_string(), Arc::new("respawn.topic".to_string()), @@ -1418,3 +1426,152 @@ mod access_enforcement { assert!(!gated("tool.v1.response.describe.foo")); } } + +// ── Dispatch-failed signal tests (astrid.v1.dispatch.failed) ──────────── + +/// A (non-`NotSupported`) interceptor invocation error must publish a +/// structured `astrid.v1.dispatch.failed` signal carrying the failed topic, +/// capsule, error text, the best-effort `request_id` correlation from the +/// failed event's own payload, and the failed event's principal — so a +/// per-principal subscriber (react) can fail the owning request loudly +/// instead of hanging until an outer timeout (the 2026-07-28 wedge mode). +#[tokio::test] +async fn failed_dispatch_emits_dispatch_failed_signal() { + let (mut capsule, _) = MockCapsule::new("llm-mock", "llm.v1.request.generate.mock"); + capsule.error_override = Some("wasm trap: cannot enter component instance".to_string()); + + let mut registry = CapsuleRegistry::new(); + registry.register(Box::new(capsule)).unwrap(); + let registry = Arc::new(RwLock::new(registry)); + + let bus = Arc::new(EventBus::with_capacity(64)); + let mut failed_rx = bus.subscribe_as("dispatch-failed-listener"); + let dispatcher = EventDispatcher::new(Arc::clone(®istry), Arc::clone(&bus)); + let handle = tokio::spawn(dispatcher.run()); + tokio::task::yield_now().await; + + // Shaped like react's `llm.v1.request.generate.*` publish: principal- + // tagged, request_id in the payload. + let msg = astrid_events::ipc::IpcMessage::new( + Topic::from_raw("llm.v1.request.generate.mock"), + IpcPayload::Custom { + data: serde_json::json!({ + "request_id": "11111111-2222-3333-4444-555555555555" + }), + }, + uuid::Uuid::nil(), + ) + .with_principal("alice"); + bus.publish(AstridEvent::Ipc { + metadata: astrid_events::EventMetadata::new("test"), + message: msg, + }); + + // Await the signal (skipping the original event echoing back to us). + let deadline = std::time::Instant::now() + Duration::from_secs(2); + let mut signal = None; + while std::time::Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(200), failed_rx.recv()).await { + Ok(Some(event)) => { + if let AstridEvent::Ipc { message, .. } = &*event + && message.topic.to_string() == DISPATCH_FAILED_TOPIC + { + signal = Some(message.clone()); + break; + } + }, + Ok(None) => break, + Err(_elapsed) => {}, + } + } + let signal = signal.expect("a dispatch.failed signal must be published for a failed invocation"); + assert_eq!( + signal.principal.as_deref(), + Some("alice"), + "the signal must carry the failed event's principal" + ); + let IpcPayload::Custom { data } = &signal.payload else { + panic!("dispatch.failed payload must be Custom, got {:?}", signal.payload); + }; + assert_eq!(data["failed_topic"], "llm.v1.request.generate.mock"); + assert_eq!(data["capsule_id"], "llm-mock"); + assert_eq!(data["action"], "test_action"); + assert_eq!(data["request_id"], "11111111-2222-3333-4444-555555555555"); + assert!( + data["error"] + .as_str() + .expect("error string") + .contains("cannot enter component instance"), + "error text must propagate" + ); + + handle.abort(); +} + +/// Loop guard: a failure while HANDLING `astrid.v1.dispatch.failed` itself +/// must NOT mint another signal — otherwise a persistently-failing handler +/// recurses forever. Exactly the one originally-published event is seen. +#[tokio::test] +async fn dispatch_failed_handling_failure_is_not_resignalled() { + let (mut capsule, invoked) = MockCapsule::new("broken-failed-handler", DISPATCH_FAILED_TOPIC); + capsule.error_override = Some("handler is broken".to_string()); + + let mut registry = CapsuleRegistry::new(); + registry.register(Box::new(capsule)).unwrap(); + let registry = Arc::new(RwLock::new(registry)); + + let bus = Arc::new(EventBus::with_capacity(64)); + let mut rx = bus.subscribe_as("loop-guard-listener"); + let dispatcher = EventDispatcher::new(Arc::clone(®istry), Arc::clone(&bus)); + let handle = tokio::spawn(dispatcher.run()); + tokio::task::yield_now().await; + + publish_ipc(&bus, DISPATCH_FAILED_TOPIC); + tokio::time::sleep(Duration::from_millis(300)).await; + assert!(invoked.load(Ordering::SeqCst), "handler must have run (and failed)"); + + // Count dispatch.failed events on the bus: exactly the one we published. + let mut seen = 0; + while let Ok(Some(event)) = + tokio::time::timeout(Duration::from_millis(100), rx.recv()).await + { + if let AstridEvent::Ipc { message, .. } = &*event + && message.topic.to_string() == DISPATCH_FAILED_TOPIC + { + seen += 1; + } + } + assert_eq!( + seen, 1, + "a failure handling dispatch.failed must not be re-signalled" + ); + + handle.abort(); +} + +/// `extract_correlation` finds ids at the top level and one object level deep +/// (tagged-enum payload shapes), and gives up cleanly on garbage/oversized +/// payloads. +#[test] +fn extract_correlation_shapes() { + let top = serde_json::to_vec(&serde_json::json!({ + "request_id": "r-1", "session_id": "s-1" + })) + .unwrap(); + assert_eq!( + extract_correlation(&top), + (Some("r-1".to_string()), Some("s-1".to_string())) + ); + + let nested = serde_json::to_vec(&serde_json::json!({ + "type": "llm_request", + "inner": { "request_id": "r-2" } + })) + .unwrap(); + assert_eq!(extract_correlation(&nested), (Some("r-2".to_string()), None)); + + assert_eq!(extract_correlation(b"not json"), (None, None)); + + let oversized = vec![b' '; 600 * 1024]; + assert_eq!(extract_correlation(&oversized), (None, None)); +} diff --git a/crates/astrid-capsule/src/engine/wasm/mod.rs b/crates/astrid-capsule/src/engine/wasm/mod.rs index 73baa9c5e..67bcc4fbf 100644 --- a/crates/astrid-capsule/src/engine/wasm/mod.rs +++ b/crates/astrid-capsule/src/engine/wasm/mod.rs @@ -2148,9 +2148,13 @@ impl ExecutionEngine for WasmEngine { // busy), distinct from a slow guest call. let pool_wait_ms = checkout_start.elapsed().as_millis() as u64; let typed_instance = checkout.instance(); - let result: CapsuleResult = { + // ── Phase 1: SET ────────────────────────────────────── + // + // Host-side only — nothing here enters the guest, so an early exit + // (error or future-drop) in this phase leaves the instance clean and + // `PoolCheckout::drop` returns it to the pool un-armed. + { let s = checkout.store_mut(); - // ── Phase 1: SET ────────────────────────────────────── let applied_profile: Arc = invocation_profile.clone().unwrap_or_else(|| { Arc::new(astrid_core::profile::PrincipalProfile::default_ref().clone()) @@ -2256,14 +2260,23 @@ impl ExecutionEngine for WasmEngine { } } - // ── Phase 2: CALL ───────────────────────────────────── - // - // Cancellation safety: the `call_async` future below may be - // dropped by the dispatcher (e.g. tokio task abort). Dropping it - // drops `checkout`, whose `Drop` synchronously runs Phase 3 CLEAR - // *before* the wasm fiber is torn down and returns the instance to - // the pool, so the next lease observes `caller_context = None` and - // every `invocation_*` field cleared. + } + + // ── Phase 2: CALL ───────────────────────────────────── + // + // Cancellation safety: the `call_async` future below may be + // dropped by the dispatcher (e.g. tokio task abort). Dropping it + // drops `checkout`, whose `Drop` synchronously runs — and because the + // checkout is ARMED across the guest call, a future-drop mid-guest + // DISCARDS the instance rather than returning it: wasmtime poisons a + // component instance on any non-clean exit (trap, cancel, panic), and + // a returned poisoned instance would be re-leased first (LIFO) and + // permanently capture the capsule's invocation path (the 2026-07-28 + // all-principals agent wedge). Disarm happens below only on a clean + // exit; see `PoolCheckout::arm`/`disarm`. + checkout.arm(); + let result: CapsuleResult = { + let s = checkout.store_mut(); let typed_lookup = typed_instance .get_typed_func::<(String, Vec), (HookTriggerResult,)>( &mut *s, @@ -2282,6 +2295,16 @@ impl ExecutionEngine for WasmEngine { ))), } }; + match &result { + // Clean guest exit — the instance is re-enterable; return it. + Ok(_) => checkout.disarm(), + // Export lookup failed BEFORE the guest was entered — clean. + Err(CapsuleError::UnsupportedEntryPoint(_)) => checkout.disarm(), + // Trap or host-error propagated out of `call_async`: the component + // instance is poisoned (`cannot enter component instance` on every + // later call). Stay armed — `PoolCheckout::drop` discards it. + Err(_) => {}, + } // Per-invocation CPU measurement: fuel counts DOWN from the seed, so // `seed - remaining` is the exact deterministic instruction count for // this call. Read while `checkout` is still alive (the `s` borrow above diff --git a/crates/astrid-capsule/src/engine/wasm/pool.rs b/crates/astrid-capsule/src/engine/wasm/pool.rs index 211c6eca4..5d04e43c6 100644 --- a/crates/astrid-capsule/src/engine/wasm/pool.rs +++ b/crates/astrid-capsule/src/engine/wasm/pool.rs @@ -165,11 +165,12 @@ pub(super) struct CapsuleInstancePool { reset_resources_on_return: bool, /// On-demand instance factory for lazy growth. builder: Arc, - /// Whether a checkout that finds no warm instance may build one. `false` - /// for the size-1 `host_process` carve-out (`max == min_idle == 1`): its - /// single instance is always warm, so this is belt-and-suspenders — if a - /// build were ever reached it would mint a *second* Store and violate the - /// carve-out, so we fail closed instead. + /// Whether this pool can hold more instances than its warm set + /// (`max > min_idle`). Gates the idle evictor and the log level of a + /// checkout-time build. NOTE: a checkout that finds no warm instance + /// builds one regardless — the held permit already guarantees the total + /// stays ≤ `max` — which is how a size-1 carve-out recovers after an + /// armed discard destroyed its only instance (see [`PoolCheckout`]). allow_grow: bool, /// Idle-eviction timer; aborted on drop. `None` when the pool cannot grow /// (`max == min_idle`) — `available` can then never exceed `min_idle`, so @@ -221,10 +222,11 @@ impl CapsuleInstancePool { /// Lease an instance, awaiting a permit if `max` are already in use. /// /// With a permit in hand the pool is below `max`, so this pops a warm - /// instance or — when none is warm — builds a fresh one (lazy grow). - /// Returns `None` if the semaphore is closed (capsule unloading), if a - /// lazy build fails, or if a non-growable pool somehow finds no warm - /// instance — all treated by the caller as "not invocable". + /// instance or — when none is warm — builds a fresh one (lazy grow, or a + /// rebuild after an armed discard; the permit guarantees the total never + /// exceeds `max`, so this is sound even for the size-1 carve-out). + /// Returns `None` if the semaphore is closed (capsule unloading) or if a + /// build fails — both treated by the caller as "not invocable". pub(super) async fn checkout(&self) -> Option { let permit = Arc::clone(&self.permits).acquire_owned().await.ok()?; // Pop the most-recently-returned instance (the BACK — return pushes @@ -240,10 +242,22 @@ impl CapsuleInstancePool { let pooled = match warm { Some(pooled) => pooled, None => { + // Build a replacement. Sound for EVERY pool, including the + // size-1 `host_process` carve-out: we hold a permit, so total + // instances = in-flight-under-other-permits ≤ max - 1, and + // building keeps the count ≤ max. For the carve-out this + // branch was historically unreachable ("its instance is + // always warm") and failed closed against minting a second + // Store — but an armed discard (see [`PoolCheckout`]) can now + // legitimately leave `available` empty at count 0, and + // refusing to rebuild would leave the capsule permanently + // un-invocable ("no capsule instance available" forever). The + // permit argument above is exactly the no-second-Store + // guarantee the fail-closed branch existed to protect. if !self.allow_grow { - // Unreachable for a size-1 carve-out (its instance is - // always warm); fail closed rather than mint a second Store. - return None; + tracing::info!( + "rebuilding discarded instance for non-growable pool" + ); } match self.builder.build().await { Ok(pooled) => pooled, @@ -258,6 +272,7 @@ impl CapsuleInstancePool { pooled: Some(pooled), available: Arc::clone(&self.available), reset_resources_on_return: self.reset_resources_on_return, + armed: false, _permit: permit, }) } @@ -333,12 +348,39 @@ fn drain_excess(queue: &mut VecDeque, min_idle: usize) -> Vec { /// Folding the clear into the return guarantees the next lease of this /// instance observes a clean `HostState`, and that no instance (or permit) is /// leaked on an error path. +/// +/// ## Armed discard: a non-clean guest exit poisons the instance +/// +/// Wasmtime **poisons a component instance on any trap**: the instance's +/// reentrance flag stays cleared, and every later call on it fails with +/// `cannot enter component instance`. The same applies when the `call_async` +/// future is dropped mid-guest (caller cancellation) or the invocation +/// panics. Returning such an instance to the pool is fatal in combination +/// with LIFO checkout: the poisoned instance is the most-recently-returned, +/// so it is re-leased first, fails instantly, is returned again, and +/// permanently captures the capsule's entire invocation path (the 2026-07-28 +/// all-principals agent wedge — one epoch-interrupted LLM stream poisoned +/// `openai-compat` until a daemon restart). +/// +/// The caller therefore [`arm`](Self::arm)s the checkout immediately before +/// entering the guest and [`disarm`](Self::disarm)s it only on a clean exit +/// (an `Ok` return from the guest call, or an export-lookup failure where the +/// guest was never entered). `Drop` of an **armed** checkout DISCARDS the +/// instance — dropping its Store — instead of returning it; the released +/// permit lets a later checkout lazily rebuild a fresh instance, so one bad +/// invocation costs exactly one instance, never the capsule. pub(super) struct PoolCheckout { pooled: Option, available: Arc>>, /// Mirrors [`CapsuleInstancePool::reset_resources_on_return`]; copied at /// checkout so the drop path needs no back-pointer to the pool. reset_resources_on_return: bool, + /// Whether the guest call is (still) in a possibly-non-clean state. Set by + /// [`arm`](Self::arm) just before guest entry, cleared by + /// [`disarm`](Self::disarm) on clean exit. While `true`, `Drop` discards + /// the instance instead of returning it — covering the trap path, the + /// panic-unwind path, and future-drop on caller cancellation alike. + armed: bool, _permit: OwnedSemaphorePermit, } @@ -354,11 +396,51 @@ impl PoolCheckout { pub(super) fn store_mut(&mut self) -> &mut Store { &mut self.pooled.as_mut().expect("active checkout").store } + + /// Mark the guest call as in flight: from now until [`disarm`](Self::disarm), + /// dropping this checkout DISCARDS the instance instead of returning it. + /// + /// Call immediately before entering the guest (`call_async`). Everything + /// before guest entry (the SET phase) is host-side only and cannot poison + /// the instance, so an early exit there still returns the instance. + pub(super) fn arm(&mut self) { + self.armed = true; + } + + /// Mark the guest call as cleanly exited: dropping this checkout returns + /// the instance to the pool again. + /// + /// Only call when the instance is provably re-enterable — an `Ok` return + /// from the guest call, or an export lookup that failed *before* the guest + /// was entered. Any trap or host-error propagated out of `call_async` + /// leaves the instance poisoned (`cannot enter component instance` on + /// every later call) and MUST stay armed. + pub(super) fn disarm(&mut self) { + self.armed = false; + } } impl Drop for PoolCheckout { fn drop(&mut self) { if let Some(mut pooled) = self.pooled.take() { + // Armed = the guest call did not exit cleanly (trap, panic-unwind, + // or future-drop mid-call). Wasmtime poisons the component + // instance on any of those — every later call on it would fail + // with `cannot enter component instance` — and LIFO checkout would + // re-lease it FIRST, permanently capturing the capsule's + // invocation path (the 2026-07-28 agent wedge). Discard it: drop + // the Store (also closing any orphaned resources it still holds) + // and let the released permit lazily rebuild a replacement on a + // later checkout. + if self.armed { + tracing::warn!( + capsule_id = %pooled.store.data().capsule_id, + "discarding pooled instance after non-clean guest exit \ + (trap/cancel/panic poisons the component instance)" + ); + drop(pooled); + return; + } // Phase 3: CLEAR. Reset every per-invocation field before the // instance returns to the pool so the next lease starts clean. // Mirrors the old `ClearOnDrop` guard from the single-Store path. @@ -707,4 +789,84 @@ mod tests { drop(c2); cancel.cancel(); } + + /// An ARMED checkout (the guest call did not exit cleanly — trap, panic, + /// or future-drop) must be DISCARDED on drop, never returned; a later + /// checkout must rebuild a fresh instance rather than block or fail. + /// + /// Regression test for the 2026-07-28 all-principals agent wedge: an + /// epoch-interrupt trap poisoned an `openai-compat` instance, the return + /// path pushed it back, and LIFO checkout re-leased the poisoned instance + /// to every subsequent LLM invocation (`cannot enter component instance`) + /// until a daemon restart. + #[tokio::test(flavor = "multi_thread")] + async fn armed_checkout_is_discarded_and_pool_rebuilds() { + let cancel = CancellationToken::new(); + let pool = empty_pool(2, 1, &cancel).await; + + // Lease the single warm instance and simulate a non-clean guest exit. + let mut c1 = pool.checkout().await.expect("warm instance"); + c1.arm(); + drop(c1); + + // The instance must NOT have been returned to the warm set. + assert_eq!( + pool.available.lock().expect("pool mutex").len(), + 0, + "an armed checkout must be discarded, not returned" + ); + + // The pool must recover: a later checkout rebuilds a fresh instance. + let c2 = tokio::time::timeout(Duration::from_millis(1000), pool.checkout()) + .await + .expect("checkout after a discard must not block") + .expect("pool must rebuild a fresh instance after a discard"); + drop(c2); + // The rebuilt instance exited cleanly (never armed) → returned. + assert_eq!(pool.available.lock().expect("pool mutex").len(), 1); + cancel.cancel(); + } + + /// `arm` followed by `disarm` (a clean guest exit) keeps the pre-existing + /// return-to-pool behavior byte-identical. + #[tokio::test(flavor = "multi_thread")] + async fn disarmed_checkout_returns_to_pool() { + let cancel = CancellationToken::new(); + let pool = empty_pool(2, 1, &cancel).await; + + let mut c1 = pool.checkout().await.expect("warm instance"); + c1.arm(); + c1.disarm(); // clean guest exit + drop(c1); + + assert_eq!( + pool.available.lock().expect("pool mutex").len(), + 1, + "a disarmed (clean-exit) checkout must return to the pool" + ); + cancel.cancel(); + } + + /// The size-1 `host_process` carve-out must REBUILD after an armed + /// discard destroyed its only instance — refusing to build (the old + /// fail-closed branch) would leave the capsule permanently un-invocable + /// ("no capsule instance available" forever). Holding the single permit + /// guarantees the rebuild never coexists with another live Store, which + /// is the invariant the fail-closed branch existed to protect. + #[tokio::test(flavor = "multi_thread")] + async fn carveout_pool_rebuilds_after_discard() { + let cancel = CancellationToken::new(); + let pool = empty_pool(1, 1, &cancel).await; + + let mut c1 = pool.checkout().await.expect("the one instance"); + c1.arm(); + drop(c1); // discards the carve-out's only instance + + let c2 = tokio::time::timeout(Duration::from_millis(1000), pool.checkout()) + .await + .expect("carve-out checkout after a discard must not block") + .expect("carve-out must rebuild its single instance after a discard"); + drop(c2); + cancel.cancel(); + } } diff --git a/crates/astrid-kernel/src/lib.rs b/crates/astrid-kernel/src/lib.rs index 18447a4cd..9ef8f417d 100644 --- a/crates/astrid-kernel/src/lib.rs +++ b/crates/astrid-kernel/src/lib.rs @@ -1876,37 +1876,175 @@ fn spawn_capsule_health_monitor(kernel: Arc) -> tokio::task::JoinHandle< }) } +/// How long after its last `user.v1.prompt` a principal keeps receiving +/// per-principal watchdog ticks. Comfortably beyond the longest react phase +/// timeout, so a stuck turn is still being ticked when it crosses its +/// deadline; after the TTL the principal simply stops costing a tick. +const WATCHDOG_PRINCIPAL_TTL: std::time::Duration = std::time::Duration::from_secs(900); + +/// Bound on the watchdog's recently-active-principal set. Beyond it the +/// stalest entry is evicted — an evicted principal only loses the recovery +/// *backstop* until its next prompt, never correctness. +const WATCHDOG_MAX_PRINCIPALS: usize = 256; + +/// Record a principal sighting in the watchdog's active set, evicting the +/// stalest entry when the (bounded) set is full. Factored out of the task +/// loop so the eviction discipline is unit-testable. +fn watchdog_note_principal( + active: &mut std::collections::HashMap, + principal: String, + now: std::time::Instant, + cap: usize, +) { + if principal.is_empty() { + return; + } + if active.len() >= cap + && !active.contains_key(&principal) + && let Some(stalest) = active + .iter() + .min_by_key(|(_, seen)| **seen) + .map(|(k, _)| k.clone()) + { + active.remove(&stalest); + } + active.insert(principal, now); +} + /// Spawns a periodic watchdog that publishes `astrid.v1.watchdog.tick` events every 5 seconds. /// /// The `ReAct` capsule (WASM guest) cannot use async timers, so this kernel-side task /// drives timeout enforcement by waking the capsule on a fixed interval. Each tick /// causes the capsule's `handle_watchdog_tick` interceptor to run `check_phase_timeout`. +/// +/// ## Per-principal fan-out (2026-07-28 wedge follow-up) +/// +/// The tick used to be published WITHOUT a principal, so react's +/// `handle_watchdog_tick` always ran under the capsule OWNER's KV scope and +/// could only see the owner's `react.active_sessions` index — every +/// per-principal session (each gateway-minted bearer) was invisible to +/// phase-timeout recovery, and a per-principal turn that lost its LLM stream +/// hung forever instead of failing at the phase timeout. The task now watches +/// `user.v1.prompt` on the bus and re-publishes each tick once per RECENTLY +/// ACTIVE principal (a prompt within [`WATCHDOG_PRINCIPAL_TTL`]), stamped +/// with that principal, so the dispatcher runs react's tick handler under the +/// right per-(capsule, principal) scope and the existing recovery logic sees +/// that principal's sessions. The principal-less tick is kept for owner-scope +/// sessions; the active set is bounded by [`WATCHDOG_MAX_PRINCIPALS`]. fn spawn_react_watchdog(event_bus: Arc) -> tokio::task::JoinHandle<()> { + fn publish_tick(event_bus: &EventBus, principal: Option) { + let mut msg = astrid_events::ipc::IpcMessage::new( + astrid_events::ipc::Topic::from_raw("astrid.v1.watchdog.tick"), + astrid_events::ipc::IpcPayload::Custom { + data: serde_json::json!({}), + }, + uuid::Uuid::new_v4(), + ); + if let Some(principal) = principal { + msg = msg.with_principal(principal); + } + let _ = event_bus.publish(astrid_events::AstridEvent::Ipc { + metadata: astrid_events::EventMetadata::new("kernel"), + message: msg, + }); + } + tokio::spawn(async move { + // Broadcast subscription used ONLY to observe `user.v1.prompt` + // principals; drained continuously in the select loop below, so it + // cannot meaningfully lag (a lagged/lost sighting merely delays that + // principal's backstop until its next prompt). + let mut receiver = event_bus.subscribe_as("react_watchdog"); let mut interval = tokio::time::interval(std::time::Duration::from_secs(5)); // The first tick fires immediately - skip it to give capsules time to load. interval.tick().await; - loop { - interval.tick().await; - metrics::counter!(METRIC_BACKGROUND_TICKS_TOTAL, "loop" => "react_watchdog") - .increment(1); + let mut active: std::collections::HashMap = + std::collections::HashMap::new(); - let msg = astrid_events::ipc::IpcMessage::new( - astrid_events::ipc::Topic::from_raw("astrid.v1.watchdog.tick"), - astrid_events::ipc::IpcPayload::Custom { - data: serde_json::json!({}), + loop { + tokio::select! { + event = receiver.recv() => { + let Some(event) = event else { + // Bus closed — daemon shutting down; ticks are moot. + return; + }; + if let astrid_events::AstridEvent::Ipc { message, .. } = &*event + && message.topic.as_str() == "user.v1.prompt" + && let Some(principal) = message.principal.clone() + { + watchdog_note_principal( + &mut active, + principal, + std::time::Instant::now(), + WATCHDOG_MAX_PRINCIPALS, + ); + } }, - uuid::Uuid::new_v4(), - ); - let _ = event_bus.publish(astrid_events::AstridEvent::Ipc { - metadata: astrid_events::EventMetadata::new("kernel"), - message: msg, - }); + _ = interval.tick() => { + metrics::counter!(METRIC_BACKGROUND_TICKS_TOTAL, "loop" => "react_watchdog") + .increment(1); + active.retain(|_, seen| seen.elapsed() < WATCHDOG_PRINCIPAL_TTL); + + // Owner-scope tick — the pre-existing behaviour. + publish_tick(&event_bus, None); + // Per-principal fan-out for recently-active principals. + for principal in active.keys() { + publish_tick(&event_bus, Some(principal.clone())); + } + }, + } } }) } +#[cfg(test)] +mod watchdog_tests { + use super::watchdog_note_principal; + use std::collections::HashMap; + use std::time::{Duration, Instant}; + + /// The active set stays bounded: at capacity, noting a NEW principal + /// evicts the stalest entry; re-noting an existing principal only + /// refreshes its timestamp. + #[test] + fn note_principal_bounds_and_evicts_stalest() { + let mut active: HashMap = HashMap::new(); + let t0 = Instant::now(); + + watchdog_note_principal(&mut active, "alice".into(), t0, 2); + watchdog_note_principal( + &mut active, + "bob".into(), + t0.checked_add(Duration::from_secs(1)).expect("time"), + 2, + ); + // Refresh alice — she is now the NEWEST. + watchdog_note_principal( + &mut active, + "alice".into(), + t0.checked_add(Duration::from_secs(2)).expect("time"), + 2, + ); + // At cap: carol evicts the stalest (bob, not alice). + watchdog_note_principal( + &mut active, + "carol".into(), + t0.checked_add(Duration::from_secs(3)).expect("time"), + 2, + ); + + assert_eq!(active.len(), 2, "set must stay at cap"); + assert!(active.contains_key("alice"), "refreshed entry survives"); + assert!(active.contains_key("carol"), "new entry admitted"); + assert!(!active.contains_key("bob"), "stalest entry evicted"); + + // Empty principals are never tracked. + watchdog_note_principal(&mut active, String::new(), t0, 2); + assert_eq!(active.len(), 2); + } +} + // --------------------------------------------------------------------------- // Boot validation // ---------------------------------------------------------------------------