crabcode 0.1.6__py3-none-win_amd64.whl

crabcode_core/team/models.py

@@ -0,0 +1,191 @@
+"""Core data models for Agent Teams.
+
+Defines Team, Teammate, Message, Task and related types.
+"""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+def _now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _new_id() -> str:
+    return str(uuid.uuid4())
+
+
+# ---------------------------------------------------------------------------
+# Enums
+# ---------------------------------------------------------------------------
+
+
+class TeamState(str, Enum):
+    FORMING = "forming"
+    ACTIVE = "active"
+    IDLE = "idle"
+    SHUTDOWN = "shutdown"
+
+
+class TeammateRole(str, Enum):
+    LEAD = "lead"
+    WORKER = "worker"
+    RESEARCHER = "researcher"
+    REVIEWER = "reviewer"
+
+
+class TeammateState(str, Enum):
+    IDLE = "idle"
+    BUSY = "busy"
+    READY = "ready"
+    CANCELLING = "cancelling"
+    SHUTDOWN = "shutdown"
+
+    @classmethod
+    def terminal_states(cls) -> set[TeammateState]:
+        return {cls.IDLE, cls.READY, cls.SHUTDOWN}
+
+    def can_transition_to(self, target: TeammateState) -> bool:
+        """Validate state machine transitions."""
+        _ALLOWED: dict[TeammateState, set[TeammateState]] = {
+            TeammateState.IDLE: {TeammateState.BUSY, TeammateState.SHUTDOWN},
+            TeammateState.BUSY: {TeammateState.READY, TeammateState.CANCELLING, TeammateState.SHUTDOWN},
+            TeammateState.READY: {TeammateState.BUSY, TeammateState.SHUTDOWN},
+            TeammateState.CANCELLING: {TeammateState.READY, TeammateState.SHUTDOWN},
+            TeammateState.SHUTDOWN: set(),  # terminal
+        }
+        return target in _ALLOWED.get(self, set())
+
+
+class TaskStatus(str, Enum):
+    PENDING = "pending"
+    CLAIMED = "claimed"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class BridgePolicy(str, Enum):
+    ALLOW_ALL = "allow_all"
+    ALLOW_TAGGED = "allow_tagged"
+    DENY = "deny"
+
+
+# ---------------------------------------------------------------------------
+# Models
+# ---------------------------------------------------------------------------
+
+
+class TeamConfig(BaseModel):
+    """Configuration for a team."""
+
+    name: str = Field(min_length=1)
+    max_teammates: int = Field(default=8, gt=0)
+    inbox_dir: str | None = None  # override default inbox storage path
+    backpressure_queue_size: int = Field(default=100, gt=0)
+    max_message_size_bytes: int = Field(default=10_000, gt=0)  # 10KB per message
+    bridge_policy: BridgePolicy = BridgePolicy.DENY
+
+
+class TeammateInfo(BaseModel):
+    """Information about a teammate within a team."""
+
+    agent_id: str
+    role: TeammateRole = TeammateRole.WORKER
+    state: TeammateState = TeammateState.IDLE
+    model_profile: str | None = None
+    name: str | None = None
+    joined_at: str = Field(default_factory=_now_iso)
+    updated_at: str = Field(default_factory=_now_iso)
+
+    def transition_to(
+        self,
+        target: TeammateState,
+        *,
+        guard: bool = False,
+        force: bool = False,
+    ) -> bool:
+        """Attempt a state transition.
+
+        Args:
+            guard: If True, skip if already in a terminal state (prevents
+                   race conditions during cleanup).
+            force: If True, bypass validation entirely (used in recovery).
+        """
+        if guard and self.state in TeammateState.terminal_states():
+            return False
+        if force or self.state.can_transition_to(target):
+            self.state = target
+            self.updated_at = _now_iso()
+            return True
+        return False
+
+
+class TeamMessage(BaseModel):
+    """A message between teammates."""
+
+    id: str = Field(default_factory=_new_id)
+    from_agent: str = ""  # agent_id of sender
+    to_agent: str = ""  # agent_id of recipient (empty = broadcast)
+    text: str = ""
+    timestamp: str = Field(default_factory=_now_iso)
+    read: bool = False
+    msg_type: str = "text"  # text | task | system | delivery_receipt
+
+    def size_bytes(self) -> int:
+        return len(self.text.encode("utf-8"))
+
+
+class TaskItem(BaseModel):
+    """A task on the shared task board."""
+
+    id: str = Field(default_factory=_new_id)
+    description: str = ""
+    assignee: str | None = None  # agent_id
+    status: TaskStatus = TaskStatus.PENDING
+    result: str = ""
+    created_at: str = Field(default_factory=_now_iso)
+    updated_at: str = Field(default_factory=_now_iso)
+
+    def claim(self, agent_id: str) -> bool:
+        """Atomically claim this task if it's still pending."""
+        if self.status != TaskStatus.PENDING:
+            return False
+        self.status = TaskStatus.CLAIMED
+        self.assignee = agent_id
+        self.updated_at = _now_iso()
+        return True
+
+    def complete(self, result: str = "") -> bool:
+        if self.status != TaskStatus.CLAIMED:
+            return False
+        self.status = TaskStatus.COMPLETED
+        self.result = result
+        self.updated_at = _now_iso()
+        return True
+
+    def fail(self, reason: str = "") -> bool:
+        if self.status != TaskStatus.CLAIMED:
+            return False
+        self.status = TaskStatus.FAILED
+        self.result = reason
+        self.updated_at = _now_iso()
+        return True
+
+
+class CrossTeamMessage(BaseModel):
+    """A message sent between teams via a bridge."""
+
+    id: str = Field(default_factory=_new_id)
+    from_team: str = ""
+    from_agent: str = ""
+    to_team: str = ""
+    to_agent: str = ""  # empty = broadcast within target team
+    text: str = ""
+    timestamp: str = Field(default_factory=_now_iso)
+    bridge_policy: BridgePolicy = BridgePolicy.ALLOW_TAGGED

