claude-team-mcp 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

claude_team/poller.py

@@ -0,0 +1,245 @@
+"""Background poller for worker state snapshots."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Literal, Protocol
+
+from . import events
+from .idle_detection import detect_worker_idle
+
+logger = logging.getLogger("claude-team-poller")
+
+WorkerState = Literal["idle", "active"]
+
+
+# Minimal registry interface used by WorkerPoller.
+class _RegistryLike(Protocol):
+    def list_all(self) -> list["_SessionLike"]:
+        ...
+
+
+# Minimal session interface used by WorkerPoller.
+class _SessionLike(Protocol):
+    session_id: str
+    agent_type: Literal["claude", "codex"]
+    project_path: str
+    claude_session_id: str | None
+    output_path: Path | None
+    message_count: int | None
+    last_message_count: int | None
+    last_message_timestamp: float | None
+    pid: int | None
+    is_idle: bool
+
+    def to_dict(self) -> dict:
+        ...
+
+
+# Snapshot of a worker at a point in time.
+@dataclass(frozen=True)
+class _WorkerSnapshot:
+    session_id: str
+    state: WorkerState
+    info: dict
+
+
+def _isoformat_zulu(value: datetime) -> str:
+    # Format timestamps with a Z suffix for UTC.
+    return value.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+def _sanitize_for_json(obj: object) -> object:
+    # Recursively sanitize an object for JSON serialization.
+    # Removes non-serializable types like asyncio Futures, methods, etc.
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return obj
+    if callable(obj):
+        # Skip methods, functions, lambdas
+        return None
+    if isinstance(obj, dict):
+        return {str(k): _sanitize_for_json(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_sanitize_for_json(item) for item in obj]
+    if isinstance(obj, Path):
+        return str(obj)
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    # For any other type, try to convert to string, else skip
+    try:
+        return str(obj)
+    except Exception:
+        return None
+
+
+def _build_snapshot(registry: _RegistryLike) -> dict[str, _WorkerSnapshot]:
+    # Capture current worker states from the registry.
+    snapshots: dict[str, _WorkerSnapshot] = {}
+    for session in registry.list_all():
+        info = _sanitize_for_json(session.to_dict())
+        is_idle, _ = detect_worker_idle(session, idle_threshold_seconds=300)
+        info["is_idle"] = is_idle
+        state: WorkerState = "idle" if is_idle else "active"
+        snapshots[session.session_id] = _WorkerSnapshot(session.session_id, state, info)
+    return snapshots
+
+
+def _snapshot_payload(snapshots: dict[str, _WorkerSnapshot]) -> dict:
+    # Build a full snapshot payload for persistence.
+    workers = []
+    for snapshot in snapshots.values():
+        payload = dict(snapshot.info)
+        payload["state"] = snapshot.state
+        workers.append(payload)
+    return {"count": len(workers), "workers": workers}
+
+
+def _transition_payload(snapshot: _WorkerSnapshot, previous_state: WorkerState | None) -> dict:
+    # Build transition data payload for a worker event.
+    payload = dict(snapshot.info)
+    payload["state"] = snapshot.state
+    payload["previous_state"] = previous_state
+    return payload
+
+
+def _closed_payload(snapshot: _WorkerSnapshot) -> dict:
+    # Build payload for a worker_closed event using the last known state.
+    payload = dict(snapshot.info)
+    payload["state"] = "closed"
+    payload["previous_state"] = snapshot.state
+    return payload
+
+
+def _build_transition_events(
+    previous: dict[str, _WorkerSnapshot],
+    current: dict[str, _WorkerSnapshot],
+    timestamp: str,
+) -> list[events.WorkerEvent]:
+    # Compare snapshot sets and emit lifecycle transition events.
+    results: list[events.WorkerEvent] = []
+    previous_ids = set(previous)
+    current_ids = set(current)
+
+    # New sessions -> worker_started events.
+    for session_id in current_ids - previous_ids:
+        snapshot = current[session_id]
+        results.append(events.WorkerEvent(
+            ts=timestamp,
+            type="worker_started",
+            worker_id=session_id,
+            data=_transition_payload(snapshot, None),
+        ))
+
+    # Removed sessions -> worker_closed events.
+    for session_id in previous_ids - current_ids:
+        snapshot = previous[session_id]
+        results.append(events.WorkerEvent(
+            ts=timestamp,
+            type="worker_closed",
+            worker_id=session_id,
+            data=_closed_payload(snapshot),
+        ))
+
+    # Existing sessions -> idle/active transitions.
+    for session_id in previous_ids & current_ids:
+        before = previous[session_id]
+        after = current[session_id]
+        if before.state == after.state:
+            continue
+        event_type = "worker_idle" if after.state == "idle" else "worker_active"
+        results.append(events.WorkerEvent(
+            ts=timestamp,
+            type=event_type,
+            worker_id=session_id,
+            data=_transition_payload(after, before.state),
+        ))
+
+    return results
+
+
+class WorkerPoller:
+    """Background poller that snapshots worker state and logs transitions."""
+
+    def __init__(
+        self,
+        registry: _RegistryLike,
+        poll_interval_seconds: int = 60,
+        snapshot_interval_seconds: int = 300,
+    ) -> None:
+        """Initialize the poller with registry and polling cadence."""
+        self._registry = registry
+        self._poll_interval_seconds = poll_interval_seconds
+        self._snapshot_interval_seconds = snapshot_interval_seconds
+        self._stop_event = asyncio.Event()
+        self._task: asyncio.Task | None = None
+        self._last_snapshot: dict[str, _WorkerSnapshot] = {}
+        self._last_snapshot_event_at: float | None = None
+
+    def start(self) -> None:
+        """Start the background polling task."""
+        if self._task and not self._task.done():
+            return
+        self._stop_event.clear()
+        self._task = asyncio.create_task(self._run(), name="worker-poller")
+
+    async def stop(self) -> None:
+        """Stop the background polling task."""
+        if not self._task:
+            return
+        self._stop_event.set()
+        await self._task
+        self._task = None
+
+    async def _run(self) -> None:
+        # Poll until stop is requested, logging events along the way.
+        while not self._stop_event.is_set():
+            try:
+                self._poll_once()
+            except Exception as exc:  # pragma: no cover - defensive logging
+                logger.exception("Worker poller failed: %s", exc)
+            await self._wait_for_next_tick()
+
+    async def _wait_for_next_tick(self) -> None:
+        # Wait for either the next poll interval or a stop request.
+        try:
+            await asyncio.wait_for(self._stop_event.wait(), timeout=self._poll_interval_seconds)
+        except asyncio.TimeoutError:
+            return
+
+    def _poll_once(self) -> None:
+        # Capture a snapshot, diff it, and persist any resulting events.
+        # Timestamp both for events and snapshot cadence.
+        now_iso = _isoformat_zulu(datetime.now(timezone.utc))
+        now_monotonic = time.monotonic()
+        # Snapshot current registry and compute transitions.
+        current_snapshot = _build_snapshot(self._registry)
+        transitions = _build_transition_events(self._last_snapshot, current_snapshot, now_iso)
+
+        # Emit periodic full snapshot for recovery.
+        if self._should_emit_snapshot(now_monotonic):
+            transitions.append(events.WorkerEvent(
+                ts=now_iso,
+                type="snapshot",
+                worker_id=None,
+                data=_snapshot_payload(current_snapshot),
+            ))
+            self._last_snapshot_event_at = now_monotonic
+
+        # Persist any events in a single batch.
+        if transitions:
+            events.append_events(transitions)
+
+        # Update the in-memory snapshot for the next diff.
+        self._last_snapshot = current_snapshot
+
+    def _should_emit_snapshot(self, now_monotonic: float) -> bool:
+        # Decide whether it's time to emit a full snapshot event.
+        last = self._last_snapshot_event_at
+        if last is None:
+            return True
+        return (now_monotonic - last) >= self._snapshot_interval_seconds

claude_team_mcp/idle_detection.py

@@ -372,6 +372,7 @@ class SessionInfo:
 
     jsonl_path: Path
     session_id: str
+    agent_type: str = "claude"
 
 
 async def wait_for_any_idle(
@@ -403,7 +404,11 @@ async def wait_for_any_idle(
 
     while time.time() - start < timeout:
         for session in sessions:
-            if is_idle(session.jsonl_path, session.session_id):
+            if session.agent_type == "codex":
+                idle = is_codex_idle(session.jsonl_path)
+            else:
+                idle = is_idle(session.jsonl_path, session.session_id)
+            if idle:
                 return {
                     "idle_session_id": session.session_id,
                     "idle": True,
@@ -453,7 +458,11 @@ async def wait_for_all_idle(
         working_sessions = []
 
         for session in sessions:
-            if is_idle(session.jsonl_path, session.session_id):
+            if session.agent_type == "codex":
+                idle = is_codex_idle(session.jsonl_path)
+            else:
+                idle = is_idle(session.jsonl_path, session.session_id)
+            if idle:
                 idle_sessions.append(session.session_id)
             else:
                 working_sessions.append(session.session_id)
@@ -474,7 +483,11 @@ async def wait_for_all_idle(
     idle_sessions = []
     working_sessions = []
     for session in sessions:
-        if is_idle(session.jsonl_path, session.session_id):
+        if session.agent_type == "codex":
+            idle = is_codex_idle(session.jsonl_path)
+        else:
+            idle = is_idle(session.jsonl_path, session.session_id)
+        if idle:
             idle_sessions.append(session.session_id)
         else:
             working_sessions.append(session.session_id)

claude_team_mcp/iterm_utils.py

@@ -705,75 +705,6 @@ async def start_agent_in_session(
         )
 
 
-async def start_claude_in_session(
-    session: "ItermSession",
-    project_path: str,
-    dangerously_skip_permissions: bool = False,
-    env: Optional[dict[str, str]] = None,
-    shell_ready_timeout: float = 10.0,
-    claude_ready_timeout: float = 30.0,
-    stop_hook_marker_id: Optional[str] = None,
-) -> None:
-    """
-    Start Claude Code in an existing iTerm2 session.
-
-    Convenience wrapper around start_agent_in_session() for Claude.
-    Changes to the project directory and launches Claude Code in a single
-    atomic command (cd && claude). Waits for shell readiness before sending
-    the command, then waits for Claude's startup banner to appear.
-
-    The command used to launch Claude Code can be overridden by setting
-    the CLAUDE_TEAM_COMMAND environment variable (defaults to "claude").
-    This is useful for running alternative Claude CLI implementations
-    like "happy" or for testing purposes.
-
-    Args:
-        session: iTerm2 session to use
-        project_path: Directory to run Claude in
-        dangerously_skip_permissions: If True, start with --dangerously-skip-permissions
-        env: Optional dict of environment variables to set before running claude
-        shell_ready_timeout: Max seconds to wait for shell prompt
-        claude_ready_timeout: Max seconds to wait for Claude to start and show banner
-        stop_hook_marker_id: If provided, inject a Stop hook that logs this marker
-            to the JSONL for completion detection
-
-    Raises:
-        RuntimeError: If shell not ready or Claude fails to start within timeout
-    """
-    from .cli_backends import claude_cli
-
-    await start_agent_in_session(
-        session=session,
-        cli=claude_cli,
-        project_path=project_path,
-        dangerously_skip_permissions=dangerously_skip_permissions,
-        env=env,
-        shell_ready_timeout=shell_ready_timeout,
-        agent_ready_timeout=claude_ready_timeout,
-        stop_hook_marker_id=stop_hook_marker_id,
-    )
-
-
-# Legacy alias for backward compatibility with Claude-specific code
-# that checks for banner patterns. Uses wait_for_agent_ready with claude_cli.
-async def _wait_for_claude_ready_via_agent(
-    session: "ItermSession",
-    timeout_seconds: float = 15.0,
-    poll_interval: float = 0.2,
-    stable_count: int = 2,
-) -> bool:
-    """Internal helper - uses wait_for_agent_ready with Claude CLI."""
-    from .cli_backends import claude_cli
-
-    return await wait_for_agent_ready(
-        session=session,
-        cli=claude_cli,
-        timeout_seconds=timeout_seconds,
-        poll_interval=poll_interval,
-        stable_count=stable_count,
-    )
-
-
 # =============================================================================
 # Multi-Pane Layouts
 # =============================================================================
@@ -983,15 +914,17 @@ async def create_multi_claude_layout(
         profile_customizations=profile_customizations,
     )
 
+    from .cli_backends import claude_cli
+
     # Start Claude in all panes in parallel.
-    # Each start_claude_in_session call uses wait_for_shell_ready() internally
-    # which provides proper readiness detection, so no sleeps between starts needed.
+    # start_agent_in_session uses wait_for_shell_ready() internally, so no sleeps needed.
     async def start_claude_for_pane(pane_name: str, project_path: str) -> None:
         session = panes[pane_name]
         pane_env = project_envs.get(pane_name) if project_envs else None
         marker_id = pane_marker_ids.get(pane_name) if pane_marker_ids else None
-        await start_claude_in_session(
+        await start_agent_in_session(
             session=session,
+            cli=claude_cli,
             project_path=project_path,
             dangerously_skip_permissions=skip_permissions,
             env=pane_env,
@@ -1116,4 +1049,3 @@ async def get_window_for_session(
                     return window
     return None
 
-

claude_team_mcp/registry.py

@@ -2,7 +2,7 @@
 Session Registry for Claude Team MCP
 
 Tracks all spawned Claude Code sessions, maintaining the mapping between
-our session IDs, iTerm2 session objects, and Claude JSONL session IDs.
+our session IDs, terminal session handles, and Claude JSONL session IDs.
 """
 
 import uuid
