171 lines
8.5 KiB
Python
171 lines
8.5 KiB
Python
import json
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
from sessions import SessionStore
|
|
from resource_limits import ContextBudget, CostTracker
|
|
|
|
|
|
class SessionTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.store = SessionStore(Path(self.temp.name))
|
|
self.messages = [
|
|
{"role": "system", "content": "Follow security policy."},
|
|
{"role": "user", "content": "Fix cache. " * 300},
|
|
{"role": "assistant", "content": "Changed cache.py and tests. " * 300},
|
|
{"role": "user", "content": "What remains?"},
|
|
]
|
|
self.store.create_session("one", self.messages)
|
|
self.client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(
|
|
create=lambda **kwargs: SimpleNamespace(usage=None, choices=[SimpleNamespace(
|
|
message=SimpleNamespace(content="Fixed cache.py; tests passed; document next."))]))))
|
|
|
|
def compact(self):
|
|
return self.store.compact("one", self.client, ContextBudget(), CostTracker(), force=True)
|
|
|
|
def test_lossless_compaction_and_followup(self):
|
|
replacement = self.compact()
|
|
record = self.store.load_session("one")
|
|
self.assertEqual(record["transcript"], self.messages)
|
|
self.assertEqual(replacement[-1], self.messages[-1])
|
|
self.assertEqual(replacement[1]["role"], "user")
|
|
followup = {"role": "assistant", "content": "Documentation."}
|
|
self.store.save_messages("one", replacement + [followup])
|
|
self.assertEqual(self.store.load_session("one")["transcript"], self.messages + [followup])
|
|
self.assertEqual(len(record["compactions"]), 1)
|
|
|
|
def test_compaction_keeps_complete_tool_batch(self):
|
|
messages = self.messages[:1] + [self.messages[1], self.messages[2],
|
|
{"role": "user", "content": "Inspect"},
|
|
{"role": "assistant", "tool_calls": [{"id": "call", "function": {"name": "read_file"}}]},
|
|
{"role": "tool", "tool_call_id": "call", "content": "file"}]
|
|
self.store.create_session("tools", messages)
|
|
result = self.store.compact("tools", self.client, ContextBudget(), CostTracker(), force=True)
|
|
self.assertEqual(result[-3:], messages[-3:])
|
|
|
|
def test_failed_compaction_does_not_change_session(self):
|
|
before = self.store.export("one")
|
|
self.client.chat.completions.create = lambda **kwargs: (_ for _ in ()).throw(RuntimeError("offline"))
|
|
with self.assertRaises(RuntimeError):
|
|
self.compact()
|
|
self.assertEqual(self.store.export("one"), before)
|
|
|
|
def test_empty_or_oversized_summary_retains_original_context(self):
|
|
before = self.store.export("one")
|
|
for summary in ("", "x" * 50000):
|
|
self.client.chat.completions.create = lambda **kwargs: SimpleNamespace(
|
|
usage=None, choices=[SimpleNamespace(message=SimpleNamespace(content=summary))])
|
|
with self.assertRaises(ValueError):
|
|
self.compact()
|
|
self.assertEqual(self.store.export("one"), before)
|
|
|
|
def test_automatic_compaction_threshold_and_cost_guard(self):
|
|
self.assertIsNone(self.store.compact("one", self.client, ContextBudget(), CostTracker()))
|
|
self.client.chat.completions.create = Mock(side_effect=AssertionError("Must not call model"))
|
|
with self.assertRaises(ValueError):
|
|
self.store.compact("one", self.client, ContextBudget(), CostTracker(max_cost_usd=0), force=True)
|
|
self.client.chat.completions.create.assert_not_called()
|
|
|
|
def test_fork_archive_rename_export_delete(self):
|
|
self.compact()
|
|
child = self.store.fork("one")
|
|
self.store.update(child, title="Branch", archived=True)
|
|
self.assertEqual(len(self.store.list_sessions()), 1)
|
|
self.assertEqual(len(self.store.list_sessions(include_archived=True)), 2)
|
|
exported = json.loads(self.store.export(child))
|
|
self.assertEqual(exported["parent_id"], "one")
|
|
self.assertEqual(exported["transcript"], self.messages)
|
|
self.store.delete(child)
|
|
with self.assertRaises(ValueError):
|
|
self.store.load_session(child)
|
|
|
|
def test_crash_recovery_does_not_replay_tools(self):
|
|
call = {"role": "assistant", "tool_calls": [
|
|
{"id": "a", "function": {"name": "write_file"}},
|
|
{"id": "b", "function": {"name": "run_bash"}}]}
|
|
self.store.save_messages("one", self.messages + [call,
|
|
{"role": "tool", "tool_call_id": "a", "content": "saved"},
|
|
{"role": "user", "content": "Abort"}])
|
|
recovered = self.store.recover("one")
|
|
self.assertEqual(recovered[-2]["tool_call_id"], "b")
|
|
self.assertIn("unknown", recovered[-2]["content"])
|
|
self.assertEqual(self.store.recover("one"), recovered)
|
|
|
|
def test_legacy_session_migration(self):
|
|
record = self.store.load_session("one")
|
|
for key in ("transcript", "state", "title", "archived", "compactions"):
|
|
record.pop(key, None)
|
|
self.store._session_path("one").write_text(json.dumps(record))
|
|
self.assertEqual(self.store.load_session("one")["transcript"], self.messages)
|
|
|
|
def test_concurrent_process_cannot_mutate_session(self):
|
|
script = "from pathlib import Path; from sessions import SessionStore; " \
|
|
f"SessionStore(Path({self.temp.name!r})).update('one', title='bad')"
|
|
with self.store.lease("one"):
|
|
result = subprocess.run([sys.executable, "-c", script],
|
|
cwd=Path(__file__).parent, capture_output=True, text=True)
|
|
self.assertNotEqual(result.returncode, 0)
|
|
self.assertIn("in use", result.stderr)
|
|
self.store.update("one", title="available")
|
|
|
|
def test_ids_and_overwrites_are_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
self.store.create_session("one", [])
|
|
with self.assertRaises(ValueError):
|
|
self.store.load_session("../outside")
|
|
|
|
def test_tool_arguments_count_toward_context_budget(self):
|
|
message = {"role": "assistant", "tool_calls": [{"arguments": "x" * 4000}]}
|
|
self.assertGreaterEqual(ContextBudget().estimate_total([message]), 1000)
|
|
|
|
def test_agent_resumes_planning_state_and_checkpoints_reply(self):
|
|
from agent import agent_loop, PermissionMode
|
|
from openai.types.chat import ChatCompletionMessage
|
|
from tools.todo import todo_store
|
|
from tools.scratchpad import scratchpad
|
|
self.store.update("one", state={"scratchpad": "Document cache",
|
|
"todos": [{"id": "docs", "content": "Write docs", "status": "pending", "retries": 2}],
|
|
"calls": 3, "tool_calls": 4})
|
|
reply = ChatCompletionMessage(role="assistant", content="Document next.")
|
|
client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(
|
|
create=lambda **kwargs: SimpleNamespace(usage=None, choices=[SimpleNamespace(
|
|
message=reply, finish_reason="stop")]))))
|
|
with patch("builtins.input", side_effect=["Continue", "\\info", "\\exit"]), patch("builtins.print"):
|
|
agent_loop(client, PermissionMode.DEFAULT, Path(self.temp.name), {}, Mock(),
|
|
tool_schemas=[], memory_store=self.store, session_id="one",
|
|
resumed_messages=self.messages)
|
|
record = self.store.load_session("one")
|
|
self.assertEqual(record["state"]["calls"], 4)
|
|
self.assertEqual(record["state"]["tool_calls"], 4)
|
|
self.assertEqual(record["state"]["todos"][0]["retries"], 2)
|
|
self.assertEqual(record["transcript"][-1]["content"], "Document next.")
|
|
self.assertEqual(scratchpad.read(), "Document cache")
|
|
todo_store._items = []
|
|
scratchpad.write("")
|
|
|
|
def test_offline_cli_management(self):
|
|
agent_path = str(Path(__file__).parent / "agent.py")
|
|
|
|
def run(*args):
|
|
return subprocess.run([sys.executable, agent_path, "--memory-dir", self.temp.name, *args],
|
|
capture_output=True, text=True)
|
|
|
|
result = run("--rename", "one", "Cache task")
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
exported = run("--export", "one")
|
|
self.assertEqual(json.loads(exported.stdout)["title"], "Cache task")
|
|
self.assertNotEqual(run("--delete", "one").returncode, 0)
|
|
self.assertEqual(run("--delete", "one", "--yes").returncode, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|