dhruv-punia-bits/memory-compaction-openenv

1from __future__ import annotations2 3from fastapi.testclient import TestClient4 5from server.app import app6from server.environment import MemoryCompactionEnvironment7from server.models import ResetRequest, StepAction8 9client = TestClient(app)10 11 12def test_reset_is_deterministic_for_seed() -> None:13    first = client.post("/reset", json={"difficulty": "easy", "seed": 123}).json()14    second = client.post("/reset", json={"difficulty": "easy", "seed": 123}).json()15 16    assert first["observation"]["current_turn"]["text"] == second["observation"]["current_turn"]["text"]17    assert first["observation"]["token_budget_remaining"] == second["observation"]["token_budget_remaining"]18 19 20def test_reset_accepts_empty_body() -> None:21    response = client.post("/reset")22    assert response.status_code == 20023    body = response.json()24    assert body["observation"]["difficulty"] == "easy"25 26 27def test_runtime_metadata_and_schema_endpoints_exist() -> None:28    health = client.get("/health")29    metadata = client.get("/metadata")30    schema = client.get("/schema")31    mcp = client.post("/mcp", json={})32 33    assert health.status_code == 20034    assert health.json()["status"] == "healthy"35    assert metadata.status_code == 20036    assert metadata.json()["name"] == "memory-compaction-openenv"37    assert schema.status_code == 20038    assert "action" in schema.json()39    assert "observation" in schema.json()40    assert "state" in schema.json()41    assert mcp.status_code == 20042    assert mcp.json()["jsonrpc"] == "2.0"43 44 45def test_step_accepts_valid_action() -> None:46    client.post("/reset", json={"difficulty": "medium", "seed": 202})47    response = client.post(48        "/step",49        json={50            "operation": "replace_summary",51            "summary_text": "User trip, project, and latest updates.",52            "memory_items": [],53        },54    )55    assert response.status_code == 20056    body = response.json()57    assert 0.0 <= body["reward"] <= 1.058    assert "reward_breakdown" in body59 60 61def test_invalid_action_is_rejected() -> None:62    client.post("/reset", json={"difficulty": "easy", "seed": 101})63    response = client.post("/step", json={"operation": "invalid"})64    assert response.status_code == 42265 66 67def test_state_matches_episode_progress() -> None:68    client.post("/reset", json={"difficulty": "hard", "seed": 303})69    client.post("/step", json={"operation": "noop", "memory_items": [], "summary_text": ""})70    state = client.get("/state").json()["state"]71 72    assert state["step_count"] == 173    assert state["current_index"] == 174    assert state["difficulty"] == "hard"75 76 77def test_environment_supersedes_conflicting_memory_updates() -> None:78    env = MemoryCompactionEnvironment()79    env.reset(ResetRequest(difficulty="medium", seed=202))80    env.step(81        StepAction(82            operation="append_memory",83            memory_items=[84                {85                    "memory_id": "project",86                    "type": "plan",87                    "subject": "Sonia",88                    "predicate": "building",89                    "object": "benchmark launch",90                    "confidence": 0.9,91                    "source_turn_ids": [4],92                    "source_text": "I am building a benchmark launch with Arun.",93                    "status": "active",94                },95                {96                    "memory_id": "deadline",97                    "type": "task",98                    "subject": "benchmark launch",99                    "predicate": "deadline",100                    "object": "May 18",101                    "confidence": 0.9,102                    "source_turn_ids": [6],103                    "source_text": "The first deadline is May 18, and I track work in Linear.",104                    "status": "active",105                },106            ],107            summary_text="Initial plan and deadline.",108        )109    )110    env.step(111        StepAction(112            operation="update_memory",113            memory_items=[114                {115                    "memory_id": "deadline_v2",116                    "type": "correction",117                    "subject": "benchmark launch",118                    "predicate": "deadline",119                    "object": "May 21",120                    "confidence": 0.98,121                    "source_turn_ids": [7],122                    "source_text": "Correction: the deadline moved to May 21.",123                    "status": "active",124                    "updated_from_memory_id": "deadline",125                }126            ],127            summary_text="Deadline updated.",128        )129    )130 131    state = env.state()132    deadline = next(item for item in state.durable_memory if item.memory_id == "deadline")133    deadline_v2 = next(item for item in state.durable_memory if item.memory_id == "deadline_v2")134    assert deadline.status == "superseded"135    assert deadline_v2.status == "active"136 137 138def test_ambiguous_memory_can_be_stored_as_uncertain_without_penalty() -> None:139    env = MemoryCompactionEnvironment()140    env.reset(ResetRequest(difficulty="medium", seed=202))141    for _ in range(9):142        env.step(StepAction(operation="noop", memory_items=[], summary_text=""))143 144    response = env.step(145        StepAction(146            operation="append_memory",147            memory_items=[148                {149                    "memory_id": "presentation_plan_uncertain",150                    "type": "plan",151                    "subject": "benchmark launch",152                    "predicate": "may_present",153                    "object": "present this work next week",154                    "confidence": 0.4,155                    "source_turn_ids": [10],156                    "source_text": "I might present this work next week, but that is not confirmed yet.",157                    "status": "uncertain",158                    "requires_confirmation": True,159                }160            ],161            summary_text="Possible future presentation, not confirmed.",162        )163    )164 165    assert response.reward > 0.5166 167 168def test_terminal_task_scores_are_strictly_inside_unit_interval() -> None:169    env = MemoryCompactionEnvironment()170    env.reset(ResetRequest(difficulty="easy", seed=1101))171 172    while True:173        response = env.step(StepAction(operation="noop", memory_items=[], summary_text=""))174        if response.done:175            assert 0.0 < response.reward < 1.0176            assert 0.0 < response.reward_breakdown.total_score < 1.0177            break178