laser_sdk/agent/
session.rs

1use crate::agent::ReplayBound;
2use crate::agent::agdx::AgdxReceipt;
3use crate::agent::lease::{LeaseKey, SessionLease};
4use crate::context::{Chain, Checkpoint, ContextMessage, ContextPolicy, LastN, TokenBudget};
5use crate::context_scope::{ContextScope, ScopedMemory};
6use crate::error::LaserError;
7use crate::laser::Laser;
8use crate::provenance::AgentTopic;
9use crate::types::{ConversationId, MintUlid};
10use futures::FutureExt;
11use iggy::prelude::{Identifier, IggyExpiry, MaxTopicSize, StreamClient};
12use laser_wire::agent::{
13    AgentEnvelope, AgentErrorBody, AgentErrorCode, AgentId, Budget, OPERATION_SESSION, RecordId,
14    SdkInfo, SessionEnd, SessionStart, TaskState,
15};
16use laser_wire::content::ContentType;
17use laser_wire::control::{ControlCommand, SessionTopics};
18use laser_wire::dispatch::{DisplayType, SessionRecord, display_type};
19use laser_wire::framing::encode_named;
20use laser_wire::validate::Validate;
21use std::collections::BTreeMap;
22use std::sync::{Arc, Mutex};
23use std::time::Duration;
24
25/// The idle timeout a session start records unless it sets its own.
26pub const DEFAULT_SESSION_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
27
28/// The heartbeat interval of a process that holds session leases.
29pub const DEFAULT_SESSION_HEARTBEAT: Duration = Duration::from_secs(60);
30
31/// The default memory namespace of a [`Session`]. Conversation scoping keeps
32/// one session's memory apart from another's, so one shared namespace is the
33/// intended layout. [`SessionConfig::memory_namespace`] changes it.
34pub const DEFAULT_SESSION_MEMORY_NAMESPACE: &str = "agent.session";
35
36/// The default turn bound of [`Session::context`].
37pub const DEFAULT_SESSION_CONTEXT_TURNS: usize = 50;
38
39/// The default estimated token bound of [`Session::context`].
40pub const DEFAULT_SESSION_CONTEXT_TOKENS: usize = 4000;
41
42/// How a user key maps to a conversation: a fresh one per call, or a stable one per user.
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum SessionPolicy {
45    PerCall,
46    PerUser,
47}
48
49impl SessionPolicy {
50    /// The conversation id for `key` (random for `PerCall`, derived deterministically for `PerUser`).
51    pub fn conversation_for(&self, key: &str) -> ConversationId {
52        match self {
53            Self::PerCall => ConversationId::new(),
54            Self::PerUser => ConversationId::derive(key),
55        }
56    }
57}
58
59/// How long `agent.sessions` keeps records. Bootstrap needs one, because the
60/// session topic holds every session's lane and there is no safe default.
61/// A policy that never expires and has no size bound is refused.
62#[derive(Debug, Clone, Copy, PartialEq)]
63pub struct TopicRetention {
64    expiry: IggyExpiry,
65    max_size: MaxTopicSize,
66}
67
68impl TopicRetention {
69    /// A retention policy, refused when it would keep records forever with no
70    /// size bound.
71    pub fn new(expiry: IggyExpiry, max_size: MaxTopicSize) -> Result<Self, LaserError> {
72        let unbounded = matches!(
73            max_size,
74            MaxTopicSize::Unlimited | MaxTopicSize::ServerDefault
75        );
76        if matches!(expiry, IggyExpiry::NeverExpire | IggyExpiry::ServerDefault) && unbounded {
77            return Err(LaserError::Invalid(
78                "agent.sessions retention must set an expiry or a size bound".to_owned(),
79            ));
80        }
81        Ok(Self { expiry, max_size })
82    }
83
84    /// Expire records after `age`, with the server's default size bound. A
85    /// zero age is raised to one microsecond.
86    pub fn expire_after(age: Duration) -> Self {
87        let micros = micros(age).max(1);
88        Self {
89            expiry: IggyExpiry::ExpireDuration(iggy::prelude::IggyDuration::from(micros)),
90            max_size: MaxTopicSize::ServerDefault,
91        }
92    }
93
94    /// The message expiry.
95    pub fn expiry(&self) -> IggyExpiry {
96        self.expiry
97    }
98
99    /// The topic size bound.
100    pub fn max_size(&self) -> MaxTopicSize {
101        self.max_size
102    }
103}
104
105/// Where a stream's agents put their work. Lifecycle and state always ride the
106/// session's lane on `agent.sessions` in every layout.
107#[derive(Debug, Clone, Default, PartialEq, Eq)]
108pub enum SessionLayout {
109    /// One `agent.sessions` topic keyed by session. The default.
110    #[default]
111    Shared,
112    /// Each declared agent reads its own topic, keyed by session. On
113    /// `agent.sessions` sends, a command addressed to a declared agent goes to
114    /// that agent's topic, and a response, error, or chunk goes to its
115    /// addressee's declared topic. Lifecycle, state, broadcast records, and
116    /// records for an undeclared agent stay on the session lane. Bootstrap
117    /// creates the declared topics with the lane's partition count and
118    /// retention.
119    PerAgentTopic(BTreeMap<AgentId, String>),
120    /// Work rides `agent.sessions` on a declared partition per agent: a
121    /// command lands on its addressee's partition and a reply on its
122    /// requester's. Lifecycle, state, and records for an undeclared agent
123    /// stay on the session's partition.
124    PerAgentPartition(BTreeMap<AgentId, u32>),
125    /// Everything on one partition: bootstrap creates every agent topic with
126    /// one partition.
127    SinglePartition,
128}
129
130impl SessionLayout {
131    /// A [`PerAgentTopic`](Self::PerAgentTopic) layout from agent and topic
132    /// pairs. Takes the prelude agent id or the wire one.
133    pub fn per_agent_topic<A, T>(topics: impl IntoIterator<Item = (A, T)>) -> Self
134    where
135        A: Into<AgentId>,
136        T: Into<String>,
137    {
138        Self::PerAgentTopic(
139            topics
140                .into_iter()
141                .map(|(agent, topic)| (agent.into(), topic.into()))
142                .collect(),
143        )
144    }
145
146    /// A [`PerAgentPartition`](Self::PerAgentPartition) layout from agent and
147    /// partition pairs. Takes the prelude agent id or the wire one.
148    pub fn per_agent_partition<A>(partitions: impl IntoIterator<Item = (A, u32)>) -> Self
149    where
150        A: Into<AgentId>,
151    {
152        Self::PerAgentPartition(
153            partitions
154                .into_iter()
155                .map(|(agent, partition)| (agent.into(), partition))
156                .collect(),
157        )
158    }
159}
160
161// How often bootstrap checks whether the registration is served.
162const REGISTRATION_POLL_INTERVAL: Duration = Duration::from_millis(100);
163
164/// The layout and defaults of [`Laser::sessions`].
165#[derive(Debug, Clone, PartialEq, Eq)]
166pub struct SessionConfig {
167    stream: Option<String>,
168    layout: Option<SessionLayout>,
169    idle_timeout: Duration,
170    heartbeat: Duration,
171    register_source: bool,
172    fail_on_dead_letter: bool,
173    memory_namespace: String,
174    context_turns: usize,
175    context_tokens: usize,
176    sdk: SdkInfo,
177}
178
179impl Default for SessionConfig {
180    fn default() -> Self {
181        Self {
182            stream: None,
183            layout: None,
184            idle_timeout: DEFAULT_SESSION_IDLE_TIMEOUT,
185            heartbeat: DEFAULT_SESSION_HEARTBEAT,
186            register_source: true,
187            fail_on_dead_letter: false,
188            memory_namespace: DEFAULT_SESSION_MEMORY_NAMESPACE.to_owned(),
189            context_turns: DEFAULT_SESSION_CONTEXT_TURNS,
190            context_tokens: DEFAULT_SESSION_CONTEXT_TOKENS,
191            sdk: SdkInfo {
192                language: "rust".to_owned(),
193                version: env!("CARGO_PKG_VERSION").to_owned(),
194            },
195        }
196    }
197}
198
199impl SessionConfig {
200    /// The defaults.
201    #[must_use]
202    pub fn new() -> Self {
203        Self::default()
204    }
205
206    /// Lay sessions out on `stream` instead of the connection's default stream.
207    #[must_use]
208    pub fn stream(mut self, stream: impl Into<String>) -> Self {
209        self.stream = Some(stream.into());
210        self
211    }
212
213    /// Where agents put their work.
214    #[must_use]
215    pub fn layout(mut self, layout: SessionLayout) -> Self {
216        self.layout = Some(layout);
217        self
218    }
219
220    /// The idle timeout a session start records unless it sets its own.
221    #[must_use]
222    pub fn idle_timeout(mut self, timeout: Duration) -> Self {
223        self.idle_timeout = timeout;
224        self
225    }
226
227    /// How often a process with session leases publishes its heartbeat.
228    #[must_use]
229    pub fn heartbeat(mut self, interval: Duration) -> Self {
230        self.heartbeat = interval;
231        self
232    }
233
234    /// Whether bootstrap registers the stream as a session source on a
235    /// deployment that serves sessions. On by default.
236    #[must_use]
237    pub fn register_source(mut self, register: bool) -> Self {
238        self.register_source = register;
239        self
240    }
241
242    /// Fail a session when one of its records is dead-lettered. Off by
243    /// default: a dead letter counts as an error and never ends a session.
244    #[must_use]
245    pub fn fail_on_dead_letter(mut self, fail: bool) -> Self {
246        self.fail_on_dead_letter = fail;
247        self
248    }
249
250    /// The memory namespace [`Session::memory`] opens.
251    #[must_use]
252    pub fn memory_namespace(mut self, namespace: impl Into<String>) -> Self {
253        self.memory_namespace = namespace.into();
254        self
255    }
256
257    /// The turn bound of [`Session::context`].
258    #[must_use]
259    pub fn context_turns(mut self, turns: usize) -> Self {
260        self.context_turns = turns;
261        self
262    }
263
264    /// The estimated token bound of [`Session::context`].
265    #[must_use]
266    pub fn context_tokens(mut self, tokens: usize) -> Self {
267        self.context_tokens = tokens;
268        self
269    }
270
271    /// The SDK a session start names. Language bindings set their own.
272    #[must_use]
273    pub fn sdk(mut self, language: impl Into<String>, version: impl Into<String>) -> Self {
274        self.sdk = SdkInfo {
275            language: language.into(),
276            version: version.into(),
277        };
278        self
279    }
280
281    /// The stream sessions ride, or `None` for the connection's default stream.
282    #[must_use]
283    pub fn stream_name(&self) -> Option<&str> {
284        self.stream.as_deref()
285    }
286
287    /// The layout.
288    #[must_use]
289    pub fn layout_kind(&self) -> &SessionLayout {
290        const SHARED: &SessionLayout = &SessionLayout::Shared;
291        self.layout.as_ref().unwrap_or(SHARED)
292    }
293
294    /// The default idle timeout.
295    #[must_use]
296    pub fn idle_timeout_value(&self) -> Duration {
297        self.idle_timeout
298    }
299
300    /// The heartbeat interval.
301    #[must_use]
302    pub fn heartbeat_value(&self) -> Duration {
303        self.heartbeat
304    }
305
306    /// Whether bootstrap registers the session source.
307    #[must_use]
308    pub fn registers_source(&self) -> bool {
309        self.register_source
310    }
311
312    /// Whether a dead letter fails its session.
313    #[must_use]
314    pub fn fails_on_dead_letter(&self) -> bool {
315        self.fail_on_dead_letter
316    }
317
318    /// The memory namespace [`Session::memory`] opens.
319    #[must_use]
320    pub fn memory_namespace_name(&self) -> &str {
321        &self.memory_namespace
322    }
323
324    /// The turn bound of [`Session::context`].
325    #[must_use]
326    pub fn context_turn_bound(&self) -> usize {
327        self.context_turns
328    }
329
330    /// The estimated token bound of [`Session::context`].
331    #[must_use]
332    pub fn context_token_bound(&self) -> usize {
333        self.context_tokens
334    }
335
336    /// The SDK a session start names.
337    #[must_use]
338    pub fn sdk_info(&self) -> &SdkInfo {
339        &self.sdk
340    }
341
342    fn policy(&self) -> Box<dyn ContextPolicy> {
343        Box::new(Chain(vec![
344            Box::new(LastN(self.context_turns)),
345            Box::new(TokenBudget::new(self.context_tokens)),
346        ]))
347    }
348}
349
350/// The session id `label` derives to in `stream` under `namespace`. The same
351/// three always reach the same session, so two applications that share a
352/// stream share a session on purpose only when they pick the same namespace
353/// and label.
354pub fn derive_session_id(stream: &str, namespace: &str, label: &str) -> ConversationId {
355    ConversationId::derive(&format!("{stream}\u{1f}{namespace}\u{1f}{label}"))
356}
357
358impl Laser {
359    /// The session accessor under [`SessionConfig::default`]. Free and
360    /// synchronous, IO happens at the verbs.
361    pub fn sessions(&self) -> Sessions {
362        self.sessions_with(SessionConfig::default())
363    }
364
365    /// [`sessions`](Self::sessions) under an explicit [`SessionConfig`].
366    pub fn sessions_with(&self, config: SessionConfig) -> Sessions {
367        let laser = match &config.stream {
368            Some(stream) => self.with_default_stream(stream.clone()),
369            None => self.clone(),
370        };
371        // A declared layout holds for every later send on the stream through
372        // this connection, so agents and lenses route the same way.
373        if let (Some(layout), Some(stream)) = (&config.layout, laser.default_stream()) {
374            laser.declare_layout(stream, layout.clone());
375        }
376        Sessions {
377            laser,
378            config: Arc::new(config),
379            lane_guard: Arc::default(),
380        }
381    }
382}
383
384/// The session factory. Build it with [`Laser::sessions`] or
385/// [`Laser::sessions_with`].
386#[derive(Clone)]
387pub struct Sessions {
388    laser: Laser,
389    config: Arc<SessionConfig>,
390    lane_guard: Arc<LaneGuard>,
391}
392
393impl Sessions {
394    /// A session named `label`. Its id derives from the stream, the namespace,
395    /// and the label, so the same label reaches the same session.
396    pub fn create(&self, label: impl Into<String>) -> SessionBuilder {
397        SessionBuilder::new(self.clone(), Some(label.into()))
398    }
399
400    /// A fresh session with a new id.
401    pub fn start(&self) -> SessionBuilder {
402        SessionBuilder::new(self.clone(), None)
403    }
404
405    /// The session over an existing `conversation`, as a lens: no IO, no
406    /// lease, and no author until [`Session::as_agent`] names one.
407    pub fn open(&self, conversation: ConversationId) -> Session {
408        Session::lens(self.laser.clone(), Arc::clone(&self.config), conversation)
409            .with_lane_guard(Arc::clone(&self.lane_guard))
410    }
411
412    /// This factory's layout.
413    pub fn config(&self) -> &SessionConfig {
414        &self.config
415    }
416
417    /// Create the agent topics with [`Laser::bootstrap`], then, when the
418    /// deployment serves sessions and [`SessionConfig::register_source`] is
419    /// on, register the stream as a session source. A refused registration
420    /// is logged and reported as `registered: false`, never an error, because
421    /// an account without session administration rights relies on
422    /// provisioning to register the stream.
423    pub async fn bootstrap(
424        &self,
425        partitions: u32,
426        retention: TopicRetention,
427    ) -> Result<SessionBootstrap, LaserError> {
428        let partitions = match self.config.layout_kind() {
429            SessionLayout::SinglePartition => 1,
430            _ => partitions,
431        };
432        let declared = match self.config.layout_kind() {
433            SessionLayout::PerAgentTopic(topics) => declared_topics(topics)?,
434            _ => Vec::new(),
435        };
436        self.laser.bootstrap(partitions, retention).await?;
437        // Each declared agent topic carries a share of the session's work, so
438        // it takes the lane's partition count and retention.
439        let stream = self.stream()?.to_owned();
440        for topic in &declared {
441            crate::laser::ensure_topic_retained(
442                &self.laser.client(),
443                &stream,
444                topic,
445                partitions,
446                retention.expiry(),
447                retention.max_size(),
448            )
449            .await?;
450        }
451        self.lane_guard.observe(&self.laser).await?;
452        if !self.config.register_source || !self.laser.capabilities().await.sessions {
453            return Ok(SessionBootstrap { registered: false });
454        }
455        let stream = self.stream()?.to_owned();
456        let registered = match self
457            .laser
458            .publish_control(ControlCommand::RegisterSessionSource {
459                stream: stream.clone(),
460                topics: SessionTopics::All,
461            })
462            .await
463        {
464            Ok(()) => self.registration_visible(&stream).await,
465            Err(error) => {
466                tracing::warn!(stream = %stream, %error, "session source registration refused");
467                false
468            }
469        };
470        Ok(SessionBootstrap { registered })
471    }
472
473    /// Stop the deployment from indexing this stream's sessions. The request
474    /// goes through the authenticated managed control path and needs session
475    /// administration rights on the stream. The index for the stream is
476    /// dropped once the control record folds. A later
477    /// [`bootstrap`](Self::bootstrap) registers the stream again.
478    pub async fn unregister(&self) -> Result<(), LaserError> {
479        if !self.laser.capabilities().await.sessions {
480            return Err(LaserError::unsupported(
481                "sessions",
482                "this deployment does not index sessions",
483            ));
484        }
485        let stream = self.stream()?.to_owned();
486        self.laser
487            .publish_control(ControlCommand::RemoveSessionSource { stream })
488            .await
489    }
490
491    /// The stream the sessions ride.
492    pub fn stream(&self) -> Result<&str, LaserError> {
493        self.laser.stream_required()
494    }
495
496    pub(crate) fn laser(&self) -> &Laser {
497        &self.laser
498    }
499
500    // Wait until the deployment serves reads for the registered stream, up to
501    // the publish timeout. Registration applies after its control record
502    // folds, so reads and linked writes right after bootstrap would otherwise
503    // race it.
504    async fn registration_visible(&self, stream: &str) -> bool {
505        let deadline = tokio::time::Instant::now() + self.laser.publish_options().timeout;
506        let mut last_unavailable = None;
507        loop {
508            match self.changes(0, 1).await {
509                Err(LaserError::Session(
510                    laser_wire::session::SessionError::NotRegistered(_)
511                    | laser_wire::session::SessionError::Stale(_),
512                )) => {}
513                Ok(_) => return true,
514                Err(error) if error.is_unavailable() => {
515                    last_unavailable = Some(error.to_string());
516                }
517                Err(error) => {
518                    tracing::warn!(stream = %stream, %error, "session source registration not confirmed");
519                    return false;
520                }
521            }
522            if tokio::time::Instant::now() >= deadline {
523                tracing::warn!(stream = %stream, last_unavailable = ?last_unavailable, "session source registration not visible before the publish timeout");
524                return false;
525            }
526            tokio::time::sleep(REGISTRATION_POLL_INTERVAL).await;
527        }
528    }
529}
530
531// The distinct topics a per-agent topic layout declares, each a valid topic
532// name that is not one of the topics the layout routes around.
533fn declared_topics(topics: &BTreeMap<AgentId, String>) -> Result<Vec<String>, LaserError> {
534    let mut declared: Vec<String> = Vec::new();
535    for topic in topics.values() {
536        iggy::prelude::Identifier::named(topic)?;
537        if topic == laser_wire::topics::AGENT_SESSIONS || topic == laser_wire::topics::AGENT_CONTROL
538        {
539            return Err(LaserError::Invalid(format!(
540                "a per-agent topic layout cannot declare `{topic}`"
541            )));
542        }
543        if !declared.contains(topic) {
544            declared.push(topic.clone());
545        }
546    }
547    Ok(declared)
548}
549
550/// What [`Sessions::bootstrap`] set up.
551#[derive(Debug, Clone, Copy, PartialEq, Eq)]
552pub struct SessionBootstrap {
553    /// Whether the stream is now registered as a session source.
554    pub registered: bool,
555}
556
557/// A session about to start. Chain the optional settings, then
558/// [`begin`](Self::begin).
559#[must_use = "a session builder does nothing until begin().await"]
560pub struct SessionBuilder {
561    sessions: Sessions,
562    label: Option<String>,
563    namespace: Option<String>,
564    agent: Option<AgentId>,
565    parent: Option<ConversationId>,
566    root: Option<ConversationId>,
567    idle_timeout: Option<Duration>,
568    budget: Option<Budget>,
569    tags: Vec<String>,
570    id: Option<ConversationId>,
571}
572
573impl SessionBuilder {
574    fn new(sessions: Sessions, label: Option<String>) -> Self {
575        Self {
576            sessions,
577            label,
578            namespace: None,
579            agent: None,
580            parent: None,
581            root: None,
582            idle_timeout: None,
583            budget: None,
584            tags: Vec::new(),
585            id: None,
586        }
587    }
588
589    /// The agent that owns and writes the session. Required.
590    pub fn agent(mut self, agent: impl Into<AgentId>) -> Self {
591        let agent = agent.into();
592        self.agent = Some(agent);
593        self
594    }
595
596    /// The namespace a labeled session id derives under.
597    pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
598        self.namespace = Some(namespace.into());
599        self
600    }
601
602    /// Make this a child of `parent`, whose tree is rooted at `root` (the
603    /// parent itself when it has no parent).
604    pub fn parent(mut self, parent: ConversationId, root: ConversationId) -> Self {
605        self.parent = Some(parent);
606        self.root = Some(root);
607        self
608    }
609
610    /// Start the session under an explicit id instead of a label or a fresh id.
611    pub fn with_id(mut self, id: ConversationId) -> Self {
612        self.id = Some(id);
613        self
614    }
615
616    /// The idle timeout, instead of the configured default.
617    pub fn idle_timeout(mut self, timeout: Duration) -> Self {
618        self.idle_timeout = Some(timeout);
619        self
620    }
621
622    /// The token and cost ceiling a reader compares the session's usage with.
623    pub fn budget(mut self, budget: Budget) -> Self {
624        self.budget = Some(budget);
625        self
626    }
627
628    /// One searchable tag.
629    pub fn tag(mut self, tag: impl Into<String>) -> Self {
630        self.tags.push(tag.into());
631        self
632    }
633
634    /// The id this builder starts: the explicit id, the label's derived id,
635    /// or a fresh one on each call for an unlabeled session.
636    pub fn id(&self) -> Result<ConversationId, LaserError> {
637        if let Some(id) = self.id {
638            return Ok(id);
639        }
640        Ok(match &self.label {
641            Some(label) => derive_session_id(
642                self.sessions.stream()?,
643                self.namespace.as_deref().unwrap_or_default(),
644                label,
645            ),
646            None => ConversationId::new(),
647        })
648    }
649
650    /// Write the session start on the lane and take a lease that keeps the
651    /// session listed in this process's heartbeat.
652    pub async fn begin(self) -> Result<(Session, SessionLease), LaserError> {
653        let agent = self
654            .agent
655            .clone()
656            .ok_or_else(|| LaserError::Invalid("a session needs an agent".to_owned()))?;
657        let id = self.id()?;
658        let config = Arc::clone(&self.sessions.config);
659        let idle_timeout = self.idle_timeout.unwrap_or(config.idle_timeout);
660        let start = SessionStart {
661            label: self.label,
662            namespace: self.namespace,
663            agent: agent.clone(),
664            sdk: config.sdk.clone(),
665            parent: self.parent.map(Into::into),
666            root: self.root.map(Into::into),
667            idle_timeout_micros: Some(micros(idle_timeout)),
668            budget: self.budget,
669            tags: self.tags,
670        };
671        start.validate()?;
672        let session = Session::lens(self.sessions.laser.clone(), config, id)
673            .with_lane_guard(Arc::clone(&self.sessions.lane_guard))
674            .as_agent(agent)
675            .with_ancestry(self.parent, self.root);
676        let generation = session.stream_generation().await?;
677        let lane = session.lane()?;
678        let status = lane
679            .status(OPERATION_SESSION)
680            .with_task_state(TaskState::Working)
681            .body(encode_named(&start)?)
682            .content_type(ContentType::Cbor);
683        session.with_envelope_ancestry(status).send().await?;
684        let lease = session.lease(generation, idle_timeout);
685        Ok((session, lease))
686    }
687}
688
689fn micros(duration: Duration) -> u64 {
690    u64::try_from(duration.as_micros()).unwrap_or(u64::MAX)
691}
692
693// The first terminal a session handle wrote or tried to write. Every clone
694// shares it, so a retry repeats the same record and a different verb is
695// refused.
696#[derive(Clone)]
697struct TerminalIntent {
698    state: TaskState,
699    end: SessionEnd,
700    record: RecordId,
701}
702
703#[derive(Clone, Copy, Debug, PartialEq, Eq)]
704pub(crate) struct LaneIdentity {
705    pub stream_id: u32,
706    pub stream_generation: u64,
707    pub topic_id: u32,
708    pub topic_generation: u64,
709    pub partitions: u32,
710}
711
712#[derive(Default)]
713pub(crate) struct LaneGuard {
714    identity: Mutex<Option<LaneIdentity>>,
715    registered: std::sync::atomic::AtomicBool,
716}
717
718impl LaneGuard {
719    pub(crate) async fn observe(&self, laser: &Laser) -> Result<LaneIdentity, LaserError> {
720        let stream = Identifier::named(laser.stream_required()?)?;
721        let client = laser.client();
722        let details = client
723            .get_stream(&stream)
724            .await?
725            .ok_or_else(|| stale_lane("the session stream does not exist"))?;
726        let topic = details
727            .topics
728            .iter()
729            .find(|topic| topic.name == laser_wire::topics::AGENT_SESSIONS)
730            .ok_or_else(|| stale_lane("agent.sessions does not exist"))?;
731        let current = LaneIdentity {
732            stream_id: details.id,
733            stream_generation: details.created_at.as_micros(),
734            topic_id: topic.id,
735            topic_generation: topic.created_at.as_micros(),
736            partitions: topic.partitions_count,
737        };
738        if current.partitions == 0 {
739            return Err(stale_lane("agent.sessions has no partitions"));
740        }
741        let mut pinned = self.identity.lock().expect("session lane identity");
742        match *pinned {
743            Some(identity) if identity != current => Err(stale_lane(
744                "the session lane identity or partition count changed",
745            )),
746            Some(identity) => Ok(identity),
747            None => {
748                *pinned = Some(current);
749                Ok(current)
750            }
751        }
752    }
753
754    pub(crate) async fn check(
755        &self,
756        laser: &Laser,
757        conversation: ConversationId,
758    ) -> Result<LaneIdentity, LaserError> {
759        let identity = self.observe(laser).await?;
760        if !self.registered.load(std::sync::atomic::Ordering::Acquire)
761            && laser.capabilities().await.sessions
762        {
763            match laser.sessions().registered_lane(conversation).await {
764                Ok(view)
765                    if view
766                        == Some((
767                            identity.topic_id,
768                            identity.topic_generation,
769                            identity.partitions,
770                        )) =>
771                {
772                    self.registered
773                        .store(true, std::sync::atomic::Ordering::Release);
774                }
775                Ok(_) => {
776                    return Err(stale_lane(
777                        "the registered session lane does not match the current source",
778                    ));
779                }
780                Err(LaserError::Session(laser_wire::session::SessionError::NotRegistered(_))) => {}
781                Err(error) => return Err(error),
782            }
783        }
784        Ok(identity)
785    }
786}
787
788fn stale_lane(message: &str) -> LaserError {
789    LaserError::Session(laser_wire::session::SessionError::Stale(message.to_owned()))
790}
791
792/// One session. A cheap handle: clones share the terminal latch. Build it
793/// with [`Sessions::create`], [`Sessions::start`], or [`Sessions::open`].
794#[derive(Clone)]
795pub struct Session {
796    pub(crate) laser: Laser,
797    config: Arc<SessionConfig>,
798    scope: ContextScope,
799    pub(crate) agent: Option<AgentId>,
800    parent: Option<ConversationId>,
801    root: Option<ConversationId>,
802    terminal: Arc<Mutex<Option<TerminalIntent>>>,
803    lane_guard: Arc<LaneGuard>,
804    pub(crate) redactor: Arc<dyn Fn(&mut serde_json::Value) + Send + Sync>,
805    pub(crate) state: Arc<Mutex<crate::agent::session_ops::StateCursor>>,
806    pub(crate) current: Option<laser_wire::graph::SourceRef>,
807    // The control state of the runtime that handed this lens to a handler.
808    pub(crate) control: Option<Arc<crate::agent::control::ControlBook>>,
809    #[cfg(feature = "sign")]
810    key: Option<crate::sign::SigningKey>,
811}
812
813impl Session {
814    fn lens(laser: Laser, config: Arc<SessionConfig>, conversation: ConversationId) -> Self {
815        Self {
816            scope: laser.context(conversation),
817            laser,
818            config,
819            agent: None,
820            parent: None,
821            root: None,
822            terminal: Arc::new(Mutex::new(None)),
823            lane_guard: Arc::default(),
824            redactor: Arc::new(crate::agent::session_ops::default_redact),
825            state: Arc::default(),
826            current: None,
827            control: None,
828            #[cfg(feature = "sign")]
829            key: None,
830        }
831    }
832
833    fn with_lane_guard(mut self, guard: Arc<LaneGuard>) -> Self {
834        self.lane_guard = guard;
835        self
836    }
837
838    /// This handle signing its terminal record with `key`, so a verifying
839    /// reader can prove which agent ended the session. Feature `sign`.
840    #[cfg(feature = "sign")]
841    #[must_use]
842    pub fn signed_by(mut self, key: crate::sign::SigningKey) -> Self {
843        self.key = Some(key);
844        self
845    }
846
847    /// This handle writing as `agent`.
848    #[must_use]
849    pub fn as_agent(mut self, agent: impl Into<AgentId>) -> Self {
850        let agent = agent.into();
851        self.agent = Some(agent);
852        self
853    }
854
855    pub(crate) fn with_ancestry(
856        mut self,
857        parent: Option<ConversationId>,
858        root: Option<ConversationId>,
859    ) -> Self {
860        self.parent = parent;
861        self.root = root;
862        self
863    }
864
865    /// This session's id.
866    pub fn conversation(&self) -> ConversationId {
867        self.scope.conversation()
868    }
869
870    /// The stream this session lives in.
871    pub fn stream(&self) -> Result<&str, LaserError> {
872        self.laser.stream_required()
873    }
874
875    /// The agent this handle writes as.
876    pub fn agent(&self) -> Option<&AgentId> {
877        self.agent.as_ref()
878    }
879
880    /// The parent session, for a child session.
881    pub fn parent(&self) -> Option<ConversationId> {
882        self.parent
883    }
884
885    /// The root of this session's tree, for a child session.
886    pub fn root(&self) -> Option<ConversationId> {
887        self.root
888    }
889
890    /// The layout this session follows.
891    pub fn config(&self) -> &SessionConfig {
892        &self.config
893    }
894
895    /// The underlying [`ContextScope`], for a topic outside the session lane
896    /// or an explicit [`ContextPolicy`].
897    pub fn scope(&self) -> &ContextScope {
898        &self.scope
899    }
900
901    /// Append one typed envelope to the session lane. The envelope must name
902    /// this session as its conversation.
903    pub async fn append(&self, envelope: AgentEnvelope) -> Result<AgdxReceipt, LaserError> {
904        if ConversationId::from(envelope.conversation) != self.conversation() {
905            return Err(LaserError::Invalid(
906                "an appended envelope must belong to this session".to_owned(),
907            ));
908        }
909        self.laser
910            .agdx(
911                AgentTopic::Sessions,
912                envelope.source.clone(),
913                envelope.conversation,
914            )
915            .with_lane_guard(Arc::clone(&self.lane_guard))
916            .publish_envelope(envelope)
917            .await
918    }
919
920    /// End the session as completed. State changed since the last snapshot is
921    /// snapshotted first, so a reader can start from the final document.
922    pub async fn end(&self) -> Result<(), LaserError> {
923        self.state().snapshot_if_changed().await?;
924        self.terminate(TaskState::Completed, SessionEnd::default())
925            .await
926    }
927
928    /// End the session as failed with `error`.
929    pub async fn fail(&self, error: AgentErrorBody) -> Result<(), LaserError> {
930        self.terminate(
931            TaskState::Failed,
932            SessionEnd {
933                reason: None,
934                error: Some(error),
935            },
936        )
937        .await
938    }
939
940    /// End the session as canceled.
941    pub async fn cancel(&self) -> Result<(), LaserError> {
942        self.terminate(TaskState::Canceled, SessionEnd::default())
943            .await
944    }
945
946    /// Run `work` inside the session and end it by the outcome: completed on
947    /// success, failed on an error or a panic. A panic is re-raised after the
948    /// failure is written. When the terminal write also fails, the error from
949    /// `work` is returned. A dropped future, `panic = "abort"`, an unawaited
950    /// child task, or process death is not captured, and the session then
951    /// shows idle once its heartbeat stops.
952    pub async fn run<F, Fut, T>(self, lease: SessionLease, work: F) -> Result<T, LaserError>
953    where
954        F: FnOnce(Session) -> Fut,
955        Fut: std::future::Future<Output = Result<T, LaserError>>,
956    {
957        let outcome = std::panic::AssertUnwindSafe(work(self.clone()))
958            .catch_unwind()
959            .await;
960        let result = match outcome {
961            Ok(Ok(value)) => self.end().await.map(|()| value),
962            Ok(Err(error)) => {
963                let _ = self.fail(error_body(error.to_string())).await;
964                Err(error)
965            }
966            Err(panic) => {
967                let mut body = error_body(panic_message(panic.as_ref()));
968                body.detail = Some(BTreeMap::from([(
969                    "panic".to_owned(),
970                    laser_wire::query::Value::Bool(true),
971                )]));
972                let _ = self.fail(body).await;
973                drop(lease);
974                std::panic::resume_unwind(panic);
975            }
976        };
977        drop(lease);
978        result
979    }
980
981    // Mark a submitted session working as this handle's agent and take a lease.
982    pub(crate) async fn pick_up(&self) -> Result<SessionLease, LaserError> {
983        let generation = self.stream_generation().await?;
984        let transition = laser_wire::agent::SessionTransition {
985            actor: self.agent.clone(),
986            acknowledges: None,
987        };
988        let lane = self.lane()?;
989        let status = lane
990            .status(OPERATION_SESSION)
991            .with_task_state(TaskState::Working)
992            .body(encode_named(&transition)?)
993            .content_type(ContentType::Cbor);
994        self.with_envelope_ancestry(status).send().await?;
995        Ok(self.lease(generation, self.config.idle_timeout))
996    }
997
998    pub(crate) async fn terminate(
999        &self,
1000        state: TaskState,
1001        end: SessionEnd,
1002    ) -> Result<(), LaserError> {
1003        let intent = {
1004            let mut latch = self.terminal.lock().expect("session terminal latch");
1005            match latch.as_ref() {
1006                Some(intent) if intent.state != state => {
1007                    return Err(LaserError::Invalid(format!(
1008                        "session already ending as {}",
1009                        intent.state
1010                    )));
1011                }
1012                Some(intent) => intent.clone(),
1013                None => {
1014                    let intent = TerminalIntent {
1015                        state,
1016                        end,
1017                        record: RecordId::mint(),
1018                    };
1019                    *latch = Some(intent.clone());
1020                    intent
1021                }
1022            }
1023        };
1024        let lane = self.lane()?;
1025        let status = lane
1026            .status(OPERATION_SESSION)
1027            .with_task_state(intent.state)
1028            .with_record(intent.record)
1029            .body(encode_named(&intent.end)?)
1030            .content_type(ContentType::Cbor)
1031            .last();
1032        #[cfg(feature = "sign")]
1033        let status = match &self.key {
1034            Some(key) => status.signed_by(key),
1035            None => status,
1036        };
1037        self.with_envelope_ancestry(status).send().await.map(|_| ())
1038    }
1039
1040    pub(crate) fn lane(&self) -> Result<crate::agent::Agdx, LaserError> {
1041        let agent = self
1042            .agent
1043            .clone()
1044            .ok_or_else(|| LaserError::Invalid("this session handle has no agent".to_owned()))?;
1045        Ok(self
1046            .laser
1047            .agdx(AgentTopic::Sessions, agent, self.conversation().into())
1048            .with_lane_guard(Arc::clone(&self.lane_guard)))
1049    }
1050
1051    fn with_envelope_ancestry<'a>(
1052        &self,
1053        send: crate::agent::AgdxSend<'a>,
1054    ) -> crate::agent::AgdxSend<'a> {
1055        send.with_ancestry(self.parent.map(Into::into), self.root.map(Into::into))
1056    }
1057
1058    async fn stream_generation(&self) -> Result<u64, LaserError> {
1059        Ok(self
1060            .lane_guard
1061            .check(&self.laser, self.conversation())
1062            .await?
1063            .stream_generation)
1064    }
1065
1066    fn lease(&self, generation: u64, idle_timeout: Duration) -> SessionLease {
1067        let interval = self
1068            .config
1069            .heartbeat
1070            .min(idle_timeout / 5)
1071            .max(Duration::from_millis(10));
1072        self.laser.lease_registry().acquire(
1073            &self.laser,
1074            LeaseKey {
1075                stream: self.stream().unwrap_or_default().to_owned(),
1076                stream_generation: generation,
1077                session: self.conversation(),
1078            },
1079            self.agent.clone().expect("a leased session has an agent"),
1080            interval,
1081        )
1082    }
1083
1084    /// The model-ready context: the configured last records on the session
1085    /// lane, trimmed to the configured estimated token bound.
1086    pub async fn context(&self) -> Result<Vec<SessionTurn>, LaserError> {
1087        self.context_with(self.config.policy()).await
1088    }
1089
1090    /// [`context`](Self::context) under an explicit policy.
1091    pub async fn context_with(
1092        &self,
1093        policy: Box<dyn ContextPolicy>,
1094    ) -> Result<Vec<SessionTurn>, LaserError> {
1095        let messages = self.scope.fetch_with(lane_topics(), policy).await?;
1096        Ok(turns_of(messages))
1097    }
1098
1099    /// This session's memory in the configured namespace, scoped to the
1100    /// session.
1101    pub fn memory(&self) -> ScopedMemory {
1102        self.memory_in(self.config.memory_namespace.clone())
1103    }
1104
1105    /// [`memory`](Self::memory) in an explicit namespace.
1106    pub fn memory_in(&self, namespace: impl Into<String>) -> ScopedMemory {
1107        self.scope.memory(namespace)
1108    }
1109
1110    /// The knowledge graph `name`, the same graph [`Laser::graph`] returns.
1111    /// Feature `graph`.
1112    #[cfg(feature = "graph")]
1113    pub fn graph(&self, name: impl Into<String>) -> crate::graph::GraphHandle<'_> {
1114        self.scope.graph(name)
1115    }
1116
1117    /// Where the session lane ends right now. Persist it (it serializes) and
1118    /// hand it to [`state_at`](Self::state_at) or [`replay`](Self::replay).
1119    pub async fn checkpoint(&self) -> Result<Checkpoint, LaserError> {
1120        self.scope.checkpoint(&lane_topics()).await
1121    }
1122
1123    /// The records up to `checkpoint`.
1124    pub async fn turns_at(&self, checkpoint: Checkpoint) -> Result<Vec<SessionTurn>, LaserError> {
1125        self.turns(ReplayBound::At(checkpoint)).await
1126    }
1127
1128    /// The records appended after `checkpoint`.
1129    pub async fn turns_since(
1130        &self,
1131        checkpoint: Checkpoint,
1132    ) -> Result<Vec<SessionTurn>, LaserError> {
1133        self.turns(ReplayBound::FromCheckpoint(checkpoint)).await
1134    }
1135
1136    /// Fold the records up to `checkpoint` into `S`: state as it stood then.
1137    pub async fn state_at<S, F>(
1138        &self,
1139        checkpoint: Checkpoint,
1140        init: S,
1141        fold: F,
1142    ) -> Result<S, LaserError>
1143    where
1144        F: FnMut(S, &SessionTurn) -> S,
1145    {
1146        Ok(self.turns_at(checkpoint).await?.iter().fold(init, fold))
1147    }
1148
1149    /// Fold the records appended after `checkpoint` into `S`: bring state
1150    /// saved at that checkpoint up to date.
1151    pub async fn replay<S, F>(
1152        &self,
1153        checkpoint: Checkpoint,
1154        init: S,
1155        fold: F,
1156    ) -> Result<S, LaserError>
1157    where
1158        F: FnMut(S, &SessionTurn) -> S,
1159    {
1160        Ok(self.turns_since(checkpoint).await?.iter().fold(init, fold))
1161    }
1162
1163    async fn turns(&self, bound: ReplayBound) -> Result<Vec<SessionTurn>, LaserError> {
1164        let messages = self
1165            .scope
1166            .state(
1167                lane_topics(),
1168                bound,
1169                Vec::new(),
1170                |mut acc: Vec<ContextMessage>, message| {
1171                    acc.push(message.clone());
1172                    acc
1173                },
1174            )
1175            .await?;
1176        Ok(turns_of(messages))
1177    }
1178}
1179
1180fn lane_topics() -> Vec<AgentTopic<'static>> {
1181    vec![AgentTopic::Sessions]
1182}
1183
1184fn turns_of(messages: Vec<ContextMessage>) -> Vec<SessionTurn> {
1185    messages
1186        .into_iter()
1187        .map(|message| SessionTurn {
1188            display: match &message.envelope {
1189                Some(envelope) => display_type(
1190                    SessionRecord::Envelope(envelope),
1191                    laser_wire::topics::AGENT_SESSIONS,
1192                ),
1193                None => DisplayType::AgentMessage,
1194            },
1195            message,
1196        })
1197        .collect()
1198}
1199
1200fn error_body(message: String) -> AgentErrorBody {
1201    AgentErrorBody {
1202        code: AgentErrorCode::Internal,
1203        message: Some(message),
1204        retryable: false,
1205        detail: None,
1206    }
1207}
1208
1209fn panic_message(panic: &(dyn std::any::Any + Send)) -> String {
1210    panic
1211        .downcast_ref::<&str>()
1212        .map(|message| (*message).to_owned())
1213        .or_else(|| panic.downcast_ref::<String>().cloned())
1214        .unwrap_or_else(|| "the session work panicked".to_owned())
1215}
1216
1217/// One record read back from a [`Session`]: the message off the log plus how
1218/// a timeline shows it.
1219#[derive(Debug, Clone)]
1220pub struct SessionTurn {
1221    pub display: DisplayType,
1222    pub message: ContextMessage,
1223}
1224
1225impl SessionTurn {
1226    /// The record's body as UTF-8, lossy: the envelope body of a typed
1227    /// record, else the raw payload.
1228    pub fn text(&self) -> String {
1229        let body = match &self.message.envelope {
1230            Some(envelope) => &envelope.body,
1231            None => &self.message.payload,
1232        };
1233        String::from_utf8_lossy(body).into_owned()
1234    }
1235}
1236
1237#[cfg(test)]
1238mod tests {
1239    use super::*;
1240
1241    #[test]
1242    fn given_per_user_policy_when_deriving_for_a_key_then_should_be_stable_and_distinct() {
1243        let policy = SessionPolicy::PerUser;
1244        assert_eq!(
1245            policy.conversation_for("alice"),
1246            policy.conversation_for("alice")
1247        );
1248        assert_ne!(
1249            policy.conversation_for("alice"),
1250            policy.conversation_for("bob")
1251        );
1252    }
1253
1254    #[test]
1255    fn given_prelude_agent_ids_when_declaring_a_layout_then_should_match_the_wire_map() {
1256        let planner = crate::types::AgentId::new("planner").unwrap();
1257        let worker = crate::types::AgentId::new("worker").unwrap();
1258        let wire = |name: &str| name.parse::<AgentId>().unwrap();
1259
1260        assert_eq!(
1261            SessionLayout::per_agent_topic([
1262                (planner.clone(), "planner.inbox"),
1263                (worker.clone(), "worker.inbox")
1264            ]),
1265            SessionLayout::PerAgentTopic(BTreeMap::from([
1266                (wire("planner"), "planner.inbox".to_owned()),
1267                (wire("worker"), "worker.inbox".to_owned()),
1268            ]))
1269        );
1270        assert_eq!(
1271            SessionLayout::per_agent_partition([(planner, 0), (worker, 1)]),
1272            SessionLayout::PerAgentPartition(BTreeMap::from([
1273                (wire("planner"), 0),
1274                (wire("worker"), 1)
1275            ]))
1276        );
1277    }
1278
1279    #[test]
1280    fn given_per_call_policy_when_deriving_twice_then_should_be_unique() {
1281        let policy = SessionPolicy::PerCall;
1282        assert_ne!(policy.conversation_for("x"), policy.conversation_for("x"));
1283    }
1284
1285    #[test]
1286    fn given_labels_namespaces_and_streams_when_derived_then_should_separate_each_dimension() {
1287        let base = derive_session_id("agents", "ops", "incident");
1288        assert_eq!(base, derive_session_id("agents", "ops", "incident"));
1289        assert_ne!(base, derive_session_id("other", "ops", "incident"));
1290        assert_ne!(base, derive_session_id("agents", "desk", "incident"));
1291        assert_ne!(base, derive_session_id("agents", "ops", "refund"));
1292        assert_ne!(
1293            derive_session_id("a", "bc", "d"),
1294            derive_session_id("ab", "c", "d")
1295        );
1296    }
1297
1298    #[test]
1299    fn given_an_unbounded_retention_when_built_then_should_be_refused() {
1300        assert!(TopicRetention::new(IggyExpiry::NeverExpire, MaxTopicSize::Unlimited).is_err());
1301        assert!(TopicRetention::new(IggyExpiry::NeverExpire, MaxTopicSize::ServerDefault).is_err());
1302        assert!(
1303            TopicRetention::new(
1304                IggyExpiry::ExpireDuration(iggy::prelude::IggyDuration::from(86_400_000_000u64)),
1305                MaxTopicSize::Unlimited
1306            )
1307            .is_ok()
1308        );
1309    }
1310
1311    #[test]
1312    fn given_the_defaults_when_read_then_should_match_the_documented_values() {
1313        let config = SessionConfig::default();
1314        assert_eq!(config.idle_timeout_value(), Duration::from_secs(300));
1315        assert_eq!(config.heartbeat_value(), Duration::from_secs(60));
1316        assert!(config.registers_source());
1317        assert!(!config.fails_on_dead_letter());
1318        assert_eq!(config.sdk_info().language, "rust");
1319    }
1320}
laser_sdk/agent/session.rs

laser_sdk/agent/
session.rs
