pulse-coding-agent 0.1.0__py3-none-any.whl

pulse/telemetry/__init__.py

@@ -0,0 +1,25 @@
+from pulse.telemetry.cost_tracker import (
+    BudgetExceededError,
+    CostTracker,
+    ModelPricing,
+    UsageRecord,
+)
+from pulse.telemetry.logger import (
+    MetricEvent,
+    TelemetryLogger,
+    correlation_scope,
+    get_correlation_id,
+    set_correlation_id,
+)
+
+__all__ = [
+    "BudgetExceededError",
+    "CostTracker",
+    "MetricEvent",
+    "ModelPricing",
+    "TelemetryLogger",
+    "UsageRecord",
+    "correlation_scope",
+    "get_correlation_id",
+    "set_correlation_id",
+]

pulse/telemetry/cost_tracker.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+
+class BudgetExceededError(Exception):
+    """Raised when token or cost budget limit is exceeded."""
+
+
+@dataclass
+class ModelPricing:
+    prompt_price_per_1k: float
+    completion_price_per_1k: float
+
+
+@dataclass
+class UsageRecord:
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    estimated_cost: float
+
+
+class CostTracker:
+    """Tracks prompt/completion token usage across models (OpenAI, Gemini, OpenRouter).
+
+    Enforces maximum daily/session token and cost budgets.
+    """
+
+    DEFAULT_PRICING: ClassVar[dict[str, ModelPricing]] = {
+        "gpt-4o": ModelPricing(0.0025, 0.0100),
+        "gpt-4o-mini": ModelPricing(0.00015, 0.0006),
+        "gemini-2.0-flash": ModelPricing(0.0001, 0.0004),
+        "openrouter/auto": ModelPricing(0.0010, 0.0030),
+        "default": ModelPricing(0.0015, 0.0050),
+    }
+
+    def __init__(
+        self,
+        max_session_tokens: int | None = None,
+        max_session_cost: float | None = None,
+        pricing: dict[str, ModelPricing] | None = None,
+    ) -> None:
+        self.max_session_tokens = max_session_tokens
+        self.max_session_cost = max_session_cost
+        self.pricing = pricing or self.DEFAULT_PRICING
+
+        self._total_prompt_tokens: int = 0
+        self._total_completion_tokens: int = 0
+        self._total_cost: float = 0.0
+
+    @property
+    def total_tokens(self) -> int:
+        return self._total_prompt_tokens + self._total_completion_tokens
+
+    @property
+    def total_cost(self) -> float:
+        return round(self._total_cost, 6)
+
+    def record_usage(self, model: str, prompt_tokens: int, completion_tokens: int) -> UsageRecord:
+        added_tokens = prompt_tokens + completion_tokens
+
+        if self.max_session_tokens is not None and (self.total_tokens + added_tokens) > self.max_session_tokens:
+            raise BudgetExceededError(
+                f"Session token budget exceeded: {self.total_tokens + added_tokens} > {self.max_session_tokens}"
+            )
+
+        model_key = model.lower()
+        pricing = self.pricing.get(model_key) or self.pricing.get("default", ModelPricing(0.0015, 0.0050))
+
+        prompt_cost = (prompt_tokens / 1000.0) * pricing.prompt_price_per_1k
+        completion_cost = (completion_tokens / 1000.0) * pricing.completion_price_per_1k
+        cost = prompt_cost + completion_cost
+
+        if self.max_session_cost is not None and (self._total_cost + cost) > self.max_session_cost:
+            raise BudgetExceededError(
+                f"Session cost budget exceeded: ${self._total_cost + cost:.4f} > ${self.max_session_cost:.4f}"
+            )
+
+        self._total_prompt_tokens += prompt_tokens
+        self._total_completion_tokens += completion_tokens
+        self._total_cost += cost
+
+        return UsageRecord(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=added_tokens,
+            estimated_cost=round(cost, 6),
+        )
+
+    def reset(self) -> None:
+        self._total_prompt_tokens = 0
+        self._total_completion_tokens = 0
+        self._total_cost = 0.0

