Skip to main content

eredu_runtime/intervention/
session.rs

1//! Shared admission, immutable installation and bounded observer forwarding.
2use super::*;
3use std::sync::Arc;
4
5/// Revalidates exact retained session capabilities before shared geometry/budget
6/// traversal. Callers first validate native execution-mode facts and capture support.
7pub fn validate_session(
8    capture: &AdmittedCapturePlan,
9    plan: &AdmittedInterventionPlan,
10    discovery: &InterventionDiscovery,
11    estimator: &dyn InterventionEstimator,
12) -> Result<(), CaptureError> {
13    validate_continuation(
14        capture,
15        plan,
16        discovery,
17        estimator,
18        0,
19        CaptureUsage::default(),
20    )
21}
22
23pub(crate) fn validate_continuation(
24    capture: &AdmittedCapturePlan,
25    plan: &AdmittedInterventionPlan,
26    discovery: &InterventionDiscovery,
27    estimator: &dyn InterventionEstimator,
28    next_prediction: u64,
29    inherited: CaptureUsage,
30) -> Result<(), CaptureError> {
31    if plan.request().batch != 1 {
32        return Err(CaptureError::Unsupported(
33            "interventions require single-sequence text generation".into(),
34        ));
35    }
36    let checked = plan
37        .plan()
38        .clone()
39        .admit(discovery, plan.request(), plan.session_id())?;
40    if checked.identity() != plan.identity() {
41        return Err(CaptureError::Invalid(
42            "intervention admission differs from loaded source/session capabilities".into(),
43        ));
44    }
45    preflight_continuation(capture, &checked, estimator, next_prediction, inherited)
46}
47
48/// Installs one immutable run after validation. Both capture-only and combined
49/// plans obey the same replacement rule; empty runs keep the ordinary fast path.
50pub fn install_session(
51    slot: &mut Option<CaptureSession>,
52    capture: AdmittedCapturePlan,
53    intervention: Option<(AdmittedInterventionPlan, Arc<dyn InterventionEstimator>)>,
54) -> Result<(), CaptureError> {
55    if slot.is_some() {
56        return Err(CaptureError::Invalid(
57            "capture/intervention plan already installed".into(),
58        ));
59    }
60    if capture.is_empty()
61        && intervention
62            .as_ref()
63            .is_none_or(|(plan, _)| plan.is_empty())
64    {
65        return Ok(());
66    }
67    let mut session = CaptureSession::new(capture);
68    if let Some((plan, estimator)) = intervention {
69        session.enable_interventions(plan, estimator)?;
70    }
71    *slot = Some(session);
72    Ok(())
73}
74
75/// Borrows the existing capture owner. Backends provide native primitives and a
76/// typed error conversion; this adapter owns no completion or native recovery state.
77pub struct CaptureObserver<'a, B, F> {
78    session: &'a mut CaptureSession,
79    backend: B,
80    map_error: F,
81}
82impl<'a, B, F> CaptureObserver<'a, B, F> {
83    /// Creates an observer that forwards all value/control hooks to the shared run.
84    pub fn new(session: &'a mut CaptureSession, backend: B, map_error: F) -> Self {
85        Self {
86            session,
87            backend,
88            map_error,
89        }
90    }
91}
92impl<B, E, F> crate::ActivationObserver<B::Tensor, E> for CaptureObserver<'_, B, F>
93where
94    B: InterventionBackend,
95    F: Fn(CaptureExecutionError<B::Error>) -> E,
96{
97    fn observe(&mut self, path: &str, value: &B::Tensor) -> Result<(), E> {
98        self.session
99            .observe(&mut self.backend, path, value)
100            .map_err(&self.map_error)
101    }
102    fn intervene(&mut self, path: &str, value: &B::Tensor) -> Result<Option<B::Tensor>, E> {
103        self.session
104            .intervene(&mut self.backend, path, value)
105            .map_err(&self.map_error)
106    }
107    fn routing_control(
108        &mut self,
109        path: &str,
110        rows: u64,
111    ) -> Result<Option<eredu_nn::routing_intervention::GroupSelectionControl>, E> {
112        self.session
113            .routing_control(path, rows)
114            .map_err(|error| (self.map_error)(error.into()))
115    }
116    fn routing_applied(
117        &mut self,
118        path: &str,
119        original: Option<crate::RoutingDecision<'_, B::Tensor>>,
120        effective: crate::RoutingDecision<'_, B::Tensor>,
121    ) -> Result<(), E> {
122        self.session
123            .routing_applied(&mut self.backend, path, original, effective)
124            .map_err(&self.map_error)
125    }
126    fn routing_failed(&mut self, path: &str, message: &str) {
127        self.session.routing_failed(path, message);
128    }
129    fn finish(&mut self) -> Result<(), E> {
130        self.session
131            .finish_interventions()
132            .map_err(|error| (self.map_error)(error.into()))
133    }
134    fn observe_routing(
135        &mut self,
136        routing: crate::RoutingObservation<'_, B::Tensor>,
137    ) -> Result<(), E> {
138        let mut result = Ok(());
139        routing.for_each_tensor(|path, value| {
140            if result.is_ok() {
141                result = self.observe(&path, value);
142            }
143        });
144        result
145    }
146}