refactor(exec_bridge): manage subprocess context

codex_telegram_bridge/src/codex_telegram_bridge/exec_bridge.py

@@ -11,6 +11,7 @@ import shutil
 import time
 from collections import deque
 from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Any
 from weakref import WeakValueDictionary
@@ -56,6 +57,21 @@ async def _drain_stderr(stderr: asyncio.StreamReader | None, tail: deque[str]) -
         logger.debug("[codex][stderr] drain error: %s", e)
 
 
+@asynccontextmanager
+async def manage_subprocess(*args, **kwargs):
+    proc = await asyncio.create_subprocess_exec(*args, **kwargs)
+    try:
+        yield proc
+    finally:
+        if proc.returncode is None:
+            proc.terminate()
+            try:
+                await asyncio.wait_for(proc.wait(), timeout=2.0)
+            except asyncio.TimeoutError:
+                proc.kill()
+                await proc.wait()
+
+
 TELEGRAM_TEXT_LIMIT = TELEGRAM_HARD_LIMIT
 TELEGRAM_MARKDOWN_LIMIT = 3500
 
@@ -187,12 +203,12 @@ class CodexExecRunner:
         else:
             args.append("-")
 
-        proc = await asyncio.create_subprocess_exec(
+        async with manage_subprocess(
             *args,
             stdin=asyncio.subprocess.PIPE,
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
-        )
+        ) as proc:
             assert proc.stdin and proc.stdout and proc.stderr
             logger.debug("[codex] spawn pid=%s args=%r", proc.pid, args)
 
@@ -248,27 +264,14 @@ class CodexExecRunner:
                             saw_agent_message = True
             except asyncio.CancelledError:
                 cancelled = True
-            if proc.returncode is None:
-                proc.terminate()
             finally:
                 if cancelled:
                     task = asyncio.current_task()
                     if task is not None:
                         while task.cancelling():
                             task.uncancel()
-
+                if not cancelled:
-                try:
-                    rc = await asyncio.wait_for(proc.wait(), timeout=2.0)
-                except TimeoutError:
-                    logger.debug(
-                        "[codex] terminate timed out pid=%s, sending kill", proc.pid
-                    )
-                    if proc.returncode is None:
-                        proc.kill()
                     rc = await proc.wait()
-            else:
-                rc = await proc.wait()
-
                 await asyncio.gather(stderr_task, return_exceptions=True)
 
             if cancelled: