mentra 0.21.0

An agent runtime for tool-using LLM applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

//! Tests for [`SessionEvent::MemoryUpdated`] reaching the session stream: the
//! hook bridge is registered on the per-session runtime handle, so one
//! session's memory activity lands on that session's channel and no other's.

use std::time::Duration;

use crate::{ContentBlock, session::event::SessionEvent, test::MockRuntime};

/// Waits until the stream yields a `MemoryUpdated`, skipping everything else.
///
/// Ingest is scheduled on a detached task after a run finishes, so the event
/// arrives asynchronously; the timeout turns a wiring regression into a fast
/// failure instead of a hung test.
async fn wait_for_memory_updated(rx: &mut crate::session::SessionEventReceiver) -> (String, usize) {
    tokio::time::timeout(Duration::from_secs(10), async {
        loop {
            match rx.recv().await {
                Ok(SessionEvent::MemoryUpdated {
                    agent_id,
                    stored_records,
                }) => return (agent_id, stored_records),
                Ok(_) => continue,
                Err(error) => panic!("event stream closed before MemoryUpdated: {error}"),
            }
        }
    })
    .await
    .expect("timed out waiting for MemoryUpdated on the session stream")
}

#[tokio::test]
async fn a_turns_memory_ingest_reaches_the_session_stream() {
    let mock = MockRuntime::builder().text("done").build().unwrap();
    let mut session = mock
        .runtime()
        .create_session("memory-events", mock.model())
        .unwrap();

    let mut rx = session.subscribe();
    session
        .append_turn(vec![ContentBlock::text("remember this")])
        .await
        .unwrap();

    let (agent_id, stored_records) = wait_for_memory_updated(&mut rx).await;
    assert_eq!(
        agent_id,
        session.agent_id(),
        "the event must name the session's own agent"
    );
    assert_eq!(
        stored_records, 1,
        "a turn with content must have stored its episode"
    );
}

#[tokio::test]
async fn memory_updated_does_not_cross_between_sessions() {
    let mock = MockRuntime::builder().text("done").build().unwrap();
    let mut active = mock
        .runtime()
        .create_session("active", mock.model())
        .unwrap();
    let bystander = mock
        .runtime()
        .create_session("bystander", mock.model())
        .unwrap();

    let mut active_rx = active.subscribe();
    let mut bystander_rx = bystander.subscribe();

    active
        .append_turn(vec![ContentBlock::text("remember this")])
        .await
        .unwrap();

    // The active session's event arriving is the ordering point: a wrongly
    // shared bridge would have sent to both channels in the same call.
    let (agent_id, _) = wait_for_memory_updated(&mut active_rx).await;
    assert_eq!(agent_id, active.agent_id());

    assert!(
        bystander_rx.try_recv().is_err(),
        "another session on the same runtime must not see this agent's memory activity"
    );
}

#[cfg(feature = "store-sqlite")]
#[tokio::test]
async fn a_resumed_session_still_emits_memory_updated() {
    use crate::runtime::SqliteRuntimeStore;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::time::{SystemTime, UNIX_EPOCH};

    static NEXT_ID: AtomicU64 = AtomicU64::new(1);
    let unique = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let store_path =
        std::env::temp_dir().join(format!("mentra-memory-resume-{timestamp}-{unique}.sqlite"));
    let store = SqliteRuntimeStore::new(store_path);
    let runtime_id = "memory-resume-test";

    let agent_id: String;

    // Phase 1: run one turn so there is a persisted agent to resume.
    {
        let mock = MockRuntime::builder()
            .runtime_identifier(runtime_id)
            .with_store(store.clone())
            .text("first response")
            .build()
            .unwrap();
        let mut session = mock
            .runtime()
            .create_session("memory-resume", mock.model())
            .unwrap();
        session
            .append_turn(vec![ContentBlock::text("hello")])
            .await
            .unwrap();
        agent_id = session.agent_id().to_owned();
        // mock (and its Runtime) dropped here, releasing the agent lease.
    }

    // Phase 2: resume on a fresh runtime; the resumed session's stream must
    // carry the memory events of its next turn.
    let mock2 = MockRuntime::builder()
        .runtime_identifier(runtime_id)
        .with_store(store)
        .text("second response")
        .build()
        .unwrap();
    let mut resumed = mock2.runtime().resume_session(&agent_id).unwrap();

    let mut rx = resumed.subscribe();
    resumed
        .append_turn(vec![ContentBlock::text("remember more")])
        .await
        .unwrap();

    let (event_agent, _) = wait_for_memory_updated(&mut rx).await;
    assert_eq!(
        event_agent, agent_id,
        "the resumed session must report its own agent's ingest"
    );
}