hanzoai_agent/sdk/python/hanzo_agents/agent_field_handler.py

import asyncio
import os
import signal
import threading
from datetime import datetime

import requests
from hanzo_agents.types import AgentStatus, HeartbeatData
from hanzo_agents.logger import (
    log_heartbeat,
    log_debug,
    log_warn,
    log_error,
    log_success,
    log_setup,
    log_info,
)


class HanzoAgentsHandler:
    """
    Hanzo Agents Server Communication handler for Agent class.

    This class encapsulates all Hanzo Agents server communication functionality including:
    - Agent registration with Hanzo Agents server
    - Heartbeat management (both simple and enhanced)
    - Fast lifecycle management
    - Graceful shutdown notifications
    - Signal handling for fast shutdown
    """

    def __init__(self, agent_instance):
        """
        Initialize the Hanzo Agents handler with a reference to the agent instance.

        Args:
            agent_instance: The Agent instance this handler belongs to
        """
        self.agent = agent_instance

    async def register_with_hanzo_agents_server(self, port: int):
        """Register this agent node with Hanzo Agents server"""
        # Import the callback URL resolution function
        from hanzo_agents.agent import (
            _build_callback_candidates,
            _resolve_callback_url,
            _is_running_in_container,
        )

        # Enhanced debugging for callback URL resolution
        log_debug("Starting callback URL resolution")
        log_debug(f"Original callback_url parameter: {self.agent.callback_url}")
        log_debug(
            f"AGENT_CALLBACK_URL env var: {os.environ.get('AGENT_CALLBACK_URL', 'NOT_SET')}"
        )
        log_debug(f"Port: {port}")
        log_debug(f"Running in container: {_is_running_in_container()}")
        log_debug(
            f"All env vars containing 'AGENT': {[k for k in os.environ.keys() if 'AGENT' in k.upper()]}"
        )

        # 🔥 FIX: Only resolve callback URL if not already set
        # This prevents overwriting the URL resolved in Agent.__init__()
        if not self.agent.base_url:
            self.agent.callback_candidates = _build_callback_candidates(
                self.agent.callback_url, port
            )
            if self.agent.callback_candidates:
                self.agent.base_url = self.agent.callback_candidates[0]
                log_debug(
                    f"Resolved callback URL during registration: {self.agent.base_url}"
                )
            else:
                self.agent.base_url = _resolve_callback_url(
                    self.agent.callback_url, port
                )
                log_debug(
                    f"Resolved callback URL during registration: {self.agent.base_url}"
                )
        else:
            # Update port in existing base_url if needed, but preserve internal container URLs
            import urllib.parse

            parsed = urllib.parse.urlparse(self.agent.base_url)

            # Don't modify internal or other container-specific URLs
            if "internal" in parsed.netloc:
                log_debug(
                    f"Preserving container-specific callback URL: {self.agent.base_url}"
                )
            elif parsed.port != port:
                # Update the port in the existing URL
                self.agent.base_url = f"{parsed.scheme}://{parsed.hostname}:{port}"
                log_debug(
                    f"Updated port in existing callback URL: {self.agent.base_url}"
                )
            else:
                log_debug(f"Using existing callback URL: {self.agent.base_url}")

        if not self.agent.callback_candidates:
            self.agent.callback_candidates = _build_callback_candidates(
                self.agent.base_url, port
            )
        elif (
            self.agent.base_url
            and self.agent.callback_candidates[0] != self.agent.base_url
        ):
            # Keep resolved base URL at front for clarity
            if self.agent.base_url in self.agent.callback_candidates:
                self.agent.callback_candidates.remove(self.agent.base_url)
            self.agent.callback_candidates.insert(0, self.agent.base_url)

        # Always log the resolved callback URL for debugging
        log_info(f"Final callback URL: {self.agent.base_url}")

        if self.agent.dev_mode:
            log_debug(f"Final callback URL: {self.agent.base_url}")

        try:
            log_debug(
                f"Attempting to register with Hanzo Agents server at {self.agent.hanzo_agents_server}"
            )
            discovery_payload = self.agent._build_callback_discovery_payload()

            success, payload = await self.agent.client.register_agent(
                node_id=self.agent.node_id,
                reasoners=self.agent.reasoners,
                skills=self.agent.skills,
                base_url=self.agent.base_url,
                discovery=discovery_payload,
                vc_metadata=self.agent._build_vc_metadata(),
                version=self.agent.version,
                agent_metadata=self.agent._build_agent_metadata(),
            )
            if success:
                if payload:
                    self.agent._apply_discovery_response(payload)
                log_success(
                    f"Registered node '{self.agent.node_id}' with Hanzo Agents server"
                )
                self.agent.hanzo_agents_connected = True

                # Attempt DID registration after successful Hanzo Agents registration
                if self.agent.did_manager:
                    did_success = self.agent._register_agent_with_did()
                    if not did_success and self.agent.dev_mode:
                        log_warn(
                            "DID registration failed, continuing without DID functionality"
                        )
            else:
                log_error("Registration failed")
                self.agent.hanzo_agents_connected = False

        except Exception as e:
            self.agent.hanzo_agents_connected = False
            if self.agent.dev_mode:
                log_warn(f"Hanzo Agents server not available: {e}")
                log_setup("Running in development mode - agent will work standalone")
                log_info(
                    f"To connect to Hanzo Agents server, start it at {self.agent.hanzo_agents_server}"
                )
            else:
                log_error(f"Failed to register with Hanzo Agents server: {e}")
                if (
                    isinstance(e, requests.exceptions.RequestException)
                    and e.response is not None
                ):
                    log_warn(f"Response status: {e.response.status_code}")
                    log_warn(f"Response text: {e.response.text}")
                raise

    def send_heartbeat(self):
        """Send heartbeat to Hanzo Agents server"""
        if not self.agent.hanzo_agents_connected:
            return  # Skip heartbeat if not connected to Hanzo Agents

        try:
            headers = {"Content-Type": "application/json"}
            if self.agent.api_key:
                headers["X-API-Key"] = self.agent.api_key
            response = requests.post(
                f"{self.agent.hanzo_agents_server}/v1/nodes/{self.agent.node_id}/heartbeat",
                headers=headers,
                timeout=5,
            )
            if response.status_code == 200:
                log_heartbeat("Heartbeat sent successfully")
            else:
                log_warn(
                    f"Heartbeat failed with status {response.status_code}: {response.text}"
                )
        except Exception as e:
            log_error(f"Failed to send heartbeat: {e}")

    def heartbeat_worker(
        self, interval: int = 30
    ):  # pragma: no cover - long-running thread loop
        """Background worker that sends periodic heartbeats"""
        if not self.agent.hanzo_agents_connected:
            log_heartbeat(
                "Heartbeat worker skipped - not connected to Hanzo Agents server"
            )
            return

        log_heartbeat(f"Starting heartbeat worker (interval: {interval}s)")
        while not self.agent._heartbeat_stop_event.wait(interval):
            self.send_heartbeat()
        log_heartbeat("Heartbeat worker stopped")

    def start_heartbeat(self, interval: int = 30):
        """Start the heartbeat background thread"""
        if not self.agent.hanzo_agents_connected:
            return  # Skip heartbeat if not connected to Hanzo Agents

        if (
            self.agent._heartbeat_thread is None
            or not self.agent._heartbeat_thread.is_alive()
        ):
            self.agent._heartbeat_stop_event.clear()
            self.agent._heartbeat_thread = threading.Thread(
                target=self.heartbeat_worker, args=(interval,), daemon=True
            )
            self.agent._heartbeat_thread.start()

    def stop_heartbeat(self):
        """Stop the heartbeat background thread"""
        if self.agent._heartbeat_thread and self.agent._heartbeat_thread.is_alive():
            log_debug("Stopping heartbeat worker...")
            self.agent._heartbeat_stop_event.set()
            self.agent._heartbeat_thread.join(timeout=5)

    async def send_enhanced_heartbeat(self) -> bool:
        """
        Send enhanced heartbeat with current status and MCP information.

        Returns:
            True if heartbeat was successful, False otherwise
        """
        if not self.agent.hanzo_agents_connected:
            return False

        try:
            # Get MCP server health information
            mcp_servers = self.agent.mcp_handler._get_mcp_server_health()

            # Create heartbeat data
            heartbeat_data = HeartbeatData(
                status=self.agent._current_status,
                mcp_servers=mcp_servers,
                timestamp=datetime.now().isoformat(),
            )

            # Send enhanced heartbeat
            success = await self.agent.client.send_enhanced_heartbeat(
                self.agent.node_id, heartbeat_data
            )

            if success:
                log_heartbeat(
                    f"Enhanced heartbeat sent - Status: {self.agent._current_status.value}"
                )

            return success

        except Exception as e:
            if self.agent.dev_mode:
                log_error(f"Enhanced heartbeat failed: {e}")
            return False

    async def notify_shutdown(self) -> bool:
        """
        Notify Hanzo Agents server of graceful shutdown.

        Returns:
            True if notification was successful, False otherwise
        """
        if not self.agent.hanzo_agents_connected:
            return False

        try:
            success = await self.agent.client.notify_graceful_shutdown(
                self.agent.node_id
            )
            if self.agent.dev_mode and success:
                log_success("Graceful shutdown notification sent")
            return success
        except Exception as e:
            if self.agent.dev_mode:
                log_error(f"Shutdown notification failed: {e}")
            return False

    def setup_fast_lifecycle_signal_handlers(
        self,
    ) -> None:  # pragma: no cover - requires OS signal integration
        """
        Setup signal handler for fast lifecycle status while allowing uvicorn to perform graceful shutdown.

        - Only intercepts SIGTERM to mark the agent offline and notify Hanzo Agents immediately.
        - Leaves SIGINT (Ctrl+C) to uvicorn so its shutdown hooks run and resources are cleaned up.
        """

        def signal_handler(signum: int, frame) -> None:
            """Handle SIGTERM: mark offline, notify Hanzo Agents, then re-emit the signal for default handling."""
            signal_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"

            if self.agent.dev_mode:
                log_warn(
                    f"{signal_name} received - initiating graceful shutdown via uvicorn"
                )

            # Set shutdown flag
            self.agent._shutdown_requested = True
            self.agent._current_status = AgentStatus.OFFLINE

            # Best-effort immediate notification to Hanzo Agents
            try:
                success = self.agent.client.notify_graceful_shutdown_sync(
                    self.agent.node_id
                )
                if self.agent.dev_mode:
                    state = "sent" if success else "failed"
                    log_info(f"Shutdown notification {state}")
            except Exception as e:
                if self.agent.dev_mode:
                    log_error(f"Shutdown notification error: {e}")

            # IMPORTANT: Do not perform heavy cleanup here. Let FastAPI/uvicorn shutdown events handle it.
            # Re-install default handler and re-emit the same signal so uvicorn orchestrates cleanup.
            try:
                signal.signal(signum, signal.SIG_DFL)
                os.kill(os.getpid(), signum)
            except Exception:
                # Fallback: polite exit (still allows finally blocks/atexit to run)
                import sys

                sys.exit(0)

        try:
            # Only register for SIGTERM; leave SIGINT (Ctrl+C) to uvicorn
            signal.signal(signal.SIGTERM, signal_handler)

            if self.agent.dev_mode:
                log_debug("Fast lifecycle signal handler registered (SIGTERM only)")
        except Exception as e:
            if self.agent.dev_mode:
                log_error(f"Failed to setup signal handlers: {e}")

    async def register_with_fast_lifecycle(
        self, port: int
    ) -> bool:  # pragma: no cover - fast-path relies on external coordination
        """
        Register agent with immediate status reporting for fast lifecycle.

        Args:
            port: The port the agent is running on

        Returns:
            True if registration was successful, False otherwise
        """
        from hanzo_agents.agent import _build_callback_candidates, _resolve_callback_url

        if not self.agent.base_url:
            self.agent.callback_candidates = _build_callback_candidates(
                self.agent.callback_url, port
            )
            if self.agent.callback_candidates:
                self.agent.base_url = self.agent.callback_candidates[0]
                log_debug(
                    f"Fast lifecycle - Resolved callback URL during registration: {self.agent.base_url}"
                )
            else:
                self.agent.base_url = _resolve_callback_url(
                    self.agent.callback_url, port
                )
                log_debug(
                    f"Fast lifecycle - Resolved callback URL during registration: {self.agent.base_url}"
                )
        else:
            import urllib.parse

            parsed = urllib.parse.urlparse(self.agent.base_url)
            if parsed.port != port:
                self.agent.base_url = f"{parsed.scheme}://{parsed.hostname}:{port}"
                log_debug(
                    f"Fast lifecycle - Updated port in existing callback URL: {self.agent.base_url}"
                )
            else:
                log_debug(
                    f"Fast lifecycle - Using existing callback URL: {self.agent.base_url}"
                )

        if not self.agent.callback_candidates:
            self.agent.callback_candidates = _build_callback_candidates(
                self.agent.base_url, port
            )
        elif (
            self.agent.base_url
            and self.agent.callback_candidates
            and self.agent.callback_candidates[0] != self.agent.base_url
        ):
            if self.agent.base_url in self.agent.callback_candidates:
                self.agent.callback_candidates.remove(self.agent.base_url)
            self.agent.callback_candidates.insert(0, self.agent.base_url)

        log_debug(f"Fast lifecycle - Final callback URL: {self.agent.base_url}")
        log_debug(
            f"Fast lifecycle - Original callback_url parameter: {self.agent.callback_url}"
        )
        log_debug(
            f"Fast lifecycle - AGENT_CALLBACK_URL env var: {os.environ.get('AGENT_CALLBACK_URL', 'NOT_SET')}"
        )
        log_debug(f"Fast lifecycle - Port: {port}")

        try:
            if self.agent.dev_mode:
                log_info(
                    f"Fast registration with Hanzo Agents server at {self.agent.hanzo_agents_server}"
                )
                log_info(f"Using callback URL: {self.agent.base_url}")

            # Register with STARTING status for immediate visibility
            discovery_payload = self.agent._build_callback_discovery_payload()

            success, payload = await self.agent.client.register_agent_with_status(
                node_id=self.agent.node_id,
                reasoners=self.agent.reasoners,
                skills=self.agent.skills,
                base_url=self.agent.base_url,
                status=AgentStatus.STARTING,
                discovery=discovery_payload,
                vc_metadata=self.agent._build_vc_metadata(),
                version=self.agent.version,
                agent_metadata=self.agent._build_agent_metadata(),
            )

            if success:
                if payload:
                    self.agent._apply_discovery_response(payload)
                if self.agent.dev_mode:
                    log_success(
                        f"Fast registration successful - Status: {AgentStatus.STARTING.value}"
                    )
                self.agent.hanzo_agents_connected = True

                # Attempt DID registration after successful Hanzo Agents registration
                if self.agent.did_manager:
                    did_success = self.agent._register_agent_with_did()
                    if not did_success and self.agent.dev_mode:
                        log_warn(
                            "DID registration failed, continuing without DID functionality"
                        )

                return True
            else:
                if self.agent.dev_mode:
                    log_error("Fast registration failed")
                self.agent.hanzo_agents_connected = False
                return False

        except Exception as e:
            self.agent.hanzo_agents_connected = False
            if self.agent.dev_mode:
                log_warn(f"Fast registration error: {e}")
            return False

    async def enhanced_heartbeat_loop(self, interval: int) -> None:
        """
        Background loop for sending enhanced heartbeats with status and MCP information.

        Args:
            interval: Heartbeat interval in seconds
        """
        if self.agent.dev_mode:
            log_debug(f"Enhanced heartbeat loop started (interval: {interval}s)")

        while not self.agent._shutdown_requested:
            try:
                # Send enhanced heartbeat
                success = await self.send_enhanced_heartbeat()

                if not success and self.agent.dev_mode:
                    log_warn("Enhanced heartbeat failed - retrying next cycle")

                # Wait for next heartbeat interval
                await asyncio.sleep(interval)

            except asyncio.CancelledError:
                if self.agent.dev_mode:
                    log_debug("Enhanced heartbeat loop cancelled")
                break
            except Exception as e:
                if self.agent.dev_mode:
                    log_error(f"Enhanced heartbeat loop error: {e}")
                # Continue loop even on errors
                await asyncio.sleep(interval)

        if self.agent.dev_mode:
            log_debug("Enhanced heartbeat loop stopped")