diff --git a/README.md b/README.md index 400b0da..c6c48ea 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 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: the only state they keep is an OAuth token cache, which is meant to be reused across clients; every request opens its own session. + +`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..c6c48ea 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 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: the only state they keep is an OAuth token cache, which is meant to be reused across clients; every request opens its own session. + +`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..e494b6e 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 @@ -25,6 +25,12 @@ logger = logging.getLogger(__name__) + +def _sanitize_manual_name(name: str) -> str: + """The name a manual is registered under: every non-word character becomes an underscore.""" + return re.sub(r'[^\w]', '_', name) + + class UtcpClientImplementation(UtcpClient): """REQUIRED Implementation of the `UtcpClient` interface. @@ -39,6 +45,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 +156,51 @@ 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 undone + # here before the failure is re-raised: the manuals this attempt + # registered are removed (the tool repository may be caller-supplied + # and shared), then the protocol instances it created are closed. + # Exactly what it registered — recorded as each registration + # succeeds, never inferred from repository state, which another client + # sharing the repository may have changed in the meantime. + registered_by_this_attempt: List[str] = [] + 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: + results = await client._register_each(config.manual_call_templates, registered_by_this_attempt) + failures = [result for result in results if isinstance(result, BaseException)] + if failures: + raise failures[0] + except BaseException: + # Rollback first (deregistration needs the protocols open), then + # close — in a finally, so a cancellation that lands during the + # rollback still closes what this client created and then reaches + # the caller, as a cancellation must. + try: + for name in registered_by_this_attempt: + try: + await client.deregister_manual(name) + except Exception: + logger.error(f"UtcpClient.create failed, and removing the manual '{name}' it had registered failed too", exc_info=True) + finally: + 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: @@ -117,14 +229,12 @@ async def register_manual(self, manual_call_template: CallTemplate) -> RegisterM ValueError: If manual name is already registered or communication protocol is not found. """ # Replace all non-word characters with underscore - manual_call_template.name = re.sub(r'[^\w]', '_', manual_call_template.name) + manual_call_template.name = _sanitize_manual_name(manual_call_template.name) 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 @@ -162,34 +272,55 @@ async def register_manuals(self, manual_call_templates: List[CallTemplate]) -> L Returns: A list of `RegisterManualResult` instances representing the results of the registration. """ - # Create tasks for parallel CallTemplate registration - tasks = [] - for manual_call_template in manual_call_templates: - async def try_register_manual(manual_call_template=manual_call_template): - try: - result = await self.register_manual(manual_call_template) - if result.success: - logger.info(f"Successfully registered manual '{manual_call_template.name}' with {len(result.manual.tools)} tools") - else: - logger.error(f"Error registering manual '{manual_call_template.name}': {result.errors}") - return result - except UtcpVariableNotFound as e: - raise e - except Exception as e: - logger.error(f"Error registering manual '{manual_call_template.name}': {traceback.format_exc()}") - return RegisterManualResult( - manual_call_template=manual_call_template, - manual=UtcpManual(manual_version="0.0.0", tools=[]), - success=False, - errors=[traceback.format_exc()] - ) - - tasks.append(try_register_manual()) - - # Wait for all tasks to complete and collect results - results = await asyncio.gather(*tasks) + results = await self._register_each(manual_call_templates) + 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 _register_each( + self, + manual_call_templates: List[CallTemplate], + registered_names: Optional[List[str]] = None, + ) -> List[Union[RegisterManualResult, BaseException]]: + """Register every template in parallel and wait for EVERY one to settle. + + A raised failure is returned in place of its result rather than + propagated at once, so no sibling registration is still running + underneath a caller that holds the failure — create(), for one, closes + the protocols right after, and a registration still in flight would be + using a closed one. + + When ``registered_names`` is given, the registered name of each manual + that succeeded is appended to it AS it succeeds, so a caller that + fails part-way knows exactly what it registered — and only that. + """ + async def try_register_manual(manual_call_template: CallTemplate): + try: + result = await self.register_manual(manual_call_template) + if result.success: + if registered_names is not None: + registered_names.append(manual_call_template.name) + logger.info(f"Successfully registered manual '{manual_call_template.name}' with {len(result.manual.tools)} tools") + else: + logger.error(f"Error registering manual '{manual_call_template.name}': {result.errors}") + return result + except UtcpVariableNotFound as e: + raise e + except Exception as e: + logger.error(f"Error registering manual '{manual_call_template.name}': {traceback.format_exc()}") + return RegisterManualResult( + manual_call_template=manual_call_template, + manual=UtcpManual(manual_version="0.0.0", tools=[]), + success=False, + errors=[traceback.format_exc()] + ) + + return await asyncio.gather( + *(try_register_manual(manual_call_template) for manual_call_template in manual_call_templates), + return_exceptions=True, + ) + async def deregister_manual(self, manual_name: str) -> bool: """REQUIRED Deregister a manual from the client. @@ -203,7 +334,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 +381,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 +431,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 @@ -334,7 +465,7 @@ async def get_required_variables_for_manual_and_tools(self, manual_call_template Returns: A list of required variables for the manual and its tools. """ - manual_call_template.name = re.sub(r'[^\w]', '_', manual_call_template.name) + manual_call_template.name = _sanitize_manual_name(manual_call_template.name) variables_for_CallTemplate = self.variable_substitutor.find_required_variables(CallTemplateSerializer().to_dict(manual_call_template), manual_call_template.name) if len(variables_for_CallTemplate) > 0: try: @@ -342,9 +473,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..a47b229 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 + 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..c08d513 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 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..b9217c8 --- /dev/null +++ b/core/tests/client/test_client_protocol_ownership.py @@ -0,0 +1,448 @@ +"""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.data.utcp_client_config import UtcpClientConfig +from utcp.implementations.in_mem_tool_repository import InMemToolRepository +from utcp.interfaces.communication_protocol import CommunicationProtocol +from utcp.plugins.plugin_loader import ensure_plugins_initialized +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. + + Plugins load lazily on the first ``UtcpClient.create()`` of the session and + register into whatever dict holds the class attribute at that moment. They + are loaded here first, so their registrations land in the ORIGINAL dicts + (which the fixture restores) and never in the throwaway ones. + """ + ensure_plugins_initialized() + 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"] + + +class TestAFailedCreateLeavesNoManualBehind: + + @pytest.mark.asyncio + async def test_a_manual_registered_by_the_failed_attempt_is_removed_from_a_shared_repository(self, isolated_registries): + # The caller supplies (and keeps) the repository; the second manual + # registers fine before the first one's failure surfaces. create() + # raises -- and must not leave that manual behind in the caller's repo. + CommunicationProtocol.communication_protocols["http"] = RecordingProtocol() + repo = InMemToolRepository() + + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config=UtcpClientConfig( + tool_repository=repo, + manual_call_templates=[unresolvable_http_manual("bad_manual"), http_manual("good_manual")], + )) + + assert await repo.get_manual("good_manual") is None + assert await repo.get_tool("good_manual.ping") is None + + @pytest.mark.asyncio + async def test_a_manual_that_was_in_the_repository_before_the_attempt_survives_it(self, isolated_registries): + # "existing" is already registered by an earlier client on the same + # shared repository. This attempt names it again (that registration + # fails as a duplicate, without raising) and also fails outright on + # the bad manual. The rollback must remove only what THIS attempt + # added -- never the pre-existing manual. + CommunicationProtocol.communication_protocols["http"] = RecordingProtocol() + repo = InMemToolRepository() + earlier = await UtcpClient.create(config=UtcpClientConfig(tool_repository=repo, manual_call_templates=[http_manual("existing")])) + assert await repo.get_manual("existing") is not None + + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config=UtcpClientConfig( + tool_repository=repo, + manual_call_templates=[http_manual("existing"), unresolvable_http_manual("bad_manual"), http_manual("added_by_attempt")], + )) + + assert await repo.get_manual("existing") is not None + assert await repo.get_manual("added_by_attempt") is None + await earlier.close() + + +def test_plugin_registrations_survive_the_isolated_tests_above(): + # Runs after every isolated test in this file. If any of them had been the + # session's first create() and the fixture had patched the registry before + # plugins loaded, the plugin instances would have gone into the throwaway + # dict and be missing here. + ensure_plugins_initialized() + assert "http" in CommunicationProtocol.communication_protocols + assert "mcp" in CommunicationProtocol.communication_protocol_factories + + +class TestRollbackRemovesOnlyWhatThisAttemptRegistered: + + @pytest.mark.asyncio + async def test_a_manual_another_client_registered_during_the_attempt_survives_the_rollback(self, isolated_registries): + # While attempt A is registering its batch, client B (same shared + # repository) registers "contested". A's own registration of + # "contested" then fails as a duplicate; A fails outright on the bad + # manual. The rollback must remove what A registered -- and not B's + # "contested", which merely appeared after A started. + repo = InMemToolRepository() + client_b = await UtcpClient.create(config=UtcpClientConfig(tool_repository=repo)) + + class InterleavingProtocol(RecordingProtocol): + async def register_manual(self, caller, manual_call_template): + if manual_call_template.name == "trigger": + await client_b.register_manual(http_manual("contested")) + return await super().register_manual(caller, manual_call_template) + + CommunicationProtocol.communication_protocols["http"] = InterleavingProtocol() + + with pytest.raises(UtcpVariableNotFound): + await UtcpClient.create(config=UtcpClientConfig( + tool_repository=repo, + manual_call_templates=[http_manual("trigger"), http_manual("contested"), unresolvable_http_manual("bad_manual")], + )) + + assert await repo.get_manual("trigger") is None # A's own registration: rolled back + assert await repo.get_manual("contested") is not None # B's: untouched + await client_b.close() + + +class TestCancellationDuringRollbackStillClosesOwnedProtocols: + + @pytest.mark.asyncio + async def test_cancel_while_deregistering_closes_the_protocols_and_propagates_the_cancellation(self, isolated_registries): + # The rollback's deregistration blocks; the caller cancels create() + # while it is blocked. The owned protocol must still be closed, and + # the caller must see the cancellation. + deregister_started = asyncio.Event() + release_deregister = asyncio.Event() + + class BlockingDeregisterProtocol(RecordingProtocol): + async def deregister_manual(self, caller, manual_call_template): + deregister_started.set() + await release_deregister.wait() + + made: List[RecordingProtocol] = [] + + def factory() -> RecordingProtocol: + protocol = BlockingDeregisterProtocol() + made.append(protocol) + return protocol + + CommunicationProtocol.communication_protocol_factories["http"] = factory + + creating = asyncio.create_task(UtcpClient.create(config={"manual_call_templates": [ + {"name": "good_manual", "call_template_type": "http", "url": "https://example.test/utcp", "http_method": "POST"}, + {"name": "bad_manual", "call_template_type": "http", "url": "https://example.test/${NOWHERE_TO_BE_FOUND}", "http_method": "POST"}, + ]})) + await asyncio.wait_for(deregister_started.wait(), timeout=5) + creating.cancel() + + with pytest.raises(asyncio.CancelledError): + await creating + + assert len(made) == 1 + assert made[0].closed == 1 diff --git a/core/tests/client/test_utcp_client.py b/core/tests/client/test_utcp_client.py index 934bebf..9a5d951 100644 --- a/core/tests/client/test_utcp_client.py +++ b/core/tests/client/test_utcp_client.py @@ -12,6 +12,7 @@ from utcp.implementations.utcp_client_implementation import UtcpClientImplementation from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp.utcp_client import UtcpClient +from utcp.plugins.plugin_loader import ensure_plugins_initialized from utcp.data.utcp_client_config import UtcpClientConfig from utcp.exceptions import UtcpVariableNotFound, UtcpSerializerValidationError from utcp.interfaces.concurrent_tool_repository import ConcurrentToolRepository @@ -187,7 +188,12 @@ async def sample_tools(): @pytest.fixture def isolated_communication_protocols(monkeypatch): - """Isolates the CommunicationProtocol registry for each test.""" + """Isolates the CommunicationProtocol registry for each test. + + Plugins are loaded first so that their registrations land in the original + dict (restored after the test), not in the throwaway one. + """ + ensure_plugins_initialized() monkeypatch.setattr(CommunicationProtocol, "communication_protocols", {}) diff --git a/plugins/communication_protocols/cli/src/utcp_cli/cli_call_template.py b/plugins/communication_protocols/cli/src/utcp_cli/cli_call_template.py index 378f9dc..cf2e63f 100644 --- a/plugins/communication_protocols/cli/src/utcp_cli/cli_call_template.py +++ b/plugins/communication_protocols/cli/src/utcp_cli/cli_call_template.py @@ -147,8 +147,10 @@ class CliCallTemplate(CallTemplate): commands: A list of CommandStep objects defining the commands to execute in order. Each command can contain UTCP_ARG_argname_UTCP_END placeholders that will be replaced with values from tool_args during execution. - Placeholders are shell-quoted and therefore expand to exactly one - shell token (see class docstring). + Each placeholder becomes a shell-variable reference whose value + reaches the subprocess through a per-invocation environment + variable, so it expands to exactly one shell token and cannot be + reinterpreted as shell syntax (see class docstring). env_vars: A dictionary of environment variables to set for the command's execution context. Values can be static strings or placeholders for variables from the UTCP client's variable substitutor. Always @@ -238,9 +240,11 @@ class CliCallTemplate(CallTemplate): Security Considerations: - Commands are executed in a subprocess. Ensure that the commands specified are from a trusted source. - - `tool_args` values are shell-quoted on substitution, but the - *command template itself* is not — never assemble it from - untrusted input. + - `tool_args` values never touch the command text: they reach the + subprocess through per-invocation environment variables and the + shell expands them only after it has parsed the script. The + *command template itself* has no such protection — never assemble + it from untrusted input. - The host environment is restricted; secrets are not propagated unless explicitly named in `env_vars` or `inherit_env_vars`. - Commands should use the appropriate syntax for the target platform diff --git a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py index 87990a1..f26433e 100644 --- a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py @@ -63,7 +63,7 @@ class HttpCommunicationProtocol(CommunicationProtocol): Attributes: _session: Optional aiohttp ClientSession for connection reuse. - _oauth_tokens: Cache of OAuth2 tokens keyed by the full credential configuration (``OAuth2Auth.cache_key``). + _oauth_tokens: Cache of OAuth2 tokens keyed by the full credential configuration (OAuth2Auth's ``cache_key``). _log: Logger function for debugging and error reporting. """ 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 9b79234..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.6" +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()) diff --git a/plugins/communication_protocols/websocket/src/utcp_websocket/websocket_communication_protocol.py b/plugins/communication_protocols/websocket/src/utcp_websocket/websocket_communication_protocol.py index 74ca2ee..38243b3 100644 --- a/plugins/communication_protocols/websocket/src/utcp_websocket/websocket_communication_protocol.py +++ b/plugins/communication_protocols/websocket/src/utcp_websocket/websocket_communication_protocol.py @@ -63,7 +63,7 @@ class WebSocketCommunicationProtocol(CommunicationProtocol): Attributes: _connections: Active WebSocket connections by provider key. _sessions: aiohttp ClientSessions for connection management. - _oauth_tokens: Cache of OAuth2 tokens keyed by the full credential configuration (``OAuth2Auth.cache_key``). + _oauth_tokens: Cache of OAuth2 tokens keyed by the full credential configuration (OAuth2Auth's ``cache_key``). """ def __init__(self, logger_func: Optional[Callable[[str], None]] = None): diff --git a/scripts/extract_required_docs.py b/scripts/extract_required_docs.py index af6bfbb..debff6e 100644 --- a/scripts/extract_required_docs.py +++ b/scripts/extract_required_docs.py @@ -148,44 +148,52 @@ def process_section_content(content_lines): line = line.replace('{', '\\{').replace('}', '\\}') stripped = line.strip() - # Check if this looks like a parameter/item definition (name: description) + # Check if this looks like a parameter/item definition (name: description). + # A line with a colon that is NOT one -- prose such as "(note: ...)" or a + # quoted URL -- is ordinary text and falls through to the branches below; + # it must never be dropped. + param_match = None if ':' in stripped and not stripped.endswith(':'): colon_pos = stripped.find(':') - param_name = stripped[:colon_pos].strip() - param_desc = stripped[colon_pos + 1:].strip() - - # Check if param_name looks like a parameter (no spaces, reasonable length) - if ' ' not in param_name and len(param_name) <= 50 and param_name.replace('_', '').isalnum(): - # This is likely a parameter definition - processed.append(f"- **`{param_name}`**: {param_desc}") - - # Check for continuation lines (indented more than the parameter line) - base_indent = len(line) - len(line.lstrip()) - i += 1 - while i < len(content_lines): - next_line = content_lines[i] - next_stripped = next_line.strip() - next_indent = len(next_line) - len(next_line.lstrip()) if next_stripped else 0 - - # Check if we hit a code block - if next_stripped.startswith('```'): - break - - if not next_stripped: - # Empty line - add it and continue - processed.append('') - i += 1 - elif next_indent > base_indent: - # Continuation line - add with proper spacing - processed.append(f" {next_stripped}") - i += 1 - else: - # Not a continuation, back up and break - break - continue - + candidate = stripped[:colon_pos].strip() + # A parameter name has no spaces and is a plain identifier, and a + # definition puts a space after its colon ("name: description") -- + # which is what separates it from a URL such as https://example.com + if (' ' not in candidate and len(candidate) <= 50 and candidate.replace('_', '').isalnum() + and stripped[colon_pos + 1] == ' '): + param_match = (candidate, stripped[colon_pos + 1:].strip()) + + if param_match is not None: + param_name, param_desc = param_match + processed.append(f"- **`{param_name}`**: {param_desc}") + + # Check for continuation lines (indented more than the parameter line) + base_indent = len(line) - len(line.lstrip()) + i += 1 + while i < len(content_lines): + next_line = content_lines[i] + next_stripped = next_line.strip() + next_indent = len(next_line) - len(next_line.lstrip()) if next_stripped else 0 + + # Check if we hit a code block + if next_stripped.startswith('```'): + break + + if not next_stripped: + # Empty line - add it and continue + processed.append('') + i += 1 + elif next_indent > base_indent: + # Continuation line - add with proper spacing + processed.append(f" {next_stripped}") + i += 1 + else: + # Not a continuation, back up and break + break + continue + # Check if line starts with a list marker - elif stripped.startswith(('- ', '* ', '+ ')): + if stripped.startswith(('- ', '* ', '+ ')): # This is already a markdown list item processed.append(stripped) elif stripped.startswith(('1. ', '2. ', '3. ', '4. ', '5. ', '6. ', '7. ', '8. ', '9. ')): @@ -204,9 +212,13 @@ def process_section_content(content_lines): # Parse the docstring line by line for line in lines: stripped_lower = line.strip().lower() - - # Check if this line is a section header - if stripped_lower in section_headers or stripped_lower.endswith(':'): + + # Check if this line is a section header: a known Google-style header, or a + # Title-Cased line ending in a colon ("Security Considerations:", + # "Return Values (Complex):"). A sentence that merely ends in a colon + # ("Inheritance is controlled by `x`:", "Use the following:") is content -- + # treating it as a header would title-case it, code span included. + if stripped_lower in section_headers or self._is_custom_section_header(line.strip()): # Save previous section if it exists if current_section: processed_content = process_section_content(current_section_content) @@ -231,6 +243,10 @@ def process_section_content(content_lines): if processed_content: result.append(f"\n**{current_section.title()}**\n") result.extend(processed_content) + else: + # No section header anywhere: the whole docstring is the preamble, + # which is otherwise only flushed when a header follows it. + result.extend(process_section_content(current_section_content)) # Clean up the result final_result = [] @@ -530,18 +546,117 @@ def add_cross_references_post_generation(self, text: str, current_output_file: s class_anchor = re.sub(r'[^\w\-_]', '-', class_name.lower()).strip('-') link = f"[{class_name}](./{relative_path_str}#{class_anchor})" - # Don't replace matches that are in code blocks - lines = modified_text.split('\n') + # Don't replace matches that are in code blocks or inline code + # spans. Fenced blocks are delimited per line; everything + # between fences is handed to the inline-span scanner as ONE + # segment, so a span that crosses a newline stays a span. + segments: List[str] = [] + prose: List[str] = [] in_code_block = False - for i, line in enumerate(lines): + + def flush_prose() -> None: + if prose: + segments.append(self._sub_outside_inline_code(pattern, link, '\n'.join(prose))) + prose.clear() + + for line in modified_text.split('\n'): if line.strip().startswith('```'): + flush_prose() in_code_block = not in_code_block - elif not in_code_block: - lines[i] = re.sub(pattern, link, line) - modified_text = '\n'.join(lines) - + segments.append(line) + elif in_code_block: + segments.append(line) + else: + prose.append(line) + flush_prose() + modified_text = '\n'.join(segments) + return modified_text + + @staticmethod + def _is_custom_section_header(stripped: str) -> bool: + """A custom section header is a Title-Cased line ending in a colon. + + Every word starts with a capital letter or a digit (leading punctuation such + as an opening parenthesis is skipped), and there is no code span. Length is + not a criterion: "Return Values (Complex):" and "Section 1/2:" are headers, + "Use the following:" and "def tool1():" are content. + """ + if not stripped.endswith(':') or '`' in stripped: + return False + words = stripped[:-1].split() + if not words: + return False + for word in words: + first = next((ch for ch in word if ch.isalnum()), None) + if first is None or not (first.isupper() or first.isdigit()): + return False + return True + + @staticmethod + def _split_inline_code(line: str) -> List[Tuple[str, bool]]: + """Split a line into (text, is_code) parts. + + A code span opened by a run of N backticks closes only on the next run of + exactly N backticks, so ``a`b`` is one span. An unclosed run is text. + """ + parts: List[Tuple[str, bool]] = [] + pos = 0 + text_start = 0 + while pos < len(line): + if line[pos] != '`': + pos += 1 + continue + run_end = pos + while run_end < len(line) and line[run_end] == '`': + run_end += 1 + fence = line[pos:run_end] + close = line.find(fence, run_end) + # The closing run must be exactly as long: skip longer runs + while close != -1 and close + len(fence) < len(line) and line[close + len(fence)] == '`': + skip = close + while skip < len(line) and line[skip] == '`': + skip += 1 + close = line.find(fence, skip) + if close == -1: + pos = run_end + continue + if text_start < pos: + parts.append((line[text_start:pos], False)) + span_end = close + len(fence) + parts.append((line[pos:span_end], True)) + pos = text_start = span_end + if text_start < len(line): + parts.append((line[text_start:], False)) + return parts + + @classmethod + def _sub_outside_inline_code(cls, pattern: str, replacement: str, line: str) -> str: + """Substitute only outside `...` / ``...`` spans. + + Markdown renders a code span literally, so a link inserted inside one shows up + as raw brackets instead of a link. + """ + return ''.join(text if is_code else re.sub(pattern, replacement, text) for text, is_code in cls._split_inline_code(line)) + def _render_methods(self, content: List[str], methods: List[DocEntry], file_path: str) -> None: + """Append a class's documented methods to ``content``.""" + content.extend(["#### Methods:", ""]) + for method in methods: + # Add cross-references to method signature + linked_signature = self.add_cross_references(method.signature, file_path) + docstrings = method.docstring if method.docstring else "*No method documentation available*" + content.extend( + [ + "
", + f"{linked_signature}", + "", + docstrings, + "
", + "", + ] + ) + def generate_module_markdown(self, file_path: str, file_data: Dict[str, List[DocEntry]]) -> str: """Generate markdown content for a single module/file.""" if not any(file_data.values()): @@ -632,34 +747,27 @@ def generate_module_markdown(self, file_path: str, file_data: Dict[str, List[Doc # Add methods for this class if class_entry.name in methods_by_class: - content.extend(["#### Methods:", ""]) - - for method in methods_by_class[class_entry.name]: - method_anchor = re.sub(r'[^\w\-_]', '-', f"{class_entry.name}-{method.name}".lower()).strip('-') - - # Add cross-references to method signature - linked_signature = self.add_cross_references(method.signature, file_path) - - docstrings = "" - - if method.docstring: - docstrings = method.docstring - else: - docstrings = "*No method documentation available*" - - content.extend( - [ - "
", - f"{linked_signature}", - "", - docstrings, - "
", - "", - ] - ) - + self._render_methods(content, methods_by_class[class_entry.name], file_path) + content.extend(["---", ""]) - + + # A class whose own docstring is not REQUIRED can still have REQUIRED + # methods. They are required documentation and the index counts them, + # so they are rendered under a bare class heading rather than lost. + documented_classes = {class_entry.name for class_entry in file_data['classes']} + for class_name, methods in methods_by_class.items(): + if class_name in documented_classes: + continue + class_anchor = re.sub(r'[^\w\-_]', '-', class_name.lower()).strip('-') + content.extend([ + f"### class {class_name} {{#{class_anchor}}}", + "", + "*No class documentation available*", + "", + ]) + self._render_methods(content, methods, file_path) + content.extend(["---", ""]) + # Add standalone functions if file_data['functions']: for func_entry in file_data['functions']: @@ -748,7 +856,7 @@ def generate_index_file(self, modules: Dict[str, Dict[str, List[DocEntry]]], out index_path = output_path / "index.md" target_path = Path(output_file_path) try: - relative_path = target_path.relative_to(output_path) + relative_path = target_path.relative_to(output_path).as_posix() link_path = f"./{relative_path}" except ValueError: # Fallback to simple filename if relative path calculation fails @@ -805,7 +913,7 @@ def generate_index_file(self, modules: Dict[str, Dict[str, List[DocEntry]]], out index_path = output_path / "index.md" target_path = Path(output_file_path) try: - relative_path = target_path.relative_to(output_path) + relative_path = target_path.relative_to(output_path).as_posix() link_path = f"./{relative_path}" except ValueError: # Fallback to simple filename if relative path calculation fails @@ -900,10 +1008,11 @@ def generate_docs(self, output_dir: str) -> None: # Second pass: Add cross-references and write files for file_path, (content, output_file_path) in generated_files.items(): - # Post-process content to add proper cross-references - processed_content = self.add_cross_references_post_generation(content, str(output_file_path).replace('\\', '/')) - # Also process field references - lines = processed_content.split('\n') + # Field lines were emitted as "- `name: type`" -- the backticks are a + # placeholder (see format_field_with_references), not a code span. + # Unwrap them BEFORE cross-referencing: the cross-reference pass + # leaves real code spans alone, and these must receive links. + lines = content.split('\n') processed_lines = [] for line in lines: if line.strip().startswith('- `') and ':' in line: @@ -916,8 +1025,9 @@ def generate_docs(self, output_dir: str) -> None: processed_lines.append(line) else: processed_lines.append(line) - - final_content = '\n'.join(processed_lines) + + # Post-process content to add proper cross-references + final_content = self.add_cross_references_post_generation('\n'.join(processed_lines), str(output_file_path).replace('\\', '/')) with open(output_file_path, 'w', encoding='utf-8') as f: f.write(final_content)