autobyteus 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

autobyteus/rpc/__init__.py

@@ -1,73 +0,0 @@
-# file: autobyteus/autobyteus/rpc/__init__.py
-"""
-AutoByteUs Remote Procedure Call (RPC) Framework.
-
-This package provides components for enabling communication between AutoByteUs agents
-running in separate processes or on different machines. It includes:
-- Protocol definitions for structured messaging.
-- Configuration management for RPC servers.
-- Client-side components for connecting to and interacting with remote agents.
-- Server-side components for exposing an agent's capabilities remotely.
-- High-level hosting utilities for easily serving agents.
-"""
-
-from .protocol import ProtocolMessage, MessageType, RequestType, ResponseType, ErrorCode, ErrorDetails, EventType
-from .transport_type import TransportType
-from .config import AgentServerConfig, AgentServerRegistry, default_agent_server_registry
-from .client import (
-    AbstractClientConnection,
-    StdioClientConnection,
-    SseClientConnection, 
-    ClientConnectionManager,
-    default_client_connection_manager
-)
-from .server import (
-    AgentServerEndpoint,
-    StdioServerHandler, 
-    SseServerHandler, 
-    BaseMethodHandler, 
-    DiscoverCapabilitiesHandler, 
-    InvokeMethodHandler,
-    InitiateStreamDownloadHandler # Added InitiateStreamDownloadHandler
-)
-from .hosting import serve_agent_stdio, serve_agent_http_sse, serve_single_agent_http_sse, serve_multiple_agents_http_sse # Added all hosting functions
-
-# RemoteAgentProxy is part of the agent package, but closely related
-# from ..agent.remote_agent import RemoteAgentProxy 
-
-__all__ = [
-    # Protocol
-    "ProtocolMessage",
-    "MessageType",
-    "RequestType",
-    "ResponseType",
-    "ErrorCode",
-    "ErrorDetails",
-    "EventType", 
-    # Transport
-    "TransportType",
-    # Config
-    "AgentServerConfig",
-    "AgentServerRegistry",
-    "default_agent_server_registry",
-    # Client
-    "AbstractClientConnection",
-    "StdioClientConnection",
-    "SseClientConnection", 
-    "ClientConnectionManager",
-    "default_client_connection_manager",
-    # Server
-    "AgentServerEndpoint",
-    "StdioServerHandler", 
-    "SseServerHandler",   
-    "BaseMethodHandler",  
-    "DiscoverCapabilitiesHandler", 
-    "InvokeMethodHandler",
-    "InitiateStreamDownloadHandler", # Added
-    # Hosting
-    "serve_agent_stdio", 
-    "serve_agent_http_sse", # Main alias
-    "serve_single_agent_http_sse", # Specific for one agent
-    "serve_multiple_agents_http_sse", # Specific for multiple agents
-]
-

autobyteus/rpc/client/__init__.py

@@ -1,17 +0,0 @@
-# file: autobyteus/autobyteus/rpc/client/__init__.py
-"""
-Client-side components for the AutoByteUs RPC framework.
-These components are used to connect to and interact with remote Agent Servers.
-"""
-from .abstract_client_connection import AbstractClientConnection
-from .stdio_client_connection import StdioClientConnection
-from .sse_client_connection import SseClientConnection # Added SseClientConnection
-from .client_connection_manager import ClientConnectionManager, default_client_connection_manager
-
-__all__ = [
-    "AbstractClientConnection",
-    "StdioClientConnection",
-    "SseClientConnection", # Added SseClientConnection to exports
-    "ClientConnectionManager",
-    "default_client_connection_manager",
-]

autobyteus/rpc/client/abstract_client_connection.py

