Improve local agent chat and device health detection

apps/device-host-agent/host_agent/app.py

@@ -12,6 +12,7 @@ from device.manager import DeviceManager
 from host_agent.assignment import AssignmentExecutor
 from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
 from host_agent.config import HostAgentConfig, load_host_agent_config
+from host_agent.conversation import ConversationAgent
 from host_agent.dependency_supervisor import DependencySupervisor
 from host_agent.devices import register_local_device
 from host_agent.enrollment import resolve_host_identity
@@ -33,6 +34,8 @@ from host_agent.status import AgentStatusTracker
 from host_agent.web.app import create_console_app
 from host_agent.web.auth import SessionManager
 from host_agent.web.mcp import build_mcp_server
+from runtime.executor import default_tool_registry
+from runtime.planner_config import load_config as load_planner_config
 from storage.artifact_store import ArtifactStore
 from storage.device_config import DeviceConfigStore
 from storage.task_metadata import TaskMetadataStore
@@ -241,6 +244,15 @@ def create_application(
             mcp_busy_tracker=mcp_busy_tracker,
             status_tracker=status_tracker,
         )
+        planner_config = load_planner_config()
+        conversation_agent = (
+            ConversationAgent(
+                config=planner_config,
+                tools=default_tool_registry(manager=resolved_manager),
+            )
+            if resolved_config.ai_planner_transport == "direct"
+            else None
+        )
         console_app = create_console_app(
             config=resolved_config,
             manager=resolved_manager,
@@ -260,6 +272,7 @@ def create_application(
             mcp_server=mcp_server,
             mcp_token_store=mcp_token_store,
             mcp_busy_tracker=mcp_busy_tracker,
+            conversation_agent=conversation_agent,
         )
         console_server = _EmbeddedConsoleServer(
             uvicorn.Config(

apps/device-host-agent/host_agent/config.py

@@ -136,9 +136,13 @@ def load_host_agent_config(
         ),
         ai_planner_transport=("direct" if mode == "local" else _parse_ai_planner_transport(values.get("AI_PLANNER_TRANSPORT"))),
         dependency_supervisor_enabled=_truthy(
-            values, "HOST_AGENT_DEPENDENCY_SUPERVISOR_ENABLED", False
+            values,
+            "HOST_AGENT_DEPENDENCY_SUPERVISOR_ENABLED",
+            mode == "local",
+        ),
+        appium_supervised=_truthy(
+            values, "HOST_AGENT_APPIUM_SUPERVISED", mode == "local"
         ),
-        appium_supervised=_truthy(values, "HOST_AGENT_APPIUM_SUPERVISED", False),
         appium_host=values.get("HOST_AGENT_APPIUM_HOST", "127.0.0.1").strip(),
         appium_port=_positive_int(values, "HOST_AGENT_APPIUM_PORT", 4723),
         dependency_restart_max_attempts=_positive_int(

apps/device-host-agent/host_agent/conversation.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from runtime.planner_config import PlannerConfig
+from runtime.tool_specs import ACTION_TOOL_SPECS, ToolSpec
+
+READ_TOOL_SPECS = [
+    ToolSpec("take_screenshot", "Capture the current device screen.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
+    ToolSpec("describe_screen", "Inspect the current screen and UI.", {"type": "object", "properties": {"device_id": {"type": "string"}}}),
+    ToolSpec("find_text", "Find visible text on the current screen.", {"type": "object", "required": ["query"], "properties": {"query": {"type": "string"}, "device_id": {"type": "string"}}}),
+    ToolSpec("get_ui_tree", "Read the current accessibility/UI tree.", {"type": "object", "properties": {"device_id": {"type": "string"}, "include_app_info": {"type": "boolean"}}}),
+    ToolSpec("list_devices", "List configured devices.", {"type": "object", "properties": {}}),
+    ToolSpec("device_status", "Read one device status.", {"type": "object", "required": ["device_id"], "properties": {"device_id": {"type": "string"}}}),
+]
+
+
+@dataclass(frozen=True)
+class ChatResult:
+    content: str
+    tool_calls: int
+
+
+class ConversationAgent:
+    """Small OpenAI-compatible agent loop backed by Host Agent tools."""
+
+    def __init__(
+        self,
+        *,
+        config: PlannerConfig,
+        tools: dict[str, Callable[..., Any]],
+        max_rounds: int = 8,
+    ) -> None:
+        self.config = config
+        self.tools = tools
+        self.max_rounds = max_rounds
+        self._client: Any | None = None
+
+    def chat(self, messages: list[dict[str, str]]) -> ChatResult:
+        if self.config.provider not in {"openai", "openai_compatible"}:
+            raise RuntimeError(
+                "conversation chat currently requires an OpenAI-compatible provider"
+            )
+        client = self._get_client()
+        history: list[dict[str, Any]] = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a mobile device assistant. Reply naturally when no "
+                    "device action is needed. Call a tool when the user asks you "
+                    "to inspect or operate a device. Never claim an action was "
+                    "completed unless the tool result confirms it."
+                ),
+            },
+            *messages,
+        ]
+        calls = 0
+        for _ in range(self.max_rounds):
+            response = client.chat.completions.create(
+                model=self.config.resolved_model(),
+                messages=history,
+                tools=[_openai_tool(spec) for spec in [*ACTION_TOOL_SPECS, *READ_TOOL_SPECS]],
+                tool_choice="auto",
+                parallel_tool_calls=False,
+                timeout=self.config.timeout,
+            )
+            message = response.choices[0].message
+            text = getattr(message, "content", None)
+            tool_calls = getattr(message, "tool_calls", None) or []
+            if not tool_calls:
+                return ChatResult(content=str(text or ""), tool_calls=calls)
+            history.append(_message_dict(message))
+            for tool_call in tool_calls:
+                name = tool_call.function.name
+                try:
+                    arguments = json.loads(tool_call.function.arguments or "{}")
+                    if not isinstance(arguments, dict):
+                        raise ValueError("tool arguments must be an object")
+                    arguments.pop("purpose", None)
+                    arguments.pop("expected_outcome", None)
+                    result = self.tools[name](**arguments)
+                except Exception as exc:
+                    result = {"ok": False, "error": str(exc)}
+                calls += 1
+                history.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": tool_call.id,
+                        "content": json.dumps(result, ensure_ascii=True, default=str),
+                    }
+                )
+        raise RuntimeError("conversation exceeded the maximum tool-call rounds")
+
+    def _get_client(self) -> Any:
+        if self._client is None:
+            from openai import OpenAI
+
+            kwargs: dict[str, Any] = {}
+            if self.config.api_key:
+                kwargs["api_key"] = self.config.api_key
+            if self.config.base_url:
+                kwargs["base_url"] = self.config.base_url
+            self._client = OpenAI(**kwargs)
+        return self._client
+
+
+def _openai_tool(spec: ToolSpec) -> dict[str, Any]:
+    parameters = dict(spec.parameters)
+    properties = dict(parameters.get("properties") or {})
+    properties.setdefault(
+        "device_id",
+        {"type": "string", "description": "Target device ID when needed."},
+    )
+    parameters["properties"] = properties
+    return {
+        "type": "function",
+        "function": {
+            "name": spec.name,
+            "description": spec.description,
+            "parameters": parameters,
+        },
+    }
+
+
+def _message_dict(message: Any) -> dict[str, Any]:
+    result: dict[str, Any] = {"role": "assistant"}
+    content = getattr(message, "content", None)
+    if content is not None:
+        result["content"] = content
+    calls = getattr(message, "tool_calls", None) or []
+    result["tool_calls"] = [
+        {
+            "id": call.id,
+            "type": "function",
+            "function": {
+                "name": call.function.name,
+                "arguments": call.function.arguments,
+            },
+        }
+        for call in calls
+    ]
+    return result

