Add assistant-ui data agent frontend

backend/api/server.py

@@ -8,9 +8,12 @@ from __future__ import annotations
 
 import asyncio
 import json
+import re
+import time
 from pathlib import Path
 from threading import Lock
 from typing import Any
+from urllib import error, request
 
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import FileResponse, JSONResponse
@@ -56,15 +59,42 @@ class AgentState:
         self.cwd = cwd.resolve()
         self.session_directory = session_directory
         self._lock = Lock()
-        self._agent: LocalCodingAgent | None = None
+        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._build_agent()
 
-    def _build_agent(self) -> None:
+    def _account_base(self, account_id: str | None) -> Path:
+        if not account_id:
+            return self.session_directory.parent
+        safe_id = _safe_account_id(account_id)
+        return (self.session_directory.parent / 'accounts' / safe_id).resolve()
+
+    def account_paths(self, account_id: str | None) -> dict[str, Path]:
+        if not account_id:
+            base = self.session_directory.parent
+            return {
+                'base': base,
+                'sessions': self.session_directory,
+                'scratchpad': self.session_directory,
+                'uploads': self.session_directory,
+                'outputs': self.session_directory,
+            }
+        base = self._account_base(account_id)
+        return {
+            'base': base,
+            'sessions': base / 'sessions',
+            'scratchpad': base / 'sessions',
+            'uploads': base / 'sessions',
+            'outputs': base / 'sessions',
+        }
+
+    def _build_agent(self, account_id: str | None = None) -> LocalCodingAgent:
+        paths = self.account_paths(account_id)
+        for directory in paths.values():
+            directory.mkdir(parents=True, exist_ok=True)
         permissions = AgentPermissions(
             allow_file_write=self.allow_write,
             allow_shell_commands=self.allow_shell,
@@ -72,22 +102,26 @@ class AgentState:
         runtime_config = AgentRuntimeConfig(
             cwd=self.cwd,
             permissions=permissions,
-            session_directory=self.session_directory,
+            session_directory=paths['sessions'],
+            scratchpad_root=paths['scratchpad'],
         )
         model_config = ModelConfig(
             model=self.model,
             base_url=self.base_url,
             api_key=self.api_key,
         )
-        self._agent = LocalCodingAgent(
+        return LocalCodingAgent(
             model_config=model_config,
             runtime_config=runtime_config,
         )
 
-    @property
-    def agent(self) -> LocalCodingAgent:
-        assert self._agent is not None
-        return self._agent
+    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
 
     def update(
         self,
@@ -98,10 +132,13 @@ class AgentState:
         cwd: str | None = None,
         allow_shell: bool | None = None,
         allow_write: bool | None = None,
+        account_id: str | None = None,
     ) -> None:
         with self._lock:
+            # account_id 只用于让调用方拿到对应账号的 snapshot，不参与全局配置更新。
+            del account_id
             if model is not None:
-                self.model = model
+                self.model = _normalize_chat_model_name(model)
             if base_url is not None:
                 self.base_url = base_url
             if api_key is not None:
@@ -115,17 +152,23 @@ class AgentState:
                 self.allow_shell = allow_shell
             if allow_write is not None:
                 self.allow_write = allow_write
-            self._build_agent()
+            self._agents.clear()
 
-    def snapshot(self) -> dict[str, Any]:
+    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),
-            'session_directory': str(self.session_directory),
+            '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': self.agent.active_session_id,
+            'active_session_id': agent.active_session_id,
         }
 
     def lock(self) -> Lock:
@@ -138,7 +181,10 @@ class AgentState:
 
 class ChatRequest(BaseModel):
     prompt: str = Field(min_length=1)
+    runtime_context: str | None = None
     resume_session_id: str | None = None
+    account_id: str | None = None
+    session_id: str | None = None
 
 
 class StateUpdate(BaseModel):
@@ -148,6 +194,12 @@ class StateUpdate(BaseModel):
     cwd: str | None = None
     allow_shell: bool | None = None
     allow_write: bool | None = None