@@ -10,12 +10,14 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import Literal, Optional
 
-if TYPE_CHECKING:
-    from iterm2.session import Session as ItermSession
-
-from .session_state import get_project_dir, parse_session
+from .session_state import (
+    find_codex_session_by_internal_id,
+    get_project_dir,
+    parse_session,
+)
+from .terminal_backends import TerminalSession
 
 # Type alias for supported agent types
 AgentType = Literal["claude", "codex"]
@@ -31,16 +33,16 @@ class TerminalId:
     tools to accept terminal IDs directly for recovery scenarios.
 
     Attributes:
-        terminal_type: Terminal emulator type ("iterm", "zed", "vscode", etc.)
+        backend_id: Terminal backend identifier ("iterm", "tmux", "zed", etc.)
         native_id: Terminal's native session ID (e.g., iTerm's UUID)
     """
 
-    terminal_type: str
+    backend_id: str
     native_id: str
 
     def __str__(self) -> str:
         """For display: 'iterm:DB29DB03-...'"""
-        return f"{self.terminal_type}:{self.native_id}"
+        return f"{self.backend_id}:{self.native_id}"
 
     @classmethod
     def from_string(cls, s: str) -> "TerminalId":
@@ -50,9 +52,8 @@ class TerminalId:
         Falls back to treating bare IDs as iTerm for backwards compatibility.
         """
         if ":" in s:
