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

asap/examples/rate_limiting.py

@@ -0,0 +1,137 @@
+"""Rate limiting strategies example for ASAP protocol.
+
+This module shows how rate limiting is configured and how per-sender
+and per-endpoint patterns work with the server.
+
+Patterns:
+    1. Global limit: create_app(..., rate_limit="10/second;100/minute") or ASAP_RATE_LIMIT env.
+    2. Per-sender: key_func returns sender URN when available (envelope.sender), else IP.
+    3. Per-endpoint: apply different limit strings to different routes (e.g. /asap vs /metrics).
+
+Run:
+    uv run python -m asap.examples.rate_limiting
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Sequence
+
+from asap.observability import get_logger
+from asap.transport.middleware import (
+    DEFAULT_RATE_LIMIT,
+    create_limiter,
+)
+
+logger = get_logger(__name__)
+
+
+def get_server_rate_limit_config() -> str:
+    """Return the effective rate limit string used by create_app.
+
+    create_app uses rate_limit parameter or ASAP_RATE_LIMIT env.
+    Default is "10/second;100/minute" (burst + sustained).
+
+    Returns:
+        Rate limit string (e.g. "10/second;100/minute").
+    """
+    return os.environ.get("ASAP_RATE_LIMIT", DEFAULT_RATE_LIMIT)
+
+
+def per_sender_key_concept(sender_urn: str | None, client_ip: str) -> str:
+    """Build a rate limit key for per-sender strategy (concept).
+
+    When the envelope sender is known (e.g. after auth or body parse), key by sender
+    so each agent has its own quota. Otherwise fall back to client IP.
+    The server's _get_sender_from_envelope(request) does this: it tries envelope.sender
+    then falls back to get_remote_address(request).
+
+    Args:
+        sender_urn: Sender URN from envelope, or None if not yet available.
+        client_ip: Client IP address (fallback key).
+
+    Returns:
+        Key string for the rate limiter (e.g. "urn:asap:agent:client-a" or "192.168.1.1").
+    """
+    if sender_urn:
+        return sender_urn
+    return client_ip
+
+
+def per_endpoint_limits_concept() -> dict[str, str]:
+    """Return example per-endpoint limit strings (concept).
+
+    Different routes can have different limits, e.g. strict for /asap and
+    looser for read-only /.well-known/asap/manifest.json.
+    The server currently applies one limit to the /asap route; this shows
+    the pattern for multiple routes.
+
+    Returns:
+        Map from route/path description to limit string.
+    """
+    return {
+        "asap": "10/second;100/minute",  # Burst + sustained for main endpoint
+        "metrics": "10/minute",  # Lower limit for metrics scraping
+        "manifest": "200/minute",  # Higher limit for discovery
+    }
+
+
+def run_demo() -> None:
+    """Demonstrate rate limiting config: global, per-sender key, per-endpoint pattern."""
+    # Global: server config
+    effective_limit = get_server_rate_limit_config()
+    logger.info(
+        "asap.rate_limiting.server_config",
+        rate_limit=effective_limit,
+        env_var="ASAP_RATE_LIMIT",
+    )
+
+    # Per-sender: key concept (sender URN vs IP fallback)
+    key_sender = per_sender_key_concept("urn:asap:agent:client-a", "192.168.1.1")
+    key_ip = per_sender_key_concept(None, "192.168.1.1")
+    logger.info(
+        "asap.rate_limiting.per_sender_key",
+        when_sender_known=key_sender,
+        when_fallback=key_ip,
+    )
+
+    # Per-endpoint: different limits per route
+    limits = per_endpoint_limits_concept()
+    logger.info(
+        "asap.rate_limiting.per_endpoint",
+        limits=limits,
+        message="Apply @limiter.limit(limit)(handler) per route",
+    )
+
+    # create_limiter uses _get_sender_from_envelope as key_func
+    _limiter = create_limiter(["50/minute"])
+    logger.info(
+        "asap.rate_limiting.limiter_created",
+        limits=["50/minute"],
+        key_func="get_sender_from_envelope (sender or IP)",
+    )
+
+    # rate_limit_handler is used by the server for 429 responses
+    logger.info(
+        "asap.rate_limiting.handler",
+        message="rate_limit_handler returns JSON-RPC error with Retry-After header",
+    )
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    """Parse command-line arguments for the rate limiting demo."""
+    parser = argparse.ArgumentParser(
+        description="Rate limiting strategies: per-sender, per-endpoint patterns."
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """Run the rate limiting demo."""
+    parse_args(argv)
+    run_demo()
+
+
+if __name__ == "__main__":
+    main()

asap/examples/run_demo.py

@@ -6,7 +6,7 @@ communication by sending a task request from the coordinator logic.
 
 import asyncio
 import signal
-import subprocess
+import subprocess  # nosec B404
 import sys
 import time
 from typing import Sequence
@@ -34,8 +34,15 @@ def start_process(command: Sequence[str]) -> subprocess.Popen[str]:
 
     Returns:
         Started subprocess handle.
+
+    Note:
+        This is example/demo code that only executes trusted commands
+        (sys.executable with known modules). The command is controlled
+        and not user input.
     """
