Merge pull request #3 from SmartAI/codex/audit-pythonic-code-architecture

Code refactor

Author: SmartAI

.github/workflows/ci.yml

@@ -0,0 +1,26 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.12"
+          enable-cache: true
+      - uses: actions/setup-node@v4
+        with:
+          node-version: 22
+          cache: npm
+      - run: uv sync --frozen
+      - run: uv run ruff check src tests
+      - run: uv run mypy src
+      - run: uv run pytest -q
+      - run: npm ci
+      - run: npm run check

README.md

@@ -2,9 +2,9 @@
 
 A coding-agent harness in Python: a durable, replayable session log, a `drive → turn → step` loop
 with a two-target inbox, four bounded tools, Anthropic and OpenAI-compatible providers, and a
-loopback Web UI. It is a port of [ava-cpp](../ava-cpp)'s architecture; the design documents there
-(`docs/agent-loop.md`, `docs/session-log.md`, `docs/model-layer.md`, `docs/functional-spec.md`)
-remain the specification this code implements.
+loopback Web UI. It is a port of ava-cpp's architecture; the local
+[architecture notes](docs/architecture.md) describe the boundaries and invariants this code
+implements.
 
 ## The shape
 
@@ -30,10 +30,10 @@ from pathlib import Path
 from ava.agent import Agent
 from ava.llm import Item, Role, make_text_block, provider_from_environment
 
-agent = Agent.create(provider_from_environment(), Path.cwd())
-subscription = agent.subscribe(lambda event: print(event.seq, type(event.payload).__name__))
-await agent.followup(Item(role=Role.user, blocks=[make_text_block("add a test for the parser")]))
-await agent.drive()
+async with Agent.create(provider_from_environment(), Path.cwd()) as agent:
+    subscription = agent.subscribe(lambda event: print(event.seq, type(event.payload).__name__))
+    await agent.followup(Item(role=Role.user, blocks=[make_text_block("add a test for the parser")]))
+    await agent.drive()
 ```
 
 Every model-visible fact is a durable event. The model's context, the browser's transcript, and
@@ -124,6 +124,7 @@ The checked-in browser bundle is rebuilt with `npm ci && npm run build` when tha
 ```sh
 uv run pytest
 uv run ruff check src tests
