Add OAuth key ring file support to MCP serve command

Add checkpoint store for run persistence and OAuth key ring support

Introduce InMemoryCheckpointStore and SQLiteCheckpointStore to save
and restore agent context after each tool call, enabling reliable
resume.  Add --checkpoint-store CLI flag for run and resume commands,
and wire checkpointing into AgentRunner.

Add OAuth key ring file support (--oauth-key-ring-file, --oauth-active-kid)
for MCP HTTP server, and make DPoP proof replay TTL configurable
via --oauth-dpop-replay-ttl.

Author: johnteee

teaagent/chat_agent.py

@@ -44,6 +44,7 @@ class ChatAgentConfig:
     stream: bool = False
     on_chunk: Optional[Callable[[str], None]] = None
     approval_handler: Optional[ApprovalHandler] = None
+    checkpoint_store: Any = None
 
     @classmethod
     def from_root(cls, root: str | Path, **kwargs: Any) -> 'ChatAgentConfig':
@@ -115,6 +116,7 @@ def run_chat_agent(
     task_spec: Optional[str] = None,
     depth: int = 0,
     initial_observations: Optional[list[dict[str, Any]]] = None,
+    initial_context_extra: Optional[dict[str, Any]] = None,
 ) -> RunResult:
     tool_registry = registry or build_workspace_tool_registry(config.root)
     if config.enable_subagent and depth < config.max_subagent_depth:
@@ -150,6 +152,7 @@ def run_chat_agent(
         ),
         approval_handler=config.approval_handler,
         compactor=ContextCompactor(memory_keys=('task_spec', 'memories')),
+        checkpoint_store=config.checkpoint_store,
     )
     run_id = uuid4().hex
     heartbeat: Optional[Heartbeat] = None
@@ -164,6 +167,7 @@ def run_chat_agent(
             decide=lambda context: engine.decide(with_memories(context, memories)),
             run_id=run_id,
             initial_observations=initial_observations,
+            initial_context_extra=initial_context_extra,
         )
     finally:
         if heartbeat is not None:

teaagent/checkpoint.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+_CHECKPOINT_KEYS = ('task', 'observations', 'compacted_summary', 'memory_keys')
+_SCHEMA_VERSION = 1
+
+
+class InMemoryCheckpointStore:
+    def __init__(self) -> None:
+        self._data: dict[str, dict[str, Any]] = {}
+
+    def save(self, run_id: str, context: dict[str, Any]) -> None:
+        self._data[run_id] = _extract_checkpoint(context)
+
+    def load(self, run_id: str) -> Optional[dict[str, Any]]:
+        return self._data.get(run_id)
+
+    def delete(self, run_id: str) -> None:
+        self._data.pop(run_id, None)
+
+
+class SQLiteCheckpointStore:
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._init_db()
+
+    def _connect(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(str(self._path), timeout=10)
+        conn.execute('PRAGMA journal_mode=WAL')
+        return conn
+
+    def _init_db(self) -> None:
+        with self._connect() as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS checkpoints (
+                    run_id TEXT PRIMARY KEY,
+                    context_json TEXT NOT NULL,
+                    updated_at TEXT NOT NULL,
+                    schema_version INTEGER NOT NULL DEFAULT 1
+                )
+                """
+            )
+
+    def save(self, run_id: str, context: dict[str, Any]) -> None:
+        snapshot = _extract_checkpoint(context)
+        now = datetime.now(timezone.utc).isoformat()
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO checkpoints (run_id, context_json, updated_at, schema_version)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(run_id) DO UPDATE SET
+                    context_json = excluded.context_json,
+                    updated_at = excluded.updated_at
+                """,
+                (
+                    run_id,
+                    json.dumps(snapshot, ensure_ascii=False),
+                    now,
+                    _SCHEMA_VERSION,
+                ),
+            )
+
+    def load(self, run_id: str) -> Optional[dict[str, Any]]:
+        with self._connect() as conn:
+            row = conn.execute(
+                'SELECT context_json FROM checkpoints WHERE run_id = ?',
+                (run_id,),
+            ).fetchone()
+        if row is None:
+            return None
+        return json.loads(row[0])
+
+    def delete(self, run_id: str) -> None:
+        with self._connect() as conn:
+            conn.execute('DELETE FROM checkpoints WHERE run_id = ?', (run_id,))
+
+
+def _extract_checkpoint(context: dict[str, Any]) -> dict[str, Any]:
+    return {k: context[k] for k in _CHECKPOINT_KEYS if k in context}

teaagent/cli/_agent_parsers.py

@@ -104,6 +104,12 @@ def _run(subs: argparse._SubParsersAction, handler: Callable) -> None:  # type:
         action='store_true',
         help='Also print OpenTelemetry spans to stderr (debug).',
     )
