contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/a2a/models.py

@@ -0,0 +1,246 @@
+"""
+A2A Protocol Data Models
+=========================
+Core dataclasses matching the Google A2A specification.
+
+Spec: https://google.github.io/A2A/
+"""
+
+from __future__ import annotations
+
+import secrets
+from dataclasses import dataclass, field, asdict
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+
+# ── Task States ───────────────────────────────────────────────────────
+
+class TaskState(str, Enum):
+    """A2A task lifecycle states."""
+    SUBMITTED = "submitted"
+    WORKING = "working"
+    INPUT_REQUIRED = "input-required"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+# Valid transitions enforced by TaskManager
+VALID_TRANSITIONS = {
+    TaskState.SUBMITTED: {TaskState.WORKING, TaskState.CANCELED},
+    TaskState.WORKING: {TaskState.COMPLETED, TaskState.FAILED, TaskState.INPUT_REQUIRED, TaskState.CANCELED},
+    TaskState.INPUT_REQUIRED: {TaskState.WORKING, TaskState.CANCELED},
+    # Terminal states — no outbound transitions
+    TaskState.COMPLETED: set(),
+    TaskState.FAILED: set(),
+    TaskState.CANCELED: set(),
+}
+
+
+# ── Message Parts ─────────────────────────────────────────────────────
+
+@dataclass
+class TextPart:
+    """Plain text content."""
+    text: str
+    type: str = "text"
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"type": self.type, "text": self.text}
+
+
+@dataclass
+class FilePart:
+    """File content — either inline bytes (base64) or a URI reference."""
+    name: str
+    mime_type: str
+    uri: Optional[str] = None
+    bytes_b64: Optional[str] = None
+    type: str = "file"
+
+    def to_dict(self) -> Dict[str, Any]:
+        f = {"name": self.name, "mimeType": self.mime_type}
+        if self.uri:
+            f["uri"] = self.uri
+        if self.bytes_b64:
+            f["bytes"] = self.bytes_b64
+        return {"type": self.type, "file": f}
+
+
+@dataclass
+class DataPart:
+    """Structured JSON data."""
+    data: Dict[str, Any] = field(default_factory=dict)
+    type: str = "data"
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"type": self.type, "data": self.data}
+
+
+Part = Union[TextPart, FilePart, DataPart]
+
+
+def part_from_dict(d: Dict[str, Any]) -> Part:
+    """Deserialize a Part from a dict."""
+    t = d.get("type", "text")
+    if t == "text":
+        return TextPart(text=d.get("text", ""))
+    elif t == "file":
+        f = d.get("file", {})
+        return FilePart(
+            name=f.get("name", ""),
+            mime_type=f.get("mimeType", "application/octet-stream"),
+            uri=f.get("uri"),
+            bytes_b64=f.get("bytes"),
+        )
+    elif t == "data":
+        return DataPart(data=d.get("data", {}))
+    return TextPart(text=str(d))
+
+
+# ── Message ───────────────────────────────────────────────────────────
+
+@dataclass
+class Message:
+    """A2A message — a sequence of parts sent by a user or agent."""
+    role: str  # "user" or "agent"
+    parts: List[Part] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "role": self.role,
+            "parts": [p.to_dict() for p in self.parts],
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Message:
+        return cls(
+            role=d.get("role", "user"),
+            parts=[part_from_dict(p) for p in d.get("parts", [])],
+            metadata=d.get("metadata", {}),
+        )
+
+    @classmethod
+    def text(cls, content: str, role: str = "user") -> Message:
+        """Convenience: create a simple text message."""
+        return cls(role=role, parts=[TextPart(text=content)])
+
+    def get_text(self) -> str:
+        """Extract concatenated text from all TextParts."""
+        return "\n".join(p.text for p in self.parts if isinstance(p, TextPart))
+
+
+# ── Artifact ──────────────────────────────────────────────────────────
+
+@dataclass
+class Artifact:
+    """A2A artifact — a named output produced by a task."""
+    name: str
+    parts: List[Part] = field(default_factory=list)
+    description: str = ""
+    index: int = 0
+    append: bool = False
+    last_chunk: bool = True
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        d = {
+            "name": self.name,
+            "parts": [p.to_dict() for p in self.parts],
+            "index": self.index,
+        }
+        if self.description:
+            d["description"] = self.description
+        if self.append:
+            d["append"] = True
+        if not self.last_chunk:
+            d["lastChunk"] = False
+        if self.metadata:
+            d["metadata"] = self.metadata
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Artifact:
+        return cls(
+            name=d.get("name", ""),
+            parts=[part_from_dict(p) for p in d.get("parts", [])],
+            description=d.get("description", ""),
+            index=d.get("index", 0),
+            append=d.get("append", False),
+            last_chunk=d.get("lastChunk", True),
+            metadata=d.get("metadata", {}),
+        )
+
+
+# ── Task Status ───────────────────────────────────────────────────────
+
+@dataclass
+class TaskStatus:
+    """Current status of an A2A task."""
+    state: TaskState
+    message: Optional[Message] = None
+    timestamp: str = ""
+
+    def __post_init__(self):
+        if not self.timestamp:
+            self.timestamp = datetime.now(timezone.utc).isoformat()
+
+    def to_dict(self) -> Dict[str, Any]:
+        d: Dict[str, Any] = {
+            "state": self.state.value,
+            "timestamp": self.timestamp,
+        }
+        if self.message:
+            d["message"] = self.message.to_dict()
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> TaskStatus:
+        msg = Message.from_dict(d["message"]) if d.get("message") else None
+        return cls(
+            state=TaskState(d.get("state", "submitted")),
+            message=msg,
+            timestamp=d.get("timestamp", ""),
+        )
+
+
+# ── Task ──────────────────────────────────────────────────────────────
+
+@dataclass
+class Task:
+    """A2A task — the top-level unit of agent-to-agent work."""
+    id: str = ""
+    session_id: str = ""
+    status: TaskStatus = field(default_factory=lambda: TaskStatus(state=TaskState.SUBMITTED))
+    artifacts: List[Artifact] = field(default_factory=list)
+    history: List[Message] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        if not self.id:
+            self.id = secrets.token_hex(12)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "sessionId": self.session_id,
+            "status": self.status.to_dict(),
+            "artifacts": [a.to_dict() for a in self.artifacts],
+            "history": [m.to_dict() for m in self.history],
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> Task:
+        return cls(
+            id=d.get("id", ""),
+            session_id=d.get("sessionId", ""),
+            status=TaskStatus.from_dict(d.get("status", {})),
+            artifacts=[Artifact.from_dict(a) for a in d.get("artifacts", [])],
+            history=[Message.from_dict(m) for m in d.get("history", [])],
+            metadata=d.get("metadata", {}),
+        )

