Improve WebUI streaming and Python tooling

backend/api/server.py

@@ -8,16 +8,20 @@ from __future__ import annotations
 
 import asyncio
 import json
+import queue
 import re
 import shutil
+import threading
 import time
+import venv
+from dataclasses import dataclass, replace
 from pathlib import Path
-from threading import Lock
+from threading import Lock, RLock
 from typing import Any
 from urllib import error, request
 
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field
 
@@ -150,8 +154,18 @@ WEBUI_SLASH_COMMAND_DESCRIPTIONS_ZH = {
 # Agent state holder
 # ---------------------------------------------------------------------------
 
+@dataclass(frozen=True)
+class AgentInstanceConfig:
+    cwd: Path
+    model: str
+    base_url: str
+    api_key: str
+    allow_shell: bool
+    allow_write: bool
+
+
 class AgentState:
-    """Holds the live agent instance plus a lock for serialized access."""
+    """Holds account-scoped agent instances, config, and execution locks."""
 
     def __init__(
         self,
@@ -164,15 +178,66 @@ class AgentState:
         allow_write: bool,
         session_directory: Path,
     ) -> None:
-        self.cwd = cwd.resolve()
         self.session_directory = session_directory
-        self._lock = Lock()
+        self._lock = RLock()
         self._agents: dict[str, LocalCodingAgent] = {}
-        self.model = model
-        self.base_url = base_url
-        self.api_key = api_key
-        self.allow_shell = allow_shell
-        self.allow_write = allow_write
+        self._run_locks: dict[str, Lock] = {}
+        self._account_configs: dict[str, AgentInstanceConfig] = {}
+        self._default_config = AgentInstanceConfig(
+            cwd=cwd.resolve(),
+            model=model,
+            base_url=base_url,
+            api_key=api_key,
+            allow_shell=allow_shell,
+            allow_write=allow_write,
+        )
+
+    @property
+    def cwd(self) -> Path:
+        return self._default_config.cwd
+
+    @property
+    def model(self) -> str:
+        return self._default_config.model
+
+    @property
+    def base_url(self) -> str:
+        return self._default_config.base_url
+
+    @property
+    def api_key(self) -> str:
+        return self._default_config.api_key
+
+    @property
+    def allow_shell(self) -> bool:
+        return self._default_config.allow_shell
+
+    @property
+    def allow_write(self) -> bool:
+        return self._default_config.allow_write
+
+    def _account_key(self, account_id: str | None) -> str:
+        return _safe_account_id(account_id) if account_id else '__default__'
+
+    def _config_for(self, account_id: str | None) -> AgentInstanceConfig:
+        if not account_id:
+            return self._default_config
+        key = self._account_key(account_id)
+        config = self._account_configs.get(key)
+        if config is None:
+            config = replace(self._default_config)
+            self._account_configs[key] = config
+        return config
+
+    def config_for(self, account_id: str | None) -> AgentInstanceConfig:
+        with self._lock:
+            return self._config_for(account_id)
+
+    def _set_config_for(self, account_id: str | None, config: AgentInstanceConfig) -> None:
+        if not account_id:
+            self._default_config = config
+            return
+        self._account_configs[self._account_key(account_id)] = config
 
     def _account_base(self, account_id: str | None) -> Path:
         if not account_id:
@@ -189,6 +254,7 @@ class AgentState:
                 'scratchpad': self.session_directory,
                 'uploads': self.session_directory,
                 'outputs': self.session_directory,
+                'python_env': base / 'python' / '.venv',
             }
         base = self._account_base(account_id)
         return {
@@ -197,26 +263,32 @@ class AgentState:
             'scratchpad': base / 'sessions',
             'uploads': base / 'sessions',
             'outputs': base / 'sessions',
+            'python_env': base / 'python' / '.venv',
         }
 
     def _build_agent(self, account_id: str | None = None) -> LocalCodingAgent:
+        config = self._config_for(account_id)
         paths = self.account_paths(account_id)
         for directory in paths.values():