pulse/telemetry/logger.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import contextvars
+import json
+import re
+import uuid
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from pulse.sandbox.secrets import SecretScrubber
+
+_CORRELATION_ID: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "pulse_correlation_id", default=None
+)
+_SAFE_CORRELATION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
+
+
+def set_correlation_id(value: object | None = None) -> str:
+    candidate = str(value).strip() if value is not None else ""
+    correlation_id = (
+        candidate if _SAFE_CORRELATION_ID.fullmatch(candidate) else uuid.uuid4().hex
+    )
+    _CORRELATION_ID.set(correlation_id)
+    return correlation_id
+
+
+def get_correlation_id() -> str:
+    return _CORRELATION_ID.get() or set_correlation_id()
+
+
+@contextmanager
+def correlation_scope(value: object | None = None):
+    correlation_id = str(value).strip() if value is not None else ""
+    if not _SAFE_CORRELATION_ID.fullmatch(correlation_id):
+        correlation_id = uuid.uuid4().hex
+    token = _CORRELATION_ID.set(correlation_id)
+    try:
+        yield correlation_id
+    finally:
+        _CORRELATION_ID.reset(token)
+
+
+@dataclass
+class MetricEvent:
+    schema_version: int = field(default=1, init=False)
+    timestamp: str
+    correlation_id: str
+    event_type: str
+    step: int | None
+    duration_ms: float | None
+    metadata: dict[str, Any]
+
+
+class TelemetryLogger:
+    """Lightweight structured logger for agent step telemetry and execution metrics."""
+
+    def __init__(self, log_path: Path | None = None) -> None:
+        self.log_path = log_path
+        self._scrubber = SecretScrubber()
+
+    def add_secret(self, secret: str | None) -> None:
+        if secret:
+            self._scrubber.add_secret(secret)
+
+    def _sanitize(self, value: Any) -> Any:
+        if isinstance(value, str):
+            return self._scrubber.redact(value)
+        if isinstance(value, dict):
+            return {str(key): self._sanitize(item) for key, item in value.items()}
+        if isinstance(value, (list, tuple)):
+            return [self._sanitize(item) for item in value]
+        return value
+
+    def log_event(
+        self,
+        event_type: str,
+        step: int | None = None,
+        duration_ms: float | None = None,
+        **metadata: Any,
+    ) -> MetricEvent:
+        event = MetricEvent(
+            timestamp=datetime.now(UTC).isoformat(),
+            correlation_id=get_correlation_id(),
+            event_type=event_type,
+            step=step,
+            duration_ms=duration_ms,
+            metadata=self._sanitize(metadata),
+        )
+
+        if self.log_path:
+            self.log_path.parent.mkdir(parents=True, exist_ok=True)
+            with self.log_path.open("a", encoding="utf-8") as handle:
+                handle.write(json.dumps(asdict(event), separators=(",", ":")) + "\n")
+
+        return event
+
+    def log_step_execution(
+        self, step: int, tool_name: str | None, duration_ms: float, success: bool, **kwargs: Any
+    ) -> MetricEvent:
+        return self.log_event(
+            event_type="step_execution",
+            step=step,
+            duration_ms=duration_ms,
+            tool_name=tool_name,
+            success=success,
+            **kwargs,
+        )

pulse/tool_policy.py

