coco-framework 0.1.0__py3-none-any.whl

coco/agents/middleware/_redaction.py

@@ -0,0 +1,454 @@
+"""Shared redaction utilities for middleware components."""
+
+from __future__ import annotations
+
+import hashlib
+import ipaddress
+import operator
+import re
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Literal
+from urllib.parse import urlparse
+
+from typing_extensions import TypedDict
+
+RedactionStrategy = Literal["block", "redact", "mask", "hash"]
+"""Supported strategies for handling detected sensitive values."""
+
+
+class PIIMatch(TypedDict):
+    """Represents an individual match of sensitive data."""
+
+    type: str
+    value: str
+    start: int
+    end: int
+
+
+class PIIDetectionError(Exception):
+    """Raised when configured to block on detected sensitive values."""
+
+    def __init__(self, pii_type: str, matches: Sequence[PIIMatch]) -> None:
+        """Initialize the exception with match context.
+
+        Args:
+            pii_type: Name of the detected sensitive type.
+            matches: All matches that were detected for that type.
+        """
+        self.pii_type = pii_type
+        self.matches = list(matches)
+        count = len(matches)
+        msg = f"Detected {count} instance(s) of {pii_type} in text content"
+        super().__init__(msg)
+
+
+Detector = Callable[[str], list[PIIMatch]]
+"""Callable signature for detectors that locate sensitive values."""
+
+
+def detect_email(content: str) -> list[PIIMatch]:
+    """Detect email addresses in content.
+
+    Args:
+        content: The text content to scan for email addresses.
+
+    Returns:
+        A list of detected email matches.
+    """
+    pattern = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
+    return [
+        PIIMatch(
+            type="email",
+            value=match.group(),
+            start=match.start(),
+            end=match.end(),
+        )
+        for match in re.finditer(pattern, content)
+    ]
+
+
+def detect_credit_card(content: str) -> list[PIIMatch]:
+    """Detect credit card numbers in content using Luhn validation.
+
+    Args:
+        content: The text content to scan for credit card numbers.
+
+    Returns:
+        A list of detected credit card matches.
+    """
+    pattern = r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b"
+    matches = []
+
+    for match in re.finditer(pattern, content):
+        card_number = match.group()
+        if _passes_luhn(card_number):
+            matches.append(
+                PIIMatch(
+                    type="credit_card",
+                    value=card_number,
+                    start=match.start(),
+                    end=match.end(),
+                )
+            )
+
+    return matches
+
+
+def detect_ip(content: str) -> list[PIIMatch]:
+    """Detect IPv4 or IPv6 addresses in content.
+
+    Args:
+        content: The text content to scan for IP addresses.
+
+    Returns:
+        A list of detected IP address matches.
+    """
+    matches: list[PIIMatch] = []
+    ipv4_pattern = r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b"
+
+    for match in re.finditer(ipv4_pattern, content):
+        ip_candidate = match.group()
+        try:
+            ipaddress.ip_address(ip_candidate)
+        except ValueError:
+            continue
+        matches.append(
+            PIIMatch(
+                type="ip",
+                value=ip_candidate,
+                start=match.start(),
+                end=match.end(),
+            )
+        )
+
+    return matches
+
+
+def detect_mac_address(content: str) -> list[PIIMatch]:
+    """Detect MAC addresses in content.
+
+    Args:
+        content: The text content to scan for MAC addresses.
+
+    Returns:
+        A list of detected MAC address matches.
+    """
+    pattern = r"\b([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}\b"
+    return [
+        PIIMatch(
+            type="mac_address",
+            value=match.group(),
+            start=match.start(),
+            end=match.end(),
+        )
+        for match in re.finditer(pattern, content)
+    ]
+
+
+def detect_url(content: str) -> list[PIIMatch]:
+    """Detect URLs in content using regex and stdlib validation.
+
+    Args:
+        content: The text content to scan for URLs.
+
+    Returns:
+        A list of detected URL matches.
+    """
+    matches: list[PIIMatch] = []
+
+    # Pattern 1: URLs with scheme (http:// or https://)
+    scheme_pattern = r"https?://[^\s<>\"{}|\\^`\[\]]+"
+
+    for match in re.finditer(scheme_pattern, content):
+        url = match.group()
+        result = urlparse(url)
+        if result.scheme in {"http", "https"} and result.netloc:
+            matches.append(
+                PIIMatch(
+                    type="url",
+                    value=url,
+                    start=match.start(),
+                    end=match.end(),
+                )
+            )
+
+    # Pattern 2: URLs without scheme (www.example.com or example.com/path)
+    # More conservative to avoid false positives
+    bare_pattern = (
+        r"\b(?:www\.)?[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?"
+        r"(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+(?:/[^\s]*)?"
+    )
+
+    for match in re.finditer(bare_pattern, content):
+        start, end = match.start(), match.end()
+        # Skip if already matched with scheme
+        if any(m["start"] <= start < m["end"] or m["start"] < end <= m["end"] for m in matches):
+            continue
+
+        url = match.group()
+        # Only accept if it has a path or starts with www
+        # This reduces false positives like "example.com" in prose
+        if "/" in url or url.startswith("www."):
+            # Add scheme for validation (required for urlparse to work correctly)
+            test_url = f"http://{url}"
+            result = urlparse(test_url)
+            if result.netloc and "." in result.netloc:
+                matches.append(
+                    PIIMatch(
+                        type="url",
+                        value=url,
+                        start=start,
+                        end=end,
+                    )
+                )
+
+    return matches
+
+
+BUILTIN_DETECTORS: dict[str, Detector] = {
+    "email": detect_email,
+    "credit_card": detect_credit_card,
+    "ip": detect_ip,
+    "mac_address": detect_mac_address,
+    "url": detect_url,
+}
+"""Registry of built-in detectors keyed by type name."""
+
+_CARD_NUMBER_MIN_DIGITS = 13
+_CARD_NUMBER_MAX_DIGITS = 19
+
+
+def _passes_luhn(card_number: str) -> bool:
+    """Validate credit card number using the Luhn checksum."""
+    digits = [int(d) for d in card_number if d.isdigit()]
+    if not _CARD_NUMBER_MIN_DIGITS <= len(digits) <= _CARD_NUMBER_MAX_DIGITS:
+        return False
+
+    checksum = 0
+    for index, digit in enumerate(reversed(digits)):
+        value = digit
+        if index % 2 == 1:
+            value *= 2
+            if value > 9:  # noqa: PLR2004
+                value -= 9
+        checksum += value
+    return checksum % 10 == 0
+
+
+def _apply_redact_strategy(content: str, matches: list[PIIMatch]) -> str:
+    result = content
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
+        replacement = f"[REDACTED_{match['type'].upper()}]"
+        result = result[: match["start"]] + replacement + result[match["end"] :]
+    return result
+
+
+_UNMASKED_CHAR_NUMBER = 4
+_IPV4_PARTS_NUMBER = 4
+
+
+def _apply_mask_strategy(content: str, matches: list[PIIMatch]) -> str:
+    result = content
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
+        value = match["value"]
+        pii_type = match["type"]
+        if pii_type == "email":
+            parts = value.split("@")
+            if len(parts) == 2:  # noqa: PLR2004
+                domain_parts = parts[1].split(".")
+                masked = (
+                    f"{parts[0]}@****.{domain_parts[-1]}"
+                    if len(domain_parts) > 1
+                    else f"{parts[0]}@****"
+                )
+            else:
+                masked = "****"
+        elif pii_type == "credit_card":
+            digits_only = "".join(c for c in value if c.isdigit())
+            separator = "-" if "-" in value else " " if " " in value else ""
+            if separator:
+                masked = (
+                    f"****{separator}****{separator}****{separator}"
+                    f"{digits_only[-_UNMASKED_CHAR_NUMBER:]}"
+                )
+            else:
+                masked = f"************{digits_only[-_UNMASKED_CHAR_NUMBER:]}"
+        elif pii_type == "ip":
+            octets = value.split(".")
+            masked = f"*.*.*.{octets[-1]}" if len(octets) == _IPV4_PARTS_NUMBER else "****"
+        elif pii_type == "mac_address":
+            separator = ":" if ":" in value else "-"
+            masked = (
+                f"**{separator}**{separator}**{separator}**{separator}**{separator}{value[-2:]}"
+            )
+        elif pii_type == "url":
+            masked = "[MASKED_URL]"
+        else:
+            masked = (
+                f"****{value[-_UNMASKED_CHAR_NUMBER:]}"
+                if len(value) > _UNMASKED_CHAR_NUMBER
+                else "****"
+            )
+        result = result[: match["start"]] + masked + result[match["end"] :]
+    return result
+
+
+def _apply_hash_strategy(content: str, matches: list[PIIMatch]) -> str:
+    result = content
+    for match in sorted(matches, key=operator.itemgetter("start"), reverse=True):
+        digest = hashlib.sha256(match["value"].encode()).hexdigest()[:8]
+        replacement = f"<{match['type']}_hash:{digest}>"
+        result = result[: match["start"]] + replacement + result[match["end"] :]
+    return result
+
+
+def apply_strategy(
+    content: str,
+    matches: list[PIIMatch],
+    strategy: RedactionStrategy,
+) -> str:
+    """Apply the configured strategy to matches within content.
+
+    Args:
+        content: The content to apply strategy to.
+        matches: List of detected PII matches.
+        strategy: The redaction strategy to apply.
+
+    Returns:
+        The content with the strategy applied.
+
+    Raises:
+        PIIDetectionError: If the strategy is `'block'` and matches are found.
+        ValueError: If the strategy is unknown.
+    """
+    if not matches:
+        return content
+    if strategy == "redact":
+        return _apply_redact_strategy(content, matches)
+    if strategy == "mask":
+        return _apply_mask_strategy(content, matches)
+    if strategy == "hash":
+        return _apply_hash_strategy(content, matches)
+    if strategy == "block":
+        raise PIIDetectionError(matches[0]["type"], matches)
+    msg = f"Unknown redaction strategy: {strategy}"  # type: ignore[unreachable]
+    raise ValueError(msg)
+
+
+def resolve_detector(pii_type: str, detector: Detector | str | None) -> Detector:
+    """Return a callable detector for the given configuration.
+
+    Args:
+        pii_type: The PII type name.
+        detector: Optional custom detector or regex pattern. If `None`, a built-in detector
+            for the given PII type will be used.
+
+    Returns:
+        The resolved detector.
+
+    Raises:
+        ValueError: If an unknown PII type is specified without a custom detector or regex.
+    """
+    if detector is None:
+        if pii_type not in BUILTIN_DETECTORS:
+            msg = (
+                f"Unknown PII type: {pii_type}. "
+                f"Must be one of {list(BUILTIN_DETECTORS.keys())} or provide a custom detector."
+            )
+            raise ValueError(msg)
+        return BUILTIN_DETECTORS[pii_type]
+    if isinstance(detector, str):
+        pattern = re.compile(detector)
+
+        def regex_detector(content: str) -> list[PIIMatch]:
+            return [
+                PIIMatch(
+                    type=pii_type,
+                    value=match.group(),
+                    start=match.start(),
+                    end=match.end(),
+                )
+                for match in pattern.finditer(content)
+            ]
+
+        return regex_detector
+
+    # Wrap the custom callable to normalize its output.
+    # Custom detectors may return dicts with "text" instead of "value"
+    # and may omit "type".  Map them to proper PIIMatch objects so that
+    # downstream strategies (hash, mask) can access match["value"].
+    raw_detector = detector
+
+    def _normalizing_detector(content: str) -> list[PIIMatch]:
+        return [
+            PIIMatch(
+                type=m.get("type", pii_type),
+                value=m.get("value", m.get("text", "")),
+                start=m["start"],
+                end=m["end"],
+            )
+            for m in raw_detector(content)
+        ]
+
+    return _normalizing_detector
+
+
+@dataclass(frozen=True)
+class RedactionRule:
+    """Configuration for handling a single PII type."""
+
+    pii_type: str
+    strategy: RedactionStrategy = "redact"
+    detector: Detector | str | None = None
+
+    def resolve(self) -> ResolvedRedactionRule:
+        """Resolve runtime detector and return an immutable rule.
+
+        Returns:
+            The resolved redaction rule.
+        """
+        resolved_detector = resolve_detector(self.pii_type, self.detector)
+        return ResolvedRedactionRule(
+            pii_type=self.pii_type,
+            strategy=self.strategy,
+            detector=resolved_detector,
+        )
+
+
+@dataclass(frozen=True)
+class ResolvedRedactionRule:
+    """Resolved redaction rule ready for execution."""
+
+    pii_type: str
+    strategy: RedactionStrategy
+    detector: Detector
+
+    def apply(self, content: str) -> tuple[str, list[PIIMatch]]:
+        """Apply this rule to content, returning new content and matches.
+
+        Args:
+            content: The text content to scan and redact.
+
+        Returns:
+            A tuple of (updated content, list of detected matches).
+        """
+        matches = self.detector(content)
+        if not matches:
+            return content, []
+        updated = apply_strategy(content, matches, self.strategy)
+        return updated, matches
+
+
+__all__ = [
+    "PIIDetectionError",
+    "PIIMatch",
+    "RedactionRule",
+    "ResolvedRedactionRule",
+    "apply_strategy",
+    "detect_credit_card",
+    "detect_email",
+    "detect_ip",
+    "detect_mac_address",
+    "detect_url",
+]

