agentpool 2.1.9__py3-none-any.whl

acp/task/debug.py

@@ -0,0 +1,197 @@
+"""Debugging extensions for ACP task system."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from datetime import datetime
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+from acp.task.state import InMemoryMessageStateStore
+
+
+if TYPE_CHECKING:
+    import asyncio
+
+    from acp.task.state import IncomingMessage
+
+
+__all__ = ["DebugEntry", "DebuggingMessageStateStore"]
+
+
+@dataclass
+class DebugEntry:
+    """Structured debug entry for ACP message tracking."""
+
+    timestamp: str
+    direction: Literal["outgoing", "incoming"]
+    event: Literal["register", "resolve", "reject", "begin", "complete", "fail"]
+    request_id: int | None
+    method: str
+    params: Any = None
+    result: Any = None
+    error: Any = None
+    status: str | None = None
+    duration_ms: float | None = None
+
+
+class DebuggingMessageStateStore(InMemoryMessageStateStore):
+    """Enhanced message state store with structured debugging output.
+
+    Provides much richer debugging information than raw JSON logging:
+    - Request/response correlation
+    - Timing information
+    - Status tracking
+    - Structured data output
+    - Error details
+    """
+
+    def __init__(self, debug_file: str | Path | None = None) -> None:
+        super().__init__()
+        self._debug_file = Path(debug_file) if debug_file else None
+        self._request_start_times: dict[int, datetime] = {}
+
+    def register_outgoing(self, request_id: int, method: str) -> asyncio.Future[Any]:
+        """Register outgoing request with debug logging."""
+        future = super().register_outgoing(request_id, method)
+
+        # Track start time for duration calculation
+        self._request_start_times[request_id] = datetime.now()
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="outgoing",
+            event="register",
+            request_id=request_id,
+            method=method,
+        )
+        self._log_debug(entry)
+
+        return future
+
+    def resolve_outgoing(self, request_id: int, result: Any) -> None:
+        """Resolve outgoing request with debug logging."""
+        duration = self._calculate_duration(request_id)
+        super().resolve_outgoing(request_id, result)
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="outgoing",
+            event="resolve",
+            request_id=request_id,
+            method=self._get_method_for_request(request_id),
+            result=result,
+            duration_ms=duration,
+        )
+        self._log_debug(entry)
+        self._cleanup_request_timing(request_id)
+
+    def reject_outgoing(self, request_id: int, error: Any) -> None:
+        """Reject outgoing request with debug logging."""
+        duration = self._calculate_duration(request_id)
+        super().reject_outgoing(request_id, error)
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="outgoing",
+            event="reject",
+            request_id=request_id,
+            method=self._get_method_for_request(request_id),
+            error=str(error),
+            duration_ms=duration,
+        )
+        self._log_debug(entry)
+        self._cleanup_request_timing(request_id)
+
+    def reject_all_outgoing(self, error: Any) -> None:
+        """Reject all outgoing requests with debug logging."""
+        # Log for each pending request
+        for request_id in list(self._outgoing.keys()):
+            duration = self._calculate_duration(request_id)
+            entry = DebugEntry(
+                timestamp=datetime.now().isoformat(),
+                direction="outgoing",
+                event="reject",
+                request_id=request_id,
+                method=self._get_method_for_request(request_id),
+                error=f"Connection error: {error}",
+                duration_ms=duration,
+            )
+            self._log_debug(entry)
+
+        super().reject_all_outgoing(error)
+        self._request_start_times.clear()
+
+    def begin_incoming(self, method: str, params: Any) -> IncomingMessage:
+        """Begin processing incoming request with debug logging."""
+        record = super().begin_incoming(method, params)
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="incoming",
+            event="begin",
+            request_id=None,
+            method=method,
+            params=params,
+            status="pending",
+        )
+        self._log_debug(entry)
+        return record
+
+    def complete_incoming(self, record: IncomingMessage, result: Any) -> None:
+        """Complete incoming request with debug logging."""
+        super().complete_incoming(record, result)
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="incoming",
+            event="complete",
+            request_id=None,
+            method=record.method,
+            result=result,
+            status="completed",
+        )
+        self._log_debug(entry)
+
+    def fail_incoming(self, record: IncomingMessage, error: Any) -> None:
+        """Fail incoming request with debug logging."""
+        super().fail_incoming(record, error)
+        entry = DebugEntry(
+            timestamp=datetime.now().isoformat(),
+            direction="incoming",
+            event="fail",
+            request_id=None,
+            method=record.method,
+            error=str(error),
+            status="failed",
+        )
+        self._log_debug(entry)
+
+    def _log_debug(self, entry: DebugEntry) -> None:
+        """Write debug entry to file if configured."""
+        if not self._debug_file:
+            return
+
+        try:
+            # Convert to dict and filter out None values for cleaner output
+            data = {k: v for k, v in asdict(entry).items() if v is not None}
+            # Write as JSONL (one JSON object per line)
+            with self._debug_file.open("a", encoding="utf-8") as f:
+                f.write(json.dumps(data, separators=(",", ":")) + "\n")
+
+        except Exception:
+            # Don't let debug logging break the connection
+            logging.exception("Failed to write debug entry")
+
+    def _calculate_duration(self, request_id: int) -> float | None:
+        """Calculate request duration in milliseconds."""
+        start_time = self._request_start_times.get(request_id)
+        if start_time:
+            return (datetime.now() - start_time).total_seconds() * 1000
+        return None
+
+    def _cleanup_request_timing(self, request_id: int) -> None:
+        """Clean up timing tracking for completed request."""
+        self._request_start_times.pop(request_id, None)
+
+    def _get_method_for_request(self, request_id: int) -> str:
+        """Get method name for outgoing request."""
+        record = self._outgoing.get(request_id)
+        return record.method if record else "unknown"

acp/task/dispatcher.py

@@ -0,0 +1,93 @@
+"""Message dispatcher for ACP tasks."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, Protocol
+
+from acp.task import RpcTaskKind
+
+
+if TYPE_CHECKING:
+    from acp.task.queue import MessageQueue
+    from acp.task.state import MessageStateStore
+    from acp.task.supervisor import TaskSupervisor
+
+
+RequestRunner = Callable[[dict[str, Any]], Awaitable[Any]]
+NotificationRunner = Callable[[dict[str, Any]], Awaitable[None]]
+
+
+class MessageDispatcher(Protocol):
+    """Protocol for message dispatchers."""
+
+    def start(self) -> None: ...
+
+    async def stop(self) -> None: ...
+
+
+class DefaultMessageDispatcher(MessageDispatcher):
+    """Background worker that consumes RPC tasks from a broker."""
+
+    def __init__(
+        self,
+        *,
+        queue: MessageQueue,
+        supervisor: TaskSupervisor,
+        store: MessageStateStore,
+        request_runner: RequestRunner,
+        notification_runner: NotificationRunner,
+    ) -> None:
+        self._queue = queue
+        self._supervisor = supervisor
+        self._store = store
+        self._request_runner = request_runner
+        self._notification_runner = notification_runner
+        self._task: asyncio.Task[None] | None = None
+
+    def start(self) -> None:
+        if self._task is not None:
+            raise RuntimeError("dispatcher already started")
+        self._task = self._supervisor.create(self._run(), name="acp.Dispatcher.loop")
+
+    async def _run(self) -> None:
+        try:
+            async for task in self._queue:
+                try:
+                    if task.kind is RpcTaskKind.REQUEST:
+                        await self._dispatch_request(task.message)
+                    else:
+                        await self._dispatch_notification(task.message)
+                finally:
+                    self._queue.task_done()
+        except asyncio.CancelledError:
+            return
+
+    async def stop(self) -> None:
+        await self._queue.close()
+        if self._task is not None:
+            with suppress(asyncio.CancelledError):
+                await self._task
+            self._task = None
+
+    async def _dispatch_request(self, msg: dict[str, Any]) -> None:
+        record = self._store.begin_incoming(msg.get("method", ""), msg.get("params"))
+
+        async def runner() -> None:
+            try:
+                result = await self._request_runner(msg)
+            except Exception as exc:
+                self._store.fail_incoming(record, exc)
+                raise
+            else:
+                self._store.complete_incoming(record, result)
+
+        self._supervisor.create(runner(), name="acp.Dispatcher.request")
+
+    async def _dispatch_notification(self, message: dict[str, Any]) -> None:
+        async def runner() -> None:
+            await self._notification_runner(message)
+
+        self._supervisor.create(runner(), name="acp.Dispatcher.notification")