crabcode_core/team/recovery.py

@@ -0,0 +1,192 @@
+"""Crash recovery for Agent Teams.
+
+When a long-lived process has stale teammate state (for example, a teammate
+marked as "busy" that is no longer running), it can be cleaned up explicitly
+through this module. Team state is currently in memory and this helper is not
+wired into gateway startup, so it does not provide cross-process/server-restart
+recovery by itself.
+
+Key design decision: NO automatic restart of interrupted teammates.
+This prevents runaway agents burning API credits after a crash.
+The human must re-engage them manually.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Awaitable, Callable
+
+from crabcode_core.logging_utils import get_logger
+from crabcode_core.team.manager import TeamManager
+from crabcode_core.team.models import TeammateState
+
+logger = get_logger(__name__)
+
+# Cancellation retry settings
+_CANCEL_RETRIES = 3
+_CANCEL_RETRY_DELAY_S = 0.12  # 120ms
+
+
+class RecoveryInfo:
+    """Result of a recovery operation."""
+
+    def __init__(
+        self,
+        team_id: str,
+        agent_id: str,
+        old_state: str,
+        new_state: str,
+    ) -> None:
+        self.team_id = team_id
+        self.agent_id = agent_id
+        self.old_state = old_state
+        self.new_state = new_state
+
+    def __repr__(self) -> str:
+        return (
+            f"RecoveryInfo(team={self.team_id}, agent={self.agent_id[:8]}, "
+            f"{self.old_state} -> {self.new_state})"
+        )
+
+
+async def recover_teams(
+    team_manager: TeamManager,
+    *,
+    inject_fn: Callable[[str, str], Awaitable[None]] | None = None,
+) -> list[RecoveryInfo]:
+    """Recover stale teammates in the supplied in-process manager.
+
+    This is not a server-startup hook: a newly-created manager has no persisted
+    team manifest to inspect, so callers must reconstruct the runtime first.
+
+    Recovery sequence:
+    1. Find all busy teammates across all teams
+    2. Force-transition them from busy -> ready
+    3. Inject a notification into the lead session
+    4. Do NOT auto-restart agents (prevents runaway)
+
+    Args:
+        team_manager: The TeamManager to recover.
+        inject_fn: Optional async function (agent_id, notification_text) to
+                   inject recovery notifications into sessions.
+
+    Returns:
+        List of RecoveryInfo for interrupted teammates.
+    """
+    interrupted = await team_manager.recover()
+
+    if not interrupted:
+        return []
+
+    recovery_infos: list[RecoveryInfo] = []
+    recovered_states = getattr(team_manager, "_last_recovered_states", {})
+    for item in interrupted:
+        state_key = (item.get("team_id", ""), item.get("agent_id", ""))
+        info = RecoveryInfo(
+            team_id=item["team_id"],
+            agent_id=item["agent_id"],
+            # ``old_state`` was not part of the original TeamManager return
+            # shape.  Keep the BUSY default for third-party managers while
+            # preserving the more useful state when the built-in manager
+            # reports a stale CANCELLING teammate.
+            old_state=item.get(
+                "old_state",
+                recovered_states.get(state_key, TeammateState.BUSY.value),
+            ),
+            new_state=TeammateState.READY.value,
+        )
+        recovery_infos.append(info)
+
+    # Inject notification into the lead session
+    if inject_fn:
+        team_ids = sorted({item["team_id"] for item in interrupted})
+        for team_id in team_ids:
+            team_agents = [
+                item["agent_id"][:8]
+                for item in interrupted
+                if item["team_id"] == team_id
+            ]
+            notification = (
+                f"[System]: Server was restarted. The following teammates in "
+                f'team "{team_id}" were interrupted and need to be resumed: '
+                f"{', '.join(team_agents)}. "
+                f"Use TeamMessage or TeamBroadcast to tell them to continue their work."
+            )
+            # Prefer an explicitly designated lead.  Older teams may not have
+            # one (the main session is not currently represented as a
+            # teammate), so retain the historical first-member fallback.
+            get_lead_agent = getattr(team_manager, "get_lead_agent", None)
+            lead_id = (
+                get_lead_agent(team_id)
+                if callable(get_lead_agent)
+                else None
+            )
+            if lead_id is None:
+                lead_id = next(
+                    (item["agent_id"] for item in interrupted if item["team_id"] == team_id),
+                    None,
+                )
+            if lead_id is None:
+                continue
+            try:
+                await inject_fn(lead_id, notification)
+            except Exception:
+                logger.warning(
+                    "Failed to inject recovery notification for %s",
+                    lead_id,
+                    exc_info=True,
+                )
+
+    logger.info("Recovered %d interrupted teammates", len(recovery_infos))
+    return recovery_infos
+
+
+async def cancel_agent_with_retry(
+    cancel_fn: Callable[[str], Awaitable[bool]],
+    agent_id: str,
+    *,
+    retries: int = _CANCEL_RETRIES,
+    delay: float = _CANCEL_RETRY_DELAY_S,
+    force_fn: Callable[[str], Awaitable[None]] | None = None,
+) -> bool:
+    """Cancel an agent with retry. If it hasn't stopped after retries,
+    force-transition as a safety net.
+
+    Args:
+        cancel_fn: Async function that cancels an agent and returns True if successful.
+        agent_id: The agent to cancel.
+        retries: Number of cancellation attempts.
+        delay: Delay between retries in seconds.
+    """
+    for attempt in range(retries):
+        try:
+            cancelled = await cancel_fn(agent_id)
+            if cancelled:
+                return True
+        except Exception:
+            logger.debug("Cancel attempt %d failed for %s", attempt + 1, agent_id, exc_info=True)
+
+        if attempt < retries - 1:
+            await asyncio.sleep(delay)
+
+    # Give the owner a chance to force-transition stale state.  The old
+    # implementation only logged and returned False, leaving a teammate stuck
+    # in cancelling/busy forever despite its docstring promising a safety net.
+    if force_fn is not None:
+        try:
+            await force_fn(agent_id)
+        except Exception:
+            logger.debug("Force-transition callback failed for %s", agent_id, exc_info=True)
+    if force_fn is not None:
+        logger.warning(
+            "Force-transitioning agent %s after %d failed cancel attempts",
+            agent_id,
+            retries,
+        )
+    else:
+        logger.warning(
+            "Unable to cancel agent %s after %d attempts; no force callback configured",
+            agent_id,
+            retries,
+        )
+    return False

crabcode_core/text_io.py

@@ -0,0 +1,52 @@
+"""UTF-8 text I/O that preserves existing file newline conventions."""
+
+from __future__ import annotations
+
+import codecs
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class Utf8Text:
+    text: str
+    newline: str | None
+    has_bom: bool
+
+
+def read_utf8_text(path: Path) -> Utf8Text:
+    raw = path.read_bytes()
+    has_bom = raw.startswith(codecs.BOM_UTF8)
+    if has_bom:
+        raw = raw[len(codecs.BOM_UTF8):]
+    text = raw.decode("utf-8")
+    return Utf8Text(text=text, newline=_detect_newline(text), has_bom=has_bom)
+
+
+def write_utf8_text(
+    path: Path,
+    text: str,
+    *,
+    newline: str | None = None,
+    has_bom: bool = False,
+) -> int:
+    if newline is not None:
+        text = normalize_newlines(text, newline)
+    raw = text.encode("utf-8")
+    if has_bom:
+        raw = codecs.BOM_UTF8 + raw
+    return path.write_bytes(raw)
+
+
+def normalize_newlines(text: str, newline: str) -> str:
+    normalized = text.replace("\r\n", "\n").replace("\r", "\n")
+    return normalized if newline == "\n" else normalized.replace("\n", newline)
+
+
+def _detect_newline(text: str) -> str | None:
+    crlf = text.count("\r\n")
+    lf = text.count("\n") - crlf
+    cr = text.count("\r") - crlf
+    counts = [(crlf, "\r\n"), (lf, "\n"), (cr, "\r")]
+    count, newline = max(counts, key=lambda item: item[0])
+    return newline if count else None

crabcode_core/tools/__init__.py

@@ -0,0 +1,112 @@
+"""Tool system - built-in tools and tool registry."""
+
+from crabcode_core.types.tool import Tool, ToolResult, ToolContext
+
+
+def get_default_tools() -> list[Tool]:
+    """Return the default set of built-in tools."""
+    from crabcode_core.tools.agent import (
+        AgentCancelTool,
+        AgentSendInputTool,
+        AgentStatusTool,
+        AgentWaitTool,
+    )
+    from crabcode_core.tools.ask_user import AskUserTool
+    from crabcode_core.tools.apply_patch import ApplyPatchTool
+    from crabcode_core.tools.checklist import ChecklistTool
+    from crabcode_core.tools.checkpoint import CheckpointTool
+    from crabcode_core.tools.computer_use import ComputerUseTool
+    from crabcode_core.tools.bash import BashTool
+    from crabcode_core.tools.browser import BrowserTool
+    from crabcode_core.tools.image import ImageTool
+    from crabcode_core.tools.image_generate import ImageGenerateTool
+    from crabcode_core.tools.file_read import FileReadTool
+    from crabcode_core.tools.file_edit import FileEditTool
+    from crabcode_core.tools.file_write import FileWriteTool
+    from crabcode_core.tools.grep import GrepTool
+    from crabcode_core.tools.goal import CreateGoalTool, GetGoalTool, UpdateGoalTool
+    from crabcode_core.tools.glob import GlobTool
+    from crabcode_core.tools.lint import LintTool
+    from crabcode_core.tools.memory import MemoryTool
+    from crabcode_core.tools.monitor import MonitorManager, MonitorTool, TaskListTool, TaskStopTool
+    from crabcode_core.tools.peer import ListAgentsTool, SendMessageTool
+    from crabcode_core.tools.revert import RevertTool
+    from crabcode_core.tools.switch_mode import SwitchModeTool
+    from crabcode_core.tools.schedule import (
+        ScheduleCancelTool,
+        ScheduleCreateTool,
+        ScheduleListTool,
+        SchedulePauseTool,
+        ScheduleResumeTool,
+        ScheduleRunTool,
+        ScheduleStatusTool,
+    )
+    from crabcode_core.tools.team import (
+        TeamBroadcastTool,
+        TeamCreateTool,
+        TeamMessageTool,
+        TeamShutdownTool,
+        TeamSpawnTool,
+        TeamStatusTool,
+        TeamTaskAddTool,
+        TeamTaskClaimTool,
+        TeamTaskCompleteTool,
+    )
+    from crabcode_core.tools.web_search import WebSearchTool
+    from crabcode_core.tools.tool_search import ToolSearchTool
+
+    monitor_manager = MonitorManager()
+
+    return [
+        ToolSearchTool(),
+        BashTool(),
+        FileReadTool(),
+        ApplyPatchTool(),
+        FileEditTool(),
+        FileWriteTool(),
+        GrepTool(),
+        GlobTool(),
+        WebSearchTool(),
+        BrowserTool(),
+        ComputerUseTool(),
+        ImageTool(),
+        ImageGenerateTool(),
+        LintTool(),
+        MemoryTool(),
+        MonitorTool(monitor_manager),
+        TaskListTool(monitor_manager),
+        TaskStopTool(monitor_manager),
+        AgentStatusTool(),
+        AgentWaitTool(),
+        AgentCancelTool(),
+        AgentSendInputTool(),
+        ListAgentsTool(),
+        SendMessageTool(),
+        AskUserTool(),
+        CheckpointTool(),
+        ChecklistTool(),
+        CreateGoalTool(),
+        GetGoalTool(),
+        UpdateGoalTool(),
+        SwitchModeTool(),
+        TeamCreateTool(),
+        TeamSpawnTool(),
+        TeamMessageTool(),
+        TeamBroadcastTool(),
+        TeamStatusTool(),
+        TeamTaskAddTool(),
+        TeamTaskClaimTool(),
+        TeamTaskCompleteTool(),
+        TeamShutdownTool(),
+        RevertTool(),
+        ScheduleCreateTool(),
+        ScheduleListTool(),
+        ScheduleCancelTool(),
+        ScheduleStatusTool(),
+        SchedulePauseTool(),
+        ScheduleResumeTool(),
+        ScheduleRunTool(),
+    ]
+
+
+__all__ = ["Tool", "ToolResult", "ToolContext", "get_default_tools"]

crabcode_core/tools/_input_helpers.py

@@ -0,0 +1,39 @@
+"""Helpers for normalizing model-provided tool arguments."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from crabcode_core.types.message import Message, MessageRole, TextBlock
+
+
+def first_non_empty_str(tool_input: dict[str, Any], keys: tuple[str, ...]) -> str:
+    """Return the first stripped non-empty string among the given keys."""
+    for key in keys:
+        raw = tool_input.get(key)
+        if isinstance(raw, str) and raw.strip():
+            return raw.strip()
+    return ""
+
+
+def latest_user_text_for_agent_fallback(messages: list[Message]) -> str:
+    """Text from the most recent user message that contains user-typed text (TextBlock).
+
+    Skips user messages that only carry tool results, so a missing ``prompt`` on
+    Agent can still map to the user's actual question when the model omits args.
+    """
+    for msg in reversed(messages):
+        if msg.role != MessageRole.USER:
+            continue
+        if isinstance(msg.content, str):
+            t = msg.content.strip()
+            if t:
+                return t
+        parts: list[str] = []
+        for block in msg.content:
+            if isinstance(block, TextBlock):
+                parts.append(block.text)
+        text = "".join(parts).strip()
+        if text:
+            return text
+    return ""