apps/device-host-agent/host_agent/heartbeat.py

@@ -60,6 +60,8 @@ class HeartbeatSynchronizer:
             self.status_tracker.mark_host_policy(self.policy)
 
     async def sync_once(self) -> HeartbeatResponse:
+        self.probe_connected_devices()
+        self.connect_devices()
         snapshot = build_device_snapshot(self.manager)
         mcp_busy_ids = (
             self.mcp_busy_tracker.busy_device_ids()
@@ -108,9 +110,19 @@ class HeartbeatSynchronizer:
 
     def connect_devices(self) -> None:
         for device in self.manager.list_devices():
-            if device.status != "idle":
+            if device.status not in {"idle", "offline", "error"}:
                 continue
             try:
                 self.manager.connect(device.id)
             except DeviceRuntimeError:
                 continue
+
+    def probe_connected_devices(self) -> None:
+        active_id: str | None = None
+        if self.status_tracker is not None:
+            current = self.status_tracker.snapshot().get("current_assignment")
+            if isinstance(current, dict) and isinstance(current.get("device_id"), str):
+                active_id = current["device_id"]
+        for device in self.manager.list_devices():
+            if device.id != active_id and device.status == "busy":
+                self.manager.probe(device.id)

apps/device-host-agent/host_agent/web/app.py

@@ -20,6 +20,7 @@ from host_agent.client import (
     HostTaskSubmissionUnknownError,
 )
 from host_agent.config import HostAgentConfig
+from host_agent.conversation import ConversationAgent
 from host_agent.devices import register_local_device, unregister_local_device
 from host_agent.history import ConsoleHistoryStore
 from host_agent.identity import HostIdentityStore
@@ -232,6 +233,7 @@ def create_console_app(
     mcp_server: FastMCP | None = None,
     mcp_token_store: McpTokenStore | None = None,
     mcp_busy_tracker: McpBusyTracker | None = None,
+    conversation_agent: ConversationAgent | None = None,
 ) -> FastAPI:
     app = FastAPI(title="Host Agent Console")
     cookie_secure = config.console_bind_host not in _LOOPBACK_BIND_HOSTS
@@ -415,6 +417,31 @@ def create_console_app(
             }
         )
 