acp/task/queue.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from typing import TYPE_CHECKING, Protocol
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from acp.task import RpcTask
+
+
+class MessageQueue(Protocol):
+    """Protocol for message queues used in RPC task dispatch."""
+
+    async def publish(self, task: RpcTask) -> None: ...
+
+    async def close(self) -> None: ...
+
+    def task_done(self) -> None: ...
+
+    async def join(self) -> None: ...
+
+    def __aiter__(self) -> AsyncIterator[RpcTask]: ...
+
+
+class InMemoryMessageQueue:
+    """Simple in-memory broker for RPC task dispatch."""
+
+    def __init__(self, *, maxsize: int = 0) -> None:
+        self._queue: asyncio.Queue[RpcTask | None] = asyncio.Queue(maxsize=maxsize)
+        self._closed = False
+
+    async def publish(self, task: RpcTask) -> None:
+        if self._closed:
+            raise RuntimeError("mssage queue already closed")
+        await self._queue.put(task)
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await self._queue.put(None)
+
+    async def join(self) -> None:
+        await self._queue.join()
+
+    def task_done(self) -> None:
+        with suppress(ValueError):
+            self._queue.task_done()
+
+    def __aiter__(self) -> AsyncIterator[RpcTask]:
+        return _QueueIterator(self)
+
+
+class _QueueIterator:
+    def __init__(self, queue: InMemoryMessageQueue) -> None:
+        self._queue = queue
+
+    def __aiter__(self) -> _QueueIterator:
+        return self
+
+    async def __anext__(self) -> RpcTask:
+        item = await self._queue._queue.get()
+        if item is None:
+            self._queue.task_done()
+            raise StopAsyncIteration
+        return item