contextsynapse/a2a/streaming.py

@@ -0,0 +1,92 @@
+"""
+A2A SSE Streaming Transport
+=============================
+Server-Sent Events transport for real-time A2A task updates.
+Integrates with PubSubHub to stream task state changes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import AsyncGenerator, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class SSETransport:
+    """SSE subscriber that yields formatted Server-Sent Events."""
+
+    def __init__(self, pubsub=None):
+        self._pubsub = pubsub
+        self._queues: dict = {}  # connection_id → asyncio.Queue
+
+    async def subscribe(
+        self, task_id: str, timeout: float = 300.0
+    ) -> AsyncGenerator[str, None]:
+        """Subscribe to SSE events for a specific task.
+
+        Yields SSE-formatted strings: ``event: <type>\\ndata: <json>\\n\\n``
+        """
+        queue: asyncio.Queue = asyncio.Queue()
+        conn_id = f"sse_{task_id}_{id(queue)}"
+        self._queues[conn_id] = queue
+
+        # Register with PubSub if available
+        if self._pubsub:
+            async def _send(data: str):
+                try:
+                    event = json.loads(data)
+                    # Only forward events for this task
+                    if event.get("data", {}).get("task_id") == task_id:
+                        await queue.put(data)
+                except Exception:
+                    await queue.put(data)
+
+            # Use session_id="" to subscribe to all events (filter in _send)
+            await self._pubsub.subscribe("", conn_id, "", _send)
+
+        try:
+            while True:
+                try:
+                    data = await asyncio.wait_for(queue.get(), timeout=timeout)
+                    # Parse event type from data
+                    try:
+                        parsed = json.loads(data)
+                        event_type = parsed.get("event_type", "TaskStatusUpdateEvent")
+                        event_data = parsed.get("data", parsed)
+                    except (json.JSONDecodeError, TypeError):
+                        event_type = "message"
+                        event_data = {"text": str(data)}
+
+                    yield format_sse(event_type, event_data)
+
+                except asyncio.TimeoutError:
+                    # Send keepalive
+                    yield ": keepalive\n\n"
+
+        except (asyncio.CancelledError, GeneratorExit):
+            pass
+        finally:
+            self._queues.pop(conn_id, None)
+            if self._pubsub:
+                try:
+                    await self._pubsub.unsubscribe("", conn_id)
+                except Exception:
+                    pass
+
+    async def push_event(self, task_id: str, event_type: str, data: dict):
+        """Push an event to all SSE subscribers watching this task."""
+        payload = json.dumps({"event_type": event_type, "data": {**data, "task_id": task_id}})
+        for conn_id, queue in list(self._queues.items()):
+            if task_id in conn_id:
+                try:
+                    await queue.put(payload)
+                except Exception:
+                    pass
+
+
+def format_sse(event_type: str, data: dict) -> str:
+    """Format an SSE event string."""
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"

contextsynapse/a2a/task_manager.py

@@ -0,0 +1,275 @@
+"""
+A2A Task Manager
+=================
+State machine + persistence for A2A tasks.
+Each task maps to a ContextSession for graph/vector/blob access.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import secrets
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from contextsynapse.core.db import IS_POSTGRES, PH, connect, dict_cursor, run_ddl, row_to_dict
+
+from .models import (
+    Artifact,
+    Message,
+    Task,
+    TaskState,
+    TaskStatus,
+    VALID_TRANSITIONS,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class A2ATaskManager:
+    """Manages A2A task lifecycle with SQLite/PostgreSQL persistence."""
+
+    _SCHEMA = """
+    CREATE TABLE IF NOT EXISTS a2a_tasks (
+        task_id TEXT PRIMARY KEY,
+        session_id TEXT,
+        state TEXT NOT NULL DEFAULT 'submitted',
+        history TEXT DEFAULT '[]',
+        artifacts TEXT DEFAULT '[]',
+        metadata TEXT DEFAULT '{}',
+        initiator_agent_id TEXT,
+        target_agent_id TEXT,
+        created_at TEXT NOT NULL,
+        updated_at TEXT NOT NULL
+    );
+    CREATE INDEX IF NOT EXISTS idx_a2a_state ON a2a_tasks(state);
+    CREATE INDEX IF NOT EXISTS idx_a2a_session ON a2a_tasks(session_id);
+    CREATE INDEX IF NOT EXISTS idx_a2a_target ON a2a_tasks(target_agent_id);
+    """
+
+    def __init__(self, db_path: str = "contextcore_data/context.db",
+                 session_manager=None, pubsub=None):
+        self._db_path = Path(db_path)
+        if not IS_POSTGRES:
+            self._db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._conn = connect(str(self._db_path))
+        run_ddl(self._conn, self._SCHEMA)
+        self._session_manager = session_manager
+        self._pubsub = pubsub
+
+    def _exec(self, sql: str, params: tuple = ()):
+        cur = dict_cursor(self._conn)
+        cur.execute(sql, params)
+        return cur
+
+    def _commit(self):
+        self._conn.commit()
+
+    # ── CRUD ──────────────────────────────────────────────────────────
+
+    def create_task(
+        self,
+        message: Message,
+        initiator_agent_id: str = "",
+        target_agent_id: str = "",
+        session_id: str = "",
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Task:
+        """Create a new task in SUBMITTED state."""
+        task_id = secrets.token_hex(12)
+        now = datetime.now(timezone.utc).isoformat()
+
+        # Create a ContextSession if none provided
+        if not session_id and self._session_manager:
+            try:
+                session = self._session_manager.create_session(
+                    name=f"a2a_{task_id[:8]}",
+                    owner_agent_id=target_agent_id or initiator_agent_id,
+                    config={"a2a_task_id": task_id, "source": "a2a"},
+                )
+                session_id = session.session_id
+            except Exception as e:
+                logger.warning("Failed to create session for A2A task: %s", e)
+
+        task = Task(
+            id=task_id,
+            session_id=session_id,
+            status=TaskStatus(state=TaskState.SUBMITTED),
+            history=[message],
+            metadata=metadata or {},
+        )
+
+        self._exec(
+            f"INSERT INTO a2a_tasks "
+            f"(task_id, session_id, state, history, artifacts, metadata, "
+            f"initiator_agent_id, target_agent_id, created_at, updated_at) "
+            f"VALUES ({PH}, {PH}, {PH}, {PH}, {PH}, {PH}, {PH}, {PH}, {PH}, {PH})",
+            (
+                task_id, session_id, TaskState.SUBMITTED.value,
+                json.dumps([m.to_dict() for m in task.history]),
+                "[]",
+                json.dumps(task.metadata),
+                initiator_agent_id, target_agent_id, now, now,
+            ),
+        )
+        self._commit()
+        self._emit_event(task, "task_submitted")
+        return task
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        """Get a task by ID."""
+        row = row_to_dict(self._exec(
+            f"SELECT * FROM a2a_tasks WHERE task_id = {PH}", (task_id,)
+        ).fetchone())
+        return self._row_to_task(row) if row else None
+
+    def list_tasks(
+        self,
+        state: Optional[str] = None,
+        target_agent_id: Optional[str] = None,
+        limit: int = 50,
+    ) -> List[Task]:
+        """List tasks with optional filters."""
+        query = "SELECT * FROM a2a_tasks WHERE 1=1"
+        params: list = []
+        if state:
+            query += f" AND state = {PH}"
+            params.append(state)
+        if target_agent_id:
+            query += f" AND target_agent_id = {PH}"
+            params.append(target_agent_id)
+        query += f" ORDER BY updated_at DESC LIMIT {PH}"
+        params.append(limit)
+        rows = self._exec(query, tuple(params)).fetchall()
+        return [self._row_to_task(row_to_dict(r)) for r in rows]
+
+    # ── State Transitions ─────────────────────────────────────────────
+
+    def transition(
+        self,
+        task_id: str,
+        new_state: TaskState,
+        message: Optional[Message] = None,
+    ) -> Task:
+        """Transition a task to a new state (validates against state machine)."""
+        task = self.get_task(task_id)
+        if not task:
+            raise ValueError(f"Task {task_id} not found")
+
+        current = task.status.state
+        if new_state not in VALID_TRANSITIONS.get(current, set()):
+            raise ValueError(
+                f"Invalid transition: {current.value} → {new_state.value}. "
+                f"Valid: {[s.value for s in VALID_TRANSITIONS.get(current, set())]}"
+            )
+
+        now = datetime.now(timezone.utc).isoformat()
+        task.status = TaskStatus(state=new_state, message=message, timestamp=now)
+        if message:
+            task.history.append(message)
+
+        self._exec(
+            f"UPDATE a2a_tasks SET state = {PH}, history = {PH}, updated_at = {PH} WHERE task_id = {PH}",
+            (
+                new_state.value,
+                json.dumps([m.to_dict() for m in task.history]),
+                now, task_id,
+            ),
+        )
+        self._commit()
+        self._emit_event(task, f"task_{new_state.value}")
+        return task
+
+    def cancel_task(self, task_id: str) -> Task:
+        """Cancel a task."""
+        return self.transition(task_id, TaskState.CANCELED)
+
+    def complete_task(self, task_id: str, message: Optional[Message] = None) -> Task:
+        """Mark a task as completed."""
+        return self.transition(task_id, TaskState.COMPLETED, message)
+
+    def fail_task(self, task_id: str, error: str) -> Task:
+        """Mark a task as failed."""
+        msg = Message.text(f"Error: {error}", role="agent")
+        return self.transition(task_id, TaskState.FAILED, msg)
+
+    # ── Artifacts ─────────────────────────────────────────────────────
+
+    def add_artifact(self, task_id: str, artifact: Artifact) -> Task:
+        """Add an artifact to a task."""
+        task = self.get_task(task_id)
+        if not task:
+            raise ValueError(f"Task {task_id} not found")
+
+        task.artifacts.append(artifact)
+        now = datetime.now(timezone.utc).isoformat()
+        self._exec(
+            f"UPDATE a2a_tasks SET artifacts = {PH}, updated_at = {PH} WHERE task_id = {PH}",
+            (json.dumps([a.to_dict() for a in task.artifacts]), now, task_id),
+        )
+        self._commit()
+        self._emit_event(task, "task_artifact_added")
+        return task
+
+    def add_message(self, task_id: str, message: Message) -> Task:
+        """Append a message to a task's history."""
+        task = self.get_task(task_id)
+        if not task:
+            raise ValueError(f"Task {task_id} not found")
+
+        task.history.append(message)
+        now = datetime.now(timezone.utc).isoformat()
+        self._exec(
+            f"UPDATE a2a_tasks SET history = {PH}, updated_at = {PH} WHERE task_id = {PH}",
+            (json.dumps([m.to_dict() for m in task.history]), now, task_id),
+        )
+        self._commit()
+        return task
+
+    # ── Internal ──────────────────────────────────────────────────────
+
+    def _row_to_task(self, row: dict) -> Task:
+        history = [Message.from_dict(m) for m in json.loads(row["history"] or "[]")]
+        artifacts = [Artifact.from_dict(a) for a in json.loads(row["artifacts"] or "[]")]
+        meta = json.loads(row["metadata"] or "{}")
+        meta["initiator_agent_id"] = row["initiator_agent_id"]
+        meta["target_agent_id"] = row["target_agent_id"]
+        return Task(
+            id=row["task_id"],
+            session_id=row["session_id"] or "",
+            status=TaskStatus(state=TaskState(row["state"])),
+            history=history,
+            artifacts=artifacts,
+            metadata=meta,
+        )
+
+    def _emit_event(self, task: Task, event_type: str):
+        """Emit a PubSub event for task state changes."""
+        if not self._pubsub:
+            return
+        try:
+            from ..context.pubsub import ContextEvent
+            event = ContextEvent(
+                event_type=event_type,
+                session_id=task.session_id,
+                agent_id=task.metadata.get("target_agent_id", ""),
+                data={
+                    "task_id": task.id,
+                    "state": task.status.state.value,
+                    "timestamp": task.status.timestamp,
+                },
+            )
+            import asyncio
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(self._pubsub.publish(event))
+            except RuntimeError:
+                pass  # No event loop — skip async emit
+        except Exception as e:
+            logger.debug("Failed to emit A2A event: %s", e)
+
+    def close(self):
+        if self._conn:
+            self._conn.close()