axio-context-sqlite 0.5.1__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,12 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.pyc
4
+ *.pyo
5
+ .mypy_cache/
6
+ .ruff_cache/
7
+ .pytest_cache/
8
+ *.egg-info/
9
+ dist/
10
+ build/
11
+ _build/
12
+ .DS_Store
@@ -0,0 +1,145 @@
1
+ Metadata-Version: 2.4
2
+ Name: axio-context-sqlite
3
+ Version: 0.5.1
4
+ Summary: SQLite-backed context store for Axio
5
+ Project-URL: Homepage, https://github.com/axio-agent/monorepo
6
+ Project-URL: Repository, https://github.com/axio-agent/monorepo
7
+ License: MIT
8
+ Keywords: agent,ai,context,llm,sqlite,storage
9
+ Requires-Python: >=3.12
10
+ Requires-Dist: aiosqlite>=0.20
11
+ Requires-Dist: axio
12
+ Description-Content-Type: text/markdown
13
+
14
+ # axio-context-sqlite
15
+
16
+ [![PyPI](https://img.shields.io/pypi/v/axio-context-sqlite)](https://pypi.org/project/axio-context-sqlite/)
17
+ [![Python](https://img.shields.io/pypi/pyversions/axio-context-sqlite)](https://pypi.org/project/axio-context-sqlite/)
18
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
19
+
20
+ SQLite-backed persistent context store for [axio](https://github.com/axio-agent/monorepo).
21
+
22
+ ## Installation
23
+
24
+ ```bash
25
+ pip install axio-context-sqlite
26
+ ```
27
+
28
+ ## Usage
29
+
30
+ Open a connection with `connect()`, then create a `SQLiteContextStore` bound to a
31
+ session. The caller owns the connection and is responsible for closing it.
32
+
33
+ <!-- name: test_readme_usage -->
34
+ ```python
35
+ import asyncio
36
+ import tempfile
37
+ import pathlib
38
+ from axio_context_sqlite import connect, SQLiteContextStore
39
+ from axio.messages import Message
40
+ from axio.blocks import TextBlock
41
+
42
+ async def main() -> None:
43
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
44
+ try:
45
+ store = SQLiteContextStore(conn, session_id="my-session")
46
+ await store.append(Message(role="user", content=[TextBlock(text="Hello")]))
47
+ history = await store.get_history()
48
+ assert len(history) == 1
49
+ finally:
50
+ await conn.close()
51
+
52
+ asyncio.run(main())
53
+ ```
54
+
55
+ `SQLiteContextStore` implements the `axio.context.ContextStore` ABC and persists
56
+ conversation history across process restarts. Multiple sessions can coexist in
57
+ the same database file, isolated by `session_id` and `project`.
58
+
59
+ ### Agent integration
60
+
61
+ <!--
62
+ name: test_readme_agent
63
+ ```python
64
+ from axio.testing import StubTransport, make_text_response
65
+ transport = StubTransport([make_text_response("Hi!")])
66
+ ```
67
+ -->
68
+ <!-- name: test_readme_agent -->
69
+ ```python
70
+ import asyncio
71
+ import tempfile
72
+ import pathlib
73
+ from axio.agent import Agent
74
+ from axio_context_sqlite import connect, SQLiteContextStore
75
+
76
+ async def main() -> None:
77
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
78
+ try:
79
+ ctx = SQLiteContextStore(conn, session_id="main")
80
+ agent = Agent(system="You are helpful.", tools=[], transport=transport)
81
+ result = await agent.run("Hello!", ctx)
82
+ assert result == "Hi!"
83
+ finally:
84
+ await conn.close()
85
+
86
+ asyncio.run(main())
87
+ ```
88
+
89
+ ### Listing sessions
90
+
91
+ <!-- name: test_readme_list_sessions -->
92
+ ```python
93
+ import asyncio
94
+ import tempfile
95
+ import pathlib
96
+ from axio_context_sqlite import connect, SQLiteContextStore
97
+ from axio.messages import Message
98
+ from axio.blocks import TextBlock
99
+
100
+ async def main() -> None:
101
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
102
+ try:
103
+ store = SQLiteContextStore(conn, session_id="main", project="/myproject")
104
+ await store.append(Message(role="user", content=[TextBlock(text="hi")]))
105
+ sessions = await store.list_sessions()
106
+ for s in sessions:
107
+ print(s.session_id, s.preview, s.message_count)
108
+ assert len(sessions) == 1
109
+ finally:
110
+ await conn.close()
111
+
112
+ asyncio.run(main())
113
+ ```
114
+
115
+ ### Forking
116
+
117
+ `fork()` copies the current session's messages into a new session — useful for
118
+ branching conversations without affecting the original:
119
+
120
+ <!-- name: test_readme_fork -->
121
+ ```python
122
+ import asyncio
123
+ import tempfile
124
+ import pathlib
125
+ from axio_context_sqlite import connect, SQLiteContextStore
126
+ from axio.messages import Message
127
+ from axio.blocks import TextBlock
128
+
129
+ async def main() -> None:
130
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
131
+ try:
132
+ store = SQLiteContextStore(conn, session_id="main")
133
+ await store.append(Message(role="user", content=[TextBlock(text="original")]))
134
+ branch = await store.fork()
135
+ assert branch.session_id != store.session_id
136
+ assert len(await branch.get_history()) == 1
137
+ finally:
138
+ await conn.close()
139
+
140
+ asyncio.run(main())
141
+ ```
142
+
143
+ ## License
144
+
145
+ MIT
@@ -0,0 +1,132 @@
1
+ # axio-context-sqlite
2
+
3
+ [![PyPI](https://img.shields.io/pypi/v/axio-context-sqlite)](https://pypi.org/project/axio-context-sqlite/)
4
+ [![Python](https://img.shields.io/pypi/pyversions/axio-context-sqlite)](https://pypi.org/project/axio-context-sqlite/)
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
6
+
7
+ SQLite-backed persistent context store for [axio](https://github.com/axio-agent/monorepo).
8
+
9
+ ## Installation
10
+
11
+ ```bash
12
+ pip install axio-context-sqlite
13
+ ```
14
+
15
+ ## Usage
16
+
17
+ Open a connection with `connect()`, then create a `SQLiteContextStore` bound to a
18
+ session. The caller owns the connection and is responsible for closing it.
19
+
20
+ <!-- name: test_readme_usage -->
21
+ ```python
22
+ import asyncio
23
+ import tempfile
24
+ import pathlib
25
+ from axio_context_sqlite import connect, SQLiteContextStore
26
+ from axio.messages import Message
27
+ from axio.blocks import TextBlock
28
+
29
+ async def main() -> None:
30
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
31
+ try:
32
+ store = SQLiteContextStore(conn, session_id="my-session")
33
+ await store.append(Message(role="user", content=[TextBlock(text="Hello")]))
34
+ history = await store.get_history()
35
+ assert len(history) == 1
36
+ finally:
37
+ await conn.close()
38
+
39
+ asyncio.run(main())
40
+ ```
41
+
42
+ `SQLiteContextStore` implements the `axio.context.ContextStore` ABC and persists
43
+ conversation history across process restarts. Multiple sessions can coexist in
44
+ the same database file, isolated by `session_id` and `project`.
45
+
46
+ ### Agent integration
47
+
48
+ <!--
49
+ name: test_readme_agent
50
+ ```python
51
+ from axio.testing import StubTransport, make_text_response
52
+ transport = StubTransport([make_text_response("Hi!")])
53
+ ```
54
+ -->
55
+ <!-- name: test_readme_agent -->
56
+ ```python
57
+ import asyncio
58
+ import tempfile
59
+ import pathlib
60
+ from axio.agent import Agent
61
+ from axio_context_sqlite import connect, SQLiteContextStore
62
+
63
+ async def main() -> None:
64
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
65
+ try:
66
+ ctx = SQLiteContextStore(conn, session_id="main")
67
+ agent = Agent(system="You are helpful.", tools=[], transport=transport)
68
+ result = await agent.run("Hello!", ctx)
69
+ assert result == "Hi!"
70
+ finally:
71
+ await conn.close()
72
+
73
+ asyncio.run(main())
74
+ ```
75
+
76
+ ### Listing sessions
77
+
78
+ <!-- name: test_readme_list_sessions -->
79
+ ```python
80
+ import asyncio
81
+ import tempfile
82
+ import pathlib
83
+ from axio_context_sqlite import connect, SQLiteContextStore
84
+ from axio.messages import Message
85
+ from axio.blocks import TextBlock
86
+
87
+ async def main() -> None:
88
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
89
+ try:
90
+ store = SQLiteContextStore(conn, session_id="main", project="/myproject")
91
+ await store.append(Message(role="user", content=[TextBlock(text="hi")]))
92
+ sessions = await store.list_sessions()
93
+ for s in sessions:
94
+ print(s.session_id, s.preview, s.message_count)
95
+ assert len(sessions) == 1
96
+ finally:
97
+ await conn.close()
98
+
99
+ asyncio.run(main())
100
+ ```
101
+
102
+ ### Forking
103
+
104
+ `fork()` copies the current session's messages into a new session — useful for
105
+ branching conversations without affecting the original:
106
+
107
+ <!-- name: test_readme_fork -->
108
+ ```python
109
+ import asyncio
110
+ import tempfile
111
+ import pathlib
112
+ from axio_context_sqlite import connect, SQLiteContextStore
113
+ from axio.messages import Message
114
+ from axio.blocks import TextBlock
115
+
116
+ async def main() -> None:
117
+ conn = await connect(pathlib.Path(tempfile.mkdtemp()) / "chat.db")
118
+ try:
119
+ store = SQLiteContextStore(conn, session_id="main")
120
+ await store.append(Message(role="user", content=[TextBlock(text="original")]))
121
+ branch = await store.fork()
122
+ assert branch.session_id != store.session_id
123
+ assert len(await branch.get_history()) == 1
124
+ finally:
125
+ await conn.close()
126
+
127
+ asyncio.run(main())
128
+ ```
129
+
130
+ ## License
131
+
132
+ MIT
@@ -0,0 +1,49 @@
1
+ [project]
2
+ name = "axio-context-sqlite"
3
+ version = "0.5.1"
4
+ description = "SQLite-backed context store for Axio"
5
+ readme = "README.md"
6
+ requires-python = ">=3.12"
7
+ license = {text = "MIT"}
8
+ keywords = ["llm", "agent", "ai", "sqlite", "context", "storage"]
9
+ dependencies = ["axio", "aiosqlite>=0.20"]
10
+
11
+ [project.urls]
12
+ Homepage = "https://github.com/axio-agent/monorepo"
13
+ Repository = "https://github.com/axio-agent/monorepo"
14
+
15
+ [build-system]
16
+ requires = ["hatchling"]
17
+ build-backend = "hatchling.build"
18
+
19
+ [tool.hatch.build.targets.wheel]
20
+ packages = ["src/axio_context_sqlite"]
21
+ sources = ["src"]
22
+
23
+ [tool.pytest.ini_options]
24
+ asyncio_mode = "auto"
25
+ testpaths = ["tests", "README.md"]
26
+
27
+ [tool.ruff]
28
+ line-length = 119
29
+ target-version = "py312"
30
+
31
+ [tool.ruff.lint]
32
+ select = ["E", "F", "I", "UP"]
33
+
34
+ [tool.mypy]
35
+ strict = true
36
+ python_version = "3.12"
37
+
38
+ [dependency-groups]
39
+ dev = [
40
+ "pytest>=8",
41
+ "pytest-asyncio>=0.24",
42
+ "mypy>=1.14",
43
+ "ruff>=0.9",
44
+ "pytest-cov>=7.1.0",
45
+ "markdown-pytest>=0.6.0",
46
+ ]
47
+
48
+ [tool.uv.sources]
49
+ axio = { workspace = true }
@@ -0,0 +1,5 @@
1
+ """axio-context-sqlite: SQLite-backed context store for Axio."""
2
+
3
+ from axio_context_sqlite.store import SQLiteContextStore, connect
4
+
5
+ __all__ = ["SQLiteContextStore", "connect"]
@@ -0,0 +1,216 @@
1
+ """SQLiteContextStore: persistent conversation storage backed by SQLite."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import base64
7
+ import gzip
8
+ import json
9
+ from pathlib import Path
10
+ from uuid import uuid4
11
+
12
+ import aiosqlite
13
+ from axio.context import ContextStore, SessionInfo
14
+ from axio.messages import Message
15
+
16
+ # Compress content payloads above this size (bytes of UTF-8 JSON).
17
+ COMPRESS_THRESHOLD = 512
18
+
19
+
20
+ def compress_payload(data: str) -> str:
21
+ raw = data.encode()
22
+ if len(raw) < COMPRESS_THRESHOLD:
23
+ return "plain:" + data
24
+ return "gzip:" + base64.b64encode(gzip.compress(raw, compresslevel=6)).decode()
25
+
26
+
27
+ def decompress_payload(data: str) -> str:
28
+ if data.startswith("gzip:"):
29
+ return gzip.decompress(base64.b64decode(data[5:])).decode()
30
+ if data.startswith("plain:"):
31
+ return data[6:]
32
+ # raw JSON
33
+ return data
34
+
35
+
36
+ async def connect(db_path: str | Path) -> aiosqlite.Connection:
37
+ """Open (or create) a SQLite database and initialise the schema.
38
+
39
+ The caller is responsible for closing the returned connection.
40
+ """
41
+ path = Path(db_path)
42
+ await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
43
+ conn = await aiosqlite.connect(str(path))
44
+ await conn.create_function("compress_payload", 1, compress_payload, deterministic=True)
45
+ await conn.create_function("decompress_payload", 1, decompress_payload, deterministic=True)
46
+ await conn.execute("PRAGMA journal_mode=WAL")
47
+ await conn.execute("PRAGMA busy_timeout=5000")
48
+ await conn.execute("PRAGMA synchronous=NORMAL")
49
+ await conn.execute(
50
+ "CREATE TABLE IF NOT EXISTS axio_context_messages ("
51
+ " id INTEGER PRIMARY KEY AUTOINCREMENT,"
52
+ " session_id TEXT NOT NULL,"
53
+ " project TEXT NOT NULL,"
54
+ " position INTEGER NOT NULL,"
55
+ " role TEXT NOT NULL,"
56
+ " content TEXT NOT NULL,"
57
+ " created_at TEXT NOT NULL DEFAULT (datetime('now')),"
58
+ " UNIQUE(session_id, position)"
59
+ ")"
60
+ )
61
+ await conn.execute(
62
+ "CREATE INDEX IF NOT EXISTS idx_axio_context_messages_session ON axio_context_messages(session_id)"
63
+ )
64
+ await conn.execute(
65
+ "CREATE INDEX IF NOT EXISTS idx_axio_context_messages_project ON axio_context_messages(project)"
66
+ )
67
+ await conn.execute(
68
+ "CREATE TABLE IF NOT EXISTS axio_context_tokens ("
69
+ " session_id TEXT NOT NULL,"
70
+ " project TEXT NOT NULL,"
71
+ " input_tokens INTEGER NOT NULL DEFAULT 0,"
72
+ " output_tokens INTEGER NOT NULL DEFAULT 0,"
73
+ " PRIMARY KEY(session_id, project)"
74
+ ")"
75
+ )
76
+ await conn.commit()
77
+ return conn
78
+
79
+
80
+ def _extract_preview(content_json: str, max_len: int = 80) -> str:
81
+ """Extract text preview from serialized content JSON."""
82
+ try:
83
+ blocks = json.loads(content_json)
84
+ for b in blocks:
85
+ if b.get("type") == "text":
86
+ text: str = b["text"]
87
+ return text[:max_len] + ("..." if len(text) > max_len else "")
88
+ except (json.JSONDecodeError, KeyError):
89
+ pass
90
+ return "(no preview)"
91
+
92
+
93
+ class SQLiteContextStore(ContextStore):
94
+ """Persistent conversation storage backed by SQLite.
95
+
96
+ The caller owns the connection and is responsible for closing it.
97
+ Use :func:`connect` to open a properly initialized connection.
98
+ """
99
+
100
+ def __init__(
101
+ self,
102
+ conn: aiosqlite.Connection,
103
+ session_id: str,
104
+ project: str | None = None,
105
+ db_name: str = "axio_context",
106
+ ) -> None:
107
+ self._conn = conn
108
+ self._db_name = db_name
109
+ self._session_id = session_id
110
+ self._project = project or str(Path.cwd().resolve())
111
+
112
+ @property
113
+ def session_id(self) -> str:
114
+ return self._session_id
115
+
116
+ async def append(self, message: Message) -> None:
117
+ content_json = json.dumps(message.to_dict()["content"])
118
+ await self._conn.execute(
119
+ "INSERT INTO axio_context_messages (session_id, project, position, role, content)"
120
+ "VALUES (?, ?, (SELECT COUNT(*) FROM axio_context_messages WHERE session_id = ?), ?, compress_payload(?))",
121
+ (self._session_id, self._project, self._session_id, message.role, content_json),
122
+ )
123
+ await self._conn.commit()
124
+
125
+ async def get_history(self) -> list[Message]:
126
+ async with self._conn.execute(
127
+ "SELECT role, decompress_payload(content) FROM axio_context_messages"
128
+ " WHERE session_id = ? ORDER BY position",
129
+ (self._session_id,),
130
+ ) as cursor:
131
+ rows = await cursor.fetchall()
132
+ return [Message.from_dict({"role": role, "content": json.loads(content)}) for role, content in rows]
133
+
134
+ async def clear(self) -> None:
135
+ await self._conn.execute("DELETE FROM axio_context_messages WHERE session_id = ?", (self._session_id,))
136
+ await self._conn.execute(
137
+ "DELETE FROM axio_context_tokens WHERE session_id = ? AND project = ?",
138
+ (self._session_id, self._project),
139
+ )
140
+ await self._conn.commit()
141
+
142
+ async def fork(self) -> SQLiteContextStore:
143
+ new_id = uuid4().hex
144
+ await self._conn.execute(
145
+ "INSERT INTO axio_context_messages (session_id, project, position, role, content)"
146
+ "SELECT ?, project, position, role, content FROM axio_context_messages WHERE session_id = ?",
147
+ (new_id, self._session_id),
148
+ )
149
+ await self._conn.execute(
150
+ "INSERT OR IGNORE INTO axio_context_tokens (session_id, project, input_tokens, output_tokens) "
151
+ "SELECT ?, project, input_tokens, output_tokens FROM axio_context_tokens "
152
+ "WHERE session_id = ? AND project = ?",
153
+ (new_id, self._session_id, self._project),
154
+ )
155
+ await self._conn.commit()
156
+ return SQLiteContextStore(self._conn, new_id, self._project)
157
+
158
+ async def set_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
159
+ await self._conn.execute(
160
+ "INSERT INTO axio_context_tokens (session_id, project, input_tokens, output_tokens)"
161
+ "VALUES (?, ?, ?, ?) "
162
+ "ON CONFLICT(session_id, project) DO UPDATE SET input_tokens=?, output_tokens=?",
163
+ (self._session_id, self._project, input_tokens, output_tokens, input_tokens, output_tokens),
164
+ )
165
+ await self._conn.commit()
166
+
167
+ async def add_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
168
+ await self._conn.execute(
169
+ "INSERT INTO axio_context_tokens (session_id, project, input_tokens, output_tokens)"
170
+ "VALUES (?, ?, ?, ?) "
171
+ "ON CONFLICT(session_id, project) DO UPDATE "
172
+ "SET input_tokens = input_tokens + excluded.input_tokens, "
173
+ " output_tokens = output_tokens + excluded.output_tokens",
174
+ (self._session_id, self._project, input_tokens, output_tokens),
175
+ )
176
+ await self._conn.commit()
177
+
178
+ async def get_context_tokens(self) -> tuple[int, int]:
179
+ async with self._conn.execute(
180
+ "SELECT input_tokens, output_tokens FROM axio_context_tokens WHERE session_id = ? AND project = ?",
181
+ (self._session_id, self._project),
182
+ ) as cursor:
183
+ row = await cursor.fetchone()
184
+ if row is None:
185
+ return 0, 0
186
+ return int(row[0]), int(row[1])
187
+
188
+ async def list_sessions(self) -> list[SessionInfo]:
189
+ """List all sessions for a project, newest first."""
190
+ async with self._conn.execute(
191
+ "SELECT m.session_id, COUNT(*) as cnt, "
192
+ "(SELECT decompress_payload(content) FROM axio_context_messages WHERE session_id = m.session_id "
193
+ "AND role = 'user' ORDER BY position LIMIT 1) as first_content, "
194
+ "MIN(m.created_at) as created, "
195
+ "COALESCE(ct.input_tokens, 0), COALESCE(ct.output_tokens, 0) "
196
+ "FROM axio_context_messages m "
197
+ "LEFT JOIN axio_context_tokens ct ON ct.session_id = m.session_id AND ct.project = m.project "
198
+ "WHERE m.project = ? "
199
+ "GROUP BY m.session_id ORDER BY created DESC",
200
+ (self._project,),
201
+ ) as cursor:
202
+ rows = await cursor.fetchall()
203
+ result: list[SessionInfo] = []
204
+ for session_id, count, first_content, created_at, in_tok, out_tok in rows:
205
+ preview = _extract_preview(first_content) if first_content else "(no preview)"
206
+ result.append(
207
+ SessionInfo(
208
+ session_id=session_id,
209
+ message_count=count,
210
+ preview=preview,
211
+ created_at=created_at,
212
+ input_tokens=int(in_tok),
213
+ output_tokens=int(out_tok),
214
+ )
215
+ )
216
+ return result
@@ -0,0 +1,164 @@
1
+ """Tests for SQLiteContextStore."""
2
+
3
+ from collections.abc import AsyncGenerator
4
+ from pathlib import Path
5
+
6
+ import aiosqlite
7
+ import pytest
8
+ from axio.blocks import TextBlock
9
+ from axio.messages import Message
10
+
11
+ from axio_context_sqlite import SQLiteContextStore, connect
12
+ from axio_context_sqlite.store import COMPRESS_THRESHOLD, compress_payload, decompress_payload
13
+
14
+
15
+ def _msg(role: str, text: str) -> Message:
16
+ return Message(role=role, content=[TextBlock(text=text)]) # type: ignore[arg-type]
17
+
18
+
19
+ @pytest.fixture
20
+ def db_path(tmp_path: Path) -> Path:
21
+ return tmp_path / "test.db"
22
+
23
+
24
+ @pytest.fixture
25
+ async def conn(db_path: Path) -> AsyncGenerator[aiosqlite.Connection, None]:
26
+ c = await connect(db_path)
27
+ yield c
28
+ await c.close()
29
+
30
+
31
+ @pytest.fixture
32
+ async def store(conn: aiosqlite.Connection) -> SQLiteContextStore:
33
+ return SQLiteContextStore(conn, "session-1", "test-project")
34
+
35
+
36
+ async def test_append_and_get_history(store: SQLiteContextStore) -> None:
37
+ await store.append(_msg("user", "Hello"))
38
+ await store.append(_msg("assistant", "Hi!"))
39
+ history = await store.get_history()
40
+ assert len(history) == 2
41
+ assert history[0].role == "user"
42
+ assert history[1].role == "assistant"
43
+
44
+
45
+ async def test_clear(store: SQLiteContextStore) -> None:
46
+ await store.append(_msg("user", "Hello"))
47
+ await store.clear()
48
+ assert await store.get_history() == []
49
+
50
+
51
+ async def test_fork(store: SQLiteContextStore) -> None:
52
+ await store.append(_msg("user", "Hello"))
53
+ forked = await store.fork()
54
+ await forked.append(_msg("assistant", "Hi!"))
55
+ assert len(await store.get_history()) == 1
56
+ assert len(await forked.get_history()) == 2
57
+
58
+
59
+ async def test_set_get_context_tokens(store: SQLiteContextStore) -> None:
60
+ await store.set_context_tokens(100, 50)
61
+ inp, out = await store.get_context_tokens()
62
+ assert inp == 100
63
+ assert out == 50
64
+
65
+
66
+ async def test_add_context_tokens(store: SQLiteContextStore) -> None:
67
+ await store.set_context_tokens(100, 50)
68
+ await store.add_context_tokens(20, 10)
69
+ inp, out = await store.get_context_tokens()
70
+ assert inp == 120
71
+ assert out == 60
72
+
73
+
74
+ async def test_list_sessions(db_path: Path) -> None:
75
+ c = await connect(db_path)
76
+ try:
77
+ s1 = SQLiteContextStore(c, "sess-a", "proj")
78
+ s2 = SQLiteContextStore(c, "sess-b", "proj")
79
+ await s1.append(_msg("user", "First session"))
80
+ await s2.append(_msg("user", "Second session"))
81
+ sessions = await s1.list_sessions()
82
+ assert len(sessions) == 2
83
+ ids = {s.session_id for s in sessions}
84
+ assert ids == {"sess-a", "sess-b"}
85
+ previews = {s.preview for s in sessions}
86
+ assert "First session" in previews
87
+ assert "Second session" in previews
88
+ finally:
89
+ await c.close()
90
+
91
+
92
+ class TestCompressPayload:
93
+ def test_small_stays_plain(self) -> None:
94
+ data = "hello"
95
+ result = compress_payload(data)
96
+ assert result.startswith("plain:")
97
+ assert decompress_payload(result) == data
98
+
99
+ def test_large_gets_compressed(self) -> None:
100
+ data = "x" * COMPRESS_THRESHOLD
101
+ result = compress_payload(data)
102
+ assert result.startswith("gzip:")
103
+ assert decompress_payload(result) == data
104
+
105
+ def test_legacy_no_prefix(self) -> None:
106
+ assert decompress_payload('{"foo": 1}') == '{"foo": 1}'
107
+
108
+ def test_roundtrip(self) -> None:
109
+ data = "[" + '{"type":"text","text":"a"},' * 50 + "]"
110
+ assert decompress_payload(compress_payload(data)) == data
111
+
112
+
113
+ async def test_large_message_compressed_on_disk(db_path: Path) -> None:
114
+ """Large content is stored compressed; get_history still returns original."""
115
+ big_text = "word " * 200 # well above threshold
116
+
117
+ c = await connect(db_path)
118
+ try:
119
+ s = SQLiteContextStore(c, "big-session", "proj")
120
+ await s.append(_msg("user", big_text))
121
+ finally:
122
+ await c.close()
123
+
124
+ # Inspect raw bytes on disk — should start with gzip:
125
+ async with aiosqlite.connect(str(db_path)) as raw:
126
+ async with raw.execute("SELECT content FROM axio_context_messages") as cur:
127
+ row = await cur.fetchone()
128
+ assert row is not None
129
+ assert row[0].startswith("gzip:")
130
+
131
+ # But get_history transparently decompresses
132
+ c2 = await connect(db_path)
133
+ try:
134
+ s2 = SQLiteContextStore(c2, "big-session", "proj")
135
+ history = await s2.get_history()
136
+ finally:
137
+ await c2.close()
138
+ assert len(history) == 1
139
+ block = history[0].content[0]
140
+ assert isinstance(block, TextBlock)
141
+ assert block.text == big_text
142
+
143
+
144
+ async def test_close_and_reopen(db_path: Path) -> None:
145
+ """Data persists after close and reopen."""
146
+ c = await connect(db_path)
147
+ try:
148
+ s = SQLiteContextStore(c, "persist-session", "proj")
149
+ await s.append(_msg("user", "Persistent message"))
150
+ await s.set_context_tokens(42, 7)
151
+ finally:
152
+ await c.close()
153
+
154
+ c2 = await connect(db_path)
155
+ try:
156
+ s2 = SQLiteContextStore(c2, "persist-session", "proj")
157
+ history = await s2.get_history()
158
+ assert len(history) == 1
159
+ assert history[0].role == "user"
160
+ inp, out = await s2.get_context_tokens()
161
+ assert inp == 42
162
+ assert out == 7
163
+ finally:
164
+ await c2.close()