tylor-mcp 1.0.0

package/server/storage/json_store.py

@@ -0,0 +1,359 @@
+"""
+server/storage/json_store.py — Local JSON storage for agent101 threads.
+
+Zero-infra default: all state lives in ~/.tylor/threads.json.
+No database, no cloud account, no configuration required.
+All writes are atomic (write-to-tmp then os.replace).
+
+This is the default storage backend. DynamoDB is optional for multi-machine sync.
+"""
+from __future__ import annotations
+import json
+import logging
+import os
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+
+from mcp.server.fastmcp.exceptions import ToolError
+
+logger = logging.getLogger(__name__)
+
+STORE_VERSION = "1.0"
+WARN_THRESHOLD = 400 * 1024  # 400 KB
+
+_EMPTY_STORE: dict = {"version": STORE_VERSION, "threads": [], "current_thread_id": None}
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def _default_store_path() -> Path:
+    return Path.home() / ".tylor" / "threads.json"
+
+
+class JsonStore:
+    """
+    Local JSON storage implementing the same interface as DynamoClient
+    so tools can use either backend transparently via _get_db().
+
+    Thread data structure:
+      {id, name, status, created_at, updated_at, messages[], summary,
+       sandbox_roots[], current: bool, agent_states{}, agent_outputs[]}
+    """
+
+    def __init__(self, path: Path | None = None) -> None:
+        self.path = path or _default_store_path()
+
+    # ── Internal load/save ─────────────────────────────────────────────
+
+    def _load(self) -> dict:
+        import copy
+        if not self.path.exists():
+            return copy.deepcopy(_EMPTY_STORE)
+        try:
+            data = json.loads(self.path.read_text(encoding="utf-8"))
+            data.setdefault("threads", [])
+            data.setdefault("current_thread_id", None)
+            return data
+        except (json.JSONDecodeError, OSError) as exc:
+            logger.warning("Could not load %s: %s", self.path, exc)
+            import copy
+            return copy.deepcopy(_EMPTY_STORE)
+
+    # Public alias — tests call store.load()
+    def load(self) -> dict:
+        return self._load()
+
+    def _save(self, data: dict) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        raw = json.dumps(data, indent=2, ensure_ascii=False)
+        if len(raw.encode()) > WARN_THRESHOLD:
+            logger.warning(
+                "Thread store approaching file size limit (%d KB)",
+                len(raw.encode()) // 1024,
+            )
+        tmp = self.path.with_suffix(".tmp")
+        try:
+            tmp.write_text(raw, encoding="utf-8")
+            os.replace(tmp, self.path)
+        except OSError:
+            tmp.unlink(missing_ok=True)
+            raise
+
+    def _find(self, data: dict, thread_id: str) -> dict | None:
+        for t in data["threads"]:
+            if t["id"] == thread_id:
+                return t
+        return None
+
+    def _require(self, data: dict, thread_id: str) -> dict:
+        t = self._find(data, thread_id)
+        if t is None:
+            raise KeyError(f"Thread not found: {thread_id}")
+        return t
+
+    # ── Thread CRUD (high-level API used by JsonStore tests) ───────────
+
+    def new_thread(self, name: str) -> dict:
+        data = self._load()
+        now = _now_iso()
+        thread: dict = {
+            "id": f"thread_{uuid.uuid4().hex}",
+            "name": name, "status": "active",
+            "created_at": now, "updated_at": now,
+            "messages": [], "summary": None,
+            "sandbox_roots": [], "agent_states": {}, "agent_outputs": [],
+        }
+        data["threads"].append(thread)
+        self._save(data)
+        return thread
+
+    def get_thread(self, thread_id: str) -> dict | None:
+        return self._find(self._load(), thread_id)
+
+    def update_thread(self, thread_id: str, **fields) -> dict:
+        data = self._load()
+        t = self._require(data, thread_id)
+        fields["updated_at"] = _now_iso()
+        t.update(fields)
+        self._save(data)
+        return t
+
+    def delete_thread(self, thread_id: str) -> bool:
+        data = self._load()
+        before = len(data["threads"])
+        data["threads"] = [t for t in data["threads"] if t["id"] != thread_id]
+        if len(data["threads"]) == before:
+            return False
+        self._save(data)
+        return True
+
+    def list_threads(self) -> list:
+        data = self._load()
+        return sorted(data["threads"], key=lambda t: t.get("updated_at", ""), reverse=True)
+
+    # ── DynamoClient-compatible interface ──────────────────────────────
+    # All tool code calls these methods via _get_db(). Keeping the same
+    # signatures means zero changes needed in tools/*.py
+
+    def put_item(self, sk: str, attributes: dict) -> dict:
+        """Store a raw item by SK. Used for messages, summaries, events."""
+        data = self._load()
+        now = _now_iso()
+        item = dict(attributes)
+        item["SK"] = sk
+        item.setdefault("CreatedAt", now)
+        item["UpdatedAt"] = now
+
+        # Route by SK pattern
+        if "#META" in sk:
+            thread_id = sk.split("#")[1] if sk.startswith("THREAD#") else None
+            if thread_id and thread_id != "CURRENT":
+                existing = self._find(data, thread_id)
+                if existing:
+                    existing.update({
+                        "name": item.get("Name", existing.get("name", "")),
+                        "status": item.get("Status", existing.get("status", "active")).lower(),
+                        "updated_at": now,
+                        **({"project": item["Project"]} if item.get("Project") else {}),
+                    })
+                    if "Summary" in item:
+                        existing["summary"] = item["Summary"]
+                    self._save(data)
+                    return item
+                else:
+                    thread = {
+                        "id": thread_id, "name": item.get("Name", ""),
+                        "status": item.get("Status", "active").lower(),
+                        "project": item.get("Project", ""),
+                        "created_at": item.get("CreatedAt", now),
+                        "updated_at": now, "messages": [], "summary": None,
+                        "sandbox_roots": [], "agent_states": {}, "agent_outputs": [],
+                    }
+                    data["threads"].append(thread)
+                    self._save(data)
+                    return item
+            elif thread_id == "CURRENT":
+                data["current_thread_id"] = item.get("CurrentThreadId")
+                data["current_thread_active_at"] = item.get("ActiveAt", now)
+                self._save(data)
+                return item
+        elif "#MSG#" in sk or "#SUMMARY_FAILURE" in sk or "#RECOVERY" in sk or "#SANDBOX" in sk:
+            parts = sk.split("#")
+            if len(parts) >= 2:
+                thread_id = parts[1]
+                t = self._find(data, thread_id)
+                if t:
+                    t.setdefault("messages", [])
+                    item["SK"] = sk
+                    t["messages"].append(item)
+                    t["updated_at"] = now
+                    self._save(data)
+            return item
+        elif "#SUMMARY" in sk and "#SUMMARY_FAILURE" not in sk:
+            parts = sk.split("#")
+            thread_id = parts[1] if len(parts) >= 2 else None
+            if thread_id:
+                t = self._find(data, thread_id)
+                if t:
+                    t["summary"] = item.get("Summary", "")
+                    t["summary_type"] = item.get("SummaryType", "")
+                    t["updated_at"] = now
+                    self._save(data)
+            return item
+
+        # Fallback: store in misc bucket
+        data.setdefault("misc", [])
+        data["misc"].append(item)
+        self._save(data)
+        return item
+
+    def get_item(self, thread_id: str, sk: str) -> dict | None:
+        data = self._load()
+        t = self._find(data, thread_id)
+        if not t:
+            return None
+        for msg in t.get("messages", []):
+            if msg.get("SK") == sk:
+                return msg
+        return None
+
+    def query_all(self, sk_prefix: str) -> list:
+        """Return all items whose SK starts with sk_prefix."""
+        data = self._load()
+        results = []
+        # Thread META items
+        if sk_prefix.startswith("THREAD#"):
+            parts = sk_prefix.split("#")
+            if len(parts) >= 2 and parts[1]:
+                thread_id = parts[1]
+                t = self._find(data, thread_id)
+                if t:
+                    results.extend(self._thread_to_items(t))
+            else:
+                # All threads
+                for t in data["threads"]:
+                    results.extend(self._thread_to_items(t))
+        return results
+
+    def _thread_to_items(self, t: dict) -> list:
+        """Convert thread dict to DynamoDB-style item list."""
+        items = [{
+            "SK": f"THREAD#{t['id']}#META",
+            "Name": t.get("name", ""),
+            "Status": t.get("status", "active").capitalize(),
+            "LastActivity": t.get("updated_at", ""),
+            "MessageCount": len(t.get("messages", [])),
+            "CreatedAt": t.get("created_at", ""),
+            "UpdatedAt": t.get("updated_at", ""),
+            "Project": t.get("project", ""),
+        }]
+        for msg in t.get("messages", []):
+            items.append(msg)
+        return items
+
+    def query_thread(self, thread_id: str, sk_prefix: str) -> list:
+        data = self._load()
+        t = self._find(data, thread_id)
+        if not t:
+            return []
+        return [m for m in t.get("messages", []) if m.get("SK", "").startswith(sk_prefix)]
+
+    def get_thread_meta(self, thread_id: str) -> dict | None:
+        data = self._load()
+        t = self._find(data, thread_id)
+        if not t:
+            return None
+        return {
+            "SK": f"THREAD#{thread_id}#META",
+            "Name": t.get("name", ""),
+            "Status": t.get("status", "active").capitalize(),
+            "LastActivity": t.get("updated_at", ""),
+            "MessageCount": len(t.get("messages", [])),
+            "CreatedAt": t.get("created_at", ""),
+            "UpdatedAt": t.get("updated_at", ""),
+        }
+
+    def get_current_thread_marker(self) -> dict | None:
+        data = self._load()
+        tid = data.get("current_thread_id")
+        if not tid:
+            return None
+        return {
+            "CurrentThreadId": tid,
+            "ActiveAt": data.get("current_thread_active_at", ""),
+        }
+
+    def resolve_thread_id(self, thread_id: str | None = None) -> str:
+        if thread_id:
+            return thread_id
+        data = self._load()
+        tid = data.get("current_thread_id")
+        if not tid:
+            raise ToolError("No active thread — create one with CT [name]")
+        return tid
+
+    def switch_thread(self, thread_id: str) -> dict:
+        """Make thread_id the active thread. Atomic in JSON via single save."""
+        data = self._load()
+        t = self._find(data, thread_id)
+        if not t:
+            raise ToolError(f"Thread not found: {thread_id}")
+        old_id = data.get("current_thread_id")
+        if old_id and old_id != thread_id:
+            old = self._find(data, old_id)
+            if old:
+                old["updated_at"] = _now_iso()
+        data["current_thread_id"] = thread_id
+        data["current_thread_active_at"] = _now_iso()
+        t["updated_at"] = _now_iso()
+        self._save(data)
+        return {"status": "switched", "thread_id": thread_id, "switched_at": _now_iso()}
+
+    def set_sandbox_roots(self, thread_id: str, sandbox_roots: list) -> dict:
+        data = self._load()
+        t = self._require(data, thread_id)
+        t["sandbox_roots"] = sandbox_roots
+        t["updated_at"] = _now_iso()
+        self._save(data)
+        return {"thread_id": thread_id, "sandbox_roots": sandbox_roots}
+
+    # ── Agent state (stub — works for single-machine use) ──────────────
+
+    def put_agent_output(self, thread_id: str, agent_id: str, output: str, task: str | None = None) -> dict:
+        data = self._load()
+        t = self._require(data, thread_id)
+        now = _now_iso()
+        item = {"SK": f"THREAD#{thread_id}#AGENT#{agent_id}#OUT#{now}", "ThreadId": thread_id, "AgentId": agent_id, "Output": output, "Task": task}
+        t.setdefault("agent_outputs", []).append(item)
+        t["updated_at"] = now
+        self._save(data)
+        return item
+
+    def put_agent_handoff(self, thread_id: str, agent_id: str, handoff_state: dict) -> dict:
+        data = self._load()
+        t = self._require(data, thread_id)
+        now = _now_iso()
+        item = {"SK": f"THREAD#{thread_id}#AGENT#{agent_id}#HANDOFF#{now}", "ThreadId": thread_id, "AgentId": agent_id, "HandoffState": handoff_state}
+        t.setdefault("agent_outputs", []).append(item)
+        t["updated_at"] = now
+        self._save(data)
+        return item
+
+    def put_agent_state(self, thread_id: str, agent_id: str, state: dict) -> dict:
+        data = self._load()
+        t = self._require(data, thread_id)
+        t.setdefault("agent_states", {})[agent_id] = {**state, "UpdatedAt": _now_iso()}
+        t["updated_at"] = _now_iso()
+        self._save(data)
+        sk = f"THREAD#{thread_id}#AGENT#{agent_id}#STATE"
+        return {"SK": sk, "ThreadId": thread_id, "AgentId": agent_id, **state}
+
+    def query_agent_states(self, thread_id: str) -> list:
+        data = self._load()
+        t = self._find(data, thread_id)
+        if not t:
+            return []
+        return [{"AgentId": aid, **s} for aid, s in t.get("agent_states", {}).items()]

