autoglm-gui 1.0.2__py3-none-any.whl → 1.2.0__py3-none-any.whl

AutoGLM_GUI/api/agents.py

@@ -1,13 +1,19 @@
 """Agent lifecycle and chat routes."""
 
 import json
+import queue
+import threading
+from typing import Any
 
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import ValidationError
 
 from AutoGLM_GUI.config import config
+from AutoGLM_GUI.logger import logger
+from AutoGLM_GUI.phone_agent_patches import apply_patches
 from AutoGLM_GUI.schemas import (
+    AbortRequest,
     APIAgentConfig,
     APIModelConfig,
     ChatRequest,
@@ -19,8 +25,6 @@ from AutoGLM_GUI.schemas import (
     StatusResponse,
 )
 from AutoGLM_GUI.state import (
-    agent_configs,
-    agents,
     non_blocking_takeover,
 )
 from AutoGLM_GUI.version import APP_VERSION
@@ -28,8 +32,29 @@ from phone_agent import PhoneAgent
 from phone_agent.agent import AgentConfig
 from phone_agent.model import ModelConfig
 
+# Apply monkey patches to phone_agent
+apply_patches()
+
 router = APIRouter()
 
+# Active chat sessions (device_id -> stop_event)
+# Used for aborting ongoing conversations
+_active_chats: dict[str, threading.Event] = {}
+_active_chats_lock = threading.Lock()
+
+
+def _release_device_lock_when_done(
+    device_id: str, threads: list[threading.Thread]
+) -> None:
+    """Block until threads finish, then release the device lock via manager."""
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
+
+    for thread in threads:
+        thread.join()
+
+    manager = PhoneAgentManager.get_instance()
+    manager.release_device(device_id)
+
 
 @router.post("/api/init")
 def init_agent(request: InitRequest) -> dict:
@@ -95,13 +120,20 @@ def init_agent(request: InitRequest) -> dict:
         verbose=req_agent_config.verbose,
     )
 
-    agents[device_id] = PhoneAgent(
-        model_config=model_config,
-        agent_config=agent_config,
-        takeover_callback=non_blocking_takeover,
-    )
+    # Initialize agent via PhoneAgentManager (thread-safe, transactional)
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
 