-            terminal_type, native_id = s.split(":", 1)
-            return cls(terminal_type, native_id)
-        # Assume bare ID is iTerm for backwards compatibility
+            backend_id, native_id = s.split(":", 1)
+            return cls(backend_id, native_id)
         return cls("iterm", s)
 
 
@@ -69,12 +70,12 @@ class ManagedSession:
     """
     Represents a spawned Claude Code session.
 
-    Tracks the iTerm2 session object, project path, and Claude session ID
+    Tracks terminal session metadata, project path, and Claude session ID
     discovered from the JSONL file.
     """
 
     session_id: str  # Our assigned ID (e.g., "worker-1")
-    iterm_session: "ItermSession"
+    terminal_session: TerminalSession
     project_path: str
     claude_session_id: Optional[str] = None  # Discovered from JSONL
     name: Optional[str] = None  # Optional friendly name
@@ -87,16 +88,19 @@ class ManagedSession:
     worktree_path: Optional[Path] = None  # Path to worker's git worktree if any
     main_repo_path: Optional[Path] = None  # Path to main git repo (for worktree cleanup)
 
-    # Terminal-agnostic identifier (auto-populated from iterm_session if not set)
+    # Terminal-agnostic identifier (auto-populated from terminal_session if not set)
     terminal_id: Optional[TerminalId] = None
 
     # Agent type: "claude" (default) or "codex"
     agent_type: AgentType = "claude"
 
     def __post_init__(self):