+    @app.post("/api/chat")
+    async def api_chat(
+        request: Request,
+        session: SessionState = Depends(require_csrf),
+    ) -> JSONResponse:
+        if conversation_agent is None:
+            raise HTTPException(status_code=503, detail="chat agent is not configured")
+        payload = await request.json()
+        raw_messages = payload.get("messages") if isinstance(payload, dict) else None
+        if not isinstance(raw_messages, list) or not raw_messages:
+            raise HTTPException(status_code=400, detail="messages must be a non-empty list")
+        messages = [
+            {"role": str(item["role"]), "content": str(item["content"])}
+            for item in raw_messages
+            if isinstance(item, dict)
+            and item.get("role") in {"user", "assistant"}
+            and isinstance(item.get("content"), str)
+        ]
+        if not messages:
+            raise HTTPException(status_code=400, detail="messages are invalid")
+        result = await asyncio.to_thread(conversation_agent.chat, messages)
+        return JSONResponse(
+            {"content": result.content, "tool_calls": result.tool_calls}
+        )
+
     @app.get("/devices", response_class=HTMLResponse)
     async def devices_page(
         request: Request,

apps/device-host-agent/tests/test_config.py

@@ -34,6 +34,8 @@ def test_load_host_agent_config_supports_local_mode() -> None:
     assert config.control_plane_url == ""
     assert config.enrollment_managed is False
     assert config.ai_planner_transport == "direct"
+    assert config.dependency_supervisor_enabled is True
+    assert config.appium_supervised is True
 
 
 def test_load_host_agent_config_parses_poll_and_retry_values() -> None:

apps/device-host-agent/tests/test_heartbeat.py

@@ -17,6 +17,9 @@ class ConnectableDriver:
     def connect(self) -> None:
         return None
 
+    def screenshot(self) -> bytes:
+        return b"ok"
+
 
 def _config() -> HostAgentConfig:
     return HostAgentConfig(
@@ -89,6 +92,33 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N
     assert manager.status("device-a") == "busy"
 
 
+def test_offline_device_is_retried_on_next_heartbeat_cycle() -> None:
+    class FlakyDriver(ConnectableDriver):
+        def __init__(self, fail: bool) -> None:
+            self.fail = fail
+
+        def screenshot(self) -> bytes:
+            if self.fail:
+                raise RuntimeError("WDA disconnected")
+            return b"ok"
+
+    instances: list[FlakyDriver] = []
+
+    def factory() -> FlakyDriver:
+        driver = FlakyDriver(not instances)
+        instances.append(driver)
+        return driver
+
+    manager = DeviceManager()
+    manager.register_device("device-a", factory)  # type: ignore[arg-type]
+    manager.connect("device-a")
+    sync = HeartbeatSynchronizer(manager, object(), _config())  # type: ignore[arg-type]
+    sync.probe_connected_devices()
+    assert manager.status("device-a") == "offline"
+    sync.connect_devices()
+    assert manager.status("device-a") == "busy"
+
+
 def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> None:
     manager = DeviceManager()
     manager.register_device(