diff --git a/README.md b/README.md index 400b0da..7929021 100644 --- a/README.md +++ b/README.md @@ -74,8 +74,31 @@ client = await UtcpClient.create(config={ # Call a tool result = await client.call_tool("my_api.get_data", {"id": "123"}) + +# Release the connections this client opened +await client.close() +``` + +### Protocol Lifetime: Shared vs. Per-Client Instances + +A protocol plugin registers itself in one of two registries, and the choice decides who its state belongs to. + +An instance registered with `register_communication_protocol` is **shared by every `UtcpClient` in the process**, and so is any state it keeps. That is right for state that is *meant* to be shared — a credential cache, a pooled HTTP session, a registry a decorator writes into — and wrong for state that belongs to one client. A shared instance lives as long as the process; no client closes it. + +A protocol whose state **must not be shared between clients** registers a **factory** instead: live sessions or connections keyed per manual, child processes — anything one client's use or `close()` would take away from another. Each `UtcpClient` calls the factory once — at creation, or on first use if it was registered later — so each client gets its own instance, its own connections, and its own teardown on `close()`: + +```python +from utcp.plugins.discovery import register_communication_protocol_factory + +register_communication_protocol_factory("custom_type", CustomCommunicationProtocol) ``` +This is what makes "a client per tenant / per user / per pooled connection" actually isolate them, rather than every client reaching into one shared instance. A type registered as a factory wins over the same type registered as an instance, so a plugin migrates by moving its registration from one call to the other and callers change nothing. + +`utcp-mcp` and `utcp-websocket` register this way — MCP sessions (and, for stdio, child processes) and per-manual WebSocket connections belong to the client that opened them. `utcp-http` (HTTP, SSE, streamable HTTP) and `utcp-gql` stay shared instances: what they keep is an OAuth token cache and a pooled HTTP session, both meant to be reused across clients. + +`client.close()` closes the instances the client created — every one of them, even if one fails, after which the failures are raised together as `UtcpProtocolCloseError`. If `UtcpClient.create` fails after creating them, they are closed the same way before the error is re-raised. + ## Protocol Plugins UTCP supports multiple communication protocols through dedicated plugins: diff --git a/core/README.md b/core/README.md index 400b0da..7929021 100644 --- a/core/README.md +++ b/core/README.md @@ -74,8 +74,31 @@ client = await UtcpClient.create(config={ # Call a tool result = await client.call_tool("my_api.get_data", {"id": "123"}) + +# Release the connections this client opened +await client.close() +``` + +### Protocol Lifetime: Shared vs. Per-Client Instances + +A protocol plugin registers itself in one of two registries, and the choice decides who its state belongs to. + +An instance registered with `register_communication_protocol` is **shared by every `UtcpClient` in the process**, and so is any state it keeps. That is right for state that is *meant* to be shared — a credential cache, a pooled HTTP session, a registry a decorator writes into — and wrong for state that belongs to one client. A shared instance lives as long as the process; no client closes it. + +A protocol whose state **must not be shared between clients** registers a **factory** instead: live sessions or connections keyed per manual, child processes — anything one client's use or `close()` would take away from another. Each `UtcpClient` calls the factory once — at creation, or on first use if it was registered later — so each client gets its own instance, its own connections, and its own teardown on `close()`: + +```python +from utcp.plugins.discovery import register_communication_protocol_factory + +register_communication_protocol_factory("custom_type", CustomCommunicationProtocol) ``` +This is what makes "a client per tenant / per user / per pooled connection" actually isolate them, rather than every client reaching into one shared instance. A type registered as a factory wins over the same type registered as an instance, so a plugin migrates by moving its registration from one call to the other and callers change nothing. + +`utcp-mcp` and `utcp-websocket` register this way — MCP sessions (and, for stdio, child processes) and per-manual WebSocket connections belong to the client that opened them. `utcp-http` (HTTP, SSE, streamable HTTP) and `utcp-gql` stay shared instances: what they keep is an OAuth token cache and a pooled HTTP session, both meant to be reused across clients. + +`client.close()` closes the instances the client created — every one of them, even if one fails, after which the failures are raised together as `UtcpProtocolCloseError`. If `UtcpClient.create` fails after creating them, they are closed the same way before the error is re-raised. + ## Protocol Plugins UTCP supports multiple communication protocols through dedicated plugins: diff --git a/core/pyproject.toml b/core/pyproject.toml index 04c8344..ce07a77 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp" -version = "1.1.4" +version = "1.2.0" authors = [ { name = "UTCP Contributors" }, ] diff --git a/core/src/utcp/exceptions/__init__.py b/core/src/utcp/exceptions/__init__.py index a33c4f6..d993eb5 100644 --- a/core/src/utcp/exceptions/__init__.py +++ b/core/src/utcp/exceptions/__init__.py @@ -1,7 +1,9 @@ from utcp.exceptions.utcp_variable_not_found_exception import UtcpVariableNotFound from utcp.exceptions.utcp_serializer_validation_error import UtcpSerializerValidationError +from utcp.exceptions.utcp_protocol_close_error import UtcpProtocolCloseError __all__ = [ "UtcpVariableNotFound", - "UtcpSerializerValidationError" + "UtcpSerializerValidationError", + "UtcpProtocolCloseError", ] diff --git a/core/src/utcp/exceptions/utcp_protocol_close_error.py b/core/src/utcp/exceptions/utcp_protocol_close_error.py new file mode 100644 index 0000000..f79304d --- /dev/null +++ b/core/src/utcp/exceptions/utcp_protocol_close_error.py @@ -0,0 +1,16 @@ +from typing import List + + +class UtcpProtocolCloseError(Exception): + """REQUIRED + Raised when one or more of a client's own protocol instances failed to close. + + Every instance is still asked to close before this is raised, so nothing is + left half-torn-down behind it; `failures` carries what each failing close raised. + """ + + def __init__(self, failures: List[BaseException], total: int): + self.failures = failures + self.total = total + super().__init__(f"{len(failures)} of {total} owned communication protocol(s) failed to close: " + + "; ".join(f"{type(f).__name__}: {f}" for f in failures)) diff --git a/core/src/utcp/implementations/utcp_client_implementation.py b/core/src/utcp/implementations/utcp_client_implementation.py index b88bead..c227b63 100644 --- a/core/src/utcp/implementations/utcp_client_implementation.py +++ b/core/src/utcp/implementations/utcp_client_implementation.py @@ -4,7 +4,7 @@ import os import json import asyncio -from typing import Dict, Any, List, Union, Optional, AsyncGenerator, TYPE_CHECKING +from typing import Callable, Dict, Any, List, Union, Optional, AsyncGenerator, TYPE_CHECKING from utcp.data.call_template import CallTemplate from utcp.data.call_template import CallTemplateSerializer @@ -15,7 +15,7 @@ from utcp.data.utcp_client_config import UtcpClientConfig, UtcpClientConfigSerializer from utcp.implementations.default_variable_substitutor import DefaultVariableSubstitutor from utcp.implementations.tag_search import TagAndDescriptionWordMatchStrategy -from utcp.exceptions import UtcpVariableNotFound +from utcp.exceptions import UtcpVariableNotFound, UtcpProtocolCloseError from utcp.data.register_manual_response import RegisterManualResult from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp.exceptions import UtcpSerializerValidationError @@ -39,6 +39,77 @@ def __init__( ): super().__init__(config, root_dir) self.variable_substitutor = variable_substitutor + # The protocol instances this client CREATED (from the factory registry) + # and therefore owns: what its teardown closes. They are adopted by + # `create()` after construction, inside its cleanup guard — a + # constructor cannot await, so nothing that needs closing on failure is + # created here. Shared instances stay in the process registry, looked up + # live, and are never closed on this client's behalf. + self._owned_comm_protocols: List[CommunicationProtocol] = [] + self._own_comm_protocols_by_type: Dict[str, CommunicationProtocol] = {} + + def _adopt(self, protocol_type: str, factory: Callable[[], CommunicationProtocol]) -> CommunicationProtocol: + """Create this client's own instance of a protocol and record it as owned. + + Recorded the moment it exists, so a factory that raises part-way + through a batch leaves the ones before it closable. + """ + protocol = factory() + self._owned_comm_protocols.append(protocol) + self._own_comm_protocols_by_type[protocol_type] = protocol + return protocol + + def _adopt_factory_protocols(self) -> None: + """Instantiate this client's own protocols from every factory registered so far.""" + for protocol_type, factory in CommunicationProtocol.communication_protocol_factories.items(): + self._adopt(protocol_type, factory) + + def _protocol_for(self, call_template_type: str) -> CommunicationProtocol: + """Resolve the protocol this client uses for a call template type. + + In order: this client's own instance; a factory registered since the + client was created (adopted now, so it is owned and closed like the + rest); a shared instance. Both registries are consulted live, so a + type resolves the same way whenever it was registered — and a factory + wins over a shared instance of the same type, which is what lets a + plugin migrate from instance to factory without callers changing + anything. + """ + protocol = self._own_comm_protocols_by_type.get(call_template_type) + if protocol is None: + factory = CommunicationProtocol.communication_protocol_factories.get(call_template_type) + if factory is not None: + protocol = self._adopt(call_template_type, factory) + if protocol is None: + protocol = CommunicationProtocol.communication_protocols.get(call_template_type) + if protocol is None: + available = set(self._own_comm_protocols_by_type) | set(CommunicationProtocol.communication_protocol_factories) | set(CommunicationProtocol.communication_protocols) + raise ValueError(f"No registered communication protocol of type {call_template_type} found, available types: {sorted(available)}") + return protocol + + async def _close_owned_protocols(self) -> None: + """Close every protocol this client created, waiting for all of them even when one fails.""" + results = await asyncio.gather(*(protocol.close() for protocol in self._owned_comm_protocols), return_exceptions=True) + failures = [result for result in results if isinstance(result, BaseException)] + if failures: + raise UtcpProtocolCloseError(failures, len(results)) + + async def close(self) -> None: + """REQUIRED + Close the protocol instances this client created and release their resources. + + Only what this client created. A shared instance is in use by every + other client in the process — closing it here would clear their state + too (a credential cache, a decorator's registry), which is the + cross-client damage per-client instances exist to prevent. Shared + instances live as long as the process that registered them. + + Raises: + UtcpProtocolCloseError: If one or more owned instances failed to close. + Every instance is still asked to close first. + """ + await self._close_owned_protocols() + logger.info("UTCP client closed, with the protocols it owned.") @classmethod async def create( @@ -79,16 +150,32 @@ async def create( # Create the client client = cls(config, DefaultVariableSubstitutor(), root_dir) - # Substitute variables in the config - if client.config.variables: - config_without_vars = client_config_serializer.copy(client.config) - config_without_vars.variables = None - client.config.variables = client.variable_substitutor.substitute(client.config.variables, config_without_vars) + # Everything from here on can fail, and the caller never receives a + # client it could close — so whatever this client CREATED is closed + # here before the failure is re-raised. Only what it created: the shared + # instances are in use by every other client and are not this one's. + try: + client._adopt_factory_protocols() + + # Substitute variables in the config + if client.config.variables: + config_without_vars = client_config_serializer.copy(client.config) + config_without_vars.variables = None + client.config.variables = client.variable_substitutor.substitute(client.config.variables, config_without_vars) + + # Load the manuals if any + if config.manual_call_templates: + await client.register_manuals(config.manual_call_templates) + except BaseException: + try: + await client._close_owned_protocols() + except Exception: + # The initialization failure stays the error the caller sees; + # the cleanup failure is reported rather than swallowed, because + # the caller has no client through which to retry it. + logger.error("UtcpClient.create failed, and closing the protocols it had created failed too", exc_info=True) + raise - # Load the manuals if any - if config.manual_call_templates: - await client.register_manuals(config.manual_call_templates) - return client async def register_manual(self, manual_call_template: CallTemplate) -> RegisterManualResult: @@ -121,10 +208,8 @@ async def register_manual(self, manual_call_template: CallTemplate) -> RegisterM if await self.config.tool_repository.get_manual(manual_call_template.name) is not None: raise ValueError(f"Manual {manual_call_template.name} already registered, please use a different name or deregister the existing manual") manual_call_template = self._substitute_call_template_variables(manual_call_template, manual_call_template.name) - if manual_call_template.call_template_type not in CommunicationProtocol.communication_protocols: - raise ValueError(f"No registered communication protocol of type {manual_call_template.call_template_type} found, available types: {CommunicationProtocol.communication_protocols.keys()}") - - result = await CommunicationProtocol.communication_protocols[manual_call_template.call_template_type].register_manual(self, manual_call_template) + + result = await self._protocol_for(manual_call_template.call_template_type).register_manual(self, manual_call_template) if result.success: # Determine allowed protocols: use explicit list or default to manual's own protocol @@ -186,8 +271,15 @@ async def try_register_manual(manual_call_template=manual_call_template): tasks.append(try_register_manual()) - # Wait for all tasks to complete and collect results - results = await asyncio.gather(*tasks) + # Wait for EVERY registration to settle, even when one raises. The + # caller receives the failure only once no sibling registration is + # still running underneath it — create(), for one, closes the + # protocols right after, and a registration still in flight would be + # using a closed one. + results = await asyncio.gather(*tasks, return_exceptions=True) + failures = [result for result in results if isinstance(result, BaseException)] + if failures: + raise failures[0] return [p for p in results if p is not None] async def deregister_manual(self, manual_name: str) -> bool: @@ -203,7 +295,7 @@ async def deregister_manual(self, manual_name: str) -> bool: manual_call_template = await self.config.tool_repository.get_manual_call_template(manual_name) if manual_call_template is None: return False - await CommunicationProtocol.communication_protocols[manual_call_template.call_template_type].deregister_manual(self, manual_call_template) + await self._protocol_for(manual_call_template.call_template_type).deregister_manual(self, manual_call_template) return await self.config.tool_repository.remove_manual(manual_name) async def call_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: @@ -250,7 +342,7 @@ async def call_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: f"Allowed protocols: {allowed_protocols}" ) - result = await CommunicationProtocol.communication_protocols[tool_call_template.call_template_type].call_tool(self, tool_name, tool_args, tool_call_template) + result = await self._protocol_for(tool_call_template.call_template_type).call_tool(self, tool_name, tool_args, tool_call_template) for post_processor in self.config.post_processing: result = post_processor.post_process(self, tool, tool_call_template, result) @@ -300,7 +392,7 @@ async def call_tool_streaming(self, tool_name: str, tool_args: Dict[str, Any]) - f"Allowed protocols: {allowed_protocols}" ) - async for item in CommunicationProtocol.communication_protocols[tool_call_template.call_template_type].call_tool_streaming(self, tool_name, tool_args, tool_call_template): + async for item in self._protocol_for(tool_call_template.call_template_type).call_tool_streaming(self, tool_name, tool_args, tool_call_template): for post_processor in self.config.post_processing: item = post_processor.post_process(self, tool, tool_call_template, item) yield item @@ -342,9 +434,7 @@ async def get_required_variables_for_manual_and_tools(self, manual_call_template except UtcpVariableNotFound as e: return variables_for_CallTemplate return variables_for_CallTemplate - if manual_call_template.call_template_type not in CommunicationProtocol.communication_protocols: - raise ValueError(f"CallTemplate type not supported: {manual_call_template.call_template_type}") - register_manual_result: RegisterManualResult = await CommunicationProtocol.communication_protocols[manual_call_template.call_template_type].register_manual(self, manual_call_template) + register_manual_result: RegisterManualResult = await self._protocol_for(manual_call_template.call_template_type).register_manual(self, manual_call_template) for tool in register_manual_result.manual.tools: variables_for_CallTemplate.extend(self.variable_substitutor.find_required_variables(CallTemplateSerializer().to_dict(tool.tool_call_template), manual_call_template.name)) return variables_for_CallTemplate diff --git a/core/src/utcp/interfaces/communication_protocol.py b/core/src/utcp/interfaces/communication_protocol.py index b12e6ea..56c055b 100644 --- a/core/src/utcp/interfaces/communication_protocol.py +++ b/core/src/utcp/interfaces/communication_protocol.py @@ -6,7 +6,7 @@ """ from abc import ABC, abstractmethod -from typing import Dict, Any, AsyncGenerator, TYPE_CHECKING +from typing import Callable, Dict, Any, AsyncGenerator, TYPE_CHECKING from utcp.data.register_manual_response import RegisterManualResult from utcp.data.call_template import CallTemplate if TYPE_CHECKING: @@ -24,8 +24,31 @@ class CommunicationProtocol(ABC): - Discovering available tools from providers - Managing provider lifecycle (registration/deregistration) - Executing tool calls through the appropriate protocol + + A protocol is registered in one of two registries, and the choice decides + who its state belongs to: + + - `communication_protocols` holds an INSTANCE that is shared by every + `UtcpClient` in the process, and so is any state it keeps. That is the + right home for state that is meant to be shared (a credential cache, a + pooled HTTP session, a registry a decorator writes into). The instance + lives as long as the process that registered it; no client closes it. + - `communication_protocol_factories` holds a FACTORY, for a protocol whose + state must not be shared between clients: live sessions or connections + keyed per manual, child processes — anything one client's use or + `close()` would take away from another. Each `UtcpClient` calls it once + — at creation, or on first use for a factory registered later — so each + client gets its own instance, its own connections, and its own teardown + on `close()`. That is what makes a client per tenant, per user, or per + pooled connection actually isolate them, rather than every client + reaching into one shared instance. + + A type registered as a factory wins over the same type registered as an + instance, so a plugin migrates by moving its registration from one + registry to the other and callers change nothing. """ communication_protocols: dict[str, 'CommunicationProtocol'] = {} + communication_protocol_factories: dict[str, Callable[[], 'CommunicationProtocol']] = {} @abstractmethod async def register_manual(self, caller: 'UtcpClient', manual_call_template: CallTemplate) -> RegisterManualResult: @@ -118,3 +141,18 @@ async def call_tool_streaming(self, caller: 'UtcpClient', tool_name: str, tool_a TimeoutError: If the tool call exceeds the configured timeout. """ pass + + async def close(self) -> None: + """REQUIRED + Release every connection, session, process or other resource this + protocol instance holds. + + `UtcpClient.close()` calls this on each instance the client created + from `communication_protocol_factories`, and `UtcpClient.create()` + calls it on those instances when initialization fails after they were + created. A shared instance from `communication_protocols` is never + closed on a client's behalf. + + The default releases nothing, for protocols that hold nothing. + """ + pass diff --git a/core/src/utcp/plugins/discovery.py b/core/src/utcp/plugins/discovery.py index 830214e..d7aca50 100644 --- a/core/src/utcp/plugins/discovery.py +++ b/core/src/utcp/plugins/discovery.py @@ -6,6 +6,7 @@ from utcp.interfaces.tool_post_processor import ToolPostProcessor, ToolPostProcessorConfigSerializer from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp.data.call_template import CallTemplate, CallTemplateSerializer +from typing import Callable import logging logger = logging.getLogger(__name__) @@ -82,6 +83,34 @@ def register_communication_protocol(communication_protocol_type: str, communicat logger.info("Registered communication protocol type: " + communication_protocol_type) return True +def register_communication_protocol_factory(communication_protocol_type: str, factory: Callable[[], CommunicationProtocol], override: bool = False) -> bool: + """REQUIRED + Register a communication protocol as a factory, so that every `UtcpClient` + gets its own instance of it. + + Use this instead of `register_communication_protocol` for a protocol whose + state must not be shared between clients: live sessions or connections + keyed per manual, child processes — anything one client's use or `close()` + would take away from another. (A credential cache or a pooled HTTP session + is meant to be shared and stays an instance.) Each `UtcpClient` calls the factory once — at + creation, or on first use if the factory is registered later — and that + client's `close()` tears the instance down. A type registered as a factory + wins over the same type registered as an instance. + + Args: + communication_protocol_type: The communication protocol type identifier. + factory: A callable returning a new communication protocol instance. + override: Whether to override an existing factory for this type. + + Returns: + True if the factory was registered, False otherwise. + """ + if not override and communication_protocol_type in CommunicationProtocol.communication_protocol_factories: + return False + CommunicationProtocol.communication_protocol_factories[communication_protocol_type] = factory + logger.info("Registered communication protocol factory for type: " + communication_protocol_type) + return True + def register_tool_repository(tool_repository_type: str, tool_repository: Serializer[ConcurrentToolRepository], override: bool = False) -> bool: """REQUIRED Register a tool repository implementation. diff --git a/core/src/utcp/utcp_client.py b/core/src/utcp/utcp_client.py index 6dd2172..f902d0e 100644 --- a/core/src/utcp/utcp_client.py +++ b/core/src/utcp/utcp_client.py @@ -158,3 +158,18 @@ async def get_required_variables_for_registered_tool(self, tool_name: str) -> Li A list of required variables for the tool. """ pass + + @abstractmethod + async def close(self) -> None: + """REQUIRED + Close the protocol instances this client created and release their resources. + + Closes every instance the client obtained from + `CommunicationProtocol.communication_protocol_factories`. Instances from + `CommunicationProtocol.communication_protocols` are shared by every client + in the process and are left running. + + Every owned instance is closed even if one of them fails to close; the + failures are then raised together as a `UtcpProtocolCloseError`. + """ + pass diff --git a/core/tests/client/test_client_protocol_ownership.py b/core/tests/client/test_client_protocol_ownership.py new file mode 100644 index 0000000..01f00da --- /dev/null +++ b/core/tests/client/test_client_protocol_ownership.py @@ -0,0 +1,317 @@ +"""A client owns the protocol instances it creates from the factory registry. + +Everything here goes through the public surface: the two registries on +``CommunicationProtocol``, ``UtcpClient.create``, ``register_manual``, +``call_tool`` and ``close``. +""" + +import asyncio +import logging +from typing import Any, AsyncGenerator, Dict, List + +import pytest + +from utcp.data.call_template import CallTemplate +from utcp.data.register_manual_response import RegisterManualResult +from utcp.data.tool import JsonSchema, Tool +from utcp.data.utcp_manual import UtcpManual +from utcp.exceptions import UtcpProtocolCloseError, UtcpVariableNotFound +from utcp.interfaces.communication_protocol import CommunicationProtocol +from utcp.utcp_client import UtcpClient +from utcp_http.http_call_template import HttpCallTemplate + + +class RecordingProtocol(CommunicationProtocol): + """Records which client-facing calls reached this instance, and how often it was closed.""" + + def __init__(self): + self.registered: List[str] = [] + self.called: List[str] = [] + self.closed = 0 + + async def register_manual(self, caller: UtcpClient, manual_call_template: CallTemplate) -> RegisterManualResult: + self.registered.append(manual_call_template.name) + tool = Tool( + name="ping", + description="answers with the identity of the instance that served it", + inputs=JsonSchema(type="object", properties={}), + outputs=JsonSchema(type="object", properties={}), + tags=[], + tool_call_template=manual_call_template, + ) + return RegisterManualResult( + manual_call_template=manual_call_template, + manual=UtcpManual(manual_version="1.0.0", tools=[tool]), + success=True, + errors=[], + ) + + async def deregister_manual(self, caller: UtcpClient, manual_call_template: CallTemplate) -> None: + pass + + async def call_tool(self, caller: UtcpClient, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> Any: + self.called.append(tool_name) + return {"served_by": id(self)} + + async def call_tool_streaming(self, caller: UtcpClient, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: + yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) + + async def close(self) -> None: + self.closed += 1 + + +def http_manual(name: str) -> HttpCallTemplate: + return HttpCallTemplate(name=name, url="https://example.test/utcp", http_method="POST", call_template_type="http") + + +@pytest.fixture +def isolated_registries(monkeypatch): + """Both registries start empty and are restored after the test.""" + monkeypatch.setattr(CommunicationProtocol, "communication_protocols", {}) + monkeypatch.setattr(CommunicationProtocol, "communication_protocol_factories", {}) + + +@pytest.fixture +def recording_factory(isolated_registries): + """Registers a factory for ``http`` and returns the instances it made, in order.""" + made: List[RecordingProtocol] = [] + + def factory() -> RecordingProtocol: + protocol = RecordingProtocol() + made.append(protocol) + return protocol + + CommunicationProtocol.communication_protocol_factories["http"] = factory + return made + + +class TestPerClientInstances: + + @pytest.mark.asyncio + async def test_each_client_gets_its_own_instance_and_routes_to_it(self, recording_factory): + client_a = await UtcpClient.create() + client_b = await UtcpClient.create() + assert len(recording_factory) == 2 + own_a, own_b = recording_factory + + await client_a.register_manual(http_manual("a_manual")) + await client_b.register_manual(http_manual("b_manual")) + result_a = await client_a.call_tool("a_manual.ping", {}) + result_b = await client_b.call_tool("b_manual.ping", {}) + + assert own_a.registered == ["a_manual"] and own_a.called == ["a_manual.ping"] + assert own_b.registered == ["b_manual"] and own_b.called == ["b_manual.ping"] + assert result_a["served_by"] == id(own_a) + assert result_b["served_by"] == id(own_b) + + @pytest.mark.asyncio + async def test_a_factory_wins_over_a_shared_instance_of_the_same_type(self, recording_factory): + shared = RecordingProtocol() + CommunicationProtocol.communication_protocols["http"] = shared + + client = await UtcpClient.create() + await client.register_manual(http_manual("m")) + + assert recording_factory[0].registered == ["m"] + assert shared.registered == [] + + @pytest.mark.asyncio + async def test_a_shared_instance_registered_after_the_client_exists_is_still_used(self, isolated_registries): + # Late registration in the shared registry keeps working: shared + # instances are looked up live, not snapshotted at creation. + client = await UtcpClient.create() + shared = RecordingProtocol() + CommunicationProtocol.communication_protocols["http"] = shared + + await client.register_manual(http_manual("m")) + + assert shared.registered == ["m"] + + @pytest.mark.asyncio + async def test_a_factory_registered_after_the_client_exists_is_adopted_on_first_use_and_owned(self, isolated_registries): + client = await UtcpClient.create() + made: List[RecordingProtocol] = [] + + def factory() -> RecordingProtocol: + protocol = RecordingProtocol() + made.append(protocol) + return protocol + + CommunicationProtocol.communication_protocol_factories["http"] = factory + + await client.register_manual(http_manual("m")) + await client.register_manual(http_manual("n")) + # One instance, made on first use, serving every later call... + assert len(made) == 1 + assert made[0].registered == ["m", "n"] + # ...and owned: the client's close() reaches it. + await client.close() + assert made[0].closed == 1 + + @pytest.mark.asyncio + async def test_a_type_with_no_protocol_in_either_registry_names_what_is_available(self, isolated_registries): + # "http" has a serializer (the plugin is imported) but, with the + # registries isolated, no protocol in either of them. + CommunicationProtocol.communication_protocol_factories["cli"] = RecordingProtocol + CommunicationProtocol.communication_protocols["text"] = RecordingProtocol() + client = await UtcpClient.create() + + with pytest.raises(ValueError, match=r"type http found, available types: \['cli', 'text'\]"): + await client.register_manual(http_manual("m")) + + +class TestCloseIsScopedToWhatTheClientOwns: + + @pytest.mark.asyncio + async def test_close_drains_own_instance_and_leaves_other_clients_and_shared_ones_running(self, recording_factory): + shared = RecordingProtocol() + CommunicationProtocol.communication_protocols["cli"] = shared + client_a = await UtcpClient.create() + client_b = await UtcpClient.create() + own_a, own_b = recording_factory + + await client_a.close() + + assert own_a.closed == 1 + assert own_b.closed == 0 + # Every other client in the process is still using this one. + assert shared.closed == 0 + + await client_b.close() + assert own_b.closed == 1 + assert shared.closed == 0 + + @pytest.mark.asyncio + async def test_close_waits_for_every_owned_protocol_even_when_one_fails(self, isolated_registries): + # The failing close raises IMMEDIATELY; the healthy one takes a moment. + # A first-failure-wins close would return before the healthy one ended. + class FailingProtocol(RecordingProtocol): + async def close(self) -> None: + raise RuntimeError("transport refused to close") + + class SlowProtocol(RecordingProtocol): + async def close(self) -> None: + await asyncio.sleep(0.02) + self.closed += 1 + + slow = SlowProtocol() + CommunicationProtocol.communication_protocol_factories["http"] = FailingProtocol + CommunicationProtocol.communication_protocol_factories["cli"] = lambda: slow + + client = await UtcpClient.create() + with pytest.raises(UtcpProtocolCloseError) as raised: + await client.close() + + # The failure is still surfaced, with its cause inside... + assert len(raised.value.failures) == 1 + assert isinstance(raised.value.failures[0], RuntimeError) + # ...but only after every owned protocol has finished closing. + assert slow.closed == 1 + + +class TestFailedCreateClosesWhatItCreated: + + @pytest.mark.asyncio + async def test_an_instance_made_for_a_create_that_fails_is_closed_not_orphaned(self, recording_factory): + # The factory fires, then variable substitution fails on a reference + # nothing can resolve — create() raises and the caller never gets a + # client to close. + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config={"variables": {"DERIVED": "${NOWHERE_TO_BE_FOUND}"}}) + + assert len(recording_factory) == 1 + assert recording_factory[0].closed == 1 + + @pytest.mark.asyncio + async def test_a_factory_that_raises_leaves_the_instances_before_it_closed(self, recording_factory): + def broken_factory() -> CommunicationProtocol: + raise RuntimeError("no transport available") + + CommunicationProtocol.communication_protocol_factories["cli"] = broken_factory + + with pytest.raises(RuntimeError, match="no transport available"): + await UtcpClient.create() + + assert len(recording_factory) == 1 + assert recording_factory[0].closed == 1 + + @pytest.mark.asyncio + async def test_a_shared_instance_survives_a_failed_create(self, recording_factory): + shared = RecordingProtocol() + CommunicationProtocol.communication_protocols["cli"] = shared + + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config={"variables": {"DERIVED": "${NOWHERE_TO_BE_FOUND}"}}) + + assert recording_factory[0].closed == 1 + assert shared.closed == 0 + + @pytest.mark.asyncio + async def test_a_failing_cleanup_is_reported_and_the_original_error_surfaces(self, isolated_registries, caplog): + class FailingProtocol(RecordingProtocol): + async def close(self) -> None: + raise RuntimeError("transport refused to close") + + CommunicationProtocol.communication_protocol_factories["http"] = FailingProtocol + + with caplog.at_level(logging.ERROR, logger="utcp.implementations.utcp_client_implementation"): + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config={"variables": {"DERIVED": "${NOWHERE_TO_BE_FOUND}"}}) + + reported = [record for record in caplog.records if "closing the protocols it had created failed too" in record.getMessage()] + assert len(reported) == 1 + assert "transport refused to close" in reported[0].exc_text + + +def unresolvable_http_manual(name: str) -> HttpCallTemplate: + """A template whose registration raises UtcpVariableNotFound before it reaches any protocol.""" + return HttpCallTemplate(name=name, url="https://example.test/${NOWHERE_TO_BE_FOUND}", http_method="POST", call_template_type="http") + + +class SlowRegisteringProtocol(RecordingProtocol): + """Registration takes a moment and is recorded in a shared event log, as is close().""" + + def __init__(self, events: List[str]): + super().__init__() + self.events = events + + async def register_manual(self, caller: UtcpClient, manual_call_template: CallTemplate) -> RegisterManualResult: + await asyncio.sleep(0.02) + self.events.append(f"registered {manual_call_template.name}") + return await super().register_manual(caller, manual_call_template) + + async def close(self) -> None: + await super().close() + self.events.append("closed") + + +class TestABatchRegistrationSettlesEverySiblingBeforeRaising: + + @pytest.mark.asyncio + async def test_register_manuals_raises_only_once_every_sibling_registration_has_finished(self, isolated_registries): + # The first manual fails immediately (unresolvable variable); the + # second is still registering. A first-failure-wins gather would raise + # with the second still running underneath the caller. + events: List[str] = [] + CommunicationProtocol.communication_protocols["http"] = SlowRegisteringProtocol(events) + client = await UtcpClient.create() + + with pytest.raises(UtcpVariableNotFound): + await client.register_manuals([unresolvable_http_manual("bad_manual"), http_manual("slow_manual")]) + + assert events == ["registered slow_manual"] + + @pytest.mark.asyncio + async def test_a_failed_create_closes_its_protocols_only_after_every_registration_has_finished(self, isolated_registries): + # Same batch, through create(): the owned protocol must not be closed + # while a sibling registration is still using it. + events: List[str] = [] + CommunicationProtocol.communication_protocol_factories["http"] = lambda: SlowRegisteringProtocol(events) + + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config={"manual_call_templates": [ + unresolvable_http_manual("bad_manual"), + http_manual("slow_manual"), + ]}) + + assert events == ["registered slow_manual", "closed"] diff --git a/plugins/communication_protocols/mcp/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index 9b0641b..8f6a78b 100644 --- a/plugins/communication_protocols/mcp/pyproject.toml +++ b/plugins/communication_protocols/mcp/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-mcp" -version = "1.1.4" +version = "1.2.0" authors = [ { name = "UTCP Contributors" }, ] @@ -14,7 +14,7 @@ requires-python = ">=3.11" dependencies = [ "pydantic>=2.0", "mcp>=1.12,<2", - "utcp>=1.1.4", + "utcp>=1.2.0", "mcp-use>=1.3", "langchain>=0.3.27,<0.4.0", ] diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/__init__.py b/plugins/communication_protocols/mcp/src/utcp_mcp/__init__.py index 85abb78..f81f9f7 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/__init__.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/__init__.py @@ -1,9 +1,16 @@ from utcp_mcp.mcp_communication_protocol import McpCommunicationProtocol from utcp_mcp.mcp_call_template import McpCallTemplate, McpCallTemplateSerializer -from utcp.plugins.discovery import register_communication_protocol, register_call_template +from utcp.plugins.discovery import register_communication_protocol_factory, register_call_template def register(): - register_communication_protocol("mcp", McpCommunicationProtocol()) + # A FACTORY, not an instance: this protocol holds live MCP sessions (and, + # for stdio, child processes). Shared, every client in the process would + # dial into one session cache — so a caller that creates a client per + # tenant, per user, or per pooled connection would not actually be + # isolating them, and one client's close() would drain everyone's + # sessions. One instance per client gives each its own connections and its + # own teardown. + register_communication_protocol_factory("mcp", McpCommunicationProtocol) register_call_template("mcp", McpCallTemplateSerializer()) __all__ = [ diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 3406fd6..ec2b032 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -197,10 +197,10 @@ def __init__(self): # concurrent first calls for the same server dial once instead of each # spawning a session and leaking all but the last. self._session_creations: "Dict[Tuple[int, str], asyncio.Task]" = {} - # One MCPClient per distinct server configuration. This protocol object is - # registered once per process and shared by every manual, so a single - # client would make manuals with different configurations evict each - # other's sessions, including sessions still in use by a concurrent call. + # One MCPClient per distinct server configuration. This protocol object + # serves every manual of the UtcpClient that owns it, so a single client + # would make manuals with different configurations evict each other's + # sessions, including sessions still in use by a concurrent call. self._mcp_clients: Dict[str, MCPClient] = {} # Which configuration each owner (calling UtcpClient plus manual name) # currently uses, so a client nothing references any more can be closed @@ -844,7 +844,7 @@ async def close(self) -> None: """Close all active sessions and clean up resources.""" self._log_info("Closing MCP communication protocol and cleaning up all sessions") await self._cleanup_all_sessions() - # Drain the OAuth state so this shared instance holds no credentials past + # Drain the OAuth state so this instance holds no credentials past # close(). Cancel in-flight fetches (asyncio can) and drop their entries; # a fetch that still lands finds it is no longer the current entry and, # by the caching rule in _on_oauth_fetch_done, does not repopulate the diff --git a/plugins/communication_protocols/websocket/pyproject.toml b/plugins/communication_protocols/websocket/pyproject.toml index 090df01..f708117 100644 --- a/plugins/communication_protocols/websocket/pyproject.toml +++ b/plugins/communication_protocols/websocket/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-websocket" -version = "1.1.5" +version = "1.2.0" authors = [ { name = "UTCP Contributors" }, ] @@ -14,7 +14,7 @@ requires-python = ">=3.10" dependencies = [ "pydantic>=2.0", "aiohttp>=3.8", - "utcp>=1.1.4" + "utcp>=1.2.0" ] classifiers = [ "Development Status :: 4 - Beta", diff --git a/plugins/communication_protocols/websocket/src/utcp_websocket/__init__.py b/plugins/communication_protocols/websocket/src/utcp_websocket/__init__.py index 21c5879..33194da 100644 --- a/plugins/communication_protocols/websocket/src/utcp_websocket/__init__.py +++ b/plugins/communication_protocols/websocket/src/utcp_websocket/__init__.py @@ -3,14 +3,18 @@ This plugin provides WebSocket-based real-time bidirectional communication protocol. """ -from utcp.plugins.discovery import register_communication_protocol, register_call_template +from utcp.plugins.discovery import register_communication_protocol_factory, register_call_template from utcp_websocket.websocket_communication_protocol import WebSocketCommunicationProtocol from utcp_websocket.websocket_call_template import WebSocketCallTemplate, WebSocketCallTemplateSerializer def register(): """Register the WebSocket communication protocol and call template serializer.""" - # Register WebSocket communication protocol - register_communication_protocol("websocket", WebSocketCommunicationProtocol()) + # A FACTORY, not an instance: this protocol keeps one live WebSocket per + # manual name and URL. Shared, two clients registering the same manual + # would use — and on deregistration close — each other's connection, and + # one client's close() would drop everyone's. One instance per client + # gives each its own connections and its own teardown. + register_communication_protocol_factory("websocket", WebSocketCommunicationProtocol) # Register call template serializer register_call_template("websocket", WebSocketCallTemplateSerializer())