+    p.add_argument(
+        '--checkpoint-store',
+        default=None,
+        metavar='PATH',
+        help='SQLite path for run checkpoint storage. Saves context after each tool call.',
+    )
     p.set_defaults(func=handler)
 
 
@@ -203,6 +209,12 @@ def _resume(subs: argparse._SubParsersAction, handler: Callable) -> None:  # typ
         action='store_true',
         help='Re-run the original task from scratch instead of replaying observations from the prior run.',
     )
+    p.add_argument(
+        '--checkpoint-store',
+        default=None,
+        metavar='PATH',
+        help='SQLite path for checkpoint storage. Used to restore compacted context on resume.',
+    )
     p.set_defaults(func=handler)
 
 

teaagent/cli/_handlers/_agent.py

@@ -27,9 +27,23 @@ def agent_resume_command(args: argparse.Namespace) -> int:
         return 1
 
     initial_observations: list[dict[str, Any]] = []
+    initial_context_extra: Optional[dict[str, Any]] = None
     auto_approved: Optional[str] = None
+
     if not args.fresh_restart:
-        initial_observations = store.observations_for_run(args.run_id)
+        checkpoint_path = getattr(args, 'checkpoint_store', None)
+        checkpoint = None
+        if checkpoint_path:
+            from teaagent.checkpoint import SQLiteCheckpointStore
+
+            checkpoint = SQLiteCheckpointStore(checkpoint_path).load(args.run_id)
+        if checkpoint is not None:
+            initial_observations = checkpoint.get('observations', [])
+            initial_context_extra = {
+                k: v for k, v in checkpoint.items() if k not in ('task', 'observations')
+            }
+        else:
+            initial_observations = store.observations_for_run(args.run_id)
         pending = store.pending_approval_for_run(args.run_id)
         if pending and pending['call_id'] not in args.approve_call_id:
             args.approve_call_id = list(args.approve_call_id) + [pending['call_id']]
@@ -40,6 +54,7 @@ def agent_resume_command(args: argparse.Namespace) -> int:
         original_task,
         resumed_from=args.run_id,
         initial_observations=initial_observations,
+        initial_context_extra=initial_context_extra,
         auto_approved_call_id=auto_approved,
     )
 
@@ -50,6 +65,7 @@ def _execute_agent_task(
     *,
     resumed_from: Optional[str] = None,
     initial_observations: Optional[list[dict[str, Any]]] = None,
+    initial_context_extra: Optional[dict[str, Any]] = None,
     auto_approved_call_id: Optional[str] = None,
 ) -> int:
     task_spec = None
@@ -102,6 +118,12 @@ def _execute_agent_task(
             print(f'Telemetry setup failed: {exc}', file=sys.stderr)
 
     approval_handler = cli_approval_handler if args.hitl_approval else None
+    checkpoint_store = None
+    checkpoint_path = getattr(args, 'checkpoint_store', None)
+    if checkpoint_path:
+        from teaagent.checkpoint import SQLiteCheckpointStore
+
+        checkpoint_store = SQLiteCheckpointStore(checkpoint_path)
     result = run_chat_agent(
         task=task,
         adapter=adapter,
@@ -117,10 +139,12 @@ def _execute_agent_task(
             max_subagent_depth=args.max_subagent_depth,
             heartbeat_seconds=args.heartbeat,
             approval_handler=approval_handler,
+            checkpoint_store=checkpoint_store,
         ),
         audit=audit,
         task_spec=task_spec,
         initial_observations=initial_observations,
+        initial_context_extra=initial_context_extra,
     )
     store.logger_for_result(result, audit)
     if _telemetry_sink is not None:

teaagent/cli/_handlers/_mcp.py

@@ -1,24 +1,42 @@
 from __future__ import annotations
 
 import argparse
+import json
 import sys
+from pathlib import Path
 
 from teaagent.mcp_http import is_loopback_host
 from teaagent.mcp_server import serve_mcp_stdio
+from teaagent.oauth21 import OAuth21AuthorizationServer, OAuthKeyRing
 from teaagent.workspace_tools import build_workspace_tool_registry
 
 
 def mcp_serve_command(args: argparse.Namespace) -> int:
-    from teaagent.oauth21 import OAuth21AuthorizationServer
 
     registry = build_workspace_tool_registry(args.root)
     if args.http:
         oauth_server = None
+        key_ring = None
+        if args.oauth_key_ring_file:
+            key_ring, error = _load_key_ring(
+                Path(args.oauth_key_ring_file),
+                active_kid_override=args.oauth_active_kid,
+            )
+            if error:
+                print(error, file=sys.stderr)
+                return 1
+        elif args.oauth_active_kid:
+            print(
+                '--oauth-active-kid requires --oauth-key-ring-file.',
+                file=sys.stderr,
+            )
+            return 1
         if args.oauth_issuer and args.oauth_signing_key:
             oauth_server = OAuth21AuthorizationServer(
                 signing_key=args.oauth_signing_key,
                 issuer=args.oauth_issuer,
                 token_ttl=args.oauth_token_ttl,
+                dpop_replay_ttl=args.oauth_dpop_replay_ttl,
             )
             for spec in args.oauth_client or []:
                 parts = spec.split(':', 2)
@@ -54,3 +72,41 @@ def mcp_serve_command(args: argparse.Namespace) -> int:
             oauth_server=oauth_server,
         )
     return serve_mcp_stdio(registry)
