asap-protocol 0.3.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

asap/observability/tracing.py

@@ -0,0 +1,293 @@
+"""OpenTelemetry tracing integration for ASAP protocol.
+
+This module provides distributed tracing with W3C Trace Context propagation.
+When enabled, FastAPI and httpx are auto-instrumented; custom spans cover
+handler execution and state transitions. Trace IDs can be carried in
+envelope.trace_id and envelope.extensions for cross-service correlation.
+
+Example:
+    >>> from asap.observability.tracing import configure_tracing, get_tracer
+    >>> configure_tracing(service_name="my-agent", app=app)
+    >>> tracer = get_tracer(__name__)
+    >>> with tracer.start_as_current_span("my.operation"):
+    ...     ...
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import AbstractContextManager
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry import context, trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+from asap.models.envelope import Envelope
+from asap.observability.logging import get_logger
+
+# Environment variables for zero-config (OpenTelemetry convention)
+_ENV_OTEL_SERVICE_NAME = "OTEL_SERVICE_NAME"
+_ENV_OTEL_TRACES_EXPORTER = "OTEL_TRACES_EXPORTER"
+_ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT"
+
+# Extension keys for W3C trace context in envelope
+EXTENSION_TRACE_ID = "trace_id"
+EXTENSION_SPAN_ID = "span_id"
+
+logger = get_logger(__name__)
+
+_tracer_provider: TracerProvider | None = None
+_tracer: trace.Tracer | None = None
+
+
+def configure_tracing(
+    service_name: str | None = None,
+    app: FastAPI | None = None,
+) -> None:
+    """Configure OpenTelemetry tracing and optionally instrument FastAPI and httpx.
+
+    Uses environment variables for zero-config:
+    - OTEL_SERVICE_NAME: service name (default: "asap-server")
+    - OTEL_TRACES_EXPORTER: "none" | "otlp" | "console" (default: "none")
+    - OTEL_EXPORTER_OTLP_ENDPOINT: OTLP endpoint (e.g. http://localhost:4317)
+
+    When OTEL_TRACES_EXPORTER is "none" or unset, tracing is configured but
+    no spans are exported (useful for dev without Jaeger). Set to "otlp" and
+    OTEL_EXPORTER_OTLP_ENDPOINT for Jaeger/collector.
+
+    Args:
+        service_name: Override for OTEL_SERVICE_NAME.
+        app: If provided, FastAPI and httpx are instrumented.
+    """
+    global _tracer_provider, _tracer
+
+    name = service_name or os.environ.get(_ENV_OTEL_SERVICE_NAME) or "asap-server"
+    resource = Resource.create({"service.name": name})
+    _tracer_provider = TracerProvider(resource=resource)
+    trace.set_tracer_provider(_tracer_provider)
+
+    exporter_name = os.environ.get(_ENV_OTEL_TRACES_EXPORTER, "none").strip().lower()
+    if exporter_name == "none":
+        _tracer = trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
+        if app is not None:
+            _instrument_app(app)
+        return
+
+    if exporter_name == "otlp":
+        _add_otlp_processor()
+    elif exporter_name == "console":
+        _add_console_processor()
+
+    _tracer = trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
+    if app is not None:
+        _instrument_app(app)
+
+
+def _add_otlp_processor() -> None:
+    """Add OTLP span processor if endpoint is set."""
+    global _tracer_provider
+    if _tracer_provider is None:
+        return
+    endpoint = os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT)
+    if not endpoint:
+        return
+    try:
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+        exporter: SpanExporter = OTLPSpanExporter(endpoint=endpoint, insecure=True)
+        _tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
+    except ImportError as e:
+        logger.debug("OTLP gRPC exporter not available: %s", e)
+        try:
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+                OTLPSpanExporter as OTLPSpanExporterHttp,
+            )
+
+            exporter = OTLPSpanExporterHttp(endpoint=endpoint)
+            _tracer_provider.add_span_processor(BatchSpanProcessor(exporter))
+        except ImportError as e2:
+            logger.debug("OTLP HTTP exporter not available: %s", e2)
+
+
+def _add_console_processor() -> None:
+    """Add console span processor for debugging."""
+    global _tracer_provider
+    if _tracer_provider is None:
+        return
+    try:
+        from opentelemetry.sdk.trace.export import ConsoleSpanExporter
+
+        _tracer_provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
+    except ImportError as e:
+        logger.debug("Console span exporter not available: %s", e)
+
+
+def _instrument_app(app: FastAPI) -> None:
+    """Instrument FastAPI and httpx for automatic spans."""
+    try:
+        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+        FastAPIInstrumentor.instrument_app(app)
+    except Exception as e:
+        logger.debug("Failed to instrument FastAPI: %s", e)
+    try:
+        from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+
+        HTTPXClientInstrumentor().instrument()
+    except Exception as e:
+        logger.debug("Failed to instrument httpx: %s", e)
+
+
+def reset_tracing() -> None:
+    """Reset global tracer state (for test teardown).
+
+    Clears the module-level tracer provider and tracer so that subsequent
+    tests or configure_tracing() calls start from a clean state.
+    """
+    global _tracer_provider, _tracer
+    _tracer_provider = None
+    _tracer = None
+
+
+def get_tracer(name: str | None = None) -> trace.Tracer:
+    """Return the ASAP protocol tracer for custom spans.
+
+    Args:
+        name: Optional logger/module name for span attribution.
+
+    Returns:
+        OpenTelemetry Tracer instance.
+    """
+    if _tracer is None:
+        trace.set_tracer_provider(TracerProvider())
+        return trace.get_tracer("asap.protocol", "1.0.0", schema_url=None)
+    return _tracer
+
+
+def inject_envelope_trace_context(envelope: Envelope) -> Envelope:
+    """Inject current span trace_id and span_id into envelope (W3C propagation).
+
+    Sets envelope.trace_id and envelope.extensions[trace_id, span_id] from the
+    current OpenTelemetry context so downstream services can continue the trace.
+
+    Args:
+        envelope: Response envelope to annotate.
+
+    Returns:
+        New envelope with trace_id and extensions updated (immutable).
+    """
+    span = trace.get_current_span()
+    if not span.is_recording():
+        return envelope
+
+    ctx = span.get_span_context()
+    trace_id_hex = format(ctx.trace_id, "032x")
+    span_id_hex = format(ctx.span_id, "016x")
+
+    extensions = dict(envelope.extensions or {})
+    extensions[EXTENSION_TRACE_ID] = trace_id_hex
+    extensions[EXTENSION_SPAN_ID] = span_id_hex
+
+    # Preserve existing trace_id for correlation (e.g. from request); set only if missing
+    new_trace_id = envelope.trace_id if envelope.trace_id else trace_id_hex
+    return envelope.model_copy(update={"trace_id": new_trace_id, "extensions": extensions})
+
+
+def extract_and_activate_envelope_trace_context(envelope: Envelope) -> Any | None:
+    """Extract trace context from envelope and set as current context (W3C).
+
+    If envelope has trace_id (and optionally span_id in extensions), creates
+    a non-recording span context so new spans become children of the incoming
+    trace. Caller should call context.detach(token) when request ends.
+
+    Args:
+        envelope: Incoming envelope carrying trace_id / extensions.
+
+    Returns:
+        Context token to pass to context.detach(), or None if no context set.
+    """
+    trace_id_str = envelope.trace_id
+    if not trace_id_str or len(trace_id_str) != 32:
+        return None
+
+    extensions = envelope.extensions or {}
+    span_id_str = extensions.get(EXTENSION_SPAN_ID)
+    if not span_id_str or len(span_id_str) != 16:
+        return None
+
+    try:
+        trace_id = int(trace_id_str, 16)
+        span_id = int(span_id_str, 16)
+    except ValueError:
+        return None
+
+    from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
+
+    span_context = SpanContext(
+        trace_id=trace_id,
+        span_id=span_id,
+        is_remote=True,
+        trace_flags=TraceFlags(0x01),
+    )
+    ctx = trace.set_span_in_context(NonRecordingSpan(span_context))
+    return context.attach(ctx)
+
+
+def handler_span_context(
+    payload_type: str,
+    agent_urn: str,
+    envelope_id: str | None,
+) -> AbstractContextManager[trace.Span]:
+    """Start a current span for handler execution (use as context manager).
+
+    Attributes: asap.payload_type, asap.agent.urn, asap.envelope.id.
+
+    Args:
+        payload_type: Envelope payload type.
+        agent_urn: Manifest/agent URN.
+        envelope_id: Envelope id.
+
+    Returns:
+        Span context manager (use with "with").
+    """
+    tracer = get_tracer(__name__)
+    attrs: dict[str, str] = {
+        "asap.payload_type": payload_type,
+        "asap.agent.urn": agent_urn,
+    }
+    if envelope_id:
+        attrs["asap.envelope.id"] = envelope_id
+    return tracer.start_as_current_span("asap.handler.execute", attributes=attrs)
+
+
+def state_transition_span_context(
+    from_status: str,
+    to_status: str,
+    task_id: str | None = None,
+) -> AbstractContextManager[trace.Span]:
+    """Start a current span for a state machine transition (use as context manager).
+
+    Attributes: asap.state.from, asap.state.to, asap.task.id.
+
+    Args:
+        from_status: Previous status.
+        to_status: New status.
+        task_id: Optional task id.
+
+    Returns:
+        Span context manager (use with "with").
+    """
+    tracer = get_tracer(__name__)
+    attrs: dict[str, str] = {
+        "asap.state.from": from_status,
+        "asap.state.to": to_status,
+    }
+    if task_id:
+        attrs["asap.task.id"] = task_id
+    return tracer.start_as_current_span("asap.state.transition", attributes=attrs)

asap/state/machine.py

@@ -14,7 +14,10 @@ from datetime import datetime, timezone
 from asap.errors import InvalidTransitionError
 from asap.models.entities import Task
 from asap.models.enums import TaskStatus
+from asap.observability import get_metrics
+from asap.observability.tracing import state_transition_span_context
 
+__all__ = ["TaskStatus", "can_transition", "transition", "VALID_TRANSITIONS"]
 
 # Valid state transitions mapping
 VALID_TRANSITIONS: dict[TaskStatus, set[TaskStatus]] = {
@@ -82,5 +85,15 @@ def transition(task: Task, new_status: TaskStatus) -> Task:
             from_state=task.status.value, to_state=new_status.value, details={"task_id": task.id}
         )
 
-    # Create new task instance with updated status and timestamp (immutable approach)
-    return task.model_copy(update={"status": new_status, "updated_at": datetime.now(timezone.utc)})
+    with state_transition_span_context(
+        from_status=task.status.value,
+        to_status=new_status.value,
+        task_id=task.id,
+    ):
+        get_metrics().increment_counter(
+            "asap_state_transitions_total",
+            {"from_status": task.status.value, "to_status": new_status.value},
+        )
+        return task.model_copy(
+            update={"status": new_status, "updated_at": datetime.now(timezone.utc)}
+        )

asap/state/snapshot.py

@@ -154,14 +154,10 @@ class InMemorySnapshotStore:
         with self._lock:
             task_id = snapshot.task_id
 
-            # Initialize storage for this task if needed
             if task_id not in self._snapshots:
                 self._snapshots[task_id] = {}
 
-            # Store the snapshot
             self._snapshots[task_id][snapshot.version] = snapshot
-
-            # Update latest version
             self._latest_versions[task_id] = max(
                 self._latest_versions.get(task_id, 0), snapshot.version
             )
@@ -186,13 +182,11 @@ class InMemorySnapshotStore:
                 return None
 
             if version is None:
-                # Return latest version
                 latest_version = self._latest_versions.get(task_id)
                 if latest_version is None:
                     return None
                 return self._snapshots[task_id].get(latest_version)
 
-            # Return specific version
             return self._snapshots[task_id].get(version)
 
     def list_versions(self, task_id: TaskID) -> list[int]:
@@ -243,18 +237,15 @@ class InMemorySnapshotStore:
                     return True
                 return False
 
-            # Delete specific version
             if version in self._snapshots[task_id]:
                 del self._snapshots[task_id][version]
 
-                # Update latest version if needed
                 if self._latest_versions.get(task_id) == version:
                     if self._snapshots[task_id]:
                         self._latest_versions[task_id] = max(self._snapshots[task_id].keys())
                     else:
                         del self._latest_versions[task_id]
 
-                # Clean up empty task dict
                 if not self._snapshots[task_id]:
                     del self._snapshots[task_id]
                     if task_id in self._latest_versions:

asap/testing/__init__.py

@@ -0,0 +1,31 @@
+"""ASAP testing utilities for easier test authoring.
+
+This package provides pytest fixtures, mock agents, and custom assertions
+to reduce boilerplate when testing ASAP protocol integrations.
+
+Modules:
+    fixtures: Pytest fixtures (mock_agent, mock_client, mock_snapshot_store)
+              and context managers (test_agent, test_client).
+    mocks: MockAgent for configurable mock agents with pre-set responses
+           and request recording.
+    assertions: Custom assertions (assert_envelope_valid, assert_task_completed,
+              assert_response_correlates).
+
+Example:
+    >>> from asap.testing import MockAgent, assert_envelope_valid
+    >>> from asap.testing.fixtures import mock_agent, test_client
+"""
+
+from asap.testing.assertions import (
+    assert_envelope_valid,
+    assert_response_correlates,
+    assert_task_completed,
+)
+from asap.testing.mocks import MockAgent
+
+__all__ = [
+    "MockAgent",
+    "assert_envelope_valid",
+    "assert_response_correlates",
+    "assert_task_completed",
+]

asap/testing/assertions.py

@@ -0,0 +1,108 @@
+"""Custom assertions for ASAP protocol tests.
+
+This module provides assertion helpers to validate envelopes and
+task outcomes with clear error messages.
+
+Functions:
+    assert_envelope_valid: Assert an Envelope has required fields and valid shape.
+    assert_task_completed: Assert a TaskResponse or envelope payload indicates
+                           task completion (e.g. status completed).
+    assert_response_correlates: Assert response envelope correlates to request.
+"""
+
+from typing import Any
+
+from asap.models.envelope import Envelope
+
+
+def assert_envelope_valid(
+    envelope: Envelope,
+    *,
+    require_id: bool = True,
+    require_timestamp: bool = True,
+    allowed_payload_types: list[str] | None = None,
+) -> None:
+    """Assert that an envelope has required fields and valid structure.
+
+    Args:
+        envelope: The envelope to validate.
+        require_id: If True, envelope.id must be non-empty.
+        require_timestamp: If True, envelope.timestamp must be set.
+        allowed_payload_types: If set, payload_type must be in this list.
+
+    Raises:
+        AssertionError: If any check fails.
+    """
+    assert envelope is not None, "Envelope must not be None"  # nosec B101
+    if require_id:
+        assert envelope.id, "Envelope must have a non-empty id"  # nosec B101
+    if require_timestamp:
+        assert envelope.timestamp is not None, "Envelope must have a timestamp"  # nosec B101
+    assert envelope.sender, "Envelope must have a sender"  # nosec B101
+    assert envelope.recipient, "Envelope must have a recipient"  # nosec B101
+    assert envelope.payload_type, "Envelope must have a payload_type"  # nosec B101
+    assert envelope.payload is not None, "Envelope must have a payload"  # nosec B101
+    if allowed_payload_types is not None:
+        assert envelope.payload_type in allowed_payload_types, (  # nosec B101
+            f"payload_type {envelope.payload_type!r} not in {allowed_payload_types}"
+        )
+
+
+def assert_task_completed(
+    payload: dict[str, Any] | Envelope,
+    *,
+    status_key: str = "status",
+    completed_value: str = "completed",
+) -> None:
+    """Assert that a task response indicates completion.
+
+    Accepts either a TaskResponse-like dict (with status_key) or an
+    Envelope whose payload is such a dict.
+
+    Args:
+        payload: TaskResponse payload dict or Envelope containing it.
+        status_key: Key in payload that holds status (default 'status').
+        completed_value: Value that indicates completion (default 'completed').
+
+    Raises:
+        AssertionError: If payload does not indicate completion.
+    """
+    if isinstance(payload, Envelope):
+        payload = payload.payload or {}
+    assert isinstance(payload, dict), "payload must be a dict or Envelope"  # nosec B101
+    actual = payload.get(status_key)
+    assert actual == completed_value, f"Expected task status {completed_value!r}, got {actual!r}"  # nosec B101
+
+
+def assert_response_correlates(
+    request_envelope: Envelope,
+    response_envelope: Envelope,
+    *,
+    correlation_id_field: str = "correlation_id",
+) -> None:
+    """Assert that a response envelope correlates to the request (by correlation id).
+
+    Args:
+        request_envelope: The request envelope (must have id).
+        response_envelope: The response envelope (must have correlation_id).
+        correlation_id_field: Attribute name on response for correlation (default
+            'correlation_id').
+
+    Raises:
+        AssertionError: If request id or response correlation_id is missing or
+            they do not match.
+    """
+    assert request_envelope.id, "Request envelope must have a non-empty id"  # nosec B101
+    correlation_id = getattr(response_envelope, correlation_id_field, None)
+    assert correlation_id is not None, f"Response envelope must have {correlation_id_field!r}"  # nosec B101
+    assert correlation_id == request_envelope.id, (  # nosec B101
+        f"Response {correlation_id_field!r} {correlation_id!r} does not match "
+        f"request id {request_envelope.id!r}"
+    )
+
+
+__all__ = [
+    "assert_envelope_valid",
+    "assert_task_completed",
+    "assert_response_correlates",
+]

asap/testing/fixtures.py

@@ -0,0 +1,113 @@
+"""Pytest fixtures and context managers for ASAP tests.
+
+This module provides shared fixtures and context managers to reduce
+boilerplate when testing agents, clients, and snapshot stores.
+
+Fixtures (use with pytest):
+    mock_agent: Configurable mock agent for request/response tests.
+    mock_client: ASAP client configured for testing (async; use in async tests).
+    mock_snapshot_store: In-memory snapshot store for state persistence tests.
+
+Context managers:
+    test_agent(): Sync context manager yielding a MockAgent for the scope.
+    test_client(): Async context manager yielding an ASAPClient for the scope.
+"""
+
+from contextlib import asynccontextmanager, contextmanager
+from typing import AsyncIterator, Iterator
+
+import pytest
+
+from asap.state.snapshot import InMemorySnapshotStore
+from asap.testing.mocks import MockAgent
+from asap.transport.client import ASAPClient
+
+DEFAULT_TEST_BASE_URL = "http://localhost:9999"
+
+
+@pytest.fixture
+def mock_agent() -> MockAgent:
+    """Create a fresh MockAgent for the test.
+
+    Returns:
+        A MockAgent instance (cleared between tests via fresh fixture).
+    """
+    return MockAgent()
+
+
+@pytest.fixture
+def mock_snapshot_store() -> InMemorySnapshotStore:
+    """Create an in-memory snapshot store for the test.
+
+    Returns:
+        An InMemorySnapshotStore instance (empty, isolated per test).
+    """
+    return InMemorySnapshotStore()
+
+
+@pytest.fixture
+async def mock_client() -> AsyncIterator[ASAPClient]:
+    """Provide an ASAPClient entered for the test (async fixture).
+
+    Yields an open ASAPClient pointing at DEFAULT_TEST_BASE_URL.
+    Use in async tests; the client is closed after the test.
+
+    Yields:
+        ASAPClient instance (already in async context).
+    """
+    async with ASAPClient(DEFAULT_TEST_BASE_URL) as client:
+        yield client
+
+
+@contextmanager
+def test_agent(agent_id: str = "urn:asap:agent:mock") -> Iterator[MockAgent]:
+    """Context manager that provides a MockAgent for the scope.
+
+    On exit, the agent is cleared (requests and responses reset).
+
+    Args:
+        agent_id: URN for the mock agent.
+
+    Yields:
+        A MockAgent instance.
+
+    Example:
+        >>> with test_agent() as agent:
+        ...     agent.set_response("echo", {"status": "completed"})
+        ...     out = agent.handle(request_envelope)
+    """
+    agent = MockAgent(agent_id)
+    try:
+        yield agent
+    finally:
+        agent.clear()
+
+
+@asynccontextmanager
+async def test_client(
+    base_url: str = DEFAULT_TEST_BASE_URL,
+) -> AsyncIterator[ASAPClient]:
+    """Async context manager that provides an ASAPClient for the scope.
+
+    Args:
+        base_url: Agent base URL (default: localhost:9999 for test servers).
+
+    Yields:
+        An open ASAPClient instance.
+
+    Example:
+        >>> async with test_client("http://localhost:8000") as client:
+        ...     response = await client.send(envelope)
+    """
+    async with ASAPClient(base_url) as client:
+        yield client
+
+
+__all__ = [
+    "DEFAULT_TEST_BASE_URL",
+    "mock_agent",
+    "mock_client",
+    "mock_snapshot_store",
+    "test_agent",
+    "test_client",
+]