autobyteus 1.2.0__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/rpc/server/method_handlers.py

@@ -1,259 +0,0 @@
-# file: autobyteus/autobyteus/rpc/server/method_handlers.py
-import logging
-from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple
-
-from autobyteus.rpc.protocol import ProtocolMessage, MessageType, RequestType, ResponseType, ErrorCode
-from autobyteus.rpc.server.base_method_handler import BaseMethodHandler
-from autobyteus.agent.agent import Agent 
-from autobyteus.agent.message.agent_input_user_message import AgentInputUserMessage
-from autobyteus.agent.message.inter_agent_message import InterAgentMessage 
-
-
-logger = logging.getLogger(__name__)
-
-class DiscoverCapabilitiesHandler(BaseMethodHandler):
-    """
-    Handles the 'discover_capabilities' RPC method.
-    Responds with the agent's ID, definition details, supported methods, and status,
-    obtained via the Agent's public interface/context.
-    """
-    async def handle(self,
-                     request_id: Optional[str],
-                     params: Optional[Dict[str, Any]],
-                     agent: Agent) -> ProtocolMessage: 
-        logger.debug(f"DiscoverCapabilitiesHandler: Handling request_id '{request_id}' for agent '{agent.agent_id}'.")
-        try:
-            # Define capabilities, including the new stream download initiation method
-            capabilities = {
-                "post_user_message": {
-                    "description": "Posts a message from an external user to the agent.",
-                    "params": {"agent_input_user_message": "dict (serialized AgentInputUserMessage)"}
-                },
-                "post_inter_agent_message": {
-                    "description": "Posts a message from another agent.",
-                    "params": {"inter_agent_message": "dict (serialized InterAgentMessage)"}
-                },
-                "post_tool_execution_approval": {
-                    "description": "Provides approval or denial for a pending tool execution.",
-                    "params": {"tool_invocation_id": "str", "is_approved": "bool", "reason": "Optional[str]"}
-                },
-                "get_status": {
-                    "description": "Gets the current operational status of the agent.",
-                    "params": None 
-                },
-                # No specific invoke method for REQUEST_STREAM_DOWNLOAD, it's a top-level RequestType handled by its own handler.
-                # However, if it were invoked via InvokeMethod, it would be listed here.
-                # For clarity, methods invokable via InvokeMethodHandler are listed.
-                # REQUEST_STREAM_DOWNLOAD will be handled by InitiateStreamDownloadHandler directly.
-            }
-            
-            if not agent.context: 
-                 logger.error(f"Agent '{agent.agent_id}' context is None. Cannot discover capabilities.")
-                 return ProtocolMessage.create_error_response(
-                    id=request_id,
-                    code=ErrorCode.SERVER_ERROR_CAPABILITY_DISCOVERY_FAILED,
-                    message="Agent context not available."
-                 )
-
-            # Explicitly list RequestTypes that this server (potentially) handles
-            # This part of discover_capabilities might need to be more dynamic or list top-level RPC methods
-            # rather than just methods invokable via InvokeMethodHandler.
-            # For now, assume InvokeMethodHandler is the primary way to call agent functions.
-            # The fact that REQUEST_STREAM_DOWNLOAD exists is a server capability.
-            supported_rpc_request_types = [rt.value for rt in RequestType]
-
-
-            result_data = {
-                "agent_id": agent.agent_id, 
-                "agent_name": agent.context.definition.name,
-                "agent_role": agent.context.definition.role,
-                "agent_description": agent.context.definition.description,
-                "supported_rpc_request_types": supported_rpc_request_types, # Top-level RPC methods
-                "invokable_methods_details": capabilities, # Methods callable via 'invoke_method'
-                "status": agent.get_status().value, 
-                "system_prompt_summary": agent.context.definition.system_prompt[:200] + "..." if len(agent.context.definition.system_prompt) > 200 else agent.context.definition.system_prompt,
-                "tool_names": agent.context.definition.tool_names,
-            }
-            return ProtocolMessage.create_response(
-                id=request_id,
-                result=result_data,
-                response_type=ResponseType.CAPABILITIES_RESPONSE
-            )
-        except Exception as e:
-            logger.error(f"Error in DiscoverCapabilitiesHandler for agent '{agent.agent_id}': {e}", exc_info=True)
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.SERVER_ERROR_CAPABILITY_DISCOVERY_FAILED,
-                message=f"Failed to discover capabilities: {e}"
-            )
-
-
-class InvokeMethodHandler(BaseMethodHandler):
-    """
-    Handles the 'invoke_method' RPC method.
-    It dispatches to public methods of the Agent instance based on 'method_name'.
-    """
-
-    def __init__(self):
-        # Map method names to Agent's public API methods (or wrappers around them)
-        self._method_map: Dict[str, Callable[[Optional[str], Dict[str, Any], Agent], Awaitable[ProtocolMessage]]] = {
-            "post_user_message": self._handle_post_user_message,
-            "post_inter_agent_message": self._handle_post_inter_agent_message,
-            "post_tool_execution_approval": self._handle_post_tool_execution_approval,
-            "get_status": self._handle_get_status,
-        }
-
-    async def handle(self,
-                     request_id: Optional[str],
-                     params: Optional[Dict[str, Any]],
-                     agent: Agent) -> ProtocolMessage: 
-        if not params or "method_name" not in params:
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.INVALID_PARAMS,
-                message="'method_name' is required in params for invoke_method."
-            )
-
-        method_name = params["method_name"]
-        method_params = params.get("method_params", {})
-
-        logger.debug(f"InvokeMethodHandler: Handling method '{method_name}' for agent '{agent.agent_id}' with request_id '{request_id}'.")
-
-        handler_func = self._method_map.get(method_name)
-        if not handler_func:
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.METHOD_NOT_FOUND,
-                message=f"Method '{method_name}' is not supported by this agent server for invocation via InvokeMethodHandler."
-            )
-
-        try:
-            return await handler_func(request_id, method_params, agent)
-        except Exception as e:
-            logger.error(f"Error invoking method '{method_name}' on agent '{agent.agent_id}': {e}", exc_info=True)
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.SERVER_ERROR_AGENT_PROCESSING_FAILED,
-                message=f"Error processing method '{method_name}': {e}"
-            )
-
-    async def _handle_post_user_message(self, request_id: Optional[str], params: Dict[str, Any], agent: Agent) -> ProtocolMessage:
-        serialized_msg = params.get("agent_input_user_message")
-        if not isinstance(serialized_msg, dict):
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, "'agent_input_user_message' (dict) parameter is required.")
-        
-        try:
-            user_message = AgentInputUserMessage.from_dict(serialized_msg)
-        except (ValueError, TypeError) as e:
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, f"Invalid 'agent_input_user_message' structure: {e}")
-
-        await agent.post_user_message(user_message) 
-        logger.info(f"Agent '{agent.agent_id}' (via RPC): Called agent.post_user_message().")
-        return ProtocolMessage.create_response(request_id, {"status": "User message posted to agent"}, ResponseType.ACKNOWLEDGEMENT)
-
-    async def _handle_post_inter_agent_message(self, request_id: Optional[str], params: Dict[str, Any], agent: Agent) -> ProtocolMessage:
-        serialized_msg = params.get("inter_agent_message")
-        if not isinstance(serialized_msg, dict):
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, "'inter_agent_message' (dict) parameter is required.")
-
-        try:
-            msg_type_str = serialized_msg.get("message_type")
-            if not msg_type_str: raise ValueError("message_type missing")
-            
-            inter_agent_msg = InterAgentMessage.create_with_dynamic_message_type(
-                recipient_role_name=serialized_msg.get("recipient_role_name"),
-                recipient_agent_id=serialized_msg.get("recipient_agent_id"),
-                content=serialized_msg.get("content"),
-                message_type=msg_type_str,
-                sender_agent_id=serialized_msg.get("sender_agent_id")
-            )
-        except (ValueError, TypeError) as e:
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, f"Invalid 'inter_agent_message' structure: {e}")
-
-        await agent.post_inter_agent_message(inter_agent_msg) 
-        logger.info(f"Agent '{agent.agent_id}' (via RPC): Called agent.post_inter_agent_message() from '{inter_agent_msg.sender_agent_id}'.")
-        return ProtocolMessage.create_response(request_id, {"status": "Inter-agent message posted to agent"}, ResponseType.ACKNOWLEDGEMENT)
-
-    async def _handle_post_tool_execution_approval(self, request_id: Optional[str], params: Dict[str, Any], agent: Agent) -> ProtocolMessage:
-        tool_invocation_id = params.get("tool_invocation_id")
-        is_approved = params.get("is_approved")
-        reason = params.get("reason")
-
-        if not isinstance(tool_invocation_id, str) or not tool_invocation_id:
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, "'tool_invocation_id' (str) is required.")
-        if not isinstance(is_approved, bool):
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, "'is_approved' (bool) is required.")
-        if reason is not None and not isinstance(reason, str):
-            return ProtocolMessage.create_error_response(request_id, ErrorCode.INVALID_PARAMS, "'reason' must be a string if provided.")
-
-        await agent.post_tool_execution_approval(tool_invocation_id, is_approved, reason) 
-        status_str = "approved" if is_approved else "denied"
-        logger.info(f"Agent '{agent.agent_id}' (via RPC): Called agent.post_tool_execution_approval() for id '{tool_invocation_id}' ({status_str}).")
-        return ProtocolMessage.create_response(request_id, {"status": f"Tool approval ({status_str}) posted to agent"}, ResponseType.ACKNOWLEDGEMENT)
-
-    async def _handle_get_status(self, request_id: Optional[str], params: Dict[str, Any], agent: Agent) -> ProtocolMessage:
-        status_value = agent.get_status().value 
-        logger.debug(f"Agent '{agent.agent_id}' (via RPC): Current status is '{status_value}'.")
-        return ProtocolMessage.create_response(request_id, {"status": status_value, "agent_id": agent.agent_id})
-
-
-class InitiateStreamDownloadHandler(BaseMethodHandler):
-    """
-    Handles the 'request_stream_download' RPC method.
-    Interacts with the agent to prepare a streamable resource and returns
-    information required by the client to download it, including a stream_id.
-    The SseServerHandler will augment this response with the full download URL.
-    """
-    async def handle(self,
-                     request_id: Optional[str],
-                     params: Optional[Dict[str, Any]],
-                     agent: Agent) -> ProtocolMessage:
-        logger.debug(f"InitiateStreamDownloadHandler: Handling request_id '{request_id}' for agent '{agent.agent_id}' with params: {params}.")
-        if not params:
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.INVALID_PARAMS,
-                message="Parameters are required for requesting a stream download (e.g., resource identifier)."
-            )
-
-        try:
-            # Agent must implement `prepare_resource_for_streaming`
-            # This method is conceptual; actual signature might vary.
-            # It should return (stream_id: str, metadata: Dict[str, Any])
-            # The agent becomes responsible for managing the lifecycle of this stream_id
-            # and providing data when get_stream_data(stream_id) is called by SseServerHandler.
-            if not hasattr(agent, "prepare_resource_for_streaming"):
-                logger.error(f"Agent '{agent.agent_id}' does not support 'prepare_resource_for_streaming'.")
-                return ProtocolMessage.create_error_response(
-                    id=request_id,
-                    code=ErrorCode.METHOD_NOT_FOUND,
-                    message=f"Agent '{agent.agent_id}' cannot prepare streamable resources."
-                )
-            
-            # Type hint for clarity if agent method is known
-            # stream_id, metadata = await agent.prepare_resource_for_streaming(params)
-            stream_preparation_result: Tuple[str, Dict[str, Any]] = await agent.prepare_resource_for_streaming(params)
-            stream_id, metadata = stream_preparation_result
-            
-            logger.info(f"Agent '{agent.agent_id}' prepared stream_id '{stream_id}' for request_id '{request_id}'.")
-
-            # The SseServerHandler will add the 'download_url' to this result later
-            # agent_id_on_server is also added by SseServerHandler for constructing the URL
-            result_payload = {
-                "stream_id": stream_id,
-                "metadata": metadata,
-                # "agent_id_on_server": agent.agent_id # Or the server_key if different. SseServerHandler knows this.
-            }
-            
-            return ProtocolMessage.create_response(
-                id=request_id,
-                result=result_payload,
-                response_type=ResponseType.STREAM_DOWNLOAD_READY
-            )
-        except Exception as e:
-            logger.error(f"Error in InitiateStreamDownloadHandler for agent '{agent.agent_id}': {e}", exc_info=True)
-            return ProtocolMessage.create_error_response(
-                id=request_id,
-                code=ErrorCode.SERVER_ERROR_STREAM_PREPARATION_FAILED,
-                message=f"Agent failed to prepare stream: {e}"
-            )

