"""Partners management API. A partner is an IM-connected companion driven by the chat agent loop. This router owns: partner CRUD + lifecycle, the soul library, channel config (schema-driven), asset provisioning (KB / skills / notebooks copied into the partner workspace), tool configuration, history, and the web chat entry points (HTTP / SSE / WebSocket). """ from __future__ import annotations import asyncio import base64 import binascii import json import logging from typing import Any, AsyncGenerator, Literal from uuid import uuid4 from fastapi import APIRouter, HTTPException, Query, WebSocket, WebSocketDisconnect from fastapi.responses import StreamingResponse from pydantic import BaseModel, ConfigDict, Field, ValidationError from deeptutor.core.i18n import t from deeptutor.partners.config.paths import get_partner_media_dir from deeptutor.partners.helpers import safe_filename from deeptutor.services.partners import ( get_partner_manager, slugify_partner_id, slugify_soul_id, ) from deeptutor.services.partners.manager import ( LEGACY_GLOBAL_DELIVERY_KEYS, PartnerConfig, PartnerInstance, mask_channel_secrets, strip_legacy_global_delivery, ) from deeptutor.services.partners.workspace import ( list_assets, provision_assets, read_soul, remove_asset, strip_frontmatter, write_soul, ) logger = logging.getLogger(__name__) router = APIRouter() # Per-partner async locks used to dedupe concurrent WebSocket-driven # auto-starts (start_partner short-circuits when running, but that check is # not async-safe under concurrent connections). _start_locks: dict[str, asyncio.Lock] = {} _start_locks_mutex = asyncio.Lock() async def _get_start_lock(partner_id: str) -> asyncio.Lock: async with _start_locks_mutex: lock = _start_locks.get(partner_id) if lock is None: lock = asyncio.Lock() _start_locks[partner_id] = lock return lock async def _ensure_running_partner( partner_id: str, *, allow_stopped: bool = False, ) -> PartnerInstance: mgr = get_partner_manager() instance = mgr.get_partner(partner_id) if instance and instance.running: return instance config = mgr.load_config(partner_id) if config is None: raise HTTPException(status_code=404, detail=t("api.partner_not_found")) if not allow_stopped and not mgr.auto_start_enabled(partner_id, default=False): raise HTTPException(status_code=409, detail=t("api.partner_stopped_start_required")) lock = await _get_start_lock(partner_id) async with lock: instance = mgr.get_partner(partner_id) if instance and instance.running: return instance if not allow_stopped and not mgr.auto_start_enabled(partner_id, default=False): raise HTTPException(status_code=409, detail=t("api.partner_stopped_start_required")) try: return await mgr.start_partner(partner_id, config) except RuntimeError as exc: raise HTTPException(status_code=500, detail=str(exc)) from None except Exception as exc: logger.exception("Failed to auto-start partner '%s'", partner_id) raise HTTPException(status_code=500, detail="Failed to start partner") from exc # ── Request models ───────────────────────────────────────────── class SoulSpec(BaseModel): """Where a new partner's soul comes from.""" source: Literal["default", "library", "persona", "custom"] = "default" id: str | None = None # library soul id, or persona name content: str | None = None # custom markdown class AssetSpec(BaseModel): knowledge_bases: list[str] = Field(default_factory=list) skills: list[str] = Field(default_factory=list) notebooks: list[str] = Field(default_factory=list) class CreatePartnerRequest(BaseModel): partner_id: str | None = None name: str = Field(..., min_length=1) description: str | None = None soul: SoulSpec | None = None channels: dict | None = None llm_selection: dict[str, str] | None = None backup_llm_selection: dict[str, str] | None = None language: str | None = None emoji: str | None = None color: str | None = None avatar: str | None = None enabled_tools: list[str] | None = None builtin_tools: list[str] | None = None mcp_tools: list[str] | None = None assets: AssetSpec | None = None start: bool = True class UpdatePartnerRequest(BaseModel): name: str | None = None description: str | None = None channels: dict | None = None llm_selection: dict[str, str] | None = None backup_llm_selection: dict[str, str] | None = None language: str | None = None emoji: str | None = None color: str | None = None avatar: str | None = None enabled_tools: list[str] | None = None builtin_tools: list[str] | None = None mcp_tools: list[str] | None = None class SoulUpdateBody(BaseModel): content: str class AssetAddRequest(AssetSpec): pass class ChatAttachmentRequest(BaseModel): type: str = "file" url: str = "" base64: str = "" filename: str = "" mime_type: str = "" class ChatMessageRequest(BaseModel): model_config = ConfigDict(populate_by_name=True) content: str = "" session_id: str | None = None session_key: str | None = None chat_id: str | None = None attachments: list[ChatAttachmentRequest] = Field(default_factory=list) llm_selection: dict[str, str] | None = Field(default=None, alias="llmSelection") class SessionKeyBody(BaseModel): session_key: str = Field(..., min_length=1) class SessionBranchBody(BaseModel): source_key: str = Field(..., min_length=1) new_key: str = Field(..., min_length=1) class SoulCreateRequest(BaseModel): id: str name: str content: str class SoulTemplateUpdateRequest(BaseModel): name: str | None = None content: str | None = None # ── Validation helpers ───────────────────────────────────────── def _validate_channels_payload(channels: dict) -> None: """Reject malformed channel configs at the API boundary (422).""" from deeptutor.partners.config.schema import ChannelsConfig legacy_keys = sorted(k for k in channels if k in LEGACY_GLOBAL_DELIVERY_KEYS) if legacy_keys: raise HTTPException( status_code=422, detail={ "message": ( "Delivery flags are configured per channel; remove top-level " f"channel keys: {', '.join(legacy_keys)}" ) }, ) try: ChannelsConfig(**channels) except ValidationError as exc: raise HTTPException( status_code=422, detail={"message": t("api.invalid_channels_config"), "errors": exc.errors()}, ) from None except TypeError as exc: raise HTTPException( status_code=422, detail=f"{t('api.invalid_channels_config')}: {exc}", ) from None # Inline avatars are client-resized to ~128px before upload; this cap is a # server-side backstop so config.yaml can't be bloated with raw photos. _AVATAR_MAX_CHARS = 200_000 def _validate_avatar_payload(value: str | None) -> str: avatar = (value or "").strip() if not avatar: return "" if not avatar.startswith("data:image/"): raise HTTPException(status_code=422, detail="Avatar must be a data:image/* URL") if len(avatar) > _AVATAR_MAX_CHARS: raise HTTPException( status_code=422, detail="Avatar too large — resize the image before uploading", ) return avatar def _validate_llm_selection_payload( value: dict[str, str] | None, ) -> dict[str, str] | None: """Validate a partner model selection against the shared LLM catalog.""" from deeptutor.services.config import get_model_catalog_service from deeptutor.services.model_selection import apply_llm_selection_to_catalog from deeptutor.services.partners.model_runtime import normalize_partner_llm_selection try: selection = normalize_partner_llm_selection(value) if selection: apply_llm_selection_to_catalog(get_model_catalog_service().load(), selection) return selection except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) from None def _resolve_soul_content(soul: SoulSpec | None) -> tuple[str, dict[str, str]]: """Resolve a SoulSpec into (markdown content, origin record).""" from deeptutor.services.partners.workspace import DEFAULT_SOUL if soul is None or soul.source == "default": return DEFAULT_SOUL, {"type": "default", "id": ""} if soul.source == "custom": content = (soul.content or "").strip() if not content: raise HTTPException(status_code=422, detail=t("api.soul_content_empty")) return content, {"type": "custom", "id": ""} if soul.source == "library": entry = get_partner_manager().get_soul(str(soul.id or "")) if not entry: raise HTTPException( status_code=404, detail=t("api.soul_library_not_found", name=str(soul.id)), ) return str(entry.get("content") or ""), {"type": "library", "id": str(soul.id)} # source == "persona": clone from the chat persona workspace (the # requesting user's personas first; non-admins fall back to admin # presets, mirroring chat's resolution). name = str(soul.id or "").strip() if not name: raise HTTPException(status_code=422, detail=t("api.persona_name_required")) content = _load_persona_markdown(name) if not content: raise HTTPException(status_code=404, detail=t("api.persona_not_found", name=name)) return content, {"type": "persona", "id": name} def _load_persona_markdown(name: str) -> str: from deeptutor.multi_user.context import get_current_user from deeptutor.multi_user.paths import get_admin_path_service from deeptutor.services.persona import PersonaService, get_persona_service try: detail = get_persona_service().get_detail(name) return strip_frontmatter(detail.content) except Exception: pass try: if not get_current_user().is_admin: admin_service = PersonaService( root=get_admin_path_service().get_workspace_dir() / "personas" ) return strip_frontmatter(admin_service.get_detail(name).content) except Exception: pass return "" # ── Soul template library (before /{partner_id} routes) ─────── @router.get("/souls") async def list_souls(): return get_partner_manager().list_souls() @router.post("/souls") async def create_soul(payload: SoulCreateRequest): mgr = get_partner_manager() # Slug the id server-side (authoritative): soul ids ride in ``/souls/`` # URLs, so a raw CJK / path-unsafe id (e.g. ``我的灵魂`` or ``a/b``) would be # unreachable or mis-routed. The client uses the returned ``id``. soul_id = slugify_soul_id(payload.id or payload.name) if mgr.get_soul(soul_id): raise HTTPException(status_code=409, detail=t("api.soul_already_exists", name=soul_id)) return mgr.create_soul(soul_id, payload.name, payload.content) @router.get("/souls/{soul_id}") async def get_soul(soul_id: str): soul = get_partner_manager().get_soul(soul_id) if not soul: raise HTTPException(status_code=404, detail=t("api.soul_not_found")) return soul @router.put("/souls/{soul_id}") async def update_soul(soul_id: str, payload: SoulTemplateUpdateRequest): result = get_partner_manager().update_soul(soul_id, payload.name, payload.content) if not result: raise HTTPException(status_code=404, detail=t("api.soul_not_found")) return result @router.delete("/souls/{soul_id}") async def delete_soul(soul_id: str): if not get_partner_manager().delete_soul(soul_id): raise HTTPException(status_code=404, detail=t("api.soul_not_found")) return {"id": soul_id, "deleted": True} @router.get("/soul-sources") async def soul_sources(): """Everything the create-wizard's soul step can start from.""" from deeptutor.multi_user.context import get_current_user from deeptutor.multi_user.paths import get_admin_path_service from deeptutor.services.persona import PersonaService, get_persona_service def _persona_entry(service: PersonaService, info: Any) -> dict[str, str]: # Content rides along so the wizard can preview the clone; creation # still re-resolves the persona server-side (_resolve_soul_content). try: content = strip_frontmatter(service.get_detail(info.name).content) except Exception: content = "" return {"name": info.name, "description": info.description, "content": content} personas: list[dict[str, str]] = [] seen: set[str] = set() try: service = get_persona_service() for info in service.list_personas(): personas.append(_persona_entry(service, info)) seen.add(info.name) except Exception: logger.warning("Failed to list user personas", exc_info=True) try: if not get_current_user().is_admin: admin_service = PersonaService( root=get_admin_path_service().get_workspace_dir() / "personas" ) for info in admin_service.list_personas(): if info.name not in seen: personas.append(_persona_entry(admin_service, info)) except Exception: logger.warning("Failed to list admin personas", exc_info=True) return {"library": get_partner_manager().list_souls(), "personas": personas} # ── Static catalog endpoints ─────────────────────────────────── @router.get("") async def list_partners(): return get_partner_manager().list_partners() @router.get("/recent") async def recent_partners(limit: int = 3): return get_partner_manager().get_recent_active_partners(limit=limit) @router.get("/channels/schema") async def list_channel_schemas(): """JSON-Schema metadata for every available channel (schema-driven UI).""" from deeptutor.api.routers._partners_channel_schema import all_channel_schemas return {"channels": all_channel_schemas()} @router.get("/tool-options") async def tool_options(): """The configurable tool surface for a partner. ``tools`` mirrors the user-toggleable system tools (the same pool the chat composer / settings expose); ``builtin_tools`` lists the auto-mounted built-in tools (rag / web_fetch / …) the owner can allow or deny; ``mcp_tools`` lists every configured MCP tool the partner could be allowed to load. ``read_memory`` / ``write_memory`` are excluded: partners use the mandatory ``partner_read`` / ``partner_memorize`` / ``partner_search`` tools instead, which are always on and not owner-configurable. """ from deeptutor.api.utils.tool_options import build_tool_options return await build_tool_options(exclude_builtin={"read_memory", "write_memory"}) # ── Create / read / update / lifecycle ───────────────────────── @router.post("") async def create_partner(payload: CreatePartnerRequest): mgr = get_partner_manager() partner_id = slugify_partner_id(payload.partner_id or payload.name) if mgr.partner_exists(partner_id): raise HTTPException( status_code=409, detail=t("api.partner_already_exists", name=partner_id), ) if payload.channels is not None: _validate_channels_payload(payload.channels) llm_selection = _validate_llm_selection_payload(payload.llm_selection) backup_llm_selection = _validate_llm_selection_payload(payload.backup_llm_selection) soul_content, soul_origin = _resolve_soul_content(payload.soul) config = PartnerConfig( name=payload.name.strip(), description=(payload.description or "").strip(), channels=payload.channels or {}, llm_selection=llm_selection, backup_llm_selection=backup_llm_selection, language=(payload.language or "").strip(), emoji=(payload.emoji or "").strip(), color=(payload.color or "").strip(), avatar=_validate_avatar_payload(payload.avatar), soul_origin=soul_origin, enabled_tools=payload.enabled_tools, builtin_tools=payload.builtin_tools, mcp_tools=payload.mcp_tools, ) mgr.save_config(partner_id, config, auto_start=bool(payload.start)) write_soul(partner_id, soul_content) provisioning: dict[str, Any] = {"copied": {}, "errors": []} if payload.assets is not None: provisioning = provision_assets( partner_id, knowledge_bases=payload.assets.knowledge_bases, skills=payload.assets.skills, notebooks=payload.assets.notebooks, ) if payload.start: try: instance = await mgr.start_partner(partner_id, config) result = instance.to_dict(mask_secrets=True) except Exception: logger.exception("Partner '%s' created but failed to start", partner_id) result = _stopped_partner_dict(partner_id, config) result["start_error"] = "Partner created but failed to start" else: result = _stopped_partner_dict(partner_id, config) result["provisioning"] = provisioning return result def _stopped_partner_dict( partner_id: str, cfg: PartnerConfig, *, include_secrets: bool = False, ) -> dict: if include_secrets: channels: object = strip_legacy_global_delivery(cfg.channels) else: channels = mask_channel_secrets(strip_legacy_global_delivery(cfg.channels)) return { "partner_id": partner_id, "name": cfg.name, "description": cfg.description, "channels": channels, "llm_selection": cfg.llm_selection, "backup_llm_selection": cfg.backup_llm_selection, "model": cfg.model, "language": cfg.language, "emoji": cfg.emoji, "color": cfg.color, "avatar": cfg.avatar, "soul_origin": cfg.soul_origin, "enabled_tools": cfg.enabled_tools, "builtin_tools": cfg.builtin_tools, "mcp_tools": cfg.mcp_tools, "running": False, "started_at": None, "last_reload_error": None, } @router.get("/{partner_id}") async def get_partner( partner_id: str, include_secrets: bool = Query( False, description=( "Return raw channel secrets (tokens, passwords). Required by the " "edit form; default response masks all secret-looking fields." ), ), ): mgr = get_partner_manager() instance = mgr.get_partner(partner_id) if instance: return instance.to_dict( include_secrets=include_secrets, mask_secrets=not include_secrets, ) cfg = mgr.load_config(partner_id) if cfg: return _stopped_partner_dict(partner_id, cfg, include_secrets=include_secrets) raise HTTPException(status_code=404, detail=t("api.partner_not_found")) def _apply_update(cfg: PartnerConfig, payload: UpdatePartnerRequest) -> None: if payload.name is not None: cfg.name = payload.name if payload.description is not None: cfg.description = payload.description if payload.channels is not None: cfg.channels = payload.channels if payload.language is not None: cfg.language = payload.language if payload.emoji is not None: cfg.emoji = payload.emoji if payload.color is not None: cfg.color = payload.color if payload.avatar is not None: cfg.avatar = _validate_avatar_payload(payload.avatar) if "llm_selection" in payload.model_fields_set: cfg.llm_selection = _validate_llm_selection_payload(payload.llm_selection) cfg.model = None # selection supersedes any legacy model string if "backup_llm_selection" in payload.model_fields_set: cfg.backup_llm_selection = _validate_llm_selection_payload(payload.backup_llm_selection) if "enabled_tools" in payload.model_fields_set: cfg.enabled_tools = payload.enabled_tools if "builtin_tools" in payload.model_fields_set: cfg.builtin_tools = payload.builtin_tools if "mcp_tools" in payload.model_fields_set: cfg.mcp_tools = payload.mcp_tools @router.patch("/{partner_id}") async def update_partner(partner_id: str, payload: UpdatePartnerRequest): if payload.channels is not None: _validate_channels_payload(payload.channels) mgr = get_partner_manager() instance = mgr.get_partner(partner_id) if instance: _apply_update(instance.config, payload) mgr.save_config(partner_id, instance.config) if payload.channels is not None: try: await mgr.reload_channels(partner_id) except Exception as exc: logger.exception("reload_channels failed for partner '%s'", partner_id) raise HTTPException( status_code=500, detail=( "Channels saved but failed to restart listeners " f"({type(exc).__name__}); try stopping and starting the partner." ), ) from None # LLM / tool changes need no reload: the runner resolves # llm_selection and tool config per turn from this same config object. return instance.to_dict(mask_secrets=True) cfg = mgr.load_config(partner_id) if not cfg: raise HTTPException(status_code=404, detail=t("api.partner_not_found")) _apply_update(cfg, payload) mgr.save_config(partner_id, cfg) return _stopped_partner_dict(partner_id, cfg) @router.post("/{partner_id}/start") async def start_partner(partner_id: str): instance = await _ensure_running_partner(partner_id, allow_stopped=True) # An explicit start is a persisted "run on boot" intent — so the partner # comes back in this state after a DeepTutor restart (a manual /stop clears # it; a lazy chat-driven start does NOT reach here, so it can't flip it). get_partner_manager().save_config(partner_id, instance.config, auto_start=True) return instance.to_dict(mask_secrets=True) @router.post("/{partner_id}/stop") async def stop_partner(partner_id: str): stopped = await get_partner_manager().stop_partner(partner_id) if not stopped: raise HTTPException(status_code=404, detail=t("api.partner_not_found_or_not_running")) return {"partner_id": partner_id, "stopped": True} @router.delete("/{partner_id}") async def destroy_partner(partner_id: str): destroyed = await get_partner_manager().destroy_partner(partner_id) if not destroyed: raise HTTPException(status_code=404, detail=t("api.partner_not_found")) return {"partner_id": partner_id, "destroyed": True} @router.post("/{partner_id}/channels/reload") async def reload_partner_channels(partner_id: str): mgr = get_partner_manager() instance = mgr.get_partner(partner_id) if not instance or not instance.running: raise HTTPException(status_code=404, detail=t("api.partner_not_running")) try: await mgr.reload_channels(partner_id) except Exception as exc: raise HTTPException( status_code=500, detail=f"Failed to reload channels: {type(exc).__name__}", ) from None return {"partner_id": partner_id, "reloaded": True} # ── Soul (the partner's own SOUL.md) ─────────────────────────── @router.get("/{partner_id}/soul") async def get_partner_soul(partner_id: str): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) return {"partner_id": partner_id, "content": read_soul(partner_id)} @router.put("/{partner_id}/soul") async def put_partner_soul(partner_id: str, payload: SoulUpdateBody): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) write_soul(partner_id, payload.content) return {"partner_id": partner_id, "saved": True} # ── Assets ───────────────────────────────────────────────────── @router.get("/{partner_id}/assets") async def get_partner_assets(partner_id: str): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) return list_assets(partner_id) @router.post("/{partner_id}/assets") async def add_partner_assets(partner_id: str, payload: AssetAddRequest): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) report = provision_assets( partner_id, knowledge_bases=payload.knowledge_bases, skills=payload.skills, notebooks=payload.notebooks, ) return {"partner_id": partner_id, **report, "assets": list_assets(partner_id)} @router.delete("/{partner_id}/assets/{asset_type}/{name}") async def delete_partner_asset(partner_id: str, asset_type: str, name: str): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) try: removed = remove_asset(partner_id, asset_type, name) except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) from None if not removed: raise HTTPException(status_code=404, detail="Asset not found") return {"partner_id": partner_id, "removed": True, "assets": list_assets(partner_id)} # ── History ──────────────────────────────────────────────────── @router.get("/{partner_id}/history") async def get_partner_history( partner_id: str, session_key: str | None = None, session_id: str | None = None, limit: int = 100, ): """Conversation history. Pass ``session_key`` for an exact key, or ``session_id`` for a web session (mapped through ``web_session_key``); with neither, all non-archived sessions are merged.""" mgr = get_partner_manager() if session_id and not session_key: session_key = mgr.web_session_key(partner_id, session_id=session_id) return mgr.get_history(partner_id, session_key=session_key, limit=limit) @router.get("/{partner_id}/sessions") async def get_partner_sessions(partner_id: str): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) return mgr.session_store(partner_id).list_sessions() @router.post("/{partner_id}/sessions/archive") async def archive_partner_session(partner_id: str, payload: SessionKeyBody): """Soft-archive a session (web /new) — it stays resumable, file untouched.""" mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) mgr.archive_session(partner_id, payload.session_key) return {"partner_id": partner_id, "archived": True, "session_key": payload.session_key} @router.post("/{partner_id}/sessions/resume") async def resume_partner_session(partner_id: str, payload: SessionKeyBody): """Clear a session's archived flag so the web app can continue it.""" mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) summary = mgr.resume_session(partner_id, payload.session_key) if summary is None: raise HTTPException(status_code=404, detail="Session not found") return {"partner_id": partner_id, "resumed": True, "session": summary} @router.post("/{partner_id}/sessions/delete") async def delete_partner_session(partner_id: str, payload: SessionKeyBody): mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) removed = mgr.delete_session(partner_id, payload.session_key) if not removed: raise HTTPException(status_code=404, detail="Session not found") return {"partner_id": partner_id, "deleted": True, "session_key": payload.session_key} @router.post("/{partner_id}/sessions/branch") async def branch_partner_session(partner_id: str, payload: SessionBranchBody): """Copy a session's full history into a new key and archive the source.""" mgr = get_partner_manager() if not mgr.partner_exists(partner_id): raise HTTPException(status_code=404, detail=t("api.partner_not_found")) summary = mgr.branch_session(partner_id, payload.source_key, payload.new_key) if summary is None: raise HTTPException(status_code=400, detail="Nothing to branch (source is empty)") return {"partner_id": partner_id, "branched": True, "session": summary} @router.get("/commands/palette") async def partner_command_palette(): from deeptutor.services.partners.commands import partner_command_palette return {"commands": partner_command_palette()} # ── Chat (HTTP / SSE / WebSocket) ────────────────────────────── def _sse(event: str, payload: dict[str, Any]) -> str: return f"event: {event}\ndata: {json.dumps(payload, ensure_ascii=False, default=str)}\n\n" def _resolve_http_session(payload: ChatMessageRequest) -> tuple[str, str]: explicit_session = (payload.session_id or "").strip() explicit_chat = (payload.chat_id or "").strip() if explicit_session: return explicit_session, explicit_chat or explicit_session if explicit_chat: return explicit_chat, explicit_chat session_id = uuid4().hex return session_id, session_id _PARTNER_UPLOAD_MAX_BYTES = 10 * 1024 * 1024 _PARTNER_UPLOAD_MAX_TOTAL_BYTES = 25 * 1024 * 1024 def _clean_attachment_base64(value: str) -> str: text = str(value or "").strip() if text.startswith("data:") and "," in text: return text.split(",", 1)[1] return text def _default_attachment_prompt(attachments: list[ChatAttachmentRequest]) -> str: if attachments and all(str(item.type).lower() == "image" for item in attachments): return t("Please analyze the attached image(s).") return t("Please use the attached file(s).") def _materialize_partner_attachments( partner_id: str, attachments: list[ChatAttachmentRequest], ) -> list[str]: """Persist browser-sent attachment bytes into the partner media tree.""" if not attachments: return [] media_dir = get_partner_media_dir(partner_id, "web") total_bytes = 0 media_paths: list[str] = [] for item in attachments: raw_b64 = _clean_attachment_base64(item.base64) if not raw_b64: # Partner web chat accepts uploaded bytes only. URL-only # attachments are ignored rather than fetched server-side. continue try: data = base64.b64decode(raw_b64, validate=False) except (binascii.Error, ValueError) as exc: raise HTTPException( status_code=422, detail=f"Invalid attachment data for {item.filename or 'file'}", ) from exc if len(data) > _PARTNER_UPLOAD_MAX_BYTES: raise HTTPException( status_code=413, detail=f"Attachment too large: {item.filename or 'file'}", ) if total_bytes + len(data) > _PARTNER_UPLOAD_MAX_TOTAL_BYTES: raise HTTPException(status_code=413, detail="Attachment batch too large") total_bytes += len(data) filename = safe_filename(item.filename or "attachment") or "attachment" path = media_dir / f"{uuid4().hex[:12]}_{filename}" path.write_bytes(data) media_paths.append(str(path)) return media_paths @router.post("/{partner_id}/chat") async def partner_chat_http(partner_id: str, payload: ChatMessageRequest) -> dict[str, Any]: """Send one HTTP message to a partner with persistent session context.""" content = payload.content.strip() if not content and not payload.attachments: raise HTTPException(status_code=400, detail=t("api.content_required")) await _ensure_running_partner(partner_id) media_paths = _materialize_partner_attachments(partner_id, payload.attachments) if not content and media_paths: content = _default_attachment_prompt(payload.attachments) mgr = get_partner_manager() session_id, chat_id = _resolve_http_session(payload) try: response = await mgr.send_message( partner_id, content, chat_id=chat_id, session_id=session_id, media=media_paths, session_key=payload.session_key, ) except RuntimeError as exc: raise HTTPException(status_code=400, detail=str(exc)) from None return { "partner_id": partner_id, "session_id": session_id, "content": response, } async def _partner_chat_stream( partner_id: str, payload: ChatMessageRequest, ) -> AsyncGenerator[str, None]: from deeptutor.core.stream import StreamEventType mgr = get_partner_manager() content = payload.content.strip() if not content and not payload.attachments: yield _sse("error", {"detail": t("api.content_required")}) return media_paths = _materialize_partner_attachments(partner_id, payload.attachments) if not content and media_paths: content = _default_attachment_prompt(payload.attachments) session_id, chat_id = _resolve_http_session(payload) queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() done = asyncio.Event() holder: dict[str, Any] = {} async def on_event(event: Any) -> None: if event.type == StreamEventType.THINKING and event.content: await queue.put({"event": "thinking", "payload": {"content": event.content}}) async def run() -> None: try: holder["content"] = await mgr.send_message( partner_id, content, chat_id=chat_id, session_id=session_id, media=media_paths, on_event=on_event, session_key=payload.session_key, ) except Exception as exc: # noqa: BLE001 holder["error"] = str(exc) finally: done.set() yield _sse("session", {"partner_id": partner_id, "session_id": session_id}) task = asyncio.create_task(run()) try: while not done.is_set(): try: item = await asyncio.wait_for(queue.get(), timeout=0.15) except asyncio.TimeoutError: continue yield _sse(item["event"], item["payload"]) while not queue.empty(): item = queue.get_nowait() yield _sse(item["event"], item["payload"]) if holder.get("error"): yield _sse("error", {"detail": holder["error"]}) return yield _sse("content", {"content": holder.get("content", "")}) yield _sse("done", {"partner_id": partner_id, "session_id": session_id}) finally: if not task.done(): task.cancel() @router.post("/{partner_id}/chat/execute-stream") async def partner_chat_http_stream(partner_id: str, payload: ChatMessageRequest): """Stream one HTTP message to a partner as server-sent events.""" if not payload.content.strip() and not payload.attachments: raise HTTPException(status_code=400, detail=t("api.content_required")) await _ensure_running_partner(partner_id) return StreamingResponse( _partner_chat_stream(partner_id, payload), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, ) @router.websocket("/{partner_id}/ws") async def partner_chat_ws(ws: WebSocket, partner_id: str): """Web chat socket. Client → server: ``{"content": str, "session_id"?: str, "chat_id"?: str, "attachments"?: [{"type", "filename", "mime_type", "base64"}]}``. Server → client frames: * ``{"type": "stream_event", "event": {...}}`` — every chat-loop StreamEvent (content/thinking/tool_call/progress/sources/result), letting the UI render the same live trace as product chat; * ``{"type": "content", "content": str}`` — the final reply; * ``{"type": "done"}`` / ``{"type": "error"}`` / ``{"type": "proactive"}``. """ from deeptutor.api.routers.auth import ws_auth_failed, ws_require_auth from deeptutor.multi_user.context import reset_current_user user_token = await ws_require_auth(ws) if user_token is ws_auth_failed: return mgr = get_partner_manager() disconnected = asyncio.Event() async def _safe_send(payload: dict) -> bool: try: await ws.send_json(payload) return True except (WebSocketDisconnect, RuntimeError): disconnected.set() return False await ws.accept() try: instance = await _ensure_running_partner(partner_id) except HTTPException as exc: message = str(exc.detail) await _safe_send({"type": "error", "content": message}) code = 4004 if exc.status_code == 404 else 4003 await ws.close(code=code, reason=message[:120]) return logger.info("WebSocket connected for partner '%s'", partner_id) # Web turns run on the partner instance (see LiveTurn), NOT tied to this # socket — so a refresh reattaches and replays instead of killing the turn. # The socket just drains a subscriber queue; the receive loop stays free to # process stop / attach frames concurrently. drain: dict[str, asyncio.Task | None] = {"task": None} async def _drain(queue: asyncio.Queue) -> None: while True: frame = await queue.get() if not await _safe_send(frame): return if frame.get("type") in {"done", "stopped"}: return def _start_drain(queue: asyncio.Queue) -> None: prev = drain["task"] if prev is not None and not prev.done(): prev.cancel() drain["task"] = asyncio.create_task(_drain(queue)) def _resolve_key(data: dict[str, Any]) -> str: return str( data.get("session_key") or mgr.web_session_key( partner_id, chat_id=data.get("chat_id", "web"), session_id=data.get("session_id"), ) ) async def _handle_user_messages(): while not disconnected.is_set(): try: raw = await ws.receive_text() except WebSocketDisconnect: disconnected.set() break try: data = json.loads(raw) except json.JSONDecodeError: if not await _safe_send({"type": "error", "content": "Invalid JSON"}): break continue action = data.get("action") if action == "stop": mgr.stop_web_turn(partner_id, _resolve_key(data)) continue if action == "attach": # Reconnect (a page refresh) — replay an in-flight turn so the # streaming answer the user was watching survives the reload. turn = mgr.subscribe_web_turn(partner_id, _resolve_key(data)) if turn is not None: await _safe_send({"type": "resuming"}) if turn.user_content: await _safe_send({"type": "user_echo", "content": turn.user_content}) _start_drain(turn.subscribe()) continue content = data.get("content", "").strip() try: attachments = [ ChatAttachmentRequest.model_validate(item) for item in (data.get("attachments") or []) if isinstance(item, dict) ] except ValidationError: if not await _safe_send({"type": "error", "content": "Invalid attachments"}): break continue if not content and not attachments: continue try: media_paths = _materialize_partner_attachments(partner_id, attachments) except HTTPException as exc: if not await _safe_send({"type": "error", "content": str(exc.detail)}): break continue if not content and media_paths: content = _default_attachment_prompt(attachments) try: turn = mgr.start_web_turn(partner_id, _resolve_key(data), content, media_paths) except RuntimeError as exc: if not await _safe_send({"type": "error", "content": str(exc)}): break continue _start_drain(turn.subscribe()) async def _handle_notifications(): while not disconnected.is_set(): get_task = asyncio.create_task(instance.notify_queue.get()) wait_task = asyncio.create_task(disconnected.wait()) done, pending = await asyncio.wait( {get_task, wait_task}, return_when=asyncio.FIRST_COMPLETED, ) for t in pending: t.cancel() if get_task not in done: break content = get_task.result() if not await _safe_send({"type": "proactive", "content": content}): break user_task = asyncio.create_task(_handle_user_messages()) notify_task = asyncio.create_task(_handle_notifications()) try: done, pending = await asyncio.wait( [user_task, notify_task], return_when=asyncio.FIRST_COMPLETED, ) disconnected.set() for t in pending: t.cancel() for t in done: if t.exception() and not isinstance(t.exception(), WebSocketDisconnect): logger.exception( "WebSocket task error for partner '%s'", partner_id, exc_info=t.exception(), ) except Exception: disconnected.set() user_task.cancel() notify_task.cancel() finally: # Detach from the stream only — the turn keeps running on the instance # so a reconnecting client can reattach and replay it. d = drain["task"] if d is not None and not d.done(): d.cancel() if user_token is not None: try: reset_current_user(user_token) except Exception: pass logger.info("WebSocket closed for partner '%s'", partner_id)