-            directory.mkdir(parents=True, exist_ok=True)
+            if directory.name != '.venv':
+                directory.mkdir(parents=True, exist_ok=True)
+        self._ensure_python_env(paths['python_env'])
         permissions = AgentPermissions(
-            allow_file_write=self.allow_write,
-            allow_shell_commands=self.allow_shell,
+            allow_file_write=config.allow_write,
+            allow_shell_commands=config.allow_shell,
         )
         runtime_config = AgentRuntimeConfig(
-            cwd=self.cwd,
+            cwd=config.cwd,
             permissions=permissions,
+            stream_model_responses=True,
             session_directory=paths['sessions'],
             scratchpad_root=paths['scratchpad'],
+            python_env_dir=paths['python_env'],
         )
         model_config = ModelConfig(
-            model=self.model,
-            base_url=self.base_url,
-            api_key=self.api_key,
+            model=config.model,
+            base_url=config.base_url,
+            api_key=config.api_key,
         )
         return LocalCodingAgent(
             model_config=model_config,
@@ -224,12 +296,22 @@ class AgentState:
         )
 
     def agent_for(self, account_id: str | None = None) -> LocalCodingAgent:
-        key = _safe_account_id(account_id) if account_id else '__default__'
-        agent = self._agents.get(key)
-        if agent is None:
-            agent = self._build_agent(account_id)
-            self._agents[key] = agent
-        return agent
+        with self._lock:
+            key = self._account_key(account_id)
+            agent = self._agents.get(key)
+            if agent is None:
+                agent = self._build_agent(account_id)
+                self._agents[key] = agent
+            return agent
+
+    def run_lock_for(self, account_id: str | None = None) -> Lock:
+        key = self._account_key(account_id)
+        with self._lock:
+            lock = self._run_locks.get(key)
+            if lock is None:
+                lock = Lock()
+                self._run_locks[key] = lock
+            return lock
 
     def update(
         self,
@@ -243,45 +325,56 @@ class AgentState:
         account_id: str | None = None,
     ) -> None:
         with self._lock:
-            # account_id 只用于让调用方拿到对应账号的 snapshot，不参与全局配置更新。
-            del account_id
+            config = self._config_for(account_id)
             if model is not None:
-                self.model = _normalize_chat_model_name(model)
+                config = replace(config, model=_normalize_chat_model_name(model))
             if base_url is not None:
-                self.base_url = base_url
+                config = replace(config, base_url=base_url)
             if api_key is not None:
-                self.api_key = api_key
+                config = replace(config, api_key=api_key)
             if cwd is not None:
                 resolved = Path(cwd).expanduser().resolve()
                 if not resolved.is_dir():
                     raise ValueError(f'cwd does not exist: {resolved}')
-                self.cwd = resolved
+                config = replace(config, cwd=resolved)
             if allow_shell is not None:
-                self.allow_shell = allow_shell
+                config = replace(config, allow_shell=allow_shell)
             if allow_write is not None:
-                self.allow_write = allow_write
-            self._agents.clear()
+                config = replace(config, allow_write=allow_write)
+            self._set_config_for(account_id, config)
+            self._agents.pop(self._account_key(account_id), None)
 
     def snapshot(self, account_id: str | None = None) -> dict[str, Any]:
-        agent = self.agent_for(account_id)
-        paths = self.account_paths(account_id)
-        return {
-            'model': self.model,
-            'base_url': self.base_url,
-            'cwd': str(self.cwd),
-            'account_id': _safe_account_id(account_id) if account_id else None,
-            'account_directory': str(paths['base']),
-            'session_directory': str(paths['sessions']),
-            'upload_directory': str(paths['uploads']),
+        with self._lock:
+            config = self._config_for(account_id)
+            agent = self.agent_for(account_id)
+            paths = self.account_paths(account_id)
+            return {
+                'model': config.model,
+                'base_url': config.base_url,
+                'cwd': str(config.cwd),
+                'account_id': _safe_account_id(account_id) if account_id else None,
+                'account_directory': str(paths['base']),
+                'session_directory': str(paths['sessions']),
+                'upload_directory': str(paths['uploads']),
             'output_directory': str(paths['outputs']),
-            'allow_shell': self.allow_shell,
-            'allow_write': self.allow_write,
-            'active_session_id': agent.active_session_id,
-        }
+            'python_env_directory': str(paths['python_env']),
+            'allow_shell': config.allow_shell,
+                'allow_write': config.allow_write,
+                'active_session_id': agent.active_session_id,
+            }
 
     def lock(self) -> Lock:
         return self._lock
 
