deeptutor-server 1.7.0__py3-none-any.whl

deeptutor/api/routers/workspace.py

@@ -0,0 +1,371 @@
+"""Content-workspace settings and authenticated presentation delivery."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi.responses import FileResponse
+from pydantic import BaseModel, Field
+
+from deeptutor.api.routers.auth import require_auth
+from deeptutor.multi_user.partner_access import visible_partners
+from deeptutor.multi_user.paths import user_context
+from deeptutor.services.auth import TokenPayload
+from deeptutor.services.workspace import (
+    WorkspaceError,
+    WorkspaceItem,
+    get_content_workspace_service,
+)
+from deeptutor.services.workspace.models import MAX_WORKSPACE_INSTRUCTIONS_CHARS
+from deeptutor.services.workspace.resources import WorkspaceResources
+
+settings_router = APIRouter()
+files_router = APIRouter()
+
+
+class WorkspacePathPayload(BaseModel):
+    path: str | None = None
+
+
+class CreateWorkspacePayload(BaseModel):
+    instructions: str | None = Field(default=None, max_length=MAX_WORKSPACE_INSTRUCTIONS_CHARS)
+    resources: WorkspaceResources | None = None
+    name: str = Field(min_length=1, max_length=100)
+    path: str | None = None
+
+
+class UpdateWorkspacePayload(BaseModel):
+    instructions: str | None = Field(default=None, max_length=MAX_WORKSPACE_INSTRUCTIONS_CHARS)
+    resources: WorkspaceResources | None = None
+    name: str | None = Field(default=None, min_length=1, max_length=100)
+    archived: bool | None = None
+
+
+class DataMigrationPayload(BaseModel):
+    source_workspace_id: str = ""
+    target_workspace_id: str = ""
+    features: list[str] = Field(min_length=1)
+    include_historical: bool = False
+
+
+class KnowledgeMovePayload(BaseModel):
+    source_id: str = Field(min_length=1)
+    target_workspace_id: str = ""
+
+
+async def _data_operation(function, *args, **kwargs):
+    _assert_migration_access()
+    task = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs))
+    task.add_done_callback(lambda done: done.exception() if not done.cancelled() else None)
+    try:
+        return await asyncio.shield(task)
+    except (WorkspaceError, OSError, ValueError) as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
+
+
+@settings_router.post("/knowledge-bases/preview-move")
+async def preview_knowledge_move(payload: KnowledgeMovePayload) -> dict:
+    from deeptutor.services.workspace.kb_move import preview_kb_move
+
+    return await _data_operation(preview_kb_move, payload.source_id, payload.target_workspace_id)
+
+
+@settings_router.post("/knowledge-bases/move")
+async def move_knowledge_base(payload: KnowledgeMovePayload) -> dict:
+    from deeptutor.services.workspace.kb_move import move_kb
+
+    return await _data_operation(move_kb, payload.source_id, payload.target_workspace_id)
+
+
+def _assert_migration_access():
+    from deeptutor.multi_user.learning_access import current_learning_policy
+
+    if current_learning_policy() is not None:
+        raise HTTPException(
+            status_code=403,
+            detail="Data migration and export are managed by the guardian for this learning account.",
+        )
+
+
+@settings_router.get("/data/discover")
+async def discover_workspace_data(source_workspace_id: str = "") -> dict:
+    from deeptutor.services.workspace.data_migration import discover
+
+    return await _data_operation(discover, source_workspace_id)
+
+
+@settings_router.post("/data/preview")
+async def preview_workspace_data(payload: DataMigrationPayload) -> dict:
+    from deeptutor.services.workspace.data_migration import preview
+
+    return await _data_operation(
+        preview, payload.source_workspace_id, payload.target_workspace_id, payload.features
+    )
+
+
+@settings_router.post("/data/migrate")
+async def migrate_workspace_data(payload: DataMigrationPayload) -> dict:
+    from deeptutor.services.workspace.data_migration import migrate_data
+
+    return await _data_operation(
+        migrate_data, payload.source_workspace_id, payload.target_workspace_id, payload.features
+    )
+
+
+@settings_router.post("/data/export")
+async def export_workspace_data(payload: DataMigrationPayload) -> dict:
+    from deeptutor.services.workspace.data_migration import export_data
+
+    return await _data_operation(
+        export_data,
+        payload.source_workspace_id,
+        payload.features,
+        include_historical=payload.include_historical,
+    )
+
+
+@settings_router.get("/data/operations")
+async def workspace_data_operations() -> dict:
+    _assert_migration_access()
+    from deeptutor.services.workspace.data_migration import operations
+
+    return {"operations": await asyncio.to_thread(operations)}
+
+
+@settings_router.post("/data/operations/{operation_id}/recover")
+async def recover_workspace_migration(operation_id: str) -> dict:
+    from deeptutor.services.workspace.data_migration import recover_operation
+
+    return await _data_operation(recover_operation, operation_id)
+
+
+@settings_router.get("/data/exports/{operation_id}")
+async def download_workspace_export(operation_id: str) -> FileResponse:
+    _assert_migration_access()
+    from deeptutor.services.workspace.data_migration import export_path
+
+    try:
+        return FileResponse(
+            export_path(operation_id),
+            media_type="application/zip",
+            filename=f"deeptutor-data-{operation_id[:8]}.zip",
+        )
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+
+
+@settings_router.get("/registrations")
+async def list_workspaces() -> dict:
+    return get_content_workspace_service().describe_catalog()
+
+
+class WorkspaceMigrationPayload(BaseModel):
+    path: str = Field(min_length=1)
+
+
+async def _migrate_workspaces(path: str, workspace_id: str | None = None) -> dict:
+    from deeptutor.services.session import get_session_store
+
+    service = get_content_workspace_service()
+
+    async def perform() -> dict:
+        from deeptutor.services.workspace.activity import data_activity
+        from deeptutor.services.workspace.data_migration import assert_no_pending_recovery
+
+        with data_activity(exclusive=True):
+            assert_no_pending_recovery(reject_unreadable=True)
+            return await migrate_location()
+
+    async def migrate_location() -> dict:
+        with service.maintenance():
+            active = await get_session_store().list_nonterminal_turns()
+            if active:
+                raise WorkspaceError(
+                    "Wait for running conversations to finish before migrating workspace storage."
+                )
+            if workspace_id is None:
+                return await asyncio.to_thread(service.migrate_root, path)
+            return await asyncio.to_thread(service.migrate_workspace, workspace_id, path)
+
+    # A disconnected browser must not abandon a copy midway or release the
+    # migration guard while its worker is still copying files.
+    task = asyncio.create_task(perform())
+    # Retrieve a late error even if the requesting client has disconnected.
+    task.add_done_callback(lambda done: done.exception() if not done.cancelled() else None)
+    try:
+        return await asyncio.shield(task)
+    except (WorkspaceError, OSError, ValueError) as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
+
+
+@settings_router.post("/registrations/migrate-root")
+async def migrate_workspace_root(payload: WorkspaceMigrationPayload) -> dict:
+    return await _migrate_workspaces(payload.path)
+
+
+@settings_router.post("/registrations/{workspace_id}/migrate")
+async def migrate_registered_workspace(
+    workspace_id: str, payload: WorkspaceMigrationPayload
+) -> dict:
+    return await _migrate_workspaces(payload.path, workspace_id)
+
+
+@settings_router.get("/registrations/system-snapshot")
+@settings_router.post("/registrations/system-snapshot")
+async def system_workspace_snapshot() -> dict:
+    try:
+        return get_content_workspace_service().refresh_system_snapshot()
+    except (WorkspaceError, OSError, ValueError) as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
+
+
+@settings_router.get("/resource-usage")
+async def workspace_resource_usage(kind: str, resource_id: str, skill_workspace: str = "") -> dict:
+    from deeptutor.services.workspace.resources import resource_usage
+
+    try:
+        return {"workspaces": resource_usage(kind, resource_id, skill_workspace=skill_workspace)}
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
+@settings_router.get("/resources")
+async def workspace_resource_catalog(workspace_id: str = "") -> dict:
+    from deeptutor.services.workspace.resources import resource_catalog
+
+    try:
+        return await asyncio.to_thread(resource_catalog, workspace_id=workspace_id)
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
+@settings_router.post("/registrations")
+async def create_workspace(payload: CreateWorkspacePayload) -> dict:
+    try:
+        workspace = get_content_workspace_service().create_workspace(
+            payload.name,
+            payload.path,
+            resources=payload.resources.model_dump() if payload.resources else None,
+            instructions=payload.instructions,
+        )
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    return {"workspace": workspace}
+
+
+@settings_router.patch("/registrations/{workspace_id}")
+async def update_registered_workspace(workspace_id: str, payload: UpdateWorkspacePayload) -> dict:
+    try:
+        from deeptutor.services.workspace.activity import data_activity
+        from deeptutor.services.workspace.data_migration import assert_no_pending_recovery
+
+        # Instruction edits are atomic metadata updates. A running turn owns
+        # its admitted snapshot, so these edits can coexist with its data lease.
+        changes = payload.model_dump(exclude_unset=True)
+        with data_activity(exclusive=set(changes) != {"instructions"}):
+            assert_no_pending_recovery(reject_unreadable=True)
+            workspace = get_content_workspace_service().update_workspace(workspace_id, **changes)
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+    return {"workspace": workspace}
+
+
+@settings_router.get("")
+async def get_workspace_settings() -> dict:
+    return get_content_workspace_service().describe_current()
+
+
+@settings_router.post("/validate")
+async def validate_workspace(payload: WorkspacePathPayload) -> dict:
+    return get_content_workspace_service().validate(payload.path)
+
+
+@settings_router.put("")
+async def update_workspace(payload: WorkspacePathPayload) -> dict:
+    service = get_content_workspace_service()
+    try:
+        service.set_workspace(payload.path)
+    except WorkspaceError as exc:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
+    return service.describe_current()
+
+
+def _resolve_partner_item(
+    workspace_id: str, workspace_item_id: str
+) -> tuple[Path, WorkspaceItem] | None:
+    """Resolve a presentation one of the caller's partners published.
+
+    A partner web chat runs inside a synthetic workspace under
+    ``data/partners/<id>/workspace``, so ``workspace_present`` writes its
+    manifest to that scope's private presentation store. The human caller's own
+    bindings never name that workspace, ``binding_by_id`` refuses it, and the
+    attachment sitting in the transcript 404s on click (#1267).
+
+    ``/files/outputs`` grew the same reach in #1012 — see
+    ``_resolve_partner_output`` — and this is the other artifact URL shape. It
+    does not need that endpoint's unique-match rule: a workspace id is derived
+    from the owning user id together with the physical root and re-derived
+    before a stored row is trusted, so a hit inside a partner scope belongs to
+    that partner by construction, where a relative output path can collide.
+    """
+    from deeptutor.services.workspace.context import current_workspace_id
+
+    if current_workspace_id():
+        return None
+    for partner in visible_partners():
+        partner_id = str(partner.get("partner_id") or "").strip()
+        if not partner_id:
+            continue
+        from deeptutor.services.partners.scope import partner_user
+
+        with user_context(partner_user(partner_id)):
+            try:
+                return get_content_workspace_service().resolve_published_item(
+                    workspace_id, workspace_item_id
+                )
+            except WorkspaceError:
+                continue
+    return None
+
+
+@files_router.get("/{workspace_id}/{workspace_item_id}", operation_id="read_workspace_item_get")
+@files_router.head("/{workspace_id}/{workspace_item_id}", operation_id="read_workspace_item_head")
+async def read_workspace_item(
+    workspace_id: str,
+    workspace_item_id: str,
+    _auth: TokenPayload | None = Depends(require_auth),
+) -> FileResponse:
+    try:
+        path, item = get_content_workspace_service().resolve_published_item(
+            workspace_id, workspace_item_id
+        )
+    except WorkspaceError as exc:
+        resolved = _resolve_partner_item(workspace_id, workspace_item_id)
+        if resolved is None:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND, detail="File not found"
+            ) from exc
+        path, item = resolved
+    return FileResponse(
+        path,
+        media_type=item.mime_type,
+        filename=item.filename,
+        content_disposition_type="inline",
+        headers={
+            "ETag": f'"{item.sha256}"',
+            "Cache-Control": "private, max-age=31536000, immutable",
+            "X-Content-Type-Options": "nosniff",
+            # Presented files are user/model-controlled. Opening an HTML or
+            # SVG snapshot directly on the app origin must never execute it
+            # with access to DeepTutor's authenticated origin.
+            "Content-Security-Policy": (
+                "sandbox; default-src 'none'; img-src data: blob:; "
+                "media-src 'self' blob:; style-src 'unsafe-inline'"
+            ),
+        },
+    )
+
+
+__all__ = ["files_router", "settings_router"]

deeptutor/api/run_server.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+"""
+Uvicorn Server Startup Script
+Uses Python API instead of command line to avoid Windows path parsing issues.
+"""
+
+import asyncio
+import os
+from pathlib import Path
+import sys
+
+from deeptutor.runtime.home import get_runtime_home
+
+# Windows: uvicorn defaults to SelectorEventLoop which does not support
+# asyncio.create_subprocess_exec.  Switch to ProactorEventLoop so that
+# child-process APIs (used by Math Animator renderer, etc.) work correctly.
+if sys.platform == "win32":
+    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
+
+import uvicorn
+
+# Force unbuffered output
+os.environ["PYTHONUNBUFFERED"] = "1"
+if hasattr(sys.stdout, "reconfigure"):
+    sys.stdout.reconfigure(line_buffering=True, encoding="utf-8", errors="replace")
+if hasattr(sys.stderr, "reconfigure"):
+    sys.stderr.reconfigure(line_buffering=True, encoding="utf-8", errors="replace")
+
+
+def main() -> None:
+    # Runtime workspace root owns data/user/settings and generated outputs.
+    project_root = get_runtime_home()
+    os.chdir(str(project_root))
+
+    # Get port from configuration
+    from deeptutor.logging import configure_logging
+    from deeptutor.runtime.mode import RunMode, set_mode
+    from deeptutor.services.config import (
+        HTTP_KEEP_ALIVE_TIMEOUT,
+        get_ws_max_size,
+        load_system_settings,
+    )
+    from deeptutor.services.setup import get_backend_port
+
+    set_mode(RunMode.SERVER)
+    configure_logging()
+    backend_port = get_backend_port(project_root)
+
+    # Configure reload_excludes to skip directories that shouldn't trigger reloads
+    # Use absolute paths to ensure they're properly resolved
+    reload_excludes = [
+        str(project_root / "venv"),  # Virtual environment
+        str(project_root / ".venv"),  # Virtual environment (alternative name)
+        str(project_root / "data"),  # Data directory (includes knowledge_bases, user data, logs)
+        str(project_root / "node_modules"),  # Node modules (if any at root)
+        str(project_root / "web" / "node_modules"),  # Web node modules
+        str(project_root / "web" / ".next"),  # Next.js build
+        str(project_root / ".git"),  # Git directory
+        str(project_root / "scripts"),  # Scripts directory - don't reload on launcher changes
+    ]
+
+    # Filter out non-existent directories to avoid warnings
+    reload_excludes = [d for d in reload_excludes if Path(d).exists()]
+
+    # Reload launches a supervisor plus a worker and retains file-watcher state,
+    # so keep it opt-in for local development. ws_max_size tracks the configured
+    # chat-attachment total so base64 uploads fit in one WS frame.
+    dev_reload = os.environ.get("DEEPTUTOR_DEV_RELOAD", "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+    backend_workers = max(1, int(load_system_settings().get("backend_workers") or 1))
+    if dev_reload and backend_workers > 1:
+        raise SystemExit(
+            "Development reload and backend_workers > 1 are mutually exclusive. "
+            "Set backend_workers=1 or disable DEEPTUTOR_DEV_RELOAD."
+        )
+    uvicorn.run(
+        "deeptutor.api.main:app",
+        host="0.0.0.0",
+        port=backend_port,
+        reload=dev_reload,
+        workers=backend_workers,
+        reload_excludes=reload_excludes if dev_reload else None,
+        log_level="info",
+        access_log=False,
+        # XFF is client-controlled in this deployment. Keep request.client tied
+        # to the actual backend peer for the selective access log.
+        proxy_headers=False,
+        ws_max_size=get_ws_max_size(),
+        timeout_keep_alive=HTTP_KEEP_ALIVE_TIMEOUT,
+        timeout_graceful_shutdown=1,
+    )
+
+
+if __name__ == "__main__":
+    main()

deeptutor/api/security_guard.py

@@ -0,0 +1,91 @@
+"""Opt-in request screening with fastapi-guard.
+
+Disabled by default. When ``system.json`` -> ``security_guard.enabled`` is false
+(or ``DEEPTUTOR_GUARD_ENABLED`` is unset), nothing here imports ``guard`` and the
+middleware stack is unchanged. See ``docs-for-user/SECURITY_GUARD.md``.
+
+The guard runs inside the application, in front of the routers. It screens HTTP
+requests by client address (rate limit, IP allow/deny lists, ban list) and scans
+the URL path, query string and headers for attack patterns. It does not read
+request bodies (prompts, documents and chat are free text), does not touch
+WebSocket handshakes, and does not use Redis, CORS or security headers.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fastapi import FastAPI
+
+from deeptutor.services.config import load_system_settings
+
+logger = logging.getLogger(__name__)
+
+INSTALL_HINT = (
+    "security_guard is enabled but the 'fastapi-guard' package is not installed. "
+    "Install it with `pip install 'deeptutor[guard]'` (Docker: set "
+    "DEEPTUTOR_EXTRAS=guard), or disable it with security_guard.enabled=false in "
+    "system.json / DEEPTUTOR_GUARD_ENABLED=false."
+)
+
+
+def load_guard_settings() -> dict[str, Any]:
+    """Effective ``security_guard`` settings (file plus deployment overrides)."""
+    return dict(load_system_settings()["security_guard"])
+
+
+def build_security_config(settings: dict[str, Any]) -> Any:
+    """Translate DeepTutor settings into a ``guard.SecurityConfig``."""
+    try:
+        from guard import SecurityConfig
+    except ImportError as exc:
+        raise RuntimeError(INSTALL_HINT) from exc
+    from guard_core.handlers.suspatterns_handler import ALL_DETECTION_CATEGORIES
+
+    return SecurityConfig(
+        rate_limit=settings["rate_limit"],
+        rate_limit_window=settings["rate_limit_window"],
+        trusted_proxies=tuple(settings["trusted_proxies"]),
+        exempt_ips=tuple(settings["exempt_ips"]),
+        whitelist=tuple(settings["whitelist"]) or None,
+        blacklist=tuple(settings["blacklist"]),
+        # Process-local state: each uvicorn worker keeps its own counters.
+        enable_redis=False,
+        # DeepTutor sets CORS itself (CORSMiddleware) and owns its response headers.
+        enable_cors=False,
+        security_headers={"enabled": False},
+        # Prompts, documents and chat messages are free text, not attack input.
+        detection_scan_body=False,
+        # Browsers send the page URL as Referer; guard-core 4.4.0 flags localhost and
+        # LAN addresses in it as ssrf (400), so that category stays off.
+        enabled_detection_categories=ALL_DETECTION_CATEGORIES - {"ssrf"},
+        enable_agent=False,
+    )
+
+
+def install_security_guard(app: FastAPI, settings: dict[str, Any] | None = None) -> bool:
+    """Add the guard middleware to ``app`` when enabled; return whether it was added.
+
+    Call this before ``CORSMiddleware`` is added: Starlette puts the last-added
+    middleware outermost, so CORS stays outside the guard and 429/403 responses
+    still carry CORS headers.
+    """
+    settings = load_guard_settings() if settings is None else settings
+    if not settings["enabled"]:
+        return False
+
+    config = build_security_config(settings)
+    from guard import PureASGISecurityMiddleware
+
+    app.add_middleware(PureASGISecurityMiddleware, config=config)
+    logger.info(
+        "Security guard enabled: rate_limit=%s/%ss trusted_proxies=%s exempt_ips=%s whitelist=%s blacklist=%s",
+        settings["rate_limit"],
+        settings["rate_limit_window"],
+        settings["trusted_proxies"],
+        settings["exempt_ips"],
+        len(settings["whitelist"]),
+        len(settings["blacklist"]),
+    )
+    return True

deeptutor/api/utils/http_headers.py

@@ -0,0 +1,27 @@
+"""Shared HTTP header builders.
+
+Small, boring, and easy to get wrong in ways that only show up for non-English
+users — so they live in one place rather than once per router.
+"""
+
+from __future__ import annotations
+
+from urllib.parse import quote
+
+
+def content_disposition(filename: str, *, disposition: str = "inline") -> str:
+    """Build a Content-Disposition header that survives non-ASCII filenames.
+
+    HTTP/1.1 headers are latin-1, so dropping a Chinese / accented filename
+    straight into ``filename="..."`` blows up with UnicodeEncodeError. RFC
+    6266 / RFC 5987 cover this: emit ``filename*=UTF-8''<percent-encoded>``
+    plus an ASCII fallback on ``filename=`` for legacy clients.
+    """
+    ascii_fallback = filename.encode("ascii", errors="replace").decode("ascii")
+    # Quotes / backslashes break the simple-quoted-string form; collapse them.
+    ascii_fallback = ascii_fallback.replace('"', "_").replace("\\", "_")
+    encoded = quote(filename, safe="")
+    return f"{disposition}; filename=\"{ascii_fallback}\"; filename*=UTF-8''{encoded}"
+
+
+__all__ = ["content_disposition"]

deeptutor/api/utils/progress_broadcaster.py

@@ -0,0 +1,73 @@
+"""
+Progress Broadcaster - Manages WebSocket broadcasting of knowledge base progress
+"""
+
+import asyncio
+import logging
+from typing import Optional
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressBroadcaster:
+    """Manages WebSocket broadcasting of knowledge base progress"""
+
+    _instance: Optional["ProgressBroadcaster"] = None
+    _connections: dict[str, set[WebSocket]] = {}  # kb_name -> Set[WebSocket]
+    _lock = asyncio.Lock()
+
+    @classmethod
+    def get_instance(cls) -> "ProgressBroadcaster":
+        """Get singleton instance"""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    async def connect(self, kb_name: str, websocket: WebSocket):
+        """Connect WebSocket to specified knowledge base"""
+        async with self._lock:
+            if kb_name not in self._connections:
+                self._connections[kb_name] = set()
+            self._connections[kb_name].add(websocket)
+            logger.debug(
+                f"Connected WebSocket for KB '{kb_name}' (total: {len(self._connections[kb_name])})"
+            )
+
+    async def disconnect(self, kb_name: str, websocket: WebSocket):
+        """Disconnect WebSocket connection"""
+        async with self._lock:
+            if kb_name in self._connections:
+                self._connections[kb_name].discard(websocket)
+                if not self._connections[kb_name]:
+                    del self._connections[kb_name]
+                logger.debug(f"Disconnected WebSocket for KB '{kb_name}'")
+
+    async def broadcast(self, kb_name: str, progress: dict):
+        """Broadcast progress update to all WebSocket connections for specified knowledge base"""
+        async with self._lock:
+            if kb_name not in self._connections:
+                return
+
+            # Create list of connections to remove (closed connections)
+            to_remove = []
+
+            for websocket in self._connections[kb_name]:
+                try:
+                    await websocket.send_json({"type": "progress", "data": progress})
+                except Exception as e:
+                    # Connection closed or error, mark for removal
+                    logger.debug(f"Error sending to WebSocket for KB '{kb_name}': {e}")
+                    to_remove.append(websocket)
+
+            # Remove closed connections
+            for ws in to_remove:
+                self._connections[kb_name].discard(ws)
+
+            if not self._connections[kb_name]:
+                del self._connections[kb_name]
+
+    def get_connection_count(self, kb_name: str) -> int:
+        """Get connection count for specified knowledge base"""
+        return len(self._connections.get(kb_name, set()))