Coverage for gco_mcp / mission / state.py: 100.00%
147 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Persistence backends for the Mission goal-directed iteration loop.
3This module defines the :class:`MissionStateBackend` protocol — the narrow
4interface the engine and tool wrappers depend on for loading, saving,
5listing, and deleting :class:`~mcp.mission.types.SessionState` records.
6Concrete implementations (filesystem, DynamoDB) and the
7:func:`get_backend` resolver land in follow-on slices of this file. The
8protocol is declared with :func:`typing.runtime_checkable` so tests can
9assert backend conformance with ``isinstance`` rather than relying on
10duck-typed call sites.
11"""
13from __future__ import annotations
15import contextlib
16import json
17import logging
18import os
19import tempfile
20from decimal import Decimal
21from pathlib import Path
22from typing import Any, Protocol, cast, runtime_checkable
24from . import SCHEMA_VERSION
25from .types import SessionState
27logger = logging.getLogger(__name__)
30@runtime_checkable
31class MissionStateBackend(Protocol):
32 """Storage contract for Mission session records.
34 All four methods operate on whole :class:`SessionState` payloads keyed
35 by ``session_id``. Implementations are responsible for whatever
36 serialization, atomicity, and access-control guarantees their backing
37 store provides; callers treat the interface as opaque key-value
38 storage with a list operation that returns lightweight metadata
39 rather than full session bodies.
40 """
42 def load_session(self, session_id: str) -> SessionState | None:
43 """Return the session record for ``session_id`` or ``None`` if absent.
45 Implementations return ``None`` for both unknown ``session_id`` and
46 records whose ``version`` does not match the current
47 :data:`mcp.mission.SCHEMA_VERSION`; the caller cannot distinguish
48 the two and treats both as a missing session.
49 """
50 ...
52 def save_session(self, session: SessionState) -> None:
53 """Persist ``session`` keyed by its ``session_id`` field.
55 Writes are expected to be atomic from the reader's perspective: a
56 concurrent :meth:`load_session` either sees the prior record or
57 the new one, never a partial write.
58 """
59 ...
61 def list_sessions(self, filter: dict[str, Any] | None = None) -> list[dict[str, Any]]:
62 """Return lightweight metadata for sessions matching ``filter``.
64 Each returned dict carries identifying fields (``session_id``,
65 ``status``, ``created_at``, and similar) rather than the full
66 :class:`SessionState`. ``filter`` is an implementation-defined
67 mapping; passing ``None`` lists every session the backend can
68 see.
69 """
70 ...
72 def delete_session(self, session_id: str) -> bool:
73 """Remove ``session_id`` and return ``True`` if a record was deleted.
75 Returns ``False`` when no record existed; implementations do not
76 raise on a missing key so that callers can use ``delete_session``
77 as an idempotent cleanup primitive.
78 """
79 ...
82class FilesystemBackend:
83 """JSON-on-disk implementation of :class:`MissionStateBackend`.
85 Each session is persisted as ``<root>/<session_id>.json`` with its
86 matching :class:`~mcp.mission.types.SessionState` payload; the
87 Final_Report (when present) lives alongside it as
88 ``<root>/<session_id>.report.json``. Writes go through the standard
89 "temp file in the same directory, ``fsync``, then ``os.replace``"
90 pattern so a reader concurrent with a writer always sees either the
91 prior version of the file or the new one — never a partial JSON
92 document. The temp file lives in the same directory as the final
93 target so ``os.replace`` is a same-filesystem rename and therefore
94 atomic on POSIX.
96 On POSIX systems the root directory is created (and re-asserted) at
97 mode ``0o700`` and every session and report file is written at mode
98 ``0o600`` so persisted state is unreadable to other local users.
99 Permission calls are gated on ``os.name != "nt"`` because the POSIX
100 permission model does not apply on Windows; the backend still works
101 on Windows, just without the explicit mode tightening.
102 """
104 def __init__(self, root: Path | None = None) -> None:
105 self.root = root if root is not None else Path.home() / ".gco" / "missions"
106 self._root_initialized = False
108 # ------------------------------------------------------------------ #
109 # internals
110 # ------------------------------------------------------------------ #
112 def _ensure_root(self) -> None:
113 """Create the root directory on first use, idempotently.
115 We defer the ``mkdir`` to the first write so simply constructing
116 a backend (e.g. in the resolver in :func:`get_backend`) does not
117 eagerly create ``~/.gco/missions`` on a host that ends up using
118 a different backend.
119 """
120 if self._root_initialized:
121 return
122 self.root.mkdir(parents=True, exist_ok=True)
123 if os.name != "nt":
124 with contextlib.suppress(OSError):
125 # Best-effort tightening: a directory we already own with
126 # different permissions is still safer to use than to
127 # refuse the write outright. 0o700 (owner-only) is
128 # intentional for ~/.gco/missions: session JSON contains
129 # operator-supplied directives, criteria, observations,
130 # and tool-call results that should not be readable by
131 # other local users.
132 # nosemgrep: python.lang.security.audit.insecure-file-permissions.insecure-file-permissions
133 os.chmod(self.root, 0o700)
134 self._root_initialized = True
136 def _session_path(self, session_id: str) -> Path:
137 return self.root / f"{session_id}.json"
139 def _report_path(self, session_id: str) -> Path:
140 return self.root / f"{session_id}.report.json"
142 # ------------------------------------------------------------------ #
143 # protocol methods
144 # ------------------------------------------------------------------ #
146 def load_session(self, session_id: str) -> SessionState | None:
147 """Return the persisted session or ``None`` for missing/unsupported.
149 Returns ``None`` when the file does not exist, when the root
150 directory has not been created yet, when the JSON cannot be
151 parsed, or when the on-disk ``version`` field does not match
152 :data:`mcp.mission.SCHEMA_VERSION`. Version mismatches log a
153 single warning naming the unsupported value so an operator can
154 spot stale state without having to grep the directory by hand.
155 """
156 path = self._session_path(session_id)
157 try:
158 text = path.read_text(encoding="utf-8")
159 except FileNotFoundError:
160 return None
161 except OSError:
162 return None
164 try:
165 payload = json.loads(text)
166 except ValueError:
167 return None
169 if not isinstance(payload, dict):
170 return None
172 version = payload.get("version")
173 if version != SCHEMA_VERSION:
174 logger.warning(
175 "Refusing to load Mission session %s: unsupported schema version %r",
176 session_id,
177 version,
178 )
179 return None
181 return payload # type: ignore[return-value]
183 def save_session(self, session: SessionState) -> None:
184 """Persist ``session`` atomically to ``<root>/<session_id>.json``.
186 Opens a temp file in the same directory, dumps JSON, flushes and
187 ``fsync``s, applies POSIX mode ``0o600`` (when supported), then
188 ``os.replace``s onto the final path. A failure mid-write leaves
189 the temp file behind but never replaces the existing final file,
190 so the previously-persisted state remains loadable.
192 Defense-in-depth strip. The validators in
193 :mod:`mcp.mission.validation` attach a cached
194 :class:`ast.Expression` under ``_parsed_ast`` on every
195 ``predicate`` criterion. That object is not JSON-serialisable;
196 a caller that hands a freshly-validated session straight to
197 ``save_session`` without first stripping the cache would
198 raise :class:`TypeError` at ``json.dump`` time. We strip
199 unconditionally here so every persistence path stays correct
200 regardless of which caller forgot. The strip is cheap and
201 idempotent on already-clean inputs.
202 """
203 # Local import: ``mission.validation`` is part of the same
204 # package so this isn't a cross-package edge, just a
205 # dependency-direction kept lazy to keep the eager import
206 # surface of ``state`` minimal.
207 from .validation import strip_private_fields
209 self._ensure_root()
210 session_id = session["session_id"]
211 final = self._session_path(session_id)
212 cleaned = cast("SessionState", strip_private_fields(session))
214 try:
215 tmp = tempfile.NamedTemporaryFile( # noqa: SIM115 - explicit close+replace below
216 mode="w",
217 encoding="utf-8",
218 dir=str(self.root),
219 prefix=f"{session_id}.",
220 suffix=".json.tmp",
221 delete=False,
222 )
223 try:
224 json.dump(cleaned, tmp)
225 tmp.flush()
226 os.fsync(tmp.fileno())
227 finally:
228 tmp.close()
230 if os.name != "nt":
231 with contextlib.suppress(OSError):
232 # Same rationale as in ``_ensure_root`` — proceed
233 # with the replace rather than abandoning a write
234 # we already fsynced.
235 os.chmod(tmp.name, 0o600)
237 os.replace(tmp.name, final)
238 except OSError as exc:
239 # Re-raise with the underlying message intact so callers and
240 # operators see the real cause (disk full, permission denied,
241 # etc.) rather than a wrapped abstraction.
242 raise OSError(str(exc)) from exc
244 def list_sessions(self, filter: dict[str, Any] | None = None) -> list[dict[str, Any]]:
245 """Return summary dicts for every parseable session under ``root``.
247 Each entry has the shape ``{"session_id", "status", "created_at",
248 "iteration_count"}``. Sessions whose JSON fails to parse, whose
249 version is unsupported, or which are missing required summary
250 fields are silently skipped (one debug-log line per skip) so a
251 single corrupt file cannot block listing the rest.
253 ``filter`` currently supports the ``status`` key only; callers
254 pass ``{"status": "running"}`` to narrow the list.
255 """
256 if not self.root.exists():
257 return []
259 results: list[dict[str, Any]] = []
260 for path in self.root.glob("*.json"):
261 # Skip the sibling report files — they share the directory
262 # but are not session payloads.
263 if path.name.endswith(".report.json"):
264 continue
265 try:
266 payload = json.loads(path.read_text(encoding="utf-8"))
267 except OSError, ValueError:
268 logger.debug("Skipping unreadable Mission file: %s", path)
269 continue
270 if not isinstance(payload, dict):
271 logger.debug("Skipping non-object Mission file: %s", path)
272 continue
273 if payload.get("version") != SCHEMA_VERSION:
274 logger.debug(
275 "Skipping Mission file %s with unknown version %r",
276 path,
277 payload.get("version"),
278 )
279 continue
281 summary = {
282 "session_id": payload.get("session_id", path.stem),
283 "status": payload.get("status"),
284 "created_at": payload.get("created_at"),
285 "iteration_count": len(payload.get("iterations", []) or []),
286 }
287 results.append(summary)
289 if filter and "status" in filter:
290 wanted = filter["status"]
291 results = [r for r in results if r.get("status") == wanted]
293 return results
295 def delete_session(self, session_id: str) -> bool:
296 """Remove the session JSON and any matching report file.
298 Returns ``True`` when at least one of the two files existed and
299 was removed; ``False`` when neither was present (including when
300 the root directory has never been created). The two removals are
301 independent so a stale ``.report.json`` left behind by an
302 earlier crash is still cleaned up even when the session JSON has
303 already been deleted.
304 """
305 if not self.root.exists():
306 return False
308 removed = False
309 for path in (self._session_path(session_id), self._report_path(session_id)):
310 try:
311 os.remove(path)
312 removed = True
313 except FileNotFoundError:
314 continue
315 except OSError:
316 # An unreadable-but-present file should not silently
317 # masquerade as "no record existed"; surface it.
318 raise
319 return removed
322def _to_dynamodb_item(value: Any) -> Any:
323 """Recursively convert a session payload into DynamoDB-storable types.
325 The boto3 resource API rejects ``float`` outright (``TypeError: Float
326 types are not supported``), and a Mission session carries floats in
327 ordinary places — criterion targets, observed metric values, budget
328 limits. Floats go through ``Decimal(str(x))`` so the decimal string
329 round-trips without binary artifacts; ``bool`` is checked first
330 because it is an ``int`` subclass and must stay a DynamoDB ``BOOL``.
331 """
332 if isinstance(value, bool):
333 return value
334 if isinstance(value, float):
335 return Decimal(str(value))
336 if isinstance(value, dict):
337 return {key: _to_dynamodb_item(entry) for key, entry in value.items()}
338 if isinstance(value, list):
339 return [_to_dynamodb_item(entry) for entry in value]
340 return value
343def _from_dynamodb_item(value: Any) -> Any:
344 """Recursively convert DynamoDB numbers back to plain ``int``/``float``.
346 The resource API deserializes every ``N`` as :class:`~decimal.Decimal`;
347 integral values become ``int`` and the rest ``float`` so a session
348 loaded from DynamoDB is indistinguishable from one loaded from the
349 filesystem backend's JSON (``version == SCHEMA_VERSION`` compares an
350 ``int``, reports ``json.dumps`` the payload, comparisons stay numeric).
351 """
352 if isinstance(value, Decimal):
353 return int(value) if value == value.to_integral_value() else float(value)
354 if isinstance(value, dict):
355 return {key: _from_dynamodb_item(entry) for key, entry in value.items()}
356 if isinstance(value, list):
357 return [_from_dynamodb_item(entry) for entry in value]
358 return value
361class DynamoDBBackend:
362 """DynamoDB-backed implementation of :class:`MissionStateBackend`.
364 The unit suite never touches AWS from here (every boto3-facing method
365 is excluded from coverage); the backend runs for real against the
366 Floci emulator in ``tests/test_floci_mission_state.py``, over a table
367 shaped exactly like the one ``gco/stacks/global_stack.py`` provisions.
369 Item schema mirrors the :class:`SessionState` TypedDict one-to-one:
370 the partition key is ``session_id`` and ``status`` plus ``created_at``
371 feed a ``status-index`` GSI so :meth:`list_sessions` can filter by
372 status without a full table scan. ``put_item`` is atomic by virtue
373 of DynamoDB's single-item write semantics, so the temp-file dance
374 used by :class:`FilesystemBackend` is unnecessary here. Numbers cross
375 the wire through :func:`_to_dynamodb_item` / :func:`_from_dynamodb_item`
376 so both backends hand the engine the same Python types.
378 Table-name resolution is lazy: when the constructor's ``table_name``
379 argument is ``None``, the table name is fetched from SSM at
380 ``/{project_name}/missions-table-name`` on the first call that
381 needs it (not at construction time). This matches the precedent
382 pattern in ``cli/models.py`` and lets unit tests construct a
383 ``DynamoDBBackend()`` on a host without AWS credentials without
384 triggering an SSM call. ``project_name`` is read from the
385 ``GCO_PROJECT_NAME`` environment variable, defaulting to ``"gco"``
386 so a fresh checkout (or CI run without the env var set) lines up
387 with the default project name in ``cli/config.py``.
389 The SSM lookup goes through :func:`gco.services.aws_ssm.get_ssm_parameter`,
390 the shared helper that consolidates the pattern previously duplicated
391 across ``cli/models.py``, ``cli/analytics_user_mgmt.py``, and
392 ``gco/services/health_monitor.py``. Putting the helper under
393 ``gco/services/`` (rather than ``cli/aws_client.py``) keeps
394 ``gco_mcp/`` free of the forbidden ``mcp -> cli`` import edge while
395 still letting every backend share one implementation.
396 """
398 def __init__(self, table_name: str | None = None) -> None:
399 self._table_name: str | None = table_name
400 self._table: Any = None # boto3 Table resource, lazily constructed
402 # ------------------------------------------------------------------ #
403 # internals
404 # ------------------------------------------------------------------ #
406 def _resolve_table_name(self) -> str: # pragma: no cover - boto3 / SSM
407 """Return the cached table name, fetching from SSM on first call.
409 Reads ``GCO_PROJECT_NAME`` (default ``"gco"``) to build the SSM
410 parameter path ``/{project_name}/missions-table-name``. The
411 value is cached on the instance so subsequent method calls do
412 not re-hit SSM.
413 """
414 if self._table_name is not None:
415 return self._table_name
417 from gco.services.aws_ssm import get_ssm_parameter
419 project_name = os.environ.get("GCO_PROJECT_NAME", "gco")
420 param_name = f"/{project_name}/missions-table-name"
422 self._table_name = get_ssm_parameter(param_name)
423 return self._table_name
425 def _get_table(self) -> Any: # pragma: no cover - boto3 resource
426 """Return the cached ``boto3`` Table resource, building it lazily."""
427 if self._table is not None:
428 return self._table
430 import boto3
432 self._table = boto3.resource("dynamodb").Table(self._resolve_table_name())
433 return self._table
435 # ------------------------------------------------------------------ #
436 # protocol methods
437 # ------------------------------------------------------------------ #
439 def load_session(self, session_id: str) -> SessionState | None: # pragma: no cover - DynamoDB
440 """Fetch the session via ``get_item`` keyed on ``session_id``."""
441 table = self._get_table()
442 response = table.get_item(Key={"session_id": session_id})
443 raw = response.get("Item")
444 if raw is None:
445 return None
446 item = _from_dynamodb_item(raw)
447 if item.get("version") != SCHEMA_VERSION:
448 logger.warning(
449 "Refusing to load Mission session %s: unsupported schema version %r",
450 session_id,
451 item.get("version"),
452 )
453 return None
454 return cast("SessionState", item)
456 def save_session(self, session: SessionState) -> None: # pragma: no cover - DynamoDB
457 """Persist the session via ``put_item`` (atomic single-item write).
459 Defense-in-depth strip — same rationale as
460 :meth:`FilesystemBackend.save_session`. DynamoDB serialises
461 through boto3's own type-converter, which raises
462 :class:`TypeError` on an :class:`ast.Expression` just like
463 the JSON path; stripping here keeps both backends symmetric.
464 """
465 from .validation import strip_private_fields
467 table = self._get_table()
468 table.put_item(Item=_to_dynamodb_item(strip_private_fields(session)))
470 def list_sessions(
471 self, filter: dict[str, Any] | None = None
472 ) -> list[dict[str, Any]]: # pragma: no cover - DynamoDB
473 """Return summary dicts via the ``status-index`` GSI.
475 When ``filter`` provides a ``status`` key, the call uses the GSI
476 partition key directly. With no filter (or any other filter
477 shape), this stub falls back to a table ``scan`` so that the
478 method still returns the same summary shape as
479 :meth:`FilesystemBackend.list_sessions`.
480 """
481 from boto3.dynamodb.conditions import Key
483 table = self._get_table()
484 if filter and "status" in filter:
485 response = table.query(
486 IndexName="status-index",
487 KeyConditionExpression=Key("status").eq(filter["status"]),
488 )
489 items = response.get("Items", [])
490 else:
491 response = table.scan()
492 items = response.get("Items", [])
494 return [
495 {
496 "session_id": item.get("session_id"),
497 "status": item.get("status"),
498 "created_at": item.get("created_at"),
499 "iteration_count": len(item.get("iterations", []) or []),
500 }
501 for item in _from_dynamodb_item(list(items))
502 ]
504 def delete_session(self, session_id: str) -> bool: # pragma: no cover - DynamoDB
505 """Delete the session via ``delete_item`` (idempotent).
507 Uses ``ReturnValues="ALL_OLD"`` so the call can distinguish a
508 successful deletion from a no-op on a missing key, matching the
509 :class:`FilesystemBackend` semantics where the return value
510 signals whether anything was actually removed.
511 """
512 table = self._get_table()
513 response = table.delete_item(
514 Key={"session_id": session_id},
515 ReturnValues="ALL_OLD",
516 )
517 return response.get("Attributes") is not None
520# ---------------------------------------------------------------------- #
521# resolver
522# ---------------------------------------------------------------------- #
524# Recognised values for the ``GCO_MISSION_STATE_BACKEND`` env var. Anything
525# outside this set normalises to ``"filesystem"`` — same fallback rule as
526# the ``GCO_MCP_TOOL_SEARCH`` precedent in ``gco_mcp/server.py``.
527_BACKEND_VALUES = frozenset({"filesystem", "dynamodb"})
529# Cached backend instance, populated on first call to ``get_backend()``.
530# ``GCO_MISSION_STATE_BACKEND`` is resolved once at first use and the
531# resulting instance is reused for every subsequent call. Env vars do not
532# change at runtime in practice, and a shared instance keeps the
533# ``FilesystemBackend._root_initialized`` cache hot across callers — the
534# same module-load resolution pattern used for ``GCO_MCP_TOOL_SEARCH`` in
535# ``gco_mcp/server.py``.
536_BACKEND_INSTANCE: MissionStateBackend | None = None
539def get_backend() -> MissionStateBackend:
540 """Return the configured Mission state backend, lazily constructed.
542 Reads ``GCO_MISSION_STATE_BACKEND`` on first call. Recognised values
543 are ``"filesystem"`` (default) and ``"dynamodb"``; any other value
544 logs a single warning naming the unrecognised input and falls back
545 to :class:`FilesystemBackend`, matching the unknown-value handling
546 for ``GCO_MCP_TOOL_SEARCH`` in ``gco_mcp/server.py``. The resolved
547 backend is cached at module scope so subsequent calls return the
548 same instance.
549 """
550 global _BACKEND_INSTANCE
551 if _BACKEND_INSTANCE is not None:
552 return _BACKEND_INSTANCE
554 raw = os.environ.get("GCO_MISSION_STATE_BACKEND", "filesystem").strip().lower()
555 if raw == "dynamodb":
556 _BACKEND_INSTANCE = DynamoDBBackend() # pragma: no cover - boto3 path
557 elif raw == "filesystem":
558 _BACKEND_INSTANCE = FilesystemBackend()
559 else:
560 logger.warning(
561 "Unrecognised GCO_MISSION_STATE_BACKEND value %r; falling back to filesystem",
562 raw,
563 )
564 _BACKEND_INSTANCE = FilesystemBackend()
565 return _BACKEND_INSTANCE