package/server/storage/opensearch.py

@@ -0,0 +1,194 @@
+"""
+server/storage/opensearch.py — OpenSearch vector client for agent101.
+
+Handles semantic memory: embedding via Amazon Titan Embeddings v2 (1536-dim)
+through AWS Bedrock, then k-NN search against the agent-memories index.
+
+Thread isolation is enforced at query time via a term filter on thread_id.
+Facts from other threads are NEVER returned.
+"""
+from __future__ import annotations
+import json
+import logging
+import uuid
+from datetime import datetime, timezone
+
+import boto3
+from mcp.server.fastmcp.exceptions import ToolError
+from opensearchpy import OpenSearch
+
+logger = logging.getLogger(__name__)
+
+INDEX = "agent-memories"
+TITAN_MODEL = "amazon.titan-embed-text-v2:0"
+VECTOR_DIM = 1536
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+class OpenSearchClient:
+    """
+    Semantic memory client backed by OpenSearch k-NN + Bedrock Titan Embeddings v2.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int = 9200,
+        bedrock_region: str = "us-east-1",
+        profile: str | None = None,
+    ) -> None:
+        self.host = host
+        self.port = port
+
+        self._os = OpenSearch(
+            hosts=[{"host": host, "port": port}],
+            http_compress=True,
+        )
+
+        session_kwargs: dict = {}
+        if profile:
+            session_kwargs["profile_name"] = profile
+        session = boto3.Session(**session_kwargs)
+        self._bedrock = session.client("bedrock-runtime", region_name=bedrock_region)
+
+    # ------------------------------------------------------------------
+    # Embedding helper
+    # ------------------------------------------------------------------
+
+    def _embed(self, text: str) -> list:
+        """Embed text using Titan Embeddings v2. Returns 1536-dim float list."""
+        try:
+            response = self._bedrock.invoke_model(
+                modelId=TITAN_MODEL,
+                body=json.dumps({"inputText": text}),
+                contentType="application/json",
+                accept="application/json",
+            )
+            embedding = json.loads(response["body"].read())["embedding"]
+        except Exception as exc:
+            raise ToolError(f"Bedrock embedding failed: {exc}") from exc
+
+        if len(embedding) != VECTOR_DIM:
+            raise ToolError(
+                f"Unexpected embedding dimension {len(embedding)} (expected {VECTOR_DIM})"
+            )
+        return embedding
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def index_memory(
+        self,
+        thread_id: str,
+        fact: str,
+        metadata: dict | None = None,
+    ) -> str:
+        """
+        Embed fact via Titan v2 and write to agent-memories index.
+
+        Args:
+            thread_id: Thread this fact belongs to (stored for isolation filtering).
+            fact: Plain-text fact to embed and store.
+            metadata: Optional extra fields stored alongside the fact.
+
+        Returns:
+            OpenSearch document ID.
+        """
+        embedding = self._embed(fact)
+        doc_id = uuid.uuid4().hex
+
+        doc = {
+            "thread_id": thread_id,
+            "content": fact,
+            "embedding": embedding,
+            "created_at": _now_iso(),
+        }
+        if metadata:
+            doc.update(metadata)
+
+        try:
+            self._os.index(index=INDEX, id=doc_id, body=doc, refresh=True)
+        except Exception as exc:
+            raise ToolError(f"OpenSearch index_memory failed: {exc}") from exc
+
+        logger.debug("index_memory: doc_id=%s thread=%s", doc_id, thread_id)
+        return doc_id
+
+    def search_memory(
+        self,
+        thread_id: str,
+        query: str,
+        k: int = 5,
+        type: str | None = None,
+    ) -> list:
+        """
+        k-NN search scoped to thread_id. Never returns other threads' facts.
+
+        Args:
+            thread_id: Only return facts belonging to this thread.
+            query: Natural-language query string to embed and search.
+            k: Max results to return.
+
+        Returns:
+            List of dicts with keys: id, content, thread_id, created_at, score.
+        """
+        query_vec = self._embed(query)
+
+        must = [
+            {
+                "knn": {
+                    "embedding": {
+                        "vector": query_vec,
+                        "k": k,
+                    }
+                }
+            },
+            {"term": {"thread_id": thread_id}},
+        ]
+        if type:
+            must.append({"term": {"type": type}})
+
+        os_query = {
+            "size": k,
+            "query": {
+                "bool": {
+                    "must": must
+                }
+            },
+            "_source": {"excludes": ["embedding"]},
+        }
+
+        try:
+            response = self._os.search(index=INDEX, body=os_query)
+        except Exception as exc:
+            raise ToolError(f"OpenSearch search_memory failed: {exc}") from exc
+
+        hits = response.get("hits", {}).get("hits", [])
+        results = []
+        for hit in hits:
+            src = hit.get("_source", {})
+            # Enforce isolation at result layer too (defence in depth)
+            if src.get("thread_id") != thread_id:
+                logger.warning(
+                    "search_memory: skipping result with wrong thread_id '%s' (expected '%s')",
+                    src.get("thread_id"),
+                    thread_id,
+                )
+                continue
+            results.append(
+                {
+                    "id": hit["_id"],
+                    "content": src.get("content", ""),
+                    "thread_id": src.get("thread_id"),
+                    "created_at": src.get("created_at"),
+                    "type": src.get("type"),
+                    "last_used_at": src.get("last_used_at"),
+                    "score": hit.get("_score"),
+                }
+            )
+
+        return results

package/server/storage/s3.py

@@ -0,0 +1,96 @@
+"""
+server/storage/s3.py — S3 blob storage client for agent101.
+
+Handles content >400KB that cannot fit in a single DynamoDB item.
+All blobs stored at: s3://{bucket}/{user_id}/threads/{thread_id}/{key}
+
+Thread isolation: the key path embeds thread_id — cross-thread access
+requires an explicit different thread_id, making accidental bleed impossible.
+"""
+from __future__ import annotations
+import logging
+from urllib.parse import urlparse
+
+import boto3
+from mcp.server.fastmcp.exceptions import ToolError
+
+logger = logging.getLogger(__name__)
+
+
+def _s3_path(user_id: str, thread_id: str, key: str) -> str:
+    return f"{user_id}/threads/{thread_id}/{key}"
+
+
+class S3Client:
+    """
+    Typed S3 client for agent101 blob storage.
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        user_id: str = "default",
+        profile: str | None = None,
+    ) -> None:
+        self.bucket = bucket
+        self.user_id = user_id
+
+        session_kwargs: dict = {}
+        if profile:
+            session_kwargs["profile_name"] = profile
+
+        session = boto3.Session(**session_kwargs)
+        self._s3 = session.client("s3")
+
+    def put_blob(self, thread_id: str, key: str, content: str | bytes) -> str:
+        """
+        Upload content to S3. Returns the s3:// URI.
+
+        Args:
+            thread_id: Thread this blob belongs to. Embedded in the S3 path.
+            key: Blob key within the thread (e.g. "summary", "msg_001").
+            content: String or bytes to store.
+
+        Returns:
+            s3://{bucket}/{user_id}/threads/{thread_id}/{key}
+        """
+        if isinstance(content, str):
+            content = content.encode("utf-8")
+
+        s3_key = _s3_path(self.user_id, thread_id, key)
+
+        try:
+            self._s3.put_object(
+                Bucket=self.bucket,
+                Key=s3_key,
+                Body=content,
+            )
+        except Exception as exc:
+            raise ToolError(f"S3 put_blob failed for '{key}': {exc}") from exc
+
+        uri = f"s3://{self.bucket}/{s3_key}"
+        logger.debug("put_blob: %s (%d bytes)", uri, len(content))
+        return uri
+
+    def get_blob(self, uri: str) -> bytes:
+        """
+        Download a blob by its s3:// URI. Returns raw bytes.
+
+        Args:
+            uri: s3://{bucket}/{key} URI returned by put_blob.
+
+        Returns:
+            Raw bytes content.
+        """
+        parsed = urlparse(uri)
+        if parsed.scheme != "s3":
+            raise ToolError(f"Invalid S3 URI scheme: '{uri}' (expected s3://...)")
+
+        bucket = parsed.netloc
+        key = parsed.path.lstrip("/")
+
+        try:
+            response = self._s3.get_object(Bucket=bucket, Key=key)
+            return response["Body"].read()
+        except Exception as exc:
+            raise ToolError(f"S3 get_blob failed for '{uri}': {exc}") from exc