wehub-resource-sync
1463 lines
53 KiB
Python
1463 lines
53 KiB
Python
"""Tests for SessionManager lifecycle operations."""
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import stat
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from opensquilla.session.compaction import CompactionConfig
|
|
from opensquilla.session.manager import SessionManager
|
|
from opensquilla.session.models import (
|
|
SessionContextState,
|
|
SessionIntent,
|
|
SessionStatus,
|
|
SessionSummary,
|
|
TranscriptEntry,
|
|
)
|
|
from opensquilla.session.storage import SessionStorage, StaleEpochError
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def manager():
|
|
storage = SessionStorage(":memory:")
|
|
await storage.connect()
|
|
mgr = SessionManager(storage, inject_time_prefix=False)
|
|
yield mgr
|
|
await storage.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session(manager):
|
|
node = await manager.create("agent:main:main")
|
|
assert node.session_key == "agent:main:main"
|
|
assert node.status == SessionStatus.RUNNING
|
|
assert node.session_id is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_session_returns_existing_without_touching(manager):
|
|
node = await manager.create("agent:main:main")
|
|
|
|
fetched = await manager.get_session("agent:main:main")
|
|
missing = await manager.get_session("agent:main:missing")
|
|
|
|
assert fetched is not None
|
|
assert fetched.session_key == node.session_key
|
|
assert fetched.session_id == node.session_id
|
|
assert missing is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_duplicate_raises(manager):
|
|
await manager.create("agent:main:main")
|
|
with pytest.raises(ValueError):
|
|
await manager.create("agent:main:main")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_or_create_creates(manager):
|
|
node, created = await manager.get_or_create("agent:main:main")
|
|
assert created is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_or_create_returns_existing(manager):
|
|
await manager.create("agent:main:main")
|
|
node, created = await manager.get_or_create("agent:main:main")
|
|
assert created is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_intent_continue_preserves_existing_identity_and_transcript(manager):
|
|
node = await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "hello")
|
|
|
|
applied, rotated = await manager.apply_intent("agent:main:main", SessionIntent.CONTINUE)
|
|
|
|
assert rotated is False
|
|
assert applied.session_id == node.session_id
|
|
assert len(await manager.get_transcript("agent:main:main")) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_intent_new_chat_rejects_existing_key(manager):
|
|
await manager.create("agent:main:main")
|
|
|
|
with pytest.raises(ValueError, match="session_key conflict"):
|
|
await manager.apply_intent("agent:main:main", SessionIntent.NEW_CHAT)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_intent_reset_same_key_rotates_identity_and_clears_state(
|
|
manager, tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.setenv("OPENSQUILLA_SESSION_ARCHIVE_DIR", str(tmp_path / "archives"))
|
|
node = await manager.create("agent:main:main")
|
|
old_session_id = node.session_id
|
|
node.total_tokens = 123
|
|
node.input_tokens = 10
|
|
node.output_tokens = 20
|
|
node.estimated_cost_usd = 0.42
|
|
node.total_cost_usd = 0.42
|
|
node.billed_cost_usd = 0.30
|
|
node.estimated_cost_component_usd = 0.12
|
|
node.cost_source = "mixed"
|
|
node.missing_cost_entries = 1
|
|
node.cache_read = 7
|
|
node.cache_write = 8
|
|
await manager._storage.upsert_session(node)
|
|
await manager.append_message("agent:main:main", "user", "hello")
|
|
await manager._storage.save_summary(
|
|
SessionSummary(
|
|
session_id=old_session_id,
|
|
session_key="agent:main:main",
|
|
summary_text="old summary",
|
|
)
|
|
)
|
|
await manager.save_context_state(
|
|
SessionContextState(
|
|
session_id=old_session_id,
|
|
session_key="agent:main:main",
|
|
provider="portable",
|
|
state_kind="structured_summary_v1",
|
|
payload={"user_goal": "old task"},
|
|
covered_through_id=1,
|
|
valid=True,
|
|
)
|
|
)
|
|
|
|
applied, rotated = await manager.apply_intent("agent:main:main", SessionIntent.RESET_SAME_KEY)
|
|
|
|
assert rotated is True
|
|
assert applied.session_key == "agent:main:main"
|
|
assert applied.session_id != old_session_id
|
|
assert await manager._storage.count_transcript_entries(old_session_id) == 0
|
|
assert await manager._storage.count_transcript_entries(applied.session_id) == 0
|
|
assert await manager._storage.get_all_summaries(old_session_id) == []
|
|
assert await manager.get_context_states("agent:main:main") == []
|
|
invalidated_states = await manager.get_context_states("agent:main:main", valid_only=False)
|
|
assert len(invalidated_states) == 1
|
|
assert invalidated_states[0].valid is False
|
|
assert invalidated_states[0].invalid_reason == "session_reset"
|
|
assert applied.total_tokens == 0
|
|
assert applied.input_tokens == 0
|
|
assert applied.output_tokens == 0
|
|
assert applied.estimated_cost_usd == 0.0
|
|
assert applied.total_cost_usd == 0.0
|
|
assert applied.billed_cost_usd == 0.0
|
|
assert applied.estimated_cost_component_usd == 0.0
|
|
assert applied.cost_source == "none"
|
|
assert applied.missing_cost_entries == 0
|
|
assert applied.cache_read == 0
|
|
assert applied.cache_write == 0
|
|
archive_files = list((tmp_path / "archives").glob("*.json"))
|
|
assert len(archive_files) == 1
|
|
archived = json.loads(archive_files[0].read_text(encoding="utf-8"))
|
|
assert archived["session_key"] == "agent:main:main"
|
|
assert archived["session_id"] == old_session_id
|
|
assert archived["transcript_entries"][0]["content"] == "hello"
|
|
assert archived["summaries"][0]["summary_text"] == "old summary"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skipif(os.name == "nt", reason="POSIX permission semantics only")
|
|
async def test_reset_archive_is_owner_only(manager, tmp_path, monkeypatch):
|
|
archive_root = tmp_path / "session-archive"
|
|
monkeypatch.setenv("OPENSQUILLA_SESSION_ARCHIVE_DIR", str(archive_root))
|
|
old_umask = os.umask(0o022)
|
|
try:
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "hello")
|
|
|
|
_, rotated = await manager.apply_intent("agent:main:main", SessionIntent.RESET_SAME_KEY)
|
|
assert rotated is True
|
|
finally:
|
|
os.umask(old_umask)
|
|
|
|
archives = list(archive_root.glob("*.json"))
|
|
assert len(archives) == 1
|
|
# The archive holds the full raw transcript, so it must match the
|
|
# sessions.db hardening (0600 file inside a 0700 directory).
|
|
assert stat.S_IMODE(archive_root.stat().st_mode) & 0o077 == 0
|
|
assert stat.S_IMODE(archives[0].stat().st_mode) & 0o077 == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_same_key_fences_appends_that_read_the_old_epoch(
|
|
manager, tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.setenv("OPENSQUILLA_SESSION_ARCHIVE_DIR", str(tmp_path / "archives"))
|
|
node = await manager.create("agent:main:main")
|
|
pre_epoch = await manager._storage.get_epoch("agent:main:main")
|
|
|
|
applied, rotated = await manager.apply_intent("agent:main:main", SessionIntent.RESET_SAME_KEY)
|
|
|
|
assert rotated is True
|
|
# The rotation itself bumps the epoch, so a writer holding the pre-reset
|
|
# node cannot land its entry (or roll back the reset via a stale upsert).
|
|
assert await manager._storage.get_epoch("agent:main:main") > pre_epoch
|
|
with pytest.raises(StaleEpochError):
|
|
await manager._storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=node.session_id,
|
|
session_key="agent:main:main",
|
|
role="user",
|
|
content="stale in-flight message",
|
|
),
|
|
expected_epoch=pre_epoch,
|
|
)
|
|
row = await manager._storage.get_session("agent:main:main")
|
|
assert row is not None
|
|
assert row.session_id == applied.session_id
|
|
assert [e.content for e in await manager.get_transcript("agent:main:main")] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_same_key_archive_preserves_compacted_canonical_transcript(
|
|
manager, tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.setenv("OPENSQUILLA_SESSION_ARCHIVE_DIR", str(tmp_path / "archives"))
|
|
node = await manager.create("agent:main:main")
|
|
old_session_id = node.session_id
|
|
for index in range(4):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"short summary",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
compaction_id="cmp_reset_archive",
|
|
)
|
|
|
|
canonical_before_reset = [
|
|
entry.content for entry in await manager.get_canonical_transcript("agent:main:main")
|
|
]
|
|
|
|
applied, rotated = await manager.apply_intent("agent:main:main", SessionIntent.RESET_SAME_KEY)
|
|
|
|
assert rotated is True
|
|
assert applied.session_id != old_session_id
|
|
archive_files = list((tmp_path / "archives").glob("*.json"))
|
|
assert len(archive_files) == 1
|
|
archived = json.loads(archive_files[0].read_text(encoding="utf-8"))
|
|
assert [entry["content"] for entry in archived["transcript_entries"]] == (
|
|
canonical_before_reset
|
|
)
|
|
assert archived["summaries"][0]["compaction_id"] == "cmp_reset_archive"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_intent_reset_same_key_missing_creates_session(manager):
|
|
applied, rotated = await manager.apply_intent(
|
|
"agent:main:missing", SessionIntent.RESET_SAME_KEY
|
|
)
|
|
|
|
assert rotated is True
|
|
assert applied.session_key == "agent:main:missing"
|
|
assert applied.session_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_apply_intent_reset_same_key_archive_failure_does_not_block(
|
|
manager, tmp_path, monkeypatch
|
|
):
|
|
archive_file = tmp_path / "not-a-directory"
|
|
archive_file.write_text("occupied", encoding="utf-8")
|
|
monkeypatch.setenv("OPENSQUILLA_SESSION_ARCHIVE_DIR", str(archive_file))
|
|
node = await manager.create("agent:main:main")
|
|
old_session_id = node.session_id
|
|
await manager.append_message("agent:main:main", "user", "hello")
|
|
|
|
applied, rotated = await manager.apply_intent("agent:main:main", SessionIntent.RESET_SAME_KEY)
|
|
|
|
assert rotated is True
|
|
assert applied.session_id != old_session_id
|
|
assert await manager._storage.count_transcript_entries(old_session_id) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_touches_updated_at(manager):
|
|
node = await manager.create("agent:main:main")
|
|
old_ts = node.updated_at
|
|
import asyncio
|
|
|
|
await asyncio.sleep(0.01)
|
|
resumed = await manager.resume("agent:main:main")
|
|
assert resumed.updated_at >= old_ts
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_missing_raises(manager):
|
|
with pytest.raises(KeyError):
|
|
await manager.resume("agent:main:nope")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_fields(manager):
|
|
await manager.create("agent:main:main")
|
|
updated = await manager.update("agent:main:main", model="claude-opus-4-6", channel="telegram")
|
|
assert updated.model == "claude-opus-4-6"
|
|
assert updated.channel == "telegram"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finish_sets_status(manager):
|
|
await manager.create("agent:main:main")
|
|
node = await manager.finish("agent:main:main")
|
|
assert node.status == SessionStatus.DONE
|
|
assert node.ended_at is not None
|
|
assert node.runtime_ms is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finish_failed(manager):
|
|
await manager.create("agent:main:main")
|
|
node = await manager.finish("agent:main:main", status=SessionStatus.FAILED)
|
|
assert node.status == SessionStatus.FAILED
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_message(manager):
|
|
await manager.create("agent:main:main")
|
|
entry = await manager.append_message("agent:main:main", "user", "Hello!")
|
|
assert entry.role == "user"
|
|
assert entry.content == "Hello!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_message_updates_tokens(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "hi", token_count=10)
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.total_tokens == 10
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_message_with_turn_usage_does_not_double_count_output_tokens(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
"hi",
|
|
turn_usage={"model": "gpt-test", "input_tokens": 10, "output_tokens": 3},
|
|
token_count=3,
|
|
)
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.total_tokens == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_transcript(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "msg1")
|
|
await manager.append_message("agent:main:main", "assistant", "resp1")
|
|
entries = await manager.get_transcript("agent:main:main")
|
|
assert len(entries) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_transcript_orders_same_timestamp_by_insert_id(manager):
|
|
node = await manager.create("agent:main:main")
|
|
for content in ("first", "second", "third"):
|
|
await manager._storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=node.session_id,
|
|
session_key=node.session_key,
|
|
role="user",
|
|
content=content,
|
|
created_at=12345,
|
|
)
|
|
)
|
|
|
|
entries = await manager.get_transcript("agent:main:main")
|
|
|
|
assert [entry.content for entry in entries] == ["first", "second", "third"]
|
|
assert [entry.id for entry in entries] == sorted(entry.id for entry in entries)
|
|
|
|
|
|
def test_get_transcript_query_uses_id_tiebreaker() -> None:
|
|
source = Path("src/opensquilla/session/storage.py").read_text(encoding="utf-8")
|
|
|
|
assert "ORDER BY created_at ASC, id ASC LIMIT ? OFFSET ?" in source
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_truncate_zero_removes_all_entries(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "msg1")
|
|
await manager.append_message("agent:main:main", "assistant", "resp1")
|
|
|
|
result = await manager.truncate("agent:main:main", max_messages=0)
|
|
|
|
assert result == {"truncated": True, "before_count": 2, "after_count": 0}
|
|
assert await manager.get_transcript("agent:main:main") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_creates_child(manager):
|
|
await manager.create("agent:main:main")
|
|
child = await manager.branch("agent:main:main", "agent:main:direct:u1")
|
|
assert child.parent_session_key == "agent:main:main"
|
|
assert child.spawn_depth == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_transcript(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "parent msg", token_count=5)
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
"parent reply",
|
|
turn_usage={"model": "openai/gpt-test", "input_tokens": 11, "output_tokens": 5},
|
|
token_count=5,
|
|
)
|
|
child = await manager.branch("agent:main:main", "agent:main:direct:u1", fork_transcript=True)
|
|
assert child.forked_from_parent is True
|
|
entries = await manager.get_transcript("agent:main:direct:u1")
|
|
assert len(entries) == 2
|
|
assert entries[0].content == "parent msg"
|
|
assert entries[1].turn_usage == {
|
|
"model": "openai/gpt-test",
|
|
"input_tokens": 11,
|
|
"output_tokens": 5,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_before_message_copies_only_prefix(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "A marker", token_count=5)
|
|
await manager.append_message("agent:main:main", "assistant", "ack A", token_count=5)
|
|
before_entry = await manager.append_message(
|
|
"agent:main:main", "user", "B marker", token_count=5
|
|
)
|
|
await manager.append_message("agent:main:main", "assistant", "ack B", token_count=5)
|
|
await manager.append_message("agent:main:main", "user", "C marker must not leak", token_count=5)
|
|
|
|
child = await manager.branch(
|
|
"agent:main:main",
|
|
"agent:main:direct:edited",
|
|
fork_transcript=True,
|
|
fork_before_message_id=before_entry.message_id,
|
|
)
|
|
|
|
assert child.forked_from_parent is True
|
|
child_contents = [
|
|
entry.content for entry in await manager.get_transcript("agent:main:direct:edited")
|
|
]
|
|
assert child_contents == [
|
|
"A marker",
|
|
"ack A",
|
|
]
|
|
assert [entry.content for entry in await manager.get_transcript("agent:main:main")] == [
|
|
"A marker",
|
|
"ack A",
|
|
"B marker",
|
|
"ack B",
|
|
"C marker must not leak",
|
|
]
|
|
assert await manager.get_summaries("agent:main:direct:edited") == []
|
|
assert await manager.get_context_states("agent:main:direct:edited") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_before_message_missing_id_does_not_create_child(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "A marker")
|
|
|
|
with pytest.raises(KeyError, match="Transcript message not found"):
|
|
await manager.branch(
|
|
"agent:main:main",
|
|
"agent:main:direct:missing",
|
|
fork_transcript=True,
|
|
fork_before_message_id="missing-message-id",
|
|
)
|
|
|
|
assert await manager.get_session("agent:main:direct:missing") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_transcript_copies_compaction_summaries(manager):
|
|
parent = await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "kept tail", token_count=5)
|
|
await manager._storage.save_summary(
|
|
SessionSummary(
|
|
session_id=parent.session_id,
|
|
session_key="agent:main:main",
|
|
compaction_id="cmp_branch_1",
|
|
trigger_reason="preflight_compaction",
|
|
summary_text="older compacted context",
|
|
summary_payload={"user_goal": "continue the parent task"},
|
|
summary_format="structured_v1",
|
|
summary_source="llm",
|
|
coverage_status="pass",
|
|
missing_obligations=["none"],
|
|
critical_carry_forward=["remember file path"],
|
|
tokens_before=1000,
|
|
tokens_after=300,
|
|
removed_count=3,
|
|
kept_count=1,
|
|
chunk_count=2,
|
|
flush_receipt_status="safe",
|
|
covered_through_id=123,
|
|
)
|
|
)
|
|
await manager.save_context_state(
|
|
SessionContextState(
|
|
session_id=parent.session_id,
|
|
session_key="agent:main:main",
|
|
provider="portable",
|
|
state_kind="structured_summary_v1",
|
|
payload={"user_goal": "continue the parent task"},
|
|
covered_through_id=123,
|
|
portable=True,
|
|
cacheable=True,
|
|
)
|
|
)
|
|
|
|
child = await manager.branch("agent:main:main", "agent:main:direct:u1", fork_transcript=True)
|
|
|
|
assert child.forked_from_parent is True
|
|
child_summaries = await manager.get_summaries("agent:main:direct:u1")
|
|
assert len(child_summaries) == 1
|
|
assert child_summaries[0].summary_text == "older compacted context"
|
|
assert child_summaries[0].compaction_id == "cmp_branch_1"
|
|
assert child_summaries[0].trigger_reason == "preflight_compaction"
|
|
assert child_summaries[0].summary_payload == {"user_goal": "continue the parent task"}
|
|
assert child_summaries[0].summary_format == "structured_v1"
|
|
assert child_summaries[0].summary_source == "llm"
|
|
assert child_summaries[0].coverage_status == "pass"
|
|
assert child_summaries[0].missing_obligations == ["none"]
|
|
assert child_summaries[0].critical_carry_forward == ["remember file path"]
|
|
assert child_summaries[0].tokens_before == 1000
|
|
assert child_summaries[0].tokens_after == 300
|
|
assert child_summaries[0].removed_count == 3
|
|
assert child_summaries[0].kept_count == 1
|
|
assert child_summaries[0].chunk_count == 2
|
|
assert child_summaries[0].flush_receipt_status == "safe"
|
|
assert child_summaries[0].covered_through_id == 123
|
|
assert child_summaries[0].session_id == child.session_id
|
|
assert child_summaries[0].session_key == "agent:main:direct:u1"
|
|
child_states = await manager.get_context_states("agent:main:direct:u1")
|
|
assert len(child_states) == 1
|
|
assert child_states[0].session_id == child.session_id
|
|
assert child_states[0].session_key == "agent:main:direct:u1"
|
|
assert child_states[0].payload == {"user_goal": "continue the parent task"}
|
|
assert child_states[0].portable is True
|
|
assert child_states[0].cacheable is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_transcript_copies_compacted_archive(manager):
|
|
await manager.create("agent:main:main")
|
|
for index in range(4):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"short summary",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
compaction_id="cmp_branch_archive",
|
|
trigger_reason="agent_inline_overflow",
|
|
)
|
|
parent_canonical = [
|
|
entry.content for entry in await manager.get_canonical_transcript("agent:main:main")
|
|
]
|
|
|
|
child = await manager.branch("agent:main:main", "agent:main:direct:u1", fork_transcript=True)
|
|
|
|
assert child.parent_session_key == "agent:main:main"
|
|
assert [entry.content for entry in await manager.get_transcript("agent:main:direct:u1")] == [
|
|
"latest reply"
|
|
]
|
|
assert [
|
|
entry.content for entry in await manager.get_canonical_transcript("agent:main:direct:u1")
|
|
] == parent_canonical
|
|
child_summaries = await manager.get_summaries("agent:main:direct:u1")
|
|
assert child_summaries[0].compaction_id == "cmp_branch_archive"
|
|
assert child_summaries[0].trigger_reason == "agent_inline_overflow"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_state_roundtrip_and_invalidate(manager):
|
|
node = await manager.create("agent:main:main")
|
|
state = await manager.save_context_state(
|
|
SessionContextState(
|
|
session_id=node.session_id,
|
|
session_key="agent:main:main",
|
|
provider="portable",
|
|
model=None,
|
|
state_kind="structured_summary_v1",
|
|
payload={"current_status": "summary state"},
|
|
covered_through_id=42,
|
|
portable=True,
|
|
cacheable=True,
|
|
)
|
|
)
|
|
|
|
loaded = await manager.get_context_states("agent:main:main")
|
|
|
|
assert state.id is not None
|
|
assert len(loaded) == 1
|
|
assert loaded[0].payload == {"current_status": "summary state"}
|
|
assert loaded[0].portable is True
|
|
assert loaded[0].cacheable is True
|
|
assert loaded[0].valid is True
|
|
|
|
invalidated = await manager.invalidate_context_states(
|
|
"agent:main:main",
|
|
state_kind="structured_summary_v1",
|
|
reason="provider switched",
|
|
)
|
|
|
|
assert invalidated == 1
|
|
assert await manager.get_context_states("agent:main:main") == []
|
|
invalid = await manager.get_context_states("agent:main:main", valid_only=False)
|
|
assert len(invalid) == 1
|
|
assert invalid[0].valid is False
|
|
assert invalid[0].invalid_reason == "provider switched"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_storage_migrates_legacy_summary_metadata_columns(tmp_path):
|
|
db_path = tmp_path / "legacy.db"
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE session_summaries (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
session_id TEXT NOT NULL,
|
|
session_key TEXT NOT NULL,
|
|
compaction_index INTEGER NOT NULL DEFAULT 0,
|
|
summary_text TEXT NOT NULL,
|
|
covered_through_id INTEGER NOT NULL DEFAULT 0,
|
|
created_at INTEGER NOT NULL,
|
|
schema_version INTEGER NOT NULL DEFAULT 1
|
|
)
|
|
"""
|
|
)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO session_summaries (
|
|
session_id, session_key, compaction_index, summary_text,
|
|
covered_through_id, created_at, schema_version
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
("legacy-session", "agent:main:legacy", 0, "legacy summary", 7, 123, 1),
|
|
)
|
|
|
|
storage = SessionStorage(str(db_path))
|
|
await storage.connect()
|
|
try:
|
|
async with storage.conn.execute("PRAGMA table_info(session_summaries)") as cur:
|
|
columns = {row[1] for row in await cur.fetchall()}
|
|
assert {
|
|
"compaction_id",
|
|
"trigger_reason",
|
|
"summary_payload",
|
|
"summary_format",
|
|
"summary_source",
|
|
"coverage_status",
|
|
"missing_obligations",
|
|
"critical_carry_forward",
|
|
"tokens_before",
|
|
"tokens_after",
|
|
"removed_count",
|
|
"kept_count",
|
|
"chunk_count",
|
|
"flush_receipt_status",
|
|
}.issubset(columns)
|
|
|
|
summary = await storage.get_latest_summary("legacy-session")
|
|
assert summary is not None
|
|
assert summary.summary_text == "legacy summary"
|
|
assert summary.summary_payload is None
|
|
assert summary.summary_format == "text"
|
|
assert summary.summary_source == "unknown"
|
|
assert summary.coverage_status == "unknown"
|
|
assert summary.missing_obligations is None
|
|
assert summary.critical_carry_forward is None
|
|
assert summary.compaction_id is None
|
|
assert summary.trigger_reason is None
|
|
assert summary.tokens_before is None
|
|
assert summary.tokens_after is None
|
|
assert summary.removed_count == 0
|
|
assert summary.kept_count == 0
|
|
assert summary.chunk_count == 0
|
|
assert summary.flush_receipt_status == "unknown"
|
|
finally:
|
|
await storage.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_storage_migrates_session_context_state_table(tmp_path):
|
|
db_path = tmp_path / "legacy-context-state.db"
|
|
with sqlite3.connect(db_path) as conn:
|
|
conn.execute("CREATE TABLE sessions (session_key TEXT PRIMARY KEY)")
|
|
|
|
storage = SessionStorage(str(db_path))
|
|
await storage.connect()
|
|
try:
|
|
async with storage.conn.execute("PRAGMA table_info(session_context_states)") as cur:
|
|
columns = {row[1] for row in await cur.fetchall()}
|
|
assert {
|
|
"session_id",
|
|
"session_key",
|
|
"provider",
|
|
"model",
|
|
"state_kind",
|
|
"payload",
|
|
"covered_through_id",
|
|
"created_at",
|
|
"expires_at",
|
|
"portable",
|
|
"cacheable",
|
|
"valid",
|
|
"invalid_reason",
|
|
}.issubset(columns)
|
|
finally:
|
|
await storage.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_skipped_if_over_budget(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "msg", token_count=1000)
|
|
child = await manager.branch(
|
|
"agent:main:main", "agent:main:direct:u1", fork_transcript=True, max_fork_tokens=10
|
|
)
|
|
assert child.forked_from_parent is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_branch_fork_budget_counts_compaction_summaries(manager):
|
|
parent = await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "kept", token_count=1)
|
|
await manager._storage.save_summary(
|
|
SessionSummary(
|
|
session_id=parent.session_id,
|
|
session_key="agent:main:main",
|
|
summary_text="x" * 400,
|
|
)
|
|
)
|
|
|
|
child = await manager.branch(
|
|
"agent:main:main",
|
|
"agent:main:direct:u1",
|
|
fork_transcript=True,
|
|
max_fork_tokens=10,
|
|
)
|
|
|
|
assert child.forked_from_parent is False
|
|
assert await manager.get_transcript("agent:main:direct:u1") == []
|
|
assert await manager.get_summaries("agent:main:direct:u1") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_no_op_small_context(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "hi", token_count=5)
|
|
summary = await manager.compact("agent:main:main", context_window_tokens=100_000)
|
|
assert summary == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_reduces_transcript(manager):
|
|
await manager.create("agent:main:main")
|
|
# Add many large messages
|
|
for i in range(20):
|
|
await manager.append_message("agent:main:main", "user", "x" * 500, token_count=200)
|
|
summary = await manager.compact("agent:main:main", context_window_tokens=1000)
|
|
assert summary != ""
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.compaction_count == 1
|
|
transcript = await manager.get_transcript("agent:main:main")
|
|
assert transcript
|
|
assert all(entry.role != "system" for entry in transcript)
|
|
summaries = await manager._storage.get_all_summaries(node.session_id)
|
|
assert len(summaries) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_returns_source_and_persists(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
original_contents = [entry.content for entry in await manager.get_transcript("agent:main:main")]
|
|
|
|
result = await manager.compact_with_result("agent:main:main", context_window_tokens=1000)
|
|
|
|
assert result.summary
|
|
assert result.summary_source == "fallback"
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.compaction_count == 1
|
|
transcript = await manager.get_transcript("agent:main:main")
|
|
assert transcript
|
|
assert all(entry.role != "system" for entry in transcript)
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert [summary.summary_text for summary in summaries] == [result.summary]
|
|
assert summaries[0].summary_format == "structured_v1"
|
|
assert summaries[0].summary_source == result.summary_source
|
|
assert summaries[0].coverage_status in {"unknown", "pass", "pass_with_backfill"}
|
|
assert summaries[0].summary_payload is not None
|
|
assert summaries[0].summary_payload["schema_version"] == 1
|
|
assert summaries[0].compaction_id
|
|
assert summaries[0].tokens_before == result.tokens_before
|
|
assert summaries[0].tokens_after == result.tokens_after
|
|
assert summaries[0].removed_count == result.removed_count
|
|
assert summaries[0].kept_count == len(result.kept_entries)
|
|
assert summaries[0].chunk_count == result.chunks_processed
|
|
canonical_contents = [
|
|
entry.content for entry in await manager.get_canonical_transcript("agent:main:main")
|
|
]
|
|
assert canonical_contents == original_contents
|
|
assert [entry.content for entry in transcript] == original_contents[-len(transcript) :]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_skips_rewrite_when_transcript_changes(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
original_contents = [entry.content for entry in await manager.get_transcript("agent:main:main")]
|
|
context_entries = 0
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def mutation_context():
|
|
nonlocal context_entries
|
|
context_entries += 1
|
|
if context_entries == 2:
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
"late queued followup",
|
|
token_count=3,
|
|
)
|
|
yield
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
mutation_context=mutation_context,
|
|
)
|
|
|
|
assert result.summary == ""
|
|
assert result.summary_source == "skipped"
|
|
assert result.skip_reason == "stale_preimage"
|
|
assert result.removed_count == 0
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.compaction_count == 0
|
|
assert await manager.get_summaries("agent:main:main") == []
|
|
transcript = await manager.get_transcript("agent:main:main")
|
|
assert [entry.content for entry in transcript] == original_contents + ["late queued followup"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_marks_unsafe_receipt_as_degraded_forensic(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"unsafe msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
original_contents = [entry.content for entry in await manager.get_transcript("agent:main:main")]
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
flush_receipt_status="unsafe",
|
|
)
|
|
|
|
assert result.removed_count > 0
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert summaries[0].flush_receipt_status == "degraded_forensic"
|
|
canonical_contents = [
|
|
entry.content for entry in await manager.get_canonical_transcript("agent:main:main")
|
|
]
|
|
assert canonical_contents == original_contents
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_degraded_compaction_preimage_can_be_listed_for_repair(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"repair msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
|
|
await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
flush_receipt_status="degraded_forensic",
|
|
)
|
|
|
|
pending = await manager.list_degraded_compactions(agent_id="main")
|
|
assert len(pending) == 1
|
|
assert pending[0].flush_receipt_status == "degraded_forensic"
|
|
preimage = await manager.get_compaction_preimage(pending[0])
|
|
assert preimage
|
|
assert preimage[0].content.startswith("repair msg 0")
|
|
await manager.mark_compaction_repair_status(pending[0], "repaired")
|
|
assert await manager.list_degraded_compactions(agent_id="main") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compaction_flush_status_can_be_backfilled_by_compaction_id(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"background flush msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
|
|
await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
compaction_id="cmp-bg-flush",
|
|
flush_receipt_status="degraded_forensic",
|
|
)
|
|
|
|
updated = await manager.mark_compaction_flush_receipt_status(
|
|
"agent:main:main",
|
|
"cmp-bg-flush",
|
|
"safe",
|
|
)
|
|
|
|
assert updated == 1
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert summaries[0].flush_receipt_status == "safe"
|
|
assert await manager.list_degraded_compactions(agent_id="main") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_noop_memory_flush_compaction_status_does_not_enter_repair_queue(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"noop msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
flush_receipt_status="noop_no_memory",
|
|
)
|
|
|
|
assert result.removed_count > 0
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert summaries[0].flush_receipt_status == "noop_no_memory"
|
|
assert await manager.list_degraded_compactions(agent_id="main") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_archive_only_memory_flush_compaction_status_does_not_enter_repair_queue(manager):
|
|
await manager.create("agent:main:main")
|
|
for i in range(20):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
f"archive msg {i} " + ("x" * 500),
|
|
token_count=200,
|
|
)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=1000,
|
|
flush_receipt_status="archive_only",
|
|
)
|
|
|
|
assert result.removed_count > 0
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert summaries[0].flush_receipt_status == "archive_only"
|
|
assert await manager.list_degraded_compactions(agent_id="main") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_reports_and_backfills_missing_obligations(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
(
|
|
"Goal: finish continuity work.\n"
|
|
"Constraint: do not enable coverage blocking by default.\n"
|
|
"Keep src/opensquilla/session/models.py and docs/Long Task Report.md."
|
|
),
|
|
token_count=250,
|
|
)
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
"Next I will run uv run pytest tests/test_session/test_manager.py.",
|
|
tool_calls=[{"id": "call_exec_1", "type": "function"}],
|
|
token_count=250,
|
|
)
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"tool",
|
|
(
|
|
"Command failed: uv run pytest tests/test_session/test_manager.py\n"
|
|
"Exit code 1\n"
|
|
"Error: missing summary_payload column"
|
|
),
|
|
tool_call_id="call_exec_1",
|
|
token_count=250,
|
|
)
|
|
for i in range(8):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
f"recent tail {i}",
|
|
token_count=20,
|
|
)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=500,
|
|
config=CompactionConfig(safety_margin=1.0),
|
|
)
|
|
|
|
assert result.removed_count > 0
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
summary = summaries[0]
|
|
assert summary.summary_text == result.summary
|
|
assert summary.summary_format == "structured_v1"
|
|
assert summary.summary_payload is not None
|
|
assert summary.coverage_status == "pass_with_backfill"
|
|
assert summary.missing_obligations
|
|
assert summary.critical_carry_forward
|
|
assert "src/opensquilla/session/models.py" in str(summary.summary_payload)
|
|
assert any("call_exec_1" in item for item in summary.critical_carry_forward)
|
|
assert await manager.get_transcript("agent:main:main")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_strict_coverage_blocks_destructive_rewrite(manager):
|
|
node = await manager.create("agent:main:main")
|
|
late_critical_path = "src/opensquilla/session/critical_continuity.py"
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
"Goal: preserve strict continuity. " + ("padding " * 40) + late_critical_path,
|
|
token_count=650,
|
|
)
|
|
for index in range(4):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
f"recent tail {index}",
|
|
token_count=20,
|
|
)
|
|
original_transcript = await manager.get_transcript("agent:main:main")
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=300,
|
|
config=CompactionConfig(safety_margin=1.0, coverage_blocking=True),
|
|
)
|
|
|
|
assert result.removed_count == 0
|
|
assert result.skip_reason == "coverage_blocked"
|
|
assert result.coverage_status == "fail_blocked"
|
|
assert any(late_critical_path in item for item in result.missing_obligations or [])
|
|
assert await manager.get_transcript("agent:main:main") == original_transcript
|
|
assert await manager.get_summaries("agent:main:main") == []
|
|
assert await manager.get_context_states("agent:main:main") == []
|
|
current_node = await manager._storage.get_session("agent:main:main")
|
|
assert current_node is not None
|
|
assert current_node.session_id == node.session_id
|
|
assert current_node.compaction_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_writes_portable_context_state(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"user",
|
|
"Goal: keep portable state. File src/opensquilla/session/models.py.",
|
|
token_count=250,
|
|
)
|
|
for i in range(8):
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
f"recent tail {i}",
|
|
token_count=20,
|
|
)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=300,
|
|
config=CompactionConfig(safety_margin=1.0),
|
|
)
|
|
|
|
states = await manager.get_context_states("agent:main:main")
|
|
|
|
assert result.removed_count > 0
|
|
assert len(states) == 1
|
|
state = states[0]
|
|
assert state.provider == "portable"
|
|
assert state.model is None
|
|
assert state.state_kind == "structured_summary_v1"
|
|
summaries = await manager.get_summaries("agent:main:main")
|
|
assert summaries[0].compaction_id
|
|
payload_without_correlation = dict(state.payload)
|
|
assert payload_without_correlation.pop("compaction_id") == summaries[0].compaction_id
|
|
assert payload_without_correlation == result.summary_payload
|
|
assert state.covered_through_id > 0
|
|
assert state.portable is True
|
|
assert state.cacheable is True
|
|
assert state.valid is True
|
|
assert await manager.get_transcript("agent:main:main")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_with_result_preserves_tool_metadata_for_boundary_cut(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.append_message("agent:main:main", "user", "old context", token_count=300)
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
"calling tool",
|
|
tool_calls=[{"id": "call_1", "type": "function"}],
|
|
token_count=4,
|
|
)
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"tool",
|
|
"tool result",
|
|
tool_call_id="call_1",
|
|
token_count=4,
|
|
)
|
|
await manager.append_message("agent:main:main", "user", "next question", token_count=3)
|
|
await manager.append_message("agent:main:main", "assistant", "answer", token_count=3)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=100,
|
|
config=CompactionConfig(safety_margin=1.0),
|
|
)
|
|
|
|
assert result.removed_count == 1
|
|
assert result.kept_entries[0]["role"] == "assistant"
|
|
assert result.kept_entries[0]["tool_calls"] == [{"id": "call_1", "type": "function"}]
|
|
transcript = await manager.get_transcript("agent:main:main")
|
|
assert transcript[0].role == "assistant"
|
|
assert transcript[0].tool_calls == [{"id": "call_1", "type": "function"}]
|
|
assert transcript[1].role == "tool"
|
|
assert transcript[1].tool_call_id == "call_1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_counts_tool_calls_when_token_count_is_underreported(manager):
|
|
await manager.create("agent:main:main")
|
|
large_content = "x" * 5000
|
|
await manager.append_message(
|
|
"agent:main:main",
|
|
"assistant",
|
|
"small visible answer",
|
|
tool_calls=[
|
|
{
|
|
"type": "tool_use",
|
|
"tool_use_id": "write-stale",
|
|
"name": "write_file",
|
|
"input": {"path": "index.html", "content": large_content},
|
|
}
|
|
],
|
|
token_count=1,
|
|
)
|
|
|
|
result = await manager.compact_with_result(
|
|
"agent:main:main",
|
|
context_window_tokens=50,
|
|
config=CompactionConfig(safety_margin=1.0),
|
|
)
|
|
|
|
assert result.removed_count == 1
|
|
assert result.summary
|
|
assert large_content not in result.summary
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.compaction_count == 1
|
|
|
|
|
|
def _fail_next_transcript_insert(monkeypatch: pytest.MonkeyPatch, storage: SessionStorage) -> None:
|
|
original_execute = storage.conn.execute
|
|
failed = False
|
|
|
|
def execute(sql: str, params: Any = ()):
|
|
nonlocal failed
|
|
if (
|
|
not failed
|
|
and isinstance(sql, str)
|
|
and sql.lstrip().upper().startswith("INSERT INTO TRANSCRIPT_ENTRIES")
|
|
):
|
|
failed = True
|
|
raise RuntimeError("rewrite insert failed")
|
|
return original_execute(sql, params)
|
|
|
|
monkeypatch.setattr(storage.conn, "execute", execute)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compact_rewrite_failure_keeps_session_state_atomic(
|
|
manager,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
node = await manager.create("agent:main:main")
|
|
for index in range(20):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index} " + ("x" * 500))
|
|
original_transcript = await manager.get_transcript("agent:main:main")
|
|
original_canonical_transcript = await manager.get_canonical_transcript("agent:main:main")
|
|
original_summaries = await manager.get_summaries("agent:main:main")
|
|
original_node = await manager._storage.get_session("agent:main:main")
|
|
|
|
_fail_next_transcript_insert(monkeypatch, manager._storage)
|
|
|
|
with pytest.raises(RuntimeError, match="rewrite insert failed"):
|
|
await manager.compact("agent:main:main", context_window_tokens=1000)
|
|
|
|
assert await manager.get_transcript("agent:main:main") == original_transcript
|
|
assert (
|
|
await manager.get_canonical_transcript("agent:main:main") == original_canonical_transcript
|
|
)
|
|
assert await manager.get_summaries("agent:main:main") == original_summaries
|
|
assert await manager.get_context_states("agent:main:main") == []
|
|
current_node = await manager._storage.get_session("agent:main:main")
|
|
assert current_node is not None
|
|
assert original_node is not None
|
|
assert current_node.session_id == node.session_id
|
|
assert current_node.compaction_count == original_node.compaction_count
|
|
assert current_node.updated_at == original_node.updated_at
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_compaction_result_rewrite_failure_keeps_session_state_atomic(
|
|
manager,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
):
|
|
node = await manager.create("agent:main:main")
|
|
for index in range(4):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
original_transcript = await manager.get_transcript("agent:main:main")
|
|
original_canonical_transcript = await manager.get_canonical_transcript("agent:main:main")
|
|
original_summaries = await manager.get_summaries("agent:main:main")
|
|
original_node = await manager._storage.get_session("agent:main:main")
|
|
|
|
_fail_next_transcript_insert(monkeypatch, manager._storage)
|
|
|
|
with pytest.raises(RuntimeError, match="rewrite insert failed"):
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"short summary",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
)
|
|
|
|
assert await manager.get_transcript("agent:main:main") == original_transcript
|
|
assert (
|
|
await manager.get_canonical_transcript("agent:main:main") == original_canonical_transcript
|
|
)
|
|
assert await manager.get_summaries("agent:main:main") == original_summaries
|
|
assert await manager.get_context_states("agent:main:main") == []
|
|
current_node = await manager._storage.get_session("agent:main:main")
|
|
assert current_node is not None
|
|
assert original_node is not None
|
|
assert current_node.session_id == node.session_id
|
|
assert current_node.compaction_count == original_node.compaction_count
|
|
assert current_node.updated_at == original_node.updated_at
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_compaction_result_stores_summary_out_of_band(manager):
|
|
node = await manager.create("agent:main:main")
|
|
for index in range(4):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"short summary",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
compaction_id="cmp_inline_1",
|
|
trigger_reason="agent_inline_overflow",
|
|
)
|
|
|
|
transcript = await manager.get_transcript("agent:main:main")
|
|
assert all(entry.role != "system" for entry in transcript)
|
|
assert transcript[-1].content == "latest reply"
|
|
canonical = await manager.get_canonical_transcript("agent:main:main")
|
|
assert [entry.content for entry in canonical] == [
|
|
"msg 0",
|
|
"msg 1",
|
|
"msg 2",
|
|
"latest reply",
|
|
]
|
|
async with manager._storage.conn.execute(
|
|
"SELECT compaction_id, compaction_index FROM compacted_transcript_entries "
|
|
"WHERE session_id = ? ORDER BY original_entry_id ASC",
|
|
(node.session_id,),
|
|
) as cur:
|
|
archived_rows = await cur.fetchall()
|
|
assert [(row[0], row[1]) for row in archived_rows] == [
|
|
("cmp_inline_1", 0),
|
|
("cmp_inline_1", 0),
|
|
("cmp_inline_1", 0),
|
|
]
|
|
summaries = await manager._storage.get_all_summaries(node.session_id)
|
|
assert len(summaries) == 1
|
|
assert summaries[0].summary_text == "short summary"
|
|
assert summaries[0].compaction_id == "cmp_inline_1"
|
|
assert summaries[0].trigger_reason == "agent_inline_overflow"
|
|
assert summaries[0].removed_count == 3
|
|
assert summaries[0].kept_count == 1
|
|
states = await manager.get_context_states("agent:main:main")
|
|
assert len(states) == 1
|
|
assert states[0].state_kind == "structured_summary_v1"
|
|
assert states[0].payload is not None
|
|
assert states[0].payload["compaction_id"] == "cmp_inline_1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_session_removes_compacted_transcript_archive(manager):
|
|
node = await manager.create("agent:main:main")
|
|
for index in range(4):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"short summary",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
)
|
|
assert len(await manager.get_canonical_transcript("agent:main:main")) == 4
|
|
|
|
await manager._storage.delete_session("agent:main:main")
|
|
|
|
async with manager._storage.conn.execute(
|
|
"SELECT COUNT(*) FROM compacted_transcript_entries WHERE session_id = ?",
|
|
(node.session_id,),
|
|
) as cur:
|
|
archived_count = (await cur.fetchone())[0]
|
|
assert archived_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_compaction_result_without_summary_does_not_rewrite_transcript(manager):
|
|
node = await manager.create("agent:main:main")
|
|
for index in range(5):
|
|
await manager.append_message("agent:main:main", "user", f"msg {index}", token_count=5)
|
|
|
|
original_transcript = await manager.get_transcript("agent:main:main")
|
|
original_node = await manager._storage.get_session("agent:main:main")
|
|
|
|
await manager.persist_compaction_result(
|
|
"agent:main:main",
|
|
"",
|
|
[{"role": "assistant", "content": "latest reply"}],
|
|
)
|
|
|
|
assert await manager.get_transcript("agent:main:main") == original_transcript
|
|
assert await manager.get_summaries("agent:main:main") == []
|
|
current_node = await manager._storage.get_session("agent:main:main")
|
|
assert current_node is not None
|
|
assert original_node is not None
|
|
assert current_node.session_id == node.session_id
|
|
assert current_node.compaction_count == original_node.compaction_count
|
|
assert current_node.updated_at == original_node.updated_at
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prune_stale(manager):
|
|
node = await manager.create("agent:main:main")
|
|
# force old timestamp
|
|
node.updated_at = 1
|
|
await manager._storage.upsert_session(node)
|
|
pruned = await manager.prune_stale(max_age_ms=1000)
|
|
assert pruned == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cap_entries(manager):
|
|
for i in range(10):
|
|
await manager.create(f"agent:main:direct:u{i}")
|
|
deleted = await manager.cap_entries(max_entries=5)
|
|
assert deleted == 5
|
|
remaining = await manager._storage.count_sessions()
|
|
assert remaining == 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cap_entries_cleans_related_transcript_and_summaries(manager):
|
|
session_ids: dict[str, str] = {}
|
|
for i in range(3):
|
|
key = f"agent:main:direct:u{i}"
|
|
node = await manager.create(key)
|
|
session_ids[key] = node.session_id
|
|
await manager._storage.append_transcript_entry(
|
|
TranscriptEntry(
|
|
session_id=node.session_id,
|
|
session_key=key,
|
|
role="user",
|
|
content="stale",
|
|
)
|
|
)
|
|
await manager._storage.save_summary(
|
|
SessionSummary(
|
|
session_id=node.session_id,
|
|
session_key=key,
|
|
summary_text="summary",
|
|
)
|
|
)
|
|
await manager.save_context_state(
|
|
SessionContextState(
|
|
session_id=node.session_id,
|
|
session_key=key,
|
|
provider="portable",
|
|
state_kind="structured_summary_v1",
|
|
payload={"summary": "state"},
|
|
portable=True,
|
|
cacheable=True,
|
|
)
|
|
)
|
|
deleted = await manager.cap_entries(max_entries=1)
|
|
assert deleted == 2
|
|
remaining = {session.session_key for session in await manager._storage.list_sessions(limit=10)}
|
|
removed = set(session_ids) - remaining
|
|
assert len(removed) == 2
|
|
for key in removed:
|
|
session_id = session_ids[key]
|
|
assert await manager._storage.count_transcript_entries(session_id) == 0
|
|
assert await manager._storage.get_all_summaries(session_id) == []
|
|
assert await manager.get_context_states(key, valid_only=False) == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_archive(manager):
|
|
await manager.create("agent:main:main")
|
|
await manager.archive("agent:main:main")
|
|
node = await manager._storage.get_session("agent:main:main")
|
|
assert node.status == SessionStatus.DONE
|