-    agent_configs[device_id] = (model_config, agent_config)
+    manager = PhoneAgentManager.get_instance()
+    try:
+        manager.initialize_agent(
+            device_id=device_id,
+            model_config=model_config,
+            agent_config=agent_config,
+            takeover_callback=non_blocking_takeover,
+        )
+    except Exception as e:
+        logger.error(f"Failed to initialize agent: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
 
     return {
         "success": True,
@@ -113,20 +145,29 @@ def init_agent(request: InitRequest) -> dict:
 @router.post("/api/chat", response_model=ChatResponse)
 def chat(request: ChatRequest) -> ChatResponse:
     """发送任务给 Agent 并执行。"""
+    from AutoGLM_GUI.exceptions import DeviceBusyError
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
+
     device_id = request.device_id
-    if device_id not in agents:
+    manager = PhoneAgentManager.get_instance()
+
+    # Check if agent is initialized
+    if not manager.is_initialized(device_id):
         raise HTTPException(
             status_code=400, detail="Agent not initialized. Call /api/init first."
         )
 
-    agent = agents[device_id]
-
+    # Use context manager for automatic lock management
     try:
-        result = agent.run(request.message)
-        steps = agent.step_count
-        agent.reset()
-
-        return ChatResponse(result=result, steps=steps, success=True)
+        with manager.use_agent(device_id, timeout=None) as agent:
+            result = agent.run(request.message)
+            steps = agent.step_count
+            agent.reset()
+            return ChatResponse(result=result, steps=steps, success=True)
+    except DeviceBusyError:
+        raise HTTPException(
+            status_code=409, detail=f"Device {device_id} is busy. Please wait."
+        )
     except Exception as e:
         return ChatResponse(result=str(e), steps=0, success=False)
 
@@ -134,96 +175,279 @@ def chat(request: ChatRequest) -> ChatResponse:
 @router.post("/api/chat/stream")
 def chat_stream(request: ChatRequest):
     """发送任务给 Agent 并实时推送执行进度（SSE，多设备支持）。"""
+    from AutoGLM_GUI.exceptions import AgentNotInitializedError, DeviceBusyError
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
+
     device_id = request.device_id
+    manager = PhoneAgentManager.get_instance()
 
-    if device_id not in agents:
+    # Check if agent is initialized
+    if not manager.is_initialized(device_id):
         raise HTTPException(
             status_code=400,
             detail=f"Device {device_id} not initialized. Call /api/init first.",
         )
 
-    agent = agents[device_id]
+    # Acquire device lock (non-blocking) to prevent concurrent requests
+    try:
+        manager.acquire_device(device_id, timeout=0, raise_on_timeout=True)
+    except DeviceBusyError:
+        raise HTTPException(
+            status_code=409,
+            detail=f"Device {device_id} is already processing a request. Please wait.",
+        )
+
+    try:
+        # Get the original agent to copy its config
+        original_agent = manager.get_agent(device_id)
 
-    def event_generator():
-        """SSE 事件生成器"""
+        # Get the stored configs for this device
         try:
-            step_result = agent.step(request.message)
-            while True:
-                event_data = {
-                    "type": "step",
-                    "step": agent.step_count,
-                    "thinking": step_result.thinking,
-                    "action": step_result.action,
-                    "success": step_result.success,
-                    "finished": step_result.finished,
+            model_config, agent_config = manager.get_config(device_id)
+        except AgentNotInitializedError:
+            manager.release_device(device_id)
+            raise HTTPException(
+                status_code=400,
+                detail=f"Device {device_id} config not found.",
+            )
+
+        def event_generator():
+            """SSE 事件生成器"""
+            threads: list[threading.Thread] = []
+            stop_event = threading.Event()
+
+            # Register stop_event to global mapping for abort support
+            with _active_chats_lock:
+                _active_chats[device_id] = stop_event
+
+            try:
+                # Create a queue to collect events from the agent
+                event_queue: queue.Queue[tuple[str, Any]] = queue.Queue()
+
+                # Create a callback to handle thinking chunks
+                def on_thinking_chunk(chunk: str):
+                    """Emit thinking chunks as they arrive"""
+                    if not stop_event.is_set():
+                        chunk_data = {
+                            "type": "thinking_chunk",
+                            "chunk": chunk,
+                        }
+                        event_queue.put(("thinking_chunk", chunk_data))
+
+                # Create a new agent instance
+                streaming_agent = PhoneAgent(
+                    model_config=model_config,
+                    agent_config=agent_config,
+                    takeover_callback=non_blocking_takeover,
+                )
+
+                # Copy context from original agent (thread-safe due to device lock)
+                streaming_agent._context = original_agent._context.copy()
+                streaming_agent._step_count = original_agent._step_count
+
+                # Monkey-patch the model_client.request to inject the callback
+                original_request = streaming_agent.model_client.request
+
+                def patched_request(messages, **kwargs):
+                    # Inject the on_thinking_chunk callback
+                    return original_request(
+                        messages, on_thinking_chunk=on_thinking_chunk
+                    )
+
+                streaming_agent.model_client.request = patched_request
+
+                # Early abort check (before starting any steps)
+                if stop_event.is_set():
+                    logger.info(
+                        f"[Abort] Agent for device {device_id} received abort signal before starting steps"
+                    )
+                    yield "event: aborted\n"
+                    yield 'data: {"type": "aborted", "message": "Chat aborted by user"}\n\n'
+                    return
+
+                # Run agent step in a separate thread
+                step_result: list[Any] = [None]
+                error_result: list[Any] = [None]
+
+                def run_step(is_first: bool = True, task: str | None = None):
+                    try:
+                        # Check before starting step
+                        if stop_event.is_set():
+                            logger.info(
+                                f"[Abort] Agent for device {device_id} received abort signal before step execution"
+                            )
+                            return
+
+                        if is_first:
+                            result = streaming_agent.step(task)
+                        else:
+                            result = streaming_agent.step()
+
+                        # Check after step completes
+                        if stop_event.is_set():
+                            logger.info(
+                                f"[Abort] Agent for device {device_id} received abort signal after step execution"
+                            )
+                            return
+
+                        step_result[0] = result
+                    except Exception as e:
+                        error_result[0] = e
+                    finally:
+                        event_queue.put(("step_done", None))
+
+                # Start first step
+                thread = threading.Thread(
+                    target=run_step, args=(True, request.message), daemon=True
+                )
+                thread.start()
+                threads.append(thread)
+
+                while not stop_event.is_set():
+                    # Wait for events from the queue
+                    try:
+                        event_type, event_data = event_queue.get(timeout=0.1)
+                    except queue.Empty:
+                        # Check again on timeout
+                        if stop_event.is_set():
+                            break
+                        continue
+
+                    if event_type == "thinking_chunk":
+                        yield "event: thinking_chunk\n"
+                        yield f"data: {json.dumps(event_data, ensure_ascii=False)}\n\n"
+
+                    elif event_type == "step_done":
+                        # Check for errors
+                        if error_result[0]:
+                            raise error_result[0]
+
+                        result = step_result[0]
+                        event_data = {
+                            "type": "step",
+                            "step": streaming_agent.step_count,
+                            "thinking": result.thinking,
+                            "action": result.action,
+                            "success": result.success,
+                            "finished": result.finished,
+                        }
+
+                        yield "event: step\n"
+                        yield f"data: {json.dumps(event_data, ensure_ascii=False)}\n\n"
+
+                        if result.finished:
+                            done_data = {
+                                "type": "done",
+                                "message": result.message,
+                                "steps": streaming_agent.step_count,
+                                "success": result.success,
+                            }
+                            yield "event: done\n"
+                            yield f"data: {json.dumps(done_data, ensure_ascii=False)}\n\n"
+                            break
+
+                        if (
+                            streaming_agent.step_count
+                            >= streaming_agent.agent_config.max_steps
+                        ):
+                            done_data = {
+                                "type": "done",
+                                "message": "Max steps reached",
+                                "steps": streaming_agent.step_count,
+                                "success": result.success,
+                            }
+                            yield "event: done\n"
+                            yield f"data: {json.dumps(done_data, ensure_ascii=False)}\n\n"
+                            break
+
+                        # Start next step
+                        step_result[0] = None
+                        error_result[0] = None
+                        thread = threading.Thread(
+                            target=run_step, args=(False, None), daemon=True
+                        )
+                        thread.start()
+                        threads.append(thread)
+
+                # Check if loop exited due to abort
+                if stop_event.is_set():
+                    logger.info(
+                        f"[Abort] Agent for device {device_id} event loop terminated due to abort signal"
+                    )
+                    yield "event: aborted\n"
+                    yield 'data: {"type": "aborted", "message": "Chat aborted by user"}\n\n'
+
+                # Update original agent state (thread-safe due to device lock)
+                original_agent._context = streaming_agent._context
+                original_agent._step_count = streaming_agent._step_count
+
+                original_agent.reset()
+
+            except Exception as e:
+                error_data = {
+                    "type": "error",
+                    "message": str(e),
                 }
-
-                yield "event: step\n"
-                yield f"data: {json.dumps(event_data, ensure_ascii=False)}\n\n"
-
-                if step_result.finished:
-                    done_data = {
-                        "type": "done",
-                        "message": step_result.message,
-                        "steps": agent.step_count,
-                        "success": step_result.success,
-                    }
-                    yield "event: done\n"
-                    yield f"data: {json.dumps(done_data, ensure_ascii=False)}\n\n"
-                    break
-
-                if agent.step_count >= agent.agent_config.max_steps:
-                    done_data = {
-                        "type": "done",
-                        "message": "Max steps reached",
-                        "steps": agent.step_count,
-                        "success": step_result.success,
-                    }
-                    yield "event: done\n"
-                    yield f"data: {json.dumps(done_data, ensure_ascii=False)}\n\n"
-                    break
-
-                step_result = agent.step()
-
-            agent.reset()
-
-        except Exception as e:
-            error_data = {
-                "type": "error",
-                "message": str(e),
-            }
-            yield "event: error\n"
-            yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"
-
-    return StreamingResponse(
-        event_generator(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "X-Accel-Buffering": "no",
-        },
-    )
+                yield "event: error\n"
+                yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"
+            finally:
+                # Clean up active chats mapping
+                with _active_chats_lock:
+                    _active_chats.pop(device_id, None)
+
+                # Signal all threads to stop
+                stop_event.set()
+
+                alive_threads = [thread for thread in threads if thread.is_alive()]
+                if alive_threads:
+                    # Release lock after background threads complete
+                    cleanup_thread = threading.Thread(
+                        target=_release_device_lock_when_done,
+                        args=(device_id, alive_threads),
+                        daemon=True,
+                    )
+                    cleanup_thread.start()
+                else:
+                    # Release lock immediately if no threads are alive
+                    manager.release_device(device_id)
+
+        return StreamingResponse(
+            event_generator(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+    except Exception:
+        # Release lock if exception occurs before generator starts
+        manager.release_device(device_id)
+        raise
 
 
 @router.get("/api/status", response_model=StatusResponse)
 def get_status(device_id: str | None = None) -> StatusResponse:
     """获取 Agent 状态和版本信息（多设备支持）。"""
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
+
+    manager = PhoneAgentManager.get_instance()
+
     if device_id is None:
         return StatusResponse(
             version=APP_VERSION,
-            initialized=len(agents) > 0,
+            initialized=len(manager.list_agents()) > 0,
             step_count=0,
         )
 
-    if device_id not in agents:
+    if not manager.is_initialized(device_id):
         return StatusResponse(
             version=APP_VERSION,
             initialized=False,
             step_count=0,
         )
 
-    agent = agents[device_id]
+    agent = manager.get_agent(device_id)
     return StatusResponse(
         version=APP_VERSION,
         initialized=True,
@@ -234,27 +458,38 @@ def get_status(device_id: str | None = None) -> StatusResponse:
 @router.post("/api/reset")
 def reset_agent(request: ResetRequest) -> dict:
     """重置 Agent 状态（多设备支持）。"""
+    from AutoGLM_GUI.exceptions import AgentNotInitializedError
+    from AutoGLM_GUI.phone_agent_manager import PhoneAgentManager
+
     device_id = request.device_id
+    manager = PhoneAgentManager.get_instance()
 
-    if device_id not in agents:
+    try:
+        manager.reset_agent(device_id)
+        return {
+            "success": True,
+            "device_id": device_id,
+            "message": f"Agent reset for device {device_id}",
+        }
+    except AgentNotInitializedError:
         raise HTTPException(status_code=404, detail=f"Device {device_id} not found")
 
-    agent = agents[device_id]
-    agent.reset()
 
-    if device_id in agent_configs:
-        model_config, agent_config = agent_configs[device_id]
-        agents[device_id] = PhoneAgent(
-            model_config=model_config,
-            agent_config=agent_config,
-            takeover_callback=non_blocking_takeover,
-        )
+@router.post("/api/chat/abort")
+def abort_chat(request: AbortRequest) -> dict:
+    """中断正在进行的对话流。"""
+    from AutoGLM_GUI.logger import logger
 
-    return {
-        "success": True,
-        "device_id": device_id,
-        "message": f"Agent reset for device {device_id}",
-    }
+    device_id = request.device_id
+
+    with _active_chats_lock:
+        if device_id in _active_chats:
+            logger.info(f"Aborting chat for device {device_id}")
+            _active_chats[device_id].set()  # 设置中断标志
+            return {"success": True, "message": "Abort requested"}
+        else:
+            logger.warning(f"No active chat found for device {device_id}")
+            return {"success": False, "message": "No active chat found"}
 
 
 @router.get("/api/config", response_model=ConfigResponse)