Keep plan file after session close and auto-clean test plan dirs

AgentSession.close no longer removes the per-session PLAN.md so the
plan stays available for later reference. Tests create real plan files
but rarely call close(), leaking python-agent-plans-* dirs into /tmp;
tests/plan_cleanup.py now tracks them and removes them after each test
(with an atexit safety net).

Author: beacoder

python_agent_harness/agent_session.py

@@ -539,7 +539,6 @@ def _conversation_text(self, messages: list) -> str:
     def close(self) -> None:
         self.cancel()
         self.alive = False
-        self.plan_mode.cleanup_plan_file()
         cleanup_spooled_files()
         # MCP server connections + event-loop thread (no-op when no MCP
         # servers are configured or none connected)

tests/plan_cleanup.py

@@ -0,0 +1,54 @@
+"""Automatic cleanup of /tmp plan dirs created by the test suite.
+
+``PlanMode.ensure_plan_file()`` puts PLAN.md in a fresh
+``python-agent-plans-*`` dir under the temp dir (usually /tmp).
+Production ``AgentSession.close()`` intentionally keeps the plan file
+(one unique file per session), so tests must clean up their own:
+this module is imported (for its side effects) by the test files that
+can create plan files — it records every plan file created by
+``ensure_plan_file`` and removes them after each test (via a
+``unittest.TestCase.run`` hook) with an atexit safety net.
+"""
+
+import atexit
+import os
+import unittest
+
+from python_agent_harness.planmode import PlanMode
+
+_created: list[str] = []
+
+_orig_ensure = PlanMode.ensure_plan_file
+
+
+def _tracked_ensure(self) -> str:
+    path = _orig_ensure(self)
+    _created.append(path)
+    return path
+
+
+def _cleanup_plan_files() -> None:
+    for path in _created:
+        try:
+            os.remove(path)
+            d = os.path.dirname(path)
+            if os.path.isdir(d) and not os.listdir(d):
+                os.rmdir(d)
+        except OSError:
+            pass
+    _created.clear()
+
+
+_orig_run = unittest.TestCase.run
+
+
+def _run(self, result=None):
+    try:
+        return _orig_run(self, result)
+    finally:
+        _cleanup_plan_files()
+
+
+PlanMode.ensure_plan_file = _tracked_ensure
+unittest.TestCase.run = _run
+atexit.register(_cleanup_plan_files)

tests/test_agent.py

@@ -12,6 +12,8 @@
 
 sys.path.insert(0, os.path.dirname(__file__))  # sibling fake server import
 
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
+
 from python_agent_harness.agent import AgentLoop, Supervisor, sanitize_tool_result
 from python_agent_harness.agent_session import AgentSession
 from python_agent_harness.models import Message, ToolCall, Usage

tests/test_agent_session.py

@@ -9,6 +9,7 @@
 
 sys.path.insert(0, os.path.dirname(__file__))
 
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
 from test_agent import FakeClient, RecordingSession
 
 from python_agent_harness import config
@@ -655,10 +656,11 @@ def test_switch_to_plan_registers_plan_exit(self):
 
 
 class TestClose(unittest.TestCase):
-    """close() cancels the run, marks the session dead, cleans up the
-    plan file and closes the client."""
+    """close() cancels the run, marks the session dead and closes the
+    client; the plan file is kept (one per session) for later
+    reference."""
 
-    def test_close_cancels_and_cleans_up(self):
+    def test_close_cancels_and_keeps_plan_file(self):
         session = RecordingSession()
         session.switch_to_plan()  # creates a real plan file
         plan_file = session.plan_mode.plan_file
@@ -674,7 +676,7 @@ def close(self):
         session.close()
         self.assertFalse(session.alive)
         self.assertTrue(session.cancel_event.is_set())
-        self.assertFalse(os.path.exists(plan_file))
+        self.assertTrue(os.path.exists(plan_file))
         self.assertTrue(ClosableClient.closed)
 
     def test_close_closes_subagent_client_too(self):

tests/test_cli.py

@@ -3,10 +3,15 @@
 from __future__ import annotations
 
 import os
+import sys
 import tempfile
 import unittest
 from pathlib import Path
 
+sys.path.insert(0, os.path.dirname(__file__))
+
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
+
 from python_agent_harness import cli, config
 
 

tests/test_invariants.py

@@ -53,6 +53,7 @@
 
 sys.path.insert(0, os.path.dirname(__file__))  # sibling test helpers
 
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
 from test_agent import ParallelToolSession, RecordingSession, agent_call  # noqa: E402
 
 from python_agent_harness import config  # noqa: E402

tests/test_planmode.py

@@ -1,9 +1,14 @@
 import os
 import re
+import sys
 import tempfile
 import unittest
 from unittest import mock
 
+sys.path.insert(0, os.path.dirname(__file__))
+
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
+
 from python_agent_harness.models import AgentMode
 from python_agent_harness.planmode import PlanMode, _plan_temp_dir
 

tests/test_scenarios.py

@@ -20,6 +20,10 @@
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from unittest import mock
 
+sys.path.insert(0, os.path.dirname(__file__))
+
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
+
 from python_agent_harness import config
 from python_agent_harness.agent import AgentLoop
 from python_agent_harness.agent_session import AgentSession

tests/test_subagent.py

@@ -2,9 +2,15 @@
 
 from __future__ import annotations
 
+import os
+import sys
 import tempfile
 import unittest
 
+sys.path.insert(0, os.path.dirname(__file__))
+
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
+
 from python_agent_harness.agent_session import AgentSession
 from python_agent_harness.models import Message, Usage
 from python_agent_harness.session_store import SessionStore

tests/test_tui.py

@@ -2,10 +2,14 @@
 
 import io
 import os
+import sys
 import tempfile
 import unittest
 import unittest.mock as mock
 
+sys.path.insert(0, os.path.dirname(__file__))
+
+import plan_cleanup  # noqa: F401,E402  (side-effect: auto-remove /tmp plan dirs)
 from rich.console import Console
 from rich.live import Live