Coverage for agentos/protocols/a2a_store.py: 34%
128 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1"""
2A2A Task Store — persistent task and session storage for A2A protocol.
4Backends: InMemory (default), SQLite, custom.
5"""
7from __future__ import annotations
9import json
10import sqlite3
11import threading
12import time
13from abc import ABC, abstractmethod
14from collections.abc import Iterator
15from contextlib import contextmanager
16from typing import Any
18from agentos.protocols.a2a import A2ATask, TaskState
21class A2ATaskStore(ABC):
22 """Abstract task store for A2A protocol persistence."""
24 @abstractmethod
25 def save_task(self, task: A2ATask) -> None:
26 """Insert or update a task."""
27 ...
29 @abstractmethod
30 def get_task(self, task_id: str) -> A2ATask | None:
31 """Retrieve a task by ID."""
32 ...
34 @abstractmethod
35 def list_tasks(
36 self,
37 state: TaskState | None = None,
38 limit: int = 100,
39 offset: int = 0,
40 agent: str = "",
41 ) -> list[A2ATask]:
42 """List tasks, optionally filtered by state/agent."""
43 ...
45 @abstractmethod
46 def delete_task(self, task_id: str) -> bool:
47 """Delete a task. Returns True if deleted."""
48 ...
50 @abstractmethod
51 def cleanup_terminal(
52 self,
53 max_age_seconds: float = 3600.0,
54 ) -> int:
55 """Remove terminal tasks older than max_age. Returns count."""
56 ...
58 @abstractmethod
59 def count(self, state: TaskState | None = None) -> int:
60 """Count tasks, optionally filtered by state."""
61 ...
64class InMemoryTaskStore(A2ATaskStore):
65 """Fast, non-persistent task store for development/testing."""
67 def __init__(self):
68 self._tasks: dict[str, A2ATask] = {}
69 self._lock = threading.Lock()
71 def save_task(self, task: A2ATask) -> None:
72 with self._lock:
73 self._tasks[task.task_id] = task
75 def get_task(self, task_id: str) -> A2ATask | None:
76 with self._lock:
77 return self._tasks.get(task_id)
79 def list_tasks(
80 self,
81 state: TaskState | None = None,
82 limit: int = 100,
83 offset: int = 0,
84 agent: str = "",
85 ) -> list[A2ATask]:
86 with self._lock:
87 tasks = list(self._tasks.values())
88 if state:
89 tasks = [t for t in tasks if t.state == state]
90 if agent:
91 tasks = [t for t in tasks if t.meta.get("target_agent") == agent]
92 return tasks[offset : offset + limit]
94 def delete_task(self, task_id: str) -> bool:
95 with self._lock:
96 if task_id in self._tasks:
97 del self._tasks[task_id]
98 return True
99 return False
101 def cleanup_terminal(self, max_age_seconds: float = 3600.0) -> int:
102 now = time.time()
103 with self._lock:
104 to_del = [
105 tid
106 for tid, t in self._tasks.items()
107 if t.is_terminal() and (now - t._updated) > max_age_seconds
108 ]
109 for tid in to_del:
110 del self._tasks[tid]
111 return len(to_del)
113 def count(self, state: TaskState | None = None) -> int:
114 if state is None:
115 with self._lock:
116 return len(self._tasks)
117 tasks = self.list_tasks(state=state, limit=999999)
118 return len(tasks)
121class SqliteTaskStore(A2ATaskStore):
122 """Persistent SQLite-backed task store for production use."""
124 SCHEMA = """
125 CREATE TABLE IF NOT EXISTS a2a_tasks (
126 task_id TEXT PRIMARY KEY,
127 state TEXT NOT NULL DEFAULT 'submitted',
128 input_json TEXT,
129 output_json TEXT,
130 artifacts_json TEXT DEFAULT '[]',
131 error TEXT,
132 meta_json TEXT DEFAULT '{}',
133 created REAL NOT NULL,
134 updated REAL NOT NULL,
135 agent TEXT DEFAULT ''
136 );
137 CREATE INDEX IF NOT EXISTS idx_a2a_state ON a2a_tasks(state);
138 CREATE INDEX IF NOT EXISTS idx_a2a_agent ON a2a_tasks(agent);
139 CREATE INDEX IF NOT EXISTS idx_a2a_updated ON a2a_tasks(updated);
140 """
142 def __init__(self, db_path: str = ":memory:"):
143 self.db_path = db_path
144 self._local = threading.local()
145 self._init_db()
147 def _init_db(self):
148 with self._conn() as conn:
149 conn.executescript(self.SCHEMA)
151 @contextmanager
152 def _conn(self) -> Iterator[sqlite3.Connection]:
153 if not hasattr(self._local, "conn") or self._local.conn is None:
154 conn = sqlite3.connect(self.db_path, check_same_thread=False)
155 conn.row_factory = sqlite3.Row
156 self._local.conn = conn
157 yield self._local.conn
159 def save_task(self, task: A2ATask) -> None:
160 with self._conn() as conn:
161 conn.execute(
162 """INSERT OR REPLACE INTO a2a_tasks
163 (task_id, state, input_json, output_json, artifacts_json,
164 error, meta_json, created, updated, agent)
165 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
166 (
167 task.task_id,
168 task.state.value,
169 json.dumps(task.input.to_dict()) if task.input else None,
170 json.dumps(task.output.to_dict()) if task.output else None,
171 json.dumps([a.to_dict() for a in task.artifacts]),
172 task.error,
173 json.dumps(task.meta),
174 task._created,
175 task._updated,
176 task.meta.get("target_agent", ""),
177 ),
178 )
179 conn.commit()
181 def get_task(self, task_id: str) -> A2ATask | None:
182 with self._conn() as conn:
183 row = conn.execute("SELECT * FROM a2a_tasks WHERE task_id = ?", (task_id,)).fetchone()
184 if row is None:
185 return None
186 return self._row_to_task(row)
188 def list_tasks(
189 self,
190 state: TaskState | None = None,
191 limit: int = 100,
192 offset: int = 0,
193 agent: str = "",
194 ) -> list[A2ATask]:
195 query = "SELECT * FROM a2a_tasks WHERE 1=1"
196 params: list[Any] = []
197 if state:
198 query += " AND state = ?"
199 params.append(state.value)
200 if agent:
201 query += " AND agent = ?"
202 params.append(agent)
203 query += " ORDER BY updated DESC LIMIT ? OFFSET ?"
204 params.extend([limit, offset])
206 with self._conn() as conn:
207 rows = conn.execute(query, params).fetchall()
208 return [self._row_to_task(r) for r in rows]
210 def delete_task(self, task_id: str) -> bool:
211 with self._conn() as conn:
212 cur = conn.execute("DELETE FROM a2a_tasks WHERE task_id = ?", (task_id,))
213 conn.commit()
214 return cur.rowcount > 0
216 def cleanup_terminal(self, max_age_seconds: float = 3600.0) -> int:
217 cutoff = time.time() - max_age_seconds
218 with self._conn() as conn:
219 cur = conn.execute(
220 """DELETE FROM a2a_tasks
221 WHERE state IN ('completed', 'failed', 'cancelled')
222 AND updated < ?""",
223 (cutoff,),
224 )
225 conn.commit()
226 return cur.rowcount
228 def count(self, state: TaskState | None = None) -> int:
229 query = "SELECT COUNT(*) FROM a2a_tasks"
230 params: list[Any] = []
231 if state:
232 query += " WHERE state = ?"
233 params.append(state.value)
234 with self._conn() as conn:
235 return conn.execute(query, params).fetchone()[0]
237 def _row_to_task(self, row) -> A2ATask:
238 from agentos.protocols.a2a import A2AArtifact, A2AMessage
240 task = A2ATask(
241 task_id=row["task_id"],
242 state=TaskState(row["state"]),
243 error=row["error"],
244 meta=json.loads(row["meta_json"] or "{}"),
245 _created=row["created"],
246 _updated=row["updated"],
247 )
248 if row["input_json"]:
249 task.input = A2AMessage.from_dict(json.loads(row["input_json"]))
250 if row["output_json"]:
251 task.output = A2AMessage.from_dict(json.loads(row["output_json"]))
252 task.artifacts = [
253 A2AArtifact.from_dict(a) for a in json.loads(row["artifacts_json"] or "[]")
254 ]
255 return task