feat: add proxy template origin field and implement socket utilities for log streaming

backend/app/api/proxy/router.py

@@ -5,15 +5,20 @@ from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
 from pymongo.errors import DuplicateKeyError
 from starlette.websockets import WebSocketDisconnect
 
-from app.core.config import settings
-from app.core.security import (
-    InvalidTokenError,
-    authenticate_token,
-    require_authenticated_user,
+from app.api.proxy.socket_utils import (
+    authenticate_proxy_socket,
+    proxy_log_path,
+    read_log_snapshot,
+    serialize_health_snapshot,
+    snapshot_log_state,
+    stream_proxy_logs,
+    watch_socket_disconnect,
 )
+from app.core.config import settings
+from app.core.security import InvalidTokenError, require_authenticated_user
 from app.integrations.docker import get_all_services_in_dir
 from app.integrations.docker.discovery import service_scan_root
-from app.models.health import HealthHistoryResponse, HealthSnapshot
+from app.models.health import HealthHistoryResponse
 from app.models.proxy import ProxyCreateInput, ProxyUpdateInput
 from app.models.template import ProxyTemplateCreateInput
 from app.repositories.health import (
@@ -51,7 +56,7 @@ from app.services.mitmproxy.runtime_validation import (
     backend_port,
     ensure_proxy_port_not_reserved,
 )
-from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
+from app.services.mitmproxy.templates import file_template_ids, list_file_templates
 from app.services.mitmproxy.validator import (
     ProxyCodeValidationError,
     validate_proxy_code,
@@ -152,125 +157,6 @@ async def get_proxy_or_404(proxy_port: int):
     return proxy
 
 
-async def authenticate_proxy_socket(websocket: WebSocket) -> None:
-    token = websocket.query_params.get("token")
-    if token is None:
-        authorization = websocket.headers.get("authorization", "")
-        scheme, _, value = authorization.partition(" ")
-        if scheme.lower() == "bearer" and value:
-            token = value
-
-    if not token:
-        raise InvalidTokenError
-
-    await authenticate_token(token)
-
-
-def proxy_log_path(proxy) -> Path:
-    if proxy.log_path:
-        return Path(proxy.log_path)
-    return manager.log_path_for(proxy)
-
-
-def serialize_health_snapshot(
-    snapshot: HealthSnapshot | None,
-) -> dict[str, object] | None:
-    if snapshot is None:
-        return None
-    return snapshot.model_dump(mode="json")
-
-
-async def read_log_snapshot(path: Path, limit: int = 8000) -> str:
-    snapshot, _offset, _marker = await snapshot_log_state(path, limit=limit)
-    return snapshot
-
-
-async def snapshot_log_state(
-    path: Path,
-    limit: int = 8000,
-    marker_width: int = 64,
-) -> tuple[str, int, bytes]:
-    if not await asyncio.to_thread(path.is_file):
-        return "", 0, b""
-
-    data = await asyncio.to_thread(path.read_bytes)
-    size = len(data)
-    marker = data[max(0, size - marker_width) : size]
-    snapshot = data[-limit:].decode("utf-8", errors="replace")
-    return snapshot, size, marker
-
-
-def _read_log_marker_bytes(path: Path, offset: int, width: int) -> bytes:
-    if offset <= 0 or not path.is_file():
-        return b""
-
-    start = max(0, offset - width)
-    with path.open("rb") as handle:
-        handle.seek(start)
-        return handle.read(offset - start)
-
-
-async def read_log_marker(path: Path, offset: int, width: int = 64) -> bytes:
-    return await asyncio.to_thread(_read_log_marker_bytes, path, offset, width)
-
-
-async def watch_socket_disconnect(websocket: WebSocket) -> None:
-    while True:
-        await websocket.receive_text()
-
-
-async def stream_proxy_logs(
-    websocket: WebSocket,
-    path: Path,
-    *,
-    offset: int,
-    marker: bytes,
-) -> None:
-
-    while True:
-        if not await asyncio.to_thread(path.exists):
-            await asyncio.sleep(0.4)
-            continue
-
-        try:
-            size = await asyncio.to_thread(lambda: path.stat().st_size)
-        except OSError:
-            await asyncio.sleep(0.4)
-            continue
-
-        if marker and size >= offset and await read_log_marker(path, offset) != marker:
-            snapshot, offset, marker = await snapshot_log_state(path)
-            await websocket.send_json({"type": "snapshot", "log": snapshot})
-            await asyncio.sleep(0.4)
-            continue
-
-        if size < offset:
-            snapshot, offset, marker = await snapshot_log_state(path)
-            await websocket.send_json({"type": "snapshot", "log": snapshot})
-            await asyncio.sleep(0.4)
-            continue
-
-        if size > offset:
-            def read_log_chunk(
-                chunk_offset: int = offset,
-                chunk_size: int = size,
-            ) -> str:
-                with path.open("rb") as handle:
-                    handle.seek(chunk_offset)
-                    return handle.read(chunk_size - chunk_offset).decode(
-                        "utf-8",
-                        errors="replace",
-                    )
-
-            chunk = await asyncio.to_thread(read_log_chunk)
-            offset = size
-            marker = await read_log_marker(path, offset)
-            if chunk:
-                await websocket.send_json({"type": "append", "log": chunk})
-
-        await asyncio.sleep(0.4)
-
-
 @router.get("", dependencies=auth_dependencies)
 async def list_proxies_handler():
     return await list_proxy_documents()
@@ -295,7 +181,7 @@ async def list_mitmproxy_options():
 
 @router.get("/mitmproxy/templates", dependencies=auth_dependencies)
 async def list_mitmproxy_templates():
-    return [*BUILTIN_TEMPLATES, *await list_custom_templates()]
+    return [*list_file_templates(), *await list_custom_templates()]
 
 
 @router.post(
@@ -305,6 +191,11 @@ async def list_mitmproxy_templates():
 )
 async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
     await _validate_code_or_raise(payload.code)
+    if payload.id in file_template_ids():
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail=f"Template already exists: {payload.id}",
+        )
     try:
         return await create_custom_template(payload)
     except DuplicateKeyError as error:
@@ -320,7 +211,7 @@ async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
     dependencies=auth_dependencies,
 )
 async def delete_mitmproxy_template(template_id: str):
-    if any(template.id == template_id for template in BUILTIN_TEMPLATES):
+    if template_id in file_template_ids():
         raise HTTPException(
             status_code=status.HTTP_409_CONFLICT,
             detail="Built-in templates cannot be deleted",