gemini_adk_rs/tasks/
services.rs

1//! Task-owned service plans, isolated work, and accepted effects.
2use super::{TaskOwner, TaskServicesFactory};
3use crate::{
4    live::{
5        computed::ComputedRegistry,
6        extractor::{ExtractionTrigger, TurnExtractor},
7        temporal::{PatternDetector, pattern_fires},
8        transcript::{TranscriptBuffer, TranscriptTurn},
9        watcher::WatchPredicate,
10    },
11    state::State,
12};
13use async_trait::async_trait;
14use futures_util::FutureExt;
15use serde_json::Value;
16use std::{
17    collections::{BTreeMap, BTreeSet, HashMap, VecDeque},
18    sync::Arc,
19    time::{Duration, Instant},
20};
21use tokio_util::sync::CancellationToken;
22
23/// Declarative effects rendered against accepted task state.
24pub type TaskEffects = Arc<dyn Fn(&State) -> Vec<TaskEffect> + Send + Sync>;
25/// An effect requested by task-local declarative services.
26#[derive(Clone, Debug, serde::Serialize)]
27pub enum TaskEffect {
28    /// Change one declared task-local field.
29    Set {
30        /// Field name.
31        key: String,
32        /// New value.
33        value: Value,
34    },
35    /// Add task-owned context without requesting a response.
36    Context(String),
37    /// Request a response, only while this owner remains foreground.
38    Prompt(String),
39    /// Persist an accepted note through the configured memory service.
40    Remember(String),
41    /// Observable service execution or stabilization failure.
42    Error(String),
43}
44/// A state predicate and its declarative task-local actions.
45pub struct TaskWatcher {
46    /// Observed key, including a bare computed variable name.
47    pub key: String,
48    /// Existing shared old/new value predicate.
49    pub predicate: WatchPredicate,
50    /// Effects rendered after the change is accepted.
51    pub effects: TaskEffects,
52}
53/// A task-local temporal detector and its declarative actions.
54pub struct TaskPattern {
55    /// Diagnostic name.
56    pub name: String,
57    /// Fresh detector for this activation.
58    pub detector: Box<dyn PatternDetector>,
59    /// Minimum active-task time between firings.
60    pub cooldown: Option<Duration>,
61    /// Effects rendered after the detector fires.
62    pub effects: TaskEffects,
63}
64/// Memory operations bound to a trusted session subject and declared task slots.
65/// Ingestion and remember are accepted effects. Cancellation after execution starts
66/// cannot undo storage; returned slot projections still require current ownership.
67/// Each result is a complete projection: omitted or null slots remove a previous
68/// memory-owned value, while explicit inputs and accepted local changes remain.
69#[async_trait]
70pub trait TaskMemoryService: Send + Sync {
71    /// Read declared slot projections without ingesting or persisting new data.
72    async fn project(&self) -> Result<Value, String>;
73    /// Ingest an accepted user turn. The effect ID is stable for deduplication.
74    async fn ingest(&self, effect_id: String, turn: u64, user: String) -> Result<Value, String>;
75    /// Remember an accepted note, then return declared slot projections.
76    async fn remember(&self, effect_id: String, turn: u64, note: String) -> Result<Value, String>;
77}
78/// Fresh per-activation services. Workers never receive session transport.
79#[derive(Default)]
80pub struct TaskServices {
81    /// Extractors using immutable windows and isolated state snapshots.
82    pub extractors: Vec<Arc<dyn TurnExtractor>>,
83    /// Fresh dependency-ordered computed registry.
84    pub computed: ComputedRegistry,
85    /// Declarative state watchers.
86    pub watchers: Vec<TaskWatcher>,
87    /// Fresh temporal detectors.
88    pub patterns: Vec<TaskPattern>,
89    /// Trusted memory projection and accepted-effect adapter.
90    pub memory: Option<Arc<dyn TaskMemoryService>>,
91}
92/// An observation whose task/revision was captured before routing or async work.
93#[derive(Clone, Debug)]
94pub enum TaskObservation {
95    /// Finalized user and heard model text for one owned turn.
96    Turn {
97        /// User utterance.
98        user: String,
99        /// Heard model response.
100        model: String,
101    },
102    /// Generated text before any interruption truncation.
103    GenerationComplete {
104        /// Current user utterance.
105        user: String,
106        /// Generated model response.
107        model: String,
108    },
109    /// An interruption belonging to this task.
110    Interrupted,
111    /// A timer tick for the foreground task.
112    Timer,
113}
114/// A service effect accepted by the reducer, retaining its original owner.
115#[derive(Clone, Debug, serde::Serialize)]
116pub struct AcceptedTaskEffect {
117    /// Task, revision and service event identity.
118    pub owner: TaskOwner,
119    /// Accepted effect. Hosts must recheck ownership before delivering speech.
120    pub effect: TaskEffect,
121}
122#[derive(Clone)]
123pub(super) enum ServiceEvent {
124    Initialize,
125    Observe(TaskObservation),
126    AfterTool {
127        tool: String,
128        args: Value,
129        result: Value,
130    },
131    Project,
132    Ingest {
133        user: String,
134    },
135    Remember {
136        note: String,
137    },
138}
139struct ExtractRequest {
140    extractor: Arc<dyn TurnExtractor>,
141    window: Vec<TranscriptTurn>,
142}
143enum ServiceWork {
144    Extract(Vec<ExtractRequest>),
145    Memory {
146        provider: Arc<dyn TaskMemoryService>,
147        event: ServiceEvent,
148    },
149}
150/// Runtime-scoped service work. The private ticket cannot grant another task writes.
151pub struct OwnedTaskServiceWork {
152    pub(super) owner: TaskOwner,
153    pub(super) scope: Arc<()>,
154    pub(super) cancel: CancellationToken,
155    state: State,
156    work: ServiceWork,
157}
158/// Private completion of a runtime-prepared service job.
159pub struct TaskServiceCompletion {
160    pub(super) owner: TaskOwner,
161    pub(super) scope: Arc<()>,
162    pub(super) result: ServiceResult,
163}
164pub(super) enum ServiceResult {
165    Extractions(Vec<(Arc<dyn TurnExtractor>, Result<Value, String>)>),
166    Memory(Result<Value, String>),
167    Failed(String),
168}
169impl OwnedTaskServiceWork {
170    /// Captured task ownership, also used for correlation in controlled replay.
171    pub fn owner(&self) -> &TaskOwner {
172        &self.owner
173    }
174    /// Names selected by this job's canonical extraction trigger and window checks.
175    pub fn extractor_names(&self) -> Vec<&str> {
176        match &self.work {
177            ServiceWork::Extract(requests) => requests.iter().map(|r| r.extractor.name()).collect(),
178            _ => Vec::new(),
179        }
180    }
181    /// Execute configured service providers without reaching canonical task state.
182    pub async fn execute(self) -> TaskServiceCompletion {
183        self.run(None).await
184    }
185    /// Execute due extraction work with explicit provider fixtures. Missing fixtures
186    /// become service errors; this never calls the extraction provider.
187    pub async fn execute_controlled(
188        self,
189        values: BTreeMap<String, Result<Value, String>>,
190    ) -> TaskServiceCompletion {
191        self.run(Some(values)).await
192    }
193    async fn run(
194        self,
195        mut fixtures: Option<BTreeMap<String, Result<Value, String>>>,
196    ) -> TaskServiceCompletion {
197        let owner = self.owner.clone();
198        let result = if self.cancel.is_cancelled() {
199            ServiceResult::Failed("service cancelled before execution".into())
200        } else {
201            let future = async {
202                match self.work {
203                    ServiceWork::Extract(requests) => {
204                        let mut results = Vec::new();
205                        for request in requests {
206                            if self.cancel.is_cancelled() {
207                                return ServiceResult::Failed("service cancelled".into());
208                            }
209                            let result = if let Some(values) = fixtures.as_mut() {
210                                values.remove(request.extractor.name()).unwrap_or_else(|| {
211                                    Err(format!(
212                                        "missing controlled extraction '{}'",
213                                        request.extractor.name()
214                                    ))
215                                })
216                            } else {
217                                request
218                                    .extractor
219                                    .extract_with_state(&request.window, &self.state)
220                                    .await
221                                    .map_err(|error| error.to_string())
222                            };
223                            results.push((request.extractor, result));
224                        }
225                        if fixtures.as_ref().is_some_and(|values| !values.is_empty()) {
226                            return ServiceResult::Failed(
227                                "controlled extraction supplied unused fixtures".into(),
228                            );
229                        }
230                        ServiceResult::Extractions(results)
231                    }
232                    ServiceWork::Memory { provider, event } => {
233                        let effect_id = format!(
234                            "{}/revision-{}/{}",
235                            owner.task.0, owner.revision, owner.operation.0
236                        );
237                        let sequence = owner
238                            .operation
239                            .0
240                            .strip_prefix("service-")
241                            .and_then(|v| v.parse().ok())
242                            .unwrap_or(0);
243                        let result = match event {
244                            ServiceEvent::Project => provider.project().await,
245                            ServiceEvent::Ingest { user, .. } => {
246                                provider.ingest(effect_id, sequence, user).await
247                            }
248                            ServiceEvent::Remember { note, .. } => {
249                                provider.remember(effect_id, sequence, note).await
250                            }
251                            _ => unreachable!("only memory work reaches memory provider"),
252                        };
253                        ServiceResult::Memory(result)
254                    }
255                }
256            };
257            match std::panic::AssertUnwindSafe(future).catch_unwind().await {
258                Ok(result) => result,
259                Err(_) => ServiceResult::Failed("task service panicked".into()),
260            }
261        };
262        TaskServiceCompletion {
263            owner: self.owner,
264            scope: self.scope,
265            result,
266        }
267    }
268}
269pub(super) struct PendingService {
270    pub owner: TaskOwner,
271    pub cancel: CancellationToken,
272    pub event: ServiceEvent,
273    pub before: HashMap<String, Value>,
274}
275pub(super) struct ServiceState {
276    pub plan: TaskServices,
277    pub queue: VecDeque<ServiceEvent>,
278    pub pending: Option<PendingService>,
279    pub transcript: TranscriptBuffer,
280    pub turns: u64,
281    pub errors: Vec<String>,
282    pub contexts: Vec<String>,
283    pub memory_keys: BTreeSet<String>,
284    pub memory_values: BTreeMap<String, Value>,
285    triggered: Vec<Option<Instant>>,
286    logical_now: Instant,
287    last_wall: Instant,
288}
289impl ServiceState {
290    pub(super) fn create(
291        factory: &Option<TaskServicesFactory>,
292        state: &State,
293    ) -> Result<Option<Self>, String> {
294        let Some(factory) = factory else {
295            return Ok(None);
296        };
297        let plan = factory()?;
298        for extractor in &plan.extractors {
299            if reserved_service_key(extractor.name())
300                || extractor
301                    .promotion_rules()
302                    .iter()
303                    .any(|rule| reserved_service_key(&rule.state_key))
304            {
305                return Err(format!(
306                    "task extractor '{}' writes a reserved runtime field",
307                    extractor.name()
308                ));
309            }
310            if extractor.on_complete().is_some()
311                || extractor.trigger() == ExtractionTrigger::OnPhaseChange
312            {
313                return Err(format!(
314                    "task extractor '{}' has unsupported detached completion or phase trigger",
315                    extractor.name()
316                ));
317            }
318        }
319        plan.computed
320            .validate()
321            .map_err(|error| error.to_string())?;
322        let mut queue = VecDeque::from([ServiceEvent::Initialize]);
323        if plan.memory.is_some() {
324            queue.push_back(ServiceEvent::Project);
325        }
326        let triggered = vec![None; plan.patterns.len()];
327        let now = state.clock().now();
328        Ok(Some(Self {
329            plan,
330            queue,
331            pending: None,
332            transcript: TranscriptBuffer::new(),
333            turns: 0,
334            errors: Vec::new(),
335            contexts: Vec::new(),
336            memory_keys: BTreeSet::new(),
337            memory_values: BTreeMap::new(),
338            triggered,
339            logical_now: now,
340            last_wall: now,
341        }))
342    }
343    pub(super) fn refreshing_memory(&self) -> bool {
344        self.pending
345            .as_ref()
346            .is_some_and(|pending| matches!(pending.event, ServiceEvent::Project))
347            || self
348                .queue
349                .iter()
350                .any(|event| matches!(event, ServiceEvent::Project))
351    }
352    pub(super) fn advance_clock(&mut self, state: &State, active: bool) {
353        let now = state.clock().now();
354        if active {
355            self.logical_now += now.saturating_duration_since(self.last_wall);
356        }
357        self.last_wall = now;
358    }
359    pub(super) fn prepare(
360        &mut self,
361        event: &ServiceEvent,
362        owner: TaskOwner,
363        state: State,
364        scope: Arc<()>,
365    ) -> Option<OwnedTaskServiceWork> {
366        let memory_event = matches!(
367            event,
368            ServiceEvent::Project | ServiceEvent::Ingest { .. } | ServiceEvent::Remember { .. }
369        );
370        let work = if memory_event {
371            ServiceWork::Memory {
372                provider: self.plan.memory.clone()?,
373                event: event.clone(),
374            }
375        } else {
376            let trigger = match event {
377                ServiceEvent::Observe(TaskObservation::Turn { user, model }) => {
378                    self.turns += 1;
379                    self.transcript.push_input(user);
380                    self.transcript.push_output(model);
381                    self.transcript.end_turn();
382                    ExtractionTrigger::EveryTurn
383                }
384                ServiceEvent::Observe(TaskObservation::GenerationComplete { .. }) => {
385                    ExtractionTrigger::OnGenerationComplete
386                }
387                ServiceEvent::AfterTool { tool, args, result } => {
388                    self.transcript.push_tool_call(tool.clone(), args, result);
389                    ExtractionTrigger::AfterToolCall
390                }
391                _ => return None,
392            };
393            let mut requests = Vec::new();
394            for extractor in &self.plan.extractors {
395                let due = match extractor.trigger() {
396                    ExtractionTrigger::Interval(n) => {
397                        trigger == ExtractionTrigger::EveryTurn
398                            && n > 0
399                            && self.turns.is_multiple_of(u64::from(n))
400                    }
401                    value => value == trigger,
402                };
403                if !due {
404                    continue;
405                }
406                let mut window = if trigger == ExtractionTrigger::AfterToolCall {
407                    self.transcript
408                        .snapshot_window_with_current(extractor.window_size())
409                        .turns()
410                        .to_vec()
411                } else {
412                    self.transcript.window(extractor.window_size()).to_vec()
413                };
414                if let ServiceEvent::Observe(TaskObservation::GenerationComplete { user, model }) =
415                    event
416                {
417                    window.push(TranscriptTurn {
418                        turn_number: self.turns as u32,
419                        user: user.clone(),
420                        model: model.clone(),
421                        tool_calls: Vec::new(),
422                        timestamp: self.logical_now,
423                    });
424                    let discard = window.len().saturating_sub(extractor.window_size());
425                    window.drain(..discard);
426                }
427                if !window.is_empty() && extractor.should_extract(&window) {
428                    requests.push(ExtractRequest {
429                        extractor: extractor.clone(),
430                        window,
431                    });
432                }
433            }
434            if requests.is_empty() {
435                return None;
436            }
437            ServiceWork::Extract(requests)
438        };
439        let cancel = CancellationToken::new();
440        self.pending = Some(PendingService {
441            owner: owner.clone(),
442            cancel: cancel.clone(),
443            event: event.clone(),
444            before: state.to_hashmap(),
445        });
446        Some(OwnedTaskServiceWork {
447            owner,
448            scope,
449            cancel,
450            state,
451            work,
452        })
453    }
454    pub(super) fn stabilize(
455        &mut self,
456        state: &State,
457        before: HashMap<String, Value>,
458        event: Option<&ServiceEvent>,
459    ) -> Vec<TaskEffect> {
460        let mut emitted = Vec::new();
461        let mut previous = before;
462        for pass in 0..32 {
463            self.plan.computed.recompute(state);
464            let current = state.to_hashmap();
465            let mut effects = Vec::new();
466            for watcher in &self.plan.watchers {
467                let key = if current.contains_key(&format!("derived:{}", watcher.key)) {
468                    format!("derived:{}", watcher.key)
469                } else {
470                    watcher.key.clone()
471                };
472                let old = previous.get(&key).unwrap_or(&Value::Null);
473                let new = current.get(&key).unwrap_or(&Value::Null);
474                if old != new && watcher.predicate.matches(old, new) {
475                    effects.extend((watcher.effects)(state));
476                }
477            }
478            previous = current;
479            let changed = Self::apply_effects(state, effects, &mut emitted);
480            if !changed {
481                break;
482            }
483            if pass == 31 {
484                return vec![TaskEffect::Error(
485                    "task watcher effects did not converge after 32 passes".into(),
486                )];
487            }
488        }
489        let turn = matches!(
490            event,
491            Some(ServiceEvent::Observe(TaskObservation::Turn { .. }))
492        );
493        let timer = matches!(event, Some(ServiceEvent::Observe(TaskObservation::Timer)));
494        let turn_event = turn.then_some(gemini_genai_rs::session::SessionEvent::TurnComplete);
495        let mut effects = Vec::new();
496        for (index, pattern) in self.plan.patterns.iter().enumerate() {
497            if !(turn || (timer && pattern.detector.needs_timer())) {
498                continue;
499            }
500            if pattern_fires(
501                pattern.detector.as_ref(),
502                state,
503                turn_event.as_ref(),
504                self.logical_now,
505                pattern.cooldown,
506                &mut self.triggered[index],
507            ) {
508                effects.extend((pattern.effects)(state));
509            }
510        }
511        if Self::apply_effects(state, effects, &mut emitted) {
512            // Temporal Set effects enter the same computed/watch stabilization without counting another turn.
513            emitted.extend(self.stabilize(state, previous, None));
514        }
515        emitted
516    }
517    fn apply_effects(
518        state: &State,
519        effects: Vec<TaskEffect>,
520        emitted: &mut Vec<TaskEffect>,
521    ) -> bool {
522        let mut changed = false;
523        for effect in effects {
524            if let TaskEffect::Set { key, value } = effect {
525                if reserved_service_key(&key) {
526                    emitted.push(TaskEffect::Error(format!(
527                        "task service effect uses reserved field '{key}'"
528                    )));
529                    continue;
530                }
531                if state.get_raw(&key).as_ref() != Some(&value) {
532                    match state.set(&key, value) {
533                        Ok(()) => changed = true,
534                        Err(error) => emitted.push(TaskEffect::Error(error.to_string())),
535                    }
536                }
537            } else {
538                emitted.push(effect);
539            }
540        }
541        changed
542    }
543    pub(super) fn invalidate(&mut self) {
544        if let Some(pending) = self.pending.take() {
545            pending.cancel.cancel();
546        }
547        self.queue.clear();
548    }
549}
550
551pub(super) fn reserved_service_key(key: &str) -> bool {
552    [
553        "task:",
554        "flow:",
555        "derived:",
556        "repair:",
557        "correction:",
558        "state_meta:",
559    ]
560    .iter()
561    .any(|prefix| key.starts_with(prefix))
562}
gemini_adk_rs/tasks/services.rs

gemini_adk_rs/tasks/
services.rs
