From 7be0140ee230f26b1a94124f12ffea258df2deb7 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:33:45 +0800 Subject: [PATCH 1/6] refactor(service): centralize task and artifact foundations Introduce a local ArtifactStore and route output path validation, upload handling, task submission, wait, cancel, and output responses through shared service helpers. Keep public API paths intact while reducing duplicated Task/OpenAI/File route behavior and preserving explicit local-only artifact semantics. Includes focused service and OpenAI route tests for the centralized paths. --- telefuser/entrypoints/cli/main.py | 4 +- telefuser/service/api/api_server.py | 106 ++++++- telefuser/service/api/middleware.py | 28 +- telefuser/service/api/openai/image_routes.py | 137 ++------ telefuser/service/api/openai/video_routes.py | 111 ++----- telefuser/service/api/routers/files.py | 13 +- telefuser/service/api/routers/service.py | 33 +- telefuser/service/api/routers/tasks.py | 102 +++--- .../service/api/task_application_service.py | 272 ++++++++++++++++ telefuser/service/core/artifact_store.py | 256 +++++++++++++++ telefuser/service/core/config.py | 29 ++ telefuser/service/core/container.py | 7 + telefuser/service/core/file_service.py | 293 ++++++++++++------ telefuser/service/core/task_manager.py | 27 +- telefuser/service/core/task_service.py | 15 +- telefuser/service/main.py | 19 +- telefuser/service_types.py | 1 + tests/server/test_middleware.py | 103 +++--- tests/unit/openai/_asgi_test_client.py | 38 +++ tests/unit/openai/test_image_routes.py | 28 +- tests/unit/openai/test_video_routes.py | 27 +- tests/unit/service/test_artifact_store.py | 173 +++++++++++ tests/unit/service/test_config.py | 10 + .../service/test_service_refactor_phase0.py | 175 +++++++++++ .../service/test_task_application_service.py | 205 ++++++++++++ tests/unit/service/test_task_routes.py | 54 +++- tests/unit/service/test_task_runtime.py | 15 + 27 files changed, 1791 insertions(+), 490 deletions(-) create mode 100644 telefuser/service/api/task_application_service.py create mode 100644 telefuser/service/core/artifact_store.py create mode 100644 tests/unit/openai/_asgi_test_client.py create mode 100644 tests/unit/service/test_artifact_store.py create mode 100644 tests/unit/service/test_service_refactor_phase0.py create mode 100644 tests/unit/service/test_task_application_service.py diff --git a/telefuser/entrypoints/cli/main.py b/telefuser/entrypoints/cli/main.py index 6978508..35cda0c 100644 --- a/telefuser/entrypoints/cli/main.py +++ b/telefuser/entrypoints/cli/main.py @@ -127,6 +127,8 @@ def serve( num_replicas=num_replicas, enable_latent_cache=enable_latent_cache, cache_mode=cache_mode.lower() if cache_mode is not None else None, + security_level=security_level, + skip_validation=skip_validation, ) @@ -176,7 +178,7 @@ def stream_serve( from telefuser.service.main import run_stream_server - run_stream_server(pipe_path, port, host, skip_validation=True) + run_stream_server(pipe_path, port, host, skip_validation=skip_validation, security_level=security_level) @main.command(name="validate") diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index 29f382d..0f196ff 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -5,7 +5,7 @@ import asyncio import inspect from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from urllib.parse import urlparse import httpx @@ -15,13 +15,17 @@ from telefuser.utils.logging import logger +from ..core.config import ServerConfig, server_config from ..core.file_service import FileService -from ..core.pipeline_service import PipelineService -from ..core.stream_pipeline_service import StreamPipelineService -from ..core.task_manager import TaskManager, TaskStatus +from ..core.task_manager import TaskManager from ..core.task_processor import AsyncTaskProcessor from ..core.task_service import MediaGenerationService from . import routers +from .task_application_service import TaskApplicationService + +if TYPE_CHECKING: + from ..core.pipeline_service import PipelineService + from ..core.stream_pipeline_service import StreamPipelineService class ApiServer: @@ -41,7 +45,9 @@ def __init__( enable_rate_limit: bool = True, enable_logging: bool = False, enable_openai_api: bool = True, + config: ServerConfig | None = None, ) -> None: + self.server_config = config or server_config self.app = app or FastAPI( title="TeleFuser API", description="API for video and image generation using TeleFuser framework. " @@ -56,6 +62,7 @@ def __init__( self.stream_service: StreamPipelineService | None = None self._webrtc_routes: object | None = None self.media_service: MediaGenerationService | None = None + self.task_app_service = TaskApplicationService(self) self.cache_service: Any | None = None self.max_queue_size = max_queue_size self.max_concurrent_tasks = max_concurrent_tasks @@ -65,12 +72,19 @@ def __init__( self.task_processor: AsyncTaskProcessor | None = None self._task_processor_lock = asyncio.Lock() + self._artifact_cleanup_task: asyncio.Task | None = None + self._artifact_cleanup_lock = asyncio.Lock() self._setup_routes() from .middleware import setup_middleware - setup_middleware(self.app, enable_rate_limit=enable_rate_limit, enable_logging=enable_logging) + setup_middleware( + self.app, + enable_rate_limit=enable_rate_limit, + enable_logging=enable_logging, + config=self.server_config, + ) @property def task_manager(self): @@ -122,10 +136,9 @@ def _stream_file_response(self, file_path: Path, filename: str | None = None) -> assert self.file_service is not None, "File service is not initialized" try: - resolved_path = file_path.resolve() - - # Security: ensure file is within allowed directories - if not str(resolved_path).startswith(str(self.file_service.output_video_dir.resolve())): + try: + resolved_path = self.file_service.resolve_output_file(file_path) + except ValueError: raise HTTPException(status_code=403, detail="Access to this file is not allowed") if not resolved_path.exists() or not resolved_path.is_file(): @@ -164,8 +177,6 @@ def file_stream_generator(file_path: str, chunk_size: int = 1024 * 1024): async def _validate_image_url(self, image_url: str) -> bool: """Validate image URL is accessible.""" - from ..core.config import server_config - if not image_url or not image_url.startswith("http"): return True @@ -175,7 +186,11 @@ async def _validate_image_url(self, image_url: str) -> bool: return False timeout = httpx.Timeout(connect=5.0, read=5.0) - verify = server_config.ssl_cert_path if server_config.ssl_cert_path else server_config.verify_ssl + verify = ( + self.server_config.ssl_cert_path + if self.server_config.ssl_cert_path + else self.server_config.verify_ssl + ) async with httpx.AsyncClient(verify=verify, timeout=timeout) as client: response = await client.head(image_url, follow_redirects=True) return response.status_code < 400 @@ -190,6 +205,7 @@ async def ensure_task_processor_running(self) -> None: return if self.task_processor.is_running: + await self.ensure_artifact_cleanup_running() return async with self._task_processor_lock: @@ -199,6 +215,48 @@ async def ensure_task_processor_running(self) -> None: if self.task_processor.is_running: return await self.task_processor.start() + await self.ensure_artifact_cleanup_running() + + async def ensure_artifact_cleanup_running(self) -> None: + """Ensure periodic artifact cleanup is running when file service is available.""" + if self.file_service is None: + return + if self._artifact_cleanup_task is not None and not self._artifact_cleanup_task.done(): + return + + async with self._artifact_cleanup_lock: + if self.file_service is None: + return + if self._artifact_cleanup_task is not None and not self._artifact_cleanup_task.done(): + return + self._artifact_cleanup_task = asyncio.create_task( + self._artifact_cleanup_loop(), + name="telefuser-artifact-cleanup", + ) + + async def _artifact_cleanup_loop(self) -> None: + interval = self.server_config.artifact_cleanup_interval_seconds + while True: + await asyncio.sleep(interval) + try: + self.run_artifact_cleanup() + except Exception as exc: + logger.warning(f"Periodic artifact cleanup failed: {exc}") + + def run_artifact_cleanup(self, *, now: Any | None = None) -> dict[str, Any]: + """Run one artifact cleanup pass and return cleanup stats.""" + if self.file_service is None: + return {} + + cleanup_artifacts = getattr(self.file_service, "cleanup_artifacts", None) + if not callable(cleanup_artifacts): + return {} + + cleanup_snapshot = self.task_manager.get_artifact_cleanup_snapshot() + result = cleanup_artifacts(**cleanup_snapshot, now=now) + if result.get("removed_task_ids") or result.get("removed_tmp_files"): + logger.info(f"Artifact cleanup result: {result}") + return result def get_supported_tasks(self) -> tuple[str, ...]: """Get tasks supported by the loaded pipeline contract, if available.""" @@ -273,9 +331,18 @@ def initialize_services( inference_service: PipelineService, cache_service: Any | None = None, cache_adapter: Any | None = None, + file_service: FileService | None = None, ) -> None: """Initialize file and media services.""" - self.file_service = FileService(cache_dir) + self.file_service = file_service or FileService( + cache_dir, + max_file_size=self.server_config.max_file_size, + verify_ssl=self.server_config.verify_ssl, + ssl_cert_path=self.server_config.ssl_cert_path, + artifact_retention_seconds=self.server_config.artifact_retention_seconds, + artifact_tmp_retention_seconds=self.server_config.artifact_tmp_retention_seconds, + artifact_max_total_bytes=self.server_config.artifact_max_total_bytes, + ) self.inference_service = inference_service self.cache_service = cache_service self.cache_adapter = cache_adapter @@ -303,6 +370,14 @@ def initialize_stream_service(self, stream_service: StreamPipelineService) -> No async def cleanup(self) -> None: """Cleanup resources and stop processing workers.""" + if self._artifact_cleanup_task is not None: + self._artifact_cleanup_task.cancel() + try: + await self._artifact_cleanup_task + except asyncio.CancelledError: + pass + self._artifact_cleanup_task = None + if self.task_processor is not None: await self.task_processor.stop() @@ -313,6 +388,11 @@ async def cleanup(self) -> None: await self._webrtc_routes.cleanup() if self.file_service: + try: + self.run_artifact_cleanup() + except Exception as exc: + logger.warning(f"Artifact cleanup failed: {exc}") + cleanup = getattr(self.file_service, "cleanup", None) if cleanup is not None: result = cleanup() diff --git a/telefuser/service/api/middleware.py b/telefuser/service/api/middleware.py index 0de37e1..87ae5dc 100644 --- a/telefuser/service/api/middleware.py +++ b/telefuser/service/api/middleware.py @@ -17,11 +17,11 @@ from telefuser.utils.logging import logger -from ..core.config import server_config - if TYPE_CHECKING: from telefuser.metrics import MetricRegistry + from ..core.config import ServerConfig + class RateLimitErrorResponse: """Structured response for rate limit errors.""" @@ -83,12 +83,13 @@ def __init__( requests_per_minute: int | None = None, window_size: int | None = None, limited_paths: list[str] | None = None, + enabled: bool = True, ): super().__init__(app) - config = server_config - self.requests_per_minute = requests_per_minute or config.rate_limit_requests_per_minute - self.window_size = window_size or config.rate_limit_window_size - self.limited_paths = tuple(limited_paths if limited_paths is not None else config.rate_limit_paths) + self.enabled = enabled + self.requests_per_minute = requests_per_minute or 60 + self.window_size = window_size or 60 + self.limited_paths = tuple(limited_paths if limited_paths is not None else []) self._clients: dict[str, list[float]] = {} @@ -104,7 +105,7 @@ async def dispatch(self, request: Request, call_next: callable) -> Response: response.headers["X-Request-ID"] = request_id return response - if not server_config.enable_rate_limit: + if not self.enabled: response = await call_next(request) response.headers["X-Request-ID"] = request_id return response @@ -306,6 +307,7 @@ def setup_middleware( enable_logging: bool = True, enable_metrics: bool = True, metrics_registry: MetricRegistry | None = None, + config: ServerConfig | None = None, ) -> None: """Setup middleware for FastAPI app. @@ -325,10 +327,14 @@ def setup_middleware( registry=metrics_registry, ) - if enable_rate_limit and server_config.enable_rate_limit: + if config is None: + from ..core.config import server_config as config + + if enable_rate_limit and config.enable_rate_limit: app.add_middleware( RateLimitMiddleware, - requests_per_minute=server_config.rate_limit_requests_per_minute, - window_size=server_config.rate_limit_window_size, - limited_paths=server_config.rate_limit_paths, + requests_per_minute=config.rate_limit_requests_per_minute, + window_size=config.rate_limit_window_size, + limited_paths=config.rate_limit_paths, + enabled=config.enable_rate_limit, ) diff --git a/telefuser/service/api/openai/image_routes.py b/telefuser/service/api/openai/image_routes.py index 2199c1a..9517a3d 100644 --- a/telefuser/service/api/openai/image_routes.py +++ b/telefuser/service/api/openai/image_routes.py @@ -11,23 +11,18 @@ from __future__ import annotations -import asyncio import base64 -import time -from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile -from fastapi.responses import FileResponse +from fastapi.responses import Response -from telefuser.service.api.schema import TaskRequest from telefuser.service.api.task_contract_runtime import ( - apply_task_contract_defaults, map_contract_fields, match_task_candidates, - validate_required_task_parameters, ) -from telefuser.service.core.task_manager import TaskManager, TaskStatus +from telefuser.service.core.task_manager import TaskManager +from telefuser.service_types import MediaType from telefuser.utils.logging import logger from .adapter import OpenAIRequestAdapter, OpenAIResponseAdapter @@ -77,7 +72,7 @@ async def create_image_edit( ) @router.get("/{image_id}/content") - async def get_image_content(image_id: str) -> FileResponse: + async def get_image_content(image_id: str) -> Response: """Download a generated image by its ID.""" return await routes.get_image_content(image_id) @@ -94,7 +89,6 @@ def __init__(self, api_server: ApiServer) -> None: """Initialize with ApiServer instance.""" self.api = api_server self.task_manager: TaskManager | None = getattr(api_server, "task_manager", None) - self.file_service = getattr(api_server, "file_service", None) self.media_service = getattr(api_server, "media_service", None) async def create_image_generation(self, request: ImageGenerationsRequest) -> ImageResponse: @@ -108,20 +102,21 @@ async def create_image_generation(self, request: ImageGenerationsRequest) -> Ima try: task_type = self._resolve_image_task(has_reference=False) task_request = OpenAIRequestAdapter.to_task_request(request, task_type=task_type) - self.api.validate_task_supported(task_request.task) - contract = self.api.get_task_contract(task_request.task) - apply_task_contract_defaults( + task_response = await self.api.task_app_service.submit( task_request, - task_contract=contract, explicit_fields=self._get_image_generation_explicit_fields(request), + ensure_processing=False, ) - validate_required_task_parameters(task_request, task_contract=contract) - task_id = self.task_manager.create_task(task_request) + task_id = task_response.task_id logger.info(f"Created image generation task: {task_id}") await self._ensure_processing() - result = await self._wait_for_task_completion(task_id=task_id, timeout=self.DEFAULT_TIMEOUT) + result = await self.api.task_app_service.wait_for_completion( + task_id=task_id, + timeout=self.DEFAULT_TIMEOUT, + poll_interval=self.POLL_INTERVAL, + ) if result is None: self.task_manager.cancel_task(task_id) @@ -201,11 +196,8 @@ async def create_image_edit( task_type = self._resolve_image_task(has_reference=True) task_request = OpenAIRequestAdapter.to_task_request(edit_request, task_type=task_type) - self.api.validate_task_supported(task_request.task) - contract = self.api.get_task_contract(task_request.task) - apply_task_contract_defaults( + task_response = await self.api.task_app_service.submit( task_request, - task_contract=contract, explicit_fields=self._get_image_edit_explicit_fields( prompt=prompt, image_path=image_path, @@ -217,14 +209,18 @@ async def create_image_edit( seed=seed, negative_prompt=negative_prompt, ), + ensure_processing=False, ) - validate_required_task_parameters(task_request, task_contract=contract) - task_id = self.task_manager.create_task(task_request) + task_id = task_response.task_id logger.info(f"Created image edit task: {task_id}") await self._ensure_processing() - result = await self._wait_for_task_completion(task_id=task_id, timeout=self.DEFAULT_TIMEOUT) + result = await self.api.task_app_service.wait_for_completion( + task_id=task_id, + timeout=self.DEFAULT_TIMEOUT, + poll_interval=self.POLL_INTERVAL, + ) if result is None: self.task_manager.cancel_task(task_id) @@ -252,41 +248,9 @@ async def create_image_edit( logger.exception(f"Image edit failed: {e}") raise HTTPException(status_code=500, detail=f"Image edit failed: {str(e)}") - async def get_image_content(self, image_id: str) -> FileResponse: + async def get_image_content(self, image_id: str) -> Response: """Retrieve generated image by ID.""" - if not self.task_manager: - raise HTTPException(status_code=503, detail="Task manager not initialized") - - if not self.file_service: - raise HTTPException(status_code=503, detail="File service not initialized") - - task_info = self.task_manager.get_task(image_id) - if not task_info: - raise HTTPException(status_code=404, detail=f"Image {image_id} not found") - - output_path = task_info.output_path - if not output_path: - raise HTTPException(status_code=404, detail=f"Image {image_id} has no output path") - - path = Path(output_path) - if not path.is_absolute(): - output_dir = getattr(self.file_service, "output_image_dir", None) - if output_dir: - path = output_dir / path - - if not path.exists(): - raise HTTPException(status_code=404, detail=f"Image file not found: {path}") - - suffix = path.suffix.lower() - media_types = { - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".webp": "image/webp", - } - media_type = media_types.get(suffix, "image/png") - - return FileResponse(path=str(path), media_type=media_type, filename=path.name) + return self.api.task_app_service.get_output_response(image_id, media_type=MediaType.IMAGE) async def _ensure_processing(self) -> None: """Ensure the task processor is running.""" @@ -301,57 +265,14 @@ def _get_base_url(self) -> str: return f"http://{host}:{port}" return "http://localhost:8000" - async def _wait_for_task_completion(self, task_id: str, timeout: float) -> dict[str, Any] | None: - """Wait for a task to complete with timeout.""" - start_time = time.time() - - while time.time() - start_time < timeout: - status = self.task_manager.get_task_status(task_id) - - if not status: - raise HTTPException(status_code=404, detail=f"Task {task_id} not found") - - task_status = status.get("status") - - if task_status == TaskStatus.COMPLETED.value: - return { - "output_path": status.get("output_path"), - "error": status.get("error"), - } - elif task_status == TaskStatus.FAILED.value: - error_msg = status.get("error", "Unknown error") - raise HTTPException(status_code=500, detail=f"Generation failed: {error_msg}") - elif task_status == TaskStatus.CANCELLED.value: - raise HTTPException(status_code=400, detail="Generation was cancelled") - - await asyncio.sleep(self.POLL_INTERVAL) - - raise HTTPException(status_code=504, detail=f"Generation timeout after {timeout} seconds") - async def _save_uploaded_file(self, file: UploadFile, prefix: str = "upload") -> str: """Save an uploaded file to disk.""" - import uuid - - if not self.file_service: - raise HTTPException(status_code=503, detail="File service not initialized") - - input_dir = getattr(self.file_service, "input_image_dir", None) - if not input_dir: - input_dir = Path("/tmp/telefuser/inputs") - input_dir.mkdir(parents=True, exist_ok=True) - - suffix = Path(file.filename or "image.png").suffix - filename = f"{prefix}_{uuid.uuid4().hex[:8]}{suffix}" - file_path = input_dir / filename - - content = await file.read() - await asyncio.to_thread(self._write_file_sync, file_path, content) - - return str(file_path) - - def _write_file_sync(self, file_path: Path, content: bytes) -> None: - """Write file synchronously.""" - file_path.write_bytes(content) + return await self.api.task_app_service.save_upload_file( + file, + media_type=MediaType.IMAGE, + prefix=prefix, + fallback_filename="image.png", + ) def _resolve_image_task(self, *, has_reference: bool) -> str: """Resolve the best image task supported by the current pipeline.""" diff --git a/telefuser/service/api/openai/video_routes.py b/telefuser/service/api/openai/video_routes.py index b7ff65c..69b4a0f 100644 --- a/telefuser/service/api/openai/video_routes.py +++ b/telefuser/service/api/openai/video_routes.py @@ -13,21 +13,18 @@ from __future__ import annotations -import asyncio -from pathlib import Path -from typing import TYPE_CHECKING, Any, List +from typing import TYPE_CHECKING, List from fastapi import APIRouter, File, Form, HTTPException, Query, Request, UploadFile -from fastapi.responses import FileResponse +from fastapi.responses import Response from telefuser.service.api.task_contract_runtime import ( - apply_task_contract_defaults, map_contract_fields, match_task_candidates, - validate_required_task_parameters, ) from telefuser.service.core.pipeline_contract import is_video_task from telefuser.service.core.task_manager import TaskManager, TaskStatus +from telefuser.service_types import MediaType from telefuser.utils.logging import logger from .adapter import OpenAIRequestAdapter, OpenAIResponseAdapter, is_probable_video_reference @@ -121,7 +118,7 @@ async def delete_video(video_id: str) -> VideoResponse: return await routes.delete_video(video_id) @router.get("/{video_id}/content") - async def get_video_content(video_id: str) -> FileResponse: + async def get_video_content(video_id: str) -> Response: """Download a generated video by its ID.""" return await routes.get_video_content(video_id) @@ -135,7 +132,6 @@ def __init__(self, api_server: ApiServer) -> None: """Initialize with ApiServer instance.""" self.api = api_server self.task_manager: TaskManager | None = getattr(api_server, "task_manager", None) - self.file_service = getattr(api_server, "file_service", None) async def create_video( self, @@ -149,18 +145,15 @@ async def create_video( try: task_type = self._resolve_video_task(request) task_request = OpenAIRequestAdapter.to_task_request(request, task_type=task_type) - self.api.validate_task_supported(task_request.task) - contract = self.api.get_task_contract(task_request.task) - apply_task_contract_defaults( + task_response = await self.api.task_app_service.submit( task_request, - task_contract=contract, explicit_fields=self._get_video_explicit_fields( task_type=task_type, source_fields=explicit_source_fields or set(getattr(request, "model_fields_set", set())), ), + ensure_processing=False, ) - validate_required_task_parameters(task_request, task_contract=contract) - task_id = self.task_manager.create_task(task_request) + task_id = task_response.task_id logger.info(f"Created video generation task: {task_id}") await self._ensure_processing() @@ -279,64 +272,26 @@ async def retrieve_video(self, video_id: str) -> VideoResponse: async def delete_video(self, video_id: str) -> VideoResponse: """Cancel or delete a video generation task.""" - if not self.task_manager: - raise HTTPException(status_code=503, detail="Task manager not initialized") - - task_status = self.task_manager.get_task_status(video_id) - if not task_status: + cancel_result = self.api.task_app_service.cancel_task(video_id) + if cancel_result["result"] == "not_found": raise HTTPException(status_code=404, detail=f"Video {video_id} not found") - cancelled = self.task_manager.cancel_task(video_id) + task_status = cancel_result["task_status"] or TaskStatus.CANCELLED.value + video_status = OpenAIResponseAdapter.task_status_to_video_status(task_status) return VideoResponse( id=video_id, - status="cancelled" if cancelled else "deleted", + status=video_status, model="wan-video", ) - async def get_video_content(self, video_id: str) -> FileResponse: + async def get_video_content(self, video_id: str) -> Response: """Download generated video.""" - if not self.task_manager: - raise HTTPException(status_code=503, detail="Task manager not initialized") - - if not self.file_service: - raise HTTPException(status_code=503, detail="File service not initialized") - - task_info = self.task_manager.get_task(video_id) - if not task_info: - raise HTTPException(status_code=404, detail=f"Video {video_id} not found") - - task_status = self.task_manager.get_task_status(video_id) - if task_status.get("status") != TaskStatus.COMPLETED.value: - raise HTTPException( - status_code=400, - detail=f"Video {video_id} is not ready (status: {task_status.get('status')})", - ) - - output_path = task_info.output_path or task_status.get("output_path") - if not output_path: - raise HTTPException(status_code=404, detail=f"Video {video_id} has no output path") - - path = Path(output_path) - if not path.is_absolute(): - output_dir = getattr(self.file_service, "output_video_dir", None) - if output_dir: - path = output_dir / path - - if not path.exists(): - raise HTTPException(status_code=404, detail=f"Video file not found: {path}") - - suffix = path.suffix.lower() - media_types = { - ".mp4": "video/mp4", - ".avi": "video/x-msvideo", - ".mov": "video/quicktime", - ".mkv": "video/x-matroska", - ".webm": "video/webm", - } - media_type = media_types.get(suffix, "video/mp4") - - return FileResponse(path=str(path), media_type=media_type, filename=path.name) + return self.api.task_app_service.get_output_response( + video_id, + media_type=MediaType.VIDEO, + require_completed=True, + ) async def _ensure_processing(self) -> None: """Ensure the task processor is running.""" @@ -344,30 +299,12 @@ async def _ensure_processing(self) -> None: async def _save_uploaded_file(self, file: UploadFile, prefix: str = "upload") -> str: """Save an uploaded file to disk.""" - import uuid - - if not self.file_service: - raise HTTPException(status_code=503, detail="File service not initialized") - - input_dir = getattr(self.file_service, "input_video_dir", None) - if not input_dir: - input_dir = getattr(self.file_service, "input_image_dir", None) - if not input_dir: - input_dir = Path("/tmp/telefuser/inputs") - input_dir.mkdir(parents=True, exist_ok=True) - - suffix = Path(file.filename or "input.mp4").suffix - filename = f"{prefix}_{uuid.uuid4().hex[:8]}{suffix}" - file_path = input_dir / filename - - content = await file.read() - await asyncio.to_thread(self._write_file_sync, file_path, content) - - return str(file_path) - - def _write_file_sync(self, file_path: Path, content: bytes) -> None: - """Write file synchronously.""" - file_path.write_bytes(content) + return await self.api.task_app_service.save_upload_file( + file, + media_type=MediaType.VIDEO, + prefix=prefix, + fallback_filename="input.mp4", + ) def _resolve_video_task(self, request: VideoGenerationsRequest) -> str: """Resolve the best video task supported by the current pipeline.""" diff --git a/telefuser/service/api/routers/files.py b/telefuser/service/api/routers/files.py index 5da62c9..b0d9f81 100644 --- a/telefuser/service/api/routers/files.py +++ b/telefuser/service/api/routers/files.py @@ -6,7 +6,6 @@ from __future__ import annotations -from pathlib import Path from typing import TYPE_CHECKING from fastapi import APIRouter, HTTPException @@ -31,16 +30,10 @@ async def download_file(self, file_path: str) -> StreamingResponse: assert self.api.file_service is not None, "File service is not initialized" try: - # Try video directory first - full_path = self.api.file_service.output_video_dir / file_path - if not full_path.exists(): - # Try image directory - full_path = self.api.file_service.output_image_dir / file_path - if not full_path.exists(): - # Try general output directory - full_path = self.api.file_service.output_dir / file_path - + full_path = self.api.file_service.resolve_output_file(file_path) return self.api._stream_file_response(full_path) + except ValueError: + raise HTTPException(status_code=403, detail="Access to this file is not allowed") except HTTPException: raise except Exception as e: diff --git a/telefuser/service/api/routers/service.py b/telefuser/service/api/routers/service.py index 2b18ea1..ac390b2 100644 --- a/telefuser/service/api/routers/service.py +++ b/telefuser/service/api/routers/service.py @@ -8,7 +8,8 @@ from typing import TYPE_CHECKING -from fastapi import APIRouter, HTTPException, Response +from fastapi import APIRouter, HTTPException, Response, status +from fastapi.responses import JSONResponse from telefuser.metrics import get_service_metrics @@ -65,11 +66,12 @@ async def get_metadata(self) -> dict: raise HTTPException(status_code=503, detail="No service is initialized") async def health_check(self) -> dict: - """Health check endpoint for monitoring.""" + """Liveness endpoint for monitoring.""" from datetime import datetime status = { "status": "healthy", + "ready": self._is_ready(), "timestamp": datetime.utcnow().isoformat(), "version": "1.0.0", } @@ -84,6 +86,29 @@ async def health_check(self) -> dict: return status + def _is_ready(self) -> bool: + """Return whether the initialized service can currently accept work.""" + if self.api.inference_service is not None: + if hasattr(self.api.inference_service, "pool_status"): + pool_status = self.api.inference_service.pool_status() + return bool(pool_status.get("alive_replicas", 0) > 0) + return bool(getattr(self.api.inference_service, "is_running", False)) + + if self.api.stream_service is not None: + return bool(getattr(self.api.stream_service, "is_running", False)) + + return False + + async def readiness_check(self) -> JSONResponse: + """Readiness endpoint for load balancers.""" + body = await self.health_check() + ready = bool(body.get("ready", False)) + body["status"] = "ready" if ready else "not_ready" + return JSONResponse( + status_code=status.HTTP_200_OK if ready else status.HTTP_503_SERVICE_UNAVAILABLE, + content=body, + ) + def create_router(api_server: ApiServer) -> APIRouter: """Create a new router with fresh routes for the given ApiServer instance.""" @@ -102,6 +127,10 @@ async def get_service_metadata() -> dict: async def health_check() -> dict: return await routes.health_check() + @new_router.get("/ready") + async def readiness_check() -> JSONResponse: + return await routes.readiness_check() + @new_router.get("/metrics", summary="Get Prometheus Metrics") async def get_metrics_endpoint() -> Response: """Get Prometheus-compatible metrics.""" diff --git a/telefuser/service/api/routers/tasks.py b/telefuser/service/api/routers/tasks.py index a77a762..b34f1c2 100644 --- a/telefuser/service/api/routers/tasks.py +++ b/telefuser/service/api/routers/tasks.py @@ -7,8 +7,6 @@ from __future__ import annotations -import asyncio -import uuid from pathlib import Path from typing import TYPE_CHECKING, Any @@ -16,15 +14,11 @@ from pydantic import ValidationError from telefuser.service.core.pipeline_contract import default_task_contract, validate_task_name_format -from telefuser.service_types import AspectRatio, OutputFormat, StopTaskStatus, TaskStatus +from telefuser.service_types import AspectRatio, MediaType, OutputFormat, StopTaskStatus from telefuser.utils.logging import logger from ..schema import StopTaskResponse, TaskRequest, TaskResponse -from ..task_contract_runtime import ( - apply_task_contract_defaults, - match_task_candidates, - validate_required_task_parameters, -) +from ..task_contract_runtime import match_task_candidates if TYPE_CHECKING: from ..api_server import ApiServer @@ -108,28 +102,13 @@ async def check_image_path(image_name: str) -> None: raise HTTPException(status_code=400, detail=f"{image_name} URL is not accessible") try: - self.api.validate_task_supported(message.task) - contract = self._get_task_contract(message.task) - apply_task_contract_defaults( - message, - task_contract=contract, - explicit_fields=set(getattr(message, "model_fields_set", set())), - ) - validate_required_task_parameters(message, task_contract=contract) await check_image_path("first_image_path") await check_image_path("last_image_path") - self._validate_task_inputs( - message.task, - first_image_path=message.first_image_path, - last_image_path=message.last_image_path, - ref_video_path=message.ref_video_path, + return await self.api.task_app_service.submit( + message, + explicit_fields=set(getattr(message, "model_fields_set", set())), + validate_inputs=self._validate_message_inputs, ) - task_id = self.api.task_manager.create_task(message) - message.task_id = task_id - await self.api.ensure_task_processor_running() - return TaskResponse(task_id=task_id, task_status=TaskStatus.PENDING, output_path=message.output_path) - except RuntimeError as e: - raise HTTPException(status_code=503, detail=str(e)) except HTTPException: raise except Exception as e: @@ -162,7 +141,8 @@ async def get_task_status(self, task_id: str) -> dict | None: async def stop_task(self, task_id: str) -> StopTaskResponse: """Stop/cancel a task and clean up resources.""" try: - if self.api.task_manager.cancel_task(task_id): + cancel_result = self.api.task_app_service.cancel_task(task_id) + if cancel_result["result"] == "accepted": import gc import torch @@ -170,12 +150,14 @@ async def stop_task(self, task_id: str) -> StopTaskResponse: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - logger.info(f"Task {task_id} stopped successfully.") - return StopTaskResponse(stop_status=StopTaskStatus.SUCCESS, reason="Task stopped successfully.") - else: + logger.info(f"Task {task_id} cancellation accepted.") + return StopTaskResponse(stop_status=StopTaskStatus.SUCCESS, reason="Task cancellation accepted.") + if cancel_result["result"] == "already_terminal": return StopTaskResponse( - stop_status=StopTaskStatus.DO_NOTHING, reason="Task not found or already completed." + stop_status=StopTaskStatus.DO_NOTHING, + reason=f"Task already terminal: {cancel_result['task_status']}.", ) + return StopTaskResponse(stop_status=StopTaskStatus.DO_NOTHING, reason="Task not found.") except Exception as e: logger.error(f"Error occurred while stopping task {task_id}: {str(e)}") return StopTaskResponse(stop_status=StopTaskStatus.ERROR, reason=str(e)) @@ -196,32 +178,30 @@ async def create_task_form( """Create task with file upload support.""" assert self.api.file_service is not None, "File service is not initialized" - async def save_file_async(file: UploadFile, target_dir: Path) -> str: + async def save_file_async(file: UploadFile, media_type: MediaType) -> str: if not file or not file.filename: return "" - file_extension = Path(file.filename).suffix - unique_filename = f"{uuid.uuid4()}{file_extension}" - file_path = target_dir / unique_filename - - content = await file.read() - await asyncio.to_thread(self._write_file_sync, file_path, content) - - return str(file_path) + return await self.api.task_app_service.save_upload_file( + file, + media_type=media_type, + prefix="input", + fallback_filename="input.mp4" if media_type == MediaType.VIDEO else "input.png", + ) first_image_path = "" ref_video_path = "" if first_image_file and first_image_file.filename: if self._is_video_upload(first_image_file): - ref_video_path = await save_file_async(first_image_file, self.api.file_service.input_video_dir) + ref_video_path = await save_file_async(first_image_file, MediaType.VIDEO) else: - first_image_path = await save_file_async(first_image_file, self.api.file_service.input_image_dir) + first_image_path = await save_file_async(first_image_file, MediaType.IMAGE) last_image_path = "" if last_image_file and last_image_file.filename: if self._is_video_upload(last_image_file): raise HTTPException(status_code=400, detail="last_image_file must be an image upload") - last_image_path = await save_file_async(last_image_file, self.api.file_service.input_image_dir) + last_image_path = await save_file_async(last_image_file, MediaType.IMAGE) try: resolved_task = self._resolve_form_task( @@ -254,27 +234,15 @@ async def save_file_async(file: UploadFile, target_dir: Path) -> str: task_payload["output_format"] = output_format message = TaskRequest(**task_payload) - self.api.validate_task_supported(message.task) - contract = self.api.get_task_contract(message.task) - apply_task_contract_defaults(message, task_contract=contract, explicit_fields=set(task_payload)) - validate_required_task_parameters(message, task_contract=contract) - self._validate_task_inputs( - message.task, - first_image_path=message.first_image_path, - last_image_path=message.last_image_path, - ref_video_path=message.ref_video_path, + return await self.api.task_app_service.submit( + message, + explicit_fields=set(task_payload), + validate_inputs=self._validate_message_inputs, ) - task_id = self.api.task_manager.create_task(message) - message.task_id = task_id - await self.api.ensure_task_processor_running() - - return TaskResponse(task_id=task_id, task_status=TaskStatus.PENDING, output_path=message.output_path) except ValidationError as e: raise HTTPException(status_code=422, detail=e.errors()) except HTTPException: raise - except RuntimeError as e: - raise HTTPException(status_code=503, detail=str(e)) except Exception as e: logger.error(f"Failed to create form task: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -393,6 +361,14 @@ def _validate_task_inputs( if task in {"vc", "vsr"} and not ref_video_path: raise HTTPException(status_code=400, detail=f"Task '{task}' requires ref_video_path") + def _validate_message_inputs(self, message: TaskRequest, contract: dict[str, Any] | None) -> None: + self._validate_task_inputs( + message.task, + first_image_path=message.first_image_path, + last_image_path=message.last_image_path, + ref_video_path=message.ref_video_path, + ) + def _collect_available_inputs( self, *, @@ -420,12 +396,6 @@ def _is_video_upload(self, file: UploadFile) -> bool: suffix = Path(file.filename or "").suffix.lower() return content_type.startswith("video/") or suffix in _VIDEO_FILE_EXTENSIONS - def _write_file_sync(self, file_path: Path, content: bytes) -> None: - """Write file synchronously.""" - file_path.parent.mkdir(parents=True, exist_ok=True) - with open(file_path, "wb") as buffer: - buffer.write(content) - def setup_routes(api_server: ApiServer) -> APIRouter: """Setup routes with ApiServer instance.""" diff --git a/telefuser/service/api/task_application_service.py b/telefuser/service/api/task_application_service.py new file mode 100644 index 0000000..bb5550e --- /dev/null +++ b/telefuser/service/api/task_application_service.py @@ -0,0 +1,272 @@ +"""Canonical task submission service for HTTP-facing APIs.""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from fastapi import HTTPException +from fastapi.responses import FileResponse, Response, StreamingResponse + +from telefuser.service.core.pipeline_contract import infer_media_type_for_task +from telefuser.service_types import MediaType, TaskStatus + +from .schema import TaskRequest, TaskResponse +from .task_contract_runtime import apply_task_contract_defaults, validate_required_task_parameters + +if TYPE_CHECKING: + from .api_server import ApiServer + + +class TaskApplicationService: + """Create tasks through a single validated submission path.""" + + def __init__(self, api_server: ApiServer) -> None: + self.api = api_server + + async def submit( + self, + message: TaskRequest, + *, + explicit_fields: set[str], + validate_inputs: Callable[[TaskRequest, dict[str, Any] | None], None] | None = None, + ensure_processing: bool = True, + ) -> TaskResponse: + """Apply contract defaults, validate paths/required params, enqueue, and start processing.""" + try: + self.api.validate_task_supported(message.task) + contract = self.api.get_task_contract(message.task) + apply_task_contract_defaults(message, task_contract=contract, explicit_fields=explicit_fields) + self.validate_output_path(message) + validate_required_task_parameters(message, task_contract=contract) + if validate_inputs is not None: + validate_inputs(message, contract) + + task_id = self.api.task_manager.create_task(message) + message.task_id = task_id + + if ensure_processing: + await self.api.ensure_task_processor_running() + + return TaskResponse(task_id=task_id, task_status=TaskStatus.PENDING, output_path=message.output_path) + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + + def validate_output_path(self, message: TaskRequest) -> None: + file_service = self.api.file_service + if file_service is None: + return + if getattr(type(file_service), "get_output_path", None) is None: + return + + resolver = getattr(file_service, "get_output_path", None) + if not callable(resolver): + return + + try: + resolver(message.output_path, media_type=infer_media_type_for_task(message.task), task_id=message.task_id) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + + async def wait_for_completion( + self, + task_id: str, + *, + timeout: float, + poll_interval: float = 0.5, + ) -> dict[str, Any]: + """Wait until a task reaches a terminal state and return completed task status.""" + deadline = time.monotonic() + timeout + + while True: + status = self.api.task_manager.get_task_status(task_id) + if not status: + raise HTTPException(status_code=404, detail=f"Task {task_id} not found") + + task_status = status.get("status") + if task_status == TaskStatus.COMPLETED.value: + return status + if task_status == TaskStatus.FAILED.value: + error_msg = status.get("error", "Unknown error") + raise HTTPException(status_code=500, detail=f"Generation failed: {error_msg}") + if task_status == TaskStatus.CANCELLED.value: + raise HTTPException(status_code=400, detail="Generation was cancelled") + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise HTTPException(status_code=504, detail=f"Generation timeout after {timeout} seconds") + + await asyncio.sleep(min(poll_interval, remaining)) + + def get_output_response( + self, + task_id: str, + *, + media_type: MediaType | str, + require_completed: bool = False, + ) -> Response: + """Return a validated streaming response for a task output artifact.""" + task_manager = getattr(self.api, "task_manager", None) + if not task_manager: + raise HTTPException(status_code=503, detail="Task manager not initialized") + + file_service = getattr(self.api, "file_service", None) + if not file_service: + raise HTTPException(status_code=503, detail="File service not initialized") + + media_label = "Image" if str(media_type) == MediaType.IMAGE.value else "Video" + task_info = task_manager.get_task(task_id) + if not task_info: + raise HTTPException(status_code=404, detail=f"{media_label} {task_id} not found") + + task_status = task_manager.get_task_status(task_id) or {} + if require_completed and task_status.get("status") != TaskStatus.COMPLETED.value: + raise HTTPException( + status_code=400, + detail=f"{media_label} {task_id} is not ready (status: {task_status.get('status')})", + ) + + output_path = task_info.output_path or task_status.get("output_path") + if not output_path: + raise HTTPException(status_code=404, detail=f"{media_label} {task_id} has no output path") + + path = self.resolve_output_file(file_service, output_path, media_type=media_type) + if not path.exists() or not path.is_file(): + raise HTTPException(status_code=404, detail=f"{media_label} file not found: {path}") + + stream_response = self.stream_file_response(path) + if stream_response is not None: + return stream_response + + response_media_type = "image/png" if str(media_type) == MediaType.IMAGE.value else "video/mp4" + return FileResponse(path=str(path), media_type=response_media_type, filename=path.name) + + def cancel_task(self, task_id: str) -> dict[str, Any]: + """Cancel a task and distinguish accepted, terminal, and missing outcomes.""" + task_manager = getattr(self.api, "task_manager", None) + if not task_manager: + raise HTTPException(status_code=503, detail="Task manager not initialized") + + task_status = task_manager.get_task_status(task_id) + if not task_status: + return {"result": "not_found", "task_status": None} + + current_status = task_status.get("status") + if current_status in { + TaskStatus.COMPLETED.value, + TaskStatus.FAILED.value, + TaskStatus.CANCELLED.value, + }: + return {"result": "already_terminal", "task_status": current_status} + + if task_manager.cancel_task(task_id): + return {"result": "accepted", "task_status": TaskStatus.CANCELLED.value} + + task_status = task_manager.get_task_status(task_id) + if not task_status: + return {"result": "not_found", "task_status": None} + return {"result": "already_terminal", "task_status": task_status.get("status")} + + async def save_upload_file( + self, + file: Any, + *, + media_type: MediaType | str, + prefix: str = "upload", + fallback_filename: str = "upload.bin", + ) -> str: + """Save an upload through FileService, with a bounded fallback for lightweight mocks.""" + file_service = getattr(self.api, "file_service", None) + if not file_service: + raise HTTPException(status_code=503, detail="File service not initialized") + + writer = self.get_declared_method(file_service, "save_upload_file") + if callable(writer): + path = await writer( + file, + media_type=media_type, + prefix=prefix, + fallback_filename=fallback_filename, + ) + return str(path) + + input_dir = self.input_dir_for_media(file_service, media_type) + suffix = Path(getattr(file, "filename", None) or fallback_filename).suffix + file_path = input_dir / f"{prefix}_{uuid.uuid4().hex[:8]}{suffix}" + await self.write_upload_stream( + file, + file_path, + max_file_size=getattr(file_service, "max_file_size", None), + ) + return str(file_path) + + def input_dir_for_media(self, file_service: Any, media_type: MediaType | str) -> Path: + """Return a legacy input directory for fallback upload writes.""" + if str(media_type) == MediaType.IMAGE.value: + input_dir = getattr(file_service, "input_image_dir", None) + else: + input_dir = getattr(file_service, "input_video_dir", None) or getattr(file_service, "input_image_dir", None) + + if not input_dir: + input_dir = Path("/tmp/telefuser/inputs") + input_dir.mkdir(parents=True, exist_ok=True) + return Path(input_dir) + + async def write_upload_stream( + self, + file: Any, + file_path: Path, + *, + max_file_size: int | None = None, + ) -> None: + """Write upload content through a .part file without reading it all into memory.""" + file_path.parent.mkdir(parents=True, exist_ok=True) + part_path = file_path.with_name(f"{file_path.name}.{uuid.uuid4().hex}.part") + bytes_written = 0 + try: + with open(part_path, "wb") as buffer: + while chunk := await file.read(1024 * 1024): + bytes_written += len(chunk) + if max_file_size is not None and bytes_written > max_file_size: + raise HTTPException( + status_code=413, + detail=f"File too large: {bytes_written} bytes, max: {max_file_size} bytes", + ) + buffer.write(chunk) + part_path.replace(file_path) + except Exception: + part_path.unlink(missing_ok=True) + raise + + def resolve_output_file(self, file_service: Any, output_path: str, *, media_type: MediaType | str) -> Path: + """Resolve an output path through FileService when available.""" + resolver = self.get_declared_method(file_service, "resolve_output_file") + if callable(resolver): + try: + return resolver(output_path) + except ValueError: + raise HTTPException(status_code=403, detail="Access to this file is not allowed") + + path = Path(output_path) + if not path.is_absolute(): + output_attr = "output_image_dir" if str(media_type) == MediaType.IMAGE.value else "output_video_dir" + output_dir = getattr(file_service, output_attr, None) + if output_dir: + path = Path(output_dir) / path + return path + + def stream_file_response(self, path: Path) -> StreamingResponse | None: + responder = self.get_declared_method(self.api, "_stream_file_response") + if callable(responder): + return responder(path, filename=path.name) + return None + + @staticmethod + def get_declared_method(obj: Any, name: str) -> Any | None: + if getattr(type(obj), name, None) is None: + return None + return getattr(obj, name, None) diff --git a/telefuser/service/core/artifact_store.py b/telefuser/service/core/artifact_store.py new file mode 100644 index 0000000..fa2fa4d --- /dev/null +++ b/telefuser/service/core/artifact_store.py @@ -0,0 +1,256 @@ +"""Local artifact storage for service inputs and outputs.""" + +from __future__ import annotations + +import shutil +from collections.abc import Mapping +from datetime import datetime +from pathlib import Path +from typing import Any + +from telefuser.service_types import MediaType +from telefuser.utils.logging import logger + + +def _to_timestamp(value: datetime | float | int | None) -> float | None: + if value is None: + return None + if isinstance(value, datetime): + return value.timestamp() + return float(value) + + +class ArtifactStore: + """Path-safe local artifact store. + + The store is intentionally local-only for now. It centralizes root checks + and task-scoped output paths without claiming remote persistence semantics. + """ + + def __init__(self, root: Path) -> None: + self.root = root.resolve() + self.legacy_input_image_dir = self.root / "inputs" / "imgs" + self.legacy_input_video_dir = self.root / "inputs" / "videos" + self.legacy_input_audio_dir = self.root / "inputs" / "audios" + self.legacy_output_dir = self.root / "outputs" + self.legacy_output_video_dir = self.legacy_output_dir / "videos" + self.legacy_output_image_dir = self.legacy_output_dir / "images" + self.tasks_dir = self.root / "tasks" + + for directory in ( + self.legacy_input_image_dir, + self.legacy_input_video_dir, + self.legacy_input_audio_dir, + self.legacy_output_video_dir, + self.legacy_output_image_dir, + self.tasks_dir, + ): + directory.mkdir(parents=True, exist_ok=True) + + @staticmethod + def _is_relative_to(path: Path, root: Path) -> bool: + try: + path.relative_to(root) + return True + except ValueError: + return False + + @staticmethod + def _media_dir_name(media_type: MediaType | str) -> str: + return "images" if media_type == MediaType.IMAGE or str(media_type) == MediaType.IMAGE.value else "videos" + + def _resolve_under(self, root: Path, relative_path: str | Path) -> Path: + path = Path(relative_path) + if path.is_absolute(): + raise ValueError("Absolute paths are not allowed") + + resolved_root = root.resolve() + resolved_path = (resolved_root / path).resolve() + if not self._is_relative_to(resolved_path, resolved_root): + raise ValueError("Path escapes the service cache directory") + return resolved_path + + def _validate_task_id(self, task_id: str) -> str: + if not task_id or any(char in task_id for char in ("/", "\\")) or task_id in {".", ".."}: + raise ValueError("Invalid task_id for artifact path") + return task_id + + def task_root(self, task_id: str) -> Path: + return self._resolve_under(self.tasks_dir, self._validate_task_id(task_id)) + + def task_input_dir(self, task_id: str, media_type: MediaType | str) -> Path: + media_dir = self._media_dir_name(media_type) + path = self.task_root(task_id) / "inputs" / media_dir + path.mkdir(parents=True, exist_ok=True) + return path + + def task_output_dir(self, task_id: str, media_type: MediaType | str) -> Path: + media_dir = self._media_dir_name(media_type) + path = self.task_root(task_id) / "outputs" / media_dir + path.mkdir(parents=True, exist_ok=True) + return path + + def task_tmp_dir(self, task_id: str) -> Path: + path = self.task_root(task_id) / "tmp" + path.mkdir(parents=True, exist_ok=True) + return path + + def output_path( + self, + output_path: str | Path, + *, + media_type: MediaType | str = MediaType.VIDEO, + task_id: str | None = None, + ) -> Path: + if task_id: + return self._resolve_under(self.task_output_dir(task_id, media_type), output_path) + + root = self.legacy_output_image_dir if self._media_dir_name(media_type) == "images" else self.legacy_output_video_dir + return self._resolve_under(root, output_path) + + def resolve_output_file(self, file_path: str | Path) -> Path: + """Resolve a downloadable output artifact under legacy or task output roots.""" + path = Path(file_path) + if path.is_absolute(): + resolved_path = path.resolve() + if self._is_allowed_output_path(resolved_path): + return resolved_path + raise ValueError("Access to this file is not allowed") + + for root in (self.legacy_output_video_dir, self.legacy_output_image_dir, self.legacy_output_dir): + candidate = self._resolve_under(root, path) + if candidate.exists(): + return candidate + + for task_root in self.tasks_dir.glob("*"): + if not task_root.is_dir(): + continue + for media_dir in ("videos", "images"): + candidate = self._resolve_under(task_root / "outputs" / media_dir, path) + if candidate.exists(): + return candidate + + return self._resolve_under(self.legacy_output_dir, path) + + def _is_allowed_output_path(self, resolved_path: Path) -> bool: + if self._is_relative_to(resolved_path, self.legacy_output_dir.resolve()): + return True + + tasks_root = self.tasks_dir.resolve() + if not self._is_relative_to(resolved_path, tasks_root): + return False + relative = resolved_path.relative_to(tasks_root) + return len(relative.parts) >= 4 and relative.parts[1] == "outputs" and relative.parts[2] in {"images", "videos"} + + def cleanup( + self, + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_end_times: Mapping[str, datetime | float | int | None], + retention_seconds: int, + tmp_retention_seconds: int, + max_total_bytes: int = 0, + now: datetime | float | int | None = None, + ) -> dict[str, Any]: + """Clean expired local artifacts without touching active task directories.""" + now_ts = _to_timestamp(now) or datetime.now().timestamp() + removed_task_ids: list[str] = [] + removed_tmp_files = 0 + errors: list[str] = [] + + removed_tmp_files += self._cleanup_tmp_files(tmp_retention_seconds=tmp_retention_seconds, now_ts=now_ts) + + if retention_seconds > 0: + for task_id, end_time in terminal_task_end_times.items(): + if task_id in active_task_ids: + continue + end_ts = _to_timestamp(end_time) + if end_ts is None or now_ts - end_ts < retention_seconds: + continue + task_dir = self.task_root(task_id) + if not task_dir.exists(): + continue + try: + shutil.rmtree(task_dir) + removed_task_ids.append(task_id) + except Exception as exc: + message = f"Failed to remove artifact directory for task {task_id}: {exc}" + logger.warning(message) + errors.append(message) + + if max_total_bytes > 0: + removed_task_ids.extend( + self._cleanup_capacity( + active_task_ids=active_task_ids, + terminal_task_end_times=terminal_task_end_times, + max_total_bytes=max_total_bytes, + ) + ) + + return { + "removed_task_ids": removed_task_ids, + "removed_tmp_files": removed_tmp_files, + "errors": errors, + } + + def _cleanup_tmp_files(self, *, tmp_retention_seconds: int, now_ts: float) -> int: + if tmp_retention_seconds <= 0: + return 0 + + removed = 0 + for part_file in self.root.rglob("*.part"): + try: + if now_ts - part_file.stat().st_mtime < tmp_retention_seconds: + continue + part_file.unlink() + removed += 1 + except FileNotFoundError: + continue + except Exception as exc: + logger.warning(f"Failed to remove temporary artifact {part_file}: {exc}") + return removed + + def _cleanup_capacity( + self, + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_end_times: Mapping[str, datetime | float | int | None], + max_total_bytes: int, + ) -> list[str]: + current_size = self._tree_size(self.root) + if current_size <= max_total_bytes: + return [] + + candidates = [ + (task_id, _to_timestamp(end_time) or 0.0, self.task_root(task_id)) + for task_id, end_time in terminal_task_end_times.items() + if task_id not in active_task_ids and self.task_root(task_id).exists() + ] + candidates.sort(key=lambda item: item[1]) + + removed: list[str] = [] + for task_id, _, task_dir in candidates: + if current_size <= max_total_bytes: + break + task_size = self._tree_size(task_dir) + try: + shutil.rmtree(task_dir) + current_size -= task_size + removed.append(task_id) + except Exception as exc: + logger.warning(f"Failed to remove artifact directory for task {task_id}: {exc}") + return removed + + @staticmethod + def _tree_size(root: Path) -> int: + if not root.exists(): + return 0 + + total = 0 + for path in root.rglob("*"): + try: + if path.is_file(): + total += path.stat().st_size + except FileNotFoundError: + continue + return total diff --git a/telefuser/service/core/config.py b/telefuser/service/core/config.py index ebbc5b3..3c9309b 100644 --- a/telefuser/service/core/config.py +++ b/telefuser/service/core/config.py @@ -95,6 +95,35 @@ class ServerConfig(BaseSettings): le=1024 * 1024 * 1024, # 1GB description="Maximum file download size", ) + artifact_retention_seconds: int = Field( + default=7 * 24 * 60 * 60, + ge=0, + description="Maximum retention time for terminal task artifacts. Zero disables TTL cleanup.", + ) + artifact_tmp_retention_seconds: int = Field( + default=60 * 60, + ge=0, + description="Maximum retention time for temporary .part files. Zero disables tmp cleanup.", + ) + artifact_cleanup_interval_seconds: int = Field( + default=60 * 60, + ge=60, + description="Suggested interval for periodic artifact cleanup.", + ) + artifact_max_total_bytes: int = Field( + default=0, + ge=0, + description="Maximum local artifact bytes. Zero disables capacity cleanup.", + ) + artifact_max_task_bytes: int = Field( + default=0, + ge=0, + description="Maximum local artifact bytes per task. Zero disables per-task capacity checks.", + ) + artifact_preserve_failed_outputs: bool = Field( + default=False, + description="Whether failed task output directories should be preserved until normal retention expiry.", + ) # Rate limiting settings enable_rate_limit: bool = Field(default=True, description="Enable rate limiting middleware") diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index ab14f27..11b502a 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -101,6 +101,11 @@ def initialize_file_service(self, cache_dir: Path | None = None) -> FileService: self.file_service = FileService( cache_dir=cache_dir, max_file_size=getattr(self.config, "max_file_size", None), + verify_ssl=getattr(self.config, "verify_ssl", True), + ssl_cert_path=getattr(self.config, "ssl_cert_path", None), + artifact_retention_seconds=self.config.artifact_retention_seconds, + artifact_tmp_retention_seconds=self.config.artifact_tmp_retention_seconds, + artifact_max_total_bytes=self.config.artifact_max_total_bytes, ) return self.file_service @@ -254,6 +259,7 @@ def get_api_app(self, enable_rate_limit: bool = True) -> FastAPI: task_manager=self.task_manager, enable_rate_limit=enable_rate_limit, enable_logging=False, + config=self.config, ) if self.file_service and self.pipeline_service: @@ -262,6 +268,7 @@ def get_api_app(self, enable_rate_limit: bool = True) -> FastAPI: self.pipeline_service, cache_service=self.cache_service, cache_adapter=self.cache_adapter, # forward adapter to api_server + file_service=self.file_service, ) if self.stream_pipeline_service: diff --git a/telefuser/service/core/file_service.py b/telefuser/service/core/file_service.py index b91ba00..b07ab8e 100644 --- a/telefuser/service/core/file_service.py +++ b/telefuser/service/core/file_service.py @@ -3,8 +3,13 @@ from __future__ import annotations import asyncio +import inspect +import os import uuid +from collections.abc import Mapping +from datetime import datetime from pathlib import Path +from typing import Any from urllib.parse import urlparse import httpx @@ -12,22 +17,40 @@ from telefuser.service_types import MediaType from telefuser.utils.logging import logger +from .artifact_store import ArtifactStore + class FileService: """Service for downloading and caching media files from URLs.""" DEFAULT_MAX_FILE_SIZE = 100 * 1024 * 1024 # 100MB - def __init__(self, cache_dir: Path, max_file_size: int | None = None) -> None: - self.cache_dir = cache_dir - self.input_image_dir = cache_dir / "inputs" / "imgs" - self.input_video_dir = cache_dir / "inputs" / "videos" - self.input_audio_dir = cache_dir / "inputs" / "audios" - self.output_dir = cache_dir / "outputs" - self.output_video_dir = self.output_dir / "videos" - self.output_image_dir = self.output_dir / "images" + def __init__( + self, + cache_dir: Path, + max_file_size: int | None = None, + *, + verify_ssl: bool = True, + ssl_cert_path: str | None = None, + artifact_retention_seconds: int = 7 * 24 * 60 * 60, + artifact_tmp_retention_seconds: int = 60 * 60, + artifact_max_total_bytes: int = 0, + ) -> None: + self.artifact_store = ArtifactStore(cache_dir) + self.cache_dir = self.artifact_store.root + self.input_image_dir = self.artifact_store.legacy_input_image_dir + self.input_video_dir = self.artifact_store.legacy_input_video_dir + self.input_audio_dir = self.artifact_store.legacy_input_audio_dir + self.output_dir = self.artifact_store.legacy_output_dir + self.output_video_dir = self.artifact_store.legacy_output_video_dir + self.output_image_dir = self.artifact_store.legacy_output_image_dir self.max_file_size = max_file_size or self.DEFAULT_MAX_FILE_SIZE + self.verify_ssl = verify_ssl + self.ssl_cert_path = ssl_cert_path + self.artifact_retention_seconds = artifact_retention_seconds + self.artifact_tmp_retention_seconds = artifact_tmp_retention_seconds + self.artifact_max_total_bytes = artifact_max_total_bytes self._http_client: httpx.AsyncClient | None = None self._client_lock = asyncio.Lock() @@ -36,31 +59,26 @@ def __init__(self, cache_dir: Path, max_file_size: int | None = None) -> None: self.retry_delay = 1.0 self.max_retry_delay = 10.0 - for directory in [ - self.input_image_dir, - self.output_dir, - self.output_video_dir, - self.output_image_dir, - self.input_audio_dir, - ]: - directory.mkdir(parents=True, exist_ok=True) - async def _get_http_client(self) -> httpx.AsyncClient: """Get or create a persistent HTTP client with connection pooling.""" - from .config import server_config - async with self._client_lock: if self._http_client is None or self._http_client.is_closed: timeout = httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=5.0) limits = httpx.Limits(max_keepalive_connections=5, max_connections=10, keepalive_expiry=30.0) - verify = server_config.ssl_cert_path if server_config.ssl_cert_path else server_config.verify_ssl + verify = self.ssl_cert_path if self.ssl_cert_path else self.verify_ssl self._http_client = httpx.AsyncClient( verify=verify, timeout=timeout, limits=limits, follow_redirects=True ) return self._http_client - async def _download_with_retry(self, url: str, max_retries: int | None = None) -> httpx.Response: - """Download with exponential backoff retry logic.""" + async def _download_file_with_retry( + self, + url: str, + destination: Path, + max_size: int, + max_retries: int | None = None, + ) -> None: + """Download a remote file with retries and bounded streaming writes.""" if max_retries is None: max_retries = self.max_retries @@ -70,25 +88,30 @@ async def _download_with_retry(self, url: str, max_retries: int | None = None) - for attempt in range(max_retries): try: client = await self._get_http_client() - response = await client.get(url) - - if response.status_code == 200: - return response - elif response.status_code >= 500: - logger.warning( - f"Server error {response.status_code} for {url}, attempt {attempt + 1}/{max_retries}" - ) - last_exception = httpx.HTTPStatusError( - f"Server returned {response.status_code}", - request=response.request, - response=response, - ) - else: - raise httpx.HTTPStatusError( - f"Client error {response.status_code}", - request=response.request, - response=response, - ) + async with client.stream("GET", url) as response: + if response.status_code == 200: + content_length = response.headers.get("content-length") + if content_length is not None and int(content_length) > max_size: + raise ValueError( + f"File too large: {content_length} bytes, max: {max_size} bytes" + ) + await self._write_response_stream(response, destination, max_size) + return + if response.status_code >= 500: + logger.warning( + f"Server error {response.status_code} for {url}, attempt {attempt + 1}/{max_retries}" + ) + last_exception = httpx.HTTPStatusError( + f"Server returned {response.status_code}", + request=response.request, + response=response, + ) + else: + raise httpx.HTTPStatusError( + f"Client error {response.status_code}", + request=response.request, + response=response, + ) except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError) as e: logger.warning(f"Connection error for {url}, attempt {attempt + 1}/{max_retries}: {str(e)}") @@ -110,6 +133,35 @@ async def _download_with_retry(self, url: str, max_retries: int | None = None) - error_msg += f": {str(last_exception)}" raise httpx.ConnectError(error_msg) + async def _write_response_stream(self, response: httpx.Response, destination: Path, max_size: int) -> None: + """Write response bytes through a .part file and atomically publish on success.""" + destination.parent.mkdir(parents=True, exist_ok=True) + part_path = destination.with_name(f"{destination.name}.{uuid.uuid4().hex}.part") + bytes_written = 0 + try: + with open(part_path, "wb") as f: + async for chunk in response.aiter_bytes(chunk_size=1024 * 1024): + bytes_written += len(chunk) + if bytes_written > max_size: + raise ValueError(f"File too large: {bytes_written} bytes, max: {max_size} bytes") + f.write(chunk) + os.replace(part_path, destination) + except Exception: + part_path.unlink(missing_ok=True) + raise + + @staticmethod + def _safe_basename(raw_name: str, fallback: str) -> str: + name = Path(raw_name).name + return name or fallback + + @staticmethod + def _is_relative_to(path: Path, root: Path) -> bool: + return ArtifactStore._is_relative_to(path, root) + + def _resolve_under(self, root: Path, relative_path: str | Path) -> Path: + return self.artifact_store._resolve_under(root, relative_path) + async def download_image(self, image_url: str, max_size: int | None = None) -> Path: """Download image with retry logic and proper error handling.""" max_size = max_size or self.max_file_size @@ -119,21 +171,9 @@ async def download_image(self, image_url: str, max_size: int | None = None) -> P if not parsed_url.scheme or not parsed_url.netloc: raise ValueError(f"Invalid URL format: {image_url}") - response = await self._download_with_retry(image_url) - - content_length = len(response.content) - if content_length > max_size: - raise ValueError(f"Image too large: {content_length} bytes, max: {max_size} bytes") - - image_name = Path(parsed_url.path).name - if not image_name: - image_name = f"{uuid.uuid4()}.jpg" - - image_path = self.input_image_dir / image_name - image_path.parent.mkdir(parents=True, exist_ok=True) - - with open(image_path, "wb") as f: - f.write(response.content) + image_name = self._safe_basename(parsed_url.path, f"{uuid.uuid4()}.jpg") + image_path = self._resolve_under(self.input_image_dir, image_name) + await self._download_file_with_retry(image_url, image_path, max_size) logger.info(f"Successfully downloaded image from {image_url} to {image_path}") return image_path @@ -162,21 +202,9 @@ async def download_audio(self, audio_url: str, max_size: int | None = None) -> P if not parsed_url.scheme or not parsed_url.netloc: raise ValueError(f"Invalid URL format: {audio_url}") - response = await self._download_with_retry(audio_url) - - content_length = len(response.content) - if content_length > max_size: - raise ValueError(f"Audio too large: {content_length} bytes, max: {max_size} bytes") - - audio_name = Path(parsed_url.path).name - if not audio_name: - audio_name = f"{uuid.uuid4()}.mp3" - - audio_path = self.input_audio_dir / audio_name - audio_path.parent.mkdir(parents=True, exist_ok=True) - - with open(audio_path, "wb") as f: - f.write(response.content) + audio_name = self._safe_basename(parsed_url.path, f"{uuid.uuid4()}.mp3") + audio_path = self._resolve_under(self.input_audio_dir, audio_name) + await self._download_file_with_retry(audio_url, audio_path, max_size) logger.info(f"Successfully downloaded audio from {audio_url} to {audio_path}") return audio_path @@ -205,21 +233,9 @@ async def download_video(self, video_url: str, max_size: int | None = None) -> P if not parsed_url.scheme or not parsed_url.netloc: raise ValueError(f"Invalid URL format: {video_url}") - response = await self._download_with_retry(video_url) - - content_length = len(response.content) - if content_length > max_size: - raise ValueError(f"Video too large: {content_length} bytes, max: {max_size} bytes") - - video_name = Path(parsed_url.path).name - if not video_name: - video_name = f"{uuid.uuid4()}.mp4" - - video_path = self.input_video_dir / video_name - video_path.parent.mkdir(parents=True, exist_ok=True) - - with open(video_path, "wb") as f: - f.write(response.content) + video_name = self._safe_basename(parsed_url.path, f"{uuid.uuid4()}.mp4") + video_path = self._resolve_under(self.input_video_dir, video_name) + await self._download_file_with_retry(video_url, video_path, max_size) logger.info(f"Successfully downloaded video from {video_url} to {video_path}") return video_path @@ -241,25 +257,106 @@ async def download_video(self, video_url: str, max_size: int | None = None) -> P def save_uploaded_file(self, file_content: bytes, filename: str) -> Path: """Save uploaded file content to disk.""" + if len(file_content) > self.max_file_size: + raise ValueError(f"File too large: {len(file_content)} bytes, max: {self.max_file_size} bytes") file_extension = Path(filename).suffix unique_filename = f"{uuid.uuid4()}{file_extension}" - file_path = self.input_image_dir / unique_filename + file_path = self._resolve_under(self.input_image_dir, unique_filename) - with open(file_path, "wb") as f: - f.write(file_content) + part_path = file_path.with_name(f"{file_path.name}.{uuid.uuid4().hex}.part") + try: + with open(part_path, "wb") as f: + f.write(file_content) + os.replace(part_path, file_path) + except Exception: + part_path.unlink(missing_ok=True) + raise + + return file_path + async def save_upload_file( + self, + file: Any, + *, + media_type: MediaType | str, + prefix: str = "upload", + fallback_filename: str = "upload.bin", + ) -> Path: + """Save an async upload stream through a bounded .part file write.""" + original_name = self._safe_basename(getattr(file, "filename", None) or fallback_filename, fallback_filename) + suffix = Path(original_name).suffix + filename = f"{prefix}_{uuid.uuid4().hex[:8]}{suffix}" + file_path = self._resolve_under(self._input_dir_for_media(media_type), filename) + await self._write_upload_stream(file, file_path, self.max_file_size) return file_path - def get_output_path(self, output_path: str, media_type: MediaType | str = MediaType.VIDEO) -> Path: + def _input_dir_for_media(self, media_type: MediaType | str) -> Path: + media_value = str(media_type) + if media_value == MediaType.IMAGE.value: + return self.input_image_dir + if media_value == MediaType.VIDEO.value: + return self.input_video_dir + if media_value == "audio": + return self.input_audio_dir + raise ValueError(f"Unsupported media type: {media_type}") + + async def _write_upload_stream(self, file: Any, destination: Path, max_size: int) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + part_path = destination.with_name(f"{destination.name}.{uuid.uuid4().hex}.part") + bytes_written = 0 + try: + with open(part_path, "wb") as buffer: + while True: + chunk = await self._read_upload_chunk(file) + if not chunk: + break + bytes_written += len(chunk) + if bytes_written > max_size: + raise ValueError(f"File too large: {bytes_written} bytes, max: {max_size} bytes") + buffer.write(chunk) + os.replace(part_path, destination) + except Exception: + part_path.unlink(missing_ok=True) + raise + + @staticmethod + async def _read_upload_chunk(file: Any) -> bytes: + reader = getattr(file, "read") + chunk = reader(1024 * 1024) + if inspect.isawaitable(chunk): + chunk = await chunk + return chunk + + def get_output_path( + self, + output_path: str, + media_type: MediaType | str = MediaType.VIDEO, + *, + task_id: str | None = None, + ) -> Path: """Get the full output path for a file.""" - path = Path(output_path) - if path.is_absolute(): - return path - - if media_type == MediaType.IMAGE: - return self.output_image_dir / output_path - else: - return self.output_video_dir / output_path + return self.artifact_store.output_path(output_path, media_type=media_type, task_id=task_id) + + def resolve_output_file(self, file_path: str | Path) -> Path: + """Resolve a downloadable output file within the configured cache root.""" + return self.artifact_store.resolve_output_file(file_path) + + def cleanup_artifacts( + self, + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_end_times: Mapping[str, datetime | float | int | None], + now: datetime | float | int | None = None, + ) -> dict[str, Any]: + """Clean expired local artifacts according to configured retention settings.""" + return self.artifact_store.cleanup( + active_task_ids=active_task_ids, + terminal_task_end_times=terminal_task_end_times, + retention_seconds=self.artifact_retention_seconds, + tmp_retention_seconds=self.artifact_tmp_retention_seconds, + max_total_bytes=self.artifact_max_total_bytes, + now=now, + ) async def cleanup(self) -> None: """Cleanup resources including HTTP client.""" diff --git a/telefuser/service/core/task_manager.py b/telefuser/service/core/task_manager.py index 8a3144d..0d8e585 100644 --- a/telefuser/service/core/task_manager.py +++ b/telefuser/service/core/task_manager.py @@ -7,25 +7,14 @@ from collections import OrderedDict from dataclasses import dataclass, field from datetime import datetime -from enum import Enum from typing import Any from telefuser.metrics import get_service_metrics from telefuser.service.core.pipeline_contract import infer_media_type_for_task +from telefuser.service_types import TaskStatus from telefuser.utils.logging import logger -class TaskStatus(Enum): - """Task status enumeration.""" - - PENDING = "pending" - PROCESSING = "processing" - STREAMING = "streaming" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - _ACTIVE_STATUSES = frozenset({TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.STREAMING}) @@ -239,6 +228,20 @@ def get_active_task_count(self) -> int: with self._lock: return sum(1 for t in self._tasks.values() if t.status in _ACTIVE_STATUSES) + def get_artifact_cleanup_snapshot(self) -> dict[str, Any]: + """Return active and terminal task state needed by artifact cleanup.""" + with self._lock: + active_task_ids = {task_id for task_id, task in self._tasks.items() if task.status in _ACTIVE_STATUSES} + terminal_task_end_times = { + task_id: task.end_time or task.start_time + for task_id, task in self._tasks.items() + if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED) + } + return { + "active_task_ids": active_task_ids, + "terminal_task_end_times": terminal_task_end_times, + } + def get_pending_task_count(self) -> int: """Get count of pending tasks.""" with self._lock: diff --git a/telefuser/service/core/task_service.py b/telefuser/service/core/task_service.py index 5279e98..6b95670 100644 --- a/telefuser/service/core/task_service.py +++ b/telefuser/service/core/task_service.py @@ -4,7 +4,7 @@ import threading from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any from telefuser.service_types import MediaType, PipelineRunStatus, TaskStatus, TaskType from telefuser.utils.logging import logger @@ -13,7 +13,9 @@ from ..media.media_base import AudioHandler, ImageHandler, VideoHandler from .file_service import FileService from .pipeline_contract import infer_media_type_for_task -from .pipeline_service import PipelineService + +if TYPE_CHECKING: + from .pipeline_service import PipelineService # Media handlers _image_handler = ImageHandler() @@ -110,7 +112,14 @@ async def update_audio_path(audio_name: str, message: TaskRequest, task_data: di # Determine media type and set appropriate output path media_type = infer_media_type_for_task(message.task) - actual_save_path = self.file_service.get_output_path(message.output_path, media_type=media_type) + try: + actual_save_path = self.file_service.get_output_path( + message.output_path, + media_type=media_type, + task_id=message.task_id, + ) + except TypeError: + actual_save_path = self.file_service.get_output_path(message.output_path, media_type=media_type) task_data["output_path"] = str(actual_save_path) # Best-effort: every step degrades silently to "no cache" on failure. diff --git a/telefuser/service/main.py b/telefuser/service/main.py index fed85c5..4f5aeff 100644 --- a/telefuser/service/main.py +++ b/telefuser/service/main.py @@ -24,12 +24,13 @@ def _run( print(TELEFUSER_LOGO) logger.info(f"Starting TeleFuser {label}...") - if not server_config.validate(): + config = container.config + if not config.validate(): raise RuntimeError("Invalid server configuration") app = container.get_api_app(enable_rate_limit=enable_rate_limit) - logger.info(f"Starting {label} on {server_config.host}:{server_config.port}") - uvicorn.run(app, host=server_config.host, port=server_config.port, log_level="warning") + logger.info(f"Starting {label} on {config.host}:{config.port}") + uvicorn.run(app, host=config.host, port=config.port, log_level="warning") except KeyboardInterrupt: logger.info(f"{label.capitalize()} interrupted by user") @@ -54,6 +55,8 @@ def run_server( num_replicas: int = 1, enable_latent_cache: bool | None = None, cache_mode: str | None = None, + security_level: str | None = None, + skip_validation: bool = False, ) -> None: """Run the TeleFuser server with dependency injection container.""" server_config.host = host @@ -61,6 +64,10 @@ def run_server( server_config.num_replicas = num_replicas server_config.enable_latent_cache = enable_latent_cache server_config.cache_mode = cache_mode + if security_level is not None: + from .security.security_validator import SecurityLevel + + server_config.security_level = SecurityLevel[security_level.upper()] container = ServiceContainer.create( config=server_config, @@ -72,6 +79,7 @@ def run_server( parallelism=parallelism, task=task, cache_dir=Path(cache_dir) if cache_dir else None, + skip_validation=skip_validation, ): raise RuntimeError("Failed to initialize services") @@ -85,6 +93,7 @@ def run_stream_server( host: str, enable_rate_limit: bool = True, skip_validation: bool = False, + security_level: str | None = None, ) -> None: """Run the TeleFuser stream server. @@ -93,6 +102,10 @@ def run_stream_server( """ server_config.host = host server_config.port = port + if security_level is not None: + from .security.security_validator import SecurityLevel + + server_config.security_level = SecurityLevel[security_level.upper()] container = ServiceContainer.create(config=server_config) diff --git a/telefuser/service_types.py b/telefuser/service_types.py index cfb4a5f..f2c98cc 100644 --- a/telefuser/service_types.py +++ b/telefuser/service_types.py @@ -55,6 +55,7 @@ class TaskStatus(_StringEnum): PENDING = "pending" PROCESSING = "processing" + STREAMING = "streaming" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" diff --git a/tests/server/test_middleware.py b/tests/server/test_middleware.py index f53550a..3540a74 100644 --- a/tests/server/test_middleware.py +++ b/tests/server/test_middleware.py @@ -5,12 +5,15 @@ """ # Set environment before imports +import asyncio +import json import os from unittest.mock import MagicMock, patch import pytest from fastapi import FastAPI -from starlette.testclient import TestClient +from starlette.requests import Request +from starlette.responses import JSONResponse os.environ["TELEFUSER_SECURITY_LEVEL"] = "NONE" @@ -22,6 +25,24 @@ from telefuser.service.core.config import ServerConfig, server_config +def _request(path: str = "/test", *, headers: list[tuple[bytes, bytes]] | None = None) -> Request: + return Request( + { + "type": "http", + "method": "GET", + "path": path, + "scheme": "http", + "server": ("testserver", 80), + "headers": headers or [], + "client": ("127.0.0.1", 12345), + } + ) + + +async def _ok_call_next(request: Request) -> JSONResponse: + return JSONResponse({"status": "ok"}) + + class TestRateLimitMiddleware: """Test rate limiting middleware.""" @@ -43,17 +64,11 @@ def test_endpoint(): # Clear any state middleware._clients = {} - @app.middleware("http") - async def rate_limit_middleware(request, call_next): - return await middleware.dispatch(request, call_next) - - client = TestClient(app) - # Should allow 5 requests for i in range(5): - response = client.get("/test") + response = asyncio.run(middleware.dispatch(_request("/test"), _ok_call_next)) assert response.status_code == 200 - assert response.json() == {"status": "ok"} + assert json.loads(response.body) == {"status": "ok"} def test_rate_limit_only_applies_to_limited_paths(self): """Test that paths outside the whitelist are not rate limited.""" @@ -76,15 +91,9 @@ def test_endpoint(): ) middleware._clients = {} - @app.middleware("http") - async def rate_limit_middleware(request, call_next): - return await middleware.dispatch(request, call_next) - - client = TestClient(app) - # /health is not in limited_paths, so it should never be rate limited. for i in range(20): - response = client.get("/health") + response = asyncio.run(middleware.dispatch(_request("/health"), _ok_call_next)) assert response.status_code == 200, f"Health check failed on request {i + 1}" def test_rate_limit_headers_present(self): @@ -103,13 +112,7 @@ def test_endpoint(): ) middleware._clients = {} - @app.middleware("http") - async def rate_limit_middleware(request, call_next): - return await middleware.dispatch(request, call_next) - - client = TestClient(app) - - response = client.get("/test") + response = asyncio.run(middleware.dispatch(_request("/test"), _ok_call_next)) assert response.status_code == 200 # Check headers @@ -133,17 +136,7 @@ def test_endpoint(): limited_paths=["/test"], ) - # Test client ID extraction - from starlette.requests import Request - - scope = { - "type": "http", - "method": "GET", - "path": "/test", - "headers": [(b"x-forwarded-for", b"192.168.1.1")], - "client": ("127.0.0.1", 12345), - } - request = Request(scope) + request = _request(headers=[(b"x-forwarded-for", b"192.168.1.1")]) client_id = middleware._get_client_id(request) assert client_id == "192.168.1.1" @@ -159,17 +152,17 @@ def test_rate_limit_get_client_id_from_client(self): limited_paths=["/test"], ) - # Test client ID extraction without X-Forwarded-For - from starlette.requests import Request - - scope = { - "type": "http", - "method": "GET", - "path": "/test", - "headers": [], - "client": ("10.0.0.1", 12345), - } - request = Request(scope) + request = Request( + { + "type": "http", + "method": "GET", + "path": "/test", + "scheme": "http", + "server": ("testserver", 80), + "headers": [], + "client": ("10.0.0.1", 12345), + } + ) client_id = middleware._get_client_id(request) assert client_id == "10.0.0.1" @@ -250,10 +243,8 @@ def test_logging_middleware_logs_requests(self, mock_logger): def test_endpoint(): return {"status": "ok"} - app.add_middleware(LoggingMiddleware) - - client = TestClient(app) - response = client.get("/test") + middleware = LoggingMiddleware(app) + response = asyncio.run(middleware.dispatch(_request("/test"), _ok_call_next)) assert response.status_code == 200 # Verify logger was called @@ -274,11 +265,7 @@ def test_endpoint(): setup_middleware(app, enable_rate_limit=False, enable_logging=True) - client = TestClient(app) - - # Requests should succeed - response = client.get("/test") - assert response.status_code == 200 + assert any(item.cls is LoggingMiddleware for item in app.user_middleware) def test_setup_middleware_no_middleware(self): """Test setup with no middleware.""" @@ -290,11 +277,9 @@ def test_endpoint(): setup_middleware(app, enable_rate_limit=False, enable_logging=False) - client = TestClient(app) - - # Requests should succeed - response = client.get("/test") - assert response.status_code == 200 + middleware_classes = {item.cls for item in app.user_middleware} + assert LoggingMiddleware not in middleware_classes + assert RateLimitMiddleware not in middleware_classes class TestConfigIntegration: diff --git a/tests/unit/openai/_asgi_test_client.py b/tests/unit/openai/_asgi_test_client.py new file mode 100644 index 0000000..5ae7ea6 --- /dev/null +++ b/tests/unit/openai/_asgi_test_client.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import asyncio +from typing import Any + +import httpx + + +class ASGITestClient: + """Small sync wrapper around httpx ASGITransport for route unit tests.""" + + __test__ = False + + def __init__(self, app) -> None: + self.app = app + + def __enter__(self) -> ASGITestClient: + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + return None + + def request(self, method: str, url: str, **kwargs: Any) -> httpx.Response: + async def _request() -> httpx.Response: + transport = httpx.ASGITransport(app=self.app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client: + return await client.request(method, url, **kwargs) + + return asyncio.run(_request()) + + def get(self, url: str, **kwargs: Any) -> httpx.Response: + return self.request("GET", url, **kwargs) + + def post(self, url: str, **kwargs: Any) -> httpx.Response: + return self.request("POST", url, **kwargs) + + def delete(self, url: str, **kwargs: Any) -> httpx.Response: + return self.request("DELETE", url, **kwargs) diff --git a/tests/unit/openai/test_image_routes.py b/tests/unit/openai/test_image_routes.py index 3c54caf..7f4839d 100644 --- a/tests/unit/openai/test_image_routes.py +++ b/tests/unit/openai/test_image_routes.py @@ -5,15 +5,19 @@ Avoids testing implementation details like internal helper methods. """ +import asyncio from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI -from fastapi.testclient import TestClient +from fastapi import HTTPException -from telefuser.service.api.openai.image_routes import create_router +from telefuser.service.api.openai.image_routes import ImageRoutes, create_router +from telefuser.service.api.task_application_service import TaskApplicationService from telefuser.service.core.task_manager import TaskStatus +from ._asgi_test_client import ASGITestClient as TestClient + @pytest.fixture def client(tmp_path): @@ -56,10 +60,12 @@ def client(tmp_path): server.file_service.output_image_dir = tmp_path server.file_service.input_image_dir = tmp_path / "input" server.file_service.input_image_dir.mkdir(exist_ok=True) + server.task_app_service = TaskApplicationService(server) # Create app with routes app = FastAPI() app.state.task_manager = task_manager + app.state.server = server app.include_router(create_router(server)) with TestClient(app) as client: @@ -128,10 +134,11 @@ class TestImageContent: def test_download_image_success(self, client): """Download existing image.""" - response = client.get("/v1/images/task_123/content") + routes = ImageRoutes(client.app.state.server) + response = asyncio.run(routes.get_image_content("task_123")) - assert response.status_code == 200 - assert response.content == b"fake_image" + assert response.path + assert response.path.endswith("output.png") def test_download_image_not_found(self, tmp_path): """Download non-existent image returns 404.""" @@ -143,12 +150,11 @@ def test_download_image_not_found(self, tmp_path): server.task_manager = task_manager server.file_service = MagicMock() server.file_service.output_image_dir = tmp_path + server.task_app_service = TaskApplicationService(server) from fastapi import FastAPI - app = FastAPI() - app.include_router(create_router(server)) - - with TestClient(app) as client: - response = client.get("/v1/images/nonexistent/content") - assert response.status_code == 404 + routes = ImageRoutes(server) + with pytest.raises(HTTPException) as exc_info: + asyncio.run(routes.get_image_content("nonexistent")) + assert exc_info.value.status_code == 404 diff --git a/tests/unit/openai/test_video_routes.py b/tests/unit/openai/test_video_routes.py index f8697ca..8e82b68 100644 --- a/tests/unit/openai/test_video_routes.py +++ b/tests/unit/openai/test_video_routes.py @@ -4,15 +4,19 @@ Uses real FastAPI TestClient with mocked services. """ +import asyncio from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI -from fastapi.testclient import TestClient +from fastapi import HTTPException -from telefuser.service.api.openai.video_routes import create_router +from telefuser.service.api.openai.video_routes import VideoRoutes, create_router +from telefuser.service.api.task_application_service import TaskApplicationService from telefuser.service.core.task_manager import TaskStatus +from ._asgi_test_client import ASGITestClient as TestClient + @pytest.fixture def client(tmp_path): @@ -65,9 +69,11 @@ def client(tmp_path): server.file_service.output_video_dir = tmp_path server.file_service.input_video_dir = tmp_path / "input" server.file_service.input_video_dir.mkdir(exist_ok=True) + server.task_app_service = TaskApplicationService(server) app = FastAPI() app.state.task_manager = task_manager + app.state.server = server app.include_router(create_router(server)) with TestClient(app) as client: @@ -139,13 +145,24 @@ def test_get_video_success(self, client): data = response.json() assert data["id"] == "vid_123" + def test_delete_completed_video_reports_completed_status(self, client): + """Delete on a terminal task reports the real terminal status.""" + response = client.delete("/v1/videos/vid_123") + + assert response.status_code == 200 + data = response.json() + assert data["id"] == "vid_123" + assert data["status"] == "completed" + client.app.state.task_manager.cancel_task.assert_not_called() + class TestVideoContent: """Tests for GET /v1/videos/{id}/content.""" def test_download_video_success(self, client): """Download completed video.""" - response = client.get("/v1/videos/vid_123/content") + routes = VideoRoutes(client.app.state.server) + response = asyncio.run(routes.get_video_content("vid_123")) - assert response.status_code == 200 - assert response.content == b"fake_video" + assert response.path + assert response.path.endswith("video.mp4") diff --git a/tests/unit/service/test_artifact_store.py b/tests/unit/service/test_artifact_store.py new file mode 100644 index 0000000..a5f24b7 --- /dev/null +++ b/tests/unit/service/test_artifact_store.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from telefuser.service.api.api_server import ApiServer +from telefuser.service.core.config import ServerConfig +from telefuser.service.core.artifact_store import ArtifactStore +from telefuser.service.core.file_service import FileService +from telefuser.service.core.task_manager import TaskManager +from telefuser.service_types import MediaType + + +class _AsyncUpload: + def __init__(self, filename: str, chunks: list[bytes]) -> None: + self.filename = filename + self._chunks = list(chunks) + + async def read(self, size: int) -> bytes: + if not self._chunks: + return b"" + return self._chunks.pop(0) + + +def test_artifact_store_creates_task_scoped_output_paths(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + + output = store.output_path("result.png", media_type=MediaType.IMAGE, task_id="task-123") + + assert output == tmp_path / "tasks" / "task-123" / "outputs" / "images" / "result.png" + assert output.parent.exists() + + +def test_artifact_store_rejects_unsafe_task_output_paths(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + + with pytest.raises(ValueError, match="Absolute paths"): + store.output_path("/tmp/escape.png", media_type=MediaType.IMAGE, task_id="task-123") + + with pytest.raises(ValueError, match="escapes"): + store.output_path("../escape.png", media_type=MediaType.IMAGE, task_id="task-123") + + with pytest.raises(ValueError, match="Invalid task_id"): + store.output_path("result.png", media_type=MediaType.IMAGE, task_id="../bad") + + +def test_artifact_store_resolves_task_scoped_outputs_for_download(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + output = store.output_path("result.mp4", media_type=MediaType.VIDEO, task_id="task-123") + output.write_bytes(b"video") + + assert store.resolve_output_file("result.mp4") == output + assert store.resolve_output_file(output) == output + + +def test_file_service_uses_task_scoped_outputs_when_task_id_is_provided(tmp_path: Path) -> None: + files = FileService(tmp_path) + + output = files.get_output_path("result.png", media_type=MediaType.IMAGE, task_id="task-123") + + assert output == tmp_path / "tasks" / "task-123" / "outputs" / "images" / "result.png" + + +def test_file_service_saves_upload_stream_to_media_input_dir(tmp_path: Path) -> None: + files = FileService(tmp_path) + upload = _AsyncUpload("cat.png", [b"ca", b"t"]) + + output = asyncio.run(files.save_upload_file(upload, media_type=MediaType.IMAGE, prefix="input")) + + assert output.parent == files.input_image_dir + assert output.name.startswith("input_") + assert output.suffix == ".png" + assert output.read_bytes() == b"cat" + assert not list(output.parent.glob("*.part")) + + +def test_file_service_upload_stream_enforces_max_file_size(tmp_path: Path) -> None: + files = FileService(tmp_path, max_file_size=2) + upload = _AsyncUpload("cat.png", [b"cat"]) + + with pytest.raises(ValueError, match="File too large"): + asyncio.run(files.save_upload_file(upload, media_type=MediaType.IMAGE, prefix="input")) + + assert not list(files.input_image_dir.glob("*.part")) + + +def test_artifact_store_cleanup_removes_only_expired_terminal_tasks(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + old_output = store.output_path("old.mp4", media_type=MediaType.VIDEO, task_id="old-task") + active_output = store.output_path("active.mp4", media_type=MediaType.VIDEO, task_id="active-task") + fresh_output = store.output_path("fresh.mp4", media_type=MediaType.VIDEO, task_id="fresh-task") + old_output.write_bytes(b"old") + active_output.write_bytes(b"active") + fresh_output.write_bytes(b"fresh") + now = datetime.now() + + result = store.cleanup( + active_task_ids={"active-task"}, + terminal_task_end_times={ + "old-task": now - timedelta(seconds=120), + "active-task": now - timedelta(seconds=120), + "fresh-task": now, + }, + retention_seconds=60, + tmp_retention_seconds=0, + now=now, + ) + + assert result["removed_task_ids"] == ["old-task"] + assert not old_output.exists() + assert active_output.exists() + assert fresh_output.exists() + + +def test_artifact_store_cleanup_removes_expired_part_files(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + part_file = store.task_tmp_dir("task-123") / "upload.part" + part_file.write_bytes(b"partial") + old_ts = (datetime.now() - timedelta(seconds=120)).timestamp() + part_file.touch() + import os + + os.utime(part_file, (old_ts, old_ts)) + + result = store.cleanup( + active_task_ids={"task-123"}, + terminal_task_end_times={}, + retention_seconds=0, + tmp_retention_seconds=60, + ) + + assert result["removed_tmp_files"] == 1 + assert not part_file.exists() + + +def test_api_server_runs_artifact_cleanup_from_task_snapshot(tmp_path: Path) -> None: + task_manager = TaskManager() + config = ServerConfig(artifact_retention_seconds=1, artifact_tmp_retention_seconds=0) + server = ApiServer(task_manager=task_manager, config=config, enable_openai_api=False) + server.initialize_services(tmp_path, Mock()) + + task_id = task_manager.create_task(SimpleNamespace(output_path="old.mp4")) + output = server.file_service.get_output_path("old.mp4", media_type=MediaType.VIDEO, task_id=task_id) + output.write_bytes(b"old") + task_manager.complete_task(task_id) + task = task_manager.get_task(task_id) + + result = server.run_artifact_cleanup(now=task.end_time + timedelta(seconds=2)) + + assert result["removed_task_ids"] == [task_id] + assert not output.exists() + + +def test_api_server_artifact_cleanup_loop_starts_and_stops(tmp_path: Path) -> None: + async def _exercise() -> None: + server = ApiServer(task_manager=TaskManager(), config=ServerConfig(), enable_openai_api=False) + server.file_service = FileService(tmp_path) + + await server.ensure_artifact_cleanup_running() + + assert server._artifact_cleanup_task is not None + assert not server._artifact_cleanup_task.done() + + await server.cleanup() + + assert server._artifact_cleanup_task is None + + asyncio.run(_exercise()) diff --git a/tests/unit/service/test_config.py b/tests/unit/service/test_config.py index 8b3a917..da44c37 100644 --- a/tests/unit/service/test_config.py +++ b/tests/unit/service/test_config.py @@ -21,3 +21,13 @@ def test_server_config_ignores_unknown_fields_for_forward_compatibility() -> Non assert config.max_queue_size == 12 assert not hasattr(config, "future_option") + + +def test_artifact_retention_defaults_are_configured() -> None: + config = ServerConfig() + + assert config.artifact_retention_seconds == 7 * 24 * 60 * 60 + assert config.artifact_tmp_retention_seconds == 60 * 60 + assert config.artifact_cleanup_interval_seconds == 60 * 60 + assert config.artifact_max_total_bytes == 0 + assert config.artifact_preserve_failed_outputs is False diff --git a/tests/unit/service/test_service_refactor_phase0.py b/tests/unit/service/test_service_refactor_phase0.py new file mode 100644 index 0000000..18f9f98 --- /dev/null +++ b/tests/unit/service/test_service_refactor_phase0.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from unittest.mock import Mock + +import pytest +from click.testing import CliRunner + +from telefuser.entrypoints.cli.main import main +from telefuser.service.api.api_server import ApiServer +from telefuser.service.api.routers.service import ServiceRoutes +from telefuser.service.core.config import ServerConfig +from telefuser.service.core.file_service import FileService +from telefuser.service.core.task_manager import TaskManager, TaskStatus as CoreTaskStatus +from telefuser.service.security.security_validator import SecurityLevel +from telefuser.service_types import MediaType, TaskStatus + + +def test_task_status_uses_shared_service_enum() -> None: + assert CoreTaskStatus is TaskStatus + assert TaskStatus.STREAMING.value == "streaming" + + +def test_file_service_rejects_unsafe_output_paths(tmp_path: Path) -> None: + files = FileService(tmp_path) + + with pytest.raises(ValueError, match="Absolute paths"): + files.get_output_path("/tmp/escape.mp4", media_type=MediaType.VIDEO) + + with pytest.raises(ValueError, match="escapes"): + files.get_output_path("../escape.mp4", media_type=MediaType.VIDEO) + + +def test_file_service_allows_image_and_video_download_roots(tmp_path: Path) -> None: + files = FileService(tmp_path) + image = files.output_image_dir / "result.png" + video = files.output_video_dir / "result.mp4" + image.write_bytes(b"image") + video.write_bytes(b"video") + + assert files.resolve_output_file("result.png") == image + assert files.resolve_output_file("result.mp4") == video + + +def test_api_server_initializes_file_service_with_configured_max_file_size(tmp_path: Path) -> None: + config = ServerConfig(max_file_size=2 * 1024 * 1024) + server = ApiServer(task_manager=TaskManager(), config=config, enable_openai_api=False) + inference_service = Mock() + + server.initialize_services(tmp_path, inference_service) + + assert server.file_service is not None + assert server.file_service.max_file_size == config.max_file_size + + +def test_health_and_readiness_are_separate() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + routes = ServiceRoutes(server) + + health = asyncio.run(routes.health_check()) + ready = asyncio.run(routes.readiness_check()) + + assert health["status"] == "healthy" + assert health["ready"] is False + assert ready.status_code == 503 + assert ready.body + + +def test_readiness_passes_when_pipeline_is_running() -> None: + class RunningPipeline: + is_running = True + + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.inference_service = RunningPipeline() + routes = ServiceRoutes(server) + + health = asyncio.run(routes.health_check()) + ready = asyncio.run(routes.readiness_check()) + + assert ready.status_code == 200 + assert health["ready"] is True + assert health["pipeline_ready"] is True + + +def test_cli_serve_forwards_security_and_skip_validation(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + def fake_run_server(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr("telefuser.service.main.run_server", fake_run_server) + + result = CliRunner().invoke( + main, + [ + "serve", + "pipeline.py", + "--skip-validation", + "--security-level", + "basic", + ], + ) + + assert result.exit_code == 0 + assert captured["security_level"] == "basic" + assert captured["skip_validation"] is True + + +def test_cli_stream_serve_does_not_force_skip_validation(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + def fake_assert_safe(self, pipe_path): + captured["validated"] = pipe_path + + def fake_run_stream_server(pipe_path, port, host, **kwargs): + captured.update({"pipe_path": pipe_path, "port": port, "host": host, **kwargs}) + + monkeypatch.setattr( + "telefuser.service.security.security_validator.PipelineSecurityValidator.assert_safe", + fake_assert_safe, + ) + monkeypatch.setattr("telefuser.service.main.run_stream_server", fake_run_stream_server) + + result = CliRunner().invoke( + main, + [ + "stream-serve", + "stream_pipeline.py", + "--security-level", + "none", + ], + ) + + assert result.exit_code == 0 + assert captured["validated"] == "stream_pipeline.py" + assert captured["security_level"] == "none" + assert captured["skip_validation"] is False + + +def test_run_server_security_level_is_applied(monkeypatch: pytest.MonkeyPatch) -> None: + from telefuser.service import main as service_main + + class FakeContainer: + def __init__(self, config): + self.config = config + + def initialize_all(self, **kwargs): + return True + + def get_api_app(self, enable_rate_limit=True): + return Mock() + + async def cleanup(self): + return None + + captured = {} + + def fake_create(config=None, cache_dir=None): + captured["config"] = config + return FakeContainer(config) + + monkeypatch.setattr(service_main.ServiceContainer, "create", fake_create) + monkeypatch.setattr(service_main.uvicorn, "run", lambda *args, **kwargs: None) + + service_main.run_server( + pipe_path="pipeline.py", + task="t2v", + port=9999, + host="127.0.0.1", + security_level="basic", + skip_validation=True, + ) + + assert captured["config"].security_level is SecurityLevel.BASIC diff --git a/tests/unit/service/test_task_application_service.py b/tests/unit/service/test_task_application_service.py new file mode 100644 index 0000000..dc42207 --- /dev/null +++ b/tests/unit/service/test_task_application_service.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +from fastapi import HTTPException + +from telefuser.service.api.api_server import ApiServer +from telefuser.service.api.schema import TaskRequest +from telefuser.service.core.file_service import FileService +from telefuser.service.core.task_manager import TaskManager + + +def test_task_application_service_applies_defaults_and_creates_task(tmp_path: Path) -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + server.file_service = FileService(tmp_path) + server.validate_task_supported = lambda task: None + server.get_task_contract = lambda task: { + "parameters": { + "resolution": {"default": "480p"}, + "seed": {"default": 7}, + } + } + message = TaskRequest(task="t2v") + + response = asyncio.run( + server.task_app_service.submit( + message, + explicit_fields={"task"}, + ensure_processing=False, + ) + ) + + assert response.task_id == message.task_id + assert response.task_status.value == "pending" + assert message.resolution == "480p" + assert message.seed == 7 + assert task_manager.get_task(message.task_id).message is message + + +def test_task_application_service_rejects_unsafe_output_path(tmp_path: Path) -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.file_service = FileService(tmp_path) + server.validate_task_supported = lambda task: None + server.get_task_contract = lambda task: None + message = TaskRequest(task="t2v", output_path="../escape.mp4") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + server.task_app_service.submit( + message, + explicit_fields={"task", "output_path"}, + ensure_processing=False, + ) + ) + + assert exc_info.value.status_code == 400 + + +def test_task_application_service_runs_route_specific_input_validation() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.validate_task_supported = lambda task: None + server.get_task_contract = lambda task: None + message = TaskRequest(task="i2v") + + def validate_inputs(message: TaskRequest, contract: dict | None) -> None: + raise HTTPException(status_code=400, detail="missing input") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + server.task_app_service.submit( + message, + explicit_fields={"task"}, + validate_inputs=validate_inputs, + ensure_processing=False, + ) + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "missing input" + assert server.task_manager.get_task(message.task_id) is None + + +def test_task_application_service_waits_for_completed_task() -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + task_manager.complete_task(task_id, output_path="out.mp4") + + status = asyncio.run( + server.task_app_service.wait_for_completion( + task_id, + timeout=1, + poll_interval=0, + ) + ) + + assert status["status"] == "completed" + assert status["output_path"] == "out.mp4" + + +def test_task_application_service_wait_maps_failed_task() -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + task_manager.fail_task(task_id, "pipeline failed") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + server.task_app_service.wait_for_completion( + task_id, + timeout=1, + poll_interval=0, + ) + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Generation failed: pipeline failed" + + +def test_task_application_service_wait_times_out_pending_task() -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run( + server.task_app_service.wait_for_completion( + task_id, + timeout=0, + poll_interval=0, + ) + ) + + assert exc_info.value.status_code == 504 + + +def test_task_application_service_returns_output_response(tmp_path: Path) -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + server.file_service = FileService(tmp_path) + message = TaskRequest(task="t2i") + task_id = task_manager.create_task(message) + output_path = server.file_service.get_output_path("result.png", media_type="image", task_id=task_id) + output_path.write_bytes(b"image") + task_manager.complete_task(task_id, output_path=str(output_path)) + + response = server.task_app_service.get_output_response(task_id, media_type="image") + + assert response.headers["content-length"] == "5" + + +def test_task_application_service_rejects_unready_required_output(tmp_path: Path) -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + server.file_service = FileService(tmp_path) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + + with pytest.raises(HTTPException) as exc_info: + server.task_app_service.get_output_response( + task_id, + media_type="video", + require_completed=True, + ) + + assert exc_info.value.status_code == 400 + assert "is not ready" in exc_info.value.detail + + +def test_task_application_service_cancel_accepts_active_task() -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + + result = server.task_app_service.cancel_task(task_id) + + assert result == {"result": "accepted", "task_status": "cancelled"} + assert task_manager.get_task_status(task_id)["status"] == "cancelled" + + +def test_task_application_service_cancel_reports_terminal_task() -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + message = TaskRequest(task="t2v") + task_id = task_manager.create_task(message) + task_manager.complete_task(task_id) + + result = server.task_app_service.cancel_task(task_id) + + assert result == {"result": "already_terminal", "task_status": "completed"} + + +def test_task_application_service_cancel_reports_missing_task() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + + result = server.task_app_service.cancel_task("missing") + + assert result == {"result": "not_found", "task_status": None} diff --git a/tests/unit/service/test_task_routes.py b/tests/unit/service/test_task_routes.py index 50bf18f..76e6a98 100644 --- a/tests/unit/service/test_task_routes.py +++ b/tests/unit/service/test_task_routes.py @@ -1,12 +1,28 @@ from __future__ import annotations +import asyncio from pathlib import Path -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest from fastapi import HTTPException, UploadFile from telefuser.service.api.routers.tasks import TaskRoutes +from telefuser.service.api.schema import TaskResponse +from telefuser.service.api.task_application_service import TaskApplicationService +from telefuser.service.core.file_service import FileService + + +class _AsyncUpload: + def __init__(self, filename: str, content_type: str, chunks: list[bytes]) -> None: + self.filename = filename + self.content_type = content_type + self._chunks = list(chunks) + + async def read(self, size: int) -> bytes: + if not self._chunks: + return b"" + return self._chunks.pop(0) def _make_routes(supported_tasks: tuple[str, ...] = ()) -> TaskRoutes: @@ -25,6 +41,16 @@ def _make_routes(supported_tasks: tuple[str, ...] = ()) -> TaskRoutes: return TaskRoutes(api_server) +def _make_form_routes(tmp_path: Path, supported_tasks: tuple[str, ...]) -> TaskRoutes: + routes = _make_routes(supported_tasks) + routes.api.file_service = FileService(tmp_path) + routes.api.task_app_service = TaskApplicationService(routes.api) + routes.api.task_app_service.submit = AsyncMock( + return_value=TaskResponse(task_id="task-123", task_status="pending", output_path="output.mp4") + ) + return routes + + def test_resolve_form_task_prefers_supported_pipeline_task_order() -> None: routes = _make_routes(("t2i", "i2i", "i2v")) @@ -114,3 +140,29 @@ def test_is_video_upload_detects_video_extension_without_content_type() -> None: assert routes._is_video_upload(upload) is True finally: upload.file.close() + + +def test_create_task_form_uses_file_service_upload_for_image(tmp_path: Path) -> None: + routes = _make_form_routes(tmp_path, ("i2i",)) + upload = _AsyncUpload("cat.png", "image/png", [b"cat"]) + + asyncio.run(routes.create_task_form(first_image_file=upload, prompt="make it blue")) + + task_request = routes.api.task_app_service.submit.call_args.args[0] + input_path = Path(task_request.first_image_path) + assert task_request.task == "i2i" + assert input_path.parent == routes.api.file_service.input_image_dir + assert input_path.read_bytes() == b"cat" + + +def test_create_task_form_uses_file_service_upload_for_video(tmp_path: Path) -> None: + routes = _make_form_routes(tmp_path, ("vc",)) + upload = _AsyncUpload("clip.mp4", "video/mp4", [b"video"]) + + asyncio.run(routes.create_task_form(first_image_file=upload, prompt="continue")) + + task_request = routes.api.task_app_service.submit.call_args.args[0] + input_path = Path(task_request.ref_video_path) + assert task_request.task == "vc" + assert input_path.parent == routes.api.file_service.input_video_dir + assert input_path.read_bytes() == b"video" diff --git a/tests/unit/service/test_task_runtime.py b/tests/unit/service/test_task_runtime.py index 9da836f..b4d9306 100644 --- a/tests/unit/service/test_task_runtime.py +++ b/tests/unit/service/test_task_runtime.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace from telefuser.service.api.schema import TaskRequest, TaskResponse from telefuser.service.core.task_manager import TaskManager, TaskStatus @@ -25,6 +26,20 @@ def test_complete_task_does_not_override_cancelled_status() -> None: assert status["output_path"] == message.output_path +def test_artifact_cleanup_snapshot_splits_active_and_terminal_tasks() -> None: + task_manager = TaskManager() + active_id = task_manager.create_task(SimpleNamespace(output_path="active.mp4")) + terminal_id = task_manager.create_task(SimpleNamespace(output_path="done.mp4")) + task_manager.start_task(active_id) + task_manager.complete_task(terminal_id) + + snapshot = task_manager.get_artifact_cleanup_snapshot() + + assert snapshot["active_task_ids"] == {active_id} + assert terminal_id in snapshot["terminal_task_end_times"] + assert active_id not in snapshot["terminal_task_end_times"] + + class _ControlledMediaService: def __init__(self) -> None: self.started = asyncio.Event() From 4f7422480ef10565eed4e4020d51b5bb2c1dda23 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:34:04 +0800 Subject: [PATCH 2/6] feat(service): configure local artifact lifecycle Add artifact local root and backend boundary configuration, reject unimplemented remote backends, and enforce terminal-task cleanup by TTL, total cache size, and per-task size. Document the local backend as the only implemented artifact backend while keeping S3-compatible storage deferred. --- telefuser/service/api/api_server.py | 1 + telefuser/service/core/artifact_store.py | 34 ++++++++++++++++++ telefuser/service/core/config.py | 13 +++++++ telefuser/service/core/container.py | 9 ++++- telefuser/service/core/file_service.py | 3 ++ tests/unit/service/test_artifact_store.py | 43 +++++++++++++++++++++++ tests/unit/service/test_config.py | 41 +++++++++++++++++++++ 7 files changed, 143 insertions(+), 1 deletion(-) diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index 0f196ff..9169f26 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -342,6 +342,7 @@ def initialize_services( artifact_retention_seconds=self.server_config.artifact_retention_seconds, artifact_tmp_retention_seconds=self.server_config.artifact_tmp_retention_seconds, artifact_max_total_bytes=self.server_config.artifact_max_total_bytes, + artifact_max_task_bytes=self.server_config.artifact_max_task_bytes, ) self.inference_service = inference_service self.cache_service = cache_service diff --git a/telefuser/service/core/artifact_store.py b/telefuser/service/core/artifact_store.py index fa2fa4d..6488aba 100644 --- a/telefuser/service/core/artifact_store.py +++ b/telefuser/service/core/artifact_store.py @@ -150,6 +150,7 @@ def cleanup( retention_seconds: int, tmp_retention_seconds: int, max_total_bytes: int = 0, + max_task_bytes: int = 0, now: datetime | float | int | None = None, ) -> dict[str, Any]: """Clean expired local artifacts without touching active task directories.""" @@ -178,6 +179,16 @@ def cleanup( logger.warning(message) errors.append(message) + if max_task_bytes > 0: + removed_task_ids.extend( + self._cleanup_oversized_tasks( + active_task_ids=active_task_ids, + terminal_task_end_times=terminal_task_end_times, + max_task_bytes=max_task_bytes, + already_removed=set(removed_task_ids), + ) + ) + if max_total_bytes > 0: removed_task_ids.extend( self._cleanup_capacity( @@ -241,6 +252,29 @@ def _cleanup_capacity( logger.warning(f"Failed to remove artifact directory for task {task_id}: {exc}") return removed + def _cleanup_oversized_tasks( + self, + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_end_times: Mapping[str, datetime | float | int | None], + max_task_bytes: int, + already_removed: set[str], + ) -> list[str]: + removed: list[str] = [] + for task_id in terminal_task_end_times: + if task_id in active_task_ids or task_id in already_removed: + continue + + task_dir = self.task_root(task_id) + if not task_dir.exists() or self._tree_size(task_dir) <= max_task_bytes: + continue + try: + shutil.rmtree(task_dir) + removed.append(task_id) + except Exception as exc: + logger.warning(f"Failed to remove oversized artifact directory for task {task_id}: {exc}") + return removed + @staticmethod def _tree_size(root: Path) -> int: if not root.exists(): diff --git a/telefuser/service/core/config.py b/telefuser/service/core/config.py index 3c9309b..7e7e857 100644 --- a/telefuser/service/core/config.py +++ b/telefuser/service/core/config.py @@ -50,6 +50,14 @@ class ServerConfig(BaseSettings): # Cache settings cache_dir: str = Field(default="work_dirs/server_cache", description="Cache directory path") + artifact_storage_backend: Literal["local", "s3"] = Field( + default="local", + description="Artifact storage backend. Only local is implemented in this refactor stage.", + ) + artifact_local_root: str | None = Field( + default=None, + description="Local artifact root. Defaults to cache_dir when unset.", + ) enable_latent_cache: bool | None = Field( default=None, @@ -218,6 +226,11 @@ def validate(self) -> bool: except Exception as e: raise ValueError(f"Invalid configuration: {e}") + @property + def effective_artifact_local_root(self) -> str: + """Return the local artifact root used when no explicit cache_dir override is supplied.""" + return self.artifact_local_root or self.cache_dir + @property def effective_max_concurrent_tasks(self) -> int: """Effective task concurrency — equals num_replicas (one slot per replica).""" diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index 11b502a..7314ef4 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -95,7 +95,13 @@ def initialize_pipeline( def initialize_file_service(self, cache_dir: Path | None = None) -> FileService: """Initialize file service.""" - cache_dir = cache_dir or self._cache_dir or Path(self.config.cache_dir) + if self.config.artifact_storage_backend != "local": + raise RuntimeError( + "Only the local artifact backend is implemented. " + f"Configured backend: {self.config.artifact_storage_backend}" + ) + + cache_dir = cache_dir or self._cache_dir or Path(self.config.effective_artifact_local_root) cache_dir.mkdir(parents=True, exist_ok=True) self.file_service = FileService( @@ -106,6 +112,7 @@ def initialize_file_service(self, cache_dir: Path | None = None) -> FileService: artifact_retention_seconds=self.config.artifact_retention_seconds, artifact_tmp_retention_seconds=self.config.artifact_tmp_retention_seconds, artifact_max_total_bytes=self.config.artifact_max_total_bytes, + artifact_max_task_bytes=self.config.artifact_max_task_bytes, ) return self.file_service diff --git a/telefuser/service/core/file_service.py b/telefuser/service/core/file_service.py index b07ab8e..dbd6c1d 100644 --- a/telefuser/service/core/file_service.py +++ b/telefuser/service/core/file_service.py @@ -35,6 +35,7 @@ def __init__( artifact_retention_seconds: int = 7 * 24 * 60 * 60, artifact_tmp_retention_seconds: int = 60 * 60, artifact_max_total_bytes: int = 0, + artifact_max_task_bytes: int = 0, ) -> None: self.artifact_store = ArtifactStore(cache_dir) self.cache_dir = self.artifact_store.root @@ -51,6 +52,7 @@ def __init__( self.artifact_retention_seconds = artifact_retention_seconds self.artifact_tmp_retention_seconds = artifact_tmp_retention_seconds self.artifact_max_total_bytes = artifact_max_total_bytes + self.artifact_max_task_bytes = artifact_max_task_bytes self._http_client: httpx.AsyncClient | None = None self._client_lock = asyncio.Lock() @@ -355,6 +357,7 @@ def cleanup_artifacts( retention_seconds=self.artifact_retention_seconds, tmp_retention_seconds=self.artifact_tmp_retention_seconds, max_total_bytes=self.artifact_max_total_bytes, + max_task_bytes=self.artifact_max_task_bytes, now=now, ) diff --git a/tests/unit/service/test_artifact_store.py b/tests/unit/service/test_artifact_store.py index a5f24b7..ecf9dbb 100644 --- a/tests/unit/service/test_artifact_store.py +++ b/tests/unit/service/test_artifact_store.py @@ -138,6 +138,49 @@ def test_artifact_store_cleanup_removes_expired_part_files(tmp_path: Path) -> No assert not part_file.exists() +def test_artifact_store_cleanup_removes_oversized_terminal_tasks(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + oversized_output = store.output_path("big.mp4", media_type=MediaType.VIDEO, task_id="oversized-task") + active_output = store.output_path("active.mp4", media_type=MediaType.VIDEO, task_id="active-task") + small_output = store.output_path("small.mp4", media_type=MediaType.VIDEO, task_id="small-task") + oversized_output.write_bytes(b"12345") + active_output.write_bytes(b"12345") + small_output.write_bytes(b"12") + now = datetime.now() + + result = store.cleanup( + active_task_ids={"active-task"}, + terminal_task_end_times={ + "oversized-task": now, + "active-task": now, + "small-task": now, + }, + retention_seconds=0, + tmp_retention_seconds=0, + max_task_bytes=4, + now=now, + ) + + assert result["removed_task_ids"] == ["oversized-task"] + assert not oversized_output.exists() + assert active_output.exists() + assert small_output.exists() + + +def test_file_service_passes_max_task_bytes_to_artifact_cleanup(tmp_path: Path) -> None: + files = FileService(tmp_path, artifact_max_task_bytes=3) + output = files.get_output_path("big.mp4", media_type=MediaType.VIDEO, task_id="task-123") + output.write_bytes(b"1234") + + result = files.cleanup_artifacts( + active_task_ids=set(), + terminal_task_end_times={"task-123": datetime.now()}, + ) + + assert result["removed_task_ids"] == ["task-123"] + assert not output.exists() + + def test_api_server_runs_artifact_cleanup_from_task_snapshot(tmp_path: Path) -> None: task_manager = TaskManager() config = ServerConfig(artifact_retention_seconds=1, artifact_tmp_retention_seconds=0) diff --git a/tests/unit/service/test_config.py b/tests/unit/service/test_config.py index da44c37..2da9046 100644 --- a/tests/unit/service/test_config.py +++ b/tests/unit/service/test_config.py @@ -1,6 +1,7 @@ from __future__ import annotations from telefuser.service.core.config import ServerConfig +from telefuser.service.core.container import ServiceContainer def test_effective_max_concurrent_tasks_default_is_one() -> None: @@ -26,8 +27,48 @@ def test_server_config_ignores_unknown_fields_for_forward_compatibility() -> Non def test_artifact_retention_defaults_are_configured() -> None: config = ServerConfig() + assert config.artifact_storage_backend == "local" + assert config.artifact_local_root is None + assert config.effective_artifact_local_root == config.cache_dir assert config.artifact_retention_seconds == 7 * 24 * 60 * 60 assert config.artifact_tmp_retention_seconds == 60 * 60 assert config.artifact_cleanup_interval_seconds == 60 * 60 assert config.artifact_max_total_bytes == 0 + assert config.artifact_max_task_bytes == 0 assert config.artifact_preserve_failed_outputs is False + + +def test_artifact_local_root_can_override_cache_dir() -> None: + config = ServerConfig(cache_dir="cache-a", artifact_local_root="artifact-a") + + assert config.effective_artifact_local_root == "artifact-a" + + +def test_container_uses_artifact_local_root_when_cache_dir_is_not_explicit(tmp_path) -> None: + config = ServerConfig(cache_dir=str(tmp_path / "cache"), artifact_local_root=str(tmp_path / "artifacts")) + container = ServiceContainer.create(config=config) + + file_service = container.initialize_file_service() + + assert file_service.cache_dir == (tmp_path / "artifacts").resolve() + + +def test_container_explicit_cache_dir_takes_precedence_over_artifact_local_root(tmp_path) -> None: + config = ServerConfig(cache_dir=str(tmp_path / "cache"), artifact_local_root=str(tmp_path / "artifacts")) + container = ServiceContainer.create(config=config, cache_dir=tmp_path / "explicit") + + file_service = container.initialize_file_service() + + assert file_service.cache_dir == (tmp_path / "explicit").resolve() + + +def test_container_rejects_unimplemented_artifact_backend(tmp_path) -> None: + config = ServerConfig(artifact_storage_backend="s3", artifact_local_root=str(tmp_path / "artifacts")) + container = ServiceContainer.create(config=config) + + try: + container.initialize_file_service() + except RuntimeError as exc: + assert "Only the local artifact backend is implemented" in str(exc) + else: + raise AssertionError("Expected S3 artifact backend to be rejected") From e87f5c004fb8e350d7ad22c827b96f6bb93647c8 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:34:34 +0800 Subject: [PATCH 3/6] fix(service): harden runtime configuration boundaries Broaden default rate-limit coverage, stop trusting forwarded client IP headers by default, use timezone-aware health timestamps, and inject runtime config into pipeline and stream services. Add tests for proxy trust behavior, middleware scope, and service/runtime config propagation. --- telefuser/service/api/middleware.py | 16 ++- telefuser/service/api/routers/service.py | 4 +- telefuser/service/api/routers/webrtc.py | 14 +- telefuser/service/core/config.py | 14 +- telefuser/service/core/container.py | 6 +- telefuser/service/core/pipeline_loader.py | 16 ++- telefuser/service/core/pipeline_pool.py | 4 + telefuser/service/core/pipeline_service.py | 34 ++++- telefuser/service/core/replica_worker.py | 5 +- .../service/core/stream_pipeline_service.py | 32 ++++- tests/server/test_middleware.py | 34 ++++- .../service/test_service_refactor_phase0.py | 122 ++++++++++++++++++ 12 files changed, 258 insertions(+), 43 deletions(-) diff --git a/telefuser/service/api/middleware.py b/telefuser/service/api/middleware.py index 87ae5dc..22ce774 100644 --- a/telefuser/service/api/middleware.py +++ b/telefuser/service/api/middleware.py @@ -69,9 +69,10 @@ def to_dict(self) -> dict: class RateLimitMiddleware(BaseHTTPMiddleware): """Rate limiting middleware based on client IP with sliding window. - Operates as a whitelist: only requests whose path starts with one of - ``limited_paths`` are counted and gated. Everything else (status polls, - health checks, file downloads, stream endpoints) passes through untouched. + Only requests whose path starts with one of ``limited_paths`` are counted + and gated. Defaults are configured to cover expensive generation, artifact + download, and stream negotiation paths while leaving liveness/readiness + checks available for infrastructure probes. Uses in-memory storage. For production with multiple instances, consider using Redis-based rate limiting with fastapi-limiter. @@ -83,6 +84,7 @@ def __init__( requests_per_minute: int | None = None, window_size: int | None = None, limited_paths: list[str] | None = None, + trust_forwarded_for: bool = False, enabled: bool = True, ): super().__init__(app) @@ -90,6 +92,7 @@ def __init__( self.requests_per_minute = requests_per_minute or 60 self.window_size = window_size or 60 self.limited_paths = tuple(limited_paths if limited_paths is not None else []) + self.trust_forwarded_for = trust_forwarded_for self._clients: dict[str, list[float]] = {} @@ -147,7 +150,7 @@ async def dispatch(self, request: Request, call_next: callable) -> Response: def _get_client_id(self, request: Request) -> str: """Get client identifier from request.""" forwarded_for = request.headers.get("X-Forwarded-For") - if forwarded_for: + if self.trust_forwarded_for and forwarded_for: return forwarded_for.split(",")[0].strip() if request.client: @@ -328,7 +331,9 @@ def setup_middleware( ) if config is None: - from ..core.config import server_config as config + from ..core.config import ServerConfig + + config = ServerConfig() if enable_rate_limit and config.enable_rate_limit: app.add_middleware( @@ -336,5 +341,6 @@ def setup_middleware( requests_per_minute=config.rate_limit_requests_per_minute, window_size=config.rate_limit_window_size, limited_paths=config.rate_limit_paths, + trust_forwarded_for=config.trust_forwarded_for, enabled=config.enable_rate_limit, ) diff --git a/telefuser/service/api/routers/service.py b/telefuser/service/api/routers/service.py index ac390b2..6c6d3d8 100644 --- a/telefuser/service/api/routers/service.py +++ b/telefuser/service/api/routers/service.py @@ -67,12 +67,12 @@ async def get_metadata(self) -> dict: async def health_check(self) -> dict: """Liveness endpoint for monitoring.""" - from datetime import datetime + from datetime import UTC, datetime status = { "status": "healthy", "ready": self._is_ready(), - "timestamp": datetime.utcnow().isoformat(), + "timestamp": datetime.now(UTC).isoformat(), "version": "1.0.0", } diff --git a/telefuser/service/api/routers/webrtc.py b/telefuser/service/api/routers/webrtc.py index 5a5044d..4e1cde1 100644 --- a/telefuser/service/api/routers/webrtc.py +++ b/telefuser/service/api/routers/webrtc.py @@ -31,25 +31,25 @@ def __init__(self, api_server: ApiServer) -> None: from aiortc import RTCConfiguration, RTCIceServer - from ...core.config import server_config from ...webrtc.session_manager import WebRTCSessionManager + config = api_server.server_config ice_servers: list[RTCIceServer] = [] - for url in server_config.stun_servers: + for url in config.stun_servers: ice_servers.append(RTCIceServer(urls=url)) - if server_config.turn_server: + if config.turn_server: ice_servers.append( RTCIceServer( - urls=server_config.turn_server, - username=server_config.turn_username or "", - credential=server_config.turn_credential or "", + urls=config.turn_server, + username=config.turn_username or "", + credential=config.turn_credential or "", ) ) configuration = RTCConfiguration(iceServers=ice_servers) if ice_servers else None self._session_manager = WebRTCSessionManager( - max_sessions=server_config.webrtc_max_sessions, + max_sessions=config.webrtc_max_sessions, configuration=configuration, ) diff --git a/telefuser/service/core/config.py b/telefuser/service/core/config.py index 7e7e857..09feb78 100644 --- a/telefuser/service/core/config.py +++ b/telefuser/service/core/config.py @@ -141,17 +141,23 @@ class ServerConfig(BaseSettings): ) rate_limit_window_size: int = Field(default=60, ge=10, le=3600, description="Rate limiting window size in seconds") + trust_forwarded_for: bool = Field( + default=False, + description="Trust X-Forwarded-For for rate-limit client identity. Enable only behind trusted proxies.", + ) rate_limit_paths: list = Field( default_factory=lambda: [ "/v1/tasks/create", "/v1/tasks/form", - "/v1/images/generations", - "/v1/videos/generations", + "/v1/images", + "/v1/videos", + "/v1/files/download", + "/v1/stream", ], description=( - "Path prefixes subject to rate limiting (whitelist). " - "Requests whose URL path does not start with any of these are not rate limited." + "Path prefixes subject to rate limiting. Defaults cover expensive task creation, " + "OpenAI-compatible generation, artifact downloads, and stream negotiation paths." ), ) diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index 7314ef4..615aaf0 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -84,7 +84,7 @@ def initialize_pipeline( skip_validation: bool = False, ) -> bool: """Initialize the pipeline service.""" - self.pipeline_service = PipelineService(security_level=self.config.security_level) + self.pipeline_service = PipelineService(security_level=self.config.security_level, config=self.config) return self.pipeline_service.start_pipeline( ppl_file=pipe_path, @@ -217,6 +217,7 @@ def initialize_all( num_replicas=self.config.num_replicas, replica_device_ids=replica_device_ids, security_level_name=self.config.security_level.name, + config=self.config, task_manager=self.task_manager, ) if not pool.start_all( @@ -228,7 +229,7 @@ def initialize_all( return False self.pipeline_service = pool # duck-type compatible else: - self.pipeline_service = PipelineService(security_level=self.config.security_level) + self.pipeline_service = PipelineService(security_level=self.config.security_level, config=self.config) if not self.pipeline_service.start_pipeline( ppl_file=pipe_path, parallelism=parallelism, @@ -251,6 +252,7 @@ def initialize_stream_service( """Initialize stream pipeline service (alternative to initialize_all).""" self.stream_pipeline_service = StreamPipelineService( security_level=self.config.security_level, + config=self.config, ) return self.stream_pipeline_service.start_service( ppl_file=pipe_path, diff --git a/telefuser/service/core/pipeline_loader.py b/telefuser/service/core/pipeline_loader.py index 7c8a5b4..0182896 100644 --- a/telefuser/service/core/pipeline_loader.py +++ b/telefuser/service/core/pipeline_loader.py @@ -8,6 +8,7 @@ import hashlib import importlib.util import sys +from dataclasses import dataclass from pathlib import Path from types import ModuleType @@ -19,7 +20,14 @@ SecurityLevel, validate_with_report, ) -from .config import server_config + + +@dataclass(frozen=True) +class PipelineValidationConfig: + """Runtime validation behavior for dynamic pipeline loading.""" + + allow_unsafe_pipelines: bool = False + strict_validation: bool = True def load_pipeline_module(ppl_file: str, prefix: str = "telefuser_ppl") -> tuple[ModuleType, str]: @@ -51,6 +59,7 @@ def validate_pipeline_file( ppl_file: str, security_level: SecurityLevel, security_validator: PipelineSecurityValidator, + validation_config: PipelineValidationConfig | None = None, ) -> bool: """Validate a pipeline file for security issues. @@ -59,6 +68,7 @@ def validate_pipeline_file( if security_level == SecurityLevel.NONE: logger.warning("Security validation is disabled (SecurityLevel.NONE)") return True + validation_config = validation_config or PipelineValidationConfig() try: logger.info(f"Validating pipeline file: {ppl_file}") @@ -79,7 +89,7 @@ def validate_pipeline_file( if critical_count > 0: raise SecurityError(f"Pipeline file contains {critical_count} critical security violations.") - if getattr(server_config, "allow_unsafe_pipelines", False): + if validation_config.allow_unsafe_pipelines: logger.warning("Allowing unsafe pipeline (allow_unsafe_pipelines=True)") return True @@ -89,6 +99,6 @@ def validate_pipeline_file( raise except Exception as e: logger.error(f"Unexpected validation error: {e}") - if getattr(server_config, "strict_validation", True): + if validation_config.strict_validation: raise SecurityError(f"Validation failed: {e}") return True diff --git a/telefuser/service/core/pipeline_pool.py b/telefuser/service/core/pipeline_pool.py index 40413cc..e508185 100644 --- a/telefuser/service/core/pipeline_pool.py +++ b/telefuser/service/core/pipeline_pool.py @@ -26,6 +26,7 @@ _spawn_ctx = mp_stdlib.get_context("spawn") if TYPE_CHECKING: + from .config import ServerConfig from .task_manager import TaskManager @@ -41,11 +42,13 @@ def __init__( num_replicas: int, replica_device_ids: list[list[str]], security_level_name: str, + config: ServerConfig | None = None, task_manager: TaskManager | None = None, ) -> None: self._num_replicas = num_replicas self._replica_device_ids = replica_device_ids self._security_level_name = security_level_name + self._server_config_data = config.model_dump(mode="json") if config is not None else None self._task_manager = task_manager self._handles: list[ReplicaHandle] = [] self._available: asyncio.Queue[int] = asyncio.Queue() @@ -104,6 +107,7 @@ def _start_all_replicas( cancel_event, self._security_level_name, skip_validation, + self._server_config_data, ), daemon=False, ) diff --git a/telefuser/service/core/pipeline_service.py b/telefuser/service/core/pipeline_service.py index 4524dc3..f928660 100644 --- a/telefuser/service/core/pipeline_service.py +++ b/telefuser/service/core/pipeline_service.py @@ -18,9 +18,14 @@ SecurityError, SecurityLevel, ) -from .config import server_config +from .config import ServerConfig, server_config from .pipeline_contract import load_pipeline_contract -from .pipeline_loader import load_pipeline_module, unload_pipeline_module, validate_pipeline_file +from .pipeline_loader import ( + PipelineValidationConfig, + load_pipeline_module, + unload_pipeline_module, + validate_pipeline_file, +) from .pipeline_runner import PipelineRunner mp.set_start_method("spawn", force=True) @@ -37,8 +42,14 @@ class PipelineService: - Configurable security levels """ - def __init__(self, security_level: SecurityLevel | None = None) -> None: + def __init__( + self, + security_level: SecurityLevel | None = None, + *, + config: ServerConfig | None = None, + ) -> None: """Initialize PipelineService.""" + self.server_config = config or server_config self.is_running = False self.pipeline = None self.task: TaskType | None = None @@ -51,10 +62,14 @@ def __init__(self, security_level: SecurityLevel | None = None) -> None: self._contract = None self._declared_contract = False - self.security_level = security_level or getattr(server_config, "security_level", SecurityLevel.STRICT) + self.security_level = security_level or self.server_config.security_level self.security_validator = PipelineSecurityValidator( security_level=self.security_level, - max_file_size=getattr(server_config, "max_ppl_file_size", 1024 * 1024), + max_file_size=self.server_config.max_ppl_file_size, + ) + self.validation_config = PipelineValidationConfig( + allow_unsafe_pipelines=self.server_config.allow_unsafe_pipelines, + strict_validation=self.server_config.strict_validation, ) logger.info(f"PipelineService initialized with security_level={self.security_level.name}") @@ -65,7 +80,12 @@ def _load_pipeline_module(self, ppl_file: str) -> ModuleType: return module def _validate_pipeline_file(self, ppl_file: str) -> bool: - return validate_pipeline_file(ppl_file, self.security_level, self.security_validator) + return validate_pipeline_file( + ppl_file, + self.security_level, + self.security_validator, + validation_config=self.validation_config, + ) def start_pipeline( self, ppl_file: str, parallelism: int, task: TaskType | str, skip_validation: bool = False @@ -173,7 +193,7 @@ async def run_task_with_stop_event( if not self.is_running or self.pipeline is None or self._runner is None: raise RuntimeError("Pipeline service is not started") - timeout_s = timeout_s if timeout_s is not None else float(getattr(server_config, "task_timeout", 600)) + timeout_s = timeout_s if timeout_s is not None else float(self.server_config.task_timeout) result = await self._runner.run( task_data=task_data, diff --git a/telefuser/service/core/replica_worker.py b/telefuser/service/core/replica_worker.py index 7fad0ab..a6e4525 100644 --- a/telefuser/service/core/replica_worker.py +++ b/telefuser/service/core/replica_worker.py @@ -47,6 +47,7 @@ def _replica_main( cancel_event: mp_stdlib.Event, security_level_name: str, skip_validation: bool, + server_config_data: dict[str, Any] | None = None, ) -> None: """Entry point for a replica subprocess. @@ -60,6 +61,7 @@ def _replica_main( asyncio.set_event_loop(loop) from telefuser.service.core.pipeline_service import PipelineService + from telefuser.service.core.config import ServerConfig from telefuser.service.security.security_validator import SecurityLevel from telefuser.service_types import TaskType from telefuser.utils.logging import logger @@ -67,7 +69,8 @@ def _replica_main( logger.info(f"Replica {replica_id} started: {device_env_var}={visible_devices}, parallelism={parallelism}") security_level = SecurityLevel[security_level_name] - svc = PipelineService(security_level=security_level) + config = ServerConfig.model_validate(server_config_data) if server_config_data is not None else None + svc = PipelineService(security_level=security_level, config=config) try: success = svc.start_pipeline( diff --git a/telefuser/service/core/stream_pipeline_service.py b/telefuser/service/core/stream_pipeline_service.py index f142d79..e7095d1 100644 --- a/telefuser/service/core/stream_pipeline_service.py +++ b/telefuser/service/core/stream_pipeline_service.py @@ -29,8 +29,13 @@ SecurityError, SecurityLevel, ) -from .config import server_config -from .pipeline_loader import load_pipeline_module, unload_pipeline_module, validate_pipeline_file +from .config import ServerConfig, server_config +from .pipeline_loader import ( + PipelineValidationConfig, + load_pipeline_module, + unload_pipeline_module, + validate_pipeline_file, +) # --------------------------------------------------------------------------- # Stream mode constants @@ -77,7 +82,13 @@ class StreamPipelineService: def get_service() -> ServerPushService | BidirectionalService """ - def __init__(self, security_level: SecurityLevel | None = None) -> None: + def __init__( + self, + security_level: SecurityLevel | None = None, + *, + config: ServerConfig | None = None, + ) -> None: + self.server_config = config or server_config self.is_running = False self.service: ServerPushService | BidirectionalService | None = None self.stream_mode: str | None = None @@ -85,10 +96,14 @@ def __init__(self, security_level: SecurityLevel | None = None) -> None: self._module: ModuleType | None = None self._module_name: str | None = None - self.security_level = security_level or getattr(server_config, "security_level", SecurityLevel.STRICT) + self.security_level = security_level or self.server_config.security_level self.security_validator = PipelineSecurityValidator( security_level=self.security_level, - max_file_size=getattr(server_config, "max_ppl_file_size", 1024 * 1024), + max_file_size=self.server_config.max_ppl_file_size, + ) + self.validation_config = PipelineValidationConfig( + allow_unsafe_pipelines=self.server_config.allow_unsafe_pipelines, + strict_validation=self.server_config.strict_validation, ) logger.info(f"StreamPipelineService initialized with security_level={self.security_level.name}") @@ -102,7 +117,12 @@ def start_service(self, ppl_file: str, skip_validation: bool = False) -> bool: try: if not skip_validation: - validate_pipeline_file(ppl_file, self.security_level, self.security_validator) + validate_pipeline_file( + ppl_file, + self.security_level, + self.security_validator, + validation_config=self.validation_config, + ) else: logger.warning("Skipping security validation for pipeline file") diff --git a/tests/server/test_middleware.py b/tests/server/test_middleware.py index 3540a74..4dd5c9d 100644 --- a/tests/server/test_middleware.py +++ b/tests/server/test_middleware.py @@ -22,7 +22,7 @@ RateLimitMiddleware, setup_middleware, ) -from telefuser.service.core.config import ServerConfig, server_config +from telefuser.service.core.config import ServerConfig def _request(path: str = "/test", *, headers: list[tuple[bytes, bytes]] | None = None) -> Request: @@ -71,7 +71,7 @@ def test_endpoint(): assert json.loads(response.body) == {"status": "ok"} def test_rate_limit_only_applies_to_limited_paths(self): - """Test that paths outside the whitelist are not rate limited.""" + """Test that paths outside the configured limited paths are not rate limited.""" app = FastAPI() @app.get("/health") @@ -82,7 +82,6 @@ def health(): def test_endpoint(): return {"status": "ok"} - # Only /test is rate limited; /health passes through untouched. middleware = RateLimitMiddleware( app, requests_per_minute=10, @@ -91,7 +90,6 @@ def test_endpoint(): ) middleware._clients = {} - # /health is not in limited_paths, so it should never be rate limited. for i in range(20): response = asyncio.run(middleware.dispatch(_request("/health"), _ok_call_next)) assert response.status_code == 200, f"Health check failed on request {i + 1}" @@ -121,8 +119,8 @@ def test_endpoint(): assert "X-RateLimit-Window" in response.headers assert response.headers["X-RateLimit-Limit"] == "60" - def test_rate_limit_get_client_id_from_forwarded_for(self): - """Test client identification from X-Forwarded-For header.""" + def test_rate_limit_ignores_forwarded_for_by_default(self): + """Test client identification does not trust forwarded headers by default.""" app = FastAPI() @app.get("/test") @@ -138,6 +136,22 @@ def test_endpoint(): request = _request(headers=[(b"x-forwarded-for", b"192.168.1.1")]) + client_id = middleware._get_client_id(request) + assert client_id == "127.0.0.1" + + def test_rate_limit_can_trust_forwarded_for_when_configured(self): + """Test client identification can use X-Forwarded-For behind trusted proxies.""" + app = FastAPI() + + middleware = RateLimitMiddleware( + app, + requests_per_minute=60, + window_size=60, + limited_paths=["/test"], + trust_forwarded_for=True, + ) + request = _request(headers=[(b"x-forwarded-for", b"192.168.1.1, 10.0.0.1")]) + client_id = middleware._get_client_id(request) assert client_id == "192.168.1.1" @@ -203,6 +217,8 @@ def test_rate_limit_prefix_match(self): assert middleware._is_limited("/v1/tasks/create") is True assert middleware._is_limited("/v1/images/generations") is True + assert middleware._is_limited("/v1/images/edits") is True + assert middleware._is_limited("/v1/videos") is False assert middleware._is_limited("/v1/tasks/123/status") is False assert middleware._is_limited("/v1/files/download/x.mp4") is False assert middleware._is_limited("/v1/service/health") is False @@ -312,8 +328,14 @@ def test_default_rate_limit_config(self): assert config.enable_rate_limit is True assert config.rate_limit_requests_per_minute == 60 assert config.rate_limit_window_size == 60 + assert config.trust_forwarded_for is False assert "/v1/tasks/create" in config.rate_limit_paths assert "/v1/tasks/form" in config.rate_limit_paths + assert "/v1/images" in config.rate_limit_paths + assert "/v1/videos" in config.rate_limit_paths + assert "/v1/files/download" in config.rate_limit_paths + assert "/v1/stream" in config.rate_limit_paths + assert "/v1/service/health" not in config.rate_limit_paths def test_rate_limit_config_validation(self): """Test rate limit config validation.""" diff --git a/tests/unit/service/test_service_refactor_phase0.py b/tests/unit/service/test_service_refactor_phase0.py index 18f9f98..5d9b5fb 100644 --- a/tests/unit/service/test_service_refactor_phase0.py +++ b/tests/unit/service/test_service_refactor_phase0.py @@ -1,6 +1,9 @@ from __future__ import annotations import asyncio +import sys +import threading +import types from pathlib import Path from unittest.mock import Mock @@ -12,6 +15,8 @@ from telefuser.service.api.routers.service import ServiceRoutes from telefuser.service.core.config import ServerConfig from telefuser.service.core.file_service import FileService +from telefuser.service.core.pipeline_service import PipelineService +from telefuser.service.core.stream_pipeline_service import StreamPipelineService from telefuser.service.core.task_manager import TaskManager, TaskStatus as CoreTaskStatus from telefuser.service.security.security_validator import SecurityLevel from telefuser.service_types import MediaType, TaskStatus @@ -173,3 +178,120 @@ def fake_create(config=None, cache_dir=None): ) assert captured["config"].security_level is SecurityLevel.BASIC + + +def test_pipeline_service_uses_injected_config() -> None: + config = ServerConfig( + security_level=SecurityLevel.BASIC, + max_ppl_file_size=4096, + allow_unsafe_pipelines=True, + strict_validation=False, + ) + + service = PipelineService(config=config) + + assert service.security_level is SecurityLevel.BASIC + assert service.security_validator.max_file_size == 4096 + assert service.validation_config.allow_unsafe_pipelines is True + assert service.validation_config.strict_validation is False + + +def test_pipeline_service_task_timeout_uses_injected_config() -> None: + class Status: + value = "completed" + + class Result: + status = Status() + output_path = "output.mp4" + message = "ok" + raw = {"ok": True} + + class Runner: + def __init__(self) -> None: + self.timeout_s = None + + async def run(self, **kwargs: object) -> Result: + self.timeout_s = kwargs["timeout_s"] + return Result() + + config = ServerConfig(task_timeout=60, security_level=SecurityLevel.NONE) + service = PipelineService(config=config) + runner = Runner() + service.is_running = True + service.pipeline = object() + service._runner = runner + + result = asyncio.run(service.run_task_with_stop_event({"task_id": "task-1"}, threading.Event())) + + assert runner.timeout_s == 60.0 + assert result["status"] == "completed" + + +def test_stream_pipeline_service_uses_injected_config() -> None: + config = ServerConfig( + security_level=SecurityLevel.NONE, + max_ppl_file_size=8192, + allow_unsafe_pipelines=True, + strict_validation=False, + ) + + service = StreamPipelineService(config=config) + + assert service.security_level is SecurityLevel.NONE + assert service.security_validator.max_file_size == 8192 + assert service.validation_config.allow_unsafe_pipelines is True + assert service.validation_config.strict_validation is False + + +def test_webrtc_routes_use_api_server_config(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + class FakeRTCIceServer: + def __init__(self, urls: str, username: str | None = None, credential: str | None = None) -> None: + self.urls = urls + self.username = username + self.credential = credential + + class FakeRTCConfiguration: + def __init__(self, iceServers: list[FakeRTCIceServer]) -> None: + self.iceServers = iceServers + + class FakeWebRTCSessionManager: + def __init__(self, max_sessions: int, configuration: FakeRTCConfiguration) -> None: + captured["max_sessions"] = max_sessions + captured["configuration"] = configuration + + monkeypatch.setitem( + sys.modules, + "aiortc", + types.SimpleNamespace(RTCConfiguration=FakeRTCConfiguration, RTCIceServer=FakeRTCIceServer), + ) + monkeypatch.setitem( + sys.modules, + "telefuser.service.webrtc.session_manager", + types.SimpleNamespace(WebRTCSessionManager=FakeWebRTCSessionManager), + ) + + config = ServerConfig( + webrtc_max_sessions=3, + stun_servers=["stun:local:3478"], + turn_server="turn:local:3478", + turn_username="user", + turn_credential="secret", + ) + server = ApiServer( + task_manager=TaskManager(), + enable_openai_api=False, + config=config, + route_profile="request_response", + ) + + from telefuser.service.api.routers.webrtc import WebRTCRoutes + + WebRTCRoutes(server) + + assert captured["max_sessions"] == 3 + ice_servers = captured["configuration"].iceServers + assert [server.urls for server in ice_servers] == ["stun:local:3478", "turn:local:3478"] + assert ice_servers[1].username == "user" + assert ice_servers[1].credential == "secret" From a39bbe41215ebe9a85418de06f5dfdf7675e9ac9 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:34:47 +0800 Subject: [PATCH 4/6] refactor(service): split request and stream route profiles Separate request-response and stream route profiles so stream apps expose only stream/service routes, and make WebRTC/session close behavior observable across session owners. Add route-profile and WebRTC lifecycle tests for stream-only serving boundaries. --- telefuser/service/api/api_server.py | 38 ++--- telefuser/service/api/routers/stream.py | 12 +- telefuser/service/api/routers/webrtc.py | 8 +- telefuser/service/core/container.py | 2 + telefuser/service/webrtc/chunk_router.py | 9 +- telefuser/service/webrtc/session_manager.py | 50 ++++--- telefuser/service/webrtc/track.py | 28 +++- .../service/test_service_refactor_phase0.py | 141 +++++++++++++++++- 8 files changed, 230 insertions(+), 58 deletions(-) diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index 9169f26..2d41e15 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -46,6 +46,7 @@ def __init__( enable_logging: bool = False, enable_openai_api: bool = True, config: ServerConfig | None = None, + route_profile: str = "all", ) -> None: self.server_config = config or server_config self.app = app or FastAPI( @@ -69,6 +70,7 @@ def __init__( self.configured_max_concurrent_tasks = configured_max_concurrent_tasks or max_concurrent_tasks self._task_manager = task_manager self.enable_openai_api = enable_openai_api + self.route_profile = route_profile self.task_processor: AsyncTaskProcessor | None = None self._task_processor_lock = asyncio.Lock() @@ -94,26 +96,28 @@ def task_manager(self): def _setup_routes(self) -> None: """Setup all API routes.""" - tasks_router = routers.tasks.setup_routes(self) - files_router = routers.files.setup_routes(self) service_router = routers.service.setup_routes(self) - - self.app.include_router(tasks_router) - self.app.include_router(files_router) self.app.include_router(service_router) - stream_router = routers.setup_stream_routes(self) - self.app.include_router(stream_router) - - if routers.setup_webrtc_routes is not None: - try: - webrtc_router = routers.setup_webrtc_routes(self) - self.app.include_router(webrtc_router) - logger.info("WebRTC routes enabled at /v1/stream/webrtc") - except Exception as e: - logger.info(f"WebRTC routes not available: {e}") - - if self.enable_openai_api: + if self.route_profile in {"all", "request_response"}: + tasks_router = routers.tasks.setup_routes(self) + files_router = routers.files.setup_routes(self) + self.app.include_router(tasks_router) + self.app.include_router(files_router) + + if self.route_profile in {"all", "stream"}: + stream_router = routers.setup_stream_routes(self) + self.app.include_router(stream_router) + + if routers.setup_webrtc_routes is not None: + try: + webrtc_router = routers.setup_webrtc_routes(self) + self.app.include_router(webrtc_router) + logger.info("WebRTC routes enabled at /v1/stream/webrtc") + except Exception as e: + logger.info(f"WebRTC routes not available: {e}") + + if self.enable_openai_api and self.route_profile in {"all", "request_response"}: try: from .openai import image_routes, video_routes diff --git a/telefuser/service/api/routers/stream.py b/telefuser/service/api/routers/stream.py index a215dd9..e83b5c9 100644 --- a/telefuser/service/api/routers/stream.py +++ b/telefuser/service/api/routers/stream.py @@ -13,6 +13,8 @@ from fastapi import APIRouter, HTTPException +from telefuser.utils.logging import logger + if TYPE_CHECKING: from ..api_server import ApiServer @@ -38,12 +40,16 @@ async def close_session(self, session_id: str) -> dict: try: svc.close_session(session_id) pipeline_closed = True - except Exception: - pass + except Exception as exc: + logger.warning(f"Failed to close pipeline stream session {session_id}: {exc}") webrtc_closed = False webrtc_routes = self.api._webrtc_routes if webrtc_routes is not None: - webrtc_closed = await webrtc_routes._session_manager.close_session(session_id) + webrtc_closed = await webrtc_routes._session_manager.close_session( + session_id, + reason="stream_session_delete", + notify_pipeline=False, + ) if not pipeline_closed and not webrtc_closed: raise HTTPException(status_code=404, detail=f"Session {session_id} not found") return {"session_id": session_id, "status": "closed"} diff --git a/telefuser/service/api/routers/webrtc.py b/telefuser/service/api/routers/webrtc.py index 4e1cde1..5108ac1 100644 --- a/telefuser/service/api/routers/webrtc.py +++ b/telefuser/service/api/routers/webrtc.py @@ -128,15 +128,9 @@ async def _handle_bidirectional(self, svc, request: WebRTCOfferRequest) -> WebRT ) async def close_session(self, session_id: str) -> dict: - closed = await self._session_manager.close_session(session_id) + closed = await self._session_manager.close_session(session_id, reason="webrtc_session_delete") if not closed: raise HTTPException(status_code=404, detail=f"WebRTC session {session_id} not found") - svc = self.api.stream_service - if svc and svc.stream_mode == STREAM_MODE_BIDIRECTIONAL: - try: - svc.close_session(session_id) - except Exception: - pass return {"session_id": session_id, "status": "closed"} async def cleanup(self) -> None: diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index 615aaf0..e55e9f8 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -261,6 +261,7 @@ def initialize_stream_service( def get_api_app(self, enable_rate_limit: bool = True) -> FastAPI: """Get FastAPI application with all services initialized.""" + route_profile = "stream" if self.stream_pipeline_service and not self.pipeline_service else "request_response" api_server = ApiServer( max_queue_size=self.config.max_queue_size, max_concurrent_tasks=self.config.effective_max_concurrent_tasks, @@ -269,6 +270,7 @@ def get_api_app(self, enable_rate_limit: bool = True) -> FastAPI: enable_rate_limit=enable_rate_limit, enable_logging=False, config=self.config, + route_profile=route_profile, ) if self.file_service and self.pipeline_service: diff --git a/telefuser/service/webrtc/chunk_router.py b/telefuser/service/webrtc/chunk_router.py index 6bd1222..487eb10 100644 --- a/telefuser/service/webrtc/chunk_router.py +++ b/telefuser/service/webrtc/chunk_router.py @@ -73,6 +73,7 @@ def _route_chunk(self, chunk: dict) -> None: np_arr = np.frombuffer(raw, dtype=np.uint8) bgr = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) if bgr is None: + logger.warning(f"ChunkRouter dropped undecodable video frame: session={self._session_id}") continue rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) frame = av.VideoFrame.from_ndarray(rgb, format="rgb24") @@ -95,8 +96,8 @@ def _route_chunk(self, chunk: dict) -> None: ) try: self._dc_send(msg.model_dump_json()) - except Exception: - pass + except Exception as exc: + logger.warning(f"ChunkRouter metadata send failed: session={self._session_id} {exc}") def _send_done(self) -> None: if self._dc_send is None: @@ -107,5 +108,5 @@ def _send_done(self) -> None: ) try: self._dc_send(done.model_dump_json()) - except Exception: - pass + except Exception as exc: + logger.warning(f"ChunkRouter done send failed: session={self._session_id} {exc}") diff --git a/telefuser/service/webrtc/session_manager.py b/telefuser/service/webrtc/session_manager.py index cba931f..1a2a4e7 100644 --- a/telefuser/service/webrtc/session_manager.py +++ b/telefuser/service/webrtc/session_manager.py @@ -91,7 +91,7 @@ async def _on_state_change() -> None: state = pc.connectionState if state in ("failed", "disconnected", "closed"): logger.info(f"WebRTC connection {state}: session={session_id}") - await self.close_session(session_id) + await self.close_session(session_id, reason=f"connection_{state}") pc.addTrack(track) if audio_track is not None: @@ -178,12 +178,16 @@ def _on_datachannel(channel) -> None: def _on_message(message) -> None: try: data = json.loads(message) if isinstance(message, str) else message - except (json.JSONDecodeError, TypeError): + except (json.JSONDecodeError, TypeError) as exc: + logger.warning(f"DataChannel message decode failed: session={session_id} {exc}") return if isinstance(data, dict) and data.get("type") == "stop": - asyncio.ensure_future(self.close_session(session_id)) + asyncio.ensure_future(self.close_session(session_id, reason="client_stop")) return - on_input(session_id, data) + try: + on_input(session_id, data) + except Exception as exc: + logger.warning(f"DataChannel input callback failed: session={session_id} {exc}") router = ChunkRouter( generator=output_generator, @@ -216,7 +220,7 @@ async def _on_state_change() -> None: state = pc.connectionState if state in ("failed", "disconnected", "closed"): logger.info(f"WebRTC connection {state}: session={session_id}") - await self.close_session(session_id) + await self.close_session(session_id, reason=f"connection_{state}") # --- SDP exchange ----------------------------------------------- @@ -237,48 +241,58 @@ async def _on_state_change() -> None: # -- Session lifecycle --------------------------------------------------- - async def close_session(self, session_id: str) -> bool: + async def close_session(self, session_id: str, *, reason: str = "api", notify_pipeline: bool = True) -> bool: async with self._lock: entry = self._sessions.pop(session_id, None) if entry is None or entry is _SENTINEL: return False + logger.info(f"Closing WebRTC session: session={session_id} reason={reason}") + if isinstance(entry, _Session): entry.track.stop() if entry.track._task is not None and not entry.track._task.done(): try: await entry.track._task - except (asyncio.CancelledError, Exception): - pass + except asyncio.CancelledError: + logger.info(f"WebRTC server-push track cancelled: session={session_id}") + except Exception as exc: + logger.warning(f"WebRTC server-push track close failed: session={session_id} {exc}") elif isinstance(entry, _BidirectionalSession): for task in entry.relay_tasks: if not task.done(): task.cancel() try: await task - except (asyncio.CancelledError, Exception): - pass + except asyncio.CancelledError: + logger.info(f"WebRTC relay task cancelled: session={session_id}") + except Exception as exc: + logger.warning(f"WebRTC relay task close failed: session={session_id} {exc}") if entry.router_task is not None and not entry.router_task.done(): entry.router_task.cancel() try: await entry.router_task - except (asyncio.CancelledError, Exception): - pass + except asyncio.CancelledError: + logger.info(f"WebRTC chunk router cancelled: session={session_id}") + except Exception as exc: + logger.warning(f"WebRTC chunk router close failed: session={session_id} {exc}") if entry.output_video_track is not None: entry.output_video_track.stop() if entry.output_audio_track is not None: entry.output_audio_track.stop() - if entry.on_close is not None: + if notify_pipeline and entry.on_close is not None: try: entry.on_close(entry.session_id) - except Exception: - pass + except Exception as exc: + logger.warning(f"WebRTC pipeline close callback failed: session={session_id} {exc}") try: await asyncio.wait_for(entry.pc.close(), timeout=5.0) - except (asyncio.TimeoutError, Exception): - pass - logger.info(f"WebRTC session closed: session={session_id}") + except asyncio.TimeoutError: + logger.warning(f"Timed out closing WebRTC peer connection: session={session_id}") + except Exception as exc: + logger.warning(f"WebRTC peer connection close failed: session={session_id} {exc}") + logger.info(f"WebRTC session closed: session={session_id} reason={reason}") return True async def close_all(self) -> None: diff --git a/telefuser/service/webrtc/track.py b/telefuser/service/webrtc/track.py index ac29191..69fb5e7 100644 --- a/telefuser/service/webrtc/track.py +++ b/telefuser/service/webrtc/track.py @@ -53,13 +53,18 @@ def __init__( self._frame_count = 0 self._start_time: float | None = None self._finished = False + self.dropped_chunks = 0 - def feed(self, data: bytes) -> None: + def feed(self, data: bytes) -> bool: """Push raw PCM16 bytes (called by FrameGeneratorTrack from its consumer task).""" try: self._queue.put_nowait(data) + return True except asyncio.QueueFull: - pass + self.dropped_chunks += 1 + if self.dropped_chunks == 1 or self.dropped_chunks % 100 == 0: + logger.warning(f"AudioGeneratorTrack queue full; dropped_chunks={self.dropped_chunks}") + return False def signal_done(self) -> None: self._finished = True @@ -141,15 +146,20 @@ def __init__( self._last_frame: av.VideoFrame | None = None self._placeholder_width = 640 self._placeholder_height = 360 + self.dropped_frames = 0 - def push_frame(self, frame: av.VideoFrame) -> None: + def push_frame(self, frame: av.VideoFrame) -> bool: """Push a decoded frame directly (used by ChunkRouter in bidirectional mode).""" self._placeholder_width = frame.width self._placeholder_height = frame.height try: self._queue.put_nowait(frame) + return True except asyncio.QueueFull: - pass + self.dropped_frames += 1 + if self.dropped_frames == 1 or self.dropped_frames % 100 == 0: + logger.warning(f"FrameGeneratorTrack queue full; dropped_frames={self.dropped_frames}") + return False def signal_done(self) -> None: self._finished = True @@ -262,9 +272,10 @@ async def run(self) -> None: rgb = frame.to_ndarray(format="rgb24") self._on_chunk(self._session_id, {"type": "media", "video_frames": [rgb]}) except MediaStreamError: - pass + logger.info(f"IncomingVideoRelay ended: session={self._session_id}") except asyncio.CancelledError: - pass + logger.info(f"IncomingVideoRelay cancelled: session={self._session_id}") + raise except Exception as exc: logger.error(f"IncomingVideoRelay error: session={self._session_id} {exc}") @@ -300,8 +311,9 @@ async def run(self) -> None: }, ) except MediaStreamError: - pass + logger.info(f"IncomingAudioRelay ended: session={self._session_id}") except asyncio.CancelledError: - pass + logger.info(f"IncomingAudioRelay cancelled: session={self._session_id}") + raise except Exception as exc: logger.error(f"IncomingAudioRelay error: session={self._session_id} {exc}") diff --git a/tests/unit/service/test_service_refactor_phase0.py b/tests/unit/service/test_service_refactor_phase0.py index 5d9b5fb..8be197e 100644 --- a/tests/unit/service/test_service_refactor_phase0.py +++ b/tests/unit/service/test_service_refactor_phase0.py @@ -13,11 +13,14 @@ from telefuser.entrypoints.cli.main import main from telefuser.service.api.api_server import ApiServer from telefuser.service.api.routers.service import ServiceRoutes +from telefuser.service.api.routers.stream import StreamRoutes from telefuser.service.core.config import ServerConfig +from telefuser.service.core.container import ServiceContainer from telefuser.service.core.file_service import FileService from telefuser.service.core.pipeline_service import PipelineService from telefuser.service.core.stream_pipeline_service import StreamPipelineService -from telefuser.service.core.task_manager import TaskManager, TaskStatus as CoreTaskStatus +from telefuser.service.core.task_manager import TaskManager +from telefuser.service.core.task_manager import TaskStatus as CoreTaskStatus from telefuser.service.security.security_validator import SecurityLevel from telefuser.service_types import MediaType, TaskStatus @@ -88,6 +91,142 @@ class RunningPipeline: assert health["pipeline_ready"] is True +def test_stream_route_profile_exposes_only_stream_service_routes() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=True, route_profile="stream") + paths = {route.path for route in server.app.routes} + + assert "/v1/service/health" in paths + assert "/v1/stream/sessions/{session_id}/status" in paths + assert "/v1/tasks/create" not in paths + assert "/v1/tasks/form" not in paths + assert "/v1/files/download/{file_path:path}" not in paths + assert "/v1/images/generations" not in paths + assert "/v1/videos" not in paths + + +def test_request_response_route_profile_excludes_stream_routes() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False, route_profile="request_response") + paths = {route.path for route in server.app.routes} + + assert "/v1/service/health" in paths + assert "/v1/tasks/create" in paths + assert "/v1/files/download/{file_path:path}" in paths + assert "/v1/stream/sessions/{session_id}/status" not in paths + + +def test_container_stream_app_uses_stream_route_profile() -> None: + container = ServiceContainer.create(config=ServerConfig()) + container.stream_pipeline_service = Mock() + container.stream_pipeline_service.is_running = True + + app = container.get_api_app() + paths = {route.path for route in app.routes} + + assert "/v1/service/health" in paths + assert "/v1/stream/sessions/{session_id}/status" in paths + assert "/v1/tasks/create" not in paths + assert "/v1/files/download/{file_path:path}" not in paths + + +def test_stream_session_close_logs_pipeline_close_failure(monkeypatch: pytest.MonkeyPatch) -> None: + class RunningStreamService: + is_running = True + + def close_session(self, session_id: str) -> None: + raise RuntimeError("pipeline close failed") + + class WebRTCSessionManager: + async def close_session(self, session_id: str, **kwargs: object) -> bool: + return True + + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.stream_service = RunningStreamService() + server._webrtc_routes = Mock(_session_manager=WebRTCSessionManager()) + warnings: list[str] = [] + monkeypatch.setattr("telefuser.service.api.routers.stream.logger.warning", warnings.append) + + result = asyncio.run(StreamRoutes(server).close_session("session-123")) + + assert result == {"session_id": "session-123", "status": "closed"} + assert warnings == ["Failed to close pipeline stream session session-123: pipeline close failed"] + + +def test_stream_session_alias_closes_webrtc_without_pipeline_notify() -> None: + class RunningStreamService: + is_running = True + + def __init__(self) -> None: + self.closed_sessions: list[str] = [] + + def close_session(self, session_id: str) -> None: + self.closed_sessions.append(session_id) + + class WebRTCSessionManager: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def close_session(self, session_id: str, *, reason: str, notify_pipeline: bool) -> bool: + self.calls.append( + { + "session_id": session_id, + "reason": reason, + "notify_pipeline": notify_pipeline, + } + ) + return True + + stream_service = RunningStreamService() + webrtc_manager = WebRTCSessionManager() + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.stream_service = stream_service + server._webrtc_routes = Mock(_session_manager=webrtc_manager) + + result = asyncio.run(StreamRoutes(server).close_session("session-123")) + + assert result == {"session_id": "session-123", "status": "closed"} + assert stream_service.closed_sessions == ["session-123"] + assert webrtc_manager.calls == [ + { + "session_id": "session-123", + "reason": "stream_session_delete", + "notify_pipeline": False, + } + ] + + +def test_webrtc_delete_uses_session_manager_pipeline_owner() -> None: + from telefuser.service.api.routers.webrtc import WebRTCRoutes + + class RunningStreamService: + stream_mode = "bidirectional" + + def __init__(self) -> None: + self.closed_sessions: list[str] = [] + + def close_session(self, session_id: str) -> None: + self.closed_sessions.append(session_id) + + class WebRTCSessionManager: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def close_session(self, session_id: str, *, reason: str) -> bool: + self.calls.append({"session_id": session_id, "reason": reason}) + return True + + stream_service = RunningStreamService() + webrtc_manager = WebRTCSessionManager() + routes = WebRTCRoutes.__new__(WebRTCRoutes) + routes.api = Mock(stream_service=stream_service) + routes._session_manager = webrtc_manager + + result = asyncio.run(routes.close_session("session-123")) + + assert result == {"session_id": "session-123", "status": "closed"} + assert stream_service.closed_sessions == [] + assert webrtc_manager.calls == [{"session_id": "session-123", "reason": "webrtc_session_delete"}] + + def test_cli_serve_forwards_security_and_skip_validation(monkeypatch: pytest.MonkeyPatch) -> None: captured = {} From efe9d9e66e377a36907465899a148f4f711b3a9d Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:38:08 +0800 Subject: [PATCH 5/6] feat(service): finalize artifact metadata and docs Complete local artifact lifecycle behavior with persistent/ephemeral cleanup, failed-output preservation, local artifact ids, artifact metadata, and OpenAI response extensions. Update English and Chinese service/stream docs to describe runtime boundaries, local-only artifact semantics, artifact-id downloads, and OpenAI artifact metadata fields. Keep S3-compatible storage and persistent multi-instance queues explicitly deferred. --- docs/en/service.md | 119 ++++++++++++- docs/en/stream_server.md | 15 +- docs/zh/service.md | 117 ++++++++++++- docs/zh/stream_server.md | 15 +- telefuser/service/api/api_server.py | 2 + telefuser/service/api/openai/adapter.py | 48 +++++ telefuser/service/api/openai/image_routes.py | 28 ++- telefuser/service/api/openai/protocol.py | 4 + telefuser/service/api/openai/video_routes.py | 56 +++--- .../service/api/task_application_service.py | 49 ++++++ telefuser/service/core/artifact_store.py | 141 ++++++++++++++- telefuser/service/core/config.py | 9 +- telefuser/service/core/container.py | 2 + telefuser/service/core/file_service.py | 26 +++ telefuser/service/core/pipeline_service.py | 2 +- telefuser/service/core/replica_worker.py | 2 +- telefuser/service/core/task_manager.py | 7 +- .../service/security/security_validator.py | 20 +-- tests/unit/openai/test_image_routes.py | 3 +- tests/unit/openai/test_integration_server.py | 3 +- tests/unit/openai/test_video_routes.py | 3 +- tests/unit/service/test_artifact_store.py | 164 +++++++++++++++++- tests/unit/service/test_config.py | 10 ++ tests/unit/service/test_latent_cache_cli.py | 5 +- tests/unit/service/test_security_validator.py | 14 ++ tests/unit/service/test_service_smoke.py | 90 ++++++++++ .../service/test_task_application_service.py | 19 ++ tests/unit/service/test_task_runtime.py | 1 + 28 files changed, 893 insertions(+), 81 deletions(-) create mode 100644 tests/unit/service/test_security_validator.py create mode 100644 tests/unit/service/test_service_smoke.py diff --git a/docs/en/service.md b/docs/en/service.md index ab77b35..9116b49 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -150,7 +150,7 @@ telefuser serve /path/to/pipeline --task i2v [OPTIONS] | `--cache-dir` | `-c` | string | `work_dirs/server_cache` | Cache directory | | `--parallelism` | `-g` | int | `1` | Number of parallel workers | | `--num-replicas` | `-n` | int | `1` | Number of independent pipeline replicas (Pipeline Pool) | -| `--security-level` | | choice | `strict` | Security level: none/basic/strict/sandbox | +| `--security-level` | | choice | `strict` | Validation level: none/basic/strict/sandbox. `sandbox` is a best-effort restricted-load check, not runtime isolation. | | `--skip-validation` | | flag | `False` | Skip security validation | | `--validate-only` | | flag | `False` | Only validate without starting | @@ -186,7 +186,7 @@ telefuser validate /path/to/pipeline.py [OPTIONS] | Parameter | Default | Description | |-----------|---------|-------------| | `pipeline_file` | **Required** | Path to pipeline Python file | -| `--level` | `strict` | Security level: none/basic/strict/sandbox | +| `--level` | `strict` | Validation level: none/basic/strict/sandbox. `sandbox` is a best-effort restricted-load check, not runtime isolation. | | `--json` | `False` | Output in JSON format | ```bash @@ -258,7 +258,17 @@ telefuser serve --help | `TELEFUSER_HOST` | Server host | `127.0.0.1` | | `TELEFUSER_PORT` | Server port | `8000` | | `TELEFUSER_RATE_LIMIT_ENABLED` | Enable rate limiting | `true` | -| `TELEFUSER_RATE_LIMIT_RPM` | Requests per minute limit | `60` | +| `TELEFUSER_RATE_LIMIT_REQUESTS_PER_MINUTE` | Requests per minute limit | `60` | +| `TELEFUSER_TRUST_FORWARDED_FOR` | Trust `X-Forwarded-For` for rate-limit identity. Enable only behind trusted proxies. | `false` | +| `TELEFUSER_ARTIFACT_STORAGE_BACKEND` | Artifact backend. Only `local` is implemented currently. | `local` | +| `TELEFUSER_ARTIFACT_LOCAL_ROOT` | Local artifact root. Defaults to `TELEFUSER_CACHE_DIR`. | unset | +| `TELEFUSER_ARTIFACT_PERSISTENCE_MODE` | Local artifact retention mode: `persistent` or `ephemeral`. | `persistent` | +| `TELEFUSER_ARTIFACT_RETENTION_SECONDS` | Retention time for terminal task artifacts. `0` disables TTL cleanup. | `604800` | +| `TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS` | Retention time for temporary `.part` files. `0` disables tmp cleanup. | `3600` | +| `TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS` | Background artifact cleanup interval. | `3600` | +| `TELEFUSER_ARTIFACT_MAX_TOTAL_BYTES` | Local artifact cache size limit. `0` disables capacity cleanup. | `0` | +| `TELEFUSER_ARTIFACT_MAX_TASK_BYTES` | Per-task artifact size limit for terminal tasks. `0` disables per-task cleanup. | `0` | +| `TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS` | Preserve failed task directories from automatic artifact cleanup. | `false` | ### Configuration File Example @@ -269,9 +279,66 @@ TELEFUSER_SECURITY_LEVEL=STRICT TELEFUSER_PORT=8080 TELEFUSER_HOST=0.0.0.0 TELEFUSER_RATE_LIMIT_ENABLED=true -TELEFUSER_RATE_LIMIT_RPM=100 +TELEFUSER_RATE_LIMIT_REQUESTS_PER_MINUTE=100 ``` +### Runtime Boundaries + +The request-response service is intentionally local and single-process by default: + +- Task status, queue state, cancellation state, and rate-limit buckets are kept in process memory. They are not shared across multiple server instances and are not restored after restart. +- File inputs, generated outputs, and temporary `.part` files are managed by the local artifact store under the configured cache root. +- `artifact_storage_backend=s3` is declared as a future backend boundary, but it is not implemented yet. Setting it to `s3` fails at startup instead of silently falling back to local storage. +- `telefuser serve` exposes request-response routes only: `/v1/tasks/*`, `/v1/files/*`, `/v1/images/*`, `/v1/videos/*`, and `/v1/service/*`. +- `telefuser stream-serve` exposes stream routes only: `/v1/stream/*`, `/v1/stream/webrtc/*`, and `/v1/service/*`. It does not expose task, file-download, or OpenAI-compatible request-response routes. + +### Artifact Storage and Cleanup + +Artifacts are local files rooted at `TELEFUSER_ARTIFACT_LOCAL_ROOT` when set, otherwise `TELEFUSER_CACHE_DIR`. +Uploads and remote downloads are written through temporary `.part` files and atomically renamed after successful completion. + +Outputs are task-scoped when a task id is available: + +```text +/tasks//inputs/ +/tasks//outputs/ +/tasks//tmp/ +``` + +Only paths resolved inside the artifact root are accepted. Absolute output paths and `..` traversal are rejected. +Download routes accept validated output paths and local artifact ids. Local artifact ids use the +`local:tasks//outputs//` form and only refer to downloadable outputs in the local +artifact backend. They are not remote object-storage ids. + +Cleanup is best-effort and never removes active tasks (`pending`, `processing`, or `streaming`). Terminal tasks +(`completed`, `failed`, or `cancelled`) may be removed by TTL cleanup, total-cache capacity cleanup, or per-task size cleanup. +Temporary `.part` files are cleaned independently by `TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS`. + +`TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent` keeps terminal task artifacts until TTL or capacity cleanup removes +them. `ephemeral` removes terminal task directories on the next cleanup pass. If +`TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true`, failed task directories are protected from automatic cleanup so +partial outputs can be inspected manually. + +### Rate Limit and Proxy Headers + +Rate limiting is in-memory and keyed by client identity. By default, the server uses the direct client address and does +not trust `X-Forwarded-For`, because arbitrary clients can spoof that header. Set +`TELEFUSER_TRUST_FORWARDED_FOR=true` only when TeleFuser is behind a trusted reverse proxy that sanitizes forwarded +headers. + +The default limited paths cover expensive generation endpoints, artifact downloads, and stream negotiation: + +```text +/v1/tasks/create +/v1/tasks/form +/v1/images +/v1/videos +/v1/files/download +/v1/stream +``` + +Liveness/readiness endpoints remain unthrottled for infrastructure probes. + ### Pipeline Contract and Parameter Definitions The service does not infer pipeline capabilities from Python signatures alone. Instead, it loads an explicit @@ -463,9 +530,9 @@ When the server is running: | POST | `/tasks/create` | Create a new generation task | | POST | `/tasks/form` | Create task with file upload | | GET | `/tasks/{task_id}/status` | Get task status | -| DELETE | `/tasks/{task_id}` | Cancel a task | +| DELETE | `/tasks/{task_id}` | Request cooperative task cancellation | | GET | `/tasks/queue/status` | Get queue status | -| GET | `/files/download/{file_id}` | Download output files | +| GET | `/files/download/{file_id}` | Download an output file by validated path or local artifact id | | GET | `/service/health` | Health check | | GET | `/service/status` | Service status | | GET | `/service/metadata` | Service metadata | @@ -814,7 +881,17 @@ TeleFuser provides an **OpenAI-compatible** REST API that allows you to use stan "data": [ { "url": "http://localhost:8000/v1/images/task_xxx/content", - "revised_prompt": "a beautiful cat" + "revised_prompt": "a beautiful cat", + "file_path": "/cache/tasks/task_xxx/outputs/images/result.png", + "artifact_id": "local:tasks/task_xxx/outputs/images/result.png", + "artifact_metadata": { + "backend": "local", + "relative_path": "tasks/task_xxx/outputs/images/result.png", + "task_id": "task_xxx", + "media_type": "image", + "filename": "result.png", + "size_bytes": 123456 + } } ], "peak_memory_mb": 4096.5, @@ -875,7 +952,7 @@ response.data[0].save("sunset.png") "size": "1024x576", // Optional: video size "seed": 1024, // Optional: random seed "negative_prompt": "blurry", // Optional: negative prompt - "output_path": "/path/to/output.mp4" // Optional: custom output path + "output_path": "custom-output.mp4" // Optional: relative custom output path } ``` @@ -891,7 +968,31 @@ response.data[0].save("sunset.png") "created_at": 1699000000, "size": "1024x576", "seconds": "5", - "url": null + "url": null, + "file_path": null, + "artifact_id": null, + "artifact_metadata": null +} +``` + +Completed video responses include a content URL and local artifact metadata: + +```json +{ + "id": "vid_xxx", + "status": "completed", + "progress": 100, + "url": "http://localhost:8000/v1/videos/vid_xxx/content", + "file_path": "/cache/tasks/vid_xxx/outputs/videos/result.mp4", + "artifact_id": "local:tasks/vid_xxx/outputs/videos/result.mp4", + "artifact_metadata": { + "backend": "local", + "relative_path": "tasks/vid_xxx/outputs/videos/result.mp4", + "task_id": "vid_xxx", + "media_type": "video", + "filename": "result.mp4", + "size_bytes": 987654 + } } ``` diff --git a/docs/en/stream_server.md b/docs/en/stream_server.md index da18650..bcc1393 100644 --- a/docs/en/stream_server.md +++ b/docs/en/stream_server.md @@ -32,6 +32,10 @@ TeleFuser stream server supports two interaction modes, both available over WebR | **Server Push** | WebRTC (RTP) | Server → Client | Real-time preview, text-to-video streaming | | **Bidirectional** | WebRTC (RTP + DataChannel) | Client ↔ Server | Interactive generation, keyboard/camera control, speech-to-video | +`telefuser stream-serve` starts a stream-only app. It exposes `/v1/stream/*`, `/v1/stream/webrtc/*`, and +`/v1/service/*`; it does not expose request-response task routes, file-download routes, or OpenAI-compatible +`/v1/images` and `/v1/videos` routes. Use `telefuser serve` for batch task submission. + ### Server Push (WebRTC) ``` @@ -105,7 +109,7 @@ telefuser stream-serve [OPTIONS] |------|---------|-------------| | `--port`, `-p` | `8088` | Server port | | `--host` | `0.0.0.0` | Bind address | -| `--security-level` | `strict` | Pipeline validation level (`none`, `basic`, `strict`, `sandbox`) | +| `--security-level` | `strict` | Pipeline validation level (`none`, `basic`, `strict`, `sandbox`). `sandbox` is a best-effort restricted-load check, not runtime isolation. | | `--skip-validation` | `false` | Skip pipeline file security checks | ### Examples @@ -227,6 +231,9 @@ Server-push chunks should contain the following fields: | `/v1/stream/sessions/{session_id}` | DELETE | Both | Close session (pipeline + WebRTC) | | `/v1/stream/sessions/{session_id}/status` | GET | Both | Get session status | +The stream app also exposes service endpoints such as `/v1/service/health`, `/v1/service/ready`, +`/v1/service/metadata`, and `/v1/service/metrics/json`. + ### WebRTC: SDP Offer **POST** `/v1/stream/webrtc/offer` @@ -290,6 +297,10 @@ Response (`200 OK`): } ``` +**DELETE** `/v1/stream/sessions/{session_id}` closes the pipeline session and WebRTC session together when both +owners are present. If one side has already disappeared but the other closes successfully, the endpoint still returns +`200 OK`; pipeline-side close failures are logged as warnings instead of being silently ignored. + ### WebRTC: DataChannel Protocol (Bidirectional) In bidirectional mode, the client-created `"telefuser"` DataChannel carries JSON messages in both directions. @@ -757,7 +768,7 @@ Pipeline files are validated before loading. Use `--skip-validation` only for de | `none` | Disable validation | | `basic` | Static AST analysis | | `strict` | Static AST analysis plus import restrictions | -| `sandbox` | Full sandbox execution | +| `sandbox` | Strict checks plus a best-effort restricted-load validation step. This is not runtime isolation. | --- diff --git a/docs/zh/service.md b/docs/zh/service.md index b61a664..200f362 100644 --- a/docs/zh/service.md +++ b/docs/zh/service.md @@ -150,7 +150,7 @@ telefuser serve /path/to/pipeline --task i2v [选项] | `--cache-dir` | `-c` | string | `work_dirs/server_cache` | 缓存目录 | | `--parallelism` | `-g` | int | `1` | 并行工作进程数 | | `--num-replicas` | `-n` | int | `1` | 独立 Pipeline 副本数量（Pipeline Pool） | -| `--security-level` | | choice | `strict` | 安全级别: none/basic/strict/sandbox | +| `--security-level` | | choice | `strict` | 验证级别：none/basic/strict/sandbox。`sandbox` 是 best-effort 受限加载检查，不是运行时隔离。 | | `--skip-validation` | | flag | `False` | 跳过安全验证 | | `--validate-only` | | flag | `False` | 仅验证不启动 | @@ -186,7 +186,7 @@ telefuser validate /path/to/pipeline.py [选项] | 参数 | 默认值 | 描述 | |-----------|---------|-------------| | `pipeline_file` | **必需** | 管道 Python 文件路径 | -| `--level` | `strict` | 安全级别: none/basic/strict/sandbox | +| `--level` | `strict` | 验证级别：none/basic/strict/sandbox。`sandbox` 是 best-effort 受限加载检查，不是运行时隔离。 | | `--json` | `False` | 以 JSON 格式输出 | ```bash @@ -258,7 +258,17 @@ telefuser serve --help | `TELEFUSER_HOST` | 服务器主机 | `127.0.0.1` | | `TELEFUSER_PORT` | 服务器端口 | `8000` | | `TELEFUSER_RATE_LIMIT_ENABLED` | 启用速率限制 | `true` | -| `TELEFUSER_RATE_LIMIT_RPM` | 每分钟请求限制 | `60` | +| `TELEFUSER_RATE_LIMIT_REQUESTS_PER_MINUTE` | 每分钟请求限制 | `60` | +| `TELEFUSER_TRUST_FORWARDED_FOR` | 是否信任 `X-Forwarded-For` 作为限流身份。仅应在可信反向代理后启用。 | `false` | +| `TELEFUSER_ARTIFACT_STORAGE_BACKEND` | Artifact 后端。当前只实现 `local`。 | `local` | +| `TELEFUSER_ARTIFACT_LOCAL_ROOT` | 本地 artifact 根目录。未设置时使用 `TELEFUSER_CACHE_DIR`。 | 未设置 | +| `TELEFUSER_ARTIFACT_PERSISTENCE_MODE` | 本地 artifact 保留模式：`persistent` 或 `ephemeral`。 | `persistent` | +| `TELEFUSER_ARTIFACT_RETENTION_SECONDS` | 终态任务 artifact 保留时间。`0` 表示关闭 TTL 清理。 | `604800` | +| `TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS` | 临时 `.part` 文件保留时间。`0` 表示关闭临时文件清理。 | `3600` | +| `TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS` | 后台 artifact 清理周期。 | `3600` | +| `TELEFUSER_ARTIFACT_MAX_TOTAL_BYTES` | 本地 artifact 缓存总容量上限。`0` 表示关闭容量清理。 | `0` | +| `TELEFUSER_ARTIFACT_MAX_TASK_BYTES` | 单任务 artifact 容量上限，仅清理终态任务。`0` 表示关闭。 | `0` | +| `TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS` | 保护失败任务目录，不让自动 artifact 清理删除。 | `false` | ### 配置文件示例 @@ -269,9 +279,64 @@ TELEFUSER_SECURITY_LEVEL=STRICT TELEFUSER_PORT=8080 TELEFUSER_HOST=0.0.0.0 TELEFUSER_RATE_LIMIT_ENABLED=true -TELEFUSER_RATE_LIMIT_RPM=100 +TELEFUSER_RATE_LIMIT_REQUESTS_PER_MINUTE=100 ``` +### 运行边界 + +请求-响应服务默认是本地、单进程语义： + +- 任务状态、队列状态、取消状态和限流 bucket 保存在进程内存中。多实例之间不共享，服务重启后也不会恢复。 +- 文件输入、生成输出和临时 `.part` 文件由本地 artifact store 管理，根目录来自配置的 cache/artifact root。 +- `artifact_storage_backend=s3` 只是后续远端后端的边界声明，当前尚未实现。显式配置为 `s3` 会在启动阶段报错，不会静默退回本地存储。 +- `telefuser serve` 只暴露请求-响应路由：`/v1/tasks/*`、`/v1/files/*`、`/v1/images/*`、`/v1/videos/*` 和 `/v1/service/*`。 +- `telefuser stream-serve` 只暴露流式路由：`/v1/stream/*`、`/v1/stream/webrtc/*` 和 `/v1/service/*`。它不会暴露任务、文件下载或 OpenAI 兼容的请求-响应路由。 + +### Artifact 存储与清理 + +Artifact 是本地文件。设置 `TELEFUSER_ARTIFACT_LOCAL_ROOT` 时使用该目录，否则使用 `TELEFUSER_CACHE_DIR`。 +上传和远程下载会先写入临时 `.part` 文件，成功完成后再原子重命名。 + +当有 task id 时，输出采用 task-scoped 目录： + +```text +/tasks//inputs/ +/tasks//outputs/ +/tasks//tmp/ +``` + +服务只接受解析后仍位于 artifact root 内的路径。绝对输出路径和 `..` 路径逃逸会被拒绝。 +下载路由接受校验后的输出路径和本地 artifact id。本地 artifact id 使用 +`local:tasks//outputs//` 格式，只表示本地 artifact 后端中的可下载输出； +它不是远端对象存储 id。 + +清理是 best-effort 行为，并且不会删除 active task（`pending`、`processing`、`streaming`）。 +终态任务（`completed`、`failed`、`cancelled`）可能被 TTL、总容量上限或单任务容量上限清理。 +临时 `.part` 文件按 `TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS` 独立清理。 + +`TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent` 会保留终态任务 artifact，直到 TTL 或容量清理删除。 +`ephemeral` 会在下一次清理时删除终态任务目录。设置 +`TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true` 后，失败任务目录会跳过自动清理，便于人工检查部分输出。 + +### 限流和代理 Header + +限流状态保存在进程内存中，并按客户端身份计数。默认使用直接客户端地址，不信任 `X-Forwarded-For`， +因为普通客户端可以伪造这个 header。只有当 TeleFuser 部署在会清洗转发 header 的可信反向代理之后， +才应设置 `TELEFUSER_TRUST_FORWARDED_FOR=true`。 + +默认限流路径覆盖昂贵生成入口、artifact 下载和 stream 协商： + +```text +/v1/tasks/create +/v1/tasks/form +/v1/images +/v1/videos +/v1/files/download +/v1/stream +``` + +Liveness/readiness 端点不会被默认限流，便于基础设施探针访问。 + ### Pipeline 契约与参数定义 服务端不会只靠 Python 函数签名推断 pipeline 能力，而是优先从 example 脚本中加载显式的 pipeline contract。 @@ -461,9 +526,9 @@ http://localhost:8000/v1 | POST | `/tasks/create` | 创建新的生成任务 | | POST | `/tasks/form` | 带文件上传创建任务 | | GET | `/tasks/{task_id}/status` | 获取任务状态 | -| DELETE | `/tasks/{task_id}` | 取消任务 | +| DELETE | `/tasks/{task_id}` | 请求协作式取消任务 | | GET | `/tasks/queue/status` | 获取队列状态 | -| GET | `/files/download/{file_id}` | 下载输出文件 | +| GET | `/files/download/{file_id}` | 通过校验后的路径或本地 artifact id 下载输出文件 | | GET | `/service/health` | 健康检查 | | GET | `/service/status` | 服务状态 | | GET | `/service/metadata` | 服务元数据 | @@ -812,7 +877,17 @@ TeleFuser 提供 **OpenAI 兼容**的 REST API，让您可以使用标准的 Ope "data": [ { "url": "http://localhost:8000/v1/images/task_xxx/content", - "revised_prompt": "一只美丽的猫" + "revised_prompt": "一只美丽的猫", + "file_path": "/cache/tasks/task_xxx/outputs/images/result.png", + "artifact_id": "local:tasks/task_xxx/outputs/images/result.png", + "artifact_metadata": { + "backend": "local", + "relative_path": "tasks/task_xxx/outputs/images/result.png", + "task_id": "task_xxx", + "media_type": "image", + "filename": "result.png", + "size_bytes": 123456 + } } ], "peak_memory_mb": 4096.5, @@ -873,7 +948,7 @@ response.data[0].save("sunset.png") "size": "1024x576", // 可选：视频尺寸 "seed": 1024, // 可选：随机种子 "negative_prompt": "模糊", // 可选：反向提示词 - "output_path": "/path/to/output.mp4" // 可选：自定义输出路径 + "output_path": "custom-output.mp4" // 可选：相对自定义输出路径 } ``` @@ -889,7 +964,31 @@ response.data[0].save("sunset.png") "created_at": 1699000000, "size": "1024x576", "seconds": "5", - "url": null + "url": null, + "file_path": null, + "artifact_id": null, + "artifact_metadata": null +} +``` + +完成后的视频响应会包含内容 URL 和本地 artifact metadata： + +```json +{ + "id": "vid_xxx", + "status": "completed", + "progress": 100, + "url": "http://localhost:8000/v1/videos/vid_xxx/content", + "file_path": "/cache/tasks/vid_xxx/outputs/videos/result.mp4", + "artifact_id": "local:tasks/vid_xxx/outputs/videos/result.mp4", + "artifact_metadata": { + "backend": "local", + "relative_path": "tasks/vid_xxx/outputs/videos/result.mp4", + "task_id": "vid_xxx", + "media_type": "video", + "filename": "result.mp4", + "size_bytes": 987654 + } } ``` diff --git a/docs/zh/stream_server.md b/docs/zh/stream_server.md index 1fc647e..3469b50 100644 --- a/docs/zh/stream_server.md +++ b/docs/zh/stream_server.md @@ -32,6 +32,10 @@ TeleFuser 流式服务器支持两种交互模式，均可通过 WebRTC 使用 | **服务端推送** | WebRTC (RTP) | 服务端 → 客户端 | 实时预览、文生视频流式传输 | | **双向交互** | WebRTC (RTP + DataChannel) | 客户端 ↔ 服务端 | 交互式生成、键盘/摄像头控制、语音生成视频 | +`telefuser stream-serve` 启动的是 stream-only app。它只暴露 `/v1/stream/*`、`/v1/stream/webrtc/*` 和 +`/v1/service/*`；不会暴露请求-响应任务路由、文件下载路由，也不会暴露 OpenAI 兼容的 `/v1/images` +和 `/v1/videos` 路由。批量任务提交请使用 `telefuser serve`。 + ### 服务端推送（WebRTC） ``` @@ -105,7 +109,7 @@ telefuser stream-serve [选项] |------|--------|------| | `--port`, `-p` | `8088` | 服务端口 | | `--host` | `0.0.0.0` | 绑定地址 | -| `--security-level` | `strict` | 管线验证级别（`none`、`basic`、`strict`、`sandbox`） | +| `--security-level` | `strict` | 管线验证级别（`none`、`basic`、`strict`、`sandbox`）。`sandbox` 是 best-effort 受限加载检查，不是运行时隔离。 | | `--skip-validation` | `false` | 跳过管线文件安全检查 | ### 示例 @@ -227,6 +231,9 @@ def get_service() -> MyBidirectionalService: | `/v1/stream/sessions/{session_id}` | DELETE | 两种模式 | 关闭会话（管线 + WebRTC） | | `/v1/stream/sessions/{session_id}/status` | GET | 两种模式 | 获取会话状态 | +stream app 也会暴露 `/v1/service/health`、`/v1/service/ready`、`/v1/service/metadata` +和 `/v1/service/metrics/json` 等服务端点。 + ### WebRTC: SDP Offer **POST** `/v1/stream/webrtc/offer` @@ -290,6 +297,10 @@ def get_service() -> MyBidirectionalService: } ``` +**DELETE** `/v1/stream/sessions/{session_id}` 会在管线会话和 WebRTC 会话都存在时一起关闭两侧。 +如果其中一侧已经不存在，但另一侧关闭成功，端点仍返回 `200 OK`；管线侧关闭失败会记录 warning， +不会再静默忽略。 + ### WebRTC: DataChannel 协议（双向交互） 在双向交互模式下，客户端创建的 `"telefuser"` DataChannel 双向传输 JSON 消息。 @@ -756,7 +767,7 @@ nvidia-smi | `none` | 禁用验证 | | `basic` | 静态 AST 分析 | | `strict` | 静态 AST 分析加导入限制 | -| `sandbox` | 完整沙箱执行 | +| `sandbox` | strict 检查加 best-effort 受限加载验证步骤。它不是运行时隔离。 | --- diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index 2d41e15..b480203 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -345,6 +345,8 @@ def initialize_services( ssl_cert_path=self.server_config.ssl_cert_path, artifact_retention_seconds=self.server_config.artifact_retention_seconds, artifact_tmp_retention_seconds=self.server_config.artifact_tmp_retention_seconds, + artifact_persistence_mode=self.server_config.artifact_persistence_mode, + artifact_preserve_failed_outputs=self.server_config.artifact_preserve_failed_outputs, artifact_max_total_bytes=self.server_config.artifact_max_total_bytes, artifact_max_task_bytes=self.server_config.artifact_max_task_bytes, ) diff --git a/telefuser/service/api/openai/adapter.py b/telefuser/service/api/openai/adapter.py index 9289c8a..4b4fc4e 100644 --- a/telefuser/service/api/openai/adapter.py +++ b/telefuser/service/api/openai/adapter.py @@ -240,9 +240,11 @@ def to_image_response( task_id: str = "", peak_memory_mb: float | None = None, inference_time_s: float | None = None, + artifact_metadata: dict[str, Any] | None = None, ) -> ImageResponse: """Convert TeleFuser output to OpenAI ImageResponse.""" data_list: List[ImageResponseData] = [] + artifact_id = artifact_metadata.get("artifact_id") if artifact_metadata else None if response_format == "b64_json": try: @@ -256,6 +258,8 @@ def to_image_response( b64_json=b64_string, revised_prompt=prompt, file_path=output_path, + artifact_id=artifact_id, + artifact_metadata=artifact_metadata, ) ) except Exception as e: @@ -272,6 +276,8 @@ def to_image_response( url=url, revised_prompt=prompt, file_path=output_path, + artifact_id=artifact_id, + artifact_metadata=artifact_metadata, ) ) @@ -295,6 +301,7 @@ def to_video_response( peak_memory_mb: float | None = None, inference_time_s: float | None = None, error: Dict[str, Any] | None = None, + artifact_metadata: dict[str, Any] | None = None, ) -> VideoResponse: """Convert TeleFuser task info to OpenAI VideoResponse.""" status_map = { @@ -318,8 +325,49 @@ def to_video_response( peak_memory_mb=peak_memory_mb, inference_time_s=inference_time_s, error=error, + artifact_id=artifact_metadata.get("artifact_id") if artifact_metadata else None, + artifact_metadata=artifact_metadata, ) + @staticmethod + def progress_for_task_status(status: str) -> int: + """Return a coarse OpenAI-compatible progress value for a task status.""" + if status == "completed": + return 100 + if status == "processing": + return 50 + return 0 + + @staticmethod + def to_video_response_from_task( + task_id: str, + task_status: dict[str, Any], + message: Any | None, + *, + url: str | None = None, + artifact_metadata: dict[str, Any] | None = None, + ) -> VideoResponse: + """Convert a TaskManager status dictionary and optional task message to VideoResponse.""" + status = task_status.get("status", "pending") + response = OpenAIResponseAdapter.to_video_response( + task_id=task_id, + status=status, + prompt=getattr(message, "prompt", ""), + size=getattr(message, "resolution", ""), + seconds=getattr(message, "target_video_length", 4), + model=getattr(message, "model", None) or "wan-video", + output_path=task_status.get("output_path"), + url=url, + progress=OpenAIResponseAdapter.progress_for_task_status(status), + peak_memory_mb=task_status.get("peak_memory_mb"), + inference_time_s=task_status.get("inference_time_s"), + artifact_metadata=artifact_metadata, + ) + + if status == "failed": + response.error = {"message": task_status.get("error", "Unknown error")} + return response + @staticmethod def to_video_list_response( videos: List[VideoResponse], diff --git a/telefuser/service/api/openai/image_routes.py b/telefuser/service/api/openai/image_routes.py index 9517a3d..49cff03 100644 --- a/telefuser/service/api/openai/image_routes.py +++ b/telefuser/service/api/openai/image_routes.py @@ -126,18 +126,28 @@ async def create_image_generation(self, request: ImageGenerationsRequest) -> Ima if not output_path: raise HTTPException(status_code=500, detail="Generation completed but no output path found") + resolved_output_path = self.api.task_app_service.resolve_task_output_path( + output_path, + media_type=MediaType.IMAGE, + ) + artifact_metadata = self.api.task_app_service.get_output_metadata( + task_id, + output_path=str(resolved_output_path), + media_type=MediaType.IMAGE, + ) peak_memory_mb = result.get("peak_memory_mb") inference_time_s = result.get("inference_time_s") - base_url = self._get_base_url() + base_url = self.api.task_app_service.get_base_url() response = OpenAIResponseAdapter.to_image_response( - output_path=output_path, + output_path=str(resolved_output_path), prompt=request.prompt, response_format=request.response_format or "url", base_url=base_url, task_id=task_id, peak_memory_mb=peak_memory_mb, inference_time_s=inference_time_s, + artifact_metadata=artifact_metadata, ) logger.info(f"Image generation completed: {task_id}") @@ -230,13 +240,23 @@ async def create_image_edit( if not output_path: raise HTTPException(status_code=500, detail="No output from generation") - base_url = self._get_base_url() + resolved_output_path = self.api.task_app_service.resolve_task_output_path( + output_path, + media_type=MediaType.IMAGE, + ) + artifact_metadata = self.api.task_app_service.get_output_metadata( + task_id, + output_path=str(resolved_output_path), + media_type=MediaType.IMAGE, + ) + base_url = self.api.task_app_service.get_base_url() response = OpenAIResponseAdapter.to_image_response( - output_path=output_path, + output_path=str(resolved_output_path), prompt=prompt, response_format=response_format or "url", base_url=base_url, task_id=task_id, + artifact_metadata=artifact_metadata, ) logger.info(f"Image edit completed: {task_id}") diff --git a/telefuser/service/api/openai/protocol.py b/telefuser/service/api/openai/protocol.py index c7e17c8..02283c6 100644 --- a/telefuser/service/api/openai/protocol.py +++ b/telefuser/service/api/openai/protocol.py @@ -125,6 +125,8 @@ class ImageResponseData(BaseModel): None, description="The prompt that was used to generate the image, if any revision was made." ) file_path: str | None = Field(None, description="Local file path of the generated image (TeleFuser ext).") + artifact_id: str | None = Field(None, description="Stable output artifact id (TeleFuser ext).") + artifact_metadata: Dict[str, Any] | None = Field(None, description="Output artifact metadata (TeleFuser ext).") class ImageResponse(BaseModel): @@ -195,6 +197,8 @@ class VideoResponse(BaseModel): expires_at: int | None = Field(None, description="The Unix timestamp when the video will be deleted.") error: Dict[str, Any] | None = Field(None, description="Error information if generation failed.") file_path: str | None = Field(None, description="Local file path of the generated video.") + artifact_id: str | None = Field(None, description="Stable output artifact id (TeleFuser ext).") + artifact_metadata: Dict[str, Any] | None = Field(None, description="Output artifact metadata (TeleFuser ext).") # TeleFuser extensions peak_memory_mb: float | None = Field(None, description="Peak memory usage in MB during generation.") diff --git a/telefuser/service/api/openai/video_routes.py b/telefuser/service/api/openai/video_routes.py index 69b4a0f..1879d21 100644 --- a/telefuser/service/api/openai/video_routes.py +++ b/telefuser/service/api/openai/video_routes.py @@ -211,14 +211,15 @@ async def list_videos(self, after: str | None = None, limit: int = 20, order: st # Convert to VideoResponse objects videos: List[VideoResponse] = [] for task_id, task_data in video_tasks: - video = OpenAIResponseAdapter.to_video_response( + task_info = self.task_manager.get_task(task_id) + message = task_info.message if task_info else None + status = task_data.get("status", TaskStatus.PENDING.value) + video = OpenAIResponseAdapter.to_video_response_from_task( task_id=task_id, - status=task_data.get("status", TaskStatus.PENDING.value), - prompt=task_data.get("prompt", ""), - size=task_data.get("resolution", ""), - seconds=task_data.get("target_video_length", 4), - model=task_data.get("model", "wan-video"), - output_path=task_data.get("output_path"), + task_status=task_data, + message=message, + url=self._content_url_for_completed_video(task_id, status), + artifact_metadata=self._artifact_metadata_for_video(task_id, task_data.get("output_path")), ) videos.append(video) @@ -246,30 +247,15 @@ async def retrieve_video(self, video_id: str) -> VideoResponse: task_info = self.task_manager.get_task(video_id) message = task_info.message if task_info else None - # Calculate progress (simplified) status = task_status.get("status", TaskStatus.PENDING.value) - progress = 0 - if status == TaskStatus.COMPLETED.value: - progress = 100 - elif status == TaskStatus.PROCESSING.value: - progress = 50 - - response = OpenAIResponseAdapter.to_video_response( + return OpenAIResponseAdapter.to_video_response_from_task( task_id=video_id, - status=status, - prompt=message.prompt if message else "", - size=message.resolution if message else "", - seconds=message.target_video_length if message else 4, - model=getattr(message, "model", None) or "wan-video", - output_path=task_status.get("output_path"), - progress=progress, + task_status=task_status, + message=message, + url=self._content_url_for_completed_video(video_id, status), + artifact_metadata=self._artifact_metadata_for_video(video_id, task_status.get("output_path")), ) - if status == TaskStatus.FAILED.value: - response.error = {"message": task_status.get("error", "Unknown error")} - - return response - async def delete_video(self, video_id: str) -> VideoResponse: """Cancel or delete a video generation task.""" cancel_result = self.api.task_app_service.cancel_task(video_id) @@ -297,6 +283,22 @@ async def _ensure_processing(self) -> None: """Ensure the task processor is running.""" await self.api.ensure_task_processor_running() + def _content_url_for_completed_video(self, video_id: str, status: str) -> str | None: + """Return a video content URL only when the output is ready.""" + if status != TaskStatus.COMPLETED.value: + return None + return self.api.task_app_service.get_openai_content_url(video_id, media_type=MediaType.VIDEO) + + def _artifact_metadata_for_video(self, video_id: str, output_path: str | None) -> dict | None: + """Return artifact metadata for a video output when it exists.""" + if not output_path: + return None + return self.api.task_app_service.get_output_metadata( + video_id, + output_path=output_path, + media_type=MediaType.VIDEO, + ) + async def _save_uploaded_file(self, file: UploadFile, prefix: str = "upload") -> str: """Save an uploaded file to disk.""" return await self.api.task_app_service.save_upload_file( diff --git a/telefuser/service/api/task_application_service.py b/telefuser/service/api/task_application_service.py index bb5550e..6d53c82 100644 --- a/telefuser/service/api/task_application_service.py +++ b/telefuser/service/api/task_application_service.py @@ -145,6 +145,55 @@ def get_output_response( response_media_type = "image/png" if str(media_type) == MediaType.IMAGE.value else "video/mp4" return FileResponse(path=str(path), media_type=response_media_type, filename=path.name) + def resolve_task_output_path(self, output_path: str, *, media_type: MediaType | str) -> Path: + """Resolve a task output path through the configured file service.""" + file_service = getattr(self.api, "file_service", None) + if not file_service: + return Path(output_path) + return self.resolve_output_file(file_service, output_path, media_type=media_type) + + def get_output_metadata( + self, + task_id: str, + *, + output_path: str | None = None, + media_type: MediaType | str, + ) -> dict[str, Any] | None: + """Return artifact metadata for a task output when the file service supports it.""" + file_service = getattr(self.api, "file_service", None) + if not file_service: + return None + + if output_path is None: + task_info = self.api.task_manager.get_task(task_id) + task_status = self.api.task_manager.get_task_status(task_id) or {} + output_path = getattr(task_info, "output_path", None) or task_status.get("output_path") + if not output_path: + return None + + metadata_builder = self.get_declared_method(file_service, "artifact_metadata") + if not callable(metadata_builder): + return None + + try: + return metadata_builder(output_path, task_id=task_id, media_type=media_type) + except ValueError: + raise HTTPException(status_code=403, detail="Access to this file is not allowed") + + def get_base_url(self) -> str: + """Return the configured HTTP base URL for local content routes.""" + server_config = getattr(self.api, "server_config", None) + if server_config: + host = getattr(server_config, "host", "localhost") + port = getattr(server_config, "port", 8000) + return f"http://{host}:{port}" + return "http://localhost:8000" + + def get_openai_content_url(self, task_id: str, *, media_type: MediaType | str) -> str: + """Return the OpenAI-compatible content URL for a task output.""" + resource = "images" if str(media_type) == MediaType.IMAGE.value else "videos" + return f"{self.get_base_url()}/v1/{resource}/{task_id}/content" + def cancel_task(self, task_id: str) -> dict[str, Any]: """Cancel a task and distinguish accepted, terminal, and missing outcomes.""" task_manager = getattr(self.api, "task_manager", None) diff --git a/telefuser/service/core/artifact_store.py b/telefuser/service/core/artifact_store.py index 6488aba..3a83a1a 100644 --- a/telefuser/service/core/artifact_store.py +++ b/telefuser/service/core/artifact_store.py @@ -4,13 +4,15 @@ import shutil from collections.abc import Mapping -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from typing import Any from telefuser.service_types import MediaType from telefuser.utils.logging import logger +LOCAL_ARTIFACT_ID_PREFIX = "local:" + def _to_timestamp(value: datetime | float | int | None) -> float | None: if value is None: @@ -70,6 +72,10 @@ def _resolve_under(self, root: Path, relative_path: str | Path) -> Path: raise ValueError("Path escapes the service cache directory") return resolved_path + @staticmethod + def _utc_timestamp(value: float) -> str: + return datetime.fromtimestamp(value, tz=timezone.utc).isoformat().replace("+00:00", "Z") + def _validate_task_id(self, task_id: str) -> str: if not task_id or any(char in task_id for char in ("/", "\\")) or task_id in {".", ".."}: raise ValueError("Invalid task_id for artifact path") @@ -105,11 +111,19 @@ def output_path( if task_id: return self._resolve_under(self.task_output_dir(task_id, media_type), output_path) - root = self.legacy_output_image_dir if self._media_dir_name(media_type) == "images" else self.legacy_output_video_dir + root = ( + self.legacy_output_image_dir + if self._media_dir_name(media_type) == "images" + else self.legacy_output_video_dir + ) return self._resolve_under(root, output_path) def resolve_output_file(self, file_path: str | Path) -> Path: """Resolve a downloadable output artifact under legacy or task output roots.""" + raw_path = str(file_path) + if raw_path.startswith(LOCAL_ARTIFACT_ID_PREFIX): + return self.resolve_artifact_id(raw_path) + path = Path(file_path) if path.is_absolute(): resolved_path = path.resolve() @@ -117,6 +131,10 @@ def resolve_output_file(self, file_path: str | Path) -> Path: return resolved_path raise ValueError("Access to this file is not allowed") + root_relative_candidate = self._resolve_under(self.root, path) + if root_relative_candidate.exists() and self._is_allowed_output_path(root_relative_candidate.resolve()): + return root_relative_candidate + for root in (self.legacy_output_video_dir, self.legacy_output_image_dir, self.legacy_output_dir): candidate = self._resolve_under(root, path) if candidate.exists(): @@ -132,6 +150,58 @@ def resolve_output_file(self, file_path: str | Path) -> Path: return self._resolve_under(self.legacy_output_dir, path) + def artifact_id_for_path(self, file_path: str | Path) -> str: + """Return a stable local artifact id for a path under the artifact root.""" + raw_path = str(file_path) + if raw_path.startswith(LOCAL_ARTIFACT_ID_PREFIX): + resolved_path = self.resolve_artifact_id(raw_path) + else: + resolved_path = Path(file_path).resolve() + + if not self._is_relative_to(resolved_path, self.root): + raise ValueError("Artifact path is outside the service cache directory") + if not self._is_allowed_output_path(resolved_path): + raise ValueError("Artifact path is not a downloadable output") + return f"{LOCAL_ARTIFACT_ID_PREFIX}{resolved_path.relative_to(self.root).as_posix()}" + + def resolve_artifact_id(self, artifact_id: str) -> Path: + """Resolve a local artifact id back to an allowed output artifact path.""" + if not artifact_id.startswith(LOCAL_ARTIFACT_ID_PREFIX): + raise ValueError("Unsupported artifact backend") + + relative_path = artifact_id[len(LOCAL_ARTIFACT_ID_PREFIX) :] + if not relative_path: + raise ValueError("Artifact id is empty") + + resolved_path = self._resolve_under(self.root, relative_path) + if not self._is_allowed_output_path(resolved_path): + raise ValueError("Access to this artifact is not allowed") + return resolved_path + + def artifact_metadata( + self, + file_path: str | Path, + *, + task_id: str | None = None, + media_type: MediaType | str | None = None, + ) -> dict[str, Any]: + """Return local artifact metadata without exposing remote-storage semantics.""" + resolved_path = self.resolve_output_file(file_path) + stat = resolved_path.stat() if resolved_path.exists() else None + media_value = media_type.value if isinstance(media_type, MediaType) else media_type + + return { + "artifact_id": self.artifact_id_for_path(resolved_path), + "backend": "local", + "relative_path": resolved_path.relative_to(self.root).as_posix(), + "task_id": task_id, + "media_type": media_value, + "filename": resolved_path.name, + "size_bytes": stat.st_size if stat is not None else None, + "created_at": self._utc_timestamp(stat.st_ctime) if stat is not None else None, + "modified_at": self._utc_timestamp(stat.st_mtime) if stat is not None else None, + } + def _is_allowed_output_path(self, resolved_path: Path) -> bool: if self._is_relative_to(resolved_path, self.legacy_output_dir.resolve()): return True @@ -147,8 +217,11 @@ def cleanup( *, active_task_ids: set[str] | frozenset[str], terminal_task_end_times: Mapping[str, datetime | float | int | None], + terminal_task_statuses: Mapping[str, str] | None = None, retention_seconds: int, tmp_retention_seconds: int, + persistence_mode: str = "persistent", + preserve_failed_outputs: bool = False, max_total_bytes: int = 0, max_task_bytes: int = 0, now: datetime | float | int | None = None, @@ -158,12 +231,28 @@ def cleanup( removed_task_ids: list[str] = [] removed_tmp_files = 0 errors: list[str] = [] + protected_task_ids = self._protected_terminal_task_ids( + active_task_ids=active_task_ids, + terminal_task_statuses=terminal_task_statuses, + preserve_failed_outputs=preserve_failed_outputs, + ) removed_tmp_files += self._cleanup_tmp_files(tmp_retention_seconds=tmp_retention_seconds, now_ts=now_ts) - if retention_seconds > 0: + if persistence_mode not in {"persistent", "ephemeral"}: + raise ValueError(f"Unsupported artifact persistence mode: {persistence_mode}") + + if persistence_mode == "ephemeral": + removed_task_ids.extend( + self._cleanup_terminal_tasks( + active_task_ids=active_task_ids, + terminal_task_end_times=terminal_task_end_times, + protected_task_ids=protected_task_ids, + ) + ) + elif retention_seconds > 0: for task_id, end_time in terminal_task_end_times.items(): - if task_id in active_task_ids: + if task_id in active_task_ids or task_id in protected_task_ids: continue end_ts = _to_timestamp(end_time) if end_ts is None or now_ts - end_ts < retention_seconds: @@ -186,6 +275,7 @@ def cleanup( terminal_task_end_times=terminal_task_end_times, max_task_bytes=max_task_bytes, already_removed=set(removed_task_ids), + protected_task_ids=protected_task_ids, ) ) @@ -195,6 +285,7 @@ def cleanup( active_task_ids=active_task_ids, terminal_task_end_times=terminal_task_end_times, max_total_bytes=max_total_bytes, + protected_task_ids=protected_task_ids, ) ) @@ -221,12 +312,49 @@ def _cleanup_tmp_files(self, *, tmp_retention_seconds: int, now_ts: float) -> in logger.warning(f"Failed to remove temporary artifact {part_file}: {exc}") return removed + @staticmethod + def _protected_terminal_task_ids( + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_statuses: Mapping[str, str] | None, + preserve_failed_outputs: bool, + ) -> set[str]: + if not preserve_failed_outputs or not terminal_task_statuses: + return set() + return { + task_id + for task_id, status in terminal_task_statuses.items() + if task_id not in active_task_ids and str(status).lower() == "failed" + } + + def _cleanup_terminal_tasks( + self, + *, + active_task_ids: set[str] | frozenset[str], + terminal_task_end_times: Mapping[str, datetime | float | int | None], + protected_task_ids: set[str], + ) -> list[str]: + removed: list[str] = [] + for task_id in terminal_task_end_times: + if task_id in active_task_ids or task_id in protected_task_ids: + continue + task_dir = self.task_root(task_id) + if not task_dir.exists(): + continue + try: + shutil.rmtree(task_dir) + removed.append(task_id) + except Exception as exc: + logger.warning(f"Failed to remove artifact directory for task {task_id}: {exc}") + return removed + def _cleanup_capacity( self, *, active_task_ids: set[str] | frozenset[str], terminal_task_end_times: Mapping[str, datetime | float | int | None], max_total_bytes: int, + protected_task_ids: set[str], ) -> list[str]: current_size = self._tree_size(self.root) if current_size <= max_total_bytes: @@ -235,7 +363,7 @@ def _cleanup_capacity( candidates = [ (task_id, _to_timestamp(end_time) or 0.0, self.task_root(task_id)) for task_id, end_time in terminal_task_end_times.items() - if task_id not in active_task_ids and self.task_root(task_id).exists() + if task_id not in active_task_ids and task_id not in protected_task_ids and self.task_root(task_id).exists() ] candidates.sort(key=lambda item: item[1]) @@ -259,10 +387,11 @@ def _cleanup_oversized_tasks( terminal_task_end_times: Mapping[str, datetime | float | int | None], max_task_bytes: int, already_removed: set[str], + protected_task_ids: set[str], ) -> list[str]: removed: list[str] = [] for task_id in terminal_task_end_times: - if task_id in active_task_ids or task_id in already_removed: + if task_id in active_task_ids or task_id in already_removed or task_id in protected_task_ids: continue task_dir = self.task_root(task_id) diff --git a/telefuser/service/core/config.py b/telefuser/service/core/config.py index 09feb78..005adde 100644 --- a/telefuser/service/core/config.py +++ b/telefuser/service/core/config.py @@ -58,6 +58,13 @@ class ServerConfig(BaseSettings): default=None, description="Local artifact root. Defaults to cache_dir when unset.", ) + artifact_persistence_mode: Literal["persistent", "ephemeral"] = Field( + default="persistent", + description=( + "Local artifact retention mode. Persistent keeps terminal artifacts until retention/capacity cleanup; " + "ephemeral removes terminal artifacts on the next cleanup pass." + ), + ) enable_latent_cache: bool | None = Field( default=None, @@ -130,7 +137,7 @@ class ServerConfig(BaseSettings): ) artifact_preserve_failed_outputs: bool = Field( default=False, - description="Whether failed task output directories should be preserved until normal retention expiry.", + description="Whether failed task output directories should be preserved until manual cleanup.", ) # Rate limiting settings diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index e55e9f8..4f7a598 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -111,6 +111,8 @@ def initialize_file_service(self, cache_dir: Path | None = None) -> FileService: ssl_cert_path=getattr(self.config, "ssl_cert_path", None), artifact_retention_seconds=self.config.artifact_retention_seconds, artifact_tmp_retention_seconds=self.config.artifact_tmp_retention_seconds, + artifact_persistence_mode=self.config.artifact_persistence_mode, + artifact_preserve_failed_outputs=self.config.artifact_preserve_failed_outputs, artifact_max_total_bytes=self.config.artifact_max_total_bytes, artifact_max_task_bytes=self.config.artifact_max_task_bytes, ) diff --git a/telefuser/service/core/file_service.py b/telefuser/service/core/file_service.py index dbd6c1d..71fdb87 100644 --- a/telefuser/service/core/file_service.py +++ b/telefuser/service/core/file_service.py @@ -34,6 +34,8 @@ def __init__( ssl_cert_path: str | None = None, artifact_retention_seconds: int = 7 * 24 * 60 * 60, artifact_tmp_retention_seconds: int = 60 * 60, + artifact_persistence_mode: str = "persistent", + artifact_preserve_failed_outputs: bool = False, artifact_max_total_bytes: int = 0, artifact_max_task_bytes: int = 0, ) -> None: @@ -51,6 +53,8 @@ def __init__( self.ssl_cert_path = ssl_cert_path self.artifact_retention_seconds = artifact_retention_seconds self.artifact_tmp_retention_seconds = artifact_tmp_retention_seconds + self.artifact_persistence_mode = artifact_persistence_mode + self.artifact_preserve_failed_outputs = artifact_preserve_failed_outputs self.artifact_max_total_bytes = artifact_max_total_bytes self.artifact_max_task_bytes = artifact_max_task_bytes @@ -343,19 +347,41 @@ def resolve_output_file(self, file_path: str | Path) -> Path: """Resolve a downloadable output file within the configured cache root.""" return self.artifact_store.resolve_output_file(file_path) + def artifact_id_for_path(self, file_path: str | Path) -> str: + """Return a stable local artifact id for a file under the artifact root.""" + return self.artifact_store.artifact_id_for_path(file_path) + + def resolve_artifact_id(self, artifact_id: str) -> Path: + """Resolve a local artifact id to an output file path.""" + return self.artifact_store.resolve_artifact_id(artifact_id) + + def artifact_metadata( + self, + file_path: str | Path, + *, + task_id: str | None = None, + media_type: MediaType | str | None = None, + ) -> dict[str, Any]: + """Return local artifact metadata for an output file.""" + return self.artifact_store.artifact_metadata(file_path, task_id=task_id, media_type=media_type) + def cleanup_artifacts( self, *, active_task_ids: set[str] | frozenset[str], terminal_task_end_times: Mapping[str, datetime | float | int | None], + terminal_task_statuses: Mapping[str, str] | None = None, now: datetime | float | int | None = None, ) -> dict[str, Any]: """Clean expired local artifacts according to configured retention settings.""" return self.artifact_store.cleanup( active_task_ids=active_task_ids, terminal_task_end_times=terminal_task_end_times, + terminal_task_statuses=terminal_task_statuses, retention_seconds=self.artifact_retention_seconds, tmp_retention_seconds=self.artifact_tmp_retention_seconds, + persistence_mode=self.artifact_persistence_mode, + preserve_failed_outputs=self.artifact_preserve_failed_outputs, max_total_bytes=self.artifact_max_total_bytes, max_task_bytes=self.artifact_max_task_bytes, now=now, diff --git a/telefuser/service/core/pipeline_service.py b/telefuser/service/core/pipeline_service.py index f928660..ca7e80f 100644 --- a/telefuser/service/core/pipeline_service.py +++ b/telefuser/service/core/pipeline_service.py @@ -38,7 +38,7 @@ class PipelineService: - Static AST analysis for dangerous operations - Import restriction and verification - Content pattern matching - - Optional sandboxed execution + - Optional best-effort restricted-load validation - Configurable security levels """ diff --git a/telefuser/service/core/replica_worker.py b/telefuser/service/core/replica_worker.py index a6e4525..abb7de6 100644 --- a/telefuser/service/core/replica_worker.py +++ b/telefuser/service/core/replica_worker.py @@ -60,8 +60,8 @@ def _replica_main( loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - from telefuser.service.core.pipeline_service import PipelineService from telefuser.service.core.config import ServerConfig + from telefuser.service.core.pipeline_service import PipelineService from telefuser.service.security.security_validator import SecurityLevel from telefuser.service_types import TaskType from telefuser.utils.logging import logger diff --git a/telefuser/service/core/task_manager.py b/telefuser/service/core/task_manager.py index 0d8e585..4163e5b 100644 --- a/telefuser/service/core/task_manager.py +++ b/telefuser/service/core/task_manager.py @@ -14,7 +14,6 @@ from telefuser.service_types import TaskStatus from telefuser.utils.logging import logger - _ACTIVE_STATUSES = frozenset({TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.STREAMING}) @@ -237,9 +236,15 @@ def get_artifact_cleanup_snapshot(self) -> dict[str, Any]: for task_id, task in self._tasks.items() if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED) } + terminal_task_statuses = { + task_id: task.status.value + for task_id, task in self._tasks.items() + if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED) + } return { "active_task_ids": active_task_ids, "terminal_task_end_times": terminal_task_end_times, + "terminal_task_statuses": terminal_task_statuses, } def get_pending_task_count(self) -> int: diff --git a/telefuser/service/security/security_validator.py b/telefuser/service/security/security_validator.py index 62d2967..1385387 100644 --- a/telefuser/service/security/security_validator.py +++ b/telefuser/service/security/security_validator.py @@ -24,7 +24,7 @@ class SecurityLevel(Enum): NONE = auto() # No validation BASIC = auto() # Static AST analysis only STRICT = auto() # AST + import restriction - SANDBOX = auto() # Full sandbox execution + SANDBOX = auto() # Strict validation plus best-effort restricted load check class SecurityError(Exception): @@ -450,13 +450,13 @@ def visit_Constant(self, node: ast.Constant) -> None: class SandboxedLoader: - """Restricted Python loader that executes code in a controlled environment. + """Best-effort restricted loader for pipeline validation. WARNING: This is NOT a complete sandbox and should not be used for truly untrusted code. It's a best-effort restriction. """ - # Safe builtins that are allowed in sandboxed environment + # Safe builtins that are allowed in the restricted-load environment. SAFE_BUILTINS: dict[str, Any] = { "True": True, "False": False, @@ -566,11 +566,11 @@ def import_hook( return __import__(name, globals, locals, fromlist, level) raise ImportError( - f"Import of '{name}' is not allowed in sandboxed environment. Allowed modules: {self.allowed_modules}" + f"Import of '{name}' is not allowed in restricted-load validation. Allowed modules: {self.allowed_modules}" ) def load_module(self, source_code: str, filename: str = "") -> types.ModuleType: - """Load a module in a restricted sandbox environment.""" + """Load a module in a best-effort restricted environment.""" restricted_globals = self.create_restricted_globals() restricted_globals["__builtins__"]["__import__"] = self.import_hook @@ -581,7 +581,7 @@ def load_module(self, source_code: str, filename: str = "") -> types.Mo exec(compiled, module.__dict__) return module except Exception as e: - raise SecurityError(f"Failed to load module in sandbox: {e}") + raise SecurityError(f"Failed to load module in restricted validation: {e}") class PipelineSecurityValidator: @@ -665,7 +665,7 @@ def validate_source(self, source_code: str, filename: str = "") -> Vali if self.blocked_patterns: self._check_custom_patterns(source_code, result) - # Level 4: Sandboxed execution (optional, for strict mode) + # Level 4: Best-effort restricted load check. This is not a runtime sandbox. if self.security_level == SecurityLevel.SANDBOX: self._try_sandboxed_load(source_code, filename, result) @@ -712,7 +712,7 @@ def _check_custom_patterns(self, source_code: str, result: ValidationResult) -> logger.warning(f"Invalid regex pattern '{pattern}': {e}") def _try_sandboxed_load(self, source_code: str, filename: str, result: ValidationResult) -> None: - """Try to load code in sandboxed environment.""" + """Try to load code in a best-effort restricted environment.""" try: loader = SandboxedLoader(allowed_modules=self.allowed_imports) loader.load_module(source_code, filename) @@ -722,13 +722,13 @@ def _try_sandboxed_load(self, source_code: str, filename: str, result: Validatio line_number=0, column=0, violation_type="SANDBOX_LOAD_FAILED", - description=f"Code failed sandboxed load: {e}", + description=f"Code failed restricted-load validation: {e}", code_snippet="", severity="high", ) ) except Exception as e: - logger.debug(f"Sandboxed load raised (possibly legitimate error): {e}") + logger.debug(f"Restricted-load validation raised (possibly legitimate error): {e}") def assert_safe(self, file_path: str) -> None: """Validate file and raise SecurityError if unsafe.""" diff --git a/tests/unit/openai/test_image_routes.py b/tests/unit/openai/test_image_routes.py index 7f4839d..88abef5 100644 --- a/tests/unit/openai/test_image_routes.py +++ b/tests/unit/openai/test_image_routes.py @@ -9,8 +9,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from fastapi import FastAPI -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException from telefuser.service.api.openai.image_routes import ImageRoutes, create_router from telefuser.service.api.task_application_service import TaskApplicationService diff --git a/tests/unit/openai/test_integration_server.py b/tests/unit/openai/test_integration_server.py index 4cf2a27..b005507 100644 --- a/tests/unit/openai/test_integration_server.py +++ b/tests/unit/openai/test_integration_server.py @@ -7,11 +7,12 @@ from unittest.mock import MagicMock import pytest -from fastapi.testclient import TestClient from telefuser.service.api.api_server import ApiServer from telefuser.service.core.task_manager import TaskStatus +from ._asgi_test_client import ASGITestClient as TestClient + @pytest.fixture def api_server(tmp_path): diff --git a/tests/unit/openai/test_video_routes.py b/tests/unit/openai/test_video_routes.py index 8e82b68..a68add5 100644 --- a/tests/unit/openai/test_video_routes.py +++ b/tests/unit/openai/test_video_routes.py @@ -8,8 +8,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from fastapi import FastAPI -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException from telefuser.service.api.openai.video_routes import VideoRoutes, create_router from telefuser.service.api.task_application_service import TaskApplicationService diff --git a/tests/unit/service/test_artifact_store.py b/tests/unit/service/test_artifact_store.py index ecf9dbb..538ed4c 100644 --- a/tests/unit/service/test_artifact_store.py +++ b/tests/unit/service/test_artifact_store.py @@ -9,11 +9,11 @@ import pytest from telefuser.service.api.api_server import ApiServer -from telefuser.service.core.config import ServerConfig from telefuser.service.core.artifact_store import ArtifactStore +from telefuser.service.core.config import ServerConfig from telefuser.service.core.file_service import FileService from telefuser.service.core.task_manager import TaskManager -from telefuser.service_types import MediaType +from telefuser.service_types import MediaType, TaskStatus class _AsyncUpload: @@ -58,6 +58,52 @@ def test_artifact_store_resolves_task_scoped_outputs_for_download(tmp_path: Path assert store.resolve_output_file(output) == output +def test_artifact_store_round_trips_local_artifact_ids(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + output = store.output_path("result.png", media_type=MediaType.IMAGE, task_id="task-123") + output.write_bytes(b"image") + + artifact_id = store.artifact_id_for_path(output) + + assert artifact_id == "local:tasks/task-123/outputs/images/result.png" + assert store.resolve_artifact_id(artifact_id) == output + assert store.resolve_output_file(artifact_id) == output + assert store.resolve_output_file("tasks/task-123/outputs/images/result.png") == output + + +def test_artifact_store_rejects_artifact_ids_outside_outputs(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + input_file = store.task_input_dir("task-123", MediaType.IMAGE) / "input.png" + input_file.write_bytes(b"image") + + with pytest.raises(ValueError, match="downloadable output"): + store.artifact_id_for_path(input_file) + + with pytest.raises(ValueError, match="not allowed"): + store.resolve_artifact_id("local:tasks/task-123/inputs/images/input.png") + + with pytest.raises(ValueError, match="escapes"): + store.resolve_artifact_id("local:../escape.png") + + +def test_artifact_store_returns_local_artifact_metadata(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + output = store.output_path("result.mp4", media_type=MediaType.VIDEO, task_id="task-123") + output.write_bytes(b"video") + + metadata = store.artifact_metadata(output, task_id="task-123", media_type=MediaType.VIDEO) + + assert metadata["artifact_id"] == "local:tasks/task-123/outputs/videos/result.mp4" + assert metadata["backend"] == "local" + assert metadata["relative_path"] == "tasks/task-123/outputs/videos/result.mp4" + assert metadata["task_id"] == "task-123" + assert metadata["media_type"] == "video" + assert metadata["filename"] == "result.mp4" + assert metadata["size_bytes"] == 5 + assert metadata["created_at"].endswith("Z") + assert metadata["modified_at"].endswith("Z") + + def test_file_service_uses_task_scoped_outputs_when_task_id_is_provided(tmp_path: Path) -> None: files = FileService(tmp_path) @@ -66,6 +112,18 @@ def test_file_service_uses_task_scoped_outputs_when_task_id_is_provided(tmp_path assert output == tmp_path / "tasks" / "task-123" / "outputs" / "images" / "result.png" +def test_file_service_exposes_artifact_metadata_helpers(tmp_path: Path) -> None: + files = FileService(tmp_path) + output = files.get_output_path("result.png", media_type=MediaType.IMAGE, task_id="task-123") + output.write_bytes(b"image") + + artifact_id = files.artifact_id_for_path(output) + + assert artifact_id == "local:tasks/task-123/outputs/images/result.png" + assert files.resolve_artifact_id(artifact_id) == output + assert files.artifact_metadata(artifact_id, task_id="task-123", media_type=MediaType.IMAGE)["size_bytes"] == 5 + + def test_file_service_saves_upload_stream_to_media_input_dir(tmp_path: Path) -> None: files = FileService(tmp_path) upload = _AsyncUpload("cat.png", [b"ca", b"t"]) @@ -167,6 +225,66 @@ def test_artifact_store_cleanup_removes_oversized_terminal_tasks(tmp_path: Path) assert small_output.exists() +def test_artifact_store_ephemeral_cleanup_removes_terminal_tasks(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + completed_output = store.output_path("done.mp4", media_type=MediaType.VIDEO, task_id="completed-task") + active_output = store.output_path("active.mp4", media_type=MediaType.VIDEO, task_id="active-task") + completed_output.write_bytes(b"done") + active_output.write_bytes(b"active") + now = datetime.now() + + result = store.cleanup( + active_task_ids={"active-task"}, + terminal_task_end_times={ + "completed-task": now, + "active-task": now, + }, + terminal_task_statuses={ + "completed-task": TaskStatus.COMPLETED.value, + "active-task": TaskStatus.PROCESSING.value, + }, + retention_seconds=0, + tmp_retention_seconds=0, + persistence_mode="ephemeral", + now=now, + ) + + assert result["removed_task_ids"] == ["completed-task"] + assert not completed_output.exists() + assert active_output.exists() + + +def test_artifact_store_preserves_failed_outputs_from_automatic_cleanup(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + failed_output = store.output_path("failed.mp4", media_type=MediaType.VIDEO, task_id="failed-task") + completed_output = store.output_path("done.mp4", media_type=MediaType.VIDEO, task_id="completed-task") + failed_output.write_bytes(b"failed" * 10) + completed_output.write_bytes(b"done" * 10) + now = datetime.now() + + result = store.cleanup( + active_task_ids=set(), + terminal_task_end_times={ + "failed-task": now - timedelta(seconds=120), + "completed-task": now - timedelta(seconds=120), + }, + terminal_task_statuses={ + "failed-task": TaskStatus.FAILED.value, + "completed-task": TaskStatus.COMPLETED.value, + }, + retention_seconds=60, + tmp_retention_seconds=0, + persistence_mode="persistent", + preserve_failed_outputs=True, + max_task_bytes=1, + now=now, + ) + + assert result["removed_task_ids"] == ["completed-task"] + assert failed_output.exists() + assert not completed_output.exists() + + def test_file_service_passes_max_task_bytes_to_artifact_cleanup(tmp_path: Path) -> None: files = FileService(tmp_path, artifact_max_task_bytes=3) output = files.get_output_path("big.mp4", media_type=MediaType.VIDEO, task_id="task-123") @@ -181,6 +299,35 @@ def test_file_service_passes_max_task_bytes_to_artifact_cleanup(tmp_path: Path) assert not output.exists() +def test_file_service_passes_persistence_and_failed_preservation_to_artifact_cleanup(tmp_path: Path) -> None: + files = FileService( + tmp_path, + artifact_persistence_mode="ephemeral", + artifact_preserve_failed_outputs=True, + ) + failed_output = files.get_output_path("failed.mp4", media_type=MediaType.VIDEO, task_id="failed-task") + completed_output = files.get_output_path("done.mp4", media_type=MediaType.VIDEO, task_id="completed-task") + failed_output.write_bytes(b"failed") + completed_output.write_bytes(b"done") + now = datetime.now() + + result = files.cleanup_artifacts( + active_task_ids=set(), + terminal_task_end_times={ + "failed-task": now, + "completed-task": now, + }, + terminal_task_statuses={ + "failed-task": TaskStatus.FAILED.value, + "completed-task": TaskStatus.COMPLETED.value, + }, + ) + + assert result["removed_task_ids"] == ["completed-task"] + assert failed_output.exists() + assert not completed_output.exists() + + def test_api_server_runs_artifact_cleanup_from_task_snapshot(tmp_path: Path) -> None: task_manager = TaskManager() config = ServerConfig(artifact_retention_seconds=1, artifact_tmp_retention_seconds=0) @@ -199,6 +346,19 @@ def test_api_server_runs_artifact_cleanup_from_task_snapshot(tmp_path: Path) -> assert not output.exists() +def test_api_server_initializes_file_service_with_artifact_lifecycle_config(tmp_path: Path) -> None: + task_manager = TaskManager() + config = ServerConfig( + artifact_persistence_mode="ephemeral", + artifact_preserve_failed_outputs=True, + ) + server = ApiServer(task_manager=task_manager, config=config, enable_openai_api=False) + server.initialize_services(tmp_path, Mock()) + + assert server.file_service.artifact_persistence_mode == "ephemeral" + assert server.file_service.artifact_preserve_failed_outputs is True + + def test_api_server_artifact_cleanup_loop_starts_and_stops(tmp_path: Path) -> None: async def _exercise() -> None: server = ApiServer(task_manager=TaskManager(), config=ServerConfig(), enable_openai_api=False) diff --git a/tests/unit/service/test_config.py b/tests/unit/service/test_config.py index 2da9046..31ff898 100644 --- a/tests/unit/service/test_config.py +++ b/tests/unit/service/test_config.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest + from telefuser.service.core.config import ServerConfig from telefuser.service.core.container import ServiceContainer @@ -33,11 +35,19 @@ def test_artifact_retention_defaults_are_configured() -> None: assert config.artifact_retention_seconds == 7 * 24 * 60 * 60 assert config.artifact_tmp_retention_seconds == 60 * 60 assert config.artifact_cleanup_interval_seconds == 60 * 60 + assert config.artifact_persistence_mode == "persistent" assert config.artifact_max_total_bytes == 0 assert config.artifact_max_task_bytes == 0 assert config.artifact_preserve_failed_outputs is False +def test_artifact_persistence_mode_is_validated() -> None: + assert ServerConfig(artifact_persistence_mode="ephemeral").artifact_persistence_mode == "ephemeral" + + with pytest.raises(ValueError): + ServerConfig(artifact_persistence_mode="invalid") + + def test_artifact_local_root_can_override_cache_dir() -> None: config = ServerConfig(cache_dir="cache-a", artifact_local_root="artifact-a") diff --git a/tests/unit/service/test_latent_cache_cli.py b/tests/unit/service/test_latent_cache_cli.py index dfdc22c..611fa83 100644 --- a/tests/unit/service/test_latent_cache_cli.py +++ b/tests/unit/service/test_latent_cache_cli.py @@ -11,6 +11,7 @@ from telefuser.entrypoints.cli.main import main from telefuser.service.core.config import ServerConfig from telefuser.service.core.container import ServiceContainer +from telefuser.service_types import TaskType def test_telefuser_pyproject_does_not_vendor_cacheseek_dependencies() -> None: @@ -57,7 +58,7 @@ def fake_run_server(**kwargs): assert calls == [ { "pipe_path": "pipeline.py", - "task": "i2v", + "task": TaskType.I2V, "port": 8000, "host": "127.0.0.1", "cache_dir": "work_dirs/server_cache", @@ -65,6 +66,8 @@ def fake_run_server(**kwargs): "num_replicas": 1, "enable_latent_cache": True, "cache_mode": "read_only", + "security_level": "strict", + "skip_validation": True, } ] diff --git a/tests/unit/service/test_security_validator.py b/tests/unit/service/test_security_validator.py new file mode 100644 index 0000000..2e437d6 --- /dev/null +++ b/tests/unit/service/test_security_validator.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from telefuser.service.security.security_validator import PipelineSecurityValidator, SecurityLevel + + +def test_sandbox_level_reports_restricted_load_not_runtime_isolation() -> None: + validator = PipelineSecurityValidator(security_level=SecurityLevel.SANDBOX) + + result = validator.validate_source("raise RuntimeError('load failed')\n") + + assert result.is_safe is False + assert result.violations + assert "restricted-load validation" in result.violations[0].description + assert "sandboxed load" not in result.violations[0].description diff --git a/tests/unit/service/test_service_smoke.py b/tests/unit/service/test_service_smoke.py new file mode 100644 index 0000000..18f635d --- /dev/null +++ b/tests/unit/service/test_service_smoke.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from unittest.mock import Mock + +from telefuser.service.api.api_server import ApiServer +from telefuser.service.api.routers.files import FileRoutes +from telefuser.service.core.file_service import FileService +from telefuser.service.core.task_manager import TaskManager +from telefuser.service_types import MediaType +from tests.unit.openai._asgi_test_client import ASGITestClient + + +def _make_smoke_server(tmp_path: Path) -> ApiServer: + task_manager = TaskManager(max_queue_size=10) + server = ApiServer( + task_manager=task_manager, + enable_rate_limit=False, + enable_openai_api=True, + ) + server.file_service = FileService(tmp_path) + server.inference_service = Mock() + server.inference_service.server_metadata.return_value = { + "pipeline_file": "/test/pipeline.py", + "pipeline_name": "smoke_pipeline", + "parallelism": 1, + "task": "t2v", + "declared_pipeline_contract": True, + "supported_tasks": ["t2v", "t2i"], + "supported_media_types": ["video", "image"], + "execution_mode": "serial_single_pipeline", + "effective_max_concurrent_tasks": 1, + "entrypoints": {"get_pipeline": "get_pipeline", "run_with_file": "run_with_file"}, + "task_contracts": { + "t2v": {"required_inputs": [], "media_type": "video"}, + "t2i": {"required_inputs": [], "media_type": "image"}, + }, + } + return server + + +def test_task_create_status_and_artifact_id_download_smoke(tmp_path: Path) -> None: + server = _make_smoke_server(tmp_path) + + with ASGITestClient(server.app) as client: + create_response = client.post( + "/v1/tasks/create", + json={"task": "t2i", "prompt": "a cat", "output_format": "png"}, + ) + assert create_response.status_code == 200 + task_id = create_response.json()["task_id"] + + output_path = server.file_service.get_output_path("result.png", media_type=MediaType.IMAGE, task_id=task_id) + output_path.write_bytes(b"image") + server.task_manager.complete_task(task_id, output_path=str(output_path)) + + status_response = client.get(f"/v1/tasks/{task_id}/status") + assert status_response.status_code == 200 + assert status_response.json()["status"] == "completed" + + artifact_id = server.file_service.artifact_id_for_path(output_path) + download_response = asyncio.run(FileRoutes(server).download_file(artifact_id)) + assert download_response.headers["content-length"] == "5" + assert download_response.headers["content-disposition"] == 'attachment; filename="result.png"' + + +def test_openai_video_retrieve_includes_artifact_metadata_smoke(tmp_path: Path) -> None: + server = _make_smoke_server(tmp_path) + + with ASGITestClient(server.app) as client: + create_response = client.post( + "/v1/videos", + json={"prompt": "a cat playing", "seconds": 4, "size": "1024x576"}, + ) + assert create_response.status_code == 200 + video_id = create_response.json()["id"] + + output_path = server.file_service.get_output_path("clip.mp4", media_type=MediaType.VIDEO, task_id=video_id) + output_path.write_bytes(b"video") + server.task_manager.complete_task(video_id, output_path=str(output_path)) + + retrieve_response = client.get(f"/v1/videos/{video_id}") + assert retrieve_response.status_code == 200 + data = retrieve_response.json() + assert data["status"] == "completed" + assert data["url"].endswith(f"/v1/videos/{video_id}/content") + assert data["artifact_id"] == f"local:tasks/{video_id}/outputs/videos/clip.mp4" + assert data["artifact_metadata"]["backend"] == "local" + assert data["artifact_metadata"]["size_bytes"] == 5 diff --git a/tests/unit/service/test_task_application_service.py b/tests/unit/service/test_task_application_service.py index dc42207..a959332 100644 --- a/tests/unit/service/test_task_application_service.py +++ b/tests/unit/service/test_task_application_service.py @@ -155,6 +155,25 @@ def test_task_application_service_returns_output_response(tmp_path: Path) -> Non assert response.headers["content-length"] == "5" +def test_task_application_service_returns_output_metadata(tmp_path: Path) -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + server.file_service = FileService(tmp_path) + message = TaskRequest(task="t2i") + task_id = task_manager.create_task(message) + output_path = server.file_service.get_output_path("result.png", media_type="image", task_id=task_id) + output_path.write_bytes(b"image") + task_manager.complete_task(task_id, output_path=str(output_path)) + + metadata = server.task_app_service.get_output_metadata(task_id, media_type="image") + + assert metadata is not None + assert metadata["artifact_id"] == f"local:tasks/{task_id}/outputs/images/result.png" + assert metadata["backend"] == "local" + assert metadata["task_id"] == task_id + assert metadata["size_bytes"] == 5 + + def test_task_application_service_rejects_unready_required_output(tmp_path: Path) -> None: task_manager = TaskManager() server = ApiServer(task_manager=task_manager, enable_openai_api=False) diff --git a/tests/unit/service/test_task_runtime.py b/tests/unit/service/test_task_runtime.py index b4d9306..283291f 100644 --- a/tests/unit/service/test_task_runtime.py +++ b/tests/unit/service/test_task_runtime.py @@ -37,6 +37,7 @@ def test_artifact_cleanup_snapshot_splits_active_and_terminal_tasks() -> None: assert snapshot["active_task_ids"] == {active_id} assert terminal_id in snapshot["terminal_task_end_times"] + assert snapshot["terminal_task_statuses"] == {terminal_id: TaskStatus.COMPLETED.value} assert active_id not in snapshot["terminal_task_end_times"] From aec99c71e20b071578a730146c82c1491c03381e Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 8 Jul 2026 16:40:06 +0000 Subject: [PATCH 6/6] feat(service): support dynamic form task parameters - Collect arbitrary non-file multipart form fields into TaskRequest-compatible payloads. - Update WanVideo I2V run_with_file helpers to consume first_image_path and reject object image inputs. - Add service tests for dynamic form fields, including resolution passthrough. - Document recommended artifact cleanup profiles in English and Chinese service docs. Verification: - python -m pytest tests/unit/service/test_task_routes.py tests/integration/test_service_api.py::TestFileUpload -q - python -m pytest tests/unit/service/test_artifact_store.py tests/unit/service/test_task_application_service.py tests/unit/service/test_service_smoke.py tests/unit/openai/test_video_routes.py tests/unit/openai/test_image_routes.py -q - python -m ruff check telefuser/service/api/routers/tasks.py tests/unit/service/test_task_routes.py tests/integration/test_service_api.py - python -m ruff check examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py examples/wan_video/wan21_14b_image_to_video_h100.py examples/wan_video/wan21_14b_image_to_video_lora_h100.py examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py examples/wan_video/wan22_14b_image_to_video_distill_h100.py examples/wan_video/wan22_14b_image_to_video_lora_h100.py - python -m py_compile examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py examples/wan_video/wan21_14b_image_to_video_h100.py examples/wan_video/wan21_14b_image_to_video_lora_h100.py examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py examples/wan_video/wan22_14b_image_to_video_distill_h100.py examples/wan_video/wan22_14b_image_to_video_lora_h100.py --- docs/en/service.md | 25 +++ docs/zh/service.md | 25 +++ ...c_wan22_14b_image_to_video_distill_h100.py | 9 +- .../wan21_14b_image_to_video_h100.py | 20 ++- .../wan21_14b_image_to_video_lora_h100.py | 20 ++- ...n22_14b_image_to_video_distill_fp8_h100.py | 20 ++- .../wan22_14b_image_to_video_distill_h100.py | 20 ++- .../wan22_14b_image_to_video_lora_h100.py | 20 ++- telefuser/service/api/api_server.py | 4 +- telefuser/service/api/routers/service.py | 27 +++- telefuser/service/api/routers/tasks.py | 149 ++++++++++-------- telefuser/service/core/artifact_store.py | 8 - telefuser/service/core/file_service.py | 4 +- telefuser/service/core/task_service.py | 13 +- tests/integration/test_service_api.py | 16 ++ tests/unit/service/test_artifact_store.py | 16 +- .../service/test_latent_cache_task_service.py | 4 +- ...actor_phase0.py => test_service_routes.py} | 49 +++++- tests/unit/service/test_task_routes.py | 46 +++++- 19 files changed, 349 insertions(+), 146 deletions(-) rename tests/unit/service/{test_service_refactor_phase0.py => test_service_routes.py} (90%) diff --git a/docs/en/service.md b/docs/en/service.md index 9116b49..fe2a4bb 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -319,6 +319,31 @@ them. `ephemeral` removes terminal task directories on the next cleanup pass. If `TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true`, failed task directories are protected from automatic cleanup so partial outputs can be inspected manually. +Recommended cleanup profiles: + +```bash +# Development/debugging: keep outputs for one day and preserve failed task outputs. +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent +export TELEFUSER_ARTIFACT_RETENTION_SECONDS=86400 +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=1800 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=600 +export TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true + +# Temporary benchmarking: remove terminal task directories on the next cleanup pass. +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=ephemeral +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=600 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=300 +export TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=false + +# Long-running service: keep artifacts for seven days with bounded local storage. +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent +export TELEFUSER_ARTIFACT_RETENTION_SECONDS=604800 +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=3600 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=3600 +export TELEFUSER_ARTIFACT_MAX_TOTAL_BYTES=107374182400 +export TELEFUSER_ARTIFACT_MAX_TASK_BYTES=10737418240 +``` + ### Rate Limit and Proxy Headers Rate limiting is in-memory and keyed by client identity. By default, the server uses the direct client address and does diff --git a/docs/zh/service.md b/docs/zh/service.md index 200f362..c4adbcd 100644 --- a/docs/zh/service.md +++ b/docs/zh/service.md @@ -318,6 +318,31 @@ Artifact 是本地文件。设置 `TELEFUSER_ARTIFACT_LOCAL_ROOT` 时使用该 `ephemeral` 会在下一次清理时删除终态任务目录。设置 `TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true` 后，失败任务目录会跳过自动清理，便于人工检查部分输出。 +推荐清理配置： + +```bash +# 开发/调试：保留输出一天，并保留失败任务输出用于排查。 +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent +export TELEFUSER_ARTIFACT_RETENTION_SECONDS=86400 +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=1800 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=600 +export TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=true + +# 临时压测：任务进入终态后，在下一次清理时删除任务目录。 +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=ephemeral +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=600 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=300 +export TELEFUSER_ARTIFACT_PRESERVE_FAILED_OUTPUTS=false + +# 长期运行服务：保留 artifact 七天，并限制本地存储容量。 +export TELEFUSER_ARTIFACT_PERSISTENCE_MODE=persistent +export TELEFUSER_ARTIFACT_RETENTION_SECONDS=604800 +export TELEFUSER_ARTIFACT_TMP_RETENTION_SECONDS=3600 +export TELEFUSER_ARTIFACT_CLEANUP_INTERVAL_SECONDS=3600 +export TELEFUSER_ARTIFACT_MAX_TOTAL_BYTES=107374182400 +export TELEFUSER_ARTIFACT_MAX_TASK_BYTES=10737418240 +``` + ### 限流和代理 Header 限流状态保存在进程内存中，并按客户端身份计数。默认使用直接客户端地址，不信任 `X-Forwarded-For`， diff --git a/examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py b/examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py index 2fe5903..489b8f2 100644 --- a/examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py +++ b/examples/wan_video/async_wan22_14b_image_to_video_distill_h100.py @@ -154,8 +154,8 @@ async def run(pipe, req_id, seed, prompt, image, ppl_config=None): return {"uri": saved_uri, "wall_time_s": request_elapsed, "pipeline_time_ms": total_ms, "stage_times": stage_times} -async def run_with_file(pipe, req_id, seed, prompt, image, ppl_config=None): - """Run async pipeline and save to file. +async def run_with_file(pipe, req_id, seed, prompt, first_image_path: str, ppl_config=None): + """Run async pipeline from an input image path and return artifact metadata. Note: The async pipeline handles file saving internally through agenerate(). The output URI is returned in the result dict. @@ -165,15 +165,18 @@ async def run_with_file(pipe, req_id, seed, prompt, image, ppl_config=None): req_id: Request ID for tracking seed: Random seed prompt: Positive guidance text prompt - image: Input image + first_image_path: Input image path ppl_config: Pipeline configuration Returns: Result dict with uri, wall_time_s, pipeline_time_ms, stage_times """ + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") if ppl_config is None: ppl_config = PPL_CONFIG + image = Image.open(first_image_path).convert("RGB") result = await run(pipe, req_id, seed, prompt, image, ppl_config) logger.info(f"Video saved to: {result['uri']}") return result diff --git a/examples/wan_video/wan21_14b_image_to_video_h100.py b/examples/wan_video/wan21_14b_image_to_video_h100.py index e736259..95b6037 100755 --- a/examples/wan_video/wan21_14b_image_to_video_h100.py +++ b/examples/wan_video/wan21_14b_image_to_video_h100.py @@ -139,21 +139,27 @@ def run( def run_with_file( - pipeline, - image, - prompt, - negative_prompt, - seed, - output_path, + pipeline: Wan21VideoPipeline, + first_image_path: str, + prompt: str = "", + negative_prompt: str = "", + seed: int = PPL_CONFIG["seed"], + output_path: str = "", + resolution: str = PPL_CONFIG["resolution"], **kwargs, ): - """Run pipeline and save to file.""" + """Run pipeline from an input image path and save to file.""" + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") + + image = Image.open(first_image_path).convert("RGB") video = run( pipeline, image, prompt, negative_prompt, seed, + resolution=resolution, ) print(f"Saving video to {output_path}") save_video( diff --git a/examples/wan_video/wan21_14b_image_to_video_lora_h100.py b/examples/wan_video/wan21_14b_image_to_video_lora_h100.py index a4f8647..995c5f2 100755 --- a/examples/wan_video/wan21_14b_image_to_video_lora_h100.py +++ b/examples/wan_video/wan21_14b_image_to_video_lora_h100.py @@ -137,16 +137,20 @@ def run( def run_with_file( - pipeline, - image, - prompt, - negative_prompt, - seed, - output_path, - resolution=PPL_CONFIG["resolution"], + pipeline: Wan21VideoPipeline, + first_image_path: str, + prompt: str = "", + negative_prompt: str = "", + seed: int = PPL_CONFIG["seed"], + output_path: str = "", + resolution: str = PPL_CONFIG["resolution"], **kwargs, ): - """Run pipeline and save to file.""" + """Run pipeline from an input image path and save to file.""" + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") + + image = Image.open(first_image_path).convert("RGB") video = run( pipeline, image, diff --git a/examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py b/examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py index f0e7b28..420021a 100755 --- a/examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py +++ b/examples/wan_video/wan22_14b_image_to_video_distill_fp8_h100.py @@ -140,21 +140,27 @@ def run( def run_with_file( - pipeline, - image, - prompt, - negative_prompt, - seed, - output_path, + pipeline: Wan22VideoPipeline, + first_image_path: str, + prompt: str = "", + negative_prompt: str = "", + seed: int = PPL_CONFIG["seed"], + output_path: str = "", + resolution: str = PPL_CONFIG["resolution"], **kwargs, ): - """Run pipeline and save to file.""" + """Run pipeline from an input image path and save to file.""" + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") + + image = Image.open(first_image_path).convert("RGB") video = run( pipeline, image, prompt, negative_prompt, seed, + resolution=resolution, ) print(f"Saving video to {output_path}") save_video( diff --git a/examples/wan_video/wan22_14b_image_to_video_distill_h100.py b/examples/wan_video/wan22_14b_image_to_video_distill_h100.py index e4e92d9..b02f10d 100755 --- a/examples/wan_video/wan22_14b_image_to_video_distill_h100.py +++ b/examples/wan_video/wan22_14b_image_to_video_distill_h100.py @@ -181,16 +181,20 @@ def run( def run_with_file( - pipeline, - image, - prompt, - negative_prompt, - seed, - output_path, - resolution=PPL_CONFIG["resolution"], + pipeline: Wan22VideoPipeline, + first_image_path: str, + prompt: str = "", + negative_prompt: str = "", + seed: int = PPL_CONFIG["seed"], + output_path: str = "", + resolution: str = PPL_CONFIG["resolution"], **kwargs, ): - """Run pipeline and save to file.""" + """Run pipeline from an input image path and save to file.""" + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") + + image = Image.open(first_image_path).convert("RGB") video = run( pipeline, image, diff --git a/examples/wan_video/wan22_14b_image_to_video_lora_h100.py b/examples/wan_video/wan22_14b_image_to_video_lora_h100.py index 9418d4c..ee90660 100755 --- a/examples/wan_video/wan22_14b_image_to_video_lora_h100.py +++ b/examples/wan_video/wan22_14b_image_to_video_lora_h100.py @@ -149,16 +149,20 @@ def run(pipeline, image, prompt, negative_prompt="", seed=PPL_CONFIG["seed"], re def run_with_file( - pipeline, - image, - prompt, - negative_prompt, - seed, - output_path, - resolution=PPL_CONFIG["resolution"], + pipeline: Wan22VideoPipeline, + first_image_path: str, + prompt: str = "", + negative_prompt: str = "", + seed: int = PPL_CONFIG["seed"], + output_path: str = "", + resolution: str = PPL_CONFIG["resolution"], **kwargs, ): - """Run pipeline and save to file.""" + """Run pipeline from an input image path and save to file.""" + if not first_image_path: + raise ValueError("run_with_file requires first_image_path") + + image = Image.open(first_image_path).convert("RGB") video = run(pipeline, image, prompt, negative_prompt, seed, resolution=resolution) logger.info(f"Saving video to {output_path}") save_video( diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index b480203..71e0cd9 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -191,9 +191,7 @@ async def _validate_image_url(self, image_url: str) -> bool: timeout = httpx.Timeout(connect=5.0, read=5.0) verify = ( - self.server_config.ssl_cert_path - if self.server_config.ssl_cert_path - else self.server_config.verify_ssl + self.server_config.ssl_cert_path if self.server_config.ssl_cert_path else self.server_config.verify_ssl ) async with httpx.AsyncClient(verify=verify, timeout=timeout) as client: response = await client.head(image_url, follow_redirects=True) diff --git a/telefuser/service/api/routers/service.py b/telefuser/service/api/routers/service.py index 6c6d3d8..3231960 100644 --- a/telefuser/service/api/routers/service.py +++ b/telefuser/service/api/routers/service.py @@ -30,9 +30,10 @@ async def get_status(self) -> dict: status = self.api.task_manager.get_service_status() status["effective_max_concurrent_tasks"] = self.api.task_manager.max_concurrent_processing status["configured_max_concurrent_tasks"] = self.api.configured_max_concurrent_tasks - if hasattr(self.api.inference_service, "pool_status"): + pool_status = self._pipeline_pool_status() + if pool_status is not None: status["execution_mode"] = "concurrent_pipeline_pool" - status["pool"] = self.api.inference_service.pool_status() + status["pool"] = pool_status else: status["execution_mode"] = "serial_single_pipeline" webrtc_stats = self._webrtc_session_stats() @@ -41,6 +42,18 @@ async def get_status(self) -> dict: status["service_status"] = "active" return status + def _pipeline_pool_status(self) -> list[dict] | None: + """Return pipeline pool status when the inference service exposes a real pool.""" + if self.api.inference_service is None: + return None + pool_status_fn = getattr(self.api.inference_service, "pool_status", None) + if not callable(pool_status_fn): + return None + pool_status = pool_status_fn() + if not isinstance(pool_status, list) or not all(isinstance(replica, dict) for replica in pool_status): + return None + return pool_status + def _webrtc_session_stats(self) -> dict: """Return WebRTC session stats if available.""" routes = self.api._webrtc_routes @@ -67,12 +80,12 @@ async def get_metadata(self) -> dict: async def health_check(self) -> dict: """Liveness endpoint for monitoring.""" - from datetime import UTC, datetime + from datetime import datetime, timezone status = { "status": "healthy", "ready": self._is_ready(), - "timestamp": datetime.now(UTC).isoformat(), + "timestamp": datetime.now(timezone.utc).isoformat(), "version": "1.0.0", } @@ -89,9 +102,9 @@ async def health_check(self) -> dict: def _is_ready(self) -> bool: """Return whether the initialized service can currently accept work.""" if self.api.inference_service is not None: - if hasattr(self.api.inference_service, "pool_status"): - pool_status = self.api.inference_service.pool_status() - return bool(pool_status.get("alive_replicas", 0) > 0) + pool_status = self._pipeline_pool_status() + if pool_status is not None: + return any(replica.get("status") != "dead" for replica in pool_status) return bool(getattr(self.api.inference_service, "is_running", False)) if self.api.stream_service is not None: diff --git a/telefuser/service/api/routers/tasks.py b/telefuser/service/api/routers/tasks.py index b34f1c2..437264d 100644 --- a/telefuser/service/api/routers/tasks.py +++ b/telefuser/service/api/routers/tasks.py @@ -7,14 +7,16 @@ from __future__ import annotations +import json from pathlib import Path from typing import TYPE_CHECKING, Any -from fastapi import APIRouter, File, Form, HTTPException, UploadFile +from fastapi import APIRouter, File, HTTPException, Request, UploadFile from pydantic import ValidationError +from starlette.datastructures import UploadFile as StarletteUploadFile from telefuser.service.core.pipeline_contract import default_task_contract, validate_task_name_format -from telefuser.service_types import AspectRatio, MediaType, OutputFormat, StopTaskStatus +from telefuser.service_types import MediaType, StopTaskStatus from telefuser.utils.logging import logger from ..schema import StopTaskResponse, TaskRequest, TaskResponse @@ -40,33 +42,15 @@ async def create_task(message: TaskRequest) -> TaskResponse: @new_router.post("/form", response_model=TaskResponse) async def create_task_form( + request: Request, first_image_file: UploadFile | None = File(default=None, description="First frame image file"), last_image_file: UploadFile | None = File(default=None, description="Last frame image file"), - prompt: str = Form(default="", description="Generation prompt"), - task: str = Form(default="", description="Optional explicit task name"), - output_path: str = Form(default="", description="Custom output path"), - negative_prompt: str = Form(default="", description="Negative prompt"), - target_video_length: int = Form(default=5, description="Video length in seconds (for video tasks)"), - seed: int = Form(default=42, description="Random seed"), - aspect_ratio: AspectRatio = Form( - default=AspectRatio.RATIO_16_9, description="Aspect ratio (16:9, 9:16, 4:3, etc.)" - ), - output_format: OutputFormat = Form( - default=OutputFormat.PNG, description="Output format (png, jpg, webp for images)" - ), ) -> TaskResponse: """Create task with file upload support.""" return await routes.create_task_form( + request=request, first_image_file=first_image_file, last_image_file=last_image_file, - prompt=prompt, - task=task, - output_path=output_path, - negative_prompt=negative_prompt, - target_video_length=target_video_length, - seed=seed, - aspect_ratio=aspect_ratio, - output_format=output_format, ) @new_router.get("/queue/status", response_model=dict) @@ -164,18 +148,11 @@ async def stop_task(self, task_id: str) -> StopTaskResponse: async def create_task_form( self, + request: Request, first_image_file: UploadFile | None = None, last_image_file: UploadFile | None = None, - prompt: str | None = None, - task: str = "", - output_path: str | None = None, - negative_prompt: str | None = None, - target_video_length: int | None = None, - seed: int | None = None, - aspect_ratio: AspectRatio = AspectRatio.RATIO_16_9, - output_format: OutputFormat = OutputFormat.PNG, ) -> TaskResponse: - """Create task with file upload support.""" + """Create task with file upload support and dynamic form parameters.""" assert self.api.file_service is not None, "File service is not initialized" async def save_file_async(file: UploadFile, media_type: MediaType) -> str: @@ -189,49 +166,37 @@ async def save_file_async(file: UploadFile, media_type: MediaType) -> str: fallback_filename="input.mp4" if media_type == MediaType.VIDEO else "input.png", ) - first_image_path = "" - ref_video_path = "" - if first_image_file and first_image_file.filename: - if self._is_video_upload(first_image_file): - ref_video_path = await save_file_async(first_image_file, MediaType.VIDEO) - else: - first_image_path = await save_file_async(first_image_file, MediaType.IMAGE) - - last_image_path = "" - if last_image_file and last_image_file.filename: - if self._is_video_upload(last_image_file): - raise HTTPException(status_code=400, detail="last_image_file must be an image upload") - last_image_path = await save_file_async(last_image_file, MediaType.IMAGE) - try: - resolved_task = self._resolve_form_task( - requested_task=task, - first_image_path=first_image_path, - last_image_path=last_image_path, - ref_video_path=ref_video_path, - ) + task_payload = await self._collect_form_task_payload(request) + + first_image_path = "" + ref_video_path = "" + if first_image_file and first_image_file.filename: + if self._is_video_upload(first_image_file): + ref_video_path = await save_file_async(first_image_file, MediaType.VIDEO) + else: + first_image_path = await save_file_async(first_image_file, MediaType.IMAGE) + + last_image_path = "" + if last_image_file and last_image_file.filename: + if self._is_video_upload(last_image_file): + raise HTTPException(status_code=400, detail="last_image_file must be an image upload") + last_image_path = await save_file_async(last_image_file, MediaType.IMAGE) - task_payload: dict[str, Any] = {"task": resolved_task} - if prompt not in (None, ""): - task_payload["prompt"] = prompt - if negative_prompt not in (None, ""): - task_payload["negative_prompt"] = negative_prompt if first_image_path: task_payload["first_image_path"] = first_image_path if last_image_path: task_payload["last_image_path"] = last_image_path if ref_video_path: task_payload["ref_video_path"] = ref_video_path - if output_path not in (None, ""): - task_payload["output_path"] = output_path - if target_video_length is not None: - task_payload["target_video_length"] = target_video_length - if seed is not None: - task_payload["seed"] = seed - if aspect_ratio not in (None, ""): - task_payload["aspect_ratio"] = aspect_ratio - if output_format not in (None, ""): - task_payload["output_format"] = output_format + + resolved_task = self._resolve_form_task( + requested_task=str(task_payload.get("task") or ""), + first_image_path=first_image_path, + last_image_path=last_image_path, + ref_video_path=ref_video_path, + ) + task_payload["task"] = resolved_task message = TaskRequest(**task_payload) return await self.api.task_app_service.submit( @@ -247,6 +212,58 @@ async def save_file_async(file: UploadFile, media_type: MediaType) -> str: logger.error(f"Failed to create form task: {e}") raise HTTPException(status_code=500, detail=str(e)) + async def _collect_form_task_payload(self, request: Request) -> dict[str, Any]: + """Collect non-file multipart fields into a TaskRequest-compatible payload.""" + form = await request.form() + task_payload: dict[str, Any] = {} + file_field_names = {"first_image_file", "last_image_file"} + + for key, value in form.multi_items(): + if key in file_field_names or isinstance(value, (UploadFile, StarletteUploadFile)): + continue + if value in (None, ""): + continue + + coerced_value = self._coerce_form_value(value) + if key in task_payload: + existing_value = task_payload[key] + if isinstance(existing_value, list): + existing_value.append(coerced_value) + else: + task_payload[key] = [existing_value, coerced_value] + else: + task_payload[key] = coerced_value + + return task_payload + + def _coerce_form_value(self, value: Any) -> Any: + """Coerce multipart string values into basic JSON-compatible scalar types.""" + if not isinstance(value, str): + return value + + lowered = value.lower() + if lowered == "true": + return True + if lowered == "false": + return False + if lowered == "null": + return None + + try: + return int(value) + except ValueError: + pass + + try: + return float(value) + except ValueError: + pass + + try: + return json.loads(value) + except json.JSONDecodeError: + return value + def _resolve_form_task( self, *, diff --git a/telefuser/service/core/artifact_store.py b/telefuser/service/core/artifact_store.py index 3a83a1a..19ebae7 100644 --- a/telefuser/service/core/artifact_store.py +++ b/telefuser/service/core/artifact_store.py @@ -140,14 +140,6 @@ def resolve_output_file(self, file_path: str | Path) -> Path: if candidate.exists(): return candidate - for task_root in self.tasks_dir.glob("*"): - if not task_root.is_dir(): - continue - for media_dir in ("videos", "images"): - candidate = self._resolve_under(task_root / "outputs" / media_dir, path) - if candidate.exists(): - return candidate - return self._resolve_under(self.legacy_output_dir, path) def artifact_id_for_path(self, file_path: str | Path) -> str: diff --git a/telefuser/service/core/file_service.py b/telefuser/service/core/file_service.py index 71fdb87..8bf13f0 100644 --- a/telefuser/service/core/file_service.py +++ b/telefuser/service/core/file_service.py @@ -98,9 +98,7 @@ async def _download_file_with_retry( if response.status_code == 200: content_length = response.headers.get("content-length") if content_length is not None and int(content_length) > max_size: - raise ValueError( - f"File too large: {content_length} bytes, max: {max_size} bytes" - ) + raise ValueError(f"File too large: {content_length} bytes, max: {max_size} bytes") await self._write_response_stream(response, destination, max_size) return if response.status_code >= 500: diff --git a/telefuser/service/core/task_service.py b/telefuser/service/core/task_service.py index 6b95670..08c7066 100644 --- a/telefuser/service/core/task_service.py +++ b/telefuser/service/core/task_service.py @@ -112,14 +112,11 @@ async def update_audio_path(audio_name: str, message: TaskRequest, task_data: di # Determine media type and set appropriate output path media_type = infer_media_type_for_task(message.task) - try: - actual_save_path = self.file_service.get_output_path( - message.output_path, - media_type=media_type, - task_id=message.task_id, - ) - except TypeError: - actual_save_path = self.file_service.get_output_path(message.output_path, media_type=media_type) + actual_save_path = self.file_service.get_output_path( + message.output_path, + media_type=media_type, + task_id=message.task_id, + ) task_data["output_path"] = str(actual_save_path) # Best-effort: every step degrades silently to "no cache" on failure. diff --git a/tests/integration/test_service_api.py b/tests/integration/test_service_api.py index 4bb8b86..2404c4c 100644 --- a/tests/integration/test_service_api.py +++ b/tests/integration/test_service_api.py @@ -351,6 +351,22 @@ def test_upload_video(self, fs_client): assert response.status_code in [200, 422] + def test_upload_image_forwards_dynamic_form_fields(self, fs_client): + """Test multipart form endpoint forwards custom task parameters.""" + response = fs_client.post( + "/v1/tasks/form", + files={"first_image_file": ("test.png", b"fake image data", "image/png")}, + data={"task": "i2v", "prompt": "test", "resolution": "480p", "seed": "123", "cfg_scale": "5.5"}, + ) + + assert response.status_code == 200 + task_id = response.json()["task_id"] + + status_response = fs_client.get(f"/v1/tasks/{task_id}/status") + assert status_response.status_code == 200 + status = status_response.json() + assert status["resolution"] == "480p" + def test_upload_without_file(self, fs_client): """Test upload without file via form.""" response = fs_client.post("/v1/tasks/form", data={"prompt": "test without image"}) diff --git a/tests/unit/service/test_artifact_store.py b/tests/unit/service/test_artifact_store.py index 538ed4c..642b958 100644 --- a/tests/unit/service/test_artifact_store.py +++ b/tests/unit/service/test_artifact_store.py @@ -54,10 +54,24 @@ def test_artifact_store_resolves_task_scoped_outputs_for_download(tmp_path: Path output = store.output_path("result.mp4", media_type=MediaType.VIDEO, task_id="task-123") output.write_bytes(b"video") - assert store.resolve_output_file("result.mp4") == output + assert store.resolve_output_file("tasks/task-123/outputs/videos/result.mp4") == output assert store.resolve_output_file(output) == output +def test_artifact_store_does_not_resolve_bare_filenames_across_tasks(tmp_path: Path) -> None: + store = ArtifactStore(tmp_path) + first = store.output_path("result.png", media_type=MediaType.IMAGE, task_id="task-a") + second = store.output_path("result.png", media_type=MediaType.IMAGE, task_id="task-b") + first.write_bytes(b"first") + second.write_bytes(b"second") + + resolved = store.resolve_output_file("result.png") + + assert resolved == tmp_path / "outputs" / "result.png" + assert resolved not in {first, second} + assert not resolved.exists() + + def test_artifact_store_round_trips_local_artifact_ids(tmp_path: Path) -> None: store = ArtifactStore(tmp_path) output = store.output_path("result.png", media_type=MediaType.IMAGE, task_id="task-123") diff --git a/tests/unit/service/test_latent_cache_task_service.py b/tests/unit/service/test_latent_cache_task_service.py index 6d3357a..68477cd 100644 --- a/tests/unit/service/test_latent_cache_task_service.py +++ b/tests/unit/service/test_latent_cache_task_service.py @@ -14,7 +14,9 @@ def __init__(self, root: Path) -> None: self.output_dir = root / "outputs" self.output_dir.mkdir(parents=True, exist_ok=True) - def get_output_path(self, output_path: str, media_type): + def get_output_path(self, output_path: str, media_type, task_id: str | None = None): + if task_id is not None: + return self.output_dir / "tasks" / task_id / str(media_type) / output_path return self.output_dir / output_path diff --git a/tests/unit/service/test_service_refactor_phase0.py b/tests/unit/service/test_service_routes.py similarity index 90% rename from tests/unit/service/test_service_refactor_phase0.py rename to tests/unit/service/test_service_routes.py index 8be197e..f982f8d 100644 --- a/tests/unit/service/test_service_refactor_phase0.py +++ b/tests/unit/service/test_service_routes.py @@ -25,6 +25,11 @@ from telefuser.service_types import MediaType, TaskStatus +def _openapi_paths(app: object) -> set[str]: + openapi = app.openapi() + return set(openapi.get("paths", {})) + + def test_task_status_uses_shared_service_enum() -> None: assert CoreTaskStatus is TaskStatus assert TaskStatus.STREAMING.value == "streaming" @@ -91,26 +96,58 @@ class RunningPipeline: assert health["pipeline_ready"] is True +def test_service_status_ignores_mock_pool_status_attribute() -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.inference_service = Mock() + routes = ServiceRoutes(server) + + status = asyncio.run(routes.get_status()) + + assert status["execution_mode"] == "serial_single_pipeline" + assert "pool" not in status + + +def test_service_status_and_readiness_use_pipeline_pool_status_list() -> None: + class PoolPipeline: + is_running = True + + def pool_status(self) -> list[dict]: + return [{"id": 0, "device_ids": ["0"], "status": "idle"}] + + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.inference_service = PoolPipeline() + routes = ServiceRoutes(server) + + status = asyncio.run(routes.get_status()) + health = asyncio.run(routes.health_check()) + ready = asyncio.run(routes.readiness_check()) + + assert status["execution_mode"] == "concurrent_pipeline_pool" + assert status["pool"] == [{"id": 0, "device_ids": ["0"], "status": "idle"}] + assert health["ready"] is True + assert ready.status_code == 200 + + def test_stream_route_profile_exposes_only_stream_service_routes() -> None: server = ApiServer(task_manager=TaskManager(), enable_openai_api=True, route_profile="stream") - paths = {route.path for route in server.app.routes} + paths = _openapi_paths(server.app) assert "/v1/service/health" in paths assert "/v1/stream/sessions/{session_id}/status" in paths assert "/v1/tasks/create" not in paths assert "/v1/tasks/form" not in paths - assert "/v1/files/download/{file_path:path}" not in paths + assert "/v1/files/download/{file_path}" not in paths assert "/v1/images/generations" not in paths assert "/v1/videos" not in paths def test_request_response_route_profile_excludes_stream_routes() -> None: server = ApiServer(task_manager=TaskManager(), enable_openai_api=False, route_profile="request_response") - paths = {route.path for route in server.app.routes} + paths = _openapi_paths(server.app) assert "/v1/service/health" in paths assert "/v1/tasks/create" in paths - assert "/v1/files/download/{file_path:path}" in paths + assert "/v1/files/download/{file_path}" in paths assert "/v1/stream/sessions/{session_id}/status" not in paths @@ -120,12 +157,12 @@ def test_container_stream_app_uses_stream_route_profile() -> None: container.stream_pipeline_service.is_running = True app = container.get_api_app() - paths = {route.path for route in app.routes} + paths = _openapi_paths(app) assert "/v1/service/health" in paths assert "/v1/stream/sessions/{session_id}/status" in paths assert "/v1/tasks/create" not in paths - assert "/v1/files/download/{file_path:path}" not in paths + assert "/v1/files/download/{file_path}" not in paths def test_stream_session_close_logs_pipeline_close_failure(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/service/test_task_routes.py b/tests/unit/service/test_task_routes.py index 76e6a98..c0e028a 100644 --- a/tests/unit/service/test_task_routes.py +++ b/tests/unit/service/test_task_routes.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException, UploadFile +from starlette.datastructures import FormData from telefuser.service.api.routers.tasks import TaskRoutes from telefuser.service.api.schema import TaskResponse @@ -51,6 +52,12 @@ def _make_form_routes(tmp_path: Path, supported_tasks: tuple[str, ...]) -> TaskR return routes +def _make_form_request(data: dict[str, object]) -> Mock: + request = Mock() + request.form = AsyncMock(return_value=FormData(list(data.items()))) + return request + + def test_resolve_form_task_prefers_supported_pipeline_task_order() -> None: routes = _make_routes(("t2i", "i2i", "i2v")) @@ -146,7 +153,9 @@ def test_create_task_form_uses_file_service_upload_for_image(tmp_path: Path) -> routes = _make_form_routes(tmp_path, ("i2i",)) upload = _AsyncUpload("cat.png", "image/png", [b"cat"]) - asyncio.run(routes.create_task_form(first_image_file=upload, prompt="make it blue")) + asyncio.run( + routes.create_task_form(request=_make_form_request({"prompt": "make it blue"}), first_image_file=upload) + ) task_request = routes.api.task_app_service.submit.call_args.args[0] input_path = Path(task_request.first_image_path) @@ -159,10 +168,43 @@ def test_create_task_form_uses_file_service_upload_for_video(tmp_path: Path) -> routes = _make_form_routes(tmp_path, ("vc",)) upload = _AsyncUpload("clip.mp4", "video/mp4", [b"video"]) - asyncio.run(routes.create_task_form(first_image_file=upload, prompt="continue")) + asyncio.run(routes.create_task_form(request=_make_form_request({"prompt": "continue"}), first_image_file=upload)) task_request = routes.api.task_app_service.submit.call_args.args[0] input_path = Path(task_request.ref_video_path) assert task_request.task == "vc" assert input_path.parent == routes.api.file_service.input_video_dir assert input_path.read_bytes() == b"video" + + +def test_create_task_form_forwards_dynamic_form_parameters(tmp_path: Path) -> None: + routes = _make_form_routes(tmp_path, ("i2i",)) + upload = _AsyncUpload("cat.png", "image/png", [b"cat"]) + + asyncio.run( + routes.create_task_form( + request=_make_form_request( + { + "prompt": "make it blue", + "resolution": "480p", + "seed": "123", + "cfg_scale": "5.5", + "enable_feature": "true", + "extra_options": '{"mode": "fast"}', + } + ), + first_image_file=upload, + ) + ) + + task_request = routes.api.task_app_service.submit.call_args.args[0] + assert task_request.task == "i2i" + assert task_request.prompt == "make it blue" + assert task_request.resolution == "480p" + assert task_request.seed == 123 + assert task_request.cfg_scale == 5.5 + assert task_request.enable_feature is True + assert task_request.extra_options == {"mode": "fast"} + + explicit_fields = routes.api.task_app_service.submit.call_args.kwargs["explicit_fields"] + assert {"resolution", "cfg_scale", "enable_feature", "extra_options"}.issubset(explicit_fields)