@@ -1,124 +0,0 @@
-# file: autobyteus/autobyteus/rpc/client/abstract_client_connection.py
-import asyncio
-import logging
-from abc import ABC, abstractmethod
-from typing import AsyncIterator, Optional, Dict, Any # Added Dict, Any
-
-from autobyteus.rpc.protocol import ProtocolMessage
-
-logger = logging.getLogger(__name__)
-
-class AbstractClientConnection(ABC):
-    """
-    Abstract base class for client-side connections to an Agent Server.
-    Defines the common interface for sending requests and handling responses/events.
-    """
-
-    def __init__(self, server_id: str):
-        self.server_id = server_id
-        self._is_connected: bool = False
-        self._connection_lock = asyncio.Lock()
-
-    @property
-    def is_connected(self) -> bool:
-        """Indicates if the connection is currently active."""
-        return self._is_connected
-
-    @abstractmethod
-    async def connect(self) -> None:
-        """
-        Establishes the connection to the remote server.
-        Sets _is_connected to True on success.
-        Raises:
-            ConnectionError: If connection fails.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def close(self) -> None:
-        """
-        Closes the connection to the remote server.
-        Sets _is_connected to False.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def send_request(self, request_message: ProtocolMessage) -> ProtocolMessage:
-        """
-        Sends a request ProtocolMessage to the server and awaits a corresponding
-        response ProtocolMessage (either a result or an error).
-
-        Args:
-            request_message: The ProtocolMessage to send (must be of type REQUEST).
-
-        Returns:
-            The response ProtocolMessage from the server.
-
-        Raises:
-            ConnectionError: If not connected or if communication fails.
-            TimeoutError: If the server does not respond within a reasonable time.
-            ValueError: If request_message is not a valid request.
-        """
-        raise NotImplementedError
-
-    async def events(self) -> AsyncIterator[ProtocolMessage]: # pragma: no cover
-        """
-        (Optional for some connection types like stdio, primary for SSE)
-        Provides an asynchronous iterator for server-pushed events.
-
-        Yields:
-            ProtocolMessage objects of type EVENT from the server.
-
-        Raises:
-            ConnectionError: If not connected or if the event stream fails.
-        """
-        # Default implementation for connection types that don't support server-pushed events.
-        if False: # This makes the method an async generator while doing nothing
-            yield
-        logger.warning(f"Connection type {self.__class__.__name__} does not support server-pushed events via events().")
-        return
-
-
-    @abstractmethod
-    async def request_and_download_stream(
-        self, 
-        stream_request_params: Dict[str, Any],
-        target_agent_id: str
-    ) -> AsyncIterator[bytes]: # pragma: no cover
-        """
-        Requests a stream download from the server and then streams the data.
-        This typically involves a two-step process:
-        1. An RPC call to initiate the stream download and get a download URL.
-        2. An HTTP GET request to the download URL to fetch the stream.
-
-        Args:
-            stream_request_params: Parameters for the REQUEST_STREAM_DOWNLOAD RPC call,
-                                   e.g., `{"resource_id": "some_file"}`.
-            target_agent_id: The ID of the target agent on the server that owns/provides the stream.
-                             This is used to construct the `target_agent_id` field in the RPC request params.
-
-        Yields:
-            bytes: Chunks of the downloaded data.
-
-        Raises:
-            ConnectionError: If not connected or if communication fails.
-            ValueError: If the server response for stream initiation is invalid.
-            NotImplementedError: If the connection type does not support this.
-        """
-        # Default implementation for connection types that don't support this.
-        if False: # Makes it an async generator
-            yield b'' 
-        logger.warning(f"Connection type {self.__class__.__name__} does not support request_and_download_stream().")
-        raise NotImplementedError(f"{self.__class__.__name__} does not support stream downloads.")
-
-
-    async def __aenter__(self):
-        await self.connect()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.close()
-
-    def __repr__(self) -> str:
-        return f"<{self.__class__.__name__} server_id='{self.server_id}', connected={self.is_connected}>"
-

autobyteus/rpc/client/client_connection_manager.py

