Skip to content

Commit 9ce419f

Browse files
authored
refactor(container): split the god-module and decouple SSE finalization (#63)
service.py held four unrelated responsibilities in 681 lines. It is now: docker_client.py DockerClient Protocol + aiodocker implementation streaming.py stream_events / stream_output / stream_and_persist cleanup.py expiry, shutdown, boot reconciliation repository.py all SQL service.py session use cases + finalize_session Bugs fixed: - SSE finalization was skipped whenever the client disconnected first: the generator closed at the yield and mark_completed, the scorecard, the PR comment and the done event never ran, leaving the session RUNNING until the reaper mislabelled it TIMEOUT. Finalization is now detached from the request. - The same cancellation dropped up to batch_size-1 buffered events; the final batch is now flushed (awaited on the normal path, detached on disconnect). - The DockerClient Protocol declared container_logs as returning an AsyncIterator, i.e. a coroutine resolving to one. No implementation satisfied it, and a double written to the declared contract would have deadlocked. This was hidden by a module-wide mypy disable_error_code, now removed — the container module typechecks clean with no suppressions. - stop_container marked a user abort COMPLETED, so cancelled sessions counted as successful runs. Added ContainerStatus.CANCELLED (plain VARCHAR column, no migration needed). - Wrong-state transitions raised ExternalServiceError (500); they are conflicts (409). - mark_completed defaulted to a hard-coded 900s TTL, silently ignoring CONTAINER_TTL_SECONDS on the SSE path. The TTL is now a required argument. - The periodic cleanup leaked one aiohttp session per tick. - Skill-name traversal is checked at the Docker boundary, which the webhook path reaches without passing request validation. - DELETE now returns 204 instead of an ad-hoc 200 dict. 351 tests pass, coverage 79% -> 84%, ruff and mypy clean.
1 parent d85d300 commit 9ce419f

13 files changed

Lines changed: 799 additions & 827 deletions

File tree

apps/api/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ init_typed = true
6666

6767
# Per-module overrides for third-party lib typing issues (aiodocker, SQLAlchemy, Starlette)
6868
[[tool.mypy.overrides]]
69-
module = ["helprs.modules.container.*", "helprs.main"]
69+
module = ["helprs.main"]
7070
disable_error_code = ["arg-type", "call-overload", "attr-defined"]
7171

7272
[[tool.mypy.overrides]]

apps/api/src/helprs/main.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
)
2222
from helprs.core.exceptions import DomainError, domain_exception_handler
2323
from helprs.core.middleware import configure_logging, configure_sentry, setup_middleware
24-
from helprs.modules.container import service as containers
24+
from helprs.modules.container import cleanup as container_cleanup
25+
from helprs.modules.container.docker_client import AioDockerClient
2526
from helprs.modules.container.router import router as container_router
2627
from helprs.modules.identity.router import router as identity_router
2728
from helprs.modules.installation.router import router as installation_router
@@ -110,12 +111,16 @@ async def _run_container_cleanup(
110111
while True:
111112
await asyncio.sleep(interval_seconds)
112113
try:
113-
docker = containers.AioDockerClient()
114-
async with app.state.session_factory() as db:
115-
cleaned = await containers.cleanup_expired(db, docker, ttl_seconds=ttl)
116-
await db.commit()
117-
if cleaned:
118-
logger.info("container_cleanup_cycle", cleaned=cleaned)
114+
docker = AioDockerClient()
115+
try:
116+
async with app.state.session_factory() as db:
117+
cleaned = await container_cleanup.cleanup_expired(db, docker, ttl_seconds=ttl)
118+
await db.commit()
119+
if cleaned:
120+
logger.info("container_cleanup_cycle", cleaned=cleaned)
121+
finally:
122+
# Without this the loop leaks one aiohttp session per tick.
123+
await docker.close()
119124
except Exception:
120125
logger.exception("container_cleanup_cycle_failed")
121126
except asyncio.CancelledError:
@@ -127,7 +132,7 @@ async def _reconcile_stale_sessions(session_factory: SessionFactory) -> None:
127132
"""Mark sessions left RUNNING/PENDING by a previous run as FAILED."""
128133
try:
129134
async with session_factory() as db:
130-
await containers.reconcile_stale_sessions(db)
135+
await container_cleanup.reconcile_stale_sessions(db)
131136
await db.commit()
132137
except Exception:
133138
logger.exception("session_reconciliation_failed")
@@ -136,9 +141,9 @@ async def _reconcile_stale_sessions(session_factory: SessionFactory) -> None:
136141
async def _stop_running_containers(session_factory: SessionFactory) -> None:
137142
"""Stop every container still running before the engine goes away."""
138143
try:
139-
docker = containers.AioDockerClient()
144+
docker = AioDockerClient()
140145
async with session_factory() as db:
141-
stopped = await containers.cleanup_all_running(db, docker)
146+
stopped = await container_cleanup.cleanup_all_running(db, docker)
142147
await db.commit()
143148
if stopped:
144149
logger.info("shutdown_containers_stopped", count=stopped)
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
"""Reaping containers the happy path did not clean up.
2+
3+
Three jobs, all idempotent so several workers can run them concurrently:
4+
expiry (TTL passed), shutdown (stop everything), and boot reconciliation
5+
(rows left RUNNING by a process that died).
6+
"""
7+
8+
from datetime import UTC, datetime, timedelta
9+
10+
import structlog
11+
from sqlalchemy.ext.asyncio import AsyncSession
12+
13+
from helprs.modules.container import repository
14+
from helprs.modules.container.docker_client import DockerClient
15+
from helprs.modules.container.models import ContainerSession, ContainerStatus
16+
17+
logger = structlog.get_logger()
18+
19+
20+
async def _destroy(docker: DockerClient, cs: ContainerSession) -> None:
21+
"""Best-effort teardown: a container that is already gone is a success."""
22+
if not cs.container_id:
23+
return
24+
try:
25+
await docker.stop_container(cs.container_id)
26+
await docker.remove_container(cs.container_id, force=True)
27+
except Exception as exc:
28+
await logger.awarning("cleanup_container_failed", session_id=str(cs.id), error=str(exc))
29+
30+
31+
async def cleanup_expired(db: AsyncSession, docker: DockerClient, *, ttl_seconds: int) -> int:
32+
"""Destroy containers whose session outlived its TTL. Returns the count."""
33+
cutoff = datetime.now(UTC) - timedelta(seconds=ttl_seconds)
34+
expired = await repository.list_unfinished(db, created_before=cutoff)
35+
36+
for cs in expired:
37+
await _destroy(docker, cs)
38+
cs.status = ContainerStatus.TIMEOUT
39+
cs.completed_at = datetime.now(UTC)
40+
41+
if expired:
42+
await db.flush()
43+
await logger.ainfo("expired_sessions_cleaned", count=len(expired))
44+
return len(expired)
45+
46+
47+
async def cleanup_all_running(db: AsyncSession, docker: DockerClient) -> int:
48+
"""Stop every live session — used on graceful shutdown."""
49+
sessions = await repository.list_unfinished(db)
50+
51+
for cs in sessions:
52+
await _destroy(docker, cs)
53+
cs.status = ContainerStatus.CANCELLED
54+
cs.completed_at = datetime.now(UTC)
55+
56+
if sessions:
57+
await db.flush()
58+
await logger.ainfo("shutdown_sessions_stopped", count=len(sessions))
59+
return len(sessions)
60+
61+
62+
async def reconcile_stale_sessions(db: AsyncSession) -> int:
63+
"""Fail sessions left RUNNING by a crashed process.
64+
65+
Their containers died with the host process, so waiting for the TTL would
66+
only leave the dashboard lying for 15 minutes.
67+
"""
68+
stale = await repository.list_unfinished(db)
69+
70+
for cs in stale:
71+
cs.status = ContainerStatus.FAILED
72+
cs.completed_at = datetime.now(UTC)
73+
74+
if stale:
75+
await db.flush()
76+
await logger.ainfo("stale_sessions_reconciled", count=len(stale))
77+
return len(stale)
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
"""Docker transport for runner containers.
2+
3+
The Protocol is the seam tests plug into: a hand-written double implements it
4+
without pulling in the Docker SDK.
5+
"""
6+
7+
import base64
8+
from collections.abc import AsyncIterator
9+
from typing import Protocol
10+
11+
CLAUDE_RUNNER_IMAGE = "claude-runner:latest"
12+
13+
CONTAINER_MEMORY_BYTES = 512 * 1024 * 1024
14+
CONTAINER_NANO_CPUS = 1_000_000_000
15+
16+
# The entrypoint reads newline-terminated commands from this FIFO.
17+
_INPUT_FIFO = "/tmp/claude-input"
18+
19+
20+
class DockerClient(Protocol):
21+
"""What the container module needs from a Docker daemon."""
22+
23+
async def create_container(
24+
self,
25+
image: str,
26+
environment: dict[str, str],
27+
volumes: list[str],
28+
labels: dict[str, str],
29+
) -> str:
30+
"""Create a container with stdin enabled and return its id."""
31+
...
32+
33+
async def start_container(self, container_id: str) -> None: ...
34+
35+
async def stop_container(self, container_id: str) -> None: ...
36+
37+
async def remove_container(self, container_id: str, force: bool = False) -> None: ...
38+
39+
# Not ``async def``: this returns an async iterator, it is not a coroutine
40+
# that resolves to one. Declaring it ``async`` made every implementation
41+
# fail the Protocol and would deadlock a double written to match it.
42+
def container_logs(self, container_id: str, follow: bool = False) -> AsyncIterator[str]:
43+
"""Stream the container's stdout, line chunks as they arrive."""
44+
...
45+
46+
async def write_to_container(self, container_id: str, data: str) -> None:
47+
"""Write one message to the container's input FIFO."""
48+
...
49+
50+
async def wait_container(self, container_id: str) -> int:
51+
"""Block until the container exits; return its exit code."""
52+
...
53+
54+
async def close(self) -> None: ...
55+
56+
57+
class AioDockerClient:
58+
"""Production client wrapping aiodocker."""
59+
60+
def __init__(self) -> None:
61+
import aiodocker
62+
63+
self._docker = aiodocker.Docker()
64+
65+
async def create_container(
66+
self,
67+
image: str,
68+
environment: dict[str, str],
69+
volumes: list[str],
70+
labels: dict[str, str],
71+
) -> str:
72+
container = await self._docker.containers.create(
73+
{
74+
"Image": image,
75+
"Env": [f"{k}={v}" for k, v in environment.items()],
76+
"Labels": labels,
77+
"OpenStdin": True,
78+
"HostConfig": {
79+
"Binds": volumes,
80+
"Memory": CONTAINER_MEMORY_BYTES,
81+
"NanoCPUs": CONTAINER_NANO_CPUS,
82+
"NetworkMode": "bridge",
83+
},
84+
}
85+
)
86+
return container.id
87+
88+
async def start_container(self, container_id: str) -> None:
89+
container = await self._docker.containers.get(container_id)
90+
await container.start()
91+
92+
async def stop_container(self, container_id: str) -> None:
93+
container = await self._docker.containers.get(container_id)
94+
await container.stop()
95+
96+
async def remove_container(self, container_id: str, force: bool = False) -> None:
97+
container = await self._docker.containers.get(container_id)
98+
await container.delete(force=force)
99+
100+
async def container_logs(self, container_id: str, follow: bool = False) -> AsyncIterator[str]:
101+
"""Yield the container's stdout.
102+
103+
stdout only: Claude Code writes stream-json there, while stderr
104+
carries diagnostics that would duplicate events. aiodocker returns a
105+
list when not following and an async iterator when following, so the
106+
two cases are spelled out rather than passed a dynamic flag.
107+
"""
108+
container = await self._docker.containers.get(container_id)
109+
if follow:
110+
async for line in container.log(stdout=True, stderr=False, follow=True):
111+
yield line
112+
return
113+
for line in await container.log(stdout=True, stderr=False, follow=False):
114+
yield line
115+
116+
async def write_to_container(self, container_id: str, data: str) -> None:
117+
"""Deliver one message to the entrypoint's read loop.
118+
119+
The payload is base64-encoded so no shell metacharacter in user text
120+
can be interpreted, and newline-terminated so ``read`` returns it as a
121+
complete line.
122+
"""
123+
container = await self._docker.containers.get(container_id)
124+
encoded = base64.b64encode((data + "\n").encode()).decode()
125+
exec_obj = await container.exec(cmd=["sh", "-c", f"echo {encoded} | base64 -d > {_INPUT_FIFO}"])
126+
await exec_obj.start(detach=True)
127+
128+
async def wait_container(self, container_id: str) -> int:
129+
container = await self._docker.containers.get(container_id)
130+
result = await container.wait()
131+
return result["StatusCode"]
132+
133+
async def close(self) -> None:
134+
await self._docker.close()

apps/api/src/helprs/modules/container/models.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@ class ContainerStatus(enum.StrEnum):
1919
COMPLETED = "completed"
2020
FAILED = "failed"
2121
TIMEOUT = "timeout"
22+
# A user pressed stop, or the API shut down: neither a success nor a
23+
# failure of the skill. The column is a plain VARCHAR with no CHECK
24+
# constraint, so adding a value needs no migration.
25+
CANCELLED = "cancelled"
2226

2327

2428
class ContainerSession(Base):

apps/api/src/helprs/modules/container/repository.py

Lines changed: 45 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from sqlalchemy.ext.asyncio import AsyncSession
1414
from sqlalchemy.types import Date
1515

16-
from helprs.modules.container.models import ContainerSession, ContainerStatus
16+
from helprs.modules.container.models import ContainerSession, ContainerStatus, SessionEvent
1717

1818

1919
@dataclass(frozen=True)
@@ -38,6 +38,47 @@ class DailyCount:
3838
count: int
3939

4040

41+
UNFINISHED_STATUSES = (ContainerStatus.RUNNING, ContainerStatus.PENDING)
42+
43+
44+
async def get(session: AsyncSession, session_id: uuid.UUID) -> ContainerSession | None:
45+
result = await session.execute(select(ContainerSession).where(ContainerSession.id == session_id))
46+
return result.scalar_one_or_none()
47+
48+
49+
async def add(session: AsyncSession, container_session: ContainerSession) -> ContainerSession:
50+
session.add(container_session)
51+
await session.flush()
52+
return container_session
53+
54+
55+
async def delete(session: AsyncSession, container_session: ContainerSession) -> None:
56+
"""Delete a session; its events go with it via ON DELETE CASCADE."""
57+
await session.delete(container_session)
58+
await session.flush()
59+
60+
61+
async def list_unfinished(
62+
session: AsyncSession,
63+
*,
64+
created_before: datetime | None = None,
65+
) -> list[ContainerSession]:
66+
"""Sessions still marked RUNNING or PENDING, optionally older than a cutoff."""
67+
query = select(ContainerSession).where(ContainerSession.status.in_(UNFINISHED_STATUSES))
68+
if created_before is not None:
69+
query = query.where(ContainerSession.created_at < created_before)
70+
result = await session.execute(query)
71+
return list(result.scalars().all())
72+
73+
74+
async def list_events(session: AsyncSession, session_id: uuid.UUID) -> list[SessionEvent]:
75+
"""Every persisted event for a session, in arrival order."""
76+
result = await session.execute(
77+
select(SessionEvent).where(SessionEvent.session_id == session_id).order_by(SessionEvent.event_id)
78+
)
79+
return list(result.scalars().all())
80+
81+
4182
async def count_by_status(session: AsyncSession, installation_ids: list[uuid.UUID]) -> StatusCounts:
4283
"""Count sessions per terminal status across the given installations."""
4384
if not installation_ids:
@@ -73,7 +114,8 @@ async def count_per_day(
73114
result = await session.execute(
74115
select(
75116
cast(ContainerSession.created_at, Date).label("day"),
76-
func.count(ContainerSession.id).label("count"),
117+
# Not labelled "count": that shadows Sequence.count on the Row.
118+
func.count(ContainerSession.id).label("sessions"),
77119
)
78120
.where(
79121
ContainerSession.installation_id.in_(installation_ids),
@@ -82,7 +124,7 @@ async def count_per_day(
82124
.group_by("day")
83125
.order_by("day")
84126
)
85-
return [DailyCount(day=row.day, count=row.count) for row in result.all()]
127+
return [DailyCount(day=row.day, count=row.sessions) for row in result.all()]
86128

87129

88130
async def list_for_installation(

0 commit comments

Comments
 (0)