@@ -0,0 +1,197 @@
+"""Typed contracts and centralized authorization for model tool calls.
+
+Tool invocations are untrusted model output.  This module keeps validation and
+authorization ahead of execution, with an explicit decision that can be
+audited without recording sensitive argument values.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Protocol
+
+
+class ToolRisk(str, Enum):
+    """Impact level used to decide whether human approval is required."""
+
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+
+
+class AuthorizationDecision(str, Enum):
+    """The only authorization outcomes available to the registry."""
+
+    ALLOW = "allow"
+    ASK = "ask"
+    DENY = "deny"
+
+
+class ArgumentKind(str, Enum):
+    """Supported wire-level argument types for a tool schema."""
+
+    BOOLEAN = "boolean"
+    INTEGER = "integer"
+    NUMBER = "number"
+    STRING = "string"
+    OBJECT = "object"
+    ARRAY = "array"
+    CALLABLE = "callable"
+
+
+@dataclass(frozen=True, slots=True)
+class ToolArgument:
+    """A named, typed argument accepted by one tool."""
+
+    name: str
+    kind: ArgumentKind
+    required: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class ToolSchema:
+    """Strict schema for untrusted tool arguments.
+
+    Unknown arguments are rejected by default.  A tool may opt out only where
+    it intentionally implements a command-style sub-protocol itself.
+    """
+
+    arguments: tuple[ToolArgument, ...] = ()
+    allow_extra: bool = False
+
+    def validate(self, arguments: Mapping[str, Any]) -> str | None:
+        declared = {argument.name: argument for argument in self.arguments}
+        for argument in self.arguments:
+            if argument.required and argument.name not in arguments:
+                return f"Missing required argument '{argument.name}'."
+
+        if not self.allow_extra:
+            extras = sorted(set(arguments) - set(declared))
+            if extras:
+                return f"Unexpected argument(s): {', '.join(extras)}."
+
+        for name, value in arguments.items():
+            schema_argument = declared.get(name)
+            if schema_argument and not _matches_kind(value, schema_argument.kind):
+                return f"Argument '{name}' must be a {schema_argument.kind.value}."
+        return None
+
+
+def _matches_kind(value: Any, kind: ArgumentKind) -> bool:
+    if kind == ArgumentKind.BOOLEAN:
+        return isinstance(value, bool)
+    if kind == ArgumentKind.INTEGER:
+        return isinstance(value, int) and not isinstance(value, bool)
+    if kind == ArgumentKind.NUMBER:
+        return isinstance(value, (int, float)) and not isinstance(value, bool)
+    if kind == ArgumentKind.STRING:
+        return isinstance(value, str)
+    if kind == ArgumentKind.OBJECT:
+        return isinstance(value, Mapping)
+    if kind == ArgumentKind.ARRAY:
+        return isinstance(value, list)
+    return callable(value)
+
+
+@dataclass(frozen=True, slots=True)
+class ToolAuthorization:
+    """An inspectable authorization result, safe to put in an audit log."""
+
+    decision: AuthorizationDecision
+    reason: str
+    subject_id: str
+    capability: str
+    risk: ToolRisk
+
+
+class AuditRecorder(Protocol):
+    def record(self, action: str, file: str, detail: str) -> None: ...
+
+
+@dataclass(slots=True)
+class ToolPolicyEngine:
+    """Evaluate a capability allowlist, workspace scope, and approval policy."""
+
+    workspace: Path
+    subject_id: str = "local-user"
+    allowed_capabilities: frozenset[str] = field(default_factory=frozenset)
+    approval_risks: frozenset[ToolRisk] = field(
+        default_factory=lambda: frozenset({ToolRisk.MEDIUM, ToolRisk.HIGH})
+    )
+    audit_log: AuditRecorder | None = None
+
+    def authorize(
+        self,
+        *,
+        tool_name: str,
+        capability: str,
+        risk: ToolRisk,
+        arguments: Mapping[str, Any],
+    ) -> ToolAuthorization:
+        """Return a deterministic decision before an executor sees arguments."""
+        if capability not in self.allowed_capabilities:
+            return self._decision(
+                AuthorizationDecision.DENY,
+                f"Capability '{capability}' is not allowlisted for this runtime.",
+                capability,
+                risk,
+            )
+
+        target_error = self._validate_workspace_scope(arguments)
+        if target_error:
+            return self._decision(AuthorizationDecision.DENY, target_error, capability, risk)
+
+        if risk in self.approval_risks:
+            return self._decision(
+                AuthorizationDecision.ASK,
+                f"{risk.value.capitalize()}-risk capability '{tool_name}' requires approval.",
+                capability,
+                risk,
+            )
+        return self._decision(
+            AuthorizationDecision.ALLOW,
+            f"Capability '{tool_name}' is allowlisted for {self.subject_id}.",
+            capability,
+            risk,
+        )
+
+    def record(self, authorization: ToolAuthorization) -> None:
+        """Record a decision without persisting untrusted tool argument values."""
+        if self.audit_log:
+            self.audit_log.record(
+                f"tool-policy-{authorization.decision.value}",
+                authorization.capability,
+                (
+                    f"subject={authorization.subject_id} risk={authorization.risk.value}; "
+                    f"{authorization.reason}"
+                ),
+            )
+
+    def _decision(
+        self,
+        decision: AuthorizationDecision,
+        reason: str,
+        capability: str,
+        risk: ToolRisk,
+    ) -> ToolAuthorization:
+        return ToolAuthorization(decision, reason, self.subject_id, capability, risk)
+
+    def _validate_workspace_scope(self, arguments: Mapping[str, Any]) -> str | None:
+        """Reject paths outside the configured workspace before execution."""
+        root = self.workspace.resolve()
+        for key in ("file", "path", "source", "destination", "cwd", "working_directory"):
+            value = arguments.get(key)
+            if value is None:
+                continue
+            if not isinstance(value, str):
+                return f"Workspace argument '{key}' must be a string path."
+            candidate = Path(value)
+            resolved = candidate.resolve() if candidate.is_absolute() else (root / candidate).resolve()
+            try:
+                resolved.relative_to(root)
+            except ValueError:
+                return f"Workspace argument '{key}' escapes the authorized workspace."
+        return None

pulse/tool_registry.py

@@ -0,0 +1,163 @@
+"""Async, UI-independent registration and execution of Pulse capabilities."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable, Iterable, Mapping
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from pulse.telemetry import get_correlation_id, set_correlation_id
+from pulse.tool_policy import (
+    AuthorizationDecision,
+    ToolPolicyEngine,
+    ToolRisk,
+    ToolSchema,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ToolInvocation:
+    name: str | None = None
+    arguments: Mapping[str, Any] = field(default_factory=dict)
+    message: str = ""
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class ToolResult:
+    content: str
+    terminal: bool = True
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class Tool(Protocol):
+    """Implement this small contract, then register an instance."""
+
+    name: str
+    description: str
+    requires_permission: bool
+    capability: str
+    risk: ToolRisk
+    schema: ToolSchema | None
+
+    def matches(self, invocation: ToolInvocation) -> bool: ...
+
+    async def execute(self, invocation: ToolInvocation) -> ToolResult: ...
+
+
+PermissionChecker = Callable[[ToolInvocation, Tool], Awaitable[bool]]
+
+
+class ToolRegistry:
+    def __init__(
+        self,
+        tools: Iterable[Tool] = (),
+        *,
+        permission_checker: PermissionChecker | None = None,
+        policy_engine: ToolPolicyEngine | None = None,
+        telemetry: Any | None = None,
+    ) -> None:
+        self._tools: dict[str, Tool] = {}
+        self._permission_checker = permission_checker
+        self._policy_engine = policy_engine
+        self._telemetry = telemetry
+        for tool in tools:
+            self.register(tool)
+
+    def register(self, tool: Tool) -> None:
+        if tool.name in self._tools:
+            raise ValueError(f"Tool already registered: {tool.name}")
+        self._tools[tool.name] = tool
+
+    def discover(self) -> tuple[Tool, ...]:
+        return tuple(self._tools.values())
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def match(self, invocation: ToolInvocation) -> Tool | None:
+        if invocation.name:
+            return self.get(invocation.name)
+        return next((tool for tool in self._tools.values() if tool.matches(invocation)), None)
+
+    async def execute(self, invocation: ToolInvocation) -> ToolResult | None:
+        tool = self.match(invocation)
+        if tool is None:
+            return None
+        correlation_id = set_correlation_id(
+            invocation.metadata.get("correlation_id") or get_correlation_id()
+        )
+        started = time.perf_counter()
+        if self._telemetry:
+            self._telemetry.log_event(
+                "tool_started", tool_name=tool.name, correlation_id=correlation_id
+            )
+
+        def finish(result: ToolResult) -> ToolResult:
+            result.metadata.setdefault("correlation_id", correlation_id)
+            if self._telemetry:
+                self._telemetry.log_event(
+                    "tool_completed",
+                    duration_ms=(time.perf_counter() - started) * 1000,
+                    tool_name=tool.name,
+                    error_code=result.metadata.get("error_code"),
+                )
+            return result
+
+        schema = getattr(tool, "schema", None)
+        if schema:
+            validation_error = schema.validate(invocation.arguments)
+            if validation_error:
+                return finish(ToolResult(
+                    f"Invalid arguments for {tool.name}: {validation_error}",
+                    metadata={"error_code": "invalid_tool_arguments", "validation_error": validation_error},
+                ))
+
+        requires_approval = bool(getattr(tool, "requires_permission", False))
+        policy_requires_approval = False
+        if self._policy_engine:
+            risk = getattr(tool, "risk", ToolRisk.LOW)
+            capability = getattr(tool, "capability", None) or tool.name
+            authorization = self._policy_engine.authorize(
+                tool_name=tool.name,
+                capability=capability,
+                risk=risk,
+                arguments=invocation.arguments,
+            )
+            self._policy_engine.record(authorization)
+            if authorization.decision == AuthorizationDecision.DENY:
+                return finish(ToolResult(
+                    f"Tool denied for {tool.name}: {authorization.reason}",
+                    metadata={"error_code": "tool_policy_denied", "policy_reason": authorization.reason},
+                ))
+            policy_requires_approval = authorization.decision == AuthorizationDecision.ASK
+            requires_approval = requires_approval or policy_requires_approval
+
+        approval_denied = (
+            policy_requires_approval and self._permission_checker is None
+        ) or (
+            requires_approval
+            and self._permission_checker is not None
+            and not await self._permission_checker(invocation, tool)
+        )
+        if approval_denied:
+            return finish(ToolResult(
+                f"Permission denied for {tool.name}.",
+                metadata={"permission_denied": True, "error_code": "tool_approval_denied"},
+            ))
+        try:
+            return finish(await tool.execute(invocation))
+        except Exception:
+            if self._telemetry:
+                self._telemetry.log_event(
+                    "tool_failed",
+                    duration_ms=(time.perf_counter() - started) * 1000,
+                    tool_name=tool.name,
+                )
+            raise
+
+    async def execute_many(self, invocations: Iterable[ToolInvocation]) -> list[ToolResult | None]:
+        """Independent tools may run concurrently; callers retain input ordering."""
+        return list(await asyncio.gather(*(self.execute(invocation) for invocation in invocations)))