@@ -1,153 +0,0 @@
-# file: autobyteus/autobyteus/rpc/client/client_connection_manager.py
-import logging
-from typing import Dict, Optional
-
-from autobyteus.utils.singleton import SingletonMeta
-from autobyteus.rpc.config import AgentServerConfig, AgentServerRegistry, default_agent_server_registry
-from autobyteus.rpc.transport_type import TransportType
-from .abstract_client_connection import AbstractClientConnection
-from .stdio_client_connection import StdioClientConnection
-from .sse_client_connection import SseClientConnection # Import SseClientConnection
-
-logger = logging.getLogger(__name__)
-
-class ClientConnectionManager(metaclass=SingletonMeta):
-    """
-    Manages the creation and caching of AbstractClientConnection instances.
-    Uses AgentServerRegistry to find server configurations.
-    """
-
-    def __init__(self, agent_server_registry: AgentServerRegistry = default_agent_server_registry):
-        self._agent_server_registry: AgentServerRegistry = agent_server_registry
-        self._active_connections: Dict[str, AbstractClientConnection] = {} # Cache connections by server_id
-        logger.info("ClientConnectionManager initialized.")
-
-    async def get_connection(self, server_config_id: str) -> AbstractClientConnection:
-        """
-        Gets or creates a client connection for the given server configuration ID.
-        If a connection for this ID already exists and is connected, it's returned.
-        Otherwise, a new connection is established.
-
-        Args:
-            server_config_id: The ID of the AgentServerConfig to connect to.
-
-        Returns:
-            An AbstractClientConnection instance.
-
-        Raises:
-            ValueError: If server_config_id is not found or config is invalid.
-            ConnectionError: If establishing a new connection fails.
-        """
-        if server_config_id in self._active_connections:
-            conn = self._active_connections[server_config_id]
-            if conn.is_connected: 
-                logger.debug(f"Returning cached and connected connection for '{server_config_id}'.")
-                return conn
-            else:
-                logger.info(f"Cached connection for '{server_config_id}' found but not connected. Attempting to recreate.")
-                await self.close_connection(server_config_id)
-
-
-        server_config: Optional[AgentServerConfig] = self._agent_server_registry.get_config(server_config_id)
-        if not server_config:
-            raise ValueError(f"AgentServerConfig not found in registry for server_id: '{server_config_id}'.")
-
-        connection: AbstractClientConnection
-        if server_config.transport_type == TransportType.STDIO:
-            connection = StdioClientConnection(server_config)
-        elif server_config.transport_type == TransportType.SSE:
-            connection = SseClientConnection(server_config) # Added SSE case
-        else:
-            raise NotImplementedError(f"Unsupported transport type '{server_config.transport_type}' for server_id '{server_config_id}'.")
-
-        try:
-            await connection.connect()
-            self._active_connections[server_config_id] = connection
-            logger.info(f"Successfully established new connection for server_id '{server_config_id}' type '{server_config.transport_type}'.")
-            return connection
-        except Exception as e:
-            logger.error(f"Failed to establish connection for server_id '{server_config_id}': {e}", exc_info=True)
-            await connection.close() # Ensure cleanup on failed connect
-            if server_config_id in self._active_connections:
-                 del self._active_connections[server_config_id]
-            raise ConnectionError(f"Failed to connect to server '{server_config_id}': {e}") from e
-
-
-    async def close_connection(self, server_config_id: str) -> None:
-        """
-        Closes and removes a specific connection from the manager.
-
-        Args:
-            server_config_id: The ID of the server connection to close.
-        """
-        connection = self._active_connections.pop(server_config_id, None)
-        if connection:
-            logger.info(f"Closing connection for server_id '{server_config_id}'.")
-            try:
-                await connection.close()
-            except Exception as e:
-                logger.error(f"Error closing connection for server_id '{server_config_id}': {e}", exc_info=True)
-        else:
-            logger.debug(f"No active connection found for server_id '{server_config_id}' to close.")
-
-    async def close_all_connections(self) -> None:
-        """Closes all active connections managed by this instance."""
-        logger.info(f"Closing all {len(self._active_connections)} active connections.")
-        for server_id in list(self._active_connections.keys()):
-            await self.close_connection(server_id)
-        logger.info("All active connections have been requested to close.")
-
-default_client_connection_manager = ClientConnectionManager()
-
-if __name__ == "__main__": # pragma: no cover
-    logging.basicConfig(level=logging.DEBUG)
-    
-    # Mock server configs
-    mock_stdio_id = "mock_stdio_ccm"
-    mock_stdio_cfg = AgentServerConfig(
-        server_id=mock_stdio_id,
-        transport_type=TransportType.STDIO,
-        stdio_command=["python", "-c", "import sys, time, json; print(json.dumps({'type':'response', 'id':'1', 'result':{'status':'ok'}}), flush=True); time.sleep(0.1); sys.exit(0)"]
-    )
-    default_agent_server_registry.register_config(mock_stdio_cfg)
-
-    mock_sse_id = "mock_sse_ccm"
-    mock_sse_cfg = AgentServerConfig(
-        server_id=mock_sse_id,
-        transport_type=TransportType.SSE,
-        sse_base_url="http://localhost:12345", # Dummy URL, server won't actually run here
-        sse_request_endpoint="/rpc",
-        sse_events_endpoint="/stream"
-    )
-    default_agent_server_registry.register_config(mock_sse_cfg)
-
-
-    manager = default_client_connection_manager
-
-    async def main_ccm_test():
-        stdio_conn = None
-        sse_conn = None
-        try:
-            logger.info(f"Testing CCM with Stdio: {mock_stdio_id}")
-            stdio_conn = await manager.get_connection(mock_stdio_id)
-            logger.info(f"Stdio connection: {stdio_conn}")
-            if stdio_conn.is_connected:
-                # Test send_request on stdio_conn
-                req = ProtocolMessage.create_request(method="test", id="1")
-                resp = await stdio_conn.send_request(req)
-                logger.info(f"Stdio test response: {resp}")
-
-
-            logger.info(f"Testing CCM with SSE: {mock_sse_id}")
-            sse_conn = await manager.get_connection(mock_sse_id) # This will create SseClientConnection and its session
-            logger.info(f"SSE connection: {sse_conn}")
-            # Actual SSE requests would fail here as no server is running at localhost:12345
-
-        except Exception as e:
-            logger.error(f"Error in CCM main_ccm_test: {e}", exc_info=True)
-        finally:
-            logger.info("Closing all connections via CCM.")
-            await manager.close_all_connections()
-            default_agent_server_registry.clear() # Clean up registry for other tests if any
-
-    asyncio.run(main_ccm_test())

