axion-code 1.0.0__py3-none-any.whl

axion/runtime/tasks.py

@@ -0,0 +1,258 @@
+"""Task packet, registry, team assignment, and cron scheduling.
+
+Maps to: rust/crates/runtime/src/task_packet.rs + team/cron registries
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+
+# ---------------------------------------------------------------------------
+# Task packet
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TaskPacket:
+    """Specification for an autonomous task."""
+
+    objective: str
+    scope: str = ""
+    repo: str = ""
+    branch_policy: str = "feature-branch"
+    acceptance_tests: list[str] = field(default_factory=list)
+    commit_policy: str = "atomic"
+    reporting_contract: str = ""
+    escalation_policy: str = "alert_human"
+    priority: int = 0
+    tags: list[str] = field(default_factory=list)
+    assigned_team: str | None = None
+    cron_schedule: str | None = None
+    created_at_ms: int = field(default_factory=lambda: int(time.time() * 1000))
+
+
+@dataclass
+class TaskPacketValidationError:
+    errors: list[str]
+
+
+def validate_packet(packet: TaskPacket) -> TaskPacketValidationError | None:
+    """Validate a task packet for completeness."""
+    errors: list[str] = []
+    if not packet.objective:
+        errors.append("objective is required")
+    if not packet.scope:
+        errors.append("scope is required")
+    if packet.branch_policy not in ("feature-branch", "direct", "trunk"):
+        errors.append(f"invalid branch_policy: {packet.branch_policy}")
+    if packet.commit_policy not in ("atomic", "squash", "incremental"):
+        errors.append(f"invalid commit_policy: {packet.commit_policy}")
+    if packet.escalation_policy not in ("alert_human", "log_and_continue", "abort"):
+        errors.append(f"invalid escalation_policy: {packet.escalation_policy}")
+    return TaskPacketValidationError(errors=errors) if errors else None
+
+
+# ---------------------------------------------------------------------------
+# Task registry
+# ---------------------------------------------------------------------------
+
+class TaskStatus:
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+@dataclass
+class TaskEntry:
+    """A task in the registry with tracking metadata."""
+
+    task_id: str
+    packet: TaskPacket
+    status: str = TaskStatus.PENDING
+    worker_id: str | None = None
+    started_at_ms: int | None = None
+    completed_at_ms: int | None = None
+    result: str = ""
+    error: str = ""
+
+
+class TaskRegistry:
+    """In-memory task lifecycle registry with team assignment and scheduling.
+
+    Maps to: rust/crates/runtime/src/task_packet.rs (registry)
+    """
+
+    def __init__(self) -> None:
+        self._tasks: dict[str, TaskEntry] = {}
+        self._counter = 0
+
+    def create(self, packet: TaskPacket) -> str:
+        """Create a new task and return its ID."""
+        self._counter += 1
+        task_id = f"task-{self._counter:04d}"
+        self._tasks[task_id] = TaskEntry(task_id=task_id, packet=packet)
+        return task_id
+
+    def get(self, task_id: str) -> TaskEntry | None:
+        return self._tasks.get(task_id)
+
+    def all_tasks(self) -> list[TaskEntry]:
+        return list(self._tasks.values())
+
+    def pending_tasks(self) -> list[TaskEntry]:
+        return [t for t in self._tasks.values() if t.status == TaskStatus.PENDING]
+
+    def running_tasks(self) -> list[TaskEntry]:
+        return [t for t in self._tasks.values() if t.status == TaskStatus.RUNNING]
+
+    def start_task(self, task_id: str, worker_id: str) -> bool:
+        task = self._tasks.get(task_id)
+        if task is None or task.status != TaskStatus.PENDING:
+            return False
+        task.status = TaskStatus.RUNNING
+        task.worker_id = worker_id
+        task.started_at_ms = int(time.time() * 1000)
+        return True
+
+    def complete_task(self, task_id: str, result: str = "") -> bool:
+        task = self._tasks.get(task_id)
+        if task is None or task.status != TaskStatus.RUNNING:
+            return False
+        task.status = TaskStatus.COMPLETED
+        task.result = result
+        task.completed_at_ms = int(time.time() * 1000)
+        return True
+
+    def fail_task(self, task_id: str, error: str = "") -> bool:
+        task = self._tasks.get(task_id)
+        if task is None or task.status != TaskStatus.RUNNING:
+            return False
+        task.status = TaskStatus.FAILED
+        task.error = error
+        task.completed_at_ms = int(time.time() * 1000)
+        return True
+
+    def cancel_task(self, task_id: str) -> bool:
+        task = self._tasks.get(task_id)
+        if task is None or task.status in (TaskStatus.COMPLETED, TaskStatus.CANCELLED):
+            return False
+        task.status = TaskStatus.CANCELLED
+        task.completed_at_ms = int(time.time() * 1000)
+        return True
+
+    def remove(self, task_id: str) -> bool:
+        return self._tasks.pop(task_id, None) is not None
+
+    def summary(self) -> dict[str, int]:
+        counts: dict[str, int] = {}
+        for task in self._tasks.values():
+            counts[task.status] = counts.get(task.status, 0) + 1
+        return counts
+
+
+# ---------------------------------------------------------------------------
+# Team registry
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Team:
+    """A named group of workers for task assignment."""
+
+    name: str
+    worker_ids: list[str] = field(default_factory=list)
+    max_concurrent: int = 1
+    tags: list[str] = field(default_factory=list)
+
+
+class TeamRegistry:
+    """Registry of teams for task assignment."""
+
+    def __init__(self) -> None:
+        self._teams: dict[str, Team] = {}
+
+    def register(self, team: Team) -> None:
+        self._teams[team.name] = team
+
+    def get(self, name: str) -> Team | None:
+        return self._teams.get(name)
+
+    def all_teams(self) -> list[Team]:
+        return list(self._teams.values())
+
+    def assign_task(self, task: TaskEntry) -> str | None:
+        """Find an available team for a task based on tags."""
+        for team in self._teams.values():
+            if task.packet.assigned_team and task.packet.assigned_team != team.name:
+                continue
+            # Check tag match
+            if task.packet.tags and not any(t in team.tags for t in task.packet.tags):
+                continue
+            return team.name
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Cron registry
+# ---------------------------------------------------------------------------
+
+@dataclass
+class CronEntry:
+    """A scheduled recurring task."""
+
+    cron_id: str
+    schedule: str  # Cron expression (e.g. "*/5 * * * *")
+    packet: TaskPacket
+    enabled: bool = True
+    last_run_ms: int = 0
+    next_run_ms: int = 0
+    run_count: int = 0
+
+
+class CronRegistry:
+    """Registry of cron-scheduled tasks."""
+
+    def __init__(self) -> None:
+        self._entries: dict[str, CronEntry] = {}
+        self._counter = 0
+
+    def create(self, schedule: str, packet: TaskPacket) -> str:
+        self._counter += 1
+        cron_id = f"cron-{self._counter:04d}"
+        self._entries[cron_id] = CronEntry(
+            cron_id=cron_id, schedule=schedule, packet=packet,
+        )
+        return cron_id
+
+    def get(self, cron_id: str) -> CronEntry | None:
+        return self._entries.get(cron_id)
+
+    def all_entries(self) -> list[CronEntry]:
+        return list(self._entries.values())
+
+    def enabled_entries(self) -> list[CronEntry]:
+        return [e for e in self._entries.values() if e.enabled]
+
+    def enable(self, cron_id: str) -> bool:
+        entry = self._entries.get(cron_id)
+        if entry is None:
+            return False
+        entry.enabled = True
+        return True
+
+    def disable(self, cron_id: str) -> bool:
+        entry = self._entries.get(cron_id)
+        if entry is None:
+            return False
+        entry.enabled = False
+        return True
+
+    def remove(self, cron_id: str) -> bool:
+        return self._entries.pop(cron_id, None) is not None
+
+    def record_run(self, cron_id: str) -> None:
+        entry = self._entries.get(cron_id)
+        if entry:
+            entry.last_run_ms = int(time.time() * 1000)
+            entry.run_count += 1

axion/runtime/usage.py

@@ -0,0 +1,241 @@
+"""Token usage tracking and cost estimation.
+
+Maps to: rust/crates/runtime/src/usage.rs
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+# Default pricing (Sonnet 4 tier - $3 input / $15 output per million tokens)
+DEFAULT_INPUT_COST_PER_MILLION = 3.0
+DEFAULT_OUTPUT_COST_PER_MILLION = 15.0
+DEFAULT_CACHE_CREATION_COST_PER_MILLION = 3.75
+DEFAULT_CACHE_READ_COST_PER_MILLION = 0.3
+
+
+@dataclass
+class ModelPricing:
+    """Per-million-token pricing for cost estimation."""
+
+    input_cost_per_million: float = DEFAULT_INPUT_COST_PER_MILLION
+    output_cost_per_million: float = DEFAULT_OUTPUT_COST_PER_MILLION
+    cache_creation_cost_per_million: float = DEFAULT_CACHE_CREATION_COST_PER_MILLION
+    cache_read_cost_per_million: float = DEFAULT_CACHE_READ_COST_PER_MILLION
+
+    @classmethod
+    def default_sonnet_tier(cls) -> ModelPricing:
+        return cls()
+
+
+@dataclass
+class TokenUsage:
+    """Token counters accumulated for a conversation turn or session."""
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_creation_input_tokens: int = 0
+    cache_read_input_tokens: int = 0
+
+    def total_tokens(self) -> int:
+        return (
+            self.input_tokens
+            + self.output_tokens
+            + self.cache_creation_input_tokens
+            + self.cache_read_input_tokens
+        )
+
+    def estimate_cost_usd(self) -> UsageCostEstimate:
+        return self.estimate_cost_usd_with_pricing(ModelPricing.default_sonnet_tier())
+
+    def estimate_cost_usd_with_pricing(self, pricing: ModelPricing) -> UsageCostEstimate:
+        return UsageCostEstimate(
+            input_cost_usd=_cost_for_tokens(self.input_tokens, pricing.input_cost_per_million),
+            output_cost_usd=_cost_for_tokens(
+                self.output_tokens, pricing.output_cost_per_million
+            ),
+            cache_creation_cost_usd=_cost_for_tokens(
+                self.cache_creation_input_tokens,
+                pricing.cache_creation_cost_per_million,
+            ),
+            cache_read_cost_usd=_cost_for_tokens(
+                self.cache_read_input_tokens,
+                pricing.cache_read_cost_per_million,
+            ),
+        )
+
+    def __iadd__(self, other: TokenUsage) -> TokenUsage:
+        self.input_tokens += other.input_tokens
+        self.output_tokens += other.output_tokens
+        self.cache_creation_input_tokens += other.cache_creation_input_tokens
+        self.cache_read_input_tokens += other.cache_read_input_tokens
+        return self
+
+    def summary_lines(self, label: str, model: str | None = None) -> list[str]:
+        pricing = pricing_for_model(model) if model else None
+        cost = (
+            self.estimate_cost_usd_with_pricing(pricing)
+            if pricing
+            else self.estimate_cost_usd()
+        )
+        model_suffix = f" model={model}" if model else ""
+        return [
+            f"{label}: total_tokens={self.total_tokens()} "
+            f"input={self.input_tokens} output={self.output_tokens} "
+            f"cache_write={self.cache_creation_input_tokens} "
+            f"cache_read={self.cache_read_input_tokens} "
+            f"estimated_cost={format_usd(cost.total_cost_usd())}{model_suffix}",
+            f"  cost breakdown: input={format_usd(cost.input_cost_usd)} "
+            f"output={format_usd(cost.output_cost_usd)} "
+            f"cache_write={format_usd(cost.cache_creation_cost_usd)} "
+            f"cache_read={format_usd(cost.cache_read_cost_usd)}",
+        ]
+
+
+@dataclass
+class UsageCostEstimate:
+    """Estimated dollar cost derived from a TokenUsage sample."""
+
+    input_cost_usd: float = 0.0
+    output_cost_usd: float = 0.0
+    cache_creation_cost_usd: float = 0.0
+    cache_read_cost_usd: float = 0.0
+
+    def total_cost_usd(self) -> float:
+        return (
+            self.input_cost_usd
+            + self.output_cost_usd
+            + self.cache_creation_cost_usd
+            + self.cache_read_cost_usd
+        )
+
+
+def pricing_for_model(model: str | None) -> ModelPricing | None:
+    """Returns pricing metadata for a known model alias or family."""
+    if model is None:
+        return None
+    normalized = model.lower()
+    if "haiku" in normalized:
+        return ModelPricing(
+            input_cost_per_million=1.0,
+            output_cost_per_million=5.0,
+            cache_creation_cost_per_million=1.25,
+            cache_read_cost_per_million=0.1,
+        )
+    if "opus" in normalized:
+        return ModelPricing(
+            input_cost_per_million=5.0,   # Opus 4.6: $5/MTok input
+            output_cost_per_million=25.0,  # Opus 4.6: $25/MTok output
+            cache_creation_cost_per_million=6.25,
+            cache_read_cost_per_million=0.5,
+        )
+    if "sonnet" in normalized:
+        return ModelPricing.default_sonnet_tier()
+    # OpenAI models — ordered specific to general
+    # GPT-4.1 series (1M context)
+    if "gpt-4.1-nano" in normalized:
+        return ModelPricing(
+            input_cost_per_million=0.10, output_cost_per_million=0.40,
+            cache_creation_cost_per_million=0.05, cache_read_cost_per_million=0.025,
+        )
+    if "gpt-4.1-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=0.40, output_cost_per_million=1.60,
+            cache_creation_cost_per_million=0.20, cache_read_cost_per_million=0.10,
+        )
+    if "gpt-4.1" in normalized:
+        return ModelPricing(
+            input_cost_per_million=2.0, output_cost_per_million=8.0,
+            cache_creation_cost_per_million=1.0, cache_read_cost_per_million=0.50,
+        )
+    # GPT-4o series
+    if "gpt-4o-mini" in normalized or "4o-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=0.15, output_cost_per_million=0.60,
+            cache_creation_cost_per_million=0.075, cache_read_cost_per_million=0.075,
+        )
+    if "gpt-4o" in normalized or "4o" in normalized:
+        return ModelPricing(
+            input_cost_per_million=2.50, output_cost_per_million=10.0,
+            cache_creation_cost_per_million=1.25, cache_read_cost_per_million=1.25,
+        )
+    # o-series reasoning
+    if "o4-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=1.10, output_cost_per_million=4.40,
+            cache_creation_cost_per_million=0.55, cache_read_cost_per_million=0.275,
+        )
+    if "o3-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=1.10, output_cost_per_million=4.40,
+            cache_creation_cost_per_million=0.55, cache_read_cost_per_million=0.55,
+        )
+    if normalized in ("o1", "o3", "o1-pro"):
+        return ModelPricing(
+            input_cost_per_million=15.0, output_cost_per_million=60.0,
+            cache_creation_cost_per_million=7.5, cache_read_cost_per_million=7.5,
+        )
+    if "o1-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=1.10, output_cost_per_million=4.40,
+            cache_creation_cost_per_million=0.55, cache_read_cost_per_million=0.55,
+        )
+    # Codex (Responses API — agent-tuned coding models)
+    if "codex-mini" in normalized:
+        return ModelPricing(
+            input_cost_per_million=0.25, output_cost_per_million=2.0,
+            cache_creation_cost_per_million=0.125, cache_read_cost_per_million=0.025,
+        )
+    if "codex" in normalized:
+        return ModelPricing(
+            input_cost_per_million=1.25, output_cost_per_million=10.0,
+            cache_creation_cost_per_million=0.625, cache_read_cost_per_million=0.125,
+        )
+    # xAI
+    if "grok" in normalized:
+        return ModelPricing(
+            input_cost_per_million=5.0,
+            output_cost_per_million=15.0,
+            cache_creation_cost_per_million=2.5,
+            cache_read_cost_per_million=2.5,
+        )
+    return None
+
+
+def _cost_for_tokens(tokens: int, cost_per_million: float) -> float:
+    return (tokens / 1_000_000) * cost_per_million
+
+
+def format_usd(value: float) -> str:
+    """Format a USD value with 4 decimal places."""
+    return f"${value:.4f}"
+
+
+# ---------------------------------------------------------------------------
+# Usage tracker (accumulates across turns)
+# ---------------------------------------------------------------------------
+
+@dataclass
+class UsageTracker:
+    """Accumulates token usage across multiple turns."""
+
+    total: TokenUsage = field(default_factory=TokenUsage)
+    turn_count: int = 0
+    model: str | None = None
+
+    def record_turn(self, usage: TokenUsage) -> None:
+        self.total += usage
+        self.turn_count += 1
+
+    def summary_lines(self) -> list[str]:
+        return self.total.summary_lines("Session total", self.model)
+
+    @classmethod
+    def from_session(cls, session: Any) -> UsageTracker:
+        """Build tracker from existing session messages' usage data."""
+        tracker = cls()
+        for msg in getattr(session, "messages", []):
+            if msg.usage is not None:
+                tracker.record_turn(msg.usage)
+        return tracker