autobyteus/rpc/server/sse_server_handler.py

@@ -1,182 +0,0 @@
-    # file: autobyteus/autobyteus/rpc/server/sse_server_handler.py
-    import asyncio
-    import logging
-    import json
-    from typing import Dict, Optional, Union, Set, cast, AsyncIterator 
-    from weakref import WeakSet
-    
-    from aiohttp import web 
-    from aiohttp_sse import sse_response 
-    
-    from autobyteus.agent.agent import Agent
-    # MODIFIED: Import AgentEventStream (remains the same, but it's now the sole class)
-    from autobyteus.agent.streaming import AgentEventStream 
-    from autobyteus.agent.streaming.stream_events import StreamEvent, StreamEventType 
-    from autobyteus.rpc.protocol import ProtocolMessage, MessageType, ErrorCode, RequestType, ResponseType, EventType as RPCEventType 
-    from autobyteus.rpc.server.base_method_handler import BaseMethodHandler
-    from autobyteus.rpc.config import AgentServerConfig 
-    
-    logger = logging.getLogger(__name__)
-    
-    DEFAULT_STREAM_CHUNK_SIZE = 8192 
-    
-    class SseServerHandler:
-        def __init__(self, agents: Dict[str, Agent], method_handlers: Dict[Union[RequestType, str], BaseMethodHandler]):
-            self._agents: Dict[str, Agent] = agents
-            self._method_handlers = method_handlers
-            self._app = web.Application()
-            self._runner: Optional[web.AppRunner] = None
-            self._site: Optional[web.TCPSite] = None
-            self._active_sse_forwarding_tasks: Dict[web.StreamResponse, asyncio.Task] = {}
-            logger.info(f"SseServerHandler initialized to serve agents: {list(self._agents.keys())}.")
-    
-        def _setup_routes(self, config: AgentServerConfig):
-            # ... (method body remains the same) ...
-            rpc_request_path = config.sse_request_endpoint
-            base_events_path = config.sse_events_endpoint.rstrip('/')
-            full_events_path = f"{base_events_path}/{{agent_id_on_server}}"
-            stream_download_prefix = config.sse_stream_download_path_prefix.rstrip('/')
-            full_stream_download_path = f"{stream_download_prefix}/{{agent_id_on_server}}/{{stream_id}}"
-            self._app.router.add_post(rpc_request_path, self.handle_rpc_request)
-            self._app.router.add_get(full_events_path, self.handle_sse_events_subscription)
-            self._app.router.add_get(full_stream_download_path, self.handle_http_stream_download)
-            logger.info(f"SseServerHandler routes: POST {rpc_request_path} (RPC), GET {full_events_path} (SSE per agent), GET {full_stream_download_path} (HTTP Stream Download).")
-
-        async def handle_rpc_request(self, http_request: web.Request) -> web.Response:
-            # ... (method body remains the same) ...
-            request_id: Optional[str] = None; raw_body_str: str = ""; target_agent_id: Optional[str] = None
-            try:
-                raw_body = await http_request.read(); raw_body_str = raw_body.decode()
-                try: parsed_json_early = json.loads(raw_body_str); request_id = parsed_json_early.get("id")
-                except json.JSONDecodeError: pass
-                request_message = ProtocolMessage.from_json_str(raw_body_str); request_id = request_message.id 
-                if request_message.type != MessageType.REQUEST or not request_message.method:
-                    return self._create_json_error_response(request_id, ErrorCode.INVALID_REQUEST, "Must be REQUEST with method.", 400)
-                if request_message.params and "target_agent_id" in request_message.params:
-                    target_agent_id = str(request_message.params["target_agent_id"])
-                if not target_agent_id: return self._create_json_error_response(request_id, ErrorCode.INVALID_PARAMS, "'target_agent_id' missing in request params.", 400)
-                target_agent = self._agents.get(target_agent_id)
-                if not target_agent: return self._create_json_error_response(request_id, ErrorCode.METHOD_NOT_FOUND, f"Agent with id '{target_agent_id}' not found.", 404)
-                actual_handler_params = request_message.params
-                handler = self._method_handlers.get(request_message.method)
-                if not handler: return self._create_json_error_response(request_id, ErrorCode.METHOD_NOT_FOUND, f"RPC Method '{request_message.method}' not found.", 404)
-                logger.debug(f"SseServerHandler dispatching RPC '{request_message.method}' (ReqID: {request_id}) to '{handler.__class__.__name__}' for agent '{target_agent_id}'.")
-                response_proto = await handler.handle(request_id, actual_handler_params, target_agent) 
-                if response_proto.response_type == ResponseType.STREAM_DOWNLOAD_READY and response_proto.result and "stream_id" in response_proto.result and target_agent_id: 
-                    server_config: AgentServerConfig = http_request.app['agent_server_config']
-                    full_url_prefix = server_config.get_sse_full_stream_download_url_prefix_for_agent(target_agent_id)
-                    if full_url_prefix:
-                        stream_id = response_proto.result["stream_id"]; download_url = f"{full_url_prefix.rstrip('/')}/{stream_id}"
-                        response_proto.result["download_url"] = download_url; logger.info(f"Constructed download URL for stream_id '{stream_id}': {download_url}")
-                    else: logger.error(f"Could not construct download_url for stream_id '{response_proto.result['stream_id']}'.")
-                http_status = 200
-                if response_proto.type == MessageType.ERROR and response_proto.error:
-                    if response_proto.error.code == ErrorCode.METHOD_NOT_FOUND.value: http_status = 404
-                    elif response_proto.error.code in [ErrorCode.INVALID_REQUEST.value, ErrorCode.INVALID_PARAMS.value, ErrorCode.PARSE_ERROR.value]: http_status = 400
-                    else: http_status = 500
-                return web.json_response(response_proto.model_dump(exclude_none=True), status=http_status)
-            except json.JSONDecodeError as e: logger.error(f"Sse JSONDecodeError: {e}. Body: '{raw_body_str[:200]}'"); return self._create_json_error_response(request_id, ErrorCode.PARSE_ERROR, f"JSON parse error: {e}", 400)
-            except ValueError as e: logger.error(f"Sse Protocol validation error: {e}. Body: '{raw_body_str[:200]}'"); return self._create_json_error_response(request_id, ErrorCode.INVALID_REQUEST, f"Invalid request: {e}", 400)
-            except Exception as e: logger.error(f"Sse unexpected error: {e}", exc_info=True); return self._create_json_error_response(request_id, ErrorCode.INTERNAL_ERROR, f"Server error: {e}", 500)
-
-        def _create_json_error_response(self, req_id: Optional[str], code: ErrorCode, msg: str, http_status: int) -> web.Response:
-            # ... (method body remains the same) ...
-            err_proto = ProtocolMessage.create_error_response(req_id, code, msg)
-            return web.json_response(err_proto.model_dump(exclude_none=True), status=http_status)
-    
-        async def handle_sse_events_subscription(self, http_request: web.Request) -> web.StreamResponse:
-            # ... (method body remains the same, it instantiates AgentEventStream and calls all_events()) ...
-            agent_id_on_server = http_request.match_info.get("agent_id_on_server")
-            if not agent_id_on_server: raise web.HTTPBadRequest(text="'agent_id_on_server' path param required.")
-            target_agent = self._agents.get(agent_id_on_server)
-            if not target_agent: raise web.HTTPNotFound(text=f"Agent '{agent_id_on_server}' not found.")
-            client_addr = http_request.remote
-            logger.info(f"SSE client {client_addr} subscribing to events for agent '{target_agent.agent_id}' (server key: '{agent_id_on_server}').")
-            sse_resp = web.StreamResponse(status=200, reason='OK', headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'})
-            await sse_resp.prepare(http_request) 
-            forwarding_task = asyncio.create_task(self._stream_agent_events_to_client(sse_resp, target_agent, agent_id_on_server), name=f"sse_fwd_{target_agent.agent_id}_{client_addr}")
-            self._active_sse_forwarding_tasks[sse_resp] = forwarding_task
-            try: await forwarding_task 
-            except asyncio.CancelledError: logger.info(f"SSE task for agent '{target_agent.agent_id}' to {client_addr} cancelled.")
-            except Exception as e: logger.error(f"Error in SSE streaming for agent '{target_agent.agent_id}' to {client_addr}: {e}", exc_info=True)
-            finally: logger.info(f"SSE client {client_addr} for agent '{target_agent.agent_id}' disconnected."); self._active_sse_forwarding_tasks.pop(sse_resp, None)
-            return sse_resp
-
-        async def _stream_agent_events_to_client(self, sse_client_resp: web.StreamResponse, agent: Agent, agent_id_on_server: str):
-            # Instantiates AgentEventStream
-            event_stream_provider = AgentEventStream(agent)
-            logger.debug(f"SseServerHandler: Streaming events from agent '{agent.agent_id}' (key: {agent_id_on_server}) via AgentEventStream.all_events().")
-            try:
-                # Calls .all_events() method
-                async for agent_event_obj in event_stream_provider.all_events(): 
-                    if sse_client_resp.closed: break
-                    rpc_event_type: RPCEventType
-                    if agent_event_obj.event_type == StreamEventType.ASSISTANT_CHUNK: rpc_event_type = RPCEventType.AGENT_OUTPUT_CHUNK
-                    elif agent_event_obj.event_type == StreamEventType.ASSISTANT_FINAL_MESSAGE: rpc_event_type = RPCEventType.AGENT_FINAL_MESSAGE
-                    elif agent_event_obj.event_type == StreamEventType.TOOL_INTERACTION_LOG_ENTRY: rpc_event_type = RPCEventType.TOOL_LOG_ENTRY
-                    elif agent_event_obj.event_type == StreamEventType.AGENT_STATUS_CHANGE: rpc_event_type = RPCEventType.AGENT_STATUS_UPDATE
-                    elif agent_event_obj.event_type == StreamEventType.ERROR_EVENT: rpc_event_type = RPCEventType.AGENT_STATUS_UPDATE; logger.error(f"Unified stream error for agent '{agent.agent_id}': {agent_event_obj.data}")
-                    else: rpc_event_type = RPCEventType.AGENT_STATUS_UPDATE; logger.warning(f"Unhandled StreamEventType '{agent_event_obj.event_type}' from agent '{agent.agent_id}'.")
-                    payload_data = {"agent_id_on_server": agent_id_on_server, **agent_event_obj.data}
-                    protocol_event_msg = ProtocolMessage.create_event(event_type=rpc_event_type, payload=payload_data)
-                    try:
-                        sse_event_str = f"event: {str(protocol_event_msg.event_type.value)}\ndata: {protocol_event_msg.to_json_str()}\n\n"
-                        await sse_client_resp.write(sse_event_str.encode('utf-8'))
-                    except ConnectionResetError: logger.info(f"SSE client reset for agent '{agent.agent_id}'."); break
-                    except Exception as send_e: logger.error(f"Error sending SSE event for agent '{agent.agent_id}': {send_e}"); break 
-                if not sse_client_resp.closed: await sse_client_resp.write_eof()
-            except asyncio.CancelledError: logger.info(f"Event streaming for agent '{agent.agent_id}' (key: {agent_id_on_server}) cancelled.")
-            except Exception as e: logger.error(f"Error in event streaming for agent '{agent.agent_id}' (key: {agent_id_on_server}): {e}", exc_info=True)
-            finally: logger.debug(f"Finished event streaming from agent '{agent.agent_id}' (key: {agent_id_on_server}).")
-
-        async def handle_http_stream_download(self, http_request: web.Request) -> web.StreamResponse:
-            # ... (method body remains the same) ...
-            agent_id_on_server = http_request.match_info.get("agent_id_on_server"); stream_id = http_request.match_info.get("stream_id")
-            if not agent_id_on_server or not stream_id: raise web.HTTPBadRequest(text="'agent_id_on_server' and 'stream_id' path params required.")
-            target_agent = self._agents.get(agent_id_on_server)
-            if not target_agent: raise web.HTTPNotFound(text=f"Agent '{agent_id_on_server}' not found.")
-            logger.info(f"HTTP stream download for agent '{target_agent.agent_id}' (key: '{agent_id_on_server}'), stream_id '{stream_id}'.")
-            if not hasattr(target_agent, "get_stream_data") or not hasattr(target_agent, "cleanup_stream_resource"):
-                raise web.HTTPNotImplemented(text=f"Agent '{target_agent.agent_id}' does not support stream retrieval.")
-            try:
-                get_stream_data_method = getattr(target_agent, "get_stream_data")
-                data_iterator: AsyncIterator[bytes] = await get_stream_data_method(stream_id)
-                response = web.StreamResponse(status=200, reason="OK", headers={"Content-Type": "application/octet-stream"})
-                await response.prepare(http_request)
-                async for chunk in data_iterator:
-                    if not isinstance(chunk, bytes): logger.error(f"Agent '{target_agent.agent_id}' stream '{stream_id}' yielded non-bytes: {type(chunk)}."); break 
-                    await response.write(chunk); await asyncio.sleep(0) 
-                await response.write_eof(); logger.info(f"Streamed data for agent '{target_agent.agent_id}', stream_id '{stream_id}'."); return response
-            except FileNotFoundError: logger.warning(f"Stream '{stream_id}' not found for agent '{target_agent.agent_id}'."); raise web.HTTPNotFound(text=f"Stream resource '{stream_id}' not found.")
-            except asyncio.CancelledError: logger.info(f"HTTP stream download for '{stream_id}' (agent '{target_agent.agent_id}') cancelled."); raise 
-            except Exception as e: logger.error(f"Error during HTTP stream download for '{stream_id}' (agent '{target_agent.agent_id}'): {e}", exc_info=True); raise web.HTTPInternalServerError(text="Error serving stream data.")
-            finally:
-                try: cleanup_method = getattr(target_agent, "cleanup_stream_resource"); await cleanup_method(stream_id); logger.debug(f"Agent '{target_agent.agent_id}' cleaned up stream '{stream_id}'.")
-                except Exception as cleanup_e: logger.error(f"Error cleaning up stream '{stream_id}': {cleanup_e}", exc_info=True)
-
-        async def start_server(self, config: AgentServerConfig):
-            # ... (method body remains the same) ...
-            if not config.sse_base_url: raise ValueError("SSE base URL required.")
-            self._app['agent_server_config'] = config; self._setup_routes(config)
-            host = config.sse_base_url.host or "0.0.0.0"; port = config.sse_base_url.port or 80
-            self._runner = web.AppRunner(self._app); await self._runner.setup()
-            self._site = web.TCPSite(self._runner, host, port)
-            try: await self._site.start(); logger.info(f"SseServerHandler started for agents {list(self._agents.keys())} on http://{host}:{port}")
-            except OSError as e: logger.error(f"Failed to start SseServer on http://{host}:{port}: {e}", exc_info=True); await self.stop_server(); raise
-            except Exception as e: logger.error(f"Unexpected error starting SseServer: {e}", exc_info=True); await self.stop_server(); raise
-
-        async def stop_server(self):
-            # ... (method body remains the same) ...
-            logger.info(f"SseServerHandler for agents {list(self._agents.keys())} stopping...")
-            for task in list(self._active_sse_forwarding_tasks.values()): 
-                if task and not task.done(): task.cancel()
-            if self._active_sse_forwarding_tasks: await asyncio.gather(*self._active_sse_forwarding_tasks.values(), return_exceptions=True)
-            self._active_sse_forwarding_tasks.clear()
-            if self._site: try: await self._site.stop()
-                except Exception as e: logger.error(f"Error stopping TCPSite: {e}", exc_info=True)
-                self._site = None
-            if self._runner: try: await self._runner.cleanup()
-                except Exception as e: logger.error(f"Error cleaning AppRunner: {e}", exc_info=True)
-                self._runner = None
-            logger.info(f"SseServerHandler for agents {list(self._agents.keys())} stopped.")
-    

autobyteus/rpc/server/stdio_server_handler.py

@@ -1,151 +0,0 @@
-# file: autobyteus/autobyteus/rpc/server/stdio_server_handler.py
-import asyncio
-import sys
-import json
-import logging
-from typing import Dict, Optional, Union
-
-from autobyteus.rpc.protocol import ProtocolMessage, MessageType, ErrorCode, RequestType
-from autobyteus.rpc.server.base_method_handler import BaseMethodHandler
-from autobyteus.agent.agent import Agent # Changed from AgentRuntime
-
-logger = logging.getLogger(__name__)
-
-class StdioServerHandler:
-    """
-    Handles RPC communication over stdio for an Agent Server.
-    Reads newline-delimited JSON ProtocolMessages from stdin, dispatches them
-    to appropriate method handlers (which operate on an Agent instance),
-    and writes responses to stdout.
-    """
-
-    def __init__(self, agent: Agent, method_handlers: Dict[Union[RequestType, str], BaseMethodHandler]): # Changed runtime to agent
-        """
-        Initializes the StdioServerHandler.
-
-        Args:
-            agent: The Agent instance this server handler is serving.
-            method_handlers: A dictionary mapping method names (RequestType or str)
-                             to their handler instances.
-        """
-        self._agent = agent # Changed from _runtime to _agent
-        self._method_handlers = method_handlers
-        self._running = False
-        logger.info(f"StdioServerHandler initialized for agent '{self._agent.agent_id}'.") # Use agent.agent_id
-
-    async def listen_and_dispatch(self) -> None:
-        """
-        Starts listening on stdin for requests and dispatches them.
-        This method runs indefinitely until stdin is closed or an error occurs.
-        """
-        self._running = True
-        logger.info(f"StdioServerHandler for agent '{self._agent.agent_id}' now listening on stdin.")
-        
-        loop = asyncio.get_event_loop()
-        reader = asyncio.StreamReader(loop=loop)
-        protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
-        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
-
-        try:
-            while self._running:
-                line_bytes = await reader.readline()
-                if not line_bytes:
-                    logger.info(f"StdioServerHandler for agent '{self._agent.agent_id}': stdin EOF reached. Shutting down.")
-                    self._running = False
-                    break
-                
-                line_str = line_bytes.decode().strip()
-                if not line_str:
-                    continue
-
-                request_id: Optional[str] = None
-                try:
-                    try:
-                        parsed_json = json.loads(line_str)
-                        request_id = parsed_json.get("id")
-                    except json.JSONDecodeError:
-                        pass
-
-                    message = ProtocolMessage.from_json_str(line_str)
-                    request_id = message.id 
-
-                    if message.type == MessageType.REQUEST:
-                        response_message = await self._process_request(message)
-                    else:
-                        logger.warning(f"StdioServerHandler received non-REQUEST message type: {message.type}. Ignoring.")
-                        response_message = ProtocolMessage.create_error_response(
-                            id=request_id,
-                            code=ErrorCode.INVALID_REQUEST,
-                            message="Server only accepts REQUEST messages via stdio."
-                        )
-                    
-                    await self._send_response(response_message)
-
-                except json.JSONDecodeError as e:
-                    logger.error(f"StdioServerHandler: JSONDecodeError: {e}. Raw line: '{line_str[:200]}'")
-                    err_response = ProtocolMessage.create_error_response(
-                        id=request_id,
-                        code=ErrorCode.PARSE_ERROR,
-                        message=f"Failed to parse JSON request: {e}"
-                    )
-                    await self._send_response(err_response)
-                except ValueError as e:
-                    logger.error(f"StdioServerHandler: ProtocolMessage validation error: {e}. Raw line: '{line_str[:200]}'")
-                    err_response = ProtocolMessage.create_error_response(
-                        id=request_id,
-                        code=ErrorCode.INVALID_REQUEST,
-                        message=f"Invalid request structure: {e}"
-                    )
-                    await self._send_response(err_response)
-                except Exception as e:
-                    logger.error(f"StdioServerHandler: Unexpected error processing line: {e}. Raw line: '{line_str[:200]}'", exc_info=True)
-                    err_response = ProtocolMessage.create_error_response(
-                        id=request_id,
-                        code=ErrorCode.INTERNAL_ERROR,
-                        message=f"Internal server error: {e}"
-                    )
-                    await self._send_response(err_response)
-        
-        except asyncio.CancelledError:
-            logger.info(f"StdioServerHandler for agent '{self._agent.agent_id}' listen_and_dispatch task cancelled.")
-        except Exception as e:
-            logger.error(f"StdioServerHandler for agent '{self._agent.agent_id}' fatal error in listen_and_dispatch: {e}", exc_info=True)
-        finally:
-            self._running = False
-            logger.info(f"StdioServerHandler for agent '{self._agent.agent_id}' stopped listening.")
-
-    async def _process_request(self, request_message: ProtocolMessage) -> ProtocolMessage:
-        if not request_message.method:
-            logger.warning(f"StdioServerHandler: Request message missing 'method'. ID: {request_message.id}")
-            return ProtocolMessage.create_error_response(
-                id=request_message.id,
-                code=ErrorCode.INVALID_REQUEST,
-                message="Request message must include a 'method'."
-            )
-
-        handler = self._method_handlers.get(request_message.method)
-        if not handler:
-            logger.warning(f"StdioServerHandler: No handler found for method '{request_message.method}'. ID: {request_message.id}")
-            return ProtocolMessage.create_error_response(
-                id=request_message.id,
-                code=ErrorCode.METHOD_NOT_FOUND,
-                message=f"Method '{request_message.method}' not found."
-            )
-        
-        logger.debug(f"StdioServerHandler dispatching method '{request_message.method}' (ID: {request_message.id}) to {handler.__class__.__name__}.")
-        return await handler.handle(request_message.id, request_message.params, self._agent) # Pass self._agent
-
-    async def _send_response(self, response_message: ProtocolMessage) -> None:
-        try:
-            json_response = response_message.to_json_str()
-            loop = asyncio.get_event_loop()
-            await loop.run_in_executor(None, lambda: sys.stdout.write(json_response + '\n'))
-            await loop.run_in_executor(None, sys.stdout.flush)
-            logger.debug(f"StdioServerHandler sent response (ID: {response_message.id}, Type: {response_message.type}).")
-        except Exception as e:
-            logger.error(f"StdioServerHandler failed to send response: {e}", exc_info=True)
-
-    def stop(self):
-        logger.info(f"StdioServerHandler for agent '{self._agent.agent_id}' stop requested.")
-        self._running = False
-