autobyteus/rpc/client/sse_client_connection.py

@@ -1,306 +0,0 @@
-# file: autobyteus/autobyteus/rpc/client/sse_client_connection.py
-import asyncio
-import logging
-import json
-from typing import Optional, AsyncIterator, Dict, Any # Added Dict, Any
-
-import aiohttp
-from aiohttp_sse_client.client import EventSource # type: ignore
-
-from autobyteus.rpc.protocol import ProtocolMessage, MessageType, ErrorCode, EventType, RequestType, ResponseType # Added RequestType, ResponseType
-from .abstract_client_connection import AbstractClientConnection
-from autobyteus.rpc.config import AgentServerConfig, TransportType # Added TransportType from config
-
-logger = logging.getLogger(__name__)
-
-DEFAULT_SSE_TIMEOUT = 30.0  # seconds for an HTTP request response
-DEFAULT_SSE_RECONNECTION_TIME = 5 # seconds
-DEFAULT_STREAM_DOWNLOAD_TIMEOUT = 300.0 # 5 minutes for stream download requests
-
-class SseClientConnection(AbstractClientConnection):
-    """
-    Client connection implementation for SSE-based Agent Servers.
-    Uses aiohttp for HTTP requests and aiohttp-sse-client for consuming SSE event streams.
-    Also supports direct HTTP stream downloads.
-    """
-
-    def __init__(self, server_config: AgentServerConfig):
-        """
-        Initializes the SseClientConnection.
-
-        Args:
-            server_config: The configuration for the SSE server.
-
-        Raises:
-            ValueError: If server_config is not for SSE or base URL is missing.
-        """
-        if server_config.transport_type != TransportType.SSE: # Use Enum member
-            raise ValueError("SseClientConnection requires an AgentServerConfig with transport_type 'sse'.")
-        if not server_config.sse_base_url:
-            raise ValueError("AgentServerConfig for sse transport must have an sse_base_url.")
-
-        super().__init__(server_id=server_config.server_id)
-        self.server_config: AgentServerConfig = server_config
-        self._session: Optional[aiohttp.ClientSession] = None
-        self._event_source: Optional[EventSource] = None
-
-    async def connect(self) -> None:
-        """
-        Establishes the aiohttp ClientSession. For SSE, actual connection
-        to event stream happens when events() is iterated.
-        """
-        async with self._connection_lock:
-            if self._is_connected and self._session and not self._session.closed:
-                logger.debug(f"SseClientConnection to '{self.server_id}' session already active.")
-                return
-
-            try:
-                # Create a new session if it doesn't exist or is closed
-                if not self._session or self._session.closed:
-                    self._session = aiohttp.ClientSession()
-                    logger.info(f"SseClientConnection to '{self.server_id}': new aiohttp.ClientSession created.")
-                
-                self._is_connected = True # Mark as connected once session is ready
-                logger.info(f"SseClientConnection to '{self.server_id}' connected (session ready).")
-            except Exception as e:
-                logger.error(f"Failed to create/ensure aiohttp.ClientSession for '{self.server_id}': {e}", exc_info=True)
-                self._is_connected = False
-                if self._session and not self._session.closed:
-                    await self._session.close()
-                self._session = None
-                raise ConnectionError(f"Failed to initialize HTTP session for SSE server '{self.server_id}': {e}") from e
-
-    async def close(self) -> None:
-        """Closes the aiohttp ClientSession and any active EventSource."""
-        async with self._connection_lock:
-            if not self._is_connected and not self._session: # Check both state and session object
-                logger.debug(f"SseClientConnection to '{self.server_id}' already closed or never connected.")
-                return
-            
-            self._is_connected = False 
-
-            if self._event_source:
-                try:
-                    await self._event_source.close()
-                    logger.debug(f"EventSource for '{self.server_id}' closed.")
-                except Exception as e:
-                    logger.error(f"Error closing EventSource for '{self.server_id}': {e}", exc_info=True)
-                self._event_source = None
-
-            if self._session and not self._session.closed:
-                try:
-                    await self._session.close()
-                    logger.debug(f"aiohttp.ClientSession for '{self.server_id}' closed.")
-                except Exception as e:
-                    logger.error(f"Error closing aiohttp.ClientSession for '{self.server_id}': {e}", exc_info=True)
-            self._session = None # Ensure session object is cleared
-            
-            logger.info(f"SseClientConnection to '{self.server_id}' closed.")
-
-    async def send_request(self, request_message: ProtocolMessage) -> ProtocolMessage:
-        """Sends a request via HTTP POST and awaits a JSON response."""
-        if not self._is_connected or not self._session or self._session.closed:
-            logger.info(f"SseClientConnection '{self.server_id}' not connected or session closed/missing. Attempting to connect before send_request.")
-            await self.connect() 
-            if not self._is_connected or not self._session or self._session.closed: # Check again after connect attempt
-                raise ConnectionError(f"Failed to connect to SSE server '{self.server_id}' for send_request (connection attempt failed).")
-
-        if request_message.type != MessageType.REQUEST:
-            raise ValueError("ProtocolMessage must be of type REQUEST to be sent via send_request.")
-        if not request_message.id:
-            raise ValueError("Request ProtocolMessage must have an ID.")
-
-        request_url = self.server_config.get_sse_full_request_url()
-        if not request_url: 
-            raise ValueError("SSE request URL is not configured.")
-
-        logger.debug(f"SseClient '{self.server_id}' sending request to {request_url}: {request_message.to_json_str()}")
-        
-        try:
-            async with self._session.post(
-                request_url,
-                json=request_message.model_dump(exclude_none=True), 
-                timeout=DEFAULT_SSE_TIMEOUT
-            ) as response:
-                response_text = await response.text()
-                if response.status >= 200 and response.status < 300:
-                    try:
-                        return ProtocolMessage.from_json_str(response_text)
-                    except (json.JSONDecodeError, ValueError) as e: 
-                        logger.error(f"Error decoding JSON response from '{request_url}': {e}. Response text: {response_text[:200]}")
-                        return ProtocolMessage.create_error_response(
-                            id=request_message.id,
-                            code=ErrorCode.PARSE_ERROR,
-                            message=f"Failed to parse JSON response from server: {e}"
-                        )
-                else:
-                    logger.error(f"HTTP error from '{request_url}': {response.status} {response.reason}. Response: {response_text[:200]}")
-                    return ProtocolMessage.create_error_response(
-                        id=request_message.id,
-                        code=ErrorCode.INTERNAL_ERROR, 
-                        message=f"Server returned HTTP error {response.status}: {response.reason}. Body: {response_text[:200]}"
-                    )
-        except aiohttp.ClientConnectorError as e:
-            logger.error(f"Connection error for '{request_url}': {e}", exc_info=True)
-            await self.close() # Connection is likely broken, perform full close
-            raise ConnectionError(f"Could not connect to SSE server at '{request_url}': {e}") from e
-        except asyncio.TimeoutError:
-            logger.warning(f"Timeout sending request to '{request_url}' for ID '{request_message.id}'.")
-            # Don't necessarily close connection on timeout, server might be slow.
-            return ProtocolMessage.create_error_response(
-                id=request_message.id,
-                code=ErrorCode.SERVER_ERROR_TIMEOUT,
-                message=f"Timeout waiting for response from SSE server for request ID '{request_message.id}'."
-            )
-        except Exception as e:
-            logger.error(f"Unexpected error sending request to '{request_url}': {e}", exc_info=True)
-            if isinstance(e, (aiohttp.ClientError)): 
-                 await self.close() # If it's a client lib error, connection might be compromised
-            return ProtocolMessage.create_error_response(
-                id=request_message.id,
-                code=ErrorCode.INTERNAL_ERROR,
-                message=f"Client-side error sending request: {e}"
-            )
-
-    async def events(self) -> AsyncIterator[ProtocolMessage]:
-        """Connects to the SSE event stream and yields ProtocolMessages."""
-        if not self._is_connected or not self._session or self._session.closed:
-            logger.info(f"SseClientConnection '{self.server_id}' session not found or closed, attempting to connect before streaming events.")
-            await self.connect()
-            if not self._is_connected or not self._session or self._session.closed: # Check again
-                 raise ConnectionError(f"Not connected to SSE server '{self.server_id}' for event streaming (connection attempt failed).")
-
-        # Construct event URL. Note: SseServerHandler expects /events/{agent_id_on_server}
-        # The client config's server_id typically maps to an agent_id_on_server or is the agent_id_on_server.
-        # For multi-agent servers, the client needs to know which agent_id_on_server to subscribe to.
-        # This implies server_config.server_id is the agent_id_on_server for SSE event subscriptions.
-        base_events_url = self.server_config.get_sse_full_events_url()
-        if not base_events_url: # This is base path like /events
-             raise ValueError("SSE events base URL path is not configured.")
-        
-        # Assuming self.server_id (from config) is the key for the agent on the server.
-        events_url = f"{base_events_url.rstrip('/')}/{self.server_id}"
-        logger.info(f"SseClient '{self.server_id}' connecting to event stream at {events_url}")
-        
-        if self._event_source: # Close existing if any
-            await self._event_source.close()
-            self._event_source = None
-
-        try:
-            self._event_source = EventSource(
-                events_url,
-                session=self._session, # Use existing session
-                reconnection_time=DEFAULT_SSE_RECONNECTION_TIME,
-            )
-            await self._event_source.connect() 
-
-            async for sse_event in self._event_source:
-                try:
-                    logger.debug(f"SseClient '{self.server_id}' received SSE event data: {sse_event.data[:200]}")
-                    msg = ProtocolMessage.from_json_str(sse_event.data)
-                    if msg.type == MessageType.EVENT:
-                        yield msg
-                    else:
-                        logger.warning(f"Received non-EVENT ProtocolMessage via SSE stream: {msg.type}. Ignoring.")
-                except json.JSONDecodeError as e:
-                    logger.error(f"Error decoding JSON from SSE event data: {e}. Data: {sse_event.data[:200]}")
-                except ValueError as e: 
-                    logger.error(f"Error validating ProtocolMessage from SSE event data: {e}. Data: {sse_event.data[:200]}")
-                except Exception as e:
-                    logger.error(f"Error processing SSE event: {e}", exc_info=True)
-
-        except ConnectionRefusedError as e:
-            logger.error(f"SSE EventStream for '{self.server_id}' connection refused at {events_url}: {e}")
-            await self.close()
-            raise ConnectionError(f"SSE EventStream connection refused: {e}") from e
-        except aiohttp.ClientError as e: 
-            logger.error(f"SSE EventStream client error for '{self.server_id}' at {events_url}: {e}", exc_info=True)
-            await self.close() 
-            raise ConnectionError(f"SSE EventStream client error: {e}") from e
-        except asyncio.CancelledError:
-            logger.info(f"SSE event stream for '{self.server_id}' cancelled.")
-            if self._event_source: await self._event_source.close()
-            raise
-        except Exception as e:
-            logger.error(f"Unexpected error in SSE event stream for '{self.server_id}': {e}", exc_info=True)
-            if self._event_source: await self._event_source.close()
-            await self.close()
-            raise ConnectionError(f"Unexpected error in SSE event stream: {e}") from e
-        finally:
-            logger.info(f"SseClient '{self.server_id}' event stream iteration finished.")
-            if self._event_source: 
-                await self._event_source.close()
-                self._event_source = None
-
-
-    async def request_and_download_stream(
-        self, 
-        stream_request_params: Dict[str, Any],
-        target_agent_id: str # This is the agent_id_on_server key
-    ) -> AsyncIterator[bytes]:
-        """
-        Requests a stream download via RPC and then downloads the stream via HTTP GET.
-        """
-        if not self._is_connected or not self._session or self._session.closed:
-            logger.info(f"SseClientConnection '{self.server_id}' not connected or session closed/missing. Attempting to connect before stream download.")
-            await self.connect()
-            if not self._is_connected or not self._session or self._session.closed:
-                raise ConnectionError(f"Failed to connect to SSE server '{self.server_id}' for stream download (connection attempt failed).")
-
-        # Step 1: RPC call to initiate stream download
-        # The target_agent_id for the stream must be included in the params for the RPC call.
-        rpc_params = {**stream_request_params, "target_agent_id": target_agent_id}
-        initiate_request_msg = ProtocolMessage.create_request(
-            method=RequestType.REQUEST_STREAM_DOWNLOAD,
-            params=rpc_params
-        )
-        
-        logger.debug(f"SseClient '{self.server_id}' sending stream download initiation request for target_agent_id '{target_agent_id}': {initiate_request_msg.to_json_str()}")
-        
-        initiate_response_msg = await self.send_request(initiate_request_msg)
-
-        if initiate_response_msg.type == MessageType.ERROR or not initiate_response_msg.result:
-            err_details = initiate_response_msg.error.message if initiate_response_msg.error else "Unknown error"
-            logger.error(f"Failed to initiate stream download for target_agent_id '{target_agent_id}'. Server error: {err_details}")
-            raise ValueError(f"Server error initiating stream download: {err_details}")
-
-        if initiate_response_msg.response_type != ResponseType.STREAM_DOWNLOAD_READY:
-            logger.error(f"Unexpected response type from stream download initiation: {initiate_response_msg.response_type}")
-            raise ValueError(f"Unexpected response type: {initiate_response_msg.response_type}")
-
-        download_url = initiate_response_msg.result.get("download_url")
-        if not download_url:
-            logger.error(f"No 'download_url' in STREAM_DOWNLOAD_READY response. Result: {initiate_response_msg.result}")
-            raise ValueError("Server response did not include a download_url for the stream.")
-
-        stream_metadata = initiate_response_msg.result.get("metadata", {})
-        logger.info(f"SseClient '{self.server_id}' received download URL: {download_url}. Metadata: {stream_metadata}")
-
-        # Step 2: HTTP GET request to the download_url
-        try:
-            logger.debug(f"SseClient '{self.server_id}' starting GET request to stream from {download_url}")
-            async with self._session.get(download_url, timeout=DEFAULT_STREAM_DOWNLOAD_TIMEOUT) as response:
-                response.raise_for_status() # Raise an exception for HTTP error codes (4xx or 5xx)
-                
-                # Stream the content
-                async for chunk in response.content.iter_any(): # iter_any() or iter_chunked(chunk_size)
-                    yield chunk
-                logger.info(f"SseClient '{self.server_id}' finished streaming from {download_url}")
-
-        except aiohttp.ClientResponseError as e:
-            logger.error(f"HTTP error during stream download from '{download_url}': {e.status} {e.message}", exc_info=True)
-            # Connection might still be usable for other RPCs, don't necessarily close.
-            raise ConnectionError(f"HTTP error downloading stream: {e.status} {e.message}") from e
-        except aiohttp.ClientConnectorError as e:
-            logger.error(f"Connection error during stream download from '{download_url}': {e}", exc_info=True)
-            await self.close() # Connection likely broken
-            raise ConnectionError(f"Could not connect to download stream at '{download_url}': {e}") from e
-        except asyncio.TimeoutError:
-            logger.warning(f"Timeout downloading stream from '{download_url}'.")
-            # Don't close connection, server might be slow or stream very long with no data.
-            raise TimeoutError(f"Timeout downloading stream from '{download_url}'.")
-        except Exception as e:
-            logger.error(f"Unexpected error downloading stream from '{download_url}': {e}", exc_info=True)
-            if isinstance(e, (aiohttp.ClientError)):
-                 await self.close() # If it's a client lib error, connection might be compromised
-            raise ConnectionError(f"Client-side error downloading stream: {e}") from e