@team-agent/installer 0.1.11 → 0.2.1

package/src/team_agent/{message_store.py → message_store/core.py}

@@ -8,15 +8,16 @@ from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Any
 
+from . import agent_health as _agent_health
+from . import result_watchers as _result_watchers
+from .schema import SCHEMA_VERSION, initialize_schema, utcnow
 from team_agent.paths import runtime_dir
 from team_agent.spec import validate_result_envelope
 
 
-def utcnow() -> str:
-    return datetime.now(timezone.utc).isoformat()
-
-
 class MessageStore:
+    SCHEMA_VERSION = SCHEMA_VERSION
+
     def __init__(self, workspace: Path):
         self.workspace = workspace
         self.path = runtime_dir(workspace) / "team.db"
@@ -24,112 +25,15 @@ class MessageStore:
         self._init()
 
     def connect(self) -> sqlite3.Connection:
-        conn = sqlite3.connect(self.path)
+        conn = sqlite3.connect(self.path, timeout=30.0, isolation_level=None)
         conn.row_factory = sqlite3.Row
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA busy_timeout=30000")
         return conn
 
     def _init(self) -> None:
         with closing(self.connect()) as conn:
-            with conn:
-                conn.execute(
-                    """
-                    create table if not exists messages (
-                      message_id text primary key,
-                      task_id text,
-                      sender text,
-                      recipient text,
-                      reply_to text,
-                      requires_ack integer,
-                      status text,
-                      content text,
-                      artifact_refs text,
-                      created_at text,
-                      updated_at text,
-                      delivered_at text,
-                      acknowledged_at text,
-                      error text
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists results (
-                      result_id text primary key,
-                      task_id text not null,
-                      agent_id text not null,
-                      envelope text not null,
-                      status text not null,
-                      created_at text not null
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists scheduled_events (
-                      id integer primary key,
-                      due_at text not null,
-                      target text not null,
-                      kind text not null,
-                      payload_json text not null,
-                      status text not null,
-                      created_at text not null,
-                      fired_at text,
-                      result_json text
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists delivery_tokens (
-                      message_id text primary key,
-                      unique_token text not null,
-                      injected_at text not null,
-                      visible_at text,
-                      consumed_at text,
-                      failed_at text,
-                      failure_reason text
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists agent_health (
-                      agent_id text primary key,
-                      status text not null,
-                      last_output_at text,
-                      context_usage_pct integer,
-                      current_task_id text,
-                      updated_at text not null
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists peer_allowlist (
-                      a text not null,
-                      b text not null,
-                      created_at text not null,
-                      primary key (a, b)
-                    )
-                    """
-                )
-                conn.execute(
-                    """
-                    create table if not exists result_watchers (
-                      watcher_id text primary key,
-                      task_id text,
-                      agent_id text,
-                      message_id text,
-                      leader_id text not null,
-                      status text not null,
-                      created_at text not null,
-                      completed_at text,
-                      result_id text,
-                      notified_message_id text,
-                      error text
-                    )
-                    """
-                )
+            initialize_schema(conn)
 
     def create_message(
         self,
@@ -140,6 +44,7 @@ class MessageStore:
         reply_to: str | None = None,
         requires_ack: bool = True,
         artifact_refs: list[dict[str, Any]] | None = None,
+        owner_team_id: str | None = None,
     ) -> str:
         message_id = f"msg_{uuid.uuid4().hex[:12]}"
         now = utcnow()
@@ -147,10 +52,16 @@ class MessageStore:
             with conn:
                 conn.execute(
                     """
-                    insert into messages values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    insert into messages(
+                      message_id, owner_team_id, task_id, sender, recipient, reply_to, requires_ack,
+                      status, content, artifact_refs, created_at, updated_at,
+                      delivered_at, acknowledged_at, error, delivery_attempts
+                    )
+                    values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                     """,
                     (
                         message_id,
+                        owner_team_id,
                         task_id,
                         sender,
                         recipient,
@@ -164,6 +75,7 @@ class MessageStore:
                         None,
                         None,
                         None,
+                        0,
                     ),
                 )
         return message_id