acp/task/sender.py

@@ -0,0 +1,82 @@
+"""Sender class for sending messages to a remote peer."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+import contextlib
+from dataclasses import dataclass
+import json
+import logging
+from typing import Any
+
+import anyio
+from anyio.abc import ByteSendStream
+
+from acp.task.supervisor import TaskSupervisor
+
+
+__all__ = ["MessageSender", "SenderFactory"]
+
+
+SenderFactory = Callable[[ByteSendStream, TaskSupervisor], "MessageSender"]
+
+
+@dataclass(slots=True)
+class _PendingSend:
+    payload: bytes
+    future: asyncio.Future[None]
+
+
+class MessageSender:
+    """Async message sender that queues and transmits JSON-RPC messages."""
+
+    def __init__(
+        self,
+        writer: ByteSendStream,
+        supervisor: TaskSupervisor,
+    ) -> None:
+        self._writer = writer
+        self._queue: asyncio.Queue[_PendingSend | None] = asyncio.Queue()
+        self._closed = False
+        self._task = supervisor.create(
+            self._loop(), name="acp.Sender.loop", on_error=self._on_error
+        )
+
+    async def send(self, payload: dict[str, Any]) -> None:
+        data = (json.dumps(payload, separators=(",", ":")) + "\n").encode("utf-8")
+        future: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        await self._queue.put(_PendingSend(data, future))
+        await future
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        await self._queue.put(None)
+        if self._task is not None:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._task
+
+    async def _loop(self) -> None:
+        try:
+            while True:
+                item = await self._queue.get()
+                if item is None:
+                    return
+                try:
+                    await self._writer.send(item.payload)
+                except Exception as exc:
+                    if not item.future.done():
+                        item.future.set_exception(exc)
+                    raise
+                else:
+                    if not item.future.done():
+                        item.future.set_result(None)
+        except asyncio.CancelledError:
+            return
+        except anyio.ClosedResourceError:
+            return
+
+    def _on_error(self, task: asyncio.Task[Any], exc: BaseException) -> None:
+        logging.exception("Send loop failed", exc_info=exc)

acp/task/state.py

@@ -0,0 +1,87 @@
+"""Message State Store."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+@dataclass(slots=True)
+class OutgoingMessage:
+    """Represents an outgoing message with its request ID, method, and future."""
+
+    request_id: int
+    method: str
+    future: asyncio.Future[Any]
+
+
+@dataclass(slots=True)
+class IncomingMessage:
+    """Represents an incoming message with its status, result, and error."""
+
+    method: str
+    params: Any
+    status: str = "pending"
+    result: Any = None
+    error: Any = None
+
+
+class MessageStateStore(Protocol):
+    """Protocol for managing message states."""
+
+    def register_outgoing(self, request_id: int, method: str) -> asyncio.Future[Any]: ...
+
+    def resolve_outgoing(self, request_id: int, result: Any) -> None: ...
+
+    def reject_outgoing(self, request_id: int, error: Any) -> None: ...
+
+    def reject_all_outgoing(self, error: Any) -> None: ...
+
+    def begin_incoming(self, method: str, params: Any) -> IncomingMessage: ...
+
+    def complete_incoming(self, record: IncomingMessage, result: Any) -> None: ...
+
+    def fail_incoming(self, record: IncomingMessage, error: Any) -> None: ...
+
+
+class InMemoryMessageStateStore(MessageStateStore):
+    """In-memory implementation of MessageStateStore."""
+
+    def __init__(self) -> None:
+        self._outgoing: dict[int, OutgoingMessage] = {}
+        self._incoming: list[IncomingMessage] = []
+
+    def register_outgoing(self, request_id: int, method: str) -> asyncio.Future[Any]:
+        future: asyncio.Future[Any] = asyncio.get_running_loop().create_future()
+        self._outgoing[request_id] = OutgoingMessage(request_id, method, future)
+        return future
+
+    def resolve_outgoing(self, request_id: int, result: Any) -> None:
+        record = self._outgoing.pop(request_id, None)
+        if record and not record.future.done():
+            record.future.set_result(result)
+
+    def reject_outgoing(self, request_id: int, error: Any) -> None:
+        record = self._outgoing.pop(request_id, None)
+        if record and not record.future.done():
+            record.future.set_exception(error)
+
+    def reject_all_outgoing(self, error: Any) -> None:
+        for record in self._outgoing.values():
+            if not record.future.done():
+                record.future.set_exception(error)
+        self._outgoing.clear()
+
+    def begin_incoming(self, method: str, params: Any) -> IncomingMessage:
+        record = IncomingMessage(method=method, params=params)
+        self._incoming.append(record)
+        return record
+
+    def complete_incoming(self, record: IncomingMessage, result: Any) -> None:
+        record.status = "completed"
+        record.result = result
+
+    def fail_incoming(self, record: IncomingMessage, error: Any) -> None:
+        record.status = "failed"
+        record.error = error