+    account_id: str | None = None
+
+
+class ModelListRequest(BaseModel):
+    base_url: str | None = None
+    api_key: str | None = None
 
 
 # ---------------------------------------------------------------------------
@@ -170,8 +222,8 @@ def create_app(state: AgentState) -> FastAPI:
 
     # ------------- info ------------------------------------------------------
     @app.get('/api/state')
-    async def get_state() -> dict[str, Any]:
-        return state.snapshot()
+    async def get_state(account_id: str | None = None) -> dict[str, Any]:
+        return state.snapshot(account_id)
 
     @app.post('/api/state')
     async def post_state(payload: StateUpdate) -> dict[str, Any]:
@@ -179,7 +231,7 @@ def create_app(state: AgentState) -> FastAPI:
             state.update(**payload.model_dump(exclude_none=True))
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc))
-        return state.snapshot()
+        return state.snapshot(payload.account_id)
 
     @app.get('/api/slash-commands')
     async def list_slash_commands() -> list[dict[str, Any]]:
@@ -208,45 +260,62 @@ def create_app(state: AgentState) -> FastAPI:
             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)
+
+    @app.post('/api/models')
+    async def list_models_post(payload: ModelListRequest) -> dict[str, Any]:
+        return _list_backend_models(
+            payload.base_url or state.base_url,
+            payload.api_key or state.api_key,
+        )
+
     # ------------- sessions --------------------------------------------------
     @app.get('/api/sessions')
-    async def list_sessions() -> list[dict[str, Any]]:
-        directory = state.session_directory
+    async def list_sessions(account_id: str | None = None) -> list[dict[str, Any]]:
+        directory = state.account_paths(account_id)['sessions']
         if not directory.exists():
             return []
         results: list[dict[str, Any]] = []
-        for path in sorted(
-            directory.glob('*.json'),
-            key=lambda p: p.stat().st_mtime,
-            reverse=True,
-        ):
+        for path in sorted(_iter_session_files(directory), key=_session_mtime, reverse=True):
             try:
                 data = json.loads(path.read_text(encoding='utf-8'))
             except (OSError, json.JSONDecodeError):
                 continue
+            session_id = str(data.get('session_id', _session_id_from_path(path)))
+            usage = data.get('usage') if isinstance(data.get('usage'), dict) else {}
+            model_config = (
+                data.get('model_config')
+                if isinstance(data.get('model_config'), dict)
+                else {}
+            )
             messages = data.get('messages') or []
             preview = ''
             for msg in messages:
                 if isinstance(msg, dict) and msg.get('role') == 'user':
                     content = msg.get('content', '')
-                    if isinstance(content, str):
-                        preview = content[:120]
+                    if isinstance(content, str) and not _is_internal_message(content):
+                        preview = _strip_session_context(content)[:120]
                         break
             results.append(
                 {
-                    'session_id': data.get('session_id', path.stem),
+                    'session_id': session_id,
                     'turns': data.get('turns', 0),
                     'tool_calls': data.get('tool_calls', 0),
                     'preview': preview,
-                    'modified_at': path.stat().st_mtime,
+                    'modified_at': _session_mtime(path),
+                    'model': model_config.get('model'),
+                    'usage': usage,
                 }
             )
         return results
 
     @app.get('/api/sessions/{session_id}')
-    async def get_session(session_id: str) -> dict[str, Any]:
+    async def get_session(session_id: str, account_id: str | None = None) -> dict[str, Any]:
+        directory = state.account_paths(account_id)['sessions']
         try:
-            stored = load_agent_session(session_id, directory=state.session_directory)
+            stored = load_agent_session(session_id, directory=directory)
         except FileNotFoundError:
             raise HTTPException(status_code=404, detail='Session not found')
         return _serialize_stored_session(stored)