+    def _ensure_python_env(self, env_dir: Path) -> None:
+        python_bin = env_dir / 'bin' / 'python'
+        pip_bin = env_dir / 'bin' / 'pip'
+        if python_bin.exists() and pip_bin.exists():
+            return
+        env_dir.parent.mkdir(parents=True, exist_ok=True)
+        venv.EnvBuilder(with_pip=True, symlinks=False).create(env_dir)
+
 
 # ---------------------------------------------------------------------------
 # Request models
@@ -308,6 +401,7 @@ class StateUpdate(BaseModel):
 class ModelListRequest(BaseModel):
     base_url: str | None = None
     api_key: str | None = None
+    account_id: str | None = None
 
 
 class SessionTitleUpdate(BaseModel):
@@ -366,7 +460,8 @@ def create_app(state: AgentState) -> FastAPI:
         return commands
 
     @app.get('/api/skills')
-    async def list_skills() -> list[dict[str, Any]]:
+    async def list_skills(account_id: str | None = None) -> list[dict[str, Any]]:
+        config = state.config_for(account_id)
         return [
             {
                 'name': skill.name,
@@ -375,19 +470,21 @@ def create_app(state: AgentState) -> FastAPI:
                 'aliases': list(skill.aliases),
                 'allowed_tools': list(skill.allowed_tools),
             }
-            for skill in get_bundled_skills(state.cwd)
+            for skill in get_bundled_skills(config.cwd)
             if skill.user_invocable
         ]
 
     @app.get('/api/models')
-    async def list_models_get() -> dict[str, Any]:
-        return _list_backend_models(state.base_url, state.api_key)
+    async def list_models_get(account_id: str | None = None) -> dict[str, Any]:
+        config = state.config_for(account_id)
+        return _list_backend_models(config.base_url, config.api_key)
 
     @app.post('/api/models')
     async def list_models_post(payload: ModelListRequest) -> dict[str, Any]:
+        config = state.config_for(payload.account_id)
         return _list_backend_models(
-            payload.base_url or state.base_url,
-            payload.api_key or state.api_key,
+            payload.base_url or config.base_url,
+            payload.api_key or config.api_key,
         )
 
     # ------------- sessions --------------------------------------------------
@@ -492,6 +589,67 @@ def create_app(state: AgentState) -> FastAPI:
             )
         return _serialize_token_budget(snapshot)
 
+    def _run_chat_payload(
+        request: ChatRequest,
+        event_sink: Any | None = None,
+    ) -> dict[str, Any]:
+        run_lock = state.run_lock_for(request.account_id)
+        with run_lock:
+            agent = state.agent_for(request.account_id)
+            config = state.config_for(request.account_id)
+            session_directory = state.account_paths(request.account_id)['sessions']
+            requested_session_id = _safe_session_id(
+                request.resume_session_id or request.session_id
+            )
+            previous_title = _read_session_title(
+                session_directory,
+                requested_session_id,
+            )
+            started_at = time.perf_counter()
+            if request.resume_session_id is not None:
+                try:
+                    stored = load_agent_session(
+                        request.resume_session_id,
+                        directory=session_directory,
+                    )
+                except FileNotFoundError:
+                    raise HTTPException(
+                        status_code=404,
+                        detail='Session to resume not found',
+                    )
+                result = agent.resume(
+                    request.prompt.strip(),
+                    stored,
+                    runtime_context=request.runtime_context,
+                    event_sink=event_sink,
+                )
+            else:
+                result = agent.run(
+                    request.prompt.strip(),
+                    session_id=_safe_session_id(request.session_id),
+                    runtime_context=request.runtime_context,
+                    event_sink=event_sink,
+                )
+            elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
+            payload = _serialize_run_result(result)
+            payload['elapsed_ms'] = elapsed_ms
+            if result.session_id:
+                _annotate_last_assistant_elapsed(
+                    session_directory,
+                    result.session_id,
+                    elapsed_ms,
+                )
+                _ensure_session_title(
+                    session_directory,
+                    result.session_id,
+                    model=config.model,
+                    base_url=config.base_url,
+                    api_key=config.api_key,
+                    previous_title=previous_title,
+                )
+                _annotate_transcript_elapsed(payload, elapsed_ms)
+            return payload
+
     # ------------- chat ------------------------------------------------------
     @app.post('/api/chat')
     async def chat(request: ChatRequest) -> dict[str, Any]:
