"""19: A custom SessionStore backed by SQLite, validated with the SDK's conformance suite,
then used to resume a session 'on another host' (a different config dir)."""
import asyncio
import json
import sqlite3
import tempfile
import time
from pathlib import Path
from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, fold_session_summary, query
from claude_agent_sdk.testing import run_session_store_conformance
from demo_util import PROJECT
class SqliteSessionStore:
"""Two required methods (append, load) plus the optional ones, on one SQLite file."""
def __init__(self, path: str):
self.db = sqlite3.connect(path, check_same_thread=False)
self.db.executescript("""
CREATE TABLE IF NOT EXISTS entries (
project_key TEXT, session_id TEXT, subpath TEXT, seq INTEGER, uuid TEXT, entry TEXT,
PRIMARY KEY (project_key, session_id, subpath, seq));
CREATE TABLE IF NOT EXISTS summaries (
project_key TEXT, session_id TEXT, mtime INTEGER, data TEXT,
PRIMARY KEY (project_key, session_id));
""")
@staticmethod
def _sub(key):
return key.get("subpath") or ""
async def append(self, key, entries):
pk, sid, sub = key["project_key"], key["session_id"], self._sub(key)
with self.db:
seen = {r[0] for r in self.db.execute(
"SELECT uuid FROM entries WHERE project_key=? AND session_id=? AND subpath=? AND uuid IS NOT NULL", (pk, sid, sub))}
row = self.db.execute("SELECT COALESCE(MAX(seq), -1) FROM entries WHERE project_key=? AND session_id=? AND subpath=?",
(pk, sid, sub)).fetchone()
seq = row[0] + 1
for e in entries:
uuid = e.get("uuid")
if uuid and uuid in seen: # a retried batch may re-deliver entries: dedupe by uuid
continue
self.db.execute("INSERT INTO entries VALUES (?,?,?,?,?,?)", (pk, sid, sub, seq, uuid, json.dumps(e)))
seq += 1
if not sub: # keep the summary sidecar for the main transcript only
cur = self.db.execute("SELECT data FROM summaries WHERE project_key=? AND session_id=?", (pk, sid)).fetchone()
prev = {"session_id": sid, "mtime": 0, "data": json.loads(cur[0])} if cur else None
summary = fold_session_summary(prev, key, entries) # SDK helper: fold new entries into the summary
self.db.execute("INSERT OR REPLACE INTO summaries VALUES (?,?,?,?)",
(pk, sid, int(time.time() * 1000), json.dumps(summary["data"])))
async def load(self, key):
rows = self.db.execute("SELECT entry FROM entries WHERE project_key=? AND session_id=? AND subpath=? ORDER BY seq",
(key["project_key"], key["session_id"], self._sub(key))).fetchall()
return [json.loads(r[0]) for r in rows] if rows else None
async def list_sessions(self, project_key):
rows = self.db.execute("SELECT session_id, mtime FROM summaries WHERE project_key=? ORDER BY mtime DESC", (project_key,))
return [{"session_id": r[0], "mtime": r[1]} for r in rows]
async def list_session_summaries(self, project_key):
rows = self.db.execute("SELECT session_id, mtime, data FROM summaries WHERE project_key=? ORDER BY mtime DESC", (project_key,))
return [{"session_id": r[0], "mtime": r[1], "data": json.loads(r[2])} for r in rows]
async def delete(self, key):
with self.db:
if self._sub(key):
self.db.execute("DELETE FROM entries WHERE project_key=? AND session_id=? AND subpath=?",
(key["project_key"], key["session_id"], self._sub(key)))
else: # main key: cascade to subkeys and summary
self.db.execute("DELETE FROM entries WHERE project_key=? AND session_id=?", (key["project_key"], key["session_id"]))
self.db.execute("DELETE FROM summaries WHERE project_key=? AND session_id=?", (key["project_key"], key["session_id"]))
async def list_subkeys(self, key):
rows = self.db.execute("SELECT DISTINCT subpath FROM entries WHERE project_key=? AND session_id=? AND subpath<>''",
(key["project_key"], key["session_id"]))
return [r[0] for r in rows]
async def main():
# 1. Conformance: the SDK's own contract tests, each on a fresh empty store
print("=== conformance suite ===")
await run_session_store_conformance(lambda: SqliteSessionStore(":memory:"))
print("conformance: PASSED")
# 2. Host A writes a session into the shared store
db_path = str(Path(tempfile.mkdtemp()) / "sessions.db")
store = SqliteSessionStore(db_path)
host_a = tempfile.mkdtemp(prefix="hostA-config-")
session_id = None
async for m in query(prompt="Remember the code word OSPREY. Reply only: noted",
options=ClaudeAgentOptions(cwd=PROJECT, tools=[], max_turns=2, session_store=store,
env={"CLAUDE_CONFIG_DIR": host_a})):
if isinstance(m, ResultMessage):
session_id = m.session_id
print(f"\nhost A: session {session_id} -> {m.result!r}")
rows = store.db.execute("SELECT COUNT(*) FROM entries WHERE session_id=?", (session_id,)).fetchone()[0]
pk = store.db.execute("SELECT project_key FROM summaries LIMIT 1").fetchone()[0]
summaries = await store.list_session_summaries(pk)
print(f"store now holds {rows} transcript entries for that session (project_key={pk!r})")
print(f"summary sidecar: {json.dumps(summaries[0]['data'])[:200]}")
# 3. Host B (a different config dir, i.e. no local transcript) resumes from the store
host_b = tempfile.mkdtemp(prefix="hostB-config-")
async for m in query(prompt="What was the code word? One word.",
options=ClaudeAgentOptions(cwd=PROJECT, tools=[], max_turns=2, session_store=store,
resume=session_id, env={"CLAUDE_CONFIG_DIR": host_b})):
if isinstance(m, ResultMessage):
print(f"host B: resumed {m.session_id} -> {m.result!r}")
print(f"host B had no local transcript: {list(Path(host_b).glob('projects/*/*.jsonl'))}")
asyncio.run(main())
118 lines · Back to the folder