@@ -260,22 +329,42 @@ def create_app(state: AgentState) -> FastAPI:
 
         def _run() -> dict[str, Any]:
             with state.lock():
-                agent = state.agent
+                agent = state.agent_for(request.account_id)
+                session_directory = state.account_paths(request.account_id)['sessions']
+                started_at = time.perf_counter()
                 if request.resume_session_id is not None:
                     try:
                         stored = load_agent_session(
                             request.resume_session_id,
-                            directory=state.session_directory,
+                            directory=session_directory,
                         )
                     except FileNotFoundError:
                         raise HTTPException(
                             status_code=404,
                             detail='Session to resume not found',
                         )
-                    result = agent.resume(prompt, stored)
+                    result = agent.resume(
+                        prompt,
+                        stored,
+                        runtime_context=request.runtime_context,
+                    )
                 else:
-                    result = agent.run(prompt)
-                return _serialize_run_result(result)
+                    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,
+                    )
+                    _annotate_transcript_elapsed(payload, elapsed_ms)
+                return payload
 
         try:
             payload = await asyncio.to_thread(_run)
@@ -292,10 +381,10 @@ def create_app(state: AgentState) -> FastAPI:
         return payload
 
     @app.post('/api/clear')
-    async def clear_state() -> dict[str, Any]:
+    async def clear_state(account_id: str | None = None) -> dict[str, Any]:
         with state.lock():
-            state.agent.clear_runtime_state()
-        return state.snapshot()
+            state.agent_for(account_id).clear_runtime_state()
+        return state.snapshot(account_id)
 
     return app
 
@@ -338,3 +427,182 @@ def _serialize_stored_session(stored: StoredAgentSession) -> dict[str, Any]:
         'total_cost_usd': stored.total_cost_usd,
         'model': stored.model_config.get('model'),
     }