coco/agents/middleware/_retry.py

@@ -0,0 +1,134 @@
+"""Shared retry utilities for agent middleware.
+
+This module contains common constants, utilities, and logic used by both
+model and tool retry middleware implementations.
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Callable
+from typing import Literal
+
+from coco_core.exceptions import ModelError
+
+# Type aliases
+RetryOn = tuple[type[Exception], ...] | Callable[[Exception], bool]
+"""Type for specifying which exceptions to retry on.
+
+Can be either:
+- A tuple of exception types to retry on (based on `isinstance` checks)
+- A callable that takes an exception and returns `True` if it should be retried
+"""
+
+OnFailure = Literal["error", "continue"] | Callable[[Exception], str]
+"""Type for specifying failure handling behavior.
+
+Can be either:
+- A literal action string (`'error'` or `'continue'`)
+    - `'error'`: Re-raise the exception, stopping agent execution.
+    - `'continue'`: Inject a message with the error details, allowing the agent to continue.
+
+        For tool retries, a `ToolMessage` with the error details will be injected.
+
+        For model retries, an `AIMessage` with the error details will be returned.
+- A callable that takes an exception and returns a string for error message content
+"""
+
+
+def validate_retry_params(
+    max_retries: int,
+    initial_delay: float,
+    max_delay: float,
+    backoff_factor: float,
+) -> None:
+    """Validate retry parameters.
+
+    Args:
+        max_retries: Maximum number of retry attempts.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+        backoff_factor: Multiplier for exponential backoff.
+
+    Raises:
+        ValueError: If any parameter is invalid (negative values).
+    """
+    if max_retries < 0:
+        msg = "max_retries must be >= 0"
+        raise ValueError(msg)
+    if initial_delay < 0:
+        msg = "initial_delay must be >= 0"
+        raise ValueError(msg)
+    if max_delay < 0:
+        msg = "max_delay must be >= 0"
+        raise ValueError(msg)
+    if backoff_factor < 0:
+        msg = "backoff_factor must be >= 0"
+        raise ValueError(msg)
+
+
+def default_retry_on(exc: Exception) -> bool:
+    """Return whether an exception should be retried by default."""
+    if isinstance(exc, ModelError):
+        return exc.is_retryable
+    return True
+
+
+def should_retry_exception(
+    exc: Exception,
+    retry_on: RetryOn,
+) -> bool:
+    """Check if an exception should trigger a retry.
+
+    Args:
+        exc: The exception that occurred.
+        retry_on: Either a tuple of exception types to retry on, or a callable
+            that takes an exception and returns `True` if it should be retried.
+
+    Returns:
+        `True` if the exception should be retried, `False` otherwise.
+    """
+    if callable(retry_on):
+        return retry_on(exc)
+    return isinstance(exc, retry_on)
+
+
+def calculate_delay(
+    retry_number: int,
+    *,
+    backoff_factor: float,
+    initial_delay: float,
+    max_delay: float,
+    jitter: bool,
+) -> float:
+    """Calculate delay for a retry attempt with exponential backoff and optional jitter.
+
+    Args:
+        retry_number: The retry attempt number (0-indexed).
+        backoff_factor: Multiplier for exponential backoff.
+
+            Set to `0.0` for constant delay.
+        initial_delay: Initial delay in seconds before first retry.
+        max_delay: Maximum delay in seconds between retries.
+
+            Caps exponential backoff growth.
+        jitter: Whether to add random jitter to delay to avoid thundering herd.
+
+    Returns:
+        Delay in seconds before next retry.
+    """
+    if backoff_factor == 0.0:
+        delay = initial_delay
+    else:
+        delay = initial_delay * (backoff_factor**retry_number)
+
+    # Cap at max_delay
+    delay = min(delay, max_delay)
+
+    if jitter and delay > 0:
+        jitter_amount = delay * 0.25  # ±25% jitter
+        delay += random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
+        # Ensure delay is not negative after jitter
+        delay = max(0, delay)
+
+    return delay

coco/agents/middleware/_trace_policy.py

@@ -0,0 +1,61 @@
+"""Process-global default `TracePolicy` for agent middleware.
+
+`TracePolicy` (from coco_graph) shapes what a middleware hook span records. This module
+adds a process-wide default that individual middleware override via their own
+`trace_policy` attribute. Set it once at startup with `configure_trace_policy`; the
+default is read at span-creation time, so it applies to agents built before or after
+the call.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from coco_graph.types import TracePolicy
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_DEFAULT_TRACE_POLICY: TracePolicy | None = None
+
+
+def configure_trace_policy(policy: TracePolicy | None) -> None:
+    """Set the process-wide default `TracePolicy` for agent middleware hook spans.
+
+    Call once at startup. A middleware's own `trace_policy` overrides this wholesale
+    (no field-level merge). Pass `None` to clear the default.
+    """
+    global _DEFAULT_TRACE_POLICY  # noqa: PLW0603 - process-wide default, set once at startup
+    _DEFAULT_TRACE_POLICY = policy
+
+
+def _resolve(mw_policy: TracePolicy | None) -> TracePolicy | None:
+    """The effective policy: the middleware's own if set, else the global default."""
+    return mw_policy if mw_policy is not None else _DEFAULT_TRACE_POLICY
+
+
+def _resolved_transform(
+    mw_policy: TracePolicy | None,
+    field: Literal["process_inputs", "process_outputs"],
+) -> Callable[[Any], Any]:
+    """A trace processor that resolves the effective policy at call time.
+
+    Reading the global at call time (not build time) means `configure_trace_policy` can
+    run after `create_agent` and still apply. Passes the value through unchanged when no
+    effective policy defines `field`.
+    """
+
+    def process(value: Any) -> Any:
+        effective = _resolve(mw_policy)
+        fn = getattr(effective, field) if effective is not None else None
+        return fn(value) if fn is not None else value
+
+    return process
+
+
+def _node_trace_policy(mw_policy: TracePolicy | None) -> TracePolicy:
+    """Build the resolver-backed `TracePolicy` for a middleware's node hook."""
+    return TracePolicy(
+        process_inputs=_resolved_transform(mw_policy, "process_inputs"),
+        process_outputs=_resolved_transform(mw_policy, "process_outputs"),
+    )