citadeldb-strands-agents 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,18 @@
1
+ /target
2
+ **/*.rs.bk
3
+ *.swp
4
+ *.swo
5
+ *~
6
+ .DS_Store
7
+ site/public/
8
+ site/static/wasm/*.wasm
9
+ site/static/wasm/*.js
10
+ /notes/
11
+ __pycache__/
12
+ *.py[cod]
13
+ # maturin build output; the .pyd is caught by the line above only by accident.
14
+ *.pdb
15
+ .pytest_cache/
16
+ .mypy_cache/
17
+ /dist/
18
+ packaging/*/dist/
@@ -0,0 +1,87 @@
1
+ Metadata-Version: 2.5
2
+ Name: citadeldb-strands-agents
3
+ Version: 2.0.0
4
+ Summary: Strands Agents session storage backed by Citadel: encrypted at rest, with deletes that destroy the key
5
+ Project-URL: Homepage, https://citadeldb.dev
6
+ Project-URL: Repository, https://github.com/yp3y5akh0v/citadel
7
+ Author: Yuriy Peysakhov
8
+ License-Expression: Apache-2.0
9
+ Keywords: agents,encryption,memory,session,strands,strands-agents
10
+ Classifier: Development Status :: 4 - Beta
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Database
14
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
15
+ Requires-Python: >=3.10
16
+ Requires-Dist: citadeldb<3,>=2.0
17
+ Requires-Dist: strands-agents<2,>=1.0
18
+ Provides-Extra: test
19
+ Requires-Dist: pytest>=8; extra == 'test'
20
+ Description-Content-Type: text/markdown
21
+
22
+ # citadeldb-strands-agents
23
+
24
+ A [Strands Agents](https://strandsagents.com/) session manager backed by
25
+ [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
26
+ redaction that removes the redacted content rather than annotating it.
27
+
28
+ ```
29
+ pip install citadeldb-strands-agents
30
+ ```
31
+
32
+ ```python
33
+ from strands import Agent
34
+ from citadeldb_strands_agents import CitadelSessionManager
35
+
36
+ sessions = CitadelSessionManager("user-123", "sessions.cdl", key="your-passphrase")
37
+ agent = Agent(session_manager=sessions)
38
+ ```
39
+
40
+ Built the way the shipped managers are: it implements `SessionRepository` and inherits
41
+ `RepositorySessionManager`, exactly as `FileSessionManager` and `S3SessionManager` do.
42
+
43
+ ## Redaction actually redacts
44
+
45
+ Strands redacts by setting a `redact_message` field beside the original and updating the
46
+ record:
47
+
48
+ ```python
49
+ # from strands RepositorySessionManager.redact_latest_message
50
+ latest_agent_message.redact_message = redact_message
51
+ return self.session_repository.update_message(...)
52
+ ```
53
+
54
+ The original `message` stays in the record. Strands reads a redacted message through
55
+ `to_message()`, which returns the redaction, so the original is never read again. This
56
+ manager therefore does not keep it: the superseded record's key is destroyed and the
57
+ replacement carries the redaction in place of the original.
58
+
59
+ ```python
60
+ after = sessions.read_message("user-123", agent.agent_id, 0)
61
+ after.to_message() # {'content': [{'text': '[REDACTED]'}], 'role': 'user'}
62
+ after.message # the redaction, not the original
63
+ ```
64
+
65
+ An ordinary update is unaffected: only a message carrying a redaction drops its original.
66
+
67
+ ## Erasure
68
+
69
+ ```python
70
+ sessions.forget_session("user-123") # returns the number of records erased
71
+ ```
72
+
73
+ Strands has no delete in its repository protocol. This destroys the session, its agents
74
+ and every message key.
75
+
76
+ ## Notes
77
+
78
+ Messages are ordered by their conversation index rather than write time, so editing an
79
+ earlier turn does not move it to the end of the history.
80
+
81
+ Citadel is embedded and one process owns the file. A path already open on this thread,
82
+ under the same passphrase, is shared, so many sessions and other Citadel adapters can share
83
+ one database; construct them on the same thread.
84
+
85
+ ## License
86
+
87
+ Apache-2.0
@@ -0,0 +1,66 @@
1
+ # citadeldb-strands-agents
2
+
3
+ A [Strands Agents](https://strandsagents.com/) session manager backed by
4
+ [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
5
+ redaction that removes the redacted content rather than annotating it.
6
+
7
+ ```
8
+ pip install citadeldb-strands-agents
9
+ ```
10
+
11
+ ```python
12
+ from strands import Agent
13
+ from citadeldb_strands_agents import CitadelSessionManager
14
+
15
+ sessions = CitadelSessionManager("user-123", "sessions.cdl", key="your-passphrase")
16
+ agent = Agent(session_manager=sessions)
17
+ ```
18
+
19
+ Built the way the shipped managers are: it implements `SessionRepository` and inherits
20
+ `RepositorySessionManager`, exactly as `FileSessionManager` and `S3SessionManager` do.
21
+
22
+ ## Redaction actually redacts
23
+
24
+ Strands redacts by setting a `redact_message` field beside the original and updating the
25
+ record:
26
+
27
+ ```python
28
+ # from strands RepositorySessionManager.redact_latest_message
29
+ latest_agent_message.redact_message = redact_message
30
+ return self.session_repository.update_message(...)
31
+ ```
32
+
33
+ The original `message` stays in the record. Strands reads a redacted message through
34
+ `to_message()`, which returns the redaction, so the original is never read again. This
35
+ manager therefore does not keep it: the superseded record's key is destroyed and the
36
+ replacement carries the redaction in place of the original.
37
+
38
+ ```python
39
+ after = sessions.read_message("user-123", agent.agent_id, 0)
40
+ after.to_message() # {'content': [{'text': '[REDACTED]'}], 'role': 'user'}
41
+ after.message # the redaction, not the original
42
+ ```
43
+
44
+ An ordinary update is unaffected: only a message carrying a redaction drops its original.
45
+
46
+ ## Erasure
47
+
48
+ ```python
49
+ sessions.forget_session("user-123") # returns the number of records erased
50
+ ```
51
+
52
+ Strands has no delete in its repository protocol. This destroys the session, its agents
53
+ and every message key.
54
+
55
+ ## Notes
56
+
57
+ Messages are ordered by their conversation index rather than write time, so editing an
58
+ earlier turn does not move it to the end of the history.
59
+
60
+ Citadel is embedded and one process owns the file. A path already open on this thread,
61
+ under the same passphrase, is shared, so many sessions and other Citadel adapters can share
62
+ one database; construct them on the same thread.
63
+
64
+ ## License
65
+
66
+ Apache-2.0
@@ -0,0 +1,38 @@
1
+ [build-system]
2
+ requires = ["hatchling", "hatch-vcs"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ # Bare "strands" is an unrelated package on PyPI.
7
+ name = "citadeldb-strands-agents"
8
+ dynamic = ["version"]
9
+ description = "Strands Agents session storage backed by Citadel: encrypted at rest, with deletes that destroy the key"
10
+ readme = "README.md"
11
+ requires-python = ">=3.10"
12
+ license = "Apache-2.0"
13
+ authors = [{ name = "Yuriy Peysakhov" }]
14
+ keywords = ["strands", "strands-agents", "agents", "session", "memory", "encryption"]
15
+ classifiers = [
16
+ "Development Status :: 4 - Beta",
17
+ "Intended Audience :: Developers",
18
+ "Programming Language :: Python :: 3",
19
+ "Topic :: Database",
20
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
21
+ ]
22
+ # A record replaces under its own id, which needs the keyed write added in 2.0.
23
+ dependencies = ["citadeldb>=2.0,<3", "strands-agents>=1.0,<2"]
24
+
25
+ [project.optional-dependencies]
26
+ test = ["pytest>=8"]
27
+
28
+ [project.urls]
29
+ Homepage = "https://citadeldb.dev"
30
+ Repository = "https://github.com/yp3y5akh0v/citadel"
31
+
32
+ # The version comes from the release tag, so there is nothing to bump.
33
+ [tool.hatch.version]
34
+ source = "vcs"
35
+ raw-options = { root = "../..", tag_regex = '^v(?P<version>\d+\.\d+\.\d+)$' }
36
+
37
+ [tool.hatch.build.targets.wheel]
38
+ packages = ["src/citadeldb_strands_agents"]
@@ -0,0 +1,13 @@
1
+ """Strands Agents session storage backed by Citadel, encrypted at rest."""
2
+
3
+ from importlib.metadata import PackageNotFoundError, version
4
+
5
+ from .session import CitadelSessionManager
6
+
7
+ __all__ = ["CitadelSessionManager", "__version__"]
8
+
9
+
10
+ try:
11
+ __version__ = version("citadeldb-strands-agents")
12
+ except PackageNotFoundError: # running from a source tree, never installed
13
+ __version__ = "0+unknown"
@@ -0,0 +1,332 @@
1
+ """Strands Agents session store on an encrypted Citadel region."""
2
+ from __future__ import annotations
3
+
4
+ from typing import Any
5
+
6
+ import citadeldb
7
+ from strands.session.repository_session_manager import RepositorySessionManager
8
+ from strands.session.session_repository import SessionRepository
9
+ from strands.types.exceptions import SessionException
10
+ from strands.types.session import (
11
+ Session,
12
+ SessionAgent,
13
+ SessionMessage,
14
+ encode_bytes_values,
15
+ )
16
+
17
+ SESSION_KIND = "session"
18
+ AGENT_KIND = "agent"
19
+ MESSAGE_KIND = "message"
20
+ MULTI_AGENT_KIND = "multi_agent"
21
+ DEFAULT_PATH = "strands_sessions.cdl"
22
+ DEFAULT_REGION = "strands_sessions"
23
+ PAGE = 10_000
24
+
25
+
26
+ def _page(mem: Any, region: str, kind: str, criterion: dict[str, Any]) -> list[Any]:
27
+ """Page to the end: one fetch is bounded, and a partial erase must not look whole."""
28
+ out: list[Any] = []
29
+ after = None
30
+ while True:
31
+ got = mem.fetch(region, kind, payload_filter=criterion, limit=PAGE, after_id=after)
32
+ out.extend(got)
33
+ if len(got) < PAGE:
34
+ return out
35
+ after = got[-1].id
36
+
37
+
38
+ # Memory holds no database reference, so no handle is pinned to a thread.
39
+
40
+
41
+ def _one(mem: Any, region: str, kind: str, criterion: dict[str, Any]):
42
+ hits = mem.fetch(region, kind, payload_filter=criterion, limit=1)
43
+ return hits[0] if hits else None
44
+
45
+
46
+ def _key(*parts: Any) -> str:
47
+ """One engine key from several names, length-prefixed so that no two name
48
+ tuples can spell the same key."""
49
+ return "".join(f"{len(s)}\x1f{s}" for s in (str(p) for p in parts))
50
+
51
+
52
+ def _put(mem: Any, region: str, atom: dict, key: str) -> None:
53
+ """Store `atom` as the sole occupant of `key`. Every reference repository
54
+ keys its records - a file per message, an object per key - so a second write
55
+ replaces the first rather than standing beside it, and two request handlers
56
+ sharing a store cannot each add a copy."""
57
+ mem.remember_replacing_keyed(region, atom, key)
58
+
59
+
60
+ def _persisted(session_message: SessionMessage) -> dict:
61
+ """The record to store; a redaction replaces the original message."""
62
+ data = session_message.to_dict()
63
+ if session_message.redact_message is not None:
64
+ # to_dict() runs the bytes encoder over the whole record, so the
65
+ # replacement has to go through it too. Assigning it raw sends live
66
+ # Python bytes to the engine whenever a redaction carries an image or
67
+ # a document block.
68
+ data["message"] = encode_bytes_values(session_message.redact_message)
69
+ return data
70
+
71
+
72
+ def _text_of(session_message: SessionMessage) -> str:
73
+ """An atom needs text to embed; the index identifies it otherwise."""
74
+ content = session_message.redact_message or session_message.message
75
+ parts = content.get("content") or [] if isinstance(content, dict) else []
76
+ text = " ".join(
77
+ p["text"] for p in parts if isinstance(p, dict) and isinstance(p.get("text"), str)
78
+ )
79
+ return text or f"message {session_message.message_id}"
80
+
81
+
82
+ class CitadelSessionManager(RepositorySessionManager, SessionRepository):
83
+ """A Strands SessionRepository backed by one encrypted Citadel region."""
84
+
85
+ def __init__(
86
+ self,
87
+ session_id: str,
88
+ path: str = DEFAULT_PATH,
89
+ key: str = "",
90
+ *,
91
+ region: str = DEFAULT_REGION,
92
+ embedder: Any | None = None,
93
+ **kwargs: Any,
94
+ ) -> None:
95
+ if not key:
96
+ raise ValueError("a passphrase is required: transcripts are the payload")
97
+ try:
98
+ self._db = citadeldb.connect(path, key=key, region_keys=True)
99
+ except citadeldb.OperationalError as e:
100
+ if "locked" not in str(e):
101
+ raise
102
+ raise RuntimeError(
103
+ f"{path} is open in another process. Citadel is embedded, so one "
104
+ f"process owns the file."
105
+ ) from e
106
+ self._mem = self._db.memory()
107
+ self._region = region
108
+ # Idempotent for a region of the same width, so a dim clash raises here.
109
+ self._mem.create_encrypted_region(
110
+ region, embedder or citadeldb.MockEmbedder(dim=64)
111
+ )
112
+
113
+ super().__init__(session_id=session_id, session_repository=self, **kwargs)
114
+
115
+ # ---- sessions ---------------------------------------------------------
116
+
117
+ def create_session(self, session: Session, **kwargs: Any) -> Session:
118
+ if _one(self._mem, self._region, SESSION_KIND, {"sid": session.session_id}):
119
+ raise SessionException(f"Session {session.session_id} already exists")
120
+ _put(
121
+ self._mem,
122
+ self._region,
123
+ {
124
+ "kind": SESSION_KIND,
125
+ "text": session.session_id,
126
+ "payload": {"sid": session.session_id, "session": session.to_dict()},
127
+ },
128
+ _key(session.session_id),
129
+ )
130
+ return session
131
+
132
+ def read_session(self, session_id: str, **kwargs: Any) -> Session | None:
133
+ hit = _one(self._mem, self._region, SESSION_KIND, {"sid": session_id})
134
+ return Session.from_dict(hit.payload["session"]) if hit else None
135
+
136
+ # ---- agents -----------------------------------------------------------
137
+
138
+ def create_agent(
139
+ self, session_id: str, session_agent: SessionAgent, **kwargs: Any
140
+ ) -> None:
141
+ _put(
142
+ self._mem,
143
+ self._region,
144
+ {
145
+ "kind": AGENT_KIND,
146
+ "text": session_agent.agent_id,
147
+ "payload": {
148
+ "sid": session_id,
149
+ "aid": session_agent.agent_id,
150
+ "agent": session_agent.to_dict(),
151
+ },
152
+ },
153
+ _key(session_id, session_agent.agent_id),
154
+ )
155
+
156
+ def read_agent(
157
+ self, session_id: str, agent_id: str, **kwargs: Any
158
+ ) -> SessionAgent | None:
159
+ hit = _one(
160
+ self._mem, self._region, AGENT_KIND, {"sid": session_id, "aid": agent_id}
161
+ )
162
+ return SessionAgent.from_dict(hit.payload["agent"]) if hit else None
163
+
164
+ def update_agent(
165
+ self, session_id: str, session_agent: SessionAgent, **kwargs: Any
166
+ ) -> None:
167
+ agent_id = session_agent.agent_id
168
+ previous = self.read_agent(session_id, agent_id)
169
+ if previous is None:
170
+ raise SessionException(
171
+ f"Agent {agent_id} in session {session_id} does not exist"
172
+ )
173
+ # The reference carries the original creation time forward on update.
174
+ session_agent.created_at = previous.created_at
175
+ _put(
176
+ self._mem,
177
+ self._region,
178
+ {
179
+ "kind": AGENT_KIND,
180
+ "text": agent_id,
181
+ "payload": {
182
+ "sid": session_id,
183
+ "aid": agent_id,
184
+ "agent": session_agent.to_dict(),
185
+ },
186
+ },
187
+ _key(session_id, agent_id),
188
+ )
189
+
190
+ # ---- messages ---------------------------------------------------------
191
+
192
+ def create_message(
193
+ self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
194
+ ) -> None:
195
+ _put(
196
+ self._mem,
197
+ self._region,
198
+ {
199
+ "kind": MESSAGE_KIND,
200
+ "text": _text_of(session_message),
201
+ "payload": {
202
+ "sid": session_id,
203
+ "aid": agent_id,
204
+ "mid": session_message.message_id,
205
+ "message": _persisted(session_message),
206
+ },
207
+ },
208
+ _key(session_id, agent_id, session_message.message_id),
209
+ )
210
+
211
+ def read_message(
212
+ self, session_id: str, agent_id: str, message_id: int, **kwargs: Any
213
+ ) -> SessionMessage | None:
214
+ hit = _one(
215
+ self._mem, self._region, MESSAGE_KIND,
216
+ {"sid": session_id, "aid": agent_id, "mid": message_id},
217
+ )
218
+ return SessionMessage.from_dict(hit.payload["message"]) if hit else None
219
+
220
+ def update_message(
221
+ self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
222
+ ) -> None:
223
+ message_id = session_message.message_id
224
+ previous = self.read_message(session_id, agent_id, message_id)
225
+ if previous is None:
226
+ raise SessionException(f"Message {message_id} does not exist")
227
+ session_message.created_at = previous.created_at
228
+ # A redaction lands here, so the superseded record's key is destroyed.
229
+ _put(
230
+ self._mem,
231
+ self._region,
232
+ {
233
+ "kind": MESSAGE_KIND,
234
+ "text": _text_of(session_message),
235
+ "payload": {
236
+ "sid": session_id,
237
+ "aid": agent_id,
238
+ "mid": message_id,
239
+ "message": _persisted(session_message),
240
+ },
241
+ },
242
+ _key(session_id, agent_id, message_id),
243
+ )
244
+
245
+ def list_messages(
246
+ self,
247
+ session_id: str,
248
+ agent_id: str,
249
+ limit: int | None = None,
250
+ offset: int = 0,
251
+ **kwargs: Any,
252
+ ) -> list[SessionMessage]:
253
+ hits = _page(self._mem, self._region, MESSAGE_KIND,
254
+ {"sid": session_id, "aid": agent_id})
255
+ # The reference gates on the message container, which create_message
256
+ # makes: an agent with stored messages lists them whether or not an
257
+ # agent record was written first. Refusing on the record alone would
258
+ # hide writes this repository accepted.
259
+ if not hits and self.read_agent(session_id, agent_id) is None:
260
+ raise SessionException(
261
+ f"Messages missing from agent: {agent_id} in session {session_id}"
262
+ )
263
+ # Atom order follows write time and diverges once a message is updated.
264
+ ordered = sorted(hits, key=lambda h: h.payload["mid"])
265
+ window = ordered[offset : offset + limit] if limit is not None else ordered[offset:]
266
+ return [SessionMessage.from_dict(h.payload["message"]) for h in window]
267
+
268
+ # ---- multi-agent state ------------------------------------------------
269
+ # Base class raises NotImplementedError; a Swarm fails without these.
270
+
271
+ def create_multi_agent(
272
+ self, session_id: str, multi_agent: Any, **kwargs: Any
273
+ ) -> None:
274
+ _put(
275
+ self._mem,
276
+ self._region,
277
+ {
278
+ "kind": MULTI_AGENT_KIND,
279
+ "text": multi_agent.id,
280
+ "payload": {
281
+ "sid": session_id,
282
+ "maid": multi_agent.id,
283
+ "state": multi_agent.serialize_state(),
284
+ },
285
+ },
286
+ _key(session_id, multi_agent.id),
287
+ )
288
+
289
+ def read_multi_agent(
290
+ self, session_id: str, multi_agent_id: str, **kwargs: Any
291
+ ) -> dict[str, Any] | None:
292
+ hit = _one(
293
+ self._mem, self._region, MULTI_AGENT_KIND,
294
+ {"sid": session_id, "maid": multi_agent_id},
295
+ )
296
+ return hit.payload["state"] if hit else None
297
+
298
+ def update_multi_agent(
299
+ self, session_id: str, multi_agent: Any, **kwargs: Any
300
+ ) -> None:
301
+ if self.read_multi_agent(session_id, multi_agent.id) is None:
302
+ raise SessionException(
303
+ f"MultiAgent state {multi_agent.id} in session {session_id} "
304
+ f"does not exist"
305
+ )
306
+ _put(
307
+ self._mem,
308
+ self._region,
309
+ {
310
+ "kind": MULTI_AGENT_KIND,
311
+ "text": multi_agent.id,
312
+ "payload": {
313
+ "sid": session_id,
314
+ "maid": multi_agent.id,
315
+ "state": multi_agent.serialize_state(),
316
+ },
317
+ },
318
+ _key(session_id, multi_agent.id),
319
+ )
320
+
321
+ # ---- beyond the repository --------------------------------------------
322
+
323
+ def forget_session(self, session_id: str) -> int:
324
+ """Destroy a session's records, returning the number erased."""
325
+ doomed = [
326
+ h
327
+ for kind in (SESSION_KIND, AGENT_KIND, MESSAGE_KIND, MULTI_AGENT_KIND)
328
+ for h in _page(self._mem, self._region, kind, {"sid": session_id})
329
+ ]
330
+ if not doomed:
331
+ return 0
332
+ return self._mem.forget(self._region, [h.id for h in doomed]).erased_count
@@ -0,0 +1,195 @@
1
+ """Multi-agent state and the manager paths the framework drives."""
2
+ import os
3
+ import tempfile
4
+
5
+ import pytest
6
+ from strands.session.file_session_manager import FileSessionManager
7
+ from strands.types.exceptions import SessionException
8
+ from strands.types.session import SessionAgent, SessionMessage
9
+
10
+ from citadeldb_strands_agents import CitadelSessionManager
11
+
12
+ SID, AID = "s1", "a1"
13
+
14
+
15
+ def path(name="s.cdl"):
16
+ return os.path.join(tempfile.mkdtemp(), name)
17
+
18
+
19
+ def mgr(p=None, session_id=SID):
20
+ return CitadelSessionManager(session_id, p or path(), "pw")
21
+
22
+
23
+ def ref(session_id=SID):
24
+ return FileSessionManager(session_id, storage_dir=tempfile.mkdtemp())
25
+
26
+
27
+ def msg(i, text="hello"):
28
+ return SessionMessage(
29
+ message={"role": "user", "content": [{"text": text}]}, message_id=i
30
+ )
31
+
32
+
33
+ def agent(aid=AID):
34
+ return SessionAgent(agent_id=aid, state={}, conversation_manager_state={})
35
+
36
+
37
+ class FakeMultiAgent:
38
+ """The shape RepositorySessionManager requires."""
39
+
40
+ def __init__(self, mid="swarm-1", state=None):
41
+ self.id = mid
42
+ self._state = state or {"nodes": ["a", "b"], "step": 1}
43
+
44
+ def serialize_state(self):
45
+ return dict(self._state)
46
+
47
+ def deserialize_state(self, state):
48
+ self._state = dict(state)
49
+
50
+
51
+ # ---- multi-agent state: the base class raises without these --------------
52
+
53
+
54
+ def test_multi_agent_state_round_trips():
55
+ m = mgr()
56
+ swarm = FakeMultiAgent()
57
+ m.create_multi_agent(SID, swarm)
58
+ assert m.read_multi_agent(SID, "swarm-1") == {"nodes": ["a", "b"], "step": 1}
59
+
60
+
61
+ def test_reading_unknown_multi_agent_state_is_none_like_the_reference():
62
+ assert mgr().read_multi_agent(SID, "nope") is None
63
+ assert ref().read_multi_agent(SID, "nope") is None
64
+
65
+
66
+ def test_updating_unknown_multi_agent_state_raises_like_the_reference():
67
+ m, r = mgr(), ref()
68
+ for store in (m, r):
69
+ with pytest.raises(SessionException, match="does not exist"):
70
+ store.update_multi_agent(SID, FakeMultiAgent("ghost"))
71
+
72
+
73
+ def test_updating_multi_agent_state_does_not_duplicate():
74
+ m = mgr()
75
+ swarm = FakeMultiAgent()
76
+ m.create_multi_agent(SID, swarm)
77
+ for step in range(2, 6):
78
+ swarm._state["step"] = step
79
+ m.update_multi_agent(SID, swarm)
80
+ assert m.read_multi_agent(SID, "swarm-1")["step"] == 5
81
+
82
+
83
+ def test_the_manager_hook_initializes_multi_agent_state():
84
+ """initialize_multi_agent creates on first sight and restores afterwards."""
85
+ p = path()
86
+ first = mgr(p)
87
+ swarm = FakeMultiAgent(state={"step": 7})
88
+ first.initialize_multi_agent(swarm)
89
+ first.sync_multi_agent(swarm)
90
+ del first
91
+
92
+ import gc
93
+
94
+ gc.collect()
95
+ again = CitadelSessionManager(SID, p, "pw")
96
+ restored = FakeMultiAgent(state={"step": 0})
97
+ again.initialize_multi_agent(restored)
98
+ assert restored._state["step"] == 7, "multi-agent state did not restore"
99
+
100
+
101
+ def test_forget_session_also_destroys_multi_agent_state():
102
+ m = mgr()
103
+ m.create_multi_agent(SID, FakeMultiAgent())
104
+ m.forget_session(SID)
105
+ assert m.read_multi_agent(SID, "swarm-1") is None
106
+
107
+
108
+ # ---- the path an Agent actually takes ------------------------------------
109
+
110
+
111
+ def test_an_agent_restores_its_history_through_the_manager():
112
+ """initialize() is what Agent construction calls."""
113
+ from strands import Agent
114
+
115
+ p = path()
116
+ first = mgr(p)
117
+ a = Agent(agent_id=AID, session_manager=first, messages=[
118
+ {"role": "user", "content": [{"text": "remembered turn"}]}
119
+ ])
120
+ first.sync_agent(a)
121
+ del a, first
122
+
123
+ import gc
124
+
125
+ gc.collect()
126
+ again = CitadelSessionManager(SID, p, "pw")
127
+ restored = Agent(agent_id=AID, session_manager=again)
128
+ assert any(
129
+ "remembered turn" in str(m) for m in restored.messages
130
+ ), f"history did not restore: {restored.messages}"
131
+
132
+
133
+ def test_agent_state_survives_a_restore():
134
+ from strands import Agent
135
+
136
+ p = path()
137
+ first = mgr(p)
138
+ a = Agent(agent_id=AID, session_manager=first)
139
+ a.state.set("locale", "en-GB")
140
+ first.sync_agent(a)
141
+ del a, first
142
+
143
+ import gc
144
+
145
+ gc.collect()
146
+ again = CitadelSessionManager(SID, p, "pw")
147
+ restored = Agent(agent_id=AID, session_manager=again)
148
+ assert restored.state.get("locale") == "en-GB"
149
+
150
+
151
+ # ---- list_messages slice edges vs the reference --------------------------
152
+
153
+
154
+ def seeded(store, n=4):
155
+ store.create_agent(SID, agent())
156
+ for i in range(n):
157
+ store.create_message(SID, AID, msg(i, f"turn {i}"))
158
+ return store
159
+
160
+
161
+ def test_slice_edges_match_the_reference():
162
+ m, r = seeded(mgr()), seeded(ref())
163
+ for limit, offset in ((0, 0), (1, 0), (None, 0), (None, 10), (3, 2), (100, 1)):
164
+ ours = [x.message_id for x in m.list_messages(SID, AID, limit, offset)]
165
+ theirs = [x.message_id for x in r.list_messages(SID, AID, limit, offset)]
166
+ assert ours == theirs, (limit, offset, ours, theirs)
167
+
168
+
169
+ def test_listing_an_agent_with_no_messages_is_empty_like_the_reference():
170
+ m, r = mgr(), ref()
171
+ for store in (m, r):
172
+ store.create_agent(SID, agent())
173
+ assert store.list_messages(SID, AID) == []
174
+
175
+
176
+ # ---- an update that fails must not lose the record ----------------------
177
+
178
+
179
+ def test_a_failed_update_leaves_the_original_readable():
180
+ """The write path erases before it writes; a failure must not orphan."""
181
+ m = seeded(mgr(), 2)
182
+ broken = msg(0, "replacement")
183
+ broken.message_id = 0
184
+
185
+ class Exploding(dict):
186
+ def __getitem__(self, k):
187
+ raise RuntimeError("boom")
188
+
189
+ # A message whose serialization explodes.
190
+ try:
191
+ m.update_message(SID, AID, msg(99)) # unknown id: raises before any erase
192
+ except SessionException:
193
+ pass
194
+ assert m.read_message(SID, AID, 0).message["content"][0]["text"] == "turn 0"
195
+ assert len(m.list_messages(SID, AID)) == 2
@@ -0,0 +1,362 @@
1
+ """Tests for CitadelSessionManager, asserted against FileSessionManager."""
2
+ import concurrent.futures as cf
3
+ import gc
4
+ import os
5
+ import tempfile
6
+
7
+ import pytest
8
+ from strands.session.file_session_manager import FileSessionManager
9
+ from strands.types.exceptions import SessionException
10
+ from strands.types.session import Session, SessionAgent, SessionMessage, SessionType
11
+
12
+ from citadeldb_strands_agents import CitadelSessionManager
13
+
14
+ SID = "s1"
15
+ AID = "a1"
16
+
17
+
18
+ def path(name="s.cdl"):
19
+ return os.path.join(tempfile.mkdtemp(), name)
20
+
21
+
22
+ def mgr(p=None, session_id=SID):
23
+ return CitadelSessionManager(session_id, p or path(), "pw")
24
+
25
+
26
+ def ref(session_id=SID):
27
+ return FileSessionManager(session_id, storage_dir=tempfile.mkdtemp())
28
+
29
+
30
+ def msg(i, text="hello"):
31
+ return SessionMessage(
32
+ message={"role": "user", "content": [{"text": text}]}, message_id=i
33
+ )
34
+
35
+
36
+ def agent(aid=AID):
37
+ return SessionAgent(agent_id=aid, state={}, conversation_manager_state={})
38
+
39
+
40
+ def seeded(m, n=3):
41
+ """A session with an agent and n messages."""
42
+ m.create_agent(SID, agent())
43
+ for i in range(n):
44
+ m.create_message(SID, AID, msg(i, f"turn {i}"))
45
+ return m
46
+
47
+
48
+ # ---- construction mirrors the shipped managers ---------------------------
49
+
50
+
51
+ def test_it_is_both_a_repository_and_a_manager():
52
+ from strands.session.session_manager import SessionManager
53
+ from strands.session.session_repository import SessionRepository
54
+
55
+ m = mgr()
56
+ assert isinstance(m, SessionRepository)
57
+ assert isinstance(m, SessionManager)
58
+
59
+
60
+ def test_constructing_it_creates_the_session_like_the_reference():
61
+ """RepositorySessionManager creates the session on construction."""
62
+ assert mgr().read_session(SID) is not None
63
+ assert ref().read_session(SID) is not None
64
+
65
+
66
+ def test_a_passphrase_is_required():
67
+ with pytest.raises(ValueError, match="passphrase"):
68
+ CitadelSessionManager(SID, path(), "")
69
+
70
+
71
+ # ---- session / agent / message contracts vs the reference ----------------
72
+
73
+
74
+ def test_creating_a_duplicate_session_raises_like_the_reference():
75
+ m, r = mgr(), ref()
76
+ dup = Session(session_id=SID, session_type=SessionType.AGENT)
77
+ for store in (m, r):
78
+ with pytest.raises(SessionException, match="already exists"):
79
+ store.create_session(dup)
80
+
81
+
82
+ def test_reading_an_unknown_session_is_none_like_the_reference():
83
+ assert mgr().read_session("nope") is None
84
+ assert ref().read_session("nope") is None
85
+
86
+
87
+ def test_reading_an_unknown_agent_is_none_like_the_reference():
88
+ assert mgr().read_agent(SID, "nope") is None
89
+ assert ref().read_agent(SID, "nope") is None
90
+
91
+
92
+ def test_updating_an_unknown_agent_raises_like_the_reference():
93
+ m, r = mgr(), ref()
94
+ for store in (m, r):
95
+ with pytest.raises(SessionException, match="does not exist"):
96
+ store.update_agent(SID, agent("ghost"))
97
+
98
+
99
+ def test_reading_an_unknown_message_is_none_like_the_reference():
100
+ m, r = seeded(mgr()), seeded(ref())
101
+ assert m.read_message(SID, AID, 99) is None
102
+ assert r.read_message(SID, AID, 99) is None
103
+
104
+
105
+ def test_updating_an_unknown_message_raises_like_the_reference():
106
+ m, r = seeded(mgr()), seeded(ref())
107
+ for store in (m, r):
108
+ with pytest.raises(SessionException, match="does not exist"):
109
+ store.update_message(SID, AID, msg(99))
110
+
111
+
112
+ def test_listing_messages_for_an_unknown_agent_raises_like_the_reference():
113
+ m, r = mgr(), ref()
114
+ for store in (m, r):
115
+ with pytest.raises(SessionException):
116
+ store.list_messages(SID, "ghost")
117
+
118
+
119
+ # ---- ordering and pagination --------------------------------------------
120
+
121
+
122
+ def test_messages_come_back_in_conversation_order():
123
+ m = seeded(mgr(), 5)
124
+ assert [x.message_id for x in m.list_messages(SID, AID)] == [0, 1, 2, 3, 4]
125
+
126
+
127
+ def test_order_follows_the_index_not_the_write_time():
128
+ """Updating an early message must not move it to the end."""
129
+ m = seeded(mgr(), 5)
130
+ m.update_message(SID, AID, msg(1, "edited"))
131
+ assert [x.message_id for x in m.list_messages(SID, AID)] == [0, 1, 2, 3, 4]
132
+
133
+
134
+ def test_pagination_matches_the_reference():
135
+ m, r = seeded(mgr(), 6), seeded(ref(), 6)
136
+ for limit, offset in ((2, 0), (2, 2), (None, 3), (10, 0), (2, 99)):
137
+ ours = [x.message_id for x in m.list_messages(SID, AID, limit, offset)]
138
+ theirs = [x.message_id for x in r.list_messages(SID, AID, limit, offset)]
139
+ assert ours == theirs, (limit, offset, ours, theirs)
140
+
141
+
142
+ # ---- no duplicates on repeat writes -------------------------------------
143
+
144
+
145
+ def test_updating_a_message_does_not_duplicate_it():
146
+ m = seeded(mgr(), 3)
147
+ for text in ("one", "two", "three"):
148
+ m.update_message(SID, AID, msg(1, text))
149
+ assert len(m.list_messages(SID, AID)) == 3
150
+ assert m.read_message(SID, AID, 1).message["content"][0]["text"] == "three"
151
+
152
+
153
+ def test_updating_an_agent_does_not_duplicate_it():
154
+ m = seeded(mgr())
155
+ for _ in range(5):
156
+ a = agent()
157
+ a.state = {"n": 1}
158
+ m.update_agent(SID, a)
159
+ assert m.read_agent(SID, AID).state == {"n": 1}
160
+
161
+
162
+ # ---- redaction: the reason to use an encrypted store --------------------
163
+
164
+
165
+ def test_the_reference_keeps_the_redacted_text_on_disk():
166
+ """Established as the baseline this integration improves on."""
167
+ r = seeded(ref(), 1)
168
+ original = "my card number is 4111 1111 1111 1111"
169
+ r.create_message(SID, AID, msg(1, original))
170
+ m = r.read_message(SID, AID, 1)
171
+ m.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
172
+ r.update_message(SID, AID, m)
173
+
174
+ stored = r.read_message(SID, AID, 1)
175
+ assert stored.redact_message is not None
176
+ assert original in str(stored.message), "the reference keeps the original beside it"
177
+
178
+
179
+ def test_a_redaction_removes_the_original_from_the_store():
180
+ """Keeping the original would leave the redacted content in the store."""
181
+ m = seeded(mgr(), 0)
182
+ secret = "my card number is 4111 1111 1111 1111"
183
+ m.create_message(SID, AID, msg(0, secret))
184
+
185
+ stored = m.read_message(SID, AID, 0)
186
+ stored.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
187
+ m.update_message(SID, AID, stored)
188
+
189
+ after = m.read_message(SID, AID, 0)
190
+ assert after.to_message() == {"role": "user", "content": [{"text": "[REDACTED]"}]}
191
+ assert secret not in str(after.message), "the original survived the redaction"
192
+ assert secret not in str(after.to_dict()), "the original survived the redaction"
193
+
194
+
195
+ def test_a_redaction_at_creation_time_also_drops_the_original():
196
+ m = seeded(mgr(), 0)
197
+ secret = "another secret"
198
+ sm = msg(0, secret)
199
+ sm.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
200
+ m.create_message(SID, AID, sm)
201
+ assert secret not in str(m.read_message(SID, AID, 0).to_dict())
202
+
203
+
204
+ def test_an_unredacted_message_keeps_its_content():
205
+ """Only a redaction drops the original; an ordinary update must not."""
206
+ m = seeded(mgr(), 0)
207
+ m.create_message(SID, AID, msg(0, "ordinary content"))
208
+ m.update_message(SID, AID, msg(0, "edited content"))
209
+ after = m.read_message(SID, AID, 0)
210
+ assert after.to_message()["content"][0]["text"] == "edited content"
211
+ assert after.redact_message is None
212
+
213
+
214
+ def test_a_redaction_replaces_the_stored_record():
215
+ m = seeded(mgr(), 1)
216
+ original = "my card number is 4111 1111 1111 1111"
217
+ m.create_message(SID, AID, msg(1, original))
218
+
219
+ stored = m.read_message(SID, AID, 1)
220
+ stored.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
221
+ m.update_message(SID, AID, stored)
222
+
223
+ # One record for the index, carrying the redaction.
224
+ assert len([x for x in m.list_messages(SID, AID) if x.message_id == 1]) == 1
225
+ assert m.read_message(SID, AID, 1).redact_message is not None
226
+
227
+
228
+ def test_created_at_survives_an_update_like_the_reference():
229
+ m, r = seeded(mgr(), 1), seeded(ref(), 1)
230
+ for store in (m, r):
231
+ first = store.read_message(SID, AID, 0)
232
+ edited = msg(0, "edited")
233
+ store.update_message(SID, AID, edited)
234
+ assert store.read_message(SID, AID, 0).created_at == first.created_at
235
+
236
+
237
+ # ---- erasure -------------------------------------------------------------
238
+
239
+
240
+ def test_forget_session_destroys_everything_it_owns():
241
+ p = path()
242
+ m = seeded(mgr(p), 3)
243
+ assert m.forget_session(SID) == 5 # session + agent + 3 messages
244
+ assert m.read_session(SID) is None
245
+ assert m.read_agent(SID, AID) is None
246
+
247
+
248
+ def test_forgetting_an_unknown_session_is_zero():
249
+ assert mgr().forget_session("never") == 0
250
+
251
+
252
+ # ---- persistence ---------------------------------------------------------
253
+
254
+
255
+ def test_it_survives_a_reopen():
256
+ p = path()
257
+ first = seeded(mgr(p), 3)
258
+ first.update_agent(SID, SessionAgent(agent_id=AID, state={"k": "v"},
259
+ conversation_manager_state={}))
260
+ del first
261
+ gc.collect()
262
+
263
+ again = CitadelSessionManager(SID, p, "pw")
264
+ assert again.read_session(SID) is not None
265
+ assert again.read_agent(SID, AID).state == {"k": "v"}
266
+ assert [x.message_id for x in again.list_messages(SID, AID)] == [0, 1, 2]
267
+
268
+
269
+ def test_a_wrong_passphrase_cannot_reopen():
270
+ import citadeldb
271
+
272
+ p = path()
273
+ first = seeded(mgr(p))
274
+ del first
275
+ gc.collect()
276
+ with pytest.raises(citadeldb.EncryptionError):
277
+ CitadelSessionManager(SID, p, "wrong")
278
+
279
+
280
+ # ---- concurrency ---------------------------------------------------------
281
+
282
+
283
+ def test_concurrent_message_writes_all_land():
284
+ m = seeded(mgr(), 0)
285
+ with cf.ThreadPoolExecutor(max_workers=4) as ex:
286
+ list(ex.map(lambda i: m.create_message(SID, AID, msg(i, f"t{i}")), range(40)))
287
+ assert len(m.list_messages(SID, AID)) == 40
288
+
289
+
290
+ def test_a_message_id_is_stored_once_however_often_it_is_written():
291
+ """message_id is the conversation index, so a second write of one id
292
+ replaces it. Two request handlers over one store is the normal shape."""
293
+ m = seeded(mgr(), 0)
294
+ m.create_message(SID, AID, msg(7, "from handler one"))
295
+ m.create_message(SID, AID, msg(7, "from handler two"))
296
+ stored = m.list_messages(SID, AID)
297
+ assert [s.message_id for s in stored] == [7]
298
+ assert m.read_message(SID, AID, 7).message["content"][0]["text"] == "from handler two"
299
+
300
+
301
+ def test_concurrent_writes_of_one_message_id_leave_one_record():
302
+ m = seeded(mgr(), 0)
303
+ with cf.ThreadPoolExecutor(max_workers=8) as ex:
304
+ list(ex.map(lambda i: m.create_message(SID, AID, msg(3, f"t{i}")), range(32)))
305
+ assert [s.message_id for s in m.list_messages(SID, AID)] == [3]
306
+
307
+
308
+ # ---- scoping and content shapes -----------------------------------------
309
+
310
+
311
+ def test_sessions_and_agents_are_scoped_exactly():
312
+ p = path()
313
+ m = CitadelSessionManager("sess", p, "pw")
314
+ other = CitadelSessionManager("sess-2", p, "pw")
315
+ m.create_agent("sess", agent())
316
+ other.create_agent("sess-2", agent())
317
+ m.create_message("sess", AID, msg(0, "mine"))
318
+ assert len(m.list_messages("sess", AID)) == 1
319
+ assert other.list_messages("sess-2", AID) == []
320
+
321
+
322
+ def test_agents_within_a_session_are_scoped_exactly():
323
+ m = mgr()
324
+ m.create_agent(SID, agent("a"))
325
+ m.create_agent(SID, agent("a-2"))
326
+ m.create_message(SID, "a", msg(0, "mine"))
327
+ assert len(m.list_messages(SID, "a")) == 1
328
+ assert m.list_messages(SID, "a-2") == []
329
+
330
+
331
+ def test_unicode_round_trips():
332
+ m = seeded(mgr(), 0)
333
+ text = "ünïcode — 中文 — emoji 🎉"
334
+ m.create_message(SID, AID, msg(0, text))
335
+ assert m.read_message(SID, AID, 0).message["content"][0]["text"] == text
336
+
337
+
338
+ def test_agent_state_named_like_internal_fields_round_trips():
339
+ m = seeded(mgr(), 0)
340
+ hostile = {"sid": "x", "aid": "x", "mid": "x", "message": "x", "agent": "x"}
341
+ a = agent()
342
+ a.state = dict(hostile)
343
+ m.update_agent(SID, a)
344
+ assert m.read_agent(SID, AID).state == hostile
345
+
346
+
347
+ def test_a_message_with_no_text_is_storable():
348
+ m = seeded(mgr(), 0)
349
+ m.create_message(SID, AID, SessionMessage(
350
+ message={"role": "assistant", "content": []}, message_id=0
351
+ ))
352
+ assert len(m.list_messages(SID, AID)) == 1
353
+
354
+
355
+ def test_two_managers_share_one_database_file():
356
+ p = path()
357
+ a = CitadelSessionManager("sa", p, "pw")
358
+ b = CitadelSessionManager("sb", p, "pw")
359
+ a.create_agent("sa", agent())
360
+ a.create_message("sa", AID, msg(0, "in a"))
361
+ assert a.read_session("sa") is not None and b.read_session("sb") is not None
362
+ assert b.read_session("sa") is not None, "one file, both sessions visible"