+
+
+def _list_backend_models(base_url: str, api_key: str) -> dict[str, Any]:
+    req = request.Request(
+        _join_url(base_url, '/models'),
+        headers={
+            'Authorization': f'Bearer {api_key}',
+            'Content-Type': 'application/json',
+        },
+        method='GET',
+    )
+    try:
+        with request.urlopen(req, timeout=10) as response:
+            payload = json.loads(response.read().decode('utf-8'))
+    except error.HTTPError as exc:
+        detail = exc.read().decode('utf-8', errors='replace')
+        raise HTTPException(
+            status_code=502,
+            detail=f'HTTP {exc.code} from model backend: {detail}',
+        ) from exc
+    except (error.URLError, OSError, json.JSONDecodeError) as exc:
+        raise HTTPException(
+            status_code=502,
+            detail=f'Unable to list models from {base_url}: {exc}',
+        ) from exc
+
+    models = _normalize_model_list(payload)
+    return {'models': models, 'raw_count': len(models)}
+
+
+def _normalize_model_list(payload: Any) -> list[dict[str, str]]:
+    data = payload.get('data') if isinstance(payload, dict) else payload
+    if not isinstance(data, list):
+        return []
+    model_ids: dict[str, str] = {}
+    for item in data:
+        if isinstance(item, str):
+            normalized = _normalize_model_option(item)
+            if normalized is not None:
+                model_ids.setdefault(normalized.lower(), normalized)
+            continue
+        if not isinstance(item, dict):
+            continue
+        model_id = item.get('id') or item.get('model') or item.get('name')
+        if not isinstance(model_id, str) or not model_id:
+            continue
+        normalized = _normalize_model_option(model_id)
+        if normalized is not None:
+            model_ids.setdefault(normalized.lower(), normalized)
+    return [{'id': model_id} for model_id in sorted(model_ids.values(), key=str.lower)]
+
+
+def _normalize_model_option(model_id: str) -> str | None:
+    normalized = _normalize_chat_model_name(model_id)
+    lower = normalized.lower()
+    if lower.startswith('xiaomi/mimo-v2') and not any(
+        blocked in lower for blocked in ('audio', 'asr', 'tts', 'voiceclone', 'voicedesign')
+    ):
+        return normalized
+    known_chat_models = {
+        'xiaomi/deepseek-r1-0528',
+        'xiaomi/minimax-m2.5',
+        'xiaomi/qwen3-32b',
+    }
+    if lower in known_chat_models:
+        return normalized
+    return None
+
+
+def _normalize_chat_model_name(model: str) -> str:
+    value = model.strip()
+    lower = value.lower()
+    if lower.startswith('mimo-v2') or lower in {
+        'deepseek-r1-0528',
+        'minimax-m2.5',
+        'qwen3-32b',
+    }:
+        return f'xiaomi/{value}'
+    return value
+
+
+def _join_url(base_url: str, suffix: str) -> str:
+    return f'{base_url.rstrip("/")}/{suffix.lstrip("/")}'
+
+
+def _annotate_transcript_elapsed(payload: dict[str, Any], elapsed_ms: int) -> None:
+    transcript = payload.get('transcript')
+    if not isinstance(transcript, list):
+        return
+    for entry in reversed(transcript):
+        if not isinstance(entry, dict) or entry.get('role') != 'assistant':
+            continue
+        metadata = entry.get('metadata')
+        if not isinstance(metadata, dict):
+            metadata = {}
+            entry['metadata'] = metadata
+        metadata['elapsed_ms'] = elapsed_ms
+        return
+
+
+def _annotate_last_assistant_elapsed(
+    directory: Path,
+    session_id: str,
+    elapsed_ms: int,
+) -> None:
+    path = directory / session_id / 'session.json'
+    if not path.exists():
+        path = directory / f'{session_id}.json'
+    try:
+        data = json.loads(path.read_text(encoding='utf-8'))
+    except (OSError, json.JSONDecodeError):
+        return
+    messages = data.get('messages')
+    if not isinstance(messages, list):
+        return
+    for message in reversed(messages):
+        if not isinstance(message, dict) or message.get('role') != 'assistant':
+            continue
+        metadata = message.get('metadata')
+        if not isinstance(metadata, dict):
+            metadata = {}
+            message['metadata'] = metadata
+        metadata['elapsed_ms'] = elapsed_ms
+        try:
+            path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
+        except OSError:
+            return
+        return
+
+
+def _safe_account_id(account_id: str | None) -> str:
+    normalized = re.sub(r'[^a-zA-Z0-9._-]+', '_', (account_id or '').strip())
+    return normalized[:80] or 'default'
+
+
+def _is_internal_message(content: str) -> bool:
+    return content.lstrip().startswith('<system-reminder>')
+
+
+def _strip_session_context(content: str) -> str:
+    marker = '\n\n[当前会话目录]'
+    if marker in content:
+        return content.split(marker, 1)[0].strip()
+    marker = '\n[当前会话目录]'
+    if marker in content:
+        return content.split(marker, 1)[0].strip()
+    return content.strip()
+
+
+def _iter_session_files(directory: Path) -> list[Path]:
+    if not directory.exists():
+        return []
+    by_session_id: dict[str, Path] = {}
+    for path in directory.glob('*.json'):
+        by_session_id[_session_id_from_path(path)] = path
+    for path in directory.glob('*/session.json'):
+        # 新目录格式包含 input/output/scratchpad，优先级高于旧的平铺 json。
+        by_session_id[_session_id_from_path(path)] = path
+    return list(by_session_id.values())
+
+
+def _session_mtime(path: Path) -> float:
+    try:
+        return path.stat().st_mtime
+    except OSError:
+        return 0.0
+
+
+def _session_id_from_path(path: Path) -> str:
+    if path.name == 'session.json':
+        return path.parent.name
+    return path.stem
+
+
+def _safe_session_id(session_id: str | None) -> str | None:
+    if not session_id:
+        return None
+    normalized = re.sub(r'[^a-zA-Z0-9._-]+', '_', session_id.strip())
+    return normalized[:120] or None