cookie_agent_engine/runtime/
sessions.rs

1use std::{
2    collections::{BTreeMap, HashMap, HashSet},
3    sync::Arc,
4};
5
6use cookie_agent_protocol::{
7    AgentId, AgentUsageResult, ChildSummary, EventOrigin, GlobalUsageResult, InvocationId,
8    PermissionMode, RunSelection, SessionForkResult, SessionId, SessionMeta, SessionOrigin,
9    SessionRenameChange, SessionRenameParams, SessionRenameResult, SessionRevertResult,
10    SessionTreeUsageResult, SessionUsageResult, UsageRollup,
11};
12
13use super::{
14    Engine, EngineError, Event, SessionCommand,
15    delegation::DELEGATION_RECOVERY_SETTLE_BOUND,
16    helpers::{cwd_identity, root_id, session_depth},
17};
18use crate::policy::{self, freeze_root_agent_policy, resolve_agent};
19
20impl Engine {
21    /// The accepted selection, advanced only by a committed fallback model turn.
22    pub fn session_model_selection(&self, session: SessionId) -> Result<RunSelection, EngineError> {
23        let projection = self.inner.store.get(session)?;
24        Ok(
25            cookie_agent_protocol::SessionModelState::from_events(&projection.log.event_snapshot())
26                .selection
27                .unwrap_or(projection.meta.creation_selection),
28        )
29    }
30    /// Rebuilds the visible grant set from root logs plus the grant sets cached
31    /// by loaded trees, so it stays O(cached data) instead of O(all logs) (§4.3).
32    pub(super) fn rebuild_visible_tree_grants(&self) {
33        let invalidated = self.inner.grant_journal.invalidated_ids();
34        let mut grants = Vec::new();
35        for session in self.inner.store.root_snapshots() {
36            for event in session.log.event_snapshot().iter() {
37                if let Event::TreeApprovalGrantCommitted { grant } = &event.payload
38                    && !invalidated.contains(&grant.grant_id)
39                {
40                    grants.push(grant.clone());
41                }
42            }
43        }
44        grants.extend(
45            self.inner
46                .store
47                .loaded_tree_grants()
48                .into_iter()
49                .filter(|grant| !invalidated.contains(&grant.grant_id)),
50        );
51        self.inner.approvals.store.replace(grants);
52    }
53
54    pub fn create_session(&self, selection: RunSelection) -> Result<SessionMeta, EngineError> {
55        self.reconcile_provider_store()?;
56        let runtime = self.current_runtime();
57        if runtime.result.snapshot.models.is_empty() {
58            return Err(EngineError::NoRunnableModel);
59        }
60        let agents = runtime.agents_for_preset(selection.preset.as_deref())?;
61        let agent = resolve_agent(&agents, &selection.agent)?.clone();
62        let policy = freeze_root_agent_policy(
63            &agent,
64            agents,
65            Arc::clone(&runtime),
66            &selection.model,
67            self.inner.config.runtime.delegation.max_depth,
68            policy::ResultLimits {
69                tool_output_max_lines: self.inner.config.runtime.tool_output.max_lines,
70                tool_output_max_bytes: self.inner.config.runtime.tool_output.max_bytes,
71            },
72            self.inner.config.runtime.model_retry,
73        )?;
74        let id = SessionId::new_v7();
75        let cwd_identity = cwd_identity(self.inner.store.cwd())?;
76        // A new root begins an empty tree, so its first handle is trivially
77        // unique; it still goes through the same generation path so every
78        // session carries a handle uniformly.
79        let short_id = super::handles::generate_handle(
80            selection.agent.as_str(),
81            &HashSet::new(),
82            super::handles::random_hex,
83        )?;
84        let creation = Event::SessionCreated {
85            origin: SessionOrigin::Root,
86            short_id: Some(short_id),
87            cwd_identity: cwd_identity.clone(),
88            creation_selection: selection.clone(),
89            creation_agent: Box::new(policy.agent.clone()),
90            runtime_revision: runtime.result.snapshot.runtime_revision.clone(),
91            catalog_revision: runtime.result.snapshot.catalog_revision.clone(),
92            provider_state_revision: runtime.result.snapshot.provider_state_revision.clone(),
93            model_revision: runtime.result.snapshot.model_revision.clone(),
94            agent_revision: runtime.result.snapshot.agent_revision.clone(),
95            recipe_registry_revision: runtime.result.snapshot.recipe_registry_revision.clone(),
96            manifest_revision: runtime.current_manifest.revision.clone(),
97        };
98        self.inner.store.create(
99            id,
100            EventOrigin::new("user").expect("static event origin is valid"),
101            creation,
102        )?;
103        self.spawn_actor(id)?;
104        Ok(self.inner.store.get(id)?.metadata())
105    }
106
107    /// Top-level session listing: delegation roots only. Delegated children are
108    /// selected through their root's tree and stay reachable via
109    /// [`Engine::get_session`], [`Engine::children`], [`Engine::tree`] and
110    /// resume.
111    #[must_use]
112    pub fn list_sessions(&self) -> Vec<SessionMeta> {
113        self.inner
114            .store
115            .all_summaries()
116            .into_iter()
117            .map(|session| session.meta)
118            .filter(|meta| matches!(meta.origin, SessionOrigin::Root))
119            .collect()
120    }
121    pub fn get_session(&self, id: SessionId) -> Result<SessionMeta, EngineError> {
122        let session = self.inner.store.get(id)?;
123        #[cfg(test)]
124        let hook = self
125            .inner
126            .test_hooks
127            .read_only_reopen_hook
128            .lock()
129            .expect("read-only reopen hook lock poisoned")
130            .take();
131        #[cfg(test)]
132        if let Some(hook) = hook {
133            if let Some(reached) = hook
134                .reached
135                .lock()
136                .expect("read-only reopen reached lock poisoned")
137                .take()
138            {
139                let _ = reached.send(());
140            }
141            let _ = hook
142                .release
143                .lock()
144                .expect("read-only reopen release lock poisoned")
145                .recv();
146        }
147        Ok(session.metadata())
148    }
149
150    #[cfg(test)]
151    pub(crate) fn install_read_only_reopen_hook_for_test(
152        &self,
153    ) -> (
154        tokio::sync::oneshot::Receiver<()>,
155        std::sync::mpsc::Sender<()>,
156    ) {
157        let (reached, receiver) = tokio::sync::oneshot::channel();
158        let (release, release_receiver) = std::sync::mpsc::channel();
159        *self
160            .inner
161            .test_hooks
162            .read_only_reopen_hook
163            .lock()
164            .expect("read-only reopen hook lock poisoned") = Some(super::ReadOnlyReopenHook {
165            reached: std::sync::Mutex::new(Some(reached)),
166            release: std::sync::Mutex::new(release_receiver),
167        });
168        (receiver, release)
169    }
170    pub fn session_usage(&self, id: SessionId) -> Result<SessionUsageResult, EngineError> {
171        let catalog = self.catalog_pricing();
172        Ok(self
173            .inner
174            .store
175            .session_usage(id, &self.inner.config.runtime.pricing, &catalog)?)
176    }
177    pub fn agent_usage(&self, agent_id: AgentId) -> AgentUsageResult {
178        let mut usage = UsageRollup::default();
179        for session in self.inner.store.all_summaries() {
180            if let Some(agent) = session.agent_usage.get(&agent_id) {
181                crate::usage::merge(&mut usage, agent);
182            }
183        }
184        AgentUsageResult {
185            agent_id,
186            usage: crate::usage::with_pricing(
187                usage,
188                &self.inner.config.runtime.pricing,
189                &self.catalog_pricing(),
190            ),
191        }
192    }
193    pub fn global_usage(&self) -> GlobalUsageResult {
194        let mut usage = UsageRollup::default();
195        for session in self.inner.store.all_summaries() {
196            crate::usage::merge(&mut usage, &session.usage_rollup);
197        }
198        GlobalUsageResult {
199            usage: crate::usage::with_pricing(
200                usage,
201                &self.inner.config.runtime.pricing,
202                &self.catalog_pricing(),
203            ),
204        }
205    }
206    pub fn session_tree_usage(&self, id: SessionId) -> Result<SessionTreeUsageResult, EngineError> {
207        // A cold root's index is only a cache. Complete the tree before using
208        // its summaries, otherwise a missing or stale index can produce a
209        // plausible but incomplete rollup.
210        self.inner.store.ensure_tree_for(id)?;
211        let summaries = self.inner.store.all_summaries();
212        let mut children: HashMap<SessionId, Vec<(SessionId, InvocationId)>> = HashMap::new();
213        let mut rollups = HashMap::new();
214        for summary in summaries {
215            let session_id = summary.meta.session_id;
216            if let SessionOrigin::Delegated {
217                root_session_id,
218                parent_session_id,
219                invocation_id,
220                ..
221            } = summary.meta.origin
222                && root_session_id == id
223            {
224                let edge = (session_id, invocation_id);
225                children.entry(parent_session_id).or_default().push(edge);
226            }
227            rollups.insert(session_id, summary.usage_rollup);
228        }
229
230        let mut usage = UsageRollup::default();
231        let session_count = merge_session_tree_usage(id, &children, &rollups, &mut usage)?;
232        Ok(SessionTreeUsageResult {
233            session_id: id,
234            usage: crate::usage::with_pricing(
235                usage,
236                &self.inner.config.runtime.pricing,
237                &self.catalog_pricing(),
238            ),
239            session_count,
240        })
241    }
242    pub(super) fn catalog_pricing(
243        &self,
244    ) -> BTreeMap<cookie_agent_protocol::ModelKey, cookie_agent_models::catalog::CatalogModelCost>
245    {
246        self.inner
247            .model_manager
248            .current()
249            .models()
250            .iter()
251            .filter_map(|(model, runtime)| {
252                runtime
253                    .model
254                    .cost
255                    .as_ref()
256                    .map(|cost| (model.clone(), cost.clone()))
257            })
258            .collect()
259    }
260    pub fn delegate_targets(&self, id: SessionId) -> Result<Vec<AgentId>, EngineError> {
261        let session = self.inner.store.get(id)?;
262        let active = self
263            .inner
264            .sessions
265            .active
266            .lock()
267            .unwrap_or_else(|poisoned| poisoned.into_inner())
268            .values()
269            .find(|active| active.session == id)
270            .cloned();
271        let Some(active) = active else {
272            return Ok(Vec::new());
273        };
274        let depth = session_depth(&session.meta.origin);
275        Ok(active.policy.delegate_targets(depth))
276    }
277    pub fn children(
278        &self,
279        id: SessionId,
280    ) -> Result<Vec<cookie_agent_protocol::ChildSummary>, EngineError> {
281        // Listing is a use of the tree: complete it first so nested delegations
282        // and their summaries are in view (§3.2.2). An incomplete tree is
283        // reported, never papered over with the pre-load cache (review L14).
284        self.ensure_tree_loaded(id)?;
285        let known: HashSet<_> = self
286            .inner
287            .delegation_events
288            .entries()
289            .into_iter()
290            .map(|entry| {
291                (
292                    entry.reservation.invocation_id,
293                    entry.reservation.parent_session_id,
294                    entry.reservation.child_session_id,
295                )
296            })
297            .collect();
298        Ok(self
299            .inner
300            .store
301            .all_summaries()
302            .into_iter()
303            .filter_map(|child| match child.meta.origin {
304                SessionOrigin::Delegated {
305                    parent_session_id,
306                    invocation_id,
307                    ..
308                } if parent_session_id == id
309                    && known.contains(&(
310                        invocation_id,
311                        parent_session_id,
312                        child.meta.session_id,
313                    )) =>
314                {
315                    Some(ChildSummary {
316                        session_id: child.meta.session_id,
317                        agent: child.meta.creation_selection.agent.clone(),
318                        title: child.meta.title.clone(),
319                        title_updated_seq: child.meta.title_updated_seq,
320                        status: child.meta.status,
321                        usage: child.usage,
322                    })
323                }
324                _ => None,
325            })
326            .collect())
327    }
328
329    pub fn tree(&self, id: SessionId) -> Result<cookie_agent_protocol::SessionTree, EngineError> {
330        self.ensure_tree_loaded(id)?;
331        self.inner.store.get(id)?;
332        self.tree_summary(id)
333    }
334
335    fn tree_summary(
336        &self,
337        id: SessionId,
338    ) -> Result<cookie_agent_protocol::SessionTree, EngineError> {
339        Ok(cookie_agent_protocol::SessionTree {
340            session: self.inner.store.summary(id)?.meta,
341            children: self
342                .children(id)?
343                .into_iter()
344                .map(|child| self.tree_summary(child.session_id))
345                .collect::<Result<Vec<_>, _>>()?,
346        })
347    }
348    /// Adopts `id` and resolves the work its adoption found interrupted.
349    ///
350    /// Adoption schedules delegation recovery instead of running it inline, so
351    /// the wait below is what makes the ordering observable: when this returns,
352    /// every delegation record under the session's root that the adoption
353    /// classified as running without a live run here has been resolved, and the
354    /// root's background-delegation capacity reads settled rather than
355    /// transiently over-counted.
356    ///
357    /// The wait happens here, after the actor has replied, and never inside the
358    /// `Resume` handler: resolving a recovered delegate goes through its parent
359    /// session's mailbox, which is the very mailbox that handler occupies.
360    pub async fn resume(&self, id: SessionId) -> Result<SessionMeta, EngineError> {
361        let meta = self
362            .request(id, |reply| SessionCommand::Resume { reply })
363            .await?;
364        let root = self.inner.store.root_of(id).unwrap_or(id);
365        if !self
366            .await_delegation_recovery(root, DELEGATION_RECOVERY_SETTLE_BOUND)
367            .await
368        {
369            eprintln!(
370                "session {id} delegation recovery still settling after {DELEGATION_RECOVERY_SETTLE_BOUND:?}"
371            );
372        }
373        Ok(meta)
374    }
375    pub async fn revert_session(
376        &self,
377        session_id: SessionId,
378        through_seq: u64,
379        origin: EventOrigin,
380    ) -> Result<SessionRevertResult, EngineError> {
381        let context_id = crate::plugin::plugin_context_id();
382        let mut instructions = self
383            .inner
384            .store
385            .get(session_id)?
386            .log
387            .event_snapshot()
388            .iter()
389            .find_map(|event| match &event.payload {
390                cookie_agent_protocol::EventPayload::UserInputSubmitted { input }
391                    if event.seq > through_seq =>
392                {
393                    Some(input.clone())
394                }
395                _ => None,
396            });
397        let mut instructions_override = None;
398        for plugin in self.inner.plugins.interception_plugins(
399            cookie_agent_protocol::ExtensionInterceptionHook::SessionBeforeRevert,
400        ) {
401            let result = self
402                .inner
403                .plugins
404                .intercept_named::<_, cookie_agent_protocol::ExtensionSessionBeforeRevertResult>(
405                    &plugin,
406                    cookie_agent_protocol::PLUGIN_INTERCEPT_SESSION_BEFORE_REVERT_METHOD,
407                    &cookie_agent_protocol::ExtensionSessionBeforeRevertParams {
408                        session_id,
409                        context_id: context_id.clone(),
410                        through_seq,
411                        instructions: instructions.clone(),
412                    },
413                    Some(session_id),
414                    Some(&context_id),
415                )
416                .await;
417            match result {
418                Ok(result)
419                    if result.action
420                        == cookie_agent_protocol::ExtensionSessionBeforeRevertAction::Block =>
421                {
422                    let reason = result
423                        .reason
424                        .unwrap_or_else(|| format!("session revert blocked by {plugin}"));
425                    self.record_plugin_diagnostic(
426                        session_id,
427                        plugin,
428                        cookie_agent_protocol::PluginDiagnosticKind::HookBlocked,
429                        reason.clone(),
430                    );
431                    return Err(EngineError::SessionOperationBlocked(reason));
432                }
433                Ok(result)
434                    if result.action
435                        == cookie_agent_protocol::ExtensionSessionBeforeRevertAction::Override =>
436                {
437                    if let Some(replacement) = result
438                        .instructions_override
439                        .filter(|value| !value.trim().is_empty())
440                    {
441                        instructions = Some(replacement.clone());
442                        instructions_override = Some(replacement);
443                    } else {
444                        self.record_plugin_diagnostic(
445                            session_id,
446                            plugin,
447                            cookie_agent_protocol::PluginDiagnosticKind::InvalidModification,
448                            "empty revert instruction override".into(),
449                        );
450                    }
451                }
452                Ok(_) => {}
453                Err(error) => self.record_interception_error(session_id, plugin, error),
454            }
455        }
456        self.request(session_id, |reply| SessionCommand::Revert {
457            through_seq,
458            origin,
459            instructions_override,
460            reply,
461        })
462        .await
463    }
464    pub async fn fork_session(
465        &self,
466        session_id: SessionId,
467        through_seq: u64,
468        origin: EventOrigin,
469    ) -> Result<SessionForkResult, EngineError> {
470        let context_id = crate::plugin::plugin_context_id();
471        for plugin in self.inner.plugins.interception_plugins(
472            cookie_agent_protocol::ExtensionInterceptionHook::SessionBeforeFork,
473        ) {
474            let result = self
475                .inner
476                .plugins
477                .intercept_named::<_, cookie_agent_protocol::ExtensionSessionBeforeForkResult>(
478                    &plugin,
479                    cookie_agent_protocol::PLUGIN_INTERCEPT_SESSION_BEFORE_FORK_METHOD,
480                    &cookie_agent_protocol::ExtensionSessionBeforeForkParams {
481                        session_id,
482                        context_id: context_id.clone(),
483                        through_seq,
484                    },
485                    Some(session_id),
486                    Some(&context_id),
487                )
488                .await;
489            match result {
490                Ok(result)
491                    if result.action == cookie_agent_protocol::ExtensionAllowBlockAction::Block =>
492                {
493                    let reason = result
494                        .reason
495                        .unwrap_or_else(|| format!("session fork blocked by {plugin}"));
496                    self.record_plugin_diagnostic(
497                        session_id,
498                        plugin,
499                        cookie_agent_protocol::PluginDiagnosticKind::HookBlocked,
500                        reason.clone(),
501                    );
502                    return Err(EngineError::SessionOperationBlocked(reason));
503                }
504                Ok(_) => {}
505                Err(error) => self.record_interception_error(session_id, plugin, error),
506            }
507        }
508        let forked = self.inner.store.fork(session_id, through_seq, origin)?;
509        let producer_state = self.goal_producer_projection(forked)?;
510        if producer_state.goal.is_some()
511            || producer_state
512                .messages
513                .iter()
514                .any(|message| !message.consumed && !message.discarded)
515        {
516            self.reconcile_producers(forked).await?;
517        }
518        Ok(SessionForkResult { session_id: forked })
519    }
520    pub fn set_permission_mode(
521        &self,
522        session_id: SessionId,
523        mode: PermissionMode,
524    ) -> Result<(), EngineError> {
525        self.ensure_session_owned(session_id)?;
526        let session = self.inner.store.get(session_id)?;
527        let root = root_id(&session.meta.origin, session_id);
528        self.inner
529            .approvals
530            .permission_modes
531            .lock()
532            .unwrap_or_else(|poisoned| poisoned.into_inner())
533            .insert(root, mode);
534        Ok(())
535    }
536    pub(crate) fn permission_mode(&self, session_id: SessionId) -> PermissionMode {
537        let root = self
538            .inner
539            .store
540            .get(session_id)
541            .map(|session| root_id(&session.meta.origin, session_id))
542            .unwrap_or(session_id);
543        self.inner
544            .approvals
545            .permission_modes
546            .lock()
547            .unwrap_or_else(|poisoned| poisoned.into_inner())
548            .get(&root)
549            .copied()
550            .unwrap_or_default()
551    }
552    pub async fn rename_session(
553        &self,
554        params: SessionRenameParams,
555        origin: EventOrigin,
556    ) -> Result<SessionRenameResult, EngineError> {
557        let session_id = params.session_id;
558        let reset = matches!(params.change, SessionRenameChange::Reset);
559        let mut result = self
560            .request(session_id, |reply| SessionCommand::Rename {
561                params,
562                origin,
563                reply,
564            })
565            .await?;
566        if reset {
567            self.generate_title_after_reset(session_id).await?;
568            result.session = self.inner.store.get(session_id)?.metadata();
569        }
570        Ok(result)
571    }
572}
573
574#[cfg(test)]
575fn validated_tree_usage_edge(
576    origin: &SessionOrigin,
577    child_session_id: SessionId,
578    known: &HashSet<(InvocationId, SessionId, SessionId)>,
579) -> Option<(SessionId, (SessionId, InvocationId))> {
580    match origin {
581        SessionOrigin::Delegated {
582            parent_session_id,
583            invocation_id,
584            ..
585        } if known.contains(&(*invocation_id, *parent_session_id, child_session_id)) => {
586            Some((*parent_session_id, (child_session_id, *invocation_id)))
587        }
588        _ => None,
589    }
590}
591
592fn merge_session_tree_usage(
593    session_id: SessionId,
594    children: &HashMap<SessionId, Vec<(SessionId, InvocationId)>>,
595    rollups: &HashMap<SessionId, UsageRollup>,
596    usage: &mut UsageRollup,
597) -> Result<u64, EngineError> {
598    let mut visited = HashSet::new();
599    let mut pending = vec![(session_id, None)];
600    let mut session_count = 0_u64;
601    while let Some((current_id, incoming_invocation)) = pending.pop() {
602        if !visited.insert(current_id) {
603            let invocation_id = incoming_invocation.expect("only descendants can repeat");
604            return Err(
605                crate::delegation_events::DelegationEventError::Corrupt(invocation_id).into(),
606            );
607        }
608        if let Some(rollup) = rollups.get(&current_id) {
609            crate::usage::merge(usage, rollup);
610            session_count = session_count.saturating_add(1);
611        }
612        if let Some(descendants) = children.get(&current_id) {
613            pending.extend(
614                descendants
615                    .iter()
616                    .map(|(child_id, invocation_id)| (*child_id, Some(*invocation_id))),
617            );
618        }
619    }
620    Ok(session_count)
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626
627    #[test]
628    fn malformed_usage_tree_cycle_is_reported_as_delegation_corruption() {
629        let first = SessionId::new_v7();
630        let second = SessionId::new_v7();
631        let first_to_second = InvocationId::new_v7();
632        let second_to_first = InvocationId::new_v7();
633        let first_run = cookie_agent_protocol::RunId::new_v7();
634        let second_run = cookie_agent_protocol::RunId::new_v7();
635        let first_origin = SessionOrigin::Delegated {
636            root_session_id: first,
637            parent_session_id: second,
638            parent_run_id: second_run,
639            parent_tool_call_id: cookie_agent_protocol::ToolCallId::new_v7(),
640            invocation_id: second_to_first,
641            depth: 2,
642        };
643        let second_origin = SessionOrigin::Delegated {
644            root_session_id: first,
645            parent_session_id: first,
646            parent_run_id: first_run,
647            parent_tool_call_id: cookie_agent_protocol::ToolCallId::new_v7(),
648            invocation_id: first_to_second,
649            depth: 1,
650        };
651        let known = HashSet::from([
652            (first_to_second, first, second),
653            (second_to_first, second, first),
654        ]);
655        let mut children: HashMap<SessionId, Vec<(SessionId, InvocationId)>> = HashMap::new();
656        for (session_id, origin) in [(first, first_origin), (second, second_origin)] {
657            let (parent_id, edge) = validated_tree_usage_edge(&origin, session_id, &known)
658                .expect("matching origin and reservation triple");
659            children.entry(parent_id).or_default().push(edge);
660        }
661        let rollups = HashMap::from([
662            (first, UsageRollup::default()),
663            (second, UsageRollup::default()),
664        ]);
665        let error =
666            merge_session_tree_usage(first, &children, &rollups, &mut UsageRollup::default())
667                .expect_err("cycle must fail");
668        assert!(matches!(
669            error,
670            EngineError::DelegationEvents(
671                crate::delegation_events::DelegationEventError::Corrupt(invocation_id)
672            ) if invocation_id == second_to_first
673        ));
674    }
675}
cookie_agent_engine/runtime/sessions.rs

cookie_agent_engine/runtime/
sessions.rs