acp/task/supervisor.py

@@ -0,0 +1,93 @@
+"""Task Supervisor."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any
+
+from agentpool import log
+
+
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
+
+__all__ = ["TaskSupervisor"]
+
+logger = log.get_logger(__name__)
+ErrorHandler = Callable[[asyncio.Task[Any], BaseException], None]
+
+
+class TaskSupervisor:
+    """Track background tasks and provide graceful shutdown semantics.
+
+    Inspired by fasta2a's task manager, this supervisor keeps a registry of
+    asyncio tasks created for request handling so they can be cancelled and
+    awaited reliably when the connection closes.
+    """
+
+    def __init__(self, *, source: str) -> None:
+        self._source = source
+        self._tasks: set[asyncio.Task[Any]] = set()
+        self._closed = False
+        self._error_handlers: list[ErrorHandler] = []
+
+    def add_error_handler(self, handler: ErrorHandler) -> None:
+        self._error_handlers.append(handler)
+
+    def create(
+        self,
+        coroutine: Coroutine[Any, Any, Any],
+        *,
+        name: str | None = None,
+        on_error: ErrorHandler | None = None,
+    ) -> asyncio.Task[Any]:
+        if self._closed:
+            msg = f"TaskSupervisor for {self._source} already closed"
+            raise RuntimeError(msg)
+        task = asyncio.create_task(coroutine, name=name)
+        self._tasks.add(task)
+        task.add_done_callback(lambda t: self._on_done(t, on_error))
+        return task
+
+    def _on_done(self, task: asyncio.Task[Any], on_error: ErrorHandler | None) -> None:
+        self._tasks.discard(task)
+        if task.cancelled():
+            return
+        try:
+            task.result()
+        except Exception as exc:
+            handled = False
+            if on_error is not None:
+                try:
+                    on_error(task, exc)
+                    handled = True
+                except Exception:
+                    logger.exception(
+                        "Error in task-specific error handler",
+                        source=self._source,
+                    )
+            if not handled:
+                for handler in self._error_handlers:
+                    try:
+                        handler(task, exc)
+                        handled = True
+                    except Exception:
+                        msg = "Error in supervisor error handler"
+                        logger.exception(msg, source=self._source)
+            if not handled:
+                logger.exception("Unhandled error in task", source=self._source)
+
+    async def shutdown(self) -> None:
+        self._closed = True
+        if not self._tasks:
+            return
+        tasks = list(self._tasks)
+        for task in tasks:
+            task.cancel()
+        for task in tasks:
+            with suppress(asyncio.CancelledError):
+                await task
+        self._tasks.clear()

acp/terminal_handle.py

@@ -0,0 +1,30 @@
+"""Terminal handle implementation. NOTE: not integrated yet."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from acp.acp_requests import ACPRequests
+    from acp.schema import TerminalOutputResponse, WaitForTerminalExitResponse
+
+
+class TerminalHandle:
+    """Handle for a terminal session."""
+
+    def __init__(self, terminal_id: str, requests: ACPRequests) -> None:
+        self.terminal_id = terminal_id
+        self._requests = requests
+
+    async def current_output(self) -> TerminalOutputResponse:
+        return await self._requests.terminal_output(self.terminal_id)
+
+    async def wait_for_exit(self) -> WaitForTerminalExitResponse:
+        return await self._requests.wait_for_terminal_exit(self.terminal_id)
+
+    async def kill(self) -> None:
+        return await self._requests.kill_terminal(self.terminal_id)
+
+    async def release(self) -> None:
+        return await self._requests.release_terminal(self.terminal_id)