diff --git a/docs/experimental-v2.md b/docs/experimental-v2.md new file mode 100644 index 0000000..dada694 --- /dev/null +++ b/docs/experimental-v2.md @@ -0,0 +1,76 @@ +# Experimental Protocol v2 + +> **Experimental.** Protocol v2 is a draft. Import it from `acp.experimental` and +> expect its API and generated models to change with the upstream schema. + +The v2 runtime is separate from the stable v1 API. Its methods accept and return +generated request and response models directly: + +```python +from acp.experimental import v2 + +connection = v2.connect_to_agent(MyClient(), transport) +initialized = await connection.initialize( + v2.schema.InitializeRequest( + protocol_version=v2.PROTOCOL_VERSION, + info=v2.schema.Implementation(name="my-client", version="1.0.0"), + ) +) +session = await connection.open_session( + v2.schema.NewSessionRequest(cwd="/workspace") +) +``` + +`open_session()` returns an `ActiveSession`. A v2 prompt is accepted before the +agent finishes it, so consume session updates until `SessionStop` marks the +`running` to `idle` transition: + +```python +await session.prompt( + v2.schema.PromptRequest( + session_id=session.session_id, + prompt=[v2.schema.TextContentBlock(text="Hello")], + ) +) +stopped = await session.wait_for_idle() +``` + +## Negotiate v1 or v2 + +Use `ClientNegotiator` when the same client can speak both versions. It sends +exactly one `initialize` request and returns a version-tagged connection: + +```python +from acp.experimental import ( + ClientNegotiator, + NegotiatedV2, + V1ClientConfig, + V2ClientConfig, +) + +negotiator = ClientNegotiator( + transport, + v1=V1ClientConfig(client=v1_client, initialize=v1_initialize), + v2=V2ClientConfig(client=v2_client, initialize=v2_initialize), +) +negotiated = await negotiator.negotiate() + +if isinstance(negotiated, NegotiatedV2): + session = await negotiated.connection.open_session(v2_new_session) +else: + session = await negotiated.connection.new_session(cwd="/workspace") +``` + +Agents that serve both versions use `AgentProtocolRouter`: + +```python +from acp.experimental import AgentProtocolRouter + +router = AgentProtocolRouter(v1=v1_agent, v2=v2_agent) +await router.run() +``` + +The selected runtime remains strict after initialization: v1 messages are not +accepted by a v2 connection, and v2 messages are not translated into v1 calls. +Only the initial v2 request is reduced to the common v1 initialization fields +when an agent selects v1. diff --git a/mkdocs.yml b/mkdocs.yml index f7e1b6f..1938cf4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -12,6 +12,7 @@ nav: - Quick Start: quickstart.md - Use Cases: use-cases.md - Web Transport (HTTP/WS): web-transport.md + - Experimental Protocol v2: experimental-v2.md - Experimental Contrib: contrib.md - Releasing: releasing.md - 0.11 Migration Guide: migration-guide-0.11.md diff --git a/src/acp/agent/connection.py b/src/acp/agent/connection.py index dfe6452..f8885aa 100644 --- a/src/acp/agent/connection.py +++ b/src/acp/agent/connection.py @@ -7,7 +7,7 @@ from pydantic import TypeAdapter from .._transport import Transport -from ..connection import Connection +from ..connection import Connection, MethodHandler from ..interfaces import Agent, Client from ..meta import CLIENT_METHODS from ..schema import ( @@ -88,8 +88,7 @@ def __init__( use_unstable_protocol: bool = False, **connection_kwargs: Any, ) -> None: - agent = to_agent(self) if callable(to_agent) else to_agent - handler = build_agent_router(cast(Agent, agent), use_unstable_protocol=use_unstable_protocol) + agent, handler = self._prepare(to_agent, use_unstable_protocol=use_unstable_protocol) if isinstance(input_stream, Transport): if output_stream is not None: raise TypeError(_AGENT_CONNECTION_ERROR) @@ -100,6 +99,32 @@ def __init__( ): raise TypeError(_AGENT_CONNECTION_ERROR) self._conn = Connection(handler, input_stream, output_stream, listening=listening, **connection_kwargs) + self._notify_connected(agent) + + @classmethod + def _attach( + cls, + to_agent: Callable[[Client], Agent] | Agent, + connection: Connection, + *, + use_unstable_protocol: bool = False, + ) -> tuple[AgentSideConnection, MethodHandler]: + self = cls.__new__(cls) + agent, handler = self._prepare(to_agent, use_unstable_protocol=use_unstable_protocol) + self._conn = connection + self._notify_connected(agent) + return self, handler + + def _prepare( + self, + to_agent: Callable[[Client], Agent] | Agent, + *, + use_unstable_protocol: bool, + ) -> tuple[Agent, MethodHandler]: + agent = cast(Agent, to_agent(self) if callable(to_agent) else to_agent) + return agent, build_agent_router(agent, use_unstable_protocol=use_unstable_protocol) + + def _notify_connected(self, agent: Agent) -> None: if on_connect := getattr(agent, "on_connect", None): on_connect(self) diff --git a/src/acp/client/connection.py b/src/acp/client/connection.py index 000a5da..f37802f 100644 --- a/src/acp/client/connection.py +++ b/src/acp/client/connection.py @@ -6,7 +6,7 @@ from typing import Any, cast, final from .._transport import Transport -from ..connection import Connection +from ..connection import Connection, MethodHandler from ..exceptions import RequestError from ..interfaces import Agent, Client from ..meta import AGENT_METHODS, CLIENT_METHODS @@ -122,9 +122,7 @@ def __init__( use_unstable_protocol: bool = False, **connection_kwargs: Any, ) -> None: - client = to_client(self) if callable(to_client) else to_client - self._session_updates = _SessionUpdateTracker(cast(Client, client)) - handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol) + client, handler = self._prepare(to_client, use_unstable_protocol=use_unstable_protocol) if isinstance(input_stream, Transport): if output_stream is not None: @@ -136,6 +134,34 @@ def __init__( ): raise TypeError(_CLIENT_CONNECTION_ERROR) self._conn = Connection(handler, input_stream, output_stream, **connection_kwargs) + self._notify_connected(client) + + @classmethod + def _attach( + cls, + to_client: Callable[[Agent], Client] | Client, + connection: Connection, + *, + use_unstable_protocol: bool = False, + ) -> tuple[ClientSideConnection, MethodHandler]: + self = cls.__new__(cls) + client, handler = self._prepare(to_client, use_unstable_protocol=use_unstable_protocol) + self._conn = connection + self._notify_connected(client) + return self, handler + + def _prepare( + self, + to_client: Callable[[Agent], Client] | Client, + *, + use_unstable_protocol: bool, + ) -> tuple[Client, MethodHandler]: + client = cast(Client, to_client(self) if callable(to_client) else to_client) + self._session_updates = _SessionUpdateTracker(client) + handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol) + return client, handler + + def _notify_connected(self, client: Client) -> None: if on_connect := getattr(client, "on_connect", None): on_connect(self) diff --git a/src/acp/experimental/__init__.py b/src/acp/experimental/__init__.py index 82703aa..80c89bc 100644 --- a/src/acp/experimental/__init__.py +++ b/src/acp/experimental/__init__.py @@ -1 +1,27 @@ """Experimental ACP APIs.""" + +from . import v2 +from .negotiation import ( + AgentProtocolConnection, + AgentProtocolRouter, + ClientNegotiator, + NegotiatedClient, + NegotiatedV1, + NegotiatedV2, + UnsupportedProtocolVersionError, + V1ClientConfig, + V2ClientConfig, +) + +__all__ = [ + "AgentProtocolConnection", + "AgentProtocolRouter", + "ClientNegotiator", + "NegotiatedClient", + "NegotiatedV1", + "NegotiatedV2", + "UnsupportedProtocolVersionError", + "V1ClientConfig", + "V2ClientConfig", + "v2", +] diff --git a/src/acp/experimental/negotiation.py b/src/acp/experimental/negotiation.py new file mode 100644 index 0000000..17a54b5 --- /dev/null +++ b/src/acp/experimental/negotiation.py @@ -0,0 +1,375 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, cast + +from pydantic import BaseModel + +from acp import meta as v1_meta +from acp import schema as v1_schema +from acp.agent.connection import AgentSideConnection as V1AgentSideConnection +from acp.client.connection import ClientSideConnection as V1ClientSideConnection +from acp.connection import Connection, MethodHandler +from acp.exceptions import RequestError +from acp.interfaces import Agent as V1Agent +from acp.interfaces import Client as V1Client + +from . import v2 +from .v2._connection import open_connection +from .v2.agent import AgentFactory as V2AgentFactory +from .v2.agent import AgentSideConnection as V2AgentSideConnection +from .v2.client import ClientFactory as V2ClientFactory +from .v2.client import ClientSideConnection as V2ClientSideConnection + +__all__ = [ + "AgentProtocolConnection", + "AgentProtocolRouter", + "ClientNegotiator", + "NegotiatedClient", + "NegotiatedV1", + "NegotiatedV2", + "UnsupportedProtocolVersionError", + "V1ClientConfig", + "V2ClientConfig", +] + +V1AgentFactory = Callable[[V1Client], V1Agent] +V1ClientFactory = Callable[[V1Agent], V1Client] + + +def _dump(model: BaseModel) -> dict[str, Any]: + return model.model_dump(mode="json", by_alias=True, exclude_none=True, exclude_unset=True) + + +def _read_protocol_version(params: Any) -> int: + if not isinstance(params, dict): + raise RequestError.invalid_params({"details": "initialize params must be an object"}) + version = params.get("protocolVersion") + if isinstance(version, bool) or not isinstance(version, int) or not 0 <= version <= 0xFFFF: + raise RequestError.invalid_params({"details": "initialize.protocolVersion must be an integer from 0 to 65535"}) + return version + + +def _v2_initialize_to_v1(request: v2.schema.InitializeRequest) -> v1_schema.InitializeRequest: + capabilities = ( + v1_schema.ClientCapabilities.model_validate(_dump(request.capabilities)) + if request.capabilities is not None + else None + ) + return v1_schema.InitializeRequest( + protocol_version=v1_meta.PROTOCOL_VERSION, + client_capabilities=capabilities, + client_info=v1_schema.Implementation.model_validate(_dump(request.info)), + field_meta=request.field_meta, + ) + + +def _normalize_initialize(params: Any, selected_version: int) -> dict[str, Any]: + requested_version = _read_protocol_version(params) + if selected_version == v2.PROTOCOL_VERSION: + request = v2.schema.InitializeRequest.model_validate(params) + request.protocol_version = v2.PROTOCOL_VERSION + return _dump(request) + if requested_version >= v2.PROTOCOL_VERSION: + request = v2.schema.InitializeRequest.model_validate(params) + return _dump(_v2_initialize_to_v1(request)) + request = v1_schema.InitializeRequest.model_validate(params) + request.protocol_version = v1_meta.PROTOCOL_VERSION + return _dump(request) + + +class _SwitchingHandler: + def __init__(self) -> None: + self._handler: MethodHandler | None = None + self._failure: BaseException | None = None + self._ready = asyncio.Event() + + def bind(self, handler: MethodHandler) -> None: + if self._handler is not None or self._failure is not None: + raise RuntimeError("Protocol handler has already been resolved") + self._handler = handler + self._ready.set() + + def fail(self, error: BaseException) -> None: + if self._handler is not None: + return + self._failure = error + self._ready.set() + + async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any: + await self._ready.wait() + if self._failure is not None: + raise self._failure + if self._handler is None: + raise RuntimeError("Protocol handler was not resolved") + return await self._handler(method, params, is_notification) + + +class _AgentNegotiationHandler: + def __init__( + self, + v1_agent: V1AgentFactory | V1Agent | None, + v2_agent: V2AgentFactory | v2.Agent | None, + ) -> None: + self._v1_agent = v1_agent + self._v2_agent = v2_agent + self._connection: Connection | None = None + self._selected: MethodHandler | None = None + self._endpoint: V1AgentSideConnection | V2AgentSideConnection | None = None + self._lock = asyncio.Lock() + + def bind_connection(self, connection: Connection) -> None: + self._connection = connection + + async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any: + async with self._lock: + if self._selected is None: + return await self._initialize(method, params, is_notification) + if not is_notification and method == v2.AGENT_METHODS["initialize"]: + raise RequestError.invalid_request({"details": "ACP connections may only be initialized once"}) + handler = self._selected + return await handler(method, params, is_notification) + + async def _initialize(self, method: str, params: Any, is_notification: bool) -> Any: + if is_notification or method != v2.AGENT_METHODS["initialize"]: + raise RequestError.invalid_request({"details": "The first ACP request must be initialize"}) + requested = _read_protocol_version(params) + selected = self._select(requested) + connection = self._connection + if connection is None: + raise RuntimeError("Protocol router is not connected") + + if selected == v2.PROTOCOL_VERSION: + endpoint, handler = V2AgentSideConnection._attach(cast(Any, self._v2_agent), connection) + else: + endpoint, handler = V1AgentSideConnection._attach(cast(Any, self._v1_agent), connection) + self._endpoint = endpoint + self._selected = handler + normalized = _normalize_initialize(params, selected) + response = await handler(method, normalized, False) + parsed_version = _read_protocol_version(_dump(response) if isinstance(response, BaseModel) else response) + if parsed_version != selected: + raise RequestError.invalid_request({ + "details": f"initialize response selected protocol {parsed_version}, expected {selected}" + }) + return response + + def _select(self, requested: int) -> int: + if self._v2_agent is not None and requested >= v2.PROTOCOL_VERSION: + return v2.PROTOCOL_VERSION + if self._v1_agent is not None and requested >= v1_meta.PROTOCOL_VERSION: + return v1_meta.PROTOCOL_VERSION + supported = [ + version + for version, implementation in ( + (v1_meta.PROTOCOL_VERSION, self._v1_agent), + (v2.PROTOCOL_VERSION, self._v2_agent), + ) + if implementation is not None + ] + raise RequestError.invalid_request({ + "details": f"Unsupported ACP protocol {requested}; configured versions are {supported}" + }) + + +class AgentProtocolConnection: + def __init__(self, connection: Connection) -> None: + self._connection = connection + + async def listen(self) -> None: + await self._connection.main_loop() + + async def close(self) -> None: + await self._connection.close() + + async def __aenter__(self) -> AgentProtocolConnection: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + await self.close() + + +class AgentProtocolRouter: + """Select one strict agent runtime from the first initialize request.""" + + def __init__( + self, + *, + v1: V1AgentFactory | V1Agent | None = None, + v2: V2AgentFactory | v2.Agent | None = None, + ) -> None: + if v1 is None and v2 is None: + raise ValueError("Configure at least one ACP protocol implementation") + self._v1 = v1 + self._v2 = v2 + + def connect( + self, + input_stream: Any, + output_stream: Any = None, + *, + listening: bool = True, + **connection_kwargs: Any, + ) -> AgentProtocolConnection: + handler = _AgentNegotiationHandler(self._v1, self._v2) + connection = open_connection( + handler, + input_stream, + output_stream, + listening=listening, + **connection_kwargs, + ) + handler.bind_connection(connection) + return AgentProtocolConnection(connection) + + async def run( + self, + input_stream: Any = None, + output_stream: Any = None, + *, + stdio_buffer_limit_bytes: int = 50 * 1024 * 1024, + **connection_kwargs: Any, + ) -> None: + if input_stream is None and output_stream is None: + from acp.stdio import stdio_streams + + output_stream, input_stream = await stdio_streams(limit=stdio_buffer_limit_bytes) + connection = self.connect( + input_stream, + output_stream, + listening=False, + **connection_kwargs, + ) + try: + await connection.listen() + finally: + await asyncio.shield(connection.close()) + + +@dataclass(frozen=True, slots=True) +class V1ClientConfig: + client: V1ClientFactory | V1Client + initialize: v1_schema.InitializeRequest + + def __post_init__(self) -> None: + if self.initialize.protocol_version != v1_meta.PROTOCOL_VERSION: + raise ValueError(f"V1ClientConfig requires protocol version {v1_meta.PROTOCOL_VERSION}") + + +@dataclass(frozen=True, slots=True) +class V2ClientConfig: + client: V2ClientFactory | v2.Client + initialize: v2.schema.InitializeRequest + + def __post_init__(self) -> None: + if self.initialize.protocol_version != v2.PROTOCOL_VERSION: + raise ValueError(f"V2ClientConfig requires protocol version {v2.PROTOCOL_VERSION}") + + +@dataclass(frozen=True, slots=True) +class NegotiatedV1: + connection: V1ClientSideConnection + initialize: v1_schema.InitializeResponse + protocol_version: int = v1_meta.PROTOCOL_VERSION + + +@dataclass(frozen=True, slots=True) +class NegotiatedV2: + connection: V2ClientSideConnection + initialize: v2.schema.InitializeResponse + protocol_version: int = v2.PROTOCOL_VERSION + + +NegotiatedClient = NegotiatedV1 | NegotiatedV2 + + +class UnsupportedProtocolVersionError(ValueError): + def __init__(self, requested: int, offered: int, supported: frozenset[int]) -> None: + self.requested = requested + self.offered = offered + self.supported = supported + super().__init__(f"Agent selected ACP protocol {offered}; requested {requested}, supported {sorted(supported)}") + + +class ClientNegotiator: + """Send one initialize request and return the selected typed client.""" + + def __init__( + self, + input_stream: Any, + output_stream: Any = None, + *, + v1: V1ClientConfig | None = None, + v2: V2ClientConfig | None = None, + **connection_kwargs: Any, + ) -> None: + if v1 is None and v2 is None: + raise ValueError("Configure at least one ACP client version") + self._v1 = v1 + self._v2 = v2 + self._handler = _SwitchingHandler() + self._connection = open_connection( + self._handler, + input_stream, + output_stream, + **connection_kwargs, + ) + self._lock = asyncio.Lock() + self._resolved: NegotiatedClient | None = None + self._failure: BaseException | None = None + + async def negotiate(self) -> NegotiatedClient: + async with self._lock: + if self._resolved is not None: + return self._resolved + if self._failure is not None: + raise self._failure + try: + self._resolved = await self._negotiate_once() + except BaseException as error: + self._failure = error + self._handler.fail(error) + await self._connection.close() + raise + return self._resolved + + async def _negotiate_once(self) -> NegotiatedClient: + offered_request: BaseModel = ( + self._v2.initialize if self._v2 is not None else cast(V1ClientConfig, self._v1).initialize + ) + + response = await self._connection.send_request(v2.AGENT_METHODS["initialize"], _dump(offered_request)) + offered = _read_protocol_version(response) + requested = offered_request.protocol_version + + if offered == v2.PROTOCOL_VERSION and self._v2 is not None: + initialize = v2.schema.InitializeResponse.model_validate(response) + connection, handler = V2ClientSideConnection._attach(self._v2.client, self._connection) + connection._complete_initialization(self._v2.initialize, initialize) + self._handler.bind(handler) + return NegotiatedV2(connection, initialize) + if offered == v1_meta.PROTOCOL_VERSION and self._v1 is not None: + initialize = v1_schema.InitializeResponse.model_validate(response) + connection, handler = V1ClientSideConnection._attach(self._v1.client, self._connection) + self._handler.bind(handler) + return NegotiatedV1(connection, initialize) + supported = frozenset( + version + for version, config in ( + (v1_meta.PROTOCOL_VERSION, self._v1), + (v2.PROTOCOL_VERSION, self._v2), + ) + if config is not None + ) + raise UnsupportedProtocolVersionError(requested, offered, supported) + + async def close(self) -> None: + await self._connection.close() + + async def __aenter__(self) -> ClientNegotiator: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + await self.close() diff --git a/src/acp/experimental/v2/__init__.py b/src/acp/experimental/v2/__init__.py index c5ee37e..6bff230 100644 --- a/src/acp/experimental/v2/__init__.py +++ b/src/acp/experimental/v2/__init__.py @@ -1,5 +1,26 @@ -"""Experimental ACP protocol v2 bindings.""" +"""Experimental ACP protocol v2 API.""" +from . import schema +from .agent import AgentSideConnection, run_agent +from .client import ClientSideConnection, connect_to_agent +from .interfaces import Agent, Client from .meta import AGENT_METHODS, CLIENT_METHODS, PROTOCOL_METHODS, PROTOCOL_VERSION +from .session import ActiveSession, SessionMessage, SessionStop, SessionUpdate -__all__ = ["AGENT_METHODS", "CLIENT_METHODS", "PROTOCOL_METHODS", "PROTOCOL_VERSION"] +__all__ = [ + "AGENT_METHODS", + "CLIENT_METHODS", + "PROTOCOL_METHODS", + "PROTOCOL_VERSION", + "ActiveSession", + "Agent", + "AgentSideConnection", + "Client", + "ClientSideConnection", + "SessionMessage", + "SessionStop", + "SessionUpdate", + "connect_to_agent", + "run_agent", + "schema", +] diff --git a/src/acp/experimental/v2/_connection.py b/src/acp/experimental/v2/_connection.py new file mode 100644 index 0000000..f9915ac --- /dev/null +++ b/src/acp/experimental/v2/_connection.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +from acp._transport import Transport +from acp.connection import Connection, MethodHandler + + +def open_connection( + handler: MethodHandler, + input_stream: Any, + output_stream: Any = None, + *, + listening: bool = True, + **connection_kwargs: Any, +) -> Connection: + if isinstance(input_stream, Transport): + if output_stream is not None: + raise TypeError("A message transport cannot be combined with an output stream") + return Connection(handler, input_stream, listening=listening, **connection_kwargs) + if not isinstance(input_stream, asyncio.StreamWriter) or not isinstance(output_stream, asyncio.StreamReader): + raise TypeError("Expected an asyncio StreamWriter/StreamReader pair or a message transport") + return Connection(handler, input_stream, output_stream, listening=listening, **connection_kwargs) diff --git a/src/acp/experimental/v2/_initialization.py b/src/acp/experimental/v2/_initialization.py new file mode 100644 index 0000000..50e8b61 --- /dev/null +++ b/src/acp/experimental/v2/_initialization.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Literal + +from acp.exceptions import RequestError + +from . import schema +from .meta import PROTOCOL_VERSION + +InitializationPhase = Literal["uninitialized", "initializing", "initialized", "failed"] + + +@dataclass(frozen=True, slots=True) +class Initialization: + request: schema.InitializeRequest + response: schema.InitializeResponse + + +class InitializationState: + def __init__(self) -> None: + self._phase: InitializationPhase = "uninitialized" + self._request: schema.InitializeRequest | None = None + self._initialization: Initialization | None = None + self._failure: BaseException | None = None + self._ready = asyncio.Event() + + @property + def phase(self) -> InitializationPhase: + return self._phase + + def begin(self, request: schema.InitializeRequest) -> None: + if self._phase != "uninitialized": + raise RequestError.invalid_request({"details": "ACP v2 connections may only be initialized once"}) + if request.protocol_version != PROTOCOL_VERSION: + raise RequestError.invalid_params({ + "expectedProtocolVersion": PROTOCOL_VERSION, + "receivedProtocolVersion": request.protocol_version, + }) + self._request = request.model_copy(deep=True) + self._phase = "initializing" + + def complete(self, response: schema.InitializeResponse) -> Initialization: + if self._phase != "initializing" or self._request is None: + raise RequestError.invalid_request({"details": "ACP v2 initialization is not in progress"}) + if response.protocol_version != PROTOCOL_VERSION: + raise RequestError.invalid_request({ + "expectedProtocolVersion": PROTOCOL_VERSION, + "receivedProtocolVersion": response.protocol_version, + }) + initialization = Initialization( + request=self._request.model_copy(deep=True), + response=response.model_copy(deep=True), + ) + self._initialization = initialization + self._phase = "initialized" + self._ready.set() + return initialization + + def fail(self, error: BaseException) -> None: + if self._phase == "initialized": + return + self._phase = "failed" + self._failure = error + self._ready.set() + + async def initialized(self) -> Initialization: + if self._phase in {"uninitialized", "initializing"}: + await self._ready.wait() + if self._initialization is not None: + return self._initialization + if self._failure is not None: + raise self._failure + raise RequestError.invalid_request({"details": "ACP v2 connection has not been initialized"}) + + async def require(self, method: str) -> None: + if self._phase == "initialized": + return + if self._phase == "initializing": + await self.initialized() + return + raise RequestError.invalid_request({"details": f"ACP v2 connection must be initialized before {method!r}"}) diff --git a/src/acp/experimental/v2/_methods.py b/src/acp/experimental/v2/_methods.py new file mode 100644 index 0000000..f89507e --- /dev/null +++ b/src/acp/experimental/v2/_methods.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from pydantic import TypeAdapter + +from . import schema +from .meta import AGENT_METHODS, CLIENT_METHODS + + +@dataclass(frozen=True, slots=True) +class RequestSpec: + method: str + handler: str + request: TypeAdapter[Any] + response: TypeAdapter[Any] + empty_response: bool = False + + +@dataclass(frozen=True, slots=True) +class NotificationSpec: + method: str + handler: str + params: TypeAdapter[Any] + + +def request( + method: str, + handler: str, + request_type: Any, + response_type: Any, + *, + empty_response: bool = False, +) -> RequestSpec: + return RequestSpec( + method=method, + handler=handler, + request=TypeAdapter(request_type), + response=TypeAdapter(response_type), + empty_response=empty_response, + ) + + +def notification(method: str, handler: str, params_type: Any) -> NotificationSpec: + return NotificationSpec(method=method, handler=handler, params=TypeAdapter(params_type)) + + +SetConfigOptionRequest = ( + schema.SetSessionConfigOptionIdRequest + | schema.SetSessionConfigOptionBooleanRequest + | schema.SetSessionConfigOptionOtherRequest +) + +CreateElicitationRequest = ( + schema.CreateOtherSessionElicitationRequest + | schema.CreateOtherRequestElicitationRequest + | schema.CreateFormSessionElicitationRequest + | schema.CreateFormRequestElicitationRequest + | schema.CreateUrlSessionElicitationRequest + | schema.CreateUrlRequestElicitationRequest +) + +CreateElicitationResponse = ( + schema.AcceptElicitationResponse + | schema.DeclineElicitationResponse + | schema.CancelElicitationResponse + | schema.OtherElicitationResponse +) + + +AGENT_REQUESTS = ( + request(AGENT_METHODS["initialize"], "initialize", schema.InitializeRequest, schema.InitializeResponse), + request( + AGENT_METHODS["auth_login"], + "login", + schema.LoginAuthRequest, + schema.LoginAuthResponse, + empty_response=True, + ), + request( + AGENT_METHODS["providers_list"], "list_providers", schema.ListProvidersRequest, schema.ListProvidersResponse + ), + request( + AGENT_METHODS["providers_set"], + "set_provider", + schema.SetProviderRequest, + schema.SetProviderResponse, + empty_response=True, + ), + request( + AGENT_METHODS["providers_disable"], + "disable_provider", + schema.DisableProviderRequest, + schema.DisableProviderResponse, + empty_response=True, + ), + request(AGENT_METHODS["session_new"], "new_session", schema.NewSessionRequest, schema.NewSessionResponse), + request( + AGENT_METHODS["session_set_config_option"], + "set_config_option", + SetConfigOptionRequest, + schema.SetSessionConfigOptionResponse, + ), + request( + AGENT_METHODS["session_prompt"], + "prompt", + schema.PromptRequest, + schema.PromptResponse, + empty_response=True, + ), + request(AGENT_METHODS["mcp_message"], "message_mcp", schema.MessageMcpRequest, Any), + request(AGENT_METHODS["session_list"], "list_sessions", schema.ListSessionsRequest, schema.ListSessionsResponse), + request( + AGENT_METHODS["session_delete"], + "delete_session", + schema.DeleteSessionRequest, + schema.DeleteSessionResponse, + empty_response=True, + ), + request(AGENT_METHODS["session_fork"], "fork_session", schema.ForkSessionRequest, schema.ForkSessionResponse), + request( + AGENT_METHODS["session_resume"], "resume_session", schema.ResumeSessionRequest, schema.ResumeSessionResponse + ), + request( + AGENT_METHODS["session_close"], + "close_session", + schema.CloseSessionRequest, + schema.CloseSessionResponse, + empty_response=True, + ), + request( + AGENT_METHODS["auth_logout"], + "logout", + schema.LogoutAuthRequest, + schema.LogoutAuthResponse, + empty_response=True, + ), + request(AGENT_METHODS["nes_start"], "start_nes", schema.StartNesRequest, schema.StartNesResponse), + request(AGENT_METHODS["nes_suggest"], "suggest_nes", schema.SuggestNesRequest, schema.SuggestNesResponse), + request( + AGENT_METHODS["nes_close"], + "close_nes", + schema.CloseNesRequest, + schema.CloseNesResponse, + empty_response=True, + ), +) + +AGENT_NOTIFICATIONS = ( + notification(AGENT_METHODS["session_cancel"], "cancel", schema.CancelSessionNotification), + notification(AGENT_METHODS["mcp_message"], "message_mcp", schema.MessageMcpNotification), + notification(AGENT_METHODS["document_did_open"], "did_open", schema.DidOpenDocumentNotification), + notification(AGENT_METHODS["document_did_change"], "did_change", schema.DidChangeDocumentNotification), + notification(AGENT_METHODS["document_did_close"], "did_close", schema.DidCloseDocumentNotification), + notification(AGENT_METHODS["document_did_save"], "did_save", schema.DidSaveDocumentNotification), + notification(AGENT_METHODS["document_did_focus"], "did_focus", schema.DidFocusDocumentNotification), + notification(AGENT_METHODS["nes_accept"], "accept_nes", schema.AcceptNesNotification), + notification(AGENT_METHODS["nes_reject"], "reject_nes", schema.RejectNesNotification), +) + +CLIENT_REQUESTS = ( + request( + CLIENT_METHODS["session_request_permission"], + "request_permission", + schema.RequestPermissionRequest, + schema.RequestPermissionResponse, + ), + request(CLIENT_METHODS["mcp_connect"], "connect_mcp", schema.ConnectMcpRequest, schema.ConnectMcpResponse), + request(CLIENT_METHODS["mcp_message"], "message_mcp", schema.MessageMcpRequest, Any), + request( + CLIENT_METHODS["mcp_disconnect"], + "disconnect_mcp", + schema.DisconnectMcpRequest, + schema.DisconnectMcpResponse, + empty_response=True, + ), + request( + CLIENT_METHODS["elicitation_create"], + "create_elicitation", + CreateElicitationRequest, + CreateElicitationResponse, + ), +) + +CLIENT_NOTIFICATIONS = ( + notification(CLIENT_METHODS["session_update"], "session_update", schema.UpdateSessionNotification), + notification(CLIENT_METHODS["mcp_message"], "message_mcp", schema.MessageMcpNotification), + notification( + CLIENT_METHODS["elicitation_complete"], + "complete_elicitation", + schema.CompleteElicitationNotification, + ), +) + + +AGENT_REQUESTS_BY_METHOD = {spec.method: spec for spec in AGENT_REQUESTS} +AGENT_NOTIFICATIONS_BY_METHOD = {spec.method: spec for spec in AGENT_NOTIFICATIONS} +CLIENT_REQUESTS_BY_METHOD = {spec.method: spec for spec in CLIENT_REQUESTS} +CLIENT_NOTIFICATIONS_BY_METHOD = {spec.method: spec for spec in CLIENT_NOTIFICATIONS} diff --git a/src/acp/experimental/v2/_router.py b/src/acp/experimental/v2/_router.py new file mode 100644 index 0000000..904ba08 --- /dev/null +++ b/src/acp/experimental/v2/_router.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from acp.exceptions import RequestError + +from ._methods import NotificationSpec, RequestSpec + +ExtensionRequest = Callable[[str, Any], Awaitable[Any]] +ExtensionNotification = Callable[[str, Any], Awaitable[None]] + + +class MethodRouter: + def __init__( + self, + target: Any, + requests: tuple[RequestSpec, ...], + notifications: tuple[NotificationSpec, ...], + ) -> None: + self._target = target + self._requests = {spec.method: spec for spec in requests} + self._notifications = {spec.method: spec for spec in notifications} + + def request_spec(self, method: str) -> RequestSpec | None: + return self._requests.get(method) + + async def handle_request(self, spec: RequestSpec, params: Any) -> Any: + handler = getattr(self._target, spec.handler, None) + if handler is None: + raise RequestError.method_not_found(spec.method) + request = spec.request.validate_python(params) + response = await handler(request) + if response is None and spec.empty_response: + response = {} + return spec.response.validate_python(response) + + async def handle_notification(self, spec: NotificationSpec, params: Any) -> None: + handler = getattr(self._target, spec.handler, None) + if handler is None: + raise RequestError.method_not_found(spec.method) + await handler(spec.params.validate_python(params)) + + async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any: + if method.startswith("_"): + return await self._handle_extension(method, params, is_notification) + if is_notification: + spec = self._notifications.get(method) + if spec is None: + raise RequestError.method_not_found(method) + await self.handle_notification(spec, params) + return None + spec = self._requests.get(method) + if spec is None: + raise RequestError.method_not_found(method) + return await self.handle_request(spec, params) + + async def _handle_extension(self, method: str, params: Any, is_notification: bool) -> Any: + handler_name = "ext_notification" if is_notification else "ext_method" + handler = getattr(self._target, handler_name, None) + if handler is None: + raise RequestError.method_not_found(method) + return await handler(method, params) diff --git a/src/acp/experimental/v2/agent.py b/src/acp/experimental/v2/agent.py new file mode 100644 index 0000000..547209e --- /dev/null +++ b/src/acp/experimental/v2/agent.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any, cast + +from pydantic import BaseModel + +from acp.connection import Connection + +from . import schema +from ._connection import open_connection +from ._initialization import Initialization, InitializationState +from ._methods import ( + AGENT_NOTIFICATIONS, + AGENT_REQUESTS, + CLIENT_REQUESTS_BY_METHOD, + CreateElicitationRequest, + CreateElicitationResponse, +) +from ._router import MethodRouter +from .interfaces import Agent, Client +from .meta import CLIENT_METHODS, PROTOCOL_METHODS + +__all__ = ["AgentSideConnection", "run_agent"] + +AgentFactory = Callable[[Client], Agent] + + +def _dump(model: BaseModel) -> dict[str, Any]: + return model.model_dump(mode="json", by_alias=True, exclude_none=True, exclude_unset=True) + + +class _AgentRouter: + def __init__(self, agent: Agent, state: InitializationState) -> None: + self._router = MethodRouter(agent, AGENT_REQUESTS, AGENT_NOTIFICATIONS) + self._state = state + + async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any: + if is_notification and method == PROTOCOL_METHODS["cancel_request"]: + if self._state.phase not in {"initializing", "initialized"}: + await self._state.require(method) + return None + + initialize = self._router.request_spec("initialize") + if not is_notification and initialize is not None and method == initialize.method: + request = cast(schema.InitializeRequest, initialize.request.validate_python(params)) + self._state.begin(request) + try: + response = cast(schema.InitializeResponse, await self._router.handle_request(initialize, params)) + self._state.complete(response) + except BaseException as error: + self._state.fail(error) + raise + return response + + await self._state.require(method) + return await self._router(method, params, is_notification) + + +class AgentSideConnection: + """Strict experimental ACP v2 connection used by an agent.""" + + def __init__( + self, + to_agent: AgentFactory | Agent, + input_stream: Any, + output_stream: Any = None, + *, + listening: bool = True, + **connection_kwargs: Any, + ) -> None: + self._state = InitializationState() + agent = to_agent(self) if callable(to_agent) else to_agent + router = _AgentRouter(cast(Agent, agent), self._state) + self._conn = open_connection( + router, + input_stream, + output_stream, + listening=listening, + **connection_kwargs, + ) + if on_connect := getattr(agent, "on_connect", None): + on_connect(self) + + @classmethod + def _attach( + cls, to_agent: AgentFactory | Agent, connection: Connection + ) -> tuple[AgentSideConnection, _AgentRouter]: + self = cls.__new__(cls) + self._state = InitializationState() + self._conn = connection + agent = to_agent(self) if callable(to_agent) else to_agent + router = _AgentRouter(cast(Agent, agent), self._state) + if on_connect := getattr(agent, "on_connect", None): + on_connect(self) + return self, router + + async def wait_until_initialized(self) -> Initialization: + return await self._state.initialized() + + async def listen(self) -> None: + await self._conn.main_loop() + + async def request_permission( + self, + request: schema.RequestPermissionRequest, + ) -> schema.RequestPermissionResponse: + return await self._request( + CLIENT_METHODS["session_request_permission"], + request, + ) + + async def session_update(self, notification: schema.UpdateSessionNotification) -> None: + await self._notify(CLIENT_METHODS["session_update"], notification) + + async def connect_mcp(self, request: schema.ConnectMcpRequest) -> schema.ConnectMcpResponse: + return await self._request(CLIENT_METHODS["mcp_connect"], request) + + async def message_mcp(self, message: schema.MessageMcpRequest) -> Any: + return await self._request(CLIENT_METHODS["mcp_message"], message) + + async def notify_mcp(self, notification: schema.MessageMcpNotification) -> None: + await self._notify(CLIENT_METHODS["mcp_message"], notification) + + async def disconnect_mcp( + self, + request: schema.DisconnectMcpRequest, + ) -> schema.DisconnectMcpResponse: + return await self._request(CLIENT_METHODS["mcp_disconnect"], request) + + async def create_elicitation(self, request: CreateElicitationRequest) -> CreateElicitationResponse: + return await self._request(CLIENT_METHODS["elicitation_create"], request) + + async def complete_elicitation(self, notification: schema.CompleteElicitationNotification) -> None: + await self._notify(CLIENT_METHODS["elicitation_complete"], notification) + + async def ext_method(self, method: str, params: Any = None) -> Any: + await self._state.require(method) + return await self._conn.send_request(_extension_method(method), params) + + async def ext_notification(self, method: str, params: Any = None) -> None: + await self._state.require(method) + await self._conn.send_notification(_extension_method(method), params) + + async def close(self) -> None: + await self._conn.close() + + async def _request(self, method: str, request: BaseModel) -> Any: + await self._state.require(method) + spec = CLIENT_REQUESTS_BY_METHOD[method] + response = await self._conn.send_request(method, _dump(request)) + if response is None and spec.empty_response: + response = {} + return spec.response.validate_python(response) + + async def _notify(self, method: str, notification: BaseModel) -> None: + await self._state.require(method) + await self._conn.send_notification(method, _dump(notification)) + + async def __aenter__(self) -> AgentSideConnection: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + await self.close() + + +def _extension_method(method: str) -> str: + return method if method.startswith("_") else f"_{method}" + + +async def run_agent( + agent: AgentFactory | Agent, + input_stream: Any = None, + output_stream: Any = None, + *, + stdio_buffer_limit_bytes: int = 50 * 1024 * 1024, + **connection_kwargs: Any, +) -> None: + if input_stream is None and output_stream is None: + from acp.stdio import stdio_streams + + output_stream, input_stream = await stdio_streams(limit=stdio_buffer_limit_bytes) + connection = AgentSideConnection( + agent, + input_stream, + output_stream, + listening=False, + **connection_kwargs, + ) + try: + await connection.listen() + finally: + await asyncio.shield(connection.close()) diff --git a/src/acp/experimental/v2/client.py b/src/acp/experimental/v2/client.py new file mode 100644 index 0000000..57d775e --- /dev/null +++ b/src/acp/experimental/v2/client.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, cast + +from pydantic import BaseModel + +from acp.connection import Connection + +from . import schema +from ._connection import open_connection +from ._initialization import Initialization, InitializationState +from ._methods import ( + AGENT_REQUESTS_BY_METHOD, + CLIENT_NOTIFICATIONS, + CLIENT_REQUESTS, + SetConfigOptionRequest, +) +from ._router import MethodRouter +from .agent import _dump, _extension_method +from .interfaces import Agent, Client +from .meta import AGENT_METHODS, PROTOCOL_METHODS +from .session import ActiveSession, SessionUpdateBroker + +__all__ = ["ClientSideConnection", "connect_to_agent"] + +ClientFactory = Callable[[Agent], Client] + + +class _ClientRouter: + def __init__(self, client: Client, state: InitializationState) -> None: + self._router = MethodRouter(client, CLIENT_REQUESTS, CLIENT_NOTIFICATIONS) + self._state = state + + async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any: + if is_notification and method == PROTOCOL_METHODS["cancel_request"]: + if self._state.phase not in {"initializing", "initialized"}: + await self._state.require(method) + return None + await self._state.require(method) + return await self._router(method, params, is_notification) + + +class ClientSideConnection: + """Strict experimental ACP v2 connection used by a client.""" + + def __init__( + self, + to_client: ClientFactory | Client, + input_stream: Any, + output_stream: Any = None, + **connection_kwargs: Any, + ) -> None: + self._state = InitializationState() + client = to_client(self) if callable(to_client) else to_client + self._session_updates = SessionUpdateBroker(cast(Client, client)) + router = _ClientRouter(cast(Client, self._session_updates), self._state) + self._conn = open_connection(router, input_stream, output_stream, **connection_kwargs) + if on_connect := getattr(client, "on_connect", None): + on_connect(self) + + @classmethod + def _attach( + cls, + to_client: ClientFactory | Client, + connection: Connection, + ) -> tuple[ClientSideConnection, _ClientRouter]: + self = cls.__new__(cls) + self._state = InitializationState() + self._conn = connection + client = to_client(self) if callable(to_client) else to_client + self._session_updates = SessionUpdateBroker(cast(Client, client)) + router = _ClientRouter(cast(Client, self._session_updates), self._state) + if on_connect := getattr(client, "on_connect", None): + on_connect(self) + return self, router + + def _complete_initialization( + self, + request: schema.InitializeRequest, + response: schema.InitializeResponse, + ) -> Initialization: + self._state.begin(request) + return self._state.complete(response) + + async def wait_until_initialized(self) -> Initialization: + return await self._state.initialized() + + async def initialize(self, request: schema.InitializeRequest) -> schema.InitializeResponse: + self._state.begin(request) + try: + response = await self._conn.send_request(AGENT_METHODS["initialize"], _dump(request)) + parsed = schema.InitializeResponse.model_validate(response) + self._state.complete(parsed) + except BaseException as error: + self._state.fail(error) + await self._conn.close() + raise + return parsed + + async def login(self, request: schema.LoginAuthRequest) -> schema.LoginAuthResponse: + return await self._request(AGENT_METHODS["auth_login"], request) + + async def logout(self, request: schema.LogoutAuthRequest) -> schema.LogoutAuthResponse: + return await self._request(AGENT_METHODS["auth_logout"], request) + + async def list_providers(self, request: schema.ListProvidersRequest) -> schema.ListProvidersResponse: + return await self._request(AGENT_METHODS["providers_list"], request) + + async def set_provider(self, request: schema.SetProviderRequest) -> schema.SetProviderResponse: + return await self._request(AGENT_METHODS["providers_set"], request) + + async def disable_provider(self, request: schema.DisableProviderRequest) -> schema.DisableProviderResponse: + return await self._request(AGENT_METHODS["providers_disable"], request) + + async def new_session(self, request: schema.NewSessionRequest) -> schema.NewSessionResponse: + return await self._request(AGENT_METHODS["session_new"], request) + + async def open_session(self, request: schema.NewSessionRequest) -> ActiveSession: + self._session_updates.begin_capture() + try: + response = await self.new_session(request) + return ActiveSession(self, response, self._session_updates) + finally: + self._session_updates.end_capture() + + async def list_sessions(self, request: schema.ListSessionsRequest) -> schema.ListSessionsResponse: + return await self._request(AGENT_METHODS["session_list"], request) + + async def delete_session(self, request: schema.DeleteSessionRequest) -> schema.DeleteSessionResponse: + return await self._request(AGENT_METHODS["session_delete"], request) + + async def fork_session(self, request: schema.ForkSessionRequest) -> schema.ForkSessionResponse: + return await self._request(AGENT_METHODS["session_fork"], request) + + async def resume_session(self, request: schema.ResumeSessionRequest) -> schema.ResumeSessionResponse: + return await self._request(AGENT_METHODS["session_resume"], request) + + async def close_session(self, request: schema.CloseSessionRequest) -> schema.CloseSessionResponse: + return await self._request(AGENT_METHODS["session_close"], request) + + async def set_config_option(self, request: SetConfigOptionRequest) -> schema.SetSessionConfigOptionResponse: + return await self._request(AGENT_METHODS["session_set_config_option"], request) + + async def prompt(self, request: schema.PromptRequest) -> schema.PromptResponse: + return await self._request(AGENT_METHODS["session_prompt"], request) + + async def cancel(self, notification: schema.CancelSessionNotification) -> None: + await self._notify(AGENT_METHODS["session_cancel"], notification) + + async def message_mcp(self, message: schema.MessageMcpRequest) -> Any: + return await self._request(AGENT_METHODS["mcp_message"], message) + + async def notify_mcp(self, notification: schema.MessageMcpNotification) -> None: + await self._notify(AGENT_METHODS["mcp_message"], notification) + + async def start_nes(self, request: schema.StartNesRequest) -> schema.StartNesResponse: + return await self._request(AGENT_METHODS["nes_start"], request) + + async def suggest_nes(self, request: schema.SuggestNesRequest) -> schema.SuggestNesResponse: + return await self._request(AGENT_METHODS["nes_suggest"], request) + + async def accept_nes(self, notification: schema.AcceptNesNotification) -> None: + await self._notify(AGENT_METHODS["nes_accept"], notification) + + async def reject_nes(self, notification: schema.RejectNesNotification) -> None: + await self._notify(AGENT_METHODS["nes_reject"], notification) + + async def close_nes(self, request: schema.CloseNesRequest) -> schema.CloseNesResponse: + return await self._request(AGENT_METHODS["nes_close"], request) + + async def did_open(self, notification: schema.DidOpenDocumentNotification) -> None: + await self._notify(AGENT_METHODS["document_did_open"], notification) + + async def did_change(self, notification: schema.DidChangeDocumentNotification) -> None: + await self._notify(AGENT_METHODS["document_did_change"], notification) + + async def did_close(self, notification: schema.DidCloseDocumentNotification) -> None: + await self._notify(AGENT_METHODS["document_did_close"], notification) + + async def did_save(self, notification: schema.DidSaveDocumentNotification) -> None: + await self._notify(AGENT_METHODS["document_did_save"], notification) + + async def did_focus(self, notification: schema.DidFocusDocumentNotification) -> None: + await self._notify(AGENT_METHODS["document_did_focus"], notification) + + async def ext_method(self, method: str, params: Any = None) -> Any: + await self._state.require(method) + return await self._conn.send_request(_extension_method(method), params) + + async def ext_notification(self, method: str, params: Any = None) -> None: + await self._state.require(method) + await self._conn.send_notification(_extension_method(method), params) + + async def close(self) -> None: + await self._conn.close() + + async def _request(self, method: str, request: BaseModel) -> Any: + await self._state.require(method) + spec = AGENT_REQUESTS_BY_METHOD[method] + response = await self._conn.send_request(method, _dump(request)) + if response is None and spec.empty_response: + response = {} + return spec.response.validate_python(response) + + async def _notify(self, method: str, notification: BaseModel) -> None: + await self._state.require(method) + await self._conn.send_notification(method, _dump(notification)) + + async def __aenter__(self) -> ClientSideConnection: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + await self.close() + + +def connect_to_agent( + client: ClientFactory | Client, + input_stream: Any, + output_stream: Any = None, + **connection_kwargs: Any, +) -> ClientSideConnection: + return ClientSideConnection(client, input_stream, output_stream, **connection_kwargs) diff --git a/src/acp/experimental/v2/interfaces.py b/src/acp/experimental/v2/interfaces.py new file mode 100644 index 0000000..9f0325a --- /dev/null +++ b/src/acp/experimental/v2/interfaces.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from typing import Any, Protocol + +from . import schema +from ._methods import CreateElicitationRequest, CreateElicitationResponse, SetConfigOptionRequest + +__all__ = ["Agent", "Client"] + + +class Agent(Protocol): + async def initialize(self, request: schema.InitializeRequest) -> schema.InitializeResponse: ... + + async def login(self, request: schema.LoginAuthRequest) -> schema.LoginAuthResponse | None: ... + + async def logout(self, request: schema.LogoutAuthRequest) -> schema.LogoutAuthResponse | None: ... + + async def list_providers(self, request: schema.ListProvidersRequest) -> schema.ListProvidersResponse: ... + + async def set_provider(self, request: schema.SetProviderRequest) -> schema.SetProviderResponse | None: ... + + async def disable_provider( + self, request: schema.DisableProviderRequest + ) -> schema.DisableProviderResponse | None: ... + + async def new_session(self, request: schema.NewSessionRequest) -> schema.NewSessionResponse: ... + + async def list_sessions(self, request: schema.ListSessionsRequest) -> schema.ListSessionsResponse: ... + + async def delete_session(self, request: schema.DeleteSessionRequest) -> schema.DeleteSessionResponse | None: ... + + async def fork_session(self, request: schema.ForkSessionRequest) -> schema.ForkSessionResponse: ... + + async def resume_session(self, request: schema.ResumeSessionRequest) -> schema.ResumeSessionResponse: ... + + async def close_session(self, request: schema.CloseSessionRequest) -> schema.CloseSessionResponse | None: ... + + async def set_config_option(self, request: SetConfigOptionRequest) -> schema.SetSessionConfigOptionResponse: ... + + async def prompt(self, request: schema.PromptRequest) -> schema.PromptResponse | None: ... + + async def cancel(self, notification: schema.CancelSessionNotification) -> None: ... + + async def message_mcp(self, message: schema.MessageMcpRequest | schema.MessageMcpNotification) -> Any: ... + + async def start_nes(self, request: schema.StartNesRequest) -> schema.StartNesResponse: ... + + async def suggest_nes(self, request: schema.SuggestNesRequest) -> schema.SuggestNesResponse: ... + + async def accept_nes(self, notification: schema.AcceptNesNotification) -> None: ... + + async def reject_nes(self, notification: schema.RejectNesNotification) -> None: ... + + async def close_nes(self, request: schema.CloseNesRequest) -> schema.CloseNesResponse | None: ... + + async def did_open(self, notification: schema.DidOpenDocumentNotification) -> None: ... + + async def did_change(self, notification: schema.DidChangeDocumentNotification) -> None: ... + + async def did_close(self, notification: schema.DidCloseDocumentNotification) -> None: ... + + async def did_save(self, notification: schema.DidSaveDocumentNotification) -> None: ... + + async def did_focus(self, notification: schema.DidFocusDocumentNotification) -> None: ... + + async def ext_method(self, method: str, params: Any) -> Any: ... + + async def ext_notification(self, method: str, params: Any) -> None: ... + + def on_connect(self, connection: Client) -> None: ... + + +class Client(Protocol): + async def request_permission( + self, + request: schema.RequestPermissionRequest, + ) -> schema.RequestPermissionResponse: ... + + async def session_update(self, notification: schema.UpdateSessionNotification) -> None: ... + + async def connect_mcp(self, request: schema.ConnectMcpRequest) -> schema.ConnectMcpResponse: ... + + async def message_mcp(self, message: schema.MessageMcpRequest | schema.MessageMcpNotification) -> Any: ... + + async def disconnect_mcp( + self, + request: schema.DisconnectMcpRequest, + ) -> schema.DisconnectMcpResponse | None: ... + + async def create_elicitation(self, request: CreateElicitationRequest) -> CreateElicitationResponse: ... + + async def complete_elicitation(self, notification: schema.CompleteElicitationNotification) -> None: ... + + async def ext_method(self, method: str, params: Any) -> Any: ... + + async def ext_notification(self, method: str, params: Any) -> None: ... + + def on_connect(self, connection: Agent) -> None: ... diff --git a/src/acp/experimental/v2/session.py b/src/acp/experimental/v2/session.py new file mode 100644 index 0000000..1e14adf --- /dev/null +++ b/src/acp/experimental/v2/session.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +from . import schema + +if TYPE_CHECKING: + from .client import ClientSideConnection + from .interfaces import Client + +__all__ = ["ActiveSession", "SessionMessage", "SessionStop", "SessionUpdate"] + + +@dataclass(frozen=True, slots=True) +class SessionUpdate: + notification: schema.UpdateSessionNotification + + @property + def update(self) -> Any: + return self.notification.update + + +@dataclass(frozen=True, slots=True) +class SessionStop: + notification: schema.UpdateSessionNotification + + @property + def update(self) -> schema.IdleSessionStateUpdate: + return cast(schema.IdleSessionStateUpdate, self.notification.update) + + @property + def stop_reason(self) -> Any: + return self.update.stop_reason + + +SessionMessage = SessionUpdate | SessionStop +UpdateHandler = Callable[[schema.UpdateSessionNotification], None] + + +class SessionUpdateBroker: + def __init__(self, client: Client) -> None: + self._client = client + self._handlers: dict[str, set[UpdateHandler]] = defaultdict(set) + self._pending: dict[str, list[schema.UpdateSessionNotification]] = defaultdict(list) + self._captures = 0 + + def begin_capture(self) -> None: + self._captures += 1 + + def end_capture(self) -> None: + self._captures -= 1 + if self._captures == 0: + self._pending.clear() + + def register(self, session_id: str, handler: UpdateHandler) -> Callable[[], None]: + self._handlers[session_id].add(handler) + for notification in self._pending.pop(session_id, []): + handler(notification) + + def unregister() -> None: + handlers = self._handlers.get(session_id) + if handlers is None: + return + handlers.discard(handler) + if not handlers: + self._handlers.pop(session_id, None) + + return unregister + + async def session_update(self, notification: schema.UpdateSessionNotification) -> None: + handlers = tuple(self._handlers.get(notification.session_id, ())) + if handlers: + for handler in handlers: + handler(notification) + elif self._captures: + self._pending[notification.session_id].append(notification) + + user_handler = getattr(self._client, "session_update", None) + if user_handler is not None: + await user_handler(notification) + + def __getattr__(self, name: str) -> Any: + return getattr(self._client, name) + + +class ActiveSession: + """Route one session's updates and track the split v2 prompt lifecycle.""" + + def __init__( + self, + connection: ClientSideConnection, + response: schema.NewSessionResponse, + broker: SessionUpdateBroker, + ) -> None: + self._connection = connection + self._response = response + self._updates: asyncio.Queue[schema.UpdateSessionNotification] = asyncio.Queue() + self._unregister = broker.register(response.session_id, self._updates.put_nowait) + self._prompt_active = False + self._observed_running = False + self._disposed = False + + @property + def session_id(self) -> str: + return self._response.session_id + + @property + def new_session_response(self) -> schema.NewSessionResponse: + return self._response + + async def prompt(self, request: schema.PromptRequest) -> schema.PromptResponse: + if self._disposed: + raise RuntimeError("ActiveSession has been disposed") + if request.session_id != self.session_id: + raise ValueError(f"Prompt belongs to session {request.session_id!r}, not {self.session_id!r}") + if self._prompt_active: + raise RuntimeError("Wait for the current prompt to become idle before sending another") + + self._prompt_active = True + self._observed_running = False + try: + return await self._connection.prompt(request) + except BaseException: + self._prompt_active = False + raise + + async def next_update(self) -> SessionMessage: + if self._disposed: + raise RuntimeError("ActiveSession has been disposed") + notification = await self._updates.get() + update = notification.update + if self._prompt_active and isinstance(update, schema.RunningSessionStateUpdate): + self._observed_running = True + elif self._prompt_active and self._observed_running and isinstance(update, schema.IdleSessionStateUpdate): + self._prompt_active = False + self._observed_running = False + return SessionStop(notification) + return SessionUpdate(notification) + + async def wait_for_idle(self) -> SessionStop: + while True: + message = await self.next_update() + if isinstance(message, SessionStop): + return message + + def dispose(self) -> None: + if self._disposed: + return + self._disposed = True + self._unregister() + + async def __aenter__(self) -> ActiveSession: + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.dispose() diff --git a/tests/test_protocol_negotiation.py b/tests/test_protocol_negotiation.py new file mode 100644 index 0000000..6684743 --- /dev/null +++ b/tests/test_protocol_negotiation.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +import acp +from acp._transport import memory_transport_pair +from acp.connection import StreamDirection, StreamEvent +from acp.experimental import ( + AgentProtocolRouter, + ClientNegotiator, + NegotiatedV1, + NegotiatedV2, + V1ClientConfig, + V2ClientConfig, + v2, +) + + +class Client: + pass + + +class V1Agent: + def __init__(self) -> None: + self.initialize_calls = 0 + self.client_name: str | None = None + + async def initialize( + self, + protocol_version: int, + client_capabilities: acp.schema.ClientCapabilities | None = None, + client_info: acp.schema.Implementation | None = None, + **kwargs: Any, + ) -> acp.InitializeResponse: + self.initialize_calls += 1 + self.client_name = client_info.name if client_info is not None else None + return acp.InitializeResponse(protocol_version=protocol_version) + + async def new_session( + self, + cwd: str, + additional_directories: list[str] | None = None, + mcp_servers: list[Any] | None = None, + **kwargs: Any, + ) -> acp.NewSessionResponse: + return acp.NewSessionResponse(session_id=f"v1:{cwd}") + + +class V2Agent: + def __init__(self) -> None: + self.initialize_calls = 0 + + async def initialize(self, request: v2.schema.InitializeRequest) -> v2.schema.InitializeResponse: + self.initialize_calls += 1 + return v2.schema.InitializeResponse( + protocol_version=v2.PROTOCOL_VERSION, + info=v2.schema.Implementation(name="v2-agent", version="1.0.0"), + ) + + +def v1_config() -> V1ClientConfig: + return V1ClientConfig( + client=Client(), + initialize=acp.InitializeRequest( + protocol_version=acp.PROTOCOL_VERSION, + client_info=acp.schema.Implementation(name="v1-client", version="1.0.0"), + ), + ) + + +def v2_config() -> V2ClientConfig: + return V2ClientConfig( + client=Client(), + initialize=v2.schema.InitializeRequest( + protocol_version=v2.PROTOCOL_VERSION, + info=v2.schema.Implementation(name="v2-client", version="2.0.0"), + ), + ) + + +@pytest.mark.asyncio +async def test_negotiation_selects_v2_with_one_initialize() -> None: + client_transport, agent_transport = memory_transport_pair() + v1_agent = V1Agent() + v2_agent = V2Agent() + wire: list[StreamEvent] = [] + router = AgentProtocolRouter(v1=v1_agent, v2=v2_agent) + agent_connection = router.connect(agent_transport) + negotiator = ClientNegotiator( + client_transport, + v1=v1_config(), + v2=v2_config(), + observers=[wire.append], + ) + + try: + negotiated = await negotiator.negotiate() + repeated = await negotiator.negotiate() + + assert isinstance(negotiated, NegotiatedV2) + assert repeated is negotiated + assert negotiated.initialize.info.name == "v2-agent" + assert v1_agent.initialize_calls == 0 + assert v2_agent.initialize_calls == 1 + assert _initialize_count(wire) == 1 + finally: + await negotiator.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_negotiation_downgrades_v2_initialize_without_repeating_it() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = V1Agent() + wire: list[StreamEvent] = [] + router = AgentProtocolRouter(v1=agent) + agent_connection = router.connect(agent_transport) + negotiator = ClientNegotiator( + client_transport, + v1=v1_config(), + v2=v2_config(), + observers=[wire.append], + ) + + try: + negotiated = await negotiator.negotiate() + + assert isinstance(negotiated, NegotiatedV1) + assert negotiated.initialize.protocol_version == acp.PROTOCOL_VERSION + assert agent.initialize_calls == 1 + assert agent.client_name == "v2-client" + assert _initialize_count(wire) == 1 + + session = await negotiated.connection.new_session(cwd="/workspace") + assert session.session_id == "v1:/workspace" + finally: + await negotiator.close() + await agent_connection.close() + + +def _initialize_count(events: list[StreamEvent]) -> int: + return sum( + event.direction == StreamDirection.OUTGOING and event.message.get("method") == "initialize" for event in events + ) diff --git a/tests/test_v2_runtime.py b/tests/test_v2_runtime.py new file mode 100644 index 0000000..eb35065 --- /dev/null +++ b/tests/test_v2_runtime.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from acp._transport import memory_transport_pair +from acp.exceptions import RequestError +from acp.experimental import v2 + + +class Client: + pass + + +class Agent: + def __init__(self, *, response_version: int = v2.PROTOCOL_VERSION) -> None: + self.response_version = response_version + self.initialize_calls = 0 + + async def initialize(self, request: v2.schema.InitializeRequest) -> v2.schema.InitializeResponse: + self.initialize_calls += 1 + return v2.schema.InitializeResponse( + protocol_version=self.response_version, + info=v2.schema.Implementation(name="test-agent", version="1.0.0"), + capabilities=v2.schema.AgentCapabilities(session=v2.schema.SessionCapabilities()), + ) + + async def new_session(self, request: v2.schema.NewSessionRequest) -> v2.schema.NewSessionResponse: + return v2.schema.NewSessionResponse(session_id=f"session:{request.cwd}") + + +class SessionAgent(Agent): + def on_connect(self, connection: v2.AgentSideConnection) -> None: + self.connection = connection + + async def new_session(self, request: v2.schema.NewSessionRequest) -> v2.schema.NewSessionResponse: + response = await super().new_session(request) + await self.connection.session_update( + v2.schema.UpdateSessionNotification( + session_id=response.session_id, + update=v2.schema.IdleSessionStateUpdate(), + ) + ) + return response + + async def prompt(self, request: v2.schema.PromptRequest) -> v2.schema.PromptResponse: + await self.connection.session_update( + v2.schema.UpdateSessionNotification( + session_id=request.session_id, + update=v2.schema.RunningSessionStateUpdate(), + ) + ) + await self.connection.session_update( + v2.schema.UpdateSessionNotification( + session_id=request.session_id, + update=v2.schema.IdleSessionStateUpdate(stop_reason="end_turn"), + ) + ) + return v2.schema.PromptResponse() + + +def initialize_request(protocol_version: int = v2.PROTOCOL_VERSION) -> v2.schema.InitializeRequest: + return v2.schema.InitializeRequest( + protocol_version=protocol_version, + info=v2.schema.Implementation(name="test-client", version="1.0.0"), + ) + + +@pytest.mark.asyncio +async def test_v2_runtime_initializes_and_routes_generated_models() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = Agent() + agent_connection = v2.AgentSideConnection(agent, agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + initialized = await client_connection.initialize(initialize_request()) + session = await client_connection.new_session(v2.schema.NewSessionRequest(cwd="/workspace")) + + assert initialized.protocol_version == v2.PROTOCOL_VERSION + assert session.session_id == "session:/workspace" + assert agent.initialize_calls == 1 + assert (await agent_connection.wait_until_initialized()).request.info.name == "test-client" + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_v2_runtime_rejects_calls_before_initialize() -> None: + client_transport, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(Agent(), agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + with pytest.raises(RequestError, match="Invalid request"): + await client_connection.new_session(v2.schema.NewSessionRequest(cwd="/workspace")) + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_v2_runtime_rejects_a_different_protocol_version() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = Agent() + agent_connection = v2.AgentSideConnection(agent, agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + with pytest.raises(RequestError) as error: + await client_connection.initialize(initialize_request(protocol_version=1)) + + assert isinstance(error.value, RequestError) + assert error.value.code == -32602 + assert agent.initialize_calls == 0 + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_v2_runtime_rejects_a_mismatched_initialize_response() -> None: + client_transport, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(Agent(response_version=1), agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + with pytest.raises(RequestError) as error: + await client_connection.initialize(initialize_request()) + + assert isinstance(error.value, RequestError) + assert error.value.code == -32600 + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_active_session_completes_only_after_running_then_idle() -> None: + client_transport, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(SessionAgent(), agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + await client_connection.initialize(initialize_request()) + session = await client_connection.open_session(v2.schema.NewSessionRequest(cwd="/workspace")) + await session.prompt( + v2.schema.PromptRequest( + session_id=session.session_id, + prompt=[v2.schema.TextContentBlock(text="hello")], + ) + ) + + ready = await session.next_update() + running = await session.next_update() + stopped = await session.next_update() + + assert isinstance(ready, v2.SessionUpdate) + assert isinstance(ready.update, v2.schema.IdleSessionStateUpdate) + assert isinstance(running, v2.SessionUpdate) + assert isinstance(running.update, v2.schema.RunningSessionStateUpdate) + assert isinstance(stopped, v2.SessionStop) + assert stopped.stop_reason == "end_turn" + session.dispose() + finally: + await client_connection.close() + await agent_connection.close() + + +def test_v2_public_entry_point_is_explicit() -> None: + exported: dict[str, Any] = {name: getattr(v2, name) for name in v2.__all__} + + assert exported["PROTOCOL_VERSION"] == 2 + assert exported["schema"] is v2.schema + assert "InitializeRequest" not in v2.__all__