+
+
+def _load_key_ring(
+    path: Path, *, active_kid_override: str | None
+) -> tuple[OAuthKeyRing | None, str | None]:
+    if not path.exists():
+        return None, f'OAuth key ring file not found: {path}'
+    try:
+        payload = json.loads(path.read_text(encoding='utf-8'))
+    except (OSError, json.JSONDecodeError) as exc:
+        return None, f'Failed to read OAuth key ring file {path}: {exc}'
+    if not isinstance(payload, dict):
+        return None, 'OAuth key ring file must contain a JSON object.'
+    raw_keys = payload.get('keys')
+    if not isinstance(raw_keys, dict) or not raw_keys:
+        return None, "OAuth key ring file requires non-empty 'keys' object."
+
+    keys: dict[str, bytes] = {}
+    for kid, secret in raw_keys.items():
+        if not isinstance(kid, str) or not kid:
+            return None, 'OAuth key ring keys must use non-empty string kids.'
+        if not isinstance(secret, str) or len(secret) < 16:
+            return (
+                None,
+                f'OAuth key ring secret for kid {kid!r} must be a string >= 16 chars.',
+            )
+        keys[kid] = secret.encode('utf-8')
+
+    active_kid = active_kid_override or payload.get('active_kid')
+    if not isinstance(active_kid, str) or not active_kid:
+        return (
+            None,
+            "OAuth key ring file requires string 'active_kid' or --oauth-active-kid.",
+        )
+    if active_kid not in keys:
+        return None, f'OAuth active kid {active_kid!r} not found in key ring keys.'
+
+    return OAuthKeyRing(active_kid=active_kid, keys=keys), None

teaagent/cli/_mcp_parsers.py

@@ -57,6 +57,16 @@ def register(
         default=None,
         help='HMAC signing key for JWT access tokens (min 16 chars).',
     )
+    serve.add_argument(
+        '--oauth-key-ring-file',
+        default=None,
+        help='Path to JSON key ring file: {"active_kid":"...","keys":{"kid":"secret"}}.',
+    )
+    serve.add_argument(
+        '--oauth-active-kid',
+        default=None,
+        help='Override active key id (kid) from the key ring file.',
+    )
     serve.add_argument(
         '--oauth-client',
         action='append',
@@ -71,4 +81,10 @@ def register(
         default=3600,
         help='Access token TTL in seconds. Default 3600.',
     )
+    serve.add_argument(
+        '--oauth-dpop-replay-ttl',
+        type=int,
+        default=60,
+        help='DPoP proof replay-cache TTL in seconds. Default 60.',
+    )
     serve.set_defaults(func=handlers['serve'])

teaagent/oauth21/_server.py

@@ -43,6 +43,7 @@ def __init__(
         *,
         token_ttl: int = _DEFAULT_ACCESS_TOKEN_TTL,
         nonce_ttl: int = _NONCE_TTL_SECONDS,
+        dpop_replay_ttl: int = _PROOF_MAX_AGE_SECONDS,
         store: Optional[OAuthStore] = None,
         key_ring: Optional[OAuthKeyRing] = None,
     ) -> None:
@@ -53,6 +54,7 @@ def __init__(
         self._issuer = issuer
         self._token_ttl = token_ttl
         self._nonce_ttl = nonce_ttl
+        self._dpop_replay_ttl = dpop_replay_ttl
         self._store = store or InMemoryOAuthStore()
         self._dpop_replay_cache = DPoPReplayCache()
 
@@ -269,10 +271,10 @@ def _validate_dpop_and_extract_jkt(self, proof_jwt: str) -> str:
         _verify_dpop_signature(proof_jwt, jwk)
         iat = payload.get('iat', 0)
         now = time.time()
-        if abs(now - iat) > _PROOF_MAX_AGE_SECONDS:
+        if abs(now - iat) > self._dpop_replay_ttl:
             raise InvalidDPoPError('DPoP proof is too old')
         self._dpop_replay_cache.remember_once(
-            payload.get('jti'), iat=float(iat), now=now, ttl=_PROOF_MAX_AGE_SECONDS
+            payload.get('jti'), iat=float(iat), now=now, ttl=self._dpop_replay_ttl
         )
         return compute_jwk_thumbprint(jwk)