Coverage for gco_mcp / mission / memory.py: 100.00%
166 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Mission-memory store over the DynamoDB vector index.
3:class:`MissionMemoryStore` is the runtime client for the
4``{project}-mission-memory`` table provisioned by
5``gco/stacks/global_stack.py`` when ``mission_memory.enabled`` is set:
6one memory item per completed Mission session, searchable by directive
7similarity through the ``directive-embedding-index`` vector index.
9Three operations:
11* :meth:`MissionMemoryStore.write_memory` — embed the session's
12 directive and ``PutItem`` the memory record (verdict, lessons,
13 followups, provenance).
14* :meth:`MissionMemoryStore.search_similar` — embed a query directive
15 and ``SearchVectors`` the index for the closest past missions.
16* :meth:`MissionMemoryStore.list_memories` — ``Scan`` the base table
17 for summaries (no embedding involved), newest completion first.
19Resolution conventions, copied from their precedents:
21* **Region** follows the store convention documented in
22 ``gco/services/template_store.py``: ``DYNAMODB_REGION`` →
23 ``GLOBAL_REGION`` → ``AWS_REGION``, else the SDK default chain.
24* **Table / index names** resolve lazily from SSM
25 (``/{project}/mission-memory-table-name`` and
26 ``/{project}/mission-memory-index-name``) on first use and are cached
27 on the instance — the same shape as
28 :meth:`mcp.mission.state.DynamoDBBackend._resolve_table_name`. The
29 SSM lookup goes through :func:`gco.services.aws_ssm.get_ssm_parameter`
30 because ``gco_mcp/`` must not import ``cli/``.
32Request-shape gotcha, verified against the botocore service model:
33``SearchVectors``' ``SearchVector`` parameter is a **plain list** of
34``{"N": "..."}`` attribute values. The ``{"L": [...]}`` wrapper is only
35for writing ``directive_embedding`` on an item.
37Failure contract: infrastructure that is absent or not yet queryable —
38table or index missing, index still backfilling, SSM parameter never
39published because the feature is disabled, no credentials, endpoint
40unreachable — raises :class:`MissionMemoryUnavailableError`; everything
41else raises :class:`MissionMemoryError`. ``SearchVectors`` errors while
42the index is backfilling, so that case is deliberately part of
43"unavailable", not a hard failure. Callers on the engine path swallow
44both (memory is best-effort and must never fail a mission); the CLI
45surfaces the message.
47The runtime defaults below mirror the ``mission_memory`` block in
48``cdk.json`` — the block drives the *deployed index*, these drive the
49*vectors sent to it*, and ``tests/test_mission_memory_runtime.py``
50asserts the two stay in agreement. ``dimensions`` in particular is a
51one-way door: immutable after index creation, and query vectors must
52come from the same model at the same width.
53"""
55from __future__ import annotations
57import logging
58import os
59import time
60from collections.abc import Mapping
61from datetime import UTC, datetime
62from typing import Any
64from gco.bedrock import get_default_embedding_model_id
66from .embeddings import embed_text
68logger = logging.getLogger(__name__)
70__all__ = [
71 "DEFAULT_DIMENSIONS",
72 "DEFAULT_RETENTION_DAYS",
73 "DEFAULT_TOP_K",
74 "MEMORY_SCHEMA_VERSION",
75 "MissionMemoryError",
76 "MissionMemoryStore",
77 "MissionMemoryUnavailableError",
78]
80#: Memory-item schema version, deliberately independent of the mission
81#: session ``SCHEMA_VERSION`` — the two payloads evolve separately.
82MEMORY_SCHEMA_VERSION = "1"
84#: Runtime mirrors of the ``mission_memory`` defaults in ``cdk.json``.
85DEFAULT_DIMENSIONS = 1024
86DEFAULT_TOP_K = 3
87DEFAULT_RETENTION_DAYS = 365
89#: Region resolution order — the convention shared by every DynamoDB
90#: store in the tree (see ``gco/services/template_store.py``).
91_REGION_ENV_ORDER = ("DYNAMODB_REGION", "GLOBAL_REGION", "AWS_REGION")
93_TABLE_NAME_PARAM_SUFFIX = "mission-memory-table-name"
94_INDEX_NAME_PARAM_SUFFIX = "mission-memory-index-name"
96_UNAVAILABLE_HINT = (
97 "mission memory is unavailable — the table or index may not be "
98 "provisioned (deploy with mission_memory.enabled: true), or the "
99 "vector index may still be backfilling"
100)
103class MissionMemoryError(RuntimeError):
104 """A mission-memory operation failed."""
107class MissionMemoryUnavailableError(MissionMemoryError):
108 """The mission-memory infrastructure is absent or not yet queryable.
110 Raised for the degradable cases: table/index not provisioned, SSM
111 name parameters never published, the vector index still backfilling
112 after creation, missing credentials, or an unreachable endpoint.
113 Engine callers treat this as "no prior context"; the CLI shows the
114 message so an operator can tell "nothing similar" from "not there".
115 """
118def _resolve_region() -> str | None:
119 """Return the first configured region env var, else ``None``.
121 ``None`` lets boto3 fall through to its own default chain, matching
122 how the sibling stores behave on hosts with a configured profile.
123 """
124 for env_var in _REGION_ENV_ORDER:
125 value = os.environ.get(env_var)
126 if value:
127 return value
128 return None
131def _plain(value: Any) -> Any:
132 """Recursively convert boto3-deserialized values to JSON-friendly types.
134 ``TypeDeserializer`` yields :class:`decimal.Decimal` for every
135 ``N``; integral values become ``int`` and the rest ``float`` so the
136 CLI and prompt builders can ``json.dumps`` results directly.
137 """
138 from decimal import Decimal
140 if isinstance(value, Decimal):
141 return int(value) if value == value.to_integral_value() else float(value)
142 if isinstance(value, list):
143 return [_plain(entry) for entry in value]
144 if isinstance(value, dict):
145 return {key: _plain(entry) for key, entry in value.items()}
146 return value
149def _number_attr(value: float) -> dict[str, str]:
150 """Render one vector component as a DynamoDB number attribute value."""
151 return {"N": repr(float(value))}
154class MissionMemoryStore:
155 """Runtime client for the mission-memory table and its vector index.
157 Constructor arguments exist for tests and for callers that already
158 know the deployed names; with no arguments every name resolves
159 lazily from SSM on first use, so constructing a store on a host
160 without AWS credentials is free.
162 Args:
163 table_name: Table name override; ``None`` resolves from SSM.
164 index_name: Vector-index name override; ``None`` resolves from SSM.
165 region: Region override; ``None`` follows the
166 ``DYNAMODB_REGION`` → ``GLOBAL_REGION`` → ``AWS_REGION``
167 convention, then the SDK default chain.
168 embedding_model_id: Embedding model override; ``None`` resolves
169 the checked-in default lazily on first embed.
170 dimensions: Requested embedding width. Must match the deployed
171 index width (one-way door); the default mirrors the
172 ``mission_memory.dimensions`` default in ``cdk.json``.
173 retention_days: TTL window written on new memory items; the
174 default mirrors ``mission_memory.retention_days``.
175 """
177 def __init__(
178 self,
179 table_name: str | None = None,
180 index_name: str | None = None,
181 *,
182 region: str | None = None,
183 embedding_model_id: str | None = None,
184 dimensions: int = DEFAULT_DIMENSIONS,
185 retention_days: int = DEFAULT_RETENTION_DAYS,
186 ) -> None:
187 self._table_name = table_name
188 self._index_name = index_name
189 self._region = region
190 self._embedding_model_id = embedding_model_id
191 self._dimensions = int(dimensions)
192 self._retention_days = int(retention_days)
193 self._client: Any = None
195 # ------------------------------------------------------------------ #
196 # internals
197 # ------------------------------------------------------------------ #
199 def _resolve_ssm_name(self, suffix: str) -> str:
200 """Fetch ``/{project}/{suffix}`` from SSM, mapping absence to unavailable."""
201 from botocore.exceptions import BotoCoreError, ClientError
203 from gco.services.aws_ssm import get_ssm_parameter
205 project_name = os.environ.get("GCO_PROJECT_NAME", "gco")
206 param_name = f"/{project_name}/{suffix}"
207 try:
208 return get_ssm_parameter(param_name, region=self._region or _resolve_region())
209 except ClientError as err:
210 code = err.response.get("Error", {}).get("Code")
211 if code == "ParameterNotFound":
212 raise MissionMemoryUnavailableError(
213 f"SSM parameter {param_name} not found — {_UNAVAILABLE_HINT}"
214 ) from err
215 raise MissionMemoryError(f"SSM lookup for {param_name} failed: {err}") from err
216 except BotoCoreError as err:
217 # Covers NoCredentialsError and endpoint-unreachable faults.
218 raise MissionMemoryUnavailableError(
219 f"SSM unreachable resolving {param_name} — {_UNAVAILABLE_HINT}"
220 ) from err
222 def _resolve_table_name(self) -> str:
223 """Return the cached table name, fetching from SSM on first call."""
224 if self._table_name is None:
225 self._table_name = self._resolve_ssm_name(_TABLE_NAME_PARAM_SUFFIX)
226 return self._table_name
228 def _resolve_index_name(self) -> str:
229 """Return the cached index name, fetching from SSM on first call."""
230 if self._index_name is None:
231 self._index_name = self._resolve_ssm_name(_INDEX_NAME_PARAM_SUFFIX)
232 return self._index_name
234 def _get_client(self) -> Any:
235 """Return the cached low-level ``dynamodb`` client, building it lazily."""
236 if self._client is None:
237 import boto3
239 self._client = boto3.client("dynamodb", region_name=self._region or _resolve_region())
240 return self._client
242 def _resolved_embedding_model_id(self) -> str:
243 """Return the embedding model id, resolving the default lazily."""
244 if self._embedding_model_id is None:
245 self._embedding_model_id = get_default_embedding_model_id()
246 return self._embedding_model_id
248 def _embed(self, text: str) -> list[float]:
249 """Embed ``text`` at the configured width and guard the one-way door."""
250 vector = embed_text(
251 text,
252 model_id=self._resolved_embedding_model_id(),
253 dimensions=self._dimensions,
254 )
255 if len(vector) != self._dimensions:
256 raise MissionMemoryError(
257 f"embedding width {len(vector)} does not match the configured "
258 f"index width {self._dimensions}; the model "
259 f"{self._resolved_embedding_model_id()!r} ignored the requested "
260 "dimensions — align mission_memory.dimensions with the model's "
261 "output width (both are one-way doors on the deployed index)"
262 )
263 return vector
265 # ------------------------------------------------------------------ #
266 # operations
267 # ------------------------------------------------------------------ #
269 def write_memory(
270 self,
271 session: Mapping[str, Any],
272 verdict: str,
273 reason: str,
274 lessons: str,
275 followups: list[str],
276 ) -> None:
277 """Embed the session's directive and persist one memory item.
279 Args:
280 session: The terminal session payload — canonically a
281 :class:`mcp.mission.types.SessionState`, but any mapping
282 carrying ``session_id`` / ``directive_text`` (plus the
283 optional ``criteria`` / ``tool_allowlist`` /
284 ``iterations`` / ``created_at`` / ``ended_at`` fields)
285 works, which is what lets the backfill command replay
286 Final_Reports.
287 verdict: Terminal verdict label (``complete`` | ``terminate``).
288 reason: Terminal verdict reason.
289 lessons: The report's lessons paragraph — the sampled
290 overlay when one was produced, else the templated text.
291 followups: The report's recommended followups.
293 Raises:
294 MissionMemoryUnavailableError: Table absent / feature not
295 provisioned / no credentials / endpoint unreachable.
296 MissionMemoryError: Any other write failure.
297 EmbeddingError: Propagated from the embedding call; callers
298 on the engine path swallow it like the rest.
299 """
300 from botocore.exceptions import (
301 BotoCoreError,
302 ClientError,
303 NoCredentialsError,
304 PartialCredentialsError,
305 )
307 # Resolve the table name before embedding: on a deployment without
308 # the feature the SSM lookup raises MissionMemoryUnavailableError
309 # cheaply, so a disabled stack never pays a Bedrock embedding call
310 # just to discover there is nowhere to write.
311 table_name = self._resolve_table_name()
313 directive = str(session.get("directive_text") or "")
314 vector = self._embed(directive)
316 completed_at = str(session.get("ended_at") or datetime.now(UTC).isoformat())
317 criteria = session.get("criteria") or []
318 item: dict[str, Any] = {
319 "session_id": {"S": str(session["session_id"])},
320 "directive": {"S": directive},
321 "directive_embedding": {"L": [_number_attr(v) for v in vector]},
322 "lessons": {"S": str(lessons)},
323 "recommended_followups": {"L": [{"S": str(f)} for f in followups]},
324 "final_verdict": {"S": str(verdict)},
325 "verdict_reason": {"S": str(reason)},
326 "iteration_count": {"N": str(len(session.get("iterations") or []))},
327 "criteria_summary": {"L": [{"S": str(c.get("criterion_id", ""))} for c in criteria]},
328 "tool_allowlist": {"L": [{"S": str(t)} for t in session.get("tool_allowlist") or []]},
329 "created_at": {"S": str(session.get("created_at") or "")},
330 "completed_at": {"S": completed_at},
331 "embedding_model_id": {"S": self._resolved_embedding_model_id()},
332 "embedding_dimensions": {"N": str(len(vector))},
333 "schema_version": {"S": MEMORY_SCHEMA_VERSION},
334 "ttl": {"N": str(int(time.time()) + self._retention_days * 86400)},
335 }
337 try:
338 self._get_client().put_item(TableName=table_name, Item=item)
339 except (NoCredentialsError, PartialCredentialsError) as err:
340 raise MissionMemoryUnavailableError(
341 f"no AWS credentials — {_UNAVAILABLE_HINT}"
342 ) from err
343 except ClientError as err:
344 code = err.response.get("Error", {}).get("Code")
345 if code == "ResourceNotFoundException":
346 raise MissionMemoryUnavailableError(
347 f"table {self._table_name!r} not found — {_UNAVAILABLE_HINT}"
348 ) from err
349 raise MissionMemoryError(f"mission-memory write failed: {err}") from err
350 except BotoCoreError as err:
351 raise MissionMemoryUnavailableError(
352 f"DynamoDB unreachable — {_UNAVAILABLE_HINT}"
353 ) from err
355 def search_similar(
356 self,
357 directive: str,
358 top_k: int = DEFAULT_TOP_K,
359 final_verdict: str | None = None,
360 ) -> list[dict[str, Any]]:
361 """Return the closest past missions to ``directive``.
363 Args:
364 directive: Query text; embedded with the same model and
365 width the index was built for.
366 top_k: Number of results to request. The default mirrors
367 ``mission_memory.top_k`` in ``cdk.json``.
368 final_verdict: When set, an inline equality filter on the
369 memory item's ``final_verdict`` (``complete`` |
370 ``terminate``) — the one attribute the index declares as
371 an ``INLINE_FILTER``.
373 Returns:
374 A list of plain dicts — the index's ``INCLUDE`` projection
375 (``directive``, ``lessons``, ``recommended_followups``,
376 ``final_verdict``, ``verdict_reason``, ``iteration_count``,
377 ``completed_at``) plus ``session_id`` and a float ``score``
378 (the distance-function score reported by DynamoDB).
380 Raises:
381 MissionMemoryUnavailableError: Table/index absent, index
382 still backfilling, no credentials, or endpoint
383 unreachable.
384 MissionMemoryError: ``top_k`` < 1 or any other failure.
385 EmbeddingError: Propagated from the embedding call.
386 """
387 from botocore.exceptions import (
388 BotoCoreError,
389 ClientError,
390 NoCredentialsError,
391 PartialCredentialsError,
392 )
394 if int(top_k) < 1:
395 raise MissionMemoryError(f"top_k must be a positive integer, got {top_k!r}")
397 # Resolve names before embedding — same rationale as write_memory:
398 # a deployment without the feature fails fast on the SSM lookup
399 # instead of paying an embedding call first.
400 table_name = self._resolve_table_name()
401 index_name = self._resolve_index_name()
403 vector = self._embed(directive)
405 # Request-shape gotcha: SearchVector is a plain list of number
406 # attribute values — the {"L": ...} wrapper is a write-side shape.
407 request: dict[str, Any] = {
408 "TableName": table_name,
409 "IndexName": index_name,
410 "SearchVector": [_number_attr(v) for v in vector],
411 "TopK": int(top_k),
412 }
413 if final_verdict is not None:
414 request["SearchConditionExpression"] = "final_verdict = :final_verdict"
415 request["ExpressionAttributeValues"] = {":final_verdict": {"S": str(final_verdict)}}
417 try:
418 response = self._get_client().search_vectors(**request)
419 except (NoCredentialsError, PartialCredentialsError) as err:
420 raise MissionMemoryUnavailableError(
421 f"no AWS credentials — {_UNAVAILABLE_HINT}"
422 ) from err
423 except ClientError as err:
424 code = err.response.get("Error", {}).get("Code")
425 if code in ("ResourceNotFoundException", "ValidationException"):
426 # ValidationException is what SearchVectors answers while
427 # the index is still backfilling after creation.
428 raise MissionMemoryUnavailableError(
429 f"SearchVectors failed ({code}) — {_UNAVAILABLE_HINT}: {err}"
430 ) from err
431 raise MissionMemoryError(f"mission-memory search failed: {err}") from err
432 except BotoCoreError as err:
433 raise MissionMemoryUnavailableError(
434 f"DynamoDB unreachable — {_UNAVAILABLE_HINT}"
435 ) from err
437 from boto3.dynamodb.types import TypeDeserializer
439 deserializer = TypeDeserializer()
440 results: list[dict[str, Any]] = []
441 for entry in response.get("SearchResults", []):
442 item = {
443 key: _plain(deserializer.deserialize(value))
444 for key, value in (entry.get("Item") or {}).items()
445 }
446 # Defensive: the INCLUDE projection excludes the vector, but a
447 # recreated index with a different projection must not bloat
448 # prompts or CLI output with a thousand floats.
449 item.pop("directive_embedding", None)
450 score = entry.get("Score")
451 if score is not None:
452 item["score"] = float(score)
453 results.append(item)
454 return results
456 def list_memories(self, limit: int = 50) -> list[dict[str, Any]]:
457 """Return memory-item summaries, most recently completed first.
459 Backs ``gco mission memory list``: a paginated ``Scan`` over the
460 base table (the vector index's ``SearchVectors`` is a similarity
461 query, not a listing API) projecting only the summary fields —
462 the embedding vector never leaves DynamoDB. The table holds one
463 small item per completed mission, so a full scan is cheap at
464 this scale.
466 Args:
467 limit: Maximum summaries to return after sorting by
468 ``completed_at`` descending. Must be positive.
470 Raises:
471 MissionMemoryUnavailableError: Table absent / feature not
472 provisioned / no credentials / endpoint unreachable.
473 MissionMemoryError: ``limit`` < 1 or any other failure.
474 """
475 from botocore.exceptions import (
476 BotoCoreError,
477 ClientError,
478 NoCredentialsError,
479 PartialCredentialsError,
480 )
482 if int(limit) < 1:
483 raise MissionMemoryError(f"limit must be a positive integer, got {limit!r}")
485 table_name = self._resolve_table_name()
486 projection = (
487 "session_id, directive, final_verdict, verdict_reason, "
488 "iteration_count, created_at, completed_at, embedding_model_id"
489 )
491 from boto3.dynamodb.types import TypeDeserializer
493 deserializer = TypeDeserializer()
494 items: list[dict[str, Any]] = []
495 exclusive_start_key: dict[str, Any] | None = None
496 try:
497 while True:
498 request: dict[str, Any] = {
499 "TableName": table_name,
500 "ProjectionExpression": projection,
501 }
502 if exclusive_start_key is not None:
503 request["ExclusiveStartKey"] = exclusive_start_key
504 response = self._get_client().scan(**request)
505 items.extend(
506 {key: _plain(deserializer.deserialize(value)) for key, value in raw.items()}
507 for raw in response.get("Items", [])
508 )
509 exclusive_start_key = response.get("LastEvaluatedKey")
510 if not exclusive_start_key:
511 break
512 except (NoCredentialsError, PartialCredentialsError) as err:
513 raise MissionMemoryUnavailableError(
514 f"no AWS credentials — {_UNAVAILABLE_HINT}"
515 ) from err
516 except ClientError as err:
517 code = err.response.get("Error", {}).get("Code")
518 if code == "ResourceNotFoundException":
519 raise MissionMemoryUnavailableError(
520 f"table {self._table_name!r} not found — {_UNAVAILABLE_HINT}"
521 ) from err
522 raise MissionMemoryError(f"mission-memory list failed: {err}") from err
523 except BotoCoreError as err:
524 raise MissionMemoryUnavailableError(
525 f"DynamoDB unreachable — {_UNAVAILABLE_HINT}"
526 ) from err
528 items.sort(key=lambda item: str(item.get("completed_at") or ""), reverse=True)
529 return items[: int(limit)]