axion/runtime/workers.py

@@ -0,0 +1,186 @@
+"""Worker state machine and registry with full lifecycle management.
+
+Maps to: rust/crates/runtime/src/worker_boot.rs
+"""
+
+from __future__ import annotations
+
+import enum
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerStatus(enum.Enum):
+    SPAWNING = "spawning"
+    TRUST_REQUIRED = "trust_required"
+    READY_FOR_PROMPT = "ready_for_prompt"
+    RUNNING = "running"
+    FINISHED = "finished"
+    FAILED = "failed"
+
+
+class WorkerFailureKind(enum.Enum):
+    TRUST_GATE = "trust_gate"
+    PROMPT_DELIVERY = "prompt_delivery"
+    PROTOCOL = "protocol"
+    PROVIDER = "provider"
+
+
+class WorkerTrustResolution(enum.Enum):
+    AUTO_ALLOWLISTED = "auto_allowlisted"
+    MANUAL_APPROVAL = "manual_approval"
+
+
+class WorkerPromptTarget(enum.Enum):
+    SHELL = "shell"
+    WRONG_TARGET = "wrong_target"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class WorkerEvent:
+    """Timestamped event in a worker's lifecycle."""
+
+    seq: int
+    kind: str
+    status: WorkerStatus
+    detail: str = ""
+    payload: dict[str, Any] | None = None
+    timestamp_ms: int = field(default_factory=lambda: int(time.time() * 1000))
+
+
+@dataclass
+class Worker:
+    """Represents a single agent worker instance.
+
+    Implements the full state machine:
+    SPAWNING → TRUST_REQUIRED → READY_FOR_PROMPT → RUNNING → FINISHED/FAILED
+    """
+
+    worker_id: str = field(default_factory=lambda: f"w-{uuid.uuid4().hex[:8]}")
+    cwd: str = ""
+    status: WorkerStatus = WorkerStatus.SPAWNING
+    trust_auto_resolve: bool = False
+    trust_gate_cleared: bool = False
+    auto_recover_prompt_misdelivery: bool = True
+    prompt_delivery_attempts: int = 0
+    max_prompt_delivery_attempts: int = 3
+    events: list[WorkerEvent] = field(default_factory=list)
+
+    def transition(self, new_status: WorkerStatus, detail: str = "", **payload: Any) -> None:
+        """Transition to a new status with event logging."""
+        event = WorkerEvent(
+            seq=len(self.events) + 1,
+            kind=f"transition_{new_status.value}",
+            status=new_status,
+            detail=detail,
+            payload=payload if payload else None,
+        )
+        self.events.append(event)
+        self.status = new_status
+        logger.debug("Worker %s → %s: %s", self.worker_id, new_status.value, detail)
+
+    def resolve_trust(self, resolution: WorkerTrustResolution) -> None:
+        """Resolve the trust gate and advance to ready state."""
+        if self.status != WorkerStatus.TRUST_REQUIRED:
+            logger.warning("Cannot resolve trust in status %s", self.status.value)
+            return
+        self.trust_gate_cleared = True
+        self.transition(
+            WorkerStatus.READY_FOR_PROMPT,
+            detail=f"Trust resolved via {resolution.value}",
+        )
+
+    def deliver_prompt(self, target: WorkerPromptTarget = WorkerPromptTarget.SHELL) -> bool:
+        """Attempt to deliver the prompt to the worker."""
+        self.prompt_delivery_attempts += 1
+
+        if target == WorkerPromptTarget.WRONG_TARGET:
+            if (
+                self.auto_recover_prompt_misdelivery
+                and self.prompt_delivery_attempts < self.max_prompt_delivery_attempts
+            ):
+                self.transition(
+                    WorkerStatus.READY_FOR_PROMPT,
+                    detail=f"Prompt misdelivered (attempt {self.prompt_delivery_attempts}), retrying",
+                )
+                return False
+            self.transition(
+                WorkerStatus.FAILED,
+                detail="Prompt misdelivery exceeded max attempts",
+            )
+            return False
+
+        if target == WorkerPromptTarget.UNKNOWN:
+            self.transition(WorkerStatus.FAILED, detail="Unknown prompt target")
+            return False
+
+        self.transition(WorkerStatus.RUNNING, detail="Prompt delivered to shell")
+        return True
+
+    def finish(self, detail: str = "Completed successfully") -> None:
+        self.transition(WorkerStatus.FINISHED, detail=detail)
+
+    def fail(self, kind: WorkerFailureKind, detail: str = "") -> None:
+        self.transition(WorkerStatus.FAILED, detail=f"{kind.value}: {detail}")
+
+    def restart(self) -> None:
+        """Restart the worker from SPAWNING state."""
+        self.trust_gate_cleared = False
+        self.prompt_delivery_attempts = 0
+        self.transition(WorkerStatus.SPAWNING, detail="Restarted")
+
+    @property
+    def is_active(self) -> bool:
+        return self.status in (WorkerStatus.RUNNING, WorkerStatus.READY_FOR_PROMPT)
+
+    @property
+    def is_terminal(self) -> bool:
+        return self.status in (WorkerStatus.FINISHED, WorkerStatus.FAILED)
+
+
+class WorkerRegistry:
+    """Manages multiple worker instances.
+
+    Maps to: rust/crates/runtime/src/worker_boot.rs::WorkerRegistry
+    """
+
+    def __init__(self) -> None:
+        self._workers: dict[str, Worker] = {}
+
+    def spawn(self, cwd: str = "", **kwargs: Any) -> Worker:
+        """Create and register a new worker."""
+        worker = Worker(cwd=cwd, **kwargs)
+        self._workers[worker.worker_id] = worker
+        worker.transition(WorkerStatus.SPAWNING, detail="Worker spawned")
+        return worker
+
+    def register(self, worker: Worker) -> None:
+        self._workers[worker.worker_id] = worker
+
+    def get(self, worker_id: str) -> Worker | None:
+        return self._workers.get(worker_id)
+
+    def all_workers(self) -> list[Worker]:
+        return list(self._workers.values())
+
+    def active_workers(self) -> list[Worker]:
+        return [w for w in self._workers.values() if w.is_active]
+
+    def finished_workers(self) -> list[Worker]:
+        return [w for w in self._workers.values() if w.is_terminal]
+
+    def remove(self, worker_id: str) -> bool:
+        return self._workers.pop(worker_id, None) is not None
+
+    def summary(self) -> dict[str, int]:
+        """Count workers by status."""
+        counts: dict[str, int] = {}
+        for worker in self._workers.values():
+            counts[worker.status.value] = counts.get(worker.status.value, 0) + 1
+        return counts