@@ -177,16 +89,29 @@ class MessageStore:
                 conn.execute(
                     """
                     update messages
-                    set status = ?,
+                    set status = case
+                            when status = 'acknowledged'
+                             and ? in ('injected', 'visible', 'submitted', 'submitted_unverified', 'delivered')
+                            then status
+                            else ?
+                        end,
                         updated_at = ?,
                         delivered_at = coalesce(?, delivered_at),
                         acknowledged_at = coalesce(?, acknowledged_at),
                         error = coalesce(?, error)
                     where message_id = ?
                     """,
-                    (status, now, delivered_at, acknowledged_at, error, message_id),
+                    (status, status, now, delivered_at, acknowledged_at, error, message_id),
                 )
-                if status in {"injected", "visible", "submitted", "submitted_unverified", "injected_unverified", "failed"}:
+                if status in {
+                    "injected",
+                    "visible",
+                    "submitted",
+                    "submitted_unverified",
+                    "injected_unverified",
+                    "delivery_blocked",
+                    "failed",
+                }:
                     conn.execute(
                         """
                         insert into delivery_tokens(
@@ -203,8 +128,8 @@ class MessageStore:
                             message_id,
                             now,
                             now if status in {"visible", "submitted"} else None,
-                            now if status in {"failed", "injected_unverified"} else None,
-                            error if status in {"failed", "injected_unverified"} else None,
+                            now if status in {"failed", "injected_unverified", "delivery_blocked"} else None,
+                            error if status in {"failed", "injected_unverified", "delivery_blocked"} else None,
                         ),
                     )
                 if status == "acknowledged":
@@ -213,6 +138,22 @@ class MessageStore:
                         (now, message_id),
                     )
 
+    def defer_delivery(self, message_id: str, status: str, reason: str) -> None:
+        now = utcnow()
+        with closing(self.connect()) as conn:
+            with conn:
+                conn.execute(
+                    """
+                    update messages
+                    set status = ?,
+                        updated_at = ?,
+                        error = ?,
+                        delivery_attempts = delivery_attempts + 1
+                    where message_id = ?
+                    """,
+                    (status, now, reason, message_id),
+                )
+
     def claim_for_delivery(self, message_id: str) -> bool:
         now = utcnow()
         with closing(self.connect()) as conn:
@@ -221,23 +162,27 @@ class MessageStore:
                     """
                     update messages
                     set status = 'target_resolved',
-                        updated_at = ?
+                        updated_at = ?,
+                        delivery_attempts = delivery_attempts + 1
                     where message_id = ?
-                      and status in ('pending', 'accepted')
+                      and status in ('pending', 'accepted', 'queued_until_idle', 'queued_until_start', 'queued_stopped', 'queued_pane_missing')
                     """,
                     (now, message_id),
                 )
         return result.rowcount == 1
 
-    def fail_timeouts(self, timeout_sec: int) -> list[str]:
+    def fail_timeouts(self, timeout_sec: int, exclude_recipients: set[str] | None = None) -> list[str]:
         cutoff = datetime.now(timezone.utc) - timedelta(seconds=timeout_sec)
         failed: list[str] = []
+        exclude_recipients = exclude_recipients or set()
         with closing(self.connect()) as conn:
             with conn:
                 rows = conn.execute(
-                    "select message_id, updated_at from messages where requires_ack = 1 and status in ('pending','accepted','target_resolved','injected','visible','submitted','delivered')"
+                    "select message_id, recipient, updated_at from messages where requires_ack = 1 and status in ('pending','accepted','target_resolved','injected','visible','submitted','delivered')"
                 ).fetchall()
                 for row in rows:
+                    if row["recipient"] in exclude_recipients:
+                        continue
                     try:
                         updated = datetime.fromisoformat(row["updated_at"])
                     except ValueError:
@@ -250,22 +195,40 @@ class MessageStore:
                         )
         return failed
 
-    def messages(self) -> list[dict[str, Any]]:
+    def messages(self, owner_team_id: str | None = None) -> list[dict[str, Any]]:
         with closing(self.connect()) as conn:
-            rows = conn.execute("select * from messages order by created_at").fetchall()
+            if owner_team_id is None:
+                rows = conn.execute("select * from messages order by created_at").fetchall()
+            else:
+                rows = conn.execute(
+                    "select * from messages where owner_team_id = ? or owner_team_id is null order by created_at",
+                    (owner_team_id,),
+                ).fetchall()
         return [dict(row) for row in rows]
 
-    def inbox(self, agent_id: str, limit: int = 20) -> list[dict[str, Any]]:
+    def inbox(self, agent_id: str, limit: int = 20, owner_team_id: str | None = None) -> list[dict[str, Any]]:
         with closing(self.connect()) as conn:
-            rows = conn.execute(
-                """
-                select * from messages
-                where sender = ? or recipient = ?
-                order by created_at desc
-                limit ?
-                """,
-                (agent_id, agent_id, limit),
-            ).fetchall()
+            if owner_team_id is None:
+                rows = conn.execute(
+                    """
+                    select * from messages
+                    where sender = ? or recipient = ?
+                    order by created_at desc
+                    limit ?
+                    """,
+                    (agent_id, agent_id, limit),
+                ).fetchall()
+            else:
+                rows = conn.execute(
+                    """
+                    select * from messages
+                    where (sender = ? or recipient = ?)
+                      and (owner_team_id = ? or owner_team_id is null)
+                    order by created_at desc
+                    limit ?
+                    """,
+                    (agent_id, agent_id, owner_team_id, limit),
+                ).fetchall()
         return [dict(row) for row in reversed(rows)]
 
     def delivery_tokens(self) -> list[dict[str, Any]]:
@@ -273,58 +236,46 @@ class MessageStore:
             rows = conn.execute("select * from delivery_tokens order by injected_at").fetchall()
         return [dict(row) for row in rows]
 
-    def upsert_agent_health(
+    def add_scheduled_event(
         self,
-        agent_id: str,
-        status: str,
-        last_output_at: str | None = None,
-        context_usage_pct: int | None = None,
-        current_task_id: str | None = None,
-    ) -> None:
-        now = utcnow()
-        with closing(self.connect()) as conn:
-            with conn:
-                conn.execute(
-                    """
-                    insert into agent_health(agent_id, status, last_output_at, context_usage_pct, current_task_id, updated_at)
-                    values (?, ?, ?, ?, ?, ?)
-                    on conflict(agent_id) do update set
-                      status = excluded.status,
-                      last_output_at = coalesce(excluded.last_output_at, agent_health.last_output_at),
-                      context_usage_pct = excluded.context_usage_pct,
-                      current_task_id = excluded.current_task_id,
-                      updated_at = excluded.updated_at
-                    """,
-                    (agent_id, status, last_output_at, context_usage_pct, current_task_id, now),
-                )
-
-    def agent_health(self) -> dict[str, dict[str, Any]]:
-        with closing(self.connect()) as conn:
-            rows = conn.execute("select * from agent_health order by agent_id").fetchall()
-        return {row["agent_id"]: dict(row) for row in rows}
-
-    def add_scheduled_event(self, due_at: str, target: str, kind: str, payload: dict[str, Any]) -> int:
+        due_at: str,
+        target: str,
+        kind: str,
+        payload: dict[str, Any],
+        owner_team_id: str | None = None,
+    ) -> int:
         with closing(self.connect()) as conn:
             with conn:
                 cur = conn.execute(
                     """
-                    insert into scheduled_events(due_at, target, kind, payload_json, status, created_at)
-                    values (?, ?, ?, ?, 'pending', ?)
+                    insert into scheduled_events(owner_team_id, due_at, target, kind, payload_json, status, created_at)
+                    values (?, ?, ?, ?, ?, 'pending', ?)
                     """,
-                    (due_at, target, kind, json.dumps(payload, ensure_ascii=False), utcnow()),
+                    (owner_team_id, due_at, target, kind, json.dumps(payload, ensure_ascii=False), utcnow()),
                 )
                 return int(cur.lastrowid)
 
-    def due_scheduled_events(self, now: str | None = None) -> list[dict[str, Any]]:
+    def due_scheduled_events(self, now: str | None = None, owner_team_id: str | None = None) -> list[dict[str, Any]]:
         with closing(self.connect()) as conn:
-            rows = conn.execute(
-                """
-                select * from scheduled_events
-                where status = 'pending' and due_at <= ?
-                order by due_at, id
-                """,
-                (now or utcnow(),),
-            ).fetchall()
+            if owner_team_id is None:
+                rows = conn.execute(
+                    """
+                    select * from scheduled_events
+                    where status = 'pending' and due_at <= ?
+                    order by due_at, id
+                    """,
+                    (now or utcnow(),),
+                ).fetchall()
+            else:
+                rows = conn.execute(
+                    """
+                    select * from scheduled_events
+                    where status = 'pending' and due_at <= ?
+                      and (owner_team_id = ? or owner_team_id is null)
+                    order by due_at, id
+                    """,
+                    (now or utcnow(), owner_team_id),
+                ).fetchall()
         return [dict(row) for row in rows]
 
     def mark_scheduled_event(self, event_id: int, status: str, result: dict[str, Any]) -> None:
@@ -379,69 +330,17 @@ class MessageStore:
                 counts["by_status"][status] = n
         return counts
 
-    def create_result_watcher(
-        self,
-        task_id: str | None,
-        agent_id: str | None,
-        message_id: str | None,
-        leader_id: str = "leader",
-    ) -> str:
-        watcher_id = f"watch_{uuid.uuid4().hex[:12]}"
-        with closing(self.connect()) as conn:
-            with conn:
-                conn.execute(
-                    """
-                    insert into result_watchers(
-                      watcher_id, task_id, agent_id, message_id, leader_id, status, created_at
-                    )
-                    values (?, ?, ?, ?, ?, 'pending', ?)
-                    """,
-                    (watcher_id, task_id, agent_id, message_id, leader_id, utcnow()),
-                )
-        return watcher_id
-
-    def pending_result_watchers(self) -> list[dict[str, Any]]:
-        with closing(self.connect()) as conn:
-            rows = conn.execute(
-                "select * from result_watchers where status = 'pending' order by created_at"
-            ).fetchall()
-        return [dict(row) for row in rows]
-
-    def result_watchers(self) -> list[dict[str, Any]]:
-        with closing(self.connect()) as conn:
-            rows = conn.execute("select * from result_watchers order by created_at").fetchall()
-        return [dict(row) for row in rows]
-
-    def mark_result_watcher(
-        self,
-        watcher_id: str,
-        status: str,
-        result_id: str | None = None,
-        notified_message_id: str | None = None,
-        error: str | None = None,
-    ) -> None:
-        with closing(self.connect()) as conn:
-            with conn:
-                conn.execute(
-                    """
-                    update result_watchers
-                    set status = ?,
-                        completed_at = ?,
-                        result_id = coalesce(?, result_id),
-                        notified_message_id = coalesce(?, notified_message_id),
-                        error = coalesce(?, error)
-                    where watcher_id = ?
-                    """,
-                    (status, utcnow(), result_id, notified_message_id, error, watcher_id),
-                )
-
-    def add_result(self, envelope: dict[str, Any]) -> str:
+    def add_result(self, envelope: dict[str, Any], owner_team_id: str | None = None) -> str:
+        _ = owner_team_id
         validate_result_envelope(envelope)
         result_id = f"res_{uuid.uuid4().hex[:12]}"
         with closing(self.connect()) as conn:
             with conn:
                 conn.execute(
-                    "insert into results values (?, ?, ?, ?, ?, ?)",
+                    """
+                    insert into results(result_id, task_id, agent_id, envelope, status, created_at)
+                    values (?, ?, ?, ?, ?, ?)
+                    """,
                     (
                         result_id,
                         envelope["task_id"],
@@ -453,59 +352,90 @@ class MessageStore:
                 )
         return result_id
 
-    def acknowledge_task_messages(self, task_id: str, agent_id: str) -> list[str]:
+    def acknowledge_task_messages(self, task_id: str, agent_id: str, owner_team_id: str | None = None) -> list[str]:
         now = utcnow()
+        owner_clause = "and (owner_team_id = ? or owner_team_id is null)" if owner_team_id else ""
+        owner_args = (owner_team_id,) if owner_team_id else ()
         with closing(self.connect()) as conn:
             with conn:
                 rows = conn.execute(
-                    """
+                    f"""
                     select message_id from messages
-                    where task_id = ? and recipient = ? and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
+                    where task_id = ? and recipient = ? {owner_clause}
+                      and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
                     """,
-                    (task_id, agent_id),
+                    (task_id, agent_id, *owner_args),
                 ).fetchall()
                 ids = [row["message_id"] for row in rows]
                 conn.execute(
-                    """
+                    f"""
                     update messages
                     set status = 'acknowledged', acknowledged_at = ?, updated_at = ?
-                    where task_id = ? and recipient = ? and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
+                    where task_id = ? and recipient = ? {owner_clause}
+                      and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
                     """,
-                    (now, now, task_id, agent_id),
+                    (now, now, task_id, agent_id, *owner_args),
                 )
         return ids
 
-    def acknowledge_message(self, message_id: str, agent_id: str) -> list[str]:
+    def acknowledge_message(self, message_id: str, agent_id: str, owner_team_id: str | None = None) -> list[str]:
         now = utcnow()
+        owner_clause = "and (owner_team_id = ? or owner_team_id is null)" if owner_team_id else ""
+        owner_args = (owner_team_id,) if owner_team_id else ()
         with closing(self.connect()) as conn:
             with conn:
                 row = conn.execute(
-                    """
+                    f"""
                     select message_id from messages
-                    where message_id = ? and recipient = ? and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
+                    where message_id = ? and recipient = ? {owner_clause}
+                      and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
                     """,
-                    (message_id, agent_id),
+                    (message_id, agent_id, *owner_args),
                 ).fetchone()
                 if not row:
                     return []
                 conn.execute(
-                    """
+                    f"""
                     update messages
                     set status = 'acknowledged', acknowledged_at = ?, updated_at = ?
-                    where message_id = ? and recipient = ? and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
+                    where message_id = ? and recipient = ? {owner_clause}
+                      and status in ('pending', 'accepted', 'target_resolved', 'injected', 'visible', 'submitted', 'delivered')
                     """,
-                    (now, now, message_id, agent_id),
+                    (now, now, message_id, agent_id, *owner_args),
                 )
         return [message_id]
 
-    def results(self, uncollected_only: bool = False) -> list[dict[str, Any]]:
-        query = "select * from results order by created_at"
+    def results(self, uncollected_only: bool = False, owner_team_id: str | None = None) -> list[dict[str, Any]]:
+        _ = owner_team_id
+        clauses: list[str] = []
+        args: list[Any] = []
         if uncollected_only:
-            query = "select * from results where status not in ('collected', 'invalid') order by created_at"
+            clauses.append("status not in ('collected', 'invalid')")
+        where = " where " + " and ".join(clauses) if clauses else ""
+        query = f"select * from results{where} order by created_at"
         with closing(self.connect()) as conn:
-            rows = conn.execute(query).fetchall()
+            rows = conn.execute(query, args).fetchall()
         return [dict(row) for row in rows]
 
+    def result_by_id(self, result_id: str) -> dict[str, Any] | None:
+        with closing(self.connect()) as conn:
+            row = conn.execute("select * from results where result_id = ?", (result_id,)).fetchone()
+        return dict(row) if row else None
+
+    def latest_results(self, limit: int = 5, owner_team_id: str | None = None) -> list[dict[str, Any]]:
+        _ = owner_team_id
+        with closing(self.connect()) as conn:
+            rows = conn.execute(
+                """
+                select * from results
+                where status != 'invalid'
+                order by created_at desc
+                limit ?
+                """,
+                (limit,),
+            ).fetchall()
+        return [dict(row) for row in reversed(rows)]
+
     def mark_result_collected(self, result_id: str) -> None:
         with closing(self.connect()) as conn:
             with conn:
@@ -518,3 +448,16 @@ class MessageStore:
                     "update results set status = 'invalid' where result_id = ?",
                     (result_id,),
                 )
+
+
+
+MessageStore.upsert_agent_health = _agent_health.upsert_agent_health
+MessageStore.agent_health = _agent_health.agent_health
+MessageStore.delete_agent_health = _agent_health.delete_agent_health
+MessageStore.gc_agent_health = _agent_health.gc_agent_health
+MessageStore.create_result_watcher = _result_watchers.create_result_watcher
+MessageStore.pending_result_watchers = _result_watchers.pending_result_watchers
+MessageStore.retryable_result_watchers = _result_watchers.retryable_result_watchers
+MessageStore.result_watchers = _result_watchers.result_watchers
+MessageStore.mark_result_watcher = _result_watchers.mark_result_watcher
+MessageStore.requeue_delivery_exhausted_watchers = _result_watchers.requeue_delivery_exhausted_watchers