-    return subprocess.Popen(command, text=True)
+    # nosec B404, B603: This is example code executing trusted commands only
+    # (sys.executable with known Python modules, not user input)
+    return subprocess.Popen(command, text=True)  # nosec B404, B603
 
 
 def wait_for_ready(url: str, timeout_seconds: float) -> None:
@@ -86,12 +93,10 @@ def main() -> None:
     signal.signal(signal.SIGTERM, handle_signal)
 
     try:
-        # Start echo agent server
         echo_process = start_process(echo_command)
         wait_for_ready(ECHO_MANIFEST_URL, READY_TIMEOUT_SECONDS)
         logger.info("asap.demo.echo_ready", url=ECHO_MANIFEST_URL)
 
-        # Demonstrate communication using coordinator dispatch logic
         logger.info("asap.demo.starting_communication")
         try:
             response = asyncio.run(

asap/examples/secure_handler.py

@@ -0,0 +1,84 @@
+"""Secure handler example for ASAP protocol.
+
+This module demonstrates proper input validation and security practices
+when implementing ASAP handlers: validated payloads, FilePart URI validation,
+and sanitized logging. See docs/security.md (Handler Security) for the full guide.
+"""
+
+from __future__ import annotations
+
+from pydantic import ValidationError
+
+from asap.errors import MalformedEnvelopeError
+from asap.models.entities import Manifest
+from asap.models.envelope import Envelope
+from asap.models.ids import generate_id
+from asap.models.parts import FilePart
+from asap.models.payloads import TaskRequest, TaskResponse
+from asap.models.enums import TaskStatus
+from asap.observability import get_logger, sanitize_for_logging
+from asap.transport.handlers import Handler
+
+logger = get_logger(__name__)
+
+
+def create_secure_handler() -> Handler:
+    """Create a handler that validates input and sanitizes logs.
+
+    The handler:
+    - Parses payload with TaskRequest (raises MalformedEnvelopeError on invalid input)
+    - Validates file parts with FilePart so URIs are checked (no path traversal, no file://)
+    - Logs payload only after sanitize_for_logging() to avoid leaking secrets
+
+    Returns:
+        A handler callable (envelope, manifest) -> Envelope suitable for
+        registry.register("task.request", ...).
+    """
+
+    def secure_handler(envelope: Envelope, manifest: Manifest) -> Envelope:
+        try:
+            task_request = TaskRequest(**envelope.payload)
+        except ValidationError as e:
+            raise MalformedEnvelopeError(
+                reason="Invalid TaskRequest payload",
+                details={"validation_errors": e.errors()},
+            ) from e
+
+        logger.debug(
+            "secure_handler.request",
+            payload_type=envelope.payload_type,
+            payload=sanitize_for_logging(envelope.payload),
+        )
+
+        validated_uris: list[str] = []
+        input_data = task_request.input or {}
+        parts = input_data.get("parts") or []
+        for part in parts:
+            if isinstance(part, dict) and part.get("type") == "file":
+                try:
+                    file_part = FilePart.model_validate(part)
+                    validated_uris.append(file_part.uri)
+                except ValidationError as e:
+                    raise MalformedEnvelopeError(
+                        reason="Invalid file part (e.g. path traversal or file:// not allowed)",
+                        details={"validation_errors": e.errors()},
+                    ) from e
+
+        result = {"echoed": task_request.input, "validated_file_uris": validated_uris}
+        response_payload = TaskResponse(
+            task_id=f"task_{generate_id()}",
+            status=TaskStatus.COMPLETED,
+            result=result,
+        )
+
+        return Envelope(
+            asap_version=envelope.asap_version,
+            sender=manifest.id,
+            recipient=envelope.sender,
+            payload_type="task.response",
+            payload=response_payload.model_dump(),
+            correlation_id=envelope.id,
+            trace_id=envelope.trace_id,
+        )
+
+    return secure_handler

asap/examples/state_migration.py

@@ -0,0 +1,240 @@
+"""State migration example for ASAP protocol.
+
+This module demonstrates moving task state between agents: export state
+from one agent's SnapshotStore and import it into another, using
+StateSnapshot, StateQuery, and StateRestore.
+
+Scenario:
+    - Agent A holds task state in its SnapshotStore.
+    - We query/export the snapshot (get from store A).
+    - We save it to Agent B's store (or send StateRestore to B if B has
+      the snapshot by ID).
+
+Run:
+    uv run python -m asap.examples.state_migration
+"""
+
+from __future__ import annotations
+
+import argparse
+from datetime import datetime, timezone
+from typing import Any, Protocol, Sequence, runtime_checkable
+
+from asap.models.entities import StateSnapshot
+from asap.models.envelope import Envelope
+from asap.models.ids import generate_id
+from asap.models.payloads import StateQuery, StateRestore
+from asap.models.types import TaskID
+from asap.observability import get_logger
+from asap.state.snapshot import InMemorySnapshotStore
+
+logger = get_logger(__name__)
+
+# URNs for source and target agents in the example
+AGENT_A_ID = "urn:asap:agent:source"
+AGENT_B_ID = "urn:asap:agent:target"
+
+
+@runtime_checkable
+class SnapshotStoreLike(Protocol):
+    """Minimal protocol for get/save (for type hints in this example)."""
+
+    def get(self, task_id: TaskID, version: int | None = None) -> StateSnapshot | None: ...
+    def save(self, snapshot: StateSnapshot) -> None: ...
+
+
+def build_state_query_envelope(
+    task_id: str,
+    version: int | None = None,
+    sender_id: str = AGENT_B_ID,
+    recipient_id: str = AGENT_A_ID,
+) -> Envelope:
+    """Build an ASAP envelope containing a StateQuery (request state from an agent).
+
+    Args:
+        task_id: Task whose state to query.
+        version: Optional snapshot version; None = latest.
+        sender_id: URN of the agent requesting state (e.g. target).
+        recipient_id: URN of the agent that holds the state (e.g. source).
+
+    Returns:
+        Envelope with payload_type "state_query" and StateQuery payload.
+    """
+    payload = StateQuery(task_id=task_id, version=version)
+    return Envelope(
+        asap_version="0.1",
+        sender=sender_id,
+        recipient=recipient_id,
+        payload_type="state_query",
+        payload=payload.model_dump(),
+        trace_id=generate_id(),
+    )
+
+
+def build_state_restore_envelope(
+    task_id: str,
+    snapshot_id: str,
+    sender_id: str = AGENT_B_ID,
+    recipient_id: str = AGENT_B_ID,
+) -> Envelope:
+    """Build an ASAP envelope containing a StateRestore (restore task from snapshot).
+
+    Args:
+        task_id: Task to restore.
+        snapshot_id: Snapshot ID to restore from (must exist in recipient's store).
+        sender_id: URN of the agent requesting restore.
+        recipient_id: URN of the agent that will restore (e.g. target agent).
+
+    Returns:
+        Envelope with payload_type "state_restore" and StateRestore payload.
+    """
+    payload = StateRestore(task_id=task_id, snapshot_id=snapshot_id)
+    return Envelope(
+        asap_version="0.1",
+        sender=sender_id,
+        recipient=recipient_id,
+        payload_type="state_restore",
+        payload=payload.model_dump(),
+        trace_id=generate_id(),
+    )
+
+
+def create_snapshot(
+    task_id: str,
+    version: int,
+    data: dict[str, Any],
+    checkpoint: bool = True,
+) -> StateSnapshot:
+    """Create a StateSnapshot for the example."""
+    return StateSnapshot(
+        id=generate_id(),
+        task_id=task_id,
+        version=version,
+        data=data,
+        checkpoint=checkpoint,
+        created_at=datetime.now(timezone.utc),
+    )
+
+
+def move_state_between_agents(
+    source_store: SnapshotStoreLike,
+    target_store: SnapshotStoreLike,
+    task_id: str,
+    version: int | None = None,
+    new_task_id: str | None = None,
+) -> StateSnapshot | None:
+    """Move task state from source agent's store to target agent's store.
+
+    Exports the snapshot from source (get) and imports it into target (save).
+    If new_task_id is set, the snapshot is saved under that task_id on the target
+    (e.g. for "migrated" task identity); otherwise the same task_id is used.
+
+    Args:
+        source_store: Snapshot store of the source agent (Agent A).
+        target_store: Snapshot store of the target agent (Agent B).
+        task_id: Task ID to export from source.
+        version: Snapshot version to export; None = latest.
+        new_task_id: If set, save on target under this task_id; else use task_id.
+
+    Returns:
+        The snapshot that was saved to the target store, or None if not found on source.
+    """
+    snapshot = source_store.get(task_id, version=version)
+    if snapshot is None:
+        logger.warning(
+            "asap.state_migration.no_snapshot",
+            task_id=task_id,
+            version=version,
+        )
+        return None
+
+    target_task_id = new_task_id if new_task_id is not None else task_id
+    # New snapshot instance for target (new id, same data) so target has its own record
+    migrated = StateSnapshot(
+        id=generate_id(),
+        task_id=target_task_id,
+        version=snapshot.version,
+        data=dict(snapshot.data),
+        checkpoint=snapshot.checkpoint,
+        created_at=snapshot.created_at,
+    )
+    target_store.save(migrated)
+    logger.info(
+        "asap.state_migration.moved",
+        source_task_id=task_id,
+        target_task_id=target_task_id,
+        snapshot_version=snapshot.version,
+    )
+    return migrated
+
+
+def run_demo() -> None:
+    """Run state migration demo: save state on agent A, move to agent B, verify."""
+    source_store: InMemorySnapshotStore = InMemorySnapshotStore()
+    target_store: InMemorySnapshotStore = InMemorySnapshotStore()
+
+    task_id = generate_id()
+    snapshot = create_snapshot(
+        task_id=task_id,
+        version=1,
+        data={"step": "processing", "items_processed": 42, "agent": "A"},
+    )
+    source_store.save(snapshot)
+    logger.info(
+        "asap.state_migration.saved_on_source",
+        task_id=task_id,
+        snapshot_id=snapshot.id,
+    )
+
+    # Build StateQuery envelope (protocol-level: request state from Agent A)
+    query_envelope = build_state_query_envelope(task_id=task_id)
+    logger.info(
+        "asap.state_migration.state_query_envelope",
+        payload_type=query_envelope.payload_type,
+        task_id=query_envelope.payload.get("task_id"),
+    )
+
+    # Move state from A to B (in-process: get from source, save to target)
+    migrated = move_state_between_agents(source_store, target_store, task_id)
+    if migrated is None:
+        raise RuntimeError("State migration failed: no snapshot on source")
+
+    # Build StateRestore envelope (protocol-level: tell Agent B to restore from snapshot)
+    restore_envelope = build_state_restore_envelope(
+        task_id=task_id,
+        snapshot_id=migrated.id,
+        recipient_id=AGENT_B_ID,
+    )
+    logger.info(
+        "asap.state_migration.state_restore_envelope",
+        payload_type=restore_envelope.payload_type,
+        task_id=restore_envelope.payload.get("task_id"),
+        snapshot_id=restore_envelope.payload.get("snapshot_id"),
+    )
+
+    restored = target_store.get(migrated.task_id, version=migrated.version)
+    assert restored is not None  # nosec B101
+    assert restored.data == snapshot.data  # nosec B101
+    logger.info(
+        "asap.state_migration.demo_complete",
+        task_id=task_id,
+        target_has_state=restored is not None,
+    )
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    """Parse command-line arguments for the state migration demo."""
+    parser = argparse.ArgumentParser(
+        description="Move task state between agents (StateSnapshot, StateQuery, StateRestore)."
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """Run the state migration demo."""
+    parse_args(argv)
+    run_demo()
+
+
+if __name__ == "__main__":
+    main()

asap/examples/streaming_response.py

@@ -0,0 +1,108 @@
+"""Streaming response example for ASAP protocol.
+
+This module demonstrates how to simulate streaming task updates using
+TaskUpdate payloads: yield progress updates (and optionally a final
+TaskResponse) so clients can show real-time progress.
+
+Use case: Long-running tasks that emit TaskUpdate (progress, status_change)
+before sending a final TaskResponse.
+
+Run:
+    uv run python -m asap.examples.streaming_response
+"""
+
+from __future__ import annotations
+
+import argparse
+from typing import Any, Iterator, Sequence
+
+from asap.models.enums import TaskStatus, UpdateType
+from asap.models.ids import generate_id
+from asap.models.payloads import TaskUpdate
+from asap.observability import get_logger
+
+logger = get_logger(__name__)
+
+
+def stream_task_updates(
+    task_id: str,
+    num_chunks: int = 5,
+    chunk_message_prefix: str = "Processing chunk",
+) -> Iterator[TaskUpdate]:
+    """Yield TaskUpdate payloads to simulate streaming progress.
+
+    In a real implementation, each update would be sent over HTTP chunked
+    response, WebSocket, or Server-Sent Events. Here we just yield in-process.
+
+    Args:
+        task_id: Task identifier.
+        num_chunks: Number of progress updates to yield (1..num_chunks).
+        chunk_message_prefix: Message prefix for progress updates.
+
+    Yields:
+        TaskUpdate with progress percent and message.
+    """
+    for i in range(1, num_chunks + 1):
+        percent = (i * 100) // num_chunks
+        update = TaskUpdate(
+            task_id=task_id,
+            update_type=UpdateType.PROGRESS,
+            status=TaskStatus.WORKING,
+            progress={
+                "percent": percent,
+                "message": f"{chunk_message_prefix} {i}/{num_chunks}",
+            },
+        )
+        yield update
+        message = update.progress.get("message") if update.progress else None
+        logger.info(
+            "asap.streaming_response.update",
+            task_id=task_id,
+            percent=percent,
+            message=message,
+        )
+
+
+def run_demo(num_chunks: int = 5) -> list[dict[str, Any]]:
+    """Run streaming demo: yield TaskUpdate payloads and collect them.
+
+    Args:
+        num_chunks: Number of progress chunks to stream.
+
+    Returns:
+        List of update payload dicts (for tests or inspection).
+    """
+    task_id = generate_id()
+    updates: list[dict[str, Any]] = []
+    for update in stream_task_updates(task_id, num_chunks=num_chunks):
+        updates.append(update.model_dump())
+    logger.info(
+        "asap.streaming_response.demo_complete",
+        task_id=task_id,
+        num_updates=len(updates),
+    )
+    return updates
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    """Parse command-line arguments for the streaming demo."""
+    parser = argparse.ArgumentParser(
+        description="Streaming response example (TaskUpdate progress chunks)."
+    )
+    parser.add_argument(
+        "--chunks",
+        type=int,
+        default=5,
+        help="Number of progress chunks to stream.",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """Run the streaming response demo."""
+    args = parse_args(argv)
+    run_demo(num_chunks=args.chunks)
+
+
+if __name__ == "__main__":
+    main()