diff --git a/docs/experimental-v2.md b/docs/experimental-v2.md new file mode 100644 index 0000000..4ae26b2 --- /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. Install update handlers on the +client before opening a session because updates are independent connection +traffic: + +```python +from acp.experimental import v2 + +class MyClient: + async def session_update( + self, + notification: v2.schema.UpdateSessionNotification, + ) -> None: + handle_update(notification) + + +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.new_session( + v2.schema.NewSessionRequest(cwd="/workspace") +) +await connection.prompt( + v2.schema.PromptRequest( + session_id=session.session_id, + prompt=[v2.schema.TextContentBlock(text="Hello")], + ) +) +``` + +`session/prompt` returns when the agent accepts the prompt. It does not define a +boundary for session updates: they may arrive before, during, or after that +request, and they do not carry a prompt identifier. Applications decide how to +buffer or present them. + +Agents that serve both versions use `AgentProtocolRouter`: + +```python +from acp.experimental import AgentProtocolRouter + +router = AgentProtocolRouter( + v1=lambda connection: V1Agent(connection), + v2=lambda connection: V2Agent(connection), +) +await router.run() +``` + +The selected factory is called once per connection. Return a fresh agent from +each call to avoid sharing connection state. + +Extension method names are explicit and must include the protocol-required `_` +prefix: + +```python +result = await connection.send_extension_request("_vendor/method", {"value": 1}) +await connection.send_extension_notification("_vendor/event", {"value": 1}) +``` + +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. + +Client-side fallback is application controlled and may require opening a new +transport. Protocol-level request cancellation is not yet exposed by the +experimental runtime; `session/cancel` remains available for cancelling active +session work. 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..e4e1d31 100644 --- a/src/acp/experimental/__init__.py +++ b/src/acp/experimental/__init__.py @@ -1 +1,13 @@ """Experimental ACP APIs.""" + +from . import v2 +from .negotiation import ( + AgentProtocolConnection, + AgentProtocolRouter, +) + +__all__ = [ + "AgentProtocolConnection", + "AgentProtocolRouter", + "v2", +] diff --git a/src/acp/experimental/negotiation.py b/src/acp/experimental/negotiation.py new file mode 100644 index 0000000..c9fe5a0 --- /dev/null +++ b/src/acp/experimental/negotiation.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any + +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.connection import Connection, MethodHandler +from acp.exceptions import RequestError +from acp.interfaces import Agent as V1Agent + +from . import v2 +from .v2._connection import open_connection +from .v2.agent import AgentSideConnection as V2AgentSideConnection +from .v2.meta import AGENT_METHODS as V2_AGENT_METHODS + +__all__ = [ + "AgentProtocolConnection", + "AgentProtocolRouter", +] + +V1AgentFactory = Callable[[V1AgentSideConnection], V1Agent] +V2AgentFactory = Callable[[V2AgentSideConnection], object] + + +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 _AgentNegotiationHandler: + def __init__( + self, + v1_agent: V1AgentFactory | None, + v2_agent: V2AgentFactory | 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) + connection = self._connection + if connection is None: + raise RuntimeError("Protocol router is not connected") + + if self._v2_agent is not None and requested >= v2.PROTOCOL_VERSION: + selected = v2.PROTOCOL_VERSION + endpoint, handler = V2AgentSideConnection._attach(self._v2_agent, connection) + elif self._v1_agent is not None and requested >= v1_meta.PROTOCOL_VERSION: + selected = v1_meta.PROTOCOL_VERSION + endpoint, handler = V1AgentSideConnection._attach(self._v1_agent, connection) + else: + 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}" + }) + 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 + + +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 | None = None, + v2: V2AgentFactory | 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, + **connection_kwargs: Any, + ) -> AgentProtocolConnection: + return self._connect(input_stream, output_stream, **connection_kwargs) + + 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()) diff --git a/src/acp/experimental/v2/__init__.py b/src/acp/experimental/v2/__init__.py index c5ee37e..d0af723 100644 --- a/src/acp/experimental/v2/__init__.py +++ b/src/acp/experimental/v2/__init__.py @@ -1,5 +1,15 @@ -"""Experimental ACP protocol v2 bindings.""" +"""Experimental ACP protocol v2 API.""" -from .meta import AGENT_METHODS, CLIENT_METHODS, PROTOCOL_METHODS, PROTOCOL_VERSION +from . import schema +from .agent import AgentSideConnection, run_agent +from .client import ClientSideConnection, connect_to_agent +from .meta import PROTOCOL_VERSION -__all__ = ["AGENT_METHODS", "CLIENT_METHODS", "PROTOCOL_METHODS", "PROTOCOL_VERSION"] +__all__ = [ + "PROTOCOL_VERSION", + "AgentSideConnection", + "ClientSideConnection", + "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..5ae6efb --- /dev/null +++ b/src/acp/experimental/v2/_initialization.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import asyncio +from typing import Literal + +from acp.exceptions import RequestError + +from . import schema +from .meta import PROTOCOL_VERSION + +InitializationPhase = Literal["uninitialized", "initializing", "initialized", "failed"] + + +class InitializationState: + def __init__(self) -> None: + self._phase: InitializationPhase = "uninitialized" + 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._phase = "initializing" + + def complete(self, response: schema.InitializeResponse) -> None: + if self._phase != "initializing": + 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, + }) + self._phase = "initialized" + self._ready.set() + + def fail(self, error: BaseException) -> None: + if self._phase == "initialized": + return + self._phase = "failed" + self._failure = error + self._ready.set() + + async def initialized(self) -> None: + if self._phase in {"uninitialized", "initializing"}: + await self._ready.wait() + if self._phase == "initialized": + return + 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..43867cd --- /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"], "mcp_message", 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_session", schema.CancelSessionNotification), + notification(AGENT_METHODS["mcp_message"], "notify_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"], "mcp_message", 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"], "notify_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..57c0860 --- /dev/null +++ b/src/acp/experimental/v2/_router.py @@ -0,0 +1,65 @@ +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: + return + 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: + return None + 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 = "handle_extension_notification" if is_notification else "handle_extension_request" + handler = getattr(self._target, handler_name, None) + if handler is None: + if is_notification: + return 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..dde34b9 --- /dev/null +++ b/src/acp/experimental/v2/agent.py @@ -0,0 +1,184 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any + +from pydantic import BaseModel + +from acp.connection import Connection + +from . import schema +from ._connection import open_connection +from ._initialization import InitializationState +from ._methods import ( + AGENT_NOTIFICATIONS, + AGENT_REQUESTS, + CLIENT_REQUESTS_BY_METHOD, + CreateElicitationRequest, + CreateElicitationResponse, +) +from ._router import MethodRouter +from .meta import CLIENT_METHODS + +__all__ = ["AgentSideConnection", "run_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: object, 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: + initialize = self._router.request_spec("initialize") + if not is_notification and initialize is not None and method == initialize.method: + request: schema.InitializeRequest = initialize.request.validate_python(params) + self._state.begin(request) + try: + response: 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, + agent: object, + input_stream: Any, + output_stream: Any = None, + *, + _listening: bool = True, + **connection_kwargs: Any, + ) -> None: + self._state = InitializationState() + router = _AgentRouter(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, + agent_factory: Callable[[AgentSideConnection], object], + connection: Connection, + ) -> tuple[AgentSideConnection, _AgentRouter]: + self = cls.__new__(cls) + self._state = InitializationState() + self._conn = connection + agent = agent_factory(self) + router = _AgentRouter(agent, self._state) + return self, router + + 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 mcp_message(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 send_extension_request(self, method: str, params: Any = None) -> Any: + await self._state.require(method) + return await self._conn.send_request(_extension_method(method), params) + + async def send_extension_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: + if not method.startswith("_"): + raise ValueError("ACP extension methods must start with '_'") + return method + + +async def run_agent( + agent: object, + 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..c744b93 --- /dev/null +++ b/src/acp/experimental/v2/client.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel + +from . import schema +from ._connection import open_connection +from ._initialization import InitializationState +from ._methods import ( + AGENT_REQUESTS_BY_METHOD, + CLIENT_NOTIFICATIONS, + CLIENT_REQUESTS, + SetConfigOptionRequest, +) +from ._router import MethodRouter +from .agent import _dump, _extension_method +from .meta import AGENT_METHODS + +__all__ = ["ClientSideConnection", "connect_to_agent"] + + +class _ClientRouter: + def __init__(self, client: object, 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: + 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, + client: object, + input_stream: Any, + output_stream: Any = None, + **connection_kwargs: Any, + ) -> None: + self._state = InitializationState() + router = _ClientRouter(client, self._state) + self._conn = open_connection(router, input_stream, output_stream, **connection_kwargs) + if on_connect := getattr(client, "on_connect", None): + on_connect(self) + + 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 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_session(self, notification: schema.CancelSessionNotification) -> None: + await self._notify(AGENT_METHODS["session_cancel"], notification) + + async def mcp_message(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 send_extension_request(self, method: str, params: Any = None) -> Any: + await self._state.require(method) + return await self._conn.send_request(_extension_method(method), params) + + async def send_extension_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: object, + input_stream: Any, + output_stream: Any = None, + **connection_kwargs: Any, +) -> ClientSideConnection: + return ClientSideConnection(client, input_stream, output_stream, **connection_kwargs) diff --git a/tests/test_protocol_negotiation.py b/tests/test_protocol_negotiation.py new file mode 100644 index 0000000..8961802 --- /dev/null +++ b/tests/test_protocol_negotiation.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +import acp +from acp._transport import memory_transport_pair +from acp.connection import Connection, StreamDirection, StreamEvent +from acp.experimental import AgentProtocolRouter, v2 +from acp.experimental.v2.meta import AGENT_METHODS + + +class Client: + pass + + +class V1Client(acp.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) + + +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"), + ) + + +class UpdateClient: + def __init__(self) -> None: + self.updates: asyncio.Queue[v2.schema.UpdateSessionNotification] = asyncio.Queue() + + async def session_update(self, notification: v2.schema.UpdateSessionNotification) -> None: + await self.updates.put(notification) + + +class RoutedV2Agent(V2Agent): + def __init__(self, connection: v2.AgentSideConnection) -> None: + super().__init__() + self.connection = connection + + 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.IdleSessionStateUpdate(), + ) + ) + return v2.schema.PromptResponse() + + +def v2_initialize() -> v2.schema.InitializeRequest: + return 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_agent_protocol_router_selects_v2() -> None: + client_transport, agent_transport = memory_transport_pair() + v1_agent = V1Agent() + v2_agent = V2Agent() + wire: list[StreamEvent] = [] + agent_connection = AgentProtocolRouter(v1=lambda _: v1_agent, v2=lambda _: v2_agent).connect(agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport, observers=[wire.append]) + + try: + initialized = await client_connection.initialize(v2_initialize()) + + assert initialized.info.name == "v2-agent" + assert v1_agent.initialize_calls == 0 + assert v2_agent.initialize_calls == 1 + assert _initialize_count(wire) == 1 + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_agent_protocol_router_selects_v1() -> None: + client_transport, agent_transport = memory_transport_pair() + v1_agent = V1Agent() + v2_agent = V2Agent() + wire: list[StreamEvent] = [] + agent_connection = AgentProtocolRouter(v1=lambda _: v1_agent, v2=lambda _: v2_agent).connect(agent_transport) + client_connection = acp.connect_to_agent(V1Client(), client_transport, observers=[wire.append]) + + try: + initialized = await client_connection.initialize( + protocol_version=acp.PROTOCOL_VERSION, + client_info=acp.schema.Implementation(name="v1-client", version="1.0.0"), + ) + + assert initialized.protocol_version == acp.PROTOCOL_VERSION + assert v1_agent.initialize_calls == 1 + assert v1_agent.client_name == "v1-client" + assert v2_agent.initialize_calls == 0 + assert _initialize_count(wire) == 1 + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_agent_protocol_router_normalizes_v2_initialize_for_v1() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = V1Agent() + wire: list[StreamEvent] = [] + agent_connection = AgentProtocolRouter(v1=lambda _: agent).connect(agent_transport) + + async def ignore_incoming(method: str, params: Any, is_notification: bool) -> None: + pass + + client_connection = Connection(ignore_incoming, client_transport, observers=[wire.append]) + + try: + response = await client_connection.send_request( + AGENT_METHODS["initialize"], + v2_initialize().model_dump(mode="json", by_alias=True, exclude_none=True), + ) + + assert response["protocolVersion"] == acp.PROTOCOL_VERSION + assert agent.initialize_calls == 1 + assert agent.client_name == "v2-client" + assert _initialize_count(wire) == 1 + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_agent_protocol_router_isolates_connections() -> None: + agents: list[RoutedV2Agent] = [] + + def create_agent(connection: v2.AgentSideConnection) -> RoutedV2Agent: + agent = RoutedV2Agent(connection) + agents.append(agent) + return agent + + router = AgentProtocolRouter(v2=create_agent) + client_1_transport, agent_1_transport = memory_transport_pair() + client_2_transport, agent_2_transport = memory_transport_pair() + client_1 = UpdateClient() + client_2 = UpdateClient() + agent_connection_1 = router.connect(agent_1_transport) + agent_connection_2 = router.connect(agent_2_transport) + client_connection_1 = v2.ClientSideConnection(client_1, client_1_transport) + client_connection_2 = v2.ClientSideConnection(client_2, client_2_transport) + + try: + await client_connection_1.initialize(v2_initialize()) + await client_connection_2.initialize(v2_initialize()) + await client_connection_1.prompt( + v2.schema.PromptRequest( + session_id="session-1", + prompt=[v2.schema.TextContentBlock(text="hello")], + ) + ) + + update = await asyncio.wait_for(client_1.updates.get(), timeout=1) + assert update.session_id == "session-1" + assert client_2.updates.empty() + assert len(agents) == 2 + assert agents[0] is not agents[1] + finally: + await client_connection_1.close() + await client_connection_2.close() + await agent_connection_1.close() + await agent_connection_2.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..ee5f350 --- /dev/null +++ b/tests/test_v2_runtime.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import asyncio +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 SessionClient: + def __init__(self) -> None: + self.updates: asyncio.Queue[v2.schema.UpdateSessionNotification] = asyncio.Queue() + + async def session_update(self, notification: v2.schema.UpdateSessionNotification) -> None: + await self.updates.put(notification) + + +class Agent: + def __init__(self, *, response_version: int = v2.PROTOCOL_VERSION) -> None: + self.response_version = response_version + self.initialize_calls = 0 + self.client_name: str | None = None + + async def initialize(self, request: v2.schema.InitializeRequest) -> v2.schema.InitializeResponse: + self.initialize_calls += 1 + self.client_name = request.info.name + 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 CallableAgent(Agent): + def __call__(self, *args: Any, **kwargs: Any) -> None: + raise AssertionError("an agent object must not be treated as a factory") + + +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() + + +class ExtensionClient: + async def handle_extension_request(self, method: str, params: Any) -> Any: + return {"method": method, "params": params} + + +class ExtensionAgent: + def __init__(self) -> None: + self.notifications: asyncio.Queue[tuple[str, Any]] = asyncio.Queue() + + async def initialize(self, request: v2.schema.InitializeRequest) -> v2.schema.InitializeResponse: + return v2.schema.InitializeResponse( + protocol_version=v2.PROTOCOL_VERSION, + info=v2.schema.Implementation(name="extension-agent", version="1.0.0"), + ) + + async def handle_extension_request(self, method: str, params: Any) -> Any: + return {"method": method, "params": params} + + async def cancel_session(self, notification: v2.schema.CancelSessionNotification) -> None: + await self.notifications.put(("cancel", notification)) + + async def notify_mcp(self, notification: v2.schema.MessageMcpNotification) -> None: + await self.notifications.put(("mcp", notification)) + + +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 agent.client_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_callable_agent_is_not_treated_as_a_factory() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = CallableAgent() + agent_connection = v2.AgentSideConnection(agent, agent_transport) + client_connection = v2.ClientSideConnection(Client(), client_transport) + + try: + initialized = await client_connection.initialize(initialize_request()) + + assert initialized.protocol_version == v2.PROTOCOL_VERSION + assert agent.initialize_calls == 1 + 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_session_updates_are_delivered_independently_from_prompt() -> None: + client_transport, agent_transport = memory_transport_pair() + client = SessionClient() + 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.new_session(v2.schema.NewSessionRequest(cwd="/workspace")) + await client_connection.prompt( + v2.schema.PromptRequest( + session_id=session.session_id, + prompt=[v2.schema.TextContentBlock(text="hello")], + ) + ) + + ready = await asyncio.wait_for(client.updates.get(), timeout=1) + running = await asyncio.wait_for(client.updates.get(), timeout=1) + stopped = await asyncio.wait_for(client.updates.get(), timeout=1) + + assert isinstance(ready.update, v2.schema.IdleSessionStateUpdate) + assert isinstance(running.update, v2.schema.RunningSessionStateUpdate) + assert isinstance(stopped.update, v2.schema.IdleSessionStateUpdate) + assert stopped.update.stop_reason == "end_turn" + 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 set(exported) == { + "AgentSideConnection", + "ClientSideConnection", + "PROTOCOL_VERSION", + "connect_to_agent", + "run_agent", + "schema", + } + assert "InitializeRequest" not in v2.__all__ + + +@pytest.mark.asyncio +async def test_extension_and_notification_names_are_explicit() -> None: + client_transport, agent_transport = memory_transport_pair() + agent = ExtensionAgent() + agent_connection = v2.AgentSideConnection(agent, agent_transport) + client_connection = v2.ClientSideConnection(ExtensionClient(), client_transport) + + try: + await client_connection.initialize(initialize_request()) + + assert await client_connection.send_extension_request("_vendor/do", {"value": 1}) == { + "method": "_vendor/do", + "params": {"value": 1}, + } + assert await agent_connection.send_extension_request("_vendor/read", {"value": 2}) == { + "method": "_vendor/read", + "params": {"value": 2}, + } + with pytest.raises(ValueError, match="must start with '_'"): + await client_connection.send_extension_request("vendor/do") + + await client_connection.cancel_session(v2.schema.CancelSessionNotification(session_id="session-1")) + await client_connection.notify_mcp( + v2.schema.MessageMcpNotification(connection_id="mcp-1", method="notifications/progress") + ) + + cancel_kind, cancel = await asyncio.wait_for(agent.notifications.get(), timeout=1) + mcp_kind, mcp = await asyncio.wait_for(agent.notifications.get(), timeout=1) + assert (cancel_kind, cancel.session_id) == ("cancel", "session-1") + assert (mcp_kind, mcp.connection_id) == ("mcp", "mcp-1") + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_unhandled_notifications_are_ignored(caplog: pytest.LogCaptureFixture) -> None: + client_transport, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(Agent(), agent_transport) + client_connection = v2.ClientSideConnection(object(), client_transport) + + try: + await client_connection.initialize(initialize_request()) + await agent_connection.session_update( + v2.schema.UpdateSessionNotification( + session_id="session-1", + update=v2.schema.IdleSessionStateUpdate(), + ) + ) + await client_connection.send_extension_notification("_vendor/event") + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert not [record for record in caplog.records if record.levelno >= 40] + finally: + await client_connection.close() + await agent_connection.close() + + +@pytest.mark.asyncio +async def test_missing_request_handler_returns_method_not_found() -> None: + client_transport, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(ExtensionAgent(), agent_transport) + client_connection = v2.ClientSideConnection(object(), client_transport) + + try: + await client_connection.initialize(initialize_request()) + with pytest.raises(RequestError) as error: + await client_connection.new_session(v2.schema.NewSessionRequest(cwd="/workspace")) + + assert isinstance(error.value, RequestError) + assert error.value.code == -32601 + finally: + await client_connection.close() + await agent_connection.close() diff --git a/tests/test_v2_schema.py b/tests/test_v2_schema.py index a63709c..496e273 100644 --- a/tests/test_v2_schema.py +++ b/tests/test_v2_schema.py @@ -1,7 +1,8 @@ import pytest from pydantic import ValidationError -from acp.experimental.v2 import PROTOCOL_METHODS, PROTOCOL_VERSION +from acp.experimental.v2 import PROTOCOL_VERSION +from acp.experimental.v2.meta import PROTOCOL_METHODS from acp.experimental.v2.schema import ( AgentMessageChunk, OtherSessionUpdate,