@@ -500,58 +658,7 @@ def create_app(state: AgentState) -> FastAPI:
             raise HTTPException(status_code=400, detail='Prompt is empty')
 
         def _run() -> dict[str, Any]:
-            with state.lock():
-                agent = state.agent_for(request.account_id)
-                session_directory = state.account_paths(request.account_id)['sessions']
-                requested_session_id = _safe_session_id(
-                    request.resume_session_id or request.session_id
-                )
-                previous_title = _read_session_title(
-                    session_directory,
-                    requested_session_id,
-                )
-                started_at = time.perf_counter()
-                if request.resume_session_id is not None:
-                    try:
-                        stored = load_agent_session(
-                            request.resume_session_id,
-                            directory=session_directory,
-                        )
-                    except FileNotFoundError:
-                        raise HTTPException(
-                            status_code=404,
-                            detail='Session to resume not found',
-                        )
-                    result = agent.resume(
-                        prompt,
-                        stored,
-                        runtime_context=request.runtime_context,
-                    )
-                else:
-                    result = agent.run(
-                        prompt,
-                        session_id=_safe_session_id(request.session_id),
-                        runtime_context=request.runtime_context,
-                    )
-                elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
-                payload = _serialize_run_result(result)
-                payload['elapsed_ms'] = elapsed_ms
-                if result.session_id:
-                    _annotate_last_assistant_elapsed(
-                        session_directory,
-                        result.session_id,
-                        elapsed_ms,
-                    )
-                    _ensure_session_title(
-                        session_directory,
-                        result.session_id,
-                        model=state.model,
-                        base_url=state.base_url,
-                        api_key=state.api_key,
-                        previous_title=previous_title,
-                    )
-                    _annotate_transcript_elapsed(payload, elapsed_ms)
-                return payload
+            return _run_chat_payload(request)
 
         try:
             payload = await asyncio.to_thread(_run)
@@ -567,6 +674,53 @@ def create_app(state: AgentState) -> FastAPI:
             )
         return payload
 
+    @app.post('/api/chat/stream')
+    async def chat_stream(request: ChatRequest) -> StreamingResponse:
+        prompt = request.prompt.strip()
+        if not prompt:
+            raise HTTPException(status_code=400, detail='Prompt is empty')
+
+        event_queue: queue.Queue[dict[str, Any] | None] = queue.Queue()
+
+        def emit_event(event: dict[str, object]) -> None:
+            event_queue.put({'kind': 'event', 'event': event})
+
+        def worker() -> None:
+            try:
+                payload = _run_chat_payload(request, event_sink=emit_event)
+            except HTTPException as exc:
+                event_queue.put(
+                    {
+                        'kind': 'error',
+                        'status_code': exc.status_code,
+                        'detail': exc.detail,
+                    }
+                )
+            except Exception as exc:
+                event_queue.put(
+                    {
+                        'kind': 'error',
+                        'status_code': 500,
+                        'error': str(exc),
+                        'error_type': type(exc).__name__,
+                    }
+                )
+            else:
+                event_queue.put({'kind': 'result', 'payload': payload})
+            finally:
+                event_queue.put(None)
+
+        threading.Thread(target=worker, daemon=True).start()
+
+        def generate() -> Any:
+            while True:
+                item = event_queue.get()
+                if item is None:
+                    break
+                yield json.dumps(item, ensure_ascii=False) + '\n'
+
+        return StreamingResponse(generate(), media_type='application/x-ndjson')
+
     @app.post('/api/clear')
     async def clear_state(account_id: str | None = None) -> dict[str, Any]:
         with state.lock():