+uv run mypy src
 npm run check
 ```
 

docs/architecture.md

@@ -0,0 +1,59 @@
+# Architecture
+
+Ava has a headless core and two thin application surfaces: the CLI and the loopback Web UI.
+Dependencies point inward; provider wire formats and frontend concerns do not enter the agent or
+session packages.
+
+```text
+app ──► agent ──► session ──► llm ──► transport
+          │           │          │
+          └──► tool ──┴──► proc  └──► base
+```
+
+## Public seam
+
+`Agent` is the application boundary. Applications submit input, control the driver, subscribe to
+durable events, and read provider-neutral status or context reports through that object. Access to
+`Agent.state` is reserved for agent internals and invariant-focused tests.
+
+An agent owns its provider, durable writer, and scratchpad. Call `await agent.aclose()` when done,
+or use `async with`. The Web registry closes every chat agent during the FastAPI lifespan shutdown.
+
+## Durable input seam
+
+Acknowledging input, claiming input for a step, and deciding a turn boundary are serialized by one
+inbox gate. `followup()` and `steer()` return only after their inbox splice is durable. A step claims
+one next-turn message on its first step and all available next-step messages. The gate remains held
+from the driver's snapshot through that claim so an acknowledged message cannot be lost between
+the two operations.
+
+## Driver lifecycle
+
+The driver moves through `idle`, `running`, `pausing`, `paused`, and `aborting`. A turn contains one
+or more provider steps; a step either ends the turn or emits tool calls whose paired results become
+the next step's context. Pause is honored at a complete step boundary. Abort repairs partial model
+and tool output before it closes the turn.
+
+Only one driver may own an agent. The Web `DriveHandoff` compensates for a message acknowledged as
+the owning driver finishes, without starting a second driver.
+
+## Session authority
+
+The append-only session event stream is the authority for model context, browser replay, pending
+inbox state, and crash recovery. Physical writes are bounded frames; subscribers see durable events
+only after their bytes reach the kernel. Recovery may repair a torn final frame and append missing
+`interrupted` lifecycle closers, but never rewrites a complete logical event.
+
+The JSONL codec is intentionally explicit. Its event-by-event mapping is compatibility and
+validation code, not generic object serialization, and unknown future event kinds are preserved.
+
+## Evaluation gates
+
+Every architectural change must keep these gates green:
+
+- durable input is acknowledged and claimed exactly once;
+- every model tool call has one result after completion, abort, or recovery;
+- pause/resume and provider failure leave replayable history;
+- `ruff`, `mypy`, the Python 3.12 test suite, and the frontend build pass;
+- application code has no direct `agent.state` access;
+- provider clients close on agent replacement and shutdown.

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "fastapi>=0.115",
     "uvicorn>=0.30",
     "httpx>=0.27",
+    "pydantic>=2.10",
     "zstandard>=0.22",
 ]
 
@@ -24,6 +25,7 @@ packages = ["src/ava"]
 
 [dependency-groups]
 dev = [
+    "mypy>=1.15",
     "pytest>=8",
     "pytest-asyncio>=0.24",
     "pytest-timeout>=2.4.0",
@@ -41,3 +43,11 @@ target-version = "py312"
 [tool.ruff.lint]
 select = ["E", "F", "I", "UP", "B"]
 ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.12"
+check_untyped_defs = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true

src/ava/agent/agent.py

@@ -36,6 +36,7 @@
     sort_model_ids,
 )
 from ava.session import (
+    AttemptTiming,
     DriveError,
     EventSink,
     InboxMessage,
@@ -45,8 +46,11 @@
     TurnEnd,
     TurnEndReason,
     TurnStart,
+    Usage,
 )
 from ava.session.codec import validate_step_claimed_record
+from ava.session.compaction import estimate_context_tokens
+from ava.session.context_report import ContextReport, context_report
 from ava.session.log import Log, OpenMode
 
 
@@ -107,8 +111,21 @@ def reopen(
             agent._state.drive_state.restore_paused()
         return agent
 
-    def close(self) -> None:
-        self._state.close()
+    async def aclose(self) -> None:
+        """Close the provider and durable state. The agent must be idle."""
+        state = self._state
+        if state.drive_state.status not in (Status.idle, Status.paused):
+            raise AvaError(ErrorKind.invalid_argument, "cannot close the agent while it is running")
+        try:
+            await state.provider.aclose()
+        finally:
+            state.close()
+
+    async def __aenter__(self) -> Agent:
+        return self
+
+    async def __aexit__(self, *_: object) -> None:
+        await self.aclose()
 
     # ---- observation ----------------------------------------------------------------------
 
@@ -224,7 +241,7 @@ async def model_choices(self) -> ModelChoices:
                 available = False
         finally:
             state.model_catalog_operations -= 1
-        current = state.pending_model or state.provider.selection.model
+        current = self.current_selection().model
         models = [*listed, current, *state.provider.model_aliases.values()]
         sort_model_ids(models)
         deduplicated = list(dict.fromkeys(models))
@@ -235,21 +252,94 @@ async def model_choices(self) -> ModelChoices:
     def select_model(self, model: str) -> None:
         if not model:
             raise AvaError(ErrorKind.invalid_argument, "model id must not be empty")
-        self._state.pending_model = model
-        self._state.pending_effort = None
+        current = self.current_selection()
+        self._state.pending_selection = Selection(current.provider, model, None)
         self._state.model_revision += 1
 
     def current_selection(self) -> Selection:
-        provider = self._state.provider
-        selection = Selection(
-            provider.selection.provider, provider.selection.model, provider.selection.effort
+        pending = self._state.pending_selection
+        selected = pending or self._state.provider.selection
+        return Selection(selected.provider, selected.model, selected.effort)
+
+    def current_capabilities(self) -> ModelCapabilities:
+        """Capabilities for the pending or active model selection."""
+        return self._state.provider.capabilities(self.current_selection().model)
+
+    def context_report(self, *, prepare: bool = False) -> ContextReport:
+        """Describe the provider-neutral context without exposing runtime state."""
+        if prepare:
+            self.prepare()
+        state = self._state
+        return context_report(
+            state.session,
+            state.provider.context_window,
+            state.compaction_options.threshold_percent,
+        )
+
+    def status_snapshot(self) -> dict[str, object]:
+        """Return the stable frontend status contract from agent-owned state."""
+        state = self._state
+        selected = self.current_selection()
+        report = self.context_report()
+        used_tokens = (
+            report.estimated_tokens
+            if report.measured_input_tokens is None
+            else estimate_context_tokens(state.session)
+        )
+        capabilities = state.provider.capabilities(selected.model)
+        context_window = capabilities.context_window_tokens
+        if context_window is None and selected.model == state.provider.selection.model:
+            context_window = state.provider.context_window
+        context_window = context_window or 0
+        remaining = (
+            round(max(0, context_window - used_tokens) * 100 / context_window)
+            if context_window
+            else None
+        )
+
+        input_tokens: int | None = None
+        output_tokens: int | None = None
+        cached_read = 0
+        cache_reported = False
+        ttft_ms: int | None = None
+        for event in state.session.events:
+            match event.payload:
+                case Usage(
+                    input=input,
+                    cached_read=cached_read_tokens,
+                    cache_write=cache_write,
+                    output=output,
+                ):
+                    input_parts = (input, cached_read_tokens, cache_write)
+                    if any(value is not None for value in input_parts):
+                        input_tokens = (input_tokens or 0) + sum(value or 0 for value in input_parts)
+                    if output is not None:
+                        output_tokens = (output_tokens or 0) + output
+                    if cached_read_tokens is not None:
+                        cached_read += cached_read_tokens
+                        cache_reported = True
+                case AttemptTiming(ttft_ms=latest_ttft):
+                    # The newest attempt describes the latency the user most recently felt.
+                    ttft_ms = latest_ttft
+
+        cache_hit_percent = (
+            cached_read * 100 // input_tokens if cache_reported and input_tokens else None
         )
-        if self._state.pending_model is not None:
-            selection.model = self._state.pending_model
-            selection.effort = None
-        if self._state.pending_effort is not None:
-            selection.effort = self._state.pending_effort
-        return selection
+        return {
+            "status": self.status.value,
+            "turn_open": self.turn_open,
+            "cwd": str(self.cwd),
+            "provider": selected.provider,
+            "model": selected.model,
+            "effort": selected.effort,
+            "context_used_tokens": used_tokens,
+            "context_window_tokens": context_window,
+            "context_remaining_percent": remaining,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_hit_percent": cache_hit_percent,
+            "ttft_ms": ttft_ms,
+        }
 
     async def cycle_effort(self) -> str:
         state = self._state
@@ -280,7 +370,7 @@ async def cycle_effort(self) -> str:
         except ValueError:
             index = -1
         next_effort = efforts[0] if index < 0 or index + 1 >= len(efforts) else efforts[index + 1]
-        state.pending_effort = next_effort
+        state.pending_selection = Selection(selected.provider, selected.model, next_effort)
         return next_effort
 
     def select_effort(self, effort: str | None) -> None:
@@ -295,9 +385,10 @@ def select_effort(self, effort: str | None) -> None:
                     f"the current model does not advertise reasoning effort '{effort}'; "
                     f"choose one of {', '.join(capabilities.effort_values) or 'none'}",
                 )
-        self._state.pending_effort = effort
+        selected = self.current_selection()
+        self._state.pending_selection = Selection(selected.provider, selected.model, effort)
 
-    def reload_credentials(
+    async def reload_credentials(
         self, auth_requirement: AuthRequirement = AuthRequirement.required
     ) -> None:
         state = self._state
@@ -307,7 +398,11 @@ def reload_credentials(
             or state.model_catalog_operations
         ):
             raise AvaError(ErrorKind.invalid_argument, "cannot reload credentials while busy")
-        state.provider = provider_from_environment(None, state.provider.selection, auth_requirement)
+        old_provider = state.provider
+        new_provider = provider_from_environment(None, old_provider.selection, auth_requirement)
+        state.provider = new_provider
+        if new_provider is not old_provider:
+            await old_provider.aclose()
 
     # ---- the driver ----------------------------------------------------------------------
 
@@ -410,13 +505,15 @@ async def drive(self) -> None:
                         owns = False
                         return
                     if not drive.close_turn():
-                        error = AvaError(
+                        close_error = AvaError(
                             ErrorKind.internal,
                             "cannot close a turn while tool results remain outstanding",
                         )
-                        await self._finish_failed_turn(turn_number, TurnEndReason.completed, error)
+                        await self._finish_failed_turn(
+                            turn_number, TurnEndReason.completed, close_error
+                        )
                         owns = False
-                        raise error
+                        raise close_error
                     elapsed_ms = int((time.monotonic() - turn_started) * 1000)
                     state.append(
                         TurnEnd(

src/ava/agent/state.py

@@ -329,8 +329,7 @@ class AgentState:
     maintenance_running: bool = False
     ordinals_recovered: bool = False
     initialized: bool = False
-    pending_model: str | None = None
-    pending_effort: str | None = None
+    pending_selection: Selection | None = None
 
     @classmethod
     def create(