Skip to main content

agent_engine/
host.rs

1//! Process-global engine host: the things exactly one of which should exist
2//! per process regardless of how many sessions the process serves.
3//!
4//! `engine::setup::boot` used to construct all of this inside the session's
5//! `Runtime`. It now builds (or reuses) the [`EngineHost`] and borrows the
6//! process-global parts from it; subagent spawns borrow the same parts
7//! instead of rebuilding an HTTP client + registry + token cache per spawn.
8
9use std::sync::{Arc, OnceLock};
10use tokio::sync::RwLock;
11
12use crate::extensions::hooks::HookBus;
13use crate::runtime::RuntimeParts;
14use crate::session::{SessionActor, SessionConfig, SessionHandle, SessionId, SessionMeta};
15use crate::tools::catalog::CatalogGeneration;
16use crate::{Result, Runtime, ToolRegistry};
17
18/// Everything a `Runtime` borrows from its host rather than constructing.
19/// Cheap to clone (all `Arc`/handle types).
20#[derive(Clone)]
21pub struct HostParts {
22    pub client: reqwest::Client,
23    /// THE shared catalog.
24    pub tools: Arc<RwLock<ToolRegistry>>,
25    pub hook_bus: Arc<HookBus>,
26    /// Resolved once from config.
27    pub credential_source: crate::auth::CredentialSource,
28    /// ONE cache per process.
29    pub token_cache: crate::auth::TokenCache,
30    pub mcp_runtime: Option<Arc<crate::mcp::McpRuntimeManager>>,
31    pub extension_runtime: Option<Arc<crate::extensions::lease::ExtensionRuntimeManager>>,
32    pub capture_dir: std::path::PathBuf,
33    pub progressive_tool_disclosure: bool,
34}
35
36pub struct HostOpts {
37    /// Applied via `config::set_profile` by the FIRST boot in a process
38    /// only; a later `boot()` with a different profile is not honoured
39    /// (the installed host wins — see [`EngineHost::boot_and_install`]).
40    pub profile: Option<String>,
41    /// Echoed to callers (`EngineBoot::no_extensions`) — extension discovery
42    /// is gated by the renderer after boot, not by the host. Not read here.
43    pub no_extensions: bool,
44    /// When MCP servers without cached descriptors are discovered.
45    /// [`McpDiscovery::Inline`] (in-process modes) discovers them during
46    /// `boot()`; [`McpDiscovery::Background`] (the daemon) boots with cached
47    /// descriptors only and leaves discovery to
48    /// [`EngineHost::spawn_background_mcp_discovery`], so a hung MCP server
49    /// can never hold up readiness.
50    pub mcp_discovery: McpDiscovery,
51}
52
53pub use crate::mcp::McpDiscovery;
54
55pub struct EngineHost {
56    parts: HostParts,
57    config: arc_swap::ArcSwap<crate::SynapsConfig>,
58    command_registry: Arc<crate::skills::registry::CommandRegistry>,
59    keybind_registry: Arc<std::sync::RwLock<crate::skills::keybinds::KeybindRegistry>>,
60    ext_manager: Arc<RwLock<crate::extensions::manager::ExtensionManager>>,
61    mcp_server_count: usize,
62    /// Worker (subagent) registry template, cached by shared-catalog
63    /// generation. `worker_runtime()` clones this instead of rebuilding
64    /// `without_subagent_with_extensions` + `rebuild_schema()` per spawn.
65    worker_registry: std::sync::Mutex<Option<(CatalogGeneration, ToolRegistry)>>,
66    /// File-appender flush guard. Lives for the process (≥ any renderer), so
67    /// log lines emitted after `boot()` returns can never be dropped by an
68    /// early guard drop (the old `setup::boot` bug). Rust runs no static
69    /// destructors, so it is released explicitly by [`Self::flush_logs`] at
70    /// process exit — otherwise the tail of the log dies with the writer
71    /// thread.
72    log_guard: std::sync::Mutex<Option<tracing_appender::non_blocking::WorkerGuard>>,
73    /// Live sessions (Phase 2): id → handle. The map holds one handle per
74    /// actor so the actor outlives any single client; the actor task removes
75    /// itself here when it finishes.
76    sessions: std::sync::Mutex<std::collections::HashMap<SessionId, SessionHandle>>,
77    /// Serialises `create_session`'s resolve-then-insert so two concurrent
78    /// `--continue X` cannot both miss the live check and build two actors
79    /// on one session id.
80    create_lock: tokio::sync::Mutex<()>,
81    /// (E-P2) Single-tenancy driver lock: at most one plugin can drive one
82    /// session daemon-wide. Grants do NOT survive reload (§3 S2/S5).
83    driver_grants: std::sync::Mutex<std::collections::HashMap<String, SessionId>>,
84    /// C2: flips to `true` when extension discovery on `ext_manager` has
85    /// finished (the loader sets it; `extensions_ready()` awaits it). The
86    /// paired flag records that a loader was DISPATCHED, so a session
87    /// created between `spawn_discover_and_load` and its first lock
88    /// acquisition still waits instead of racing the walk.
89    extensions_ready: tokio::sync::watch::Sender<bool>,
90    extensions_loading: std::sync::atomic::AtomicBool,
91    /// MCP discovery deferred by [`McpDiscovery::Background`]; taken (once)
92    /// by [`Self::spawn_background_mcp_discovery`].
93    pending_mcp_discovery: std::sync::Mutex<Option<crate::mcp::PendingMcpDiscovery>>,
94}
95
96/// (E-P2) Returned by [`EngineHost::claim_driver`] when another live session
97/// already holds the driver grant for a plugin.
98#[derive(Debug)]
99#[allow(dead_code)] // fields read by Display impl + P3 handler
100pub struct DriverGrantError {
101    pub plugin: String,
102    pub held_by: SessionId,
103}
104
105impl std::fmt::Display for DriverGrantError {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        write!(
108            f,
109            "plugin '{}' already driving session {}",
110            self.plugin, self.held_by
111        )
112    }
113}
114
115static HOST: OnceLock<Arc<EngineHost>> = OnceLock::new();
116
117/// See [`EngineHost::extensions_loading_guard`].
118pub struct ExtensionsLoadingGuard {
119    host: Arc<EngineHost>,
120}
121
122impl Drop for ExtensionsLoadingGuard {
123    fn drop(&mut self) {
124        self.host.mark_extensions_ready();
125    }
126}
127
128impl EngineHost {
129    /// Every step of the old `setup::boot()` that does NOT mention a session:
130    /// profile, logging, HTTP client, registry, config, skills, MCP, extension
131    /// manager. Touches NO process static: the OpenAI routing static is
132    /// written by [`Self::install`] for the winning host only, and the global
133    /// broker is installed in [`Self::foreground_runtime`] via `apply_config`,
134    /// exactly where the old boot did it relative to session resolution.
135    pub async fn boot(opts: HostOpts) -> Result<Arc<Self>> {
136        if let Some(ref prof) = opts.profile {
137            crate::config::set_profile(Some(prof.clone()));
138        }
139        let log_guard = crate::logging::init_logging();
140
141        let client = crate::runtime::build_host_http_client()?;
142        let capture_dir = crate::runtime::trace::default_capture_dir();
143        let _ = crate::runtime::trace::sweep_expired_captures(&capture_dir);
144
145        let config = crate::config::load_config();
146        let tools = Arc::new(RwLock::new(ToolRegistry::new()));
147        // Same point the old boot's `apply_config` disabled tools: on the
148        // fresh registry, before skills/MCP register anything.
149        if !config.disabled_tools.is_empty() {
150            tools.write().await.disable(&config.disabled_tools);
151        }
152
153        // Discover plugins/skills, build command registry, register load_skill tool.
154        let (command_registry, keybind_registry) = crate::skills::register(&tools, &config).await;
155
156        // MCP (if configured): descriptor-backed tools, executed through
157        // per-session leases in every mode. With progressive disclosure they
158        // stay dormant until activated; without it they are part of every
159        // session's core. A configured server with no cached descriptors is
160        // discovered once (parallel, bounded, negative-cached): here for
161        // `Inline`, later in a background task for `Background` (daemon).
162        let (mcp_server_count, pending_mcp_discovery) = match opts.mcp_discovery {
163            McpDiscovery::Inline => (crate::mcp::setup_lazy_mcp(&tools).await, None),
164            McpDiscovery::Background => crate::mcp::setup_cached_mcp(&tools).await,
165        };
166        let mcp_runtime = Some(Arc::new(crate::mcp::McpRuntimeManager::new(
167            crate::mcp::lease::config_source_from_disk(),
168            crate::mcp::lease::DEFAULT_IDLE_MAX,
169        )));
170
171        let hook_bus = Arc::new(HookBus::new());
172        let mut ext_mgr = crate::extensions::manager::ExtensionManager::new_with_tools(
173            Arc::clone(&hook_bus),
174            Arc::clone(&tools),
175        );
176        // Per-process: tell the extension manager whether memory is exclusive
177        // (Axel) so extensions spawned later inherit the right policy.
178        let mem_binding = crate::memory_backend::MemoryBinding::from_config(&config.memory_backend);
179        ext_mgr.bind_memory_backend(mem_binding.exclusive());
180        ext_mgr.set_progressive_deferral(config.progressive_tool_disclosure);
181        let extension_runtime = if config.progressive_tool_disclosure {
182            Some(ext_mgr.extension_runtime())
183        } else {
184            None
185        };
186        let ext_manager = Arc::new(RwLock::new(ext_mgr));
187
188        let parts = HostParts {
189            client,
190            tools,
191            hook_bus,
192            credential_source: config.auth.credential_source(),
193            token_cache: crate::auth::TokenCache::new(),
194            mcp_runtime,
195            extension_runtime,
196            capture_dir,
197            progressive_tool_disclosure: config.progressive_tool_disclosure,
198        };
199
200        Ok(Arc::new(Self {
201            parts,
202            config: arc_swap::ArcSwap::from_pointee(config),
203            command_registry,
204            keybind_registry,
205            ext_manager,
206            mcp_server_count,
207            worker_registry: std::sync::Mutex::new(None),
208            log_guard: std::sync::Mutex::new(log_guard),
209            sessions: std::sync::Mutex::new(std::collections::HashMap::new()),
210            create_lock: tokio::sync::Mutex::new(()),
211            driver_grants: std::sync::Mutex::new(std::collections::HashMap::new()),
212            extensions_ready: tokio::sync::watch::channel(false).0,
213            extensions_loading: std::sync::atomic::AtomicBool::new(false),
214            pending_mcp_discovery: std::sync::Mutex::new(pending_mcp_discovery),
215        }))
216    }
217
218    /// [`McpDiscovery::Background`]: start the deferred MCP discovery in a
219    /// background task (at most once per host; `None` when nothing is
220    /// pending — `Inline` boot, every server cached, or already spawned).
221    /// The daemon calls this after it is accepting. On completion ONLY the
222    /// newly discovered servers' dormant tools are added to the shared
223    /// registry, as one atomic `try_register_batch` (a rejected batch is
224    /// logged and fails closed). Like late extension registration this
225    /// mutates the shared catalog while sessions may be live: with
226    /// progressive disclosure on the tools are dormant (searchable,
227    /// activatable); with it off they are core tools that every live session
228    /// sends from its NEXT request — one prompt-cache rewrite per live
229    /// session, once. The task resolves to the number of tools registered.
230    pub fn spawn_background_mcp_discovery(&self) -> Option<tokio::task::JoinHandle<usize>> {
231        let pending = self
232            .pending_mcp_discovery
233            .lock()
234            .unwrap_or_else(|e| e.into_inner())
235            .take()?;
236        let tools = Arc::clone(&self.parts.tools);
237        tracing::info!(
238            servers = ?pending.server_names(),
239            "MCP: discovering uncached servers in the background"
240        );
241        Some(tokio::spawn(async move { pending.run(&tools).await }))
242    }
243
244    /// Idempotent, atomic install of the process host. Exactly one host ever
245    /// wins the `OnceLock`; the winner's extension manager becomes the
246    /// OpenAI routing static in the same step. A different host arriving
247    /// later (or losing a race) is returned as `Err(rejected)` and never
248    /// replaces the installed one — nor touches any static.
249    pub fn install(host: Arc<Self>) -> std::result::Result<(), Arc<Self>> {
250        let installed = HOST.get_or_init(|| {
251            crate::runtime::openai::set_extension_manager_for_routing(Arc::clone(
252                &host.ext_manager,
253            ));
254            Arc::clone(&host)
255        });
256        if Arc::ptr_eq(installed, &host) {
257            Ok(())
258        } else {
259            Err(host)
260        }
261    }
262
263    /// The process host: the installed one, else boot + install. Under a
264    /// race the loser is discarded and the installed host returned, so a
265    /// caller never ends up on a host that `current()` does not agree with.
266    pub async fn boot_and_install(opts: HostOpts) -> Result<Arc<Self>> {
267        if let Some(h) = HOST.get() {
268            return Ok(Arc::clone(h));
269        }
270        let h = Self::boot(opts).await?;
271        Ok(match Self::install(h) {
272            Ok(()) => HOST.get().cloned().expect("just installed"),
273            Err(_rejected) => HOST.get().cloned().expect("a winner exists"),
274        })
275    }
276
277    pub fn current() -> Option<Arc<Self>> {
278        HOST.get().cloned()
279    }
280
281    /// Flush the non-blocking log writer: drops the appender guard, which
282    /// drains the channel and joins the writer thread (≤ 1 s). Idempotent.
283    /// Lines logged AFTER this call are dropped, so it belongs at the very
284    /// last exit point (`main` return, panic hook), not at session shutdown.
285    pub fn flush_logs(&self) {
286        let guard = self.log_guard.lock().unwrap_or_else(|e| e.into_inner()).take();
287        drop(guard);
288    }
289
290    /// [`Self::flush_logs`] on the installed host, if any.
291    pub fn flush_installed_logs() {
292        if let Some(h) = HOST.get() {
293            h.flush_logs();
294        }
295    }
296
297    pub fn parts(&self) -> &HostParts {
298        &self.parts
299    }
300
301    pub fn config(&self) -> arc_swap::Guard<Arc<crate::SynapsConfig>> {
302        self.config.load()
303    }
304
305    pub fn ext_manager(&self) -> &Arc<RwLock<crate::extensions::manager::ExtensionManager>> {
306        &self.ext_manager
307    }
308
309    pub fn command_registry(&self) -> &Arc<crate::skills::registry::CommandRegistry> {
310        &self.command_registry
311    }
312
313    pub fn keybind_registry(
314        &self,
315    ) -> &Arc<std::sync::RwLock<crate::skills::keybinds::KeybindRegistry>> {
316        &self.keybind_registry
317    }
318
319    pub fn mcp_server_count(&self) -> usize {
320        self.mcp_server_count
321    }
322
323    /// The foreground runtime for an interactive session: shares client,
324    /// tools, hook_bus, credential source, token cache, MCP/extension lease
325    /// managers. Everything else is fresh (exactly `Runtime::new()`'s values).
326    /// `apply_config` runs here — the ONLY place the broker is (re)installed.
327    /// `disabled_tools` is NOT re-applied: `boot()` did it once on the fresh
328    /// registry, before skills/MCP registered, exactly where the old boot did.
329    pub async fn foreground_runtime(&self) -> Result<Runtime> {
330        self.foreground_runtime_for(None).await
331    }
332
333    /// `foreground_runtime()` for a session whose cwd is not the process cwd
334    /// (daemon-hosted sessions). The cwd is set BEFORE `apply_config` so the
335    /// memory binding (`apply_memory_backend_config`, one-shot) scopes to the
336    /// session's project, not the daemon's — A-engine-merge §4. `None` =
337    /// process cwd, byte-identical to `foreground_runtime()`.
338    pub async fn foreground_runtime_for(&self, cwd: Option<std::path::PathBuf>) -> Result<Runtime> {
339        let mut runtime = Runtime::from_parts(RuntimeParts::with_reaper(self.parts.clone()));
340        runtime.set_cwd(cwd);
341        runtime.apply_config_keep_tools(&self.config());
342        Ok(runtime)
343    }
344
345    /// A worker (subagent) runtime. Shares credential source and token cache
346    /// (so NO broker re-install and NO cache eviction) and takes a CLONE of
347    /// the cached worker registry template. Does NOT share: the HTTP client
348    /// (fresh — see below), hook_bus (fresh, empty — hooks do not fire for
349    /// subagent tool calls), session_manager (fresh + reaper),
350    /// mcp_runtime/extension_runtime (`None`, as `Runtime::new()`).
351    ///
352    /// Own client: workers `block_on` a throwaway `current_thread` runtime
353    /// and hyper parks each connection's I/O driver on the runtime that
354    /// opened it. A pooled connection born on a worker would be handed to
355    /// the foreground (LIFO checkout) and die mid-stream when that worker's
356    /// runtime dropped. A per-worker pool dies with its runtime, as before.
357    pub async fn worker_runtime(&self) -> Result<Runtime> {
358        let mut host = self.parts.clone();
359        host.client = crate::runtime::build_host_http_client()?;
360        host.hook_bus = Arc::new(HookBus::new());
361        host.tools = Arc::new(RwLock::new(self.worker_registry().await));
362        host.mcp_runtime = None;
363        host.extension_runtime = None;
364        // `Runtime::new()` value. Workers never ran `apply_config`, so they
365        // never had disclosure on; their registry has no discovery/activation
366        // tools, so a projected schema would silently drop every extension
367        // tool.
368        host.progressive_tool_disclosure = false;
369        Ok(Runtime::from_parts(RuntimeParts::with_reaper(host)))
370    }
371
372    /// Cached-or-rebuilt worker registry: `without_subagent_with_extensions`
373    /// is a pure function of the shared registry's tool set, which is a pure
374    /// function of its catalog generation, so caching by generation returns
375    /// exactly what a rebuild would.
376    pub async fn worker_registry(&self) -> ToolRegistry {
377        let shared = self.parts.tools.read().await;
378        let gen = shared.catalog().generation();
379        if let Some((cached_gen, tpl)) = &*self.worker_registry.lock().unwrap() {
380            if *cached_gen == gen {
381                return tpl.clone();
382            }
383        }
384        let mut fresh = ToolRegistry::without_subagent_with_extensions(&shared);
385        drop(shared);
386        // #112: apply after the merge too — a shared registry must not
387        // reintroduce a forum tool the operator disabled. Other worker policy
388        // is preserved.
389        let disabled_forum: Vec<String> = self
390            .config()
391            .disabled_tools
392            .iter()
393            .filter(|name| matches!(name.as_str(), "forum_post" | "forum_read" | "forum_forget"))
394            .cloned()
395            .collect();
396        if !disabled_forum.is_empty() {
397            fresh.disable(&disabled_forum);
398        }
399        *self.worker_registry.lock().unwrap() = Some((gen, fresh.clone()));
400        fresh
401    }
402
403    /// Re-read config from disk into the host's snapshot. Phase 1 callers
404    /// keep `load_config()` where they use it today.
405    pub fn reload_config(&self) {
406        self.config
407            .store(Arc::new(crate::config::load_config()));
408    }
409
410    // ── sessions (Phase 2) ────────────────────────────────────────────────
411
412    /// Create a session on this host: builds the `SessionActor` (the one
413    /// `Runtime` + `ConversationState`), spawns its task, registers the
414    /// handle. The task removes itself from the map when it ends.
415    ///
416    /// **Attach-if-live:** when `cfg.continue_session` resolves to a session
417    /// that is already running on this host (`--continue X` twice,
418    /// `Attach::Create` twice, `--continue` latest while latest is live),
419    /// the existing handle is returned and NO second actor is built. Two
420    /// actors on one id would share the session file, and the second's
421    /// `spawn_session_background` would unlink the first's per-session UDS
422    /// and overwrite its registry entry. `cfg`'s other fields (cwd, model
423    /// override, …) are ignored in that case — the live session wins.
424    pub async fn create_session(self: &Arc<Self>, cfg: SessionConfig) -> Result<SessionHandle> {
425        let _serial = self.create_lock.lock().await;
426        if let Some(live) = self.live_continue_target(&cfg.continue_session) {
427            tracing::info!(session = %live.id, "create_session: target is live — attaching");
428            return Ok(live);
429        }
430        let (handle, task) = SessionActor::create(self, cfg).await?;
431        let id = handle.id.clone();
432        self.adopt_session(handle.clone());
433        let host = Arc::clone(self);
434        tokio::spawn(async move {
435            // A panicking actor must still leave the map (§6 #12): the
436            // handle would otherwise be listed forever and attach → Closed.
437            use futures::FutureExt;
438            if std::panic::AssertUnwindSafe(task.run())
439                .catch_unwind()
440                .await
441                .is_err()
442            {
443                tracing::error!(session = %id, "session actor panicked");
444            }
445            host.remove_session(&id);
446        });
447        Ok(handle)
448    }
449
450    /// Register a handle built elsewhere (daemon test factories). The
451    /// caller owns the task; `remove_session` when it ends.
452    pub fn adopt_session(&self, handle: SessionHandle) {
453        self.sessions
454            .lock()
455            .unwrap_or_else(|e| e.into_inner())
456            .insert(handle.id.clone(), handle);
457    }
458
459    /// Every registered handle (B4: THE session map; `DaemonState`
460    /// delegates here). Dead handles are dropped defensively: the actor
461    /// task drops `cmd_rx` when `run()` returns and `remove_session` runs
462    /// *after* — a listing in that window sees `is_alive() == false`. That
463    /// is an ordering window, not an invariant violation, so it is logged
464    /// at debug, never asserted.
465    pub fn session_handles(&self) -> Vec<SessionHandle> {
466        let mut map = self.sessions.lock().unwrap_or_else(|e| e.into_inner());
467        let before = map.len();
468        map.retain(|_, h| h.is_alive());
469        if before != map.len() {
470            tracing::debug!(
471                pruned = before - map.len(),
472                "session_handles: pruned ended sessions ahead of their self-removal"
473            );
474        }
475        map.values().cloned().collect()
476    }
477
478    /// Resolve a `continue_session` request against the LIVE sessions
479    /// first — actor id, current journal id (a LinkedSuccessor compaction
480    /// moves it off the actor id), or name (`--name` / `saveas`, which may
481    /// not be on disk yet) — then the way `setup::resolve_or_create_session`
482    /// will (chain → name → partial id; `Some(None)` = latest by mtime) and
483    /// match the resolved id against id *and* journal id. Anything that is
484    /// not running here returns `None` so the real boot path builds/reports.
485    ///
486    /// The daemon's `Attach::Create{continue}` goes through here too
487    /// (`host_factory` → `create_session`): a live journal is always one
488    /// actor, never a second one on the same file / UDS / registry entry.
489    fn live_continue_target(
490        &self,
491        continue_session: &Option<Option<String>>,
492    ) -> Option<SessionHandle> {
493        let query = continue_session.as_ref()?;
494        let handles: Vec<SessionHandle> = {
495            let sessions = self
496                .sessions
497                .lock()
498                .unwrap_or_else(|e| e.into_inner());
499            sessions.values().filter(|h| h.is_alive()).cloned().collect()
500        };
501        if handles.is_empty() {
502            return None;
503        }
504        let by_id = |id: &str| {
505            handles
506                .iter()
507                .find(|h| h.id.as_str() == id || h.journal_id() == id)
508                .cloned()
509        };
510        let id = match query {
511            Some(q) => {
512                if let Some(h) = by_id(q) {
513                    return Some(h);
514                }
515                if let Some(h) = handles.iter().find(|h| h.name().as_deref() == Some(q.as_str())) {
516                    return Some(h.clone());
517                }
518                // F24: resolve on disk and follow compacted_into so
519                // `--continue <old>` finds the live successor.
520                let session = crate::core::session::resolve_session(q).ok()?;
521                let resolved = crate::core::session::follow_compaction_chain(session).ok()?;
522                resolved.session.id
523            }
524            None => {
525                let session = crate::core::session::latest_session().ok()?;
526                let resolved = crate::core::session::follow_compaction_chain(session).ok()?;
527                resolved.session.id
528            }
529        };
530        by_id(&id)
531    }
532
533    /// Handle for a live session, if any.
534    pub fn attach(&self, id: &SessionId) -> Option<SessionHandle> {
535        self.sessions
536            .lock()
537            .unwrap_or_else(|e| e.into_inner())
538            .get(id)
539            .cloned()
540    }
541
542    /// Metadata of every live session, with the live cells filled
543    /// (lifecycle, clients, input_owner, awaiting_input, journal_id).
544    pub fn sessions(&self) -> Vec<SessionMeta> {
545        self.session_handles()
546            .iter()
547            .map(|h| h.meta_live())
548            .collect()
549    }
550
551    /// Drop the host's handle for a session (the actor keeps running until
552    /// its command queue closes or it receives `End`).
553    pub fn remove_session(&self, id: &SessionId) -> Option<SessionHandle> {
554        let handle = self.sessions
555            .lock()
556            .unwrap_or_else(|e| e.into_inner())
557            .remove(id);
558        // E-P2: release any driver grants held by this session.
559        if handle.is_some() {
560            self.release_all_driver_grants(id);
561        }
562        handle
563    }
564
565    // ── driver single-tenancy (E-P2, §3 S2/S5) ──────────────────────────
566
567    /// Claim the driver grant for `plugin` on behalf of `session`.
568    /// Returns `Ok(())` if the grant was acquired (or is idempotent for the
569    /// same session). Returns `Err` if another live session holds it.
570    #[allow(dead_code)] // consumed by P3 driver_arm
571    pub(crate) fn claim_driver(
572        &self,
573        plugin: &str,
574        session: &SessionId,
575    ) -> std::result::Result<(), DriverGrantError> {
576        let mut grants = self.driver_grants.lock().unwrap_or_else(|e| e.into_inner());
577        if let Some(existing) = grants.get(plugin) {
578            if existing != session {
579                return Err(DriverGrantError {
580                    plugin: plugin.into(),
581                    held_by: existing.clone(),
582                });
583            }
584            return Ok(()); // idempotent
585        }
586        grants.insert(plugin.into(), session.clone());
587        Ok(())
588    }
589
590    /// Release the driver grant for `plugin` if held by `session`.
591    #[allow(dead_code)] // consumed by P3 driver_revoke
592    pub(crate) fn release_driver(&self, plugin: &str, session: &SessionId) {
593        let mut grants = self.driver_grants.lock().unwrap_or_else(|e| e.into_inner());
594        if grants.get(plugin).is_some_and(|s| s == session) {
595            grants.remove(plugin);
596        }
597    }
598
599    /// Release all driver grants held by `session` (called from `remove_session`).
600    fn release_all_driver_grants(&self, session: &SessionId) {
601        let mut grants = self.driver_grants.lock().unwrap_or_else(|e| e.into_inner());
602        grants.retain(|_, s| s != session);
603    }
604
605    /// Clear all driver grants (called on reload — §3 S2/S5: plugin process
606    /// died with the old image).
607    #[allow(dead_code)] // consumed by reload path in P3+
608    pub(crate) fn clear_all_driver_grants(&self) {
609        self.driver_grants
610            .lock()
611            .unwrap_or_else(|e| e.into_inner())
612            .clear();
613    }
614
615    // ── shared extension host (Phase 2, C) ────────────────────────────────
616
617    /// Loader seam (C2): a discovery pass for this host's manager has been
618    /// dispatched. Called synchronously by `spawn_discover_and_load` before
619    /// its task runs, so `extensions_ready()` cannot slip through the gap.
620    pub fn note_extensions_loading(&self) {
621        self.extensions_loading
622            .store(true, std::sync::atomic::Ordering::Release);
623    }
624
625    /// Loader seam (C2): discovery finished — every `extensions_ready()`
626    /// waiter proceeds. Idempotent.
627    pub fn mark_extensions_ready(&self) {
628        self.extensions_ready.send_replace(true);
629    }
630
631    /// Loader seam (C2): `note_extensions_loading()` now, and
632    /// `mark_extensions_ready()` when the guard drops — however the loader
633    /// task ends (return, cancel, panic). Hold it inside the task so a
634    /// panicking walk cannot strand `extensions_ready()` waiters.
635    pub fn extensions_loading_guard(self: &Arc<Self>) -> ExtensionsLoadingGuard {
636        self.note_extensions_loading();
637        ExtensionsLoadingGuard {
638            host: Arc::clone(self),
639        }
640    }
641
642    /// Resolve once extension discovery on this host is known-finished, so a
643    /// session's `on_session_start` lands on subscribed extensions. Resolves
644    /// IMMEDIATELY when no loader was dispatched (hosts that never load
645    /// extensions, `--no-extensions`, tests) or discovery already completed
646    /// on the manager (in-process `discover_and_load()` callers such as
647    /// `synaps chat`); otherwise awaits `mark_extensions_ready()`.
648    ///
649    /// The watch sender is a field of the host, so "sender dropped" is NOT a
650    /// wake-up path while the host lives. Liveness comes from
651    /// [`ExtensionsLoadingGuard`] (the loader task marks ready on any exit,
652    /// including panic); callers that must not trust the loader at all
653    /// (`SessionActor::create`) additionally bound the await with
654    /// `budgets::EXTENSIONS_READY_TIMEOUT`.
655    pub async fn extensions_ready(&self) {
656        let mut rx = self.extensions_ready.subscribe();
657        if *rx.borrow() {
658            return;
659        }
660        // A running walk holds the manager write lock; this read waits it out.
661        if self.ext_manager.read().await.discovery_done().is_some() {
662            return;
663        }
664        if !self
665            .extensions_loading
666            .load(std::sync::atomic::Ordering::Acquire)
667        {
668            return;
669        }
670        while !*rx.borrow_and_update() {
671            if rx.changed().await.is_err() {
672                return;
673            }
674        }
675    }
676
677    /// C3: push one extension notification frame to every live session
678    /// (non-blocking; a session whose command queue is full drops it with a
679    /// warning — widget upserts are idempotent last-writer-wins). Frames
680    /// carry no session id today, hence broadcast. Returns how many
681    /// sessions accepted it.
682    pub async fn broadcast_extension_notification(
683        &self,
684        ext_id: &str,
685        method: &str,
686        params: serde_json::Value,
687    ) -> usize {
688        let handles: Vec<SessionHandle> = self
689            .sessions
690            .lock()
691            .unwrap_or_else(|e| e.into_inner())
692            .values()
693            .cloned()
694            .collect();
695        let mut delivered = 0;
696        for handle in handles {
697            let cmd = crate::session::SessionCommand::HostEvent(
698                crate::session::HostEvent::ExtensionNotification {
699                    extension_id: ext_id.to_string(),
700                    method: method.to_string(),
701                    params: params.clone(),
702                },
703            );
704            match handle.send(cmd).await {
705                Ok(()) => delivered += 1,
706                Err(err) => tracing::warn!(
707                    session = %handle.id,
708                    extension = %ext_id,
709                    method = %method,
710                    error = %err,
711                    "dropping extension notification for session"
712                ),
713            }
714        }
715        delivered
716    }
717}
718
719#[cfg(test)]
720mod driver_grant_tests {
721    use super::*;
722
723    /// Build a minimal EngineHost for unit testing driver grants.
724    /// Only the driver_grants, sessions, and adjacent fields are initialized;
725    /// nothing boots a runtime or loads extensions.
726    fn test_host() -> Arc<EngineHost> {
727        // Use the real boot path's struct literal so the test fails to compile
728        // if a field is added and not initialised here — belt and suspenders.
729        // We can't call boot() without full config, so we manufacture one.
730        // Instead, reach into the same fields and just test the grant logic.
731        Arc::new_cyclic(|_| {
732            // Minimal construction: only the fields we touch. Others are default.
733            // This is intentionally a helper, not the full EngineHost::boot path.
734            use std::collections::HashMap;
735            use std::sync::Mutex;
736
737            // We cannot construct EngineHost without private fields, so we
738            // test the grant logic via standalone functions that mirror the
739            // Mutex<HashMap> pattern.
740            unreachable!()
741        })
742    }
743
744    // Since EngineHost has many private fields we cannot construct in a unit
745    // test without the full boot path, test the grant logic on a standalone
746    // map that mirrors the real implementation exactly.
747
748    type Grants = std::sync::Mutex<std::collections::HashMap<String, SessionId>>;
749
750    fn claim(grants: &Grants, plugin: &str, session: &SessionId) -> std::result::Result<(), DriverGrantError> {
751        let mut g = grants.lock().unwrap();
752        if let Some(existing) = g.get(plugin) {
753            if existing != session {
754                return Err(DriverGrantError {
755                    plugin: plugin.into(),
756                    held_by: existing.clone(),
757                });
758            }
759            return Ok(());
760        }
761        g.insert(plugin.into(), session.clone());
762        Ok(())
763    }
764
765    fn release(grants: &Grants, plugin: &str, session: &SessionId) {
766        let mut g = grants.lock().unwrap();
767        if g.get(plugin).is_some_and(|s| s == session) {
768            g.remove(plugin);
769        }
770    }
771
772    fn release_all(grants: &Grants, session: &SessionId) {
773        grants.lock().unwrap().retain(|_, s| s != session);
774    }
775
776    fn clear_all(grants: &Grants) {
777        grants.lock().unwrap().clear();
778    }
779
780    fn sid(s: &str) -> SessionId {
781        SessionId::from(s)
782    }
783
784    #[test]
785    fn single_tenancy_second_session_refused() {
786        let grants: Grants = Default::default();
787        let a = sid("session-a");
788        let b = sid("session-b");
789        claim(&grants, "autonomous", &a).unwrap();
790        let err = claim(&grants, "autonomous", &b).unwrap_err();
791        assert_eq!(err.plugin, "autonomous");
792        assert_eq!(err.held_by, a);
793    }
794
795    #[test]
796    fn single_tenancy_released_on_end() {
797        let grants: Grants = Default::default();
798        let a = sid("session-a");
799        let b = sid("session-b");
800        claim(&grants, "autonomous", &a).unwrap();
801        release(&grants, "autonomous", &a);
802        // Now session B can claim it.
803        claim(&grants, "autonomous", &b).unwrap();
804    }
805
806    #[test]
807    fn single_tenancy_same_session_idempotent() {
808        let grants: Grants = Default::default();
809        let a = sid("session-a");
810        claim(&grants, "autonomous", &a).unwrap();
811        // Same session, same plugin — idempotent.
812        claim(&grants, "autonomous", &a).unwrap();
813        assert_eq!(grants.lock().unwrap().len(), 1);
814    }
815
816    #[test]
817    fn grants_cleared_on_reload_record() {
818        let grants: Grants = Default::default();
819        let a = sid("session-a");
820        let b = sid("session-b");
821        claim(&grants, "plugin1", &a).unwrap();
822        claim(&grants, "plugin2", &b).unwrap();
823        assert_eq!(grants.lock().unwrap().len(), 2);
824        // Simulated reload: all grants cleared.
825        clear_all(&grants);
826        assert!(grants.lock().unwrap().is_empty());
827        // Both can re-claim.
828        claim(&grants, "plugin1", &a).unwrap();
829        claim(&grants, "plugin2", &b).unwrap();
830    }
831
832    #[test]
833    fn release_all_for_session_leaves_other_sessions() {
834        let grants: Grants = Default::default();
835        let a = sid("session-a");
836        let b = sid("session-b");
837        claim(&grants, "plugin1", &a).unwrap();
838        claim(&grants, "plugin2", &b).unwrap();
839        release_all(&grants, &a);
840        assert_eq!(grants.lock().unwrap().len(), 1);
841        assert!(grants.lock().unwrap().contains_key("plugin2"));
842    }
843}