-        """Auto-populate terminal_id from iterm_session if not set."""
-        if self.terminal_id is None and self.iterm_session is not None:
-            self.terminal_id = TerminalId("iterm", self.iterm_session.session_id)
+        """Auto-populate terminal_id from terminal_session if not set."""
+        if self.terminal_id is None:
+            self.terminal_id = TerminalId(
+                self.terminal_session.backend_id,
+                self.terminal_session.native_id,
+            )
 
     def to_dict(self) -> dict:
         """Convert to dictionary for MCP tool responses."""
@@ -148,15 +152,21 @@ class ManagedSession:
         Get the path to this session's JSONL file.
 
         For Claude workers: uses marker-based discovery in ~/.claude/projects/.
-        For Codex workers: searches ~/.codex/sessions/ for matching session files.
+        For Codex workers: uses marker-based discovery in ~/.codex/sessions/.
 
         Returns:
             Path object, or None if session cannot be discovered
         """
         if self.agent_type == "codex":
-            # For Codex, search the sessions directory
             from .idle_detection import find_codex_session_file
 
+            # Prefer marker-based match, fall back to most recent for legacy sessions.
+            match = find_codex_session_by_internal_id(
+                self.session_id,
+                max_age_seconds=600,
+            )
+            if match:
+                return match.jsonl_path
             return find_codex_session_file(max_age_seconds=600)
         else:
             # For Claude, use marker-based discovery
@@ -205,8 +215,14 @@ class ManagedSession:
         if self.agent_type == "codex":
             from .idle_detection import find_codex_session_file, is_codex_idle
 
-            # Find the session file (will be discovered from ~/.codex/sessions/)
-            session_file = find_codex_session_file(max_age_seconds=600)
+            # Prefer marker-based match, fall back to most recent for legacy sessions.
+            match = find_codex_session_by_internal_id(
+                self.session_id,
+                max_age_seconds=600,
+            )
+            session_file = match.jsonl_path if match else None
+            if not session_file:
+                session_file = find_codex_session_file(max_age_seconds=600)
             if not session_file:
                 return False
             return is_codex_idle(session_file)
@@ -269,7 +285,7 @@ class SessionRegistry:
 
     def add(
         self,
-        iterm_session: "ItermSession",
+        terminal_session: TerminalSession,
         project_path: str,
         name: Optional[str] = None,
         session_id: Optional[str] = None,
@@ -278,7 +294,7 @@ class SessionRegistry:
         Add a new session to the registry.
 
         Args:
-            iterm_session: iTerm2 session object
+            terminal_session: Backend-agnostic terminal session handle
             project_path: Directory where Claude is running
             name: Optional friendly name
             session_id: Optional specific ID (auto-generated if not provided)
@@ -291,7 +307,7 @@ class SessionRegistry:
 
         session = ManagedSession(
             session_id=session_id,
-            iterm_session=iterm_session,
+            terminal_session=terminal_session,
             project_path=project_path,
             name=name,
         )
@@ -331,7 +347,8 @@ class SessionRegistry:
 
         Lookup order (most specific first):
         1. Internal session_id (e.g., "d875b833")
-        2. Terminal native ID (e.g., "DB29DB03-AA52-4FBF-879A-4DA2C5F9F823")
+        2. Terminal ID with backend prefix (e.g., "iterm:DB29DB03-..."),
+           or a bare iTerm ID for backwards compatibility
         3. Session name
 
         After MCP restart, internal IDs are lost until import. This method