diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..6efd355 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,48 @@ +# IDE +.idea/ +.vscode/ +*.swp +*.swo + +# Git +.git/ +.gitignore +.gitattributes + +# Node +web/node_modules/ +web/dist/ +web/electron/ +web/.nuxt/ +web/.output/ + +# Python +app/__pycache__/ +app/src/__pycache__/ +app/.venv/ +app/.uv/ +**/*.pyc +**/*.pyo +**/__pycache__/ +**/.pytest_cache/ +**/.ruff_cache/ + +# Docker +docker/data/ +docker/postgresql-16/data/ +docker/redis/data/ +docker/minio/data/ + +# Docs +docs/ +/*.md +*.drawio + +# Tests +app/tests/ +web/src/**/*.test.ts +web/src/**/*.spec.ts + +# Misc +.claude/ +.superpowers/ diff --git a/.gitignore b/.gitignore index d3abd03..799b46b 100644 --- a/.gitignore +++ b/.gitignore @@ -38,4 +38,4 @@ app/artifacts/ .playwright* .snow/ -/node_modules/ \ No newline at end of file +/node_modules/docker/.env.production diff --git a/app/.env.example b/app/.env.example index b466c14..55c8a0e 100644 --- a/app/.env.example +++ b/app/.env.example @@ -3,12 +3,22 @@ FF_DB_URI=postgresql://root:password@localhost:5432/fileflash # DATABASE_URL=postgresql://root:password@localhost:5432/fileflash APP_ENV=development +# Required when APP_ENV=production or APP_ENV=prod. +DEFAULT_ADMIN_USERNAME=administrator +DEFAULT_ADMIN_EMAIL=admin@fileflash.next +DEFAULT_ADMIN_PASSWORD=replace-with-32-bytes-or-longer-password + JWT_SECRET_KEY=please-set-at-least-32-bytes-secret-key # Optional dedicated HMAC key for token hash persistence. # Falls back to JWT_SECRET_KEY when omitted. # TOKEN_HASH_SECRET=please-set-at-least-32-bytes-and-different-from-jwt-secret ACCESS_TOKEN_EXPIRE_MINUTES=4320 REFRESH_TOKEN_EXPIRE_DAYS=7 +# Auth risk-control defaults (can be overridden per environment) +MAX_FAILED_LOGIN_ATTEMPTS=8 +ACCOUNT_LOCK_MINUTES=5 +REGISTER_RATE_LIMIT=12 +LOGIN_RATE_LIMIT=30 REDIS_URL=redis://localhost:6379/0 # RabbitMQ is optional in this stage, kept for future publisher swap. @@ -26,6 +36,7 @@ UPLOAD_CHUNK_SIZE_DEFAULT=5242880 UPLOAD_CHUNK_SIZE_MIN=1048576 UPLOAD_CHUNK_SIZE_MAX=16777216 UPLOAD_SINGLE_FILE_SIZE_MAX=5368709120 +UPLOAD_VERIFY_MERGED_OBJECT_HASH=false STARRED_ITEMS_LIMIT=20 UPLOAD_SESSION_TTL_HOURS=24 UPLOAD_TEMP_PREFIX=tmp @@ -54,9 +65,12 @@ AGENT_COMPACT_THRESHOLD=0.75 AGENT_USER_DAILY_LIMIT=50 AGENT_USER_CONCURRENT_LIMIT=2 AGENT_STAGING_TTL_SEC=86400 -AGENT_SSE_ENABLED=false +AGENT_SSE_ENABLED=true AGENT_LLM_PROVIDER=anthropic AGENT_LLM_MODEL=claude-sonnet-4-6 +AGENT_LLM_PLAN_MAX_TOKENS=8192 +# For Anthropic-compatible providers, e.g. DeepSeek: https://api.deepseek.com/anthropic +# AGENT_LLM_BASE_URL= # AGENT_LLM_API_KEY= AGENT_MCP_ENDPOINTS=[] diff --git a/app/README.md b/app/README.md index 1ca8092..6263967 100644 --- a/app/README.md +++ b/app/README.md @@ -3,26 +3,26 @@ Use one command to start backend API and file workers together: ```bash -uv run python -m fileflash.scripts.run_with_workers +uv run fileflash ``` Common options: ```bash # custom host/port -uv run python -m fileflash.scripts.run_with_workers --host 127.0.0.1 --port 8080 +uv run fileflash --host 127.0.0.1 --port 8080 # start multiple worker processes -uv run python -m fileflash.scripts.run_with_workers --worker-count 2 +uv run fileflash --worker-count 2 # API only (without workers) -uv run python -m fileflash.scripts.run_with_workers --no-worker +uv run fileflash --no-worker ``` Notes: - This runner starts `uvicorn fileflash.main:app` and `python -m fileflash.workers.consumer`. - If any subprocess exits, the runner stops all other subprocesses. -- If your environment resolves project scripts correctly, `uv run fileflash-dev` is equivalent. +- `uv run fileflash-dev` is kept as a backward-compatible alias. ## Database Migration Requirement @@ -31,6 +31,6 @@ Before starting API processes, ensure Flyway migrations are fully applied (inclu Recommended startup order: 1. Start PostgreSQL 2. Run Flyway migrate -3. Start API (`uv run fileflash`) or runner (`uv run fileflash-dev`) +3. Start backend (`uv run fileflash`; `uv run fileflash-dev` is also supported) If the schema is outdated, API startup will fail fast with an explicit compatibility error. diff --git a/app/pyproject.toml b/app/pyproject.toml index 6a8e1c7..8462a97 100644 --- a/app/pyproject.toml +++ b/app/pyproject.toml @@ -19,9 +19,10 @@ dependencies = [ "sqlalchemy>=2.0.48", "uvicorn>=0.42.0", "python-multipart>=0.0.24", + "anthropic>=0.104.1", ] [project.scripts] -fileflash = "fileflash.main:main" +fileflash = "fileflash.scripts.run_with_workers:main" fileflash-dev = "fileflash.scripts.run_with_workers:main" [dependency-groups] @@ -50,4 +51,4 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["src/fileflash"] \ No newline at end of file +packages = ["src/fileflash"] diff --git a/app/src/fileflash/agents/__init__.py b/app/src/fileflash/agents/__init__.py index b381960..98fd12d 100644 --- a/app/src/fileflash/agents/__init__.py +++ b/app/src/fileflash/agents/__init__.py @@ -14,9 +14,10 @@ ToolCall, ToolRouter, ) -from .runtime import ExecuteRunner, PlanRunner, SubagentRunner +from .runtime import AgentJobCanceled, ExecuteRunner, PlanRunner, SubagentRunner __all__ = [ + "AgentJobCanceled", "AgentEvent", "CheckpointStore", "ContextBudget", diff --git a/app/src/fileflash/agents/harness/__init__.py b/app/src/fileflash/agents/harness/__init__.py index 87e106a..acdcb08 100644 --- a/app/src/fileflash/agents/harness/__init__.py +++ b/app/src/fileflash/agents/harness/__init__.py @@ -6,6 +6,7 @@ from .policy import PolicyDecision, PolicyGuard from .prompt import PromptBuildRequest, PromptBuilder from .router import ToolCall, ToolRouter +from .tool_registry import REGISTRY, ToolContext, ToolRegistry, ToolSpec __all__ = [ "AgentEvent", @@ -20,6 +21,10 @@ "PolicyGuard", "PromptBuildRequest", "PromptBuilder", + "REGISTRY", "ToolCall", + "ToolContext", + "ToolRegistry", "ToolRouter", + "ToolSpec", ] diff --git a/app/src/fileflash/agents/harness/ask.py b/app/src/fileflash/agents/harness/ask.py new file mode 100644 index 0000000..258f4ac --- /dev/null +++ b/app/src/fileflash/agents/harness/ask.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import contextlib +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from ...repositories import AgentInboxMessageRepository +from .event_bus import AgentEventBus, AgentEventEnvelope, AgentEventStream + + +class AskTimedOut(Exception): + def __init__(self, *, ask_id: int) -> None: + super().__init__(f"Ask {ask_id} timed out") + self.ask_id = ask_id + + +class AskProtocol: + def __init__( + self, + *, + db: AsyncSession, + event_bus: AgentEventBus, + job_id: int, + ) -> None: + self._db = db + self._bus = event_bus + self._job_id = job_id + self._repo = AgentInboxMessageRepository(db) + self._waiters: dict[int, asyncio.Future[Any]] = {} + self._sub_ctx = None + self._sub_stream: AgentEventStream | None = None + self._sub_task: asyncio.Task[None] | None = None + + async def start(self) -> None: + self._sub_ctx = self._bus.subscribe(job_id=self._job_id) + self._sub_stream = await self._sub_ctx.__aenter__() + self._sub_task = asyncio.create_task(self._listen()) + + async def aclose(self) -> None: + if self._sub_task is not None: + self._sub_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._sub_task + if self._sub_ctx is not None: + await self._sub_ctx.__aexit__(None, None, None) + for future in self._waiters.values(): + if not future.done(): + future.cancel() + + async def ask( + self, + *, + prompt: str, + schema: dict[str, Any], + timeout_sec: float, + ) -> Any: + msg = await self._repo.create_ask( + job_id=self._job_id, + payload={"prompt": prompt, "schema": schema, "timeoutSec": timeout_sec}, + ) + await self._db.commit() + + ask_id = int(msg.inbox_message_id) + loop = asyncio.get_running_loop() + future: asyncio.Future[Any] = loop.create_future() + self._waiters[ask_id] = future + + await self._bus.publish( + AgentEventEnvelope( + job_id=self._job_id, + event_type="agent.ask", + payload={ + "messageId": str(ask_id), + "prompt": prompt, + "schema": schema, + "timeoutSec": timeout_sec, + }, + emitted_at=datetime.now(UTC), + ) + ) + + try: + value = await asyncio.wait_for(future, timeout=timeout_sec) + except TimeoutError as exc: + await self._repo.mark_timed_out( + inbox_message_id=ask_id, + answered_at=datetime.now(UTC), + ) + await self._db.commit() + raise AskTimedOut(ask_id=ask_id) from exc + finally: + self._waiters.pop(ask_id, None) + + await self._repo.mark_answered( + inbox_message_id=ask_id, + answered_at=datetime.now(UTC), + ) + await self._db.commit() + return value + + async def _listen(self) -> None: + assert self._sub_stream is not None + while True: + try: + envelope = await self._sub_stream.next(timeout=None) + except asyncio.CancelledError: + raise + except Exception: + continue + if envelope.event_type != "agent.inbox.reply": + continue + reply_to = envelope.payload.get("replyTo") + if reply_to is None: + continue + try: + ask_id = int(reply_to) + except (TypeError, ValueError): + continue + future = self._waiters.get(ask_id) + if future is None or future.done(): + continue + future.set_result(envelope.payload.get("value")) + + +__all__ = ["AskProtocol", "AskTimedOut"] diff --git a/app/src/fileflash/agents/harness/event_bus.py b/app/src/fileflash/agents/harness/event_bus.py new file mode 100644 index 0000000..05c73d0 --- /dev/null +++ b/app/src/fileflash/agents/harness/event_bus.py @@ -0,0 +1,194 @@ +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +from collections.abc import AsyncIterator +from contextlib import AbstractAsyncContextManager +from dataclasses import asdict, dataclass +from datetime import datetime +from typing import Any, Protocol + +from fastapi.encoders import jsonable_encoder +from redis.asyncio import Redis + +from ...core.settings import Settings, get_settings + +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class AgentEventEnvelope: + job_id: int + event_type: str + payload: dict[str, Any] + emitted_at: datetime + event_id: str | None = None + + def to_json(self) -> str: + body = jsonable_encoder(asdict(self)) + return json.dumps(body, ensure_ascii=False, separators=(",", ":")) + + @classmethod + def from_json(cls, raw: str) -> AgentEventEnvelope: + data = json.loads(raw) + return cls( + job_id=int(data["job_id"]), + event_type=str(data["event_type"]), + payload=dict(data.get("payload") or {}), + emitted_at=datetime.fromisoformat(data["emitted_at"]), + event_id=data.get("event_id"), + ) + + +class AgentEventStream(Protocol): + async def next(self, *, timeout: float | None = None) -> AgentEventEnvelope: ... + async def aclose(self) -> None: ... + + +class AgentEventBus(Protocol): + async def publish(self, envelope: AgentEventEnvelope) -> None: ... + + def subscribe( + self, + *, + job_id: int, + ) -> AbstractAsyncContextManager[AgentEventStream]: ... + + +@dataclass(slots=True) +class _InMemoryStream: + queue: asyncio.Queue[AgentEventEnvelope] + + async def next(self, *, timeout: float | None = None) -> AgentEventEnvelope: + if timeout is None: + return await self.queue.get() + return await asyncio.wait_for(self.queue.get(), timeout=timeout) + + async def aclose(self) -> None: + return None + + +class InMemoryAgentEventBus: + def __init__(self, *, buffer_size: int = 64) -> None: + self._buffer = buffer_size + self._subscribers: dict[int, list[asyncio.Queue[AgentEventEnvelope]]] = {} + + async def publish(self, envelope: AgentEventEnvelope) -> None: + queues = list(self._subscribers.get(envelope.job_id, [])) + for queue in queues: + if queue.full(): + logger.warning( + "InMemoryAgentEventBus dropped event: queue full job_id=%s", + envelope.job_id, + ) + continue + await queue.put(envelope) + + @contextlib.asynccontextmanager + async def subscribe(self, *, job_id: int) -> AsyncIterator[_InMemoryStream]: + queue: asyncio.Queue[AgentEventEnvelope] = asyncio.Queue(maxsize=self._buffer) + self._subscribers.setdefault(job_id, []).append(queue) + try: + yield _InMemoryStream(queue=queue) + finally: + subscribers = self._subscribers.get(job_id) + if subscribers is not None: + subscribers.remove(queue) + if not subscribers: + del self._subscribers[job_id] + + +class RedisAgentEventBus: + def __init__( + self, + *, + redis: Redis, + channel_prefix: str, + buffer_size: int = 64, + ) -> None: + self._redis = redis + self._channel_prefix = channel_prefix + self._buffer = buffer_size + + def _channel(self, job_id: int) -> str: + return f"{self._channel_prefix}:{job_id}:events" + + async def publish(self, envelope: AgentEventEnvelope) -> None: + await self._redis.publish(self._channel(envelope.job_id), envelope.to_json()) + + @contextlib.asynccontextmanager + async def subscribe(self, *, job_id: int) -> AsyncIterator[_RedisStream]: + pubsub = self._redis.pubsub() + channel = self._channel(job_id) + await pubsub.subscribe(channel) + stream = _RedisStream(pubsub=pubsub) + try: + yield stream + finally: + await pubsub.unsubscribe(channel) + await pubsub.aclose() + + +@dataclass(slots=True) +class _RedisStream: + pubsub: Any + + async def next(self, *, timeout: float | None = None) -> AgentEventEnvelope: + if timeout is None: + async for message in self.pubsub.listen(): + envelope = _envelope_from_redis_message(message) + if envelope is not None: + return envelope + else: + message = await self.pubsub.get_message( + ignore_subscribe_messages=True, + timeout=timeout, + ) + envelope = _envelope_from_redis_message(message) + if envelope is not None: + return envelope + raise TimeoutError("No event within timeout") + + async def aclose(self) -> None: + await self.pubsub.aclose() + + +def _envelope_from_redis_message(message: Any) -> AgentEventEnvelope | None: + if message is None: + return None + message_type = message.get("type") + if message_type not in {"message", "pmessage"}: + return None + data = message.get("data") + if isinstance(data, bytes): + data = data.decode("utf-8") + return AgentEventEnvelope.from_json(str(data)) + + +def build_agent_event_bus( + *, + settings: Settings | None = None, + redis: Redis | None = None, +) -> AgentEventBus: + cfg = settings or get_settings() + if redis is None: + if not cfg.redis_url: + return InMemoryAgentEventBus(buffer_size=cfg.agent_event_bus_buffer_size) + redis = Redis.from_url(cfg.redis_url, decode_responses=True) + return RedisAgentEventBus( + redis=redis, + channel_prefix=cfg.agent_event_channel_prefix, + buffer_size=cfg.agent_event_bus_buffer_size, + ) + + +__all__ = [ + "AgentEventBus", + "AgentEventEnvelope", + "AgentEventStream", + "InMemoryAgentEventBus", + "RedisAgentEventBus", + "build_agent_event_bus", +] diff --git a/app/src/fileflash/agents/harness/events.py b/app/src/fileflash/agents/harness/events.py index 8bf7a94..3fbfbb2 100644 --- a/app/src/fileflash/agents/harness/events.py +++ b/app/src/fileflash/agents/harness/events.py @@ -1,15 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Any +from .event_bus import AgentEventBus as EventBus +from .event_bus import AgentEventEnvelope as AgentEvent - -@dataclass(slots=True) -class AgentEvent: - event_type: str - payload: dict[str, Any] - - -class EventBus: - async def publish(self, event: AgentEvent) -> None: - raise NotImplementedError("EventBus is scaffolded only in this stage") +__all__ = ["AgentEvent", "EventBus"] diff --git a/app/src/fileflash/agents/harness/inbox.py b/app/src/fileflash/agents/harness/inbox.py new file mode 100644 index 0000000..ac7f45c --- /dev/null +++ b/app/src/fileflash/agents/harness/inbox.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models.enums import AgentInboxKind, AgentInboxStatus +from ...repositories import AgentInboxMessageRepository +from .event_bus import AgentEventBus, AgentEventEnvelope + +_INBOX_EVENT_TYPES: dict[AgentInboxKind, str] = { + AgentInboxKind.REPLY: "agent.inbox.reply", + AgentInboxKind.CONTROL_PAUSE: "agent.inbox.control", + AgentInboxKind.CONTROL_RESUME: "agent.inbox.control", + AgentInboxKind.CONTROL_APPROVE: "agent.inbox.control", + AgentInboxKind.CONTROL_DENY: "agent.inbox.control", + AgentInboxKind.CONTROL_SKIP: "agent.inbox.control", + AgentInboxKind.CONTROL_CANCEL: "agent.inbox.control", +} + + +class AgentInbox: + def __init__(self, *, db: AsyncSession, event_bus: AgentEventBus) -> None: + self._db = db + self._bus = event_bus + self._repo = AgentInboxMessageRepository(db) + + async def handle( + self, + *, + job_id: int, + kind: AgentInboxKind, + payload: dict[str, Any], + reply_to_id: int | None = None, + ): + if kind == AgentInboxKind.REPLY: + if reply_to_id is None: + raise ValueError("reply requires reply_to_id") + ask = await self._repo.get_ask(inbox_message_id=reply_to_id) + if ask is None: + raise ValueError(f"ask {reply_to_id} not found") + if int(ask.job_id) != job_id: + raise ValueError(f"ask {reply_to_id} belongs to a different job") + if ask.status != AgentInboxStatus.WAITING: + raise ValueError(f"ask {reply_to_id} is not waiting") + + msg = await self._repo.record_user_message( + job_id=job_id, + kind=kind, + payload=payload, + reply_to_id=reply_to_id, + ) + event_type = _INBOX_EVENT_TYPES[kind] + envelope_payload: dict[str, Any] = { + "kind": kind.value, + "messageId": str(msg.inbox_message_id), + } + if reply_to_id is not None: + envelope_payload["replyTo"] = str(reply_to_id) + if "value" in payload: + envelope_payload["value"] = payload["value"] + if "metadata" in payload: + envelope_payload["metadata"] = payload["metadata"] + await self._bus.publish( + AgentEventEnvelope( + job_id=job_id, + event_type=event_type, + payload=envelope_payload, + emitted_at=datetime.now(UTC), + ) + ) + return msg + + +__all__ = ["AgentInbox"] diff --git a/app/src/fileflash/agents/harness/policy.py b/app/src/fileflash/agents/harness/policy.py index 9e1e4e4..730322b 100644 --- a/app/src/fileflash/agents/harness/policy.py +++ b/app/src/fileflash/agents/harness/policy.py @@ -2,6 +2,9 @@ from dataclasses import dataclass, field +from ...schemas.agent import AgentProposedAction +from .tool_registry import REGISTRY + @dataclass(slots=True) class PolicyDecision: @@ -9,6 +12,55 @@ class PolicyDecision: reasons: list[str] = field(default_factory=list) +def classify_tool_side_effect(tool_name: str) -> str: + try: + return REGISTRY.get(tool_name).side_effect + except KeyError: + return "write" + + +def classify_tool_risk(tool_name: str) -> str: + try: + return REGISTRY.get(tool_name).risk_level + except KeyError: + return "high" + + +def normalize_action_risk(action: AgentProposedAction) -> AgentProposedAction: + risk_level = classify_tool_risk(action.tool) + requires_confirmation = action.requires_confirmation or risk_level == "high" + reason = action.confirmation_reason + if risk_level == "high" and not reason: + reason = ( + "Deleting files or folders is a high-risk action and requires explicit confirmation." + ) + return action.model_copy( + update={ + "side_effect": classify_tool_side_effect(action.tool), + "risk_level": risk_level, + "requires_confirmation": requires_confirmation, + "confirmation_reason": reason, + } + ) + + class PolicyGuard: - async def evaluate_tool_call(self, *args, **kwargs) -> PolicyDecision: - raise NotImplementedError("PolicyGuard is scaffolded only in this stage") + async def evaluate_tool_call( + self, + *, + tool_name: str, + high_risk_confirmed: bool = False, + ) -> PolicyDecision: + try: + REGISTRY.get(tool_name) + except KeyError: + return PolicyDecision( + allowed=False, + reasons=[f"Unsupported agent tool: {tool_name}"], + ) + if classify_tool_risk(tool_name) == "high" and not high_risk_confirmed: + return PolicyDecision( + allowed=False, + reasons=["High-risk delete action requires explicit user confirmation."], + ) + return PolicyDecision(allowed=True) diff --git a/app/src/fileflash/agents/harness/router.py b/app/src/fileflash/agents/harness/router.py index 1a8b3d4..2b3b639 100644 --- a/app/src/fileflash/agents/harness/router.py +++ b/app/src/fileflash/agents/harness/router.py @@ -3,6 +3,13 @@ from dataclasses import dataclass from typing import Any +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...services.file import FileService +from ...services.folder import FolderService +from .tool_registry import REGISTRY, ToolContext + @dataclass(slots=True) class ToolCall: @@ -11,5 +18,30 @@ class ToolCall: class ToolRouter: + def __init__(self, *, db: AsyncSession, user_id: int) -> None: + self.db = db + self.user_id = user_id + self.file_service = FileService(db=db) + self.folder_service = FolderService(db=db) + async def dispatch(self, call: ToolCall) -> dict[str, Any]: - raise NotImplementedError("ToolRouter is scaffolded only in this stage") + tool_name = str(call.tool_name or "").strip() + try: + spec = REGISTRY.get(tool_name) + except KeyError as exc: + raise ApiError( + status_code=400, + code=400, + message=f"Unsupported agent tool: {tool_name}", + ) from exc + + ctx = ToolContext( + db=self.db, + user_id=self.user_id, + file_service=self.file_service, + folder_service=self.folder_service, + ) + return await spec.handler(ctx, dict(call.arguments or {})) + + +__all__ = ["ToolCall", "ToolRouter"] diff --git a/app/src/fileflash/agents/harness/tool_registry.py b/app/src/fileflash/agents/harness/tool_registry.py new file mode 100644 index 0000000..c508a86 --- /dev/null +++ b/app/src/fileflash/agents/harness/tool_registry.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import re +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, Literal + +from sqlalchemy.ext.asyncio import AsyncSession + +ToolSideEffect = Literal["read", "write"] +ToolRiskLevel = Literal["low", "medium", "high"] + + +@dataclass(slots=True) +class ToolContext: + db: AsyncSession + user_id: int + file_service: Any + folder_service: Any + + +ToolHandler = Callable[[ToolContext, dict[str, Any]], Awaitable[dict[str, Any]]] +ToolAnswerFormatter = Callable[[dict[str, Any]], str | None] + + +@dataclass(frozen=True, slots=True) +class ToolSpec: + name: str + description: str + input_schema: dict[str, Any] + side_effect: ToolSideEffect + risk_level: ToolRiskLevel + requires_confirmation: bool + handler: ToolHandler + anthropic_name: str | None = None + answer_formatter: ToolAnswerFormatter | None = None + + def __post_init__(self) -> None: + if self.anthropic_name is None: + object.__setattr__(self, "anthropic_name", _to_provider_tool_name(self.name)) + + def to_anthropic_tool(self) -> dict[str, Any]: + return { + "name": self.anthropic_name, + "description": self.description, + "input_schema": self.input_schema, + "internalName": self.name, + } + + def to_planner_schema(self) -> dict[str, Any]: + return { + "tool": self.name, + "providerTool": self.anthropic_name, + "description": self.description, + "inputSchema": self.input_schema, + "sideEffect": self.side_effect, + "riskLevel": self.risk_level, + "requiresConfirmation": self.requires_confirmation, + } + + +class ToolRegistry: + def __init__(self) -> None: + self._by_name: dict[str, ToolSpec] = {} + self._by_provider_name: dict[str, ToolSpec] = {} + + def register(self, spec: ToolSpec) -> ToolSpec: + if spec.name in self._by_name: + raise ValueError(f"Agent tool already registered: {spec.name}") + provider_name = str(spec.anthropic_name or "") + if provider_name in self._by_provider_name: + raise ValueError(f"Agent provider tool already registered: {provider_name}") + self._by_name[spec.name] = spec + self._by_provider_name[provider_name] = spec + return spec + + def get(self, name: str) -> ToolSpec: + ensure_builtin_tools_registered() + return self._by_name[name] + + def get_by_provider_name(self, name: str) -> ToolSpec: + ensure_builtin_tools_registered() + return self._by_provider_name[name] + + def all(self) -> list[ToolSpec]: + ensure_builtin_tools_registered() + return list(self._by_name.values()) + + def all_names(self) -> tuple[str, ...]: + return tuple(spec.name for spec in self.all()) + + def unknown_names(self, names: list[str] | tuple[str, ...]) -> list[str]: + ensure_builtin_tools_registered() + return [name for name in names if name not in self._by_name] + + def validate_names(self, names: list[str] | tuple[str, ...]) -> tuple[str, ...]: + unknown = self.unknown_names(names) + if unknown: + raise ValueError(f"Unknown agent tools: {', '.join(sorted(unknown))}") + return tuple(names) + + def schemas_for(self, names: list[str] | tuple[str, ...]) -> list[dict[str, Any]]: + return [self.get(name).to_planner_schema() for name in names] + + def anthropic_tools_for(self, names: list[str] | tuple[str, ...]) -> list[dict[str, Any]]: + return [self.get(name).to_anthropic_tool() for name in names] + + +REGISTRY = ToolRegistry() +_BUILTINS_REGISTERED = False + + +def ensure_builtin_tools_registered() -> None: + global _BUILTINS_REGISTERED + if _BUILTINS_REGISTERED: + return + from .. import tools as _tools # noqa: F401 + _BUILTINS_REGISTERED = True + + +def _to_provider_tool_name(name: str) -> str: + value = name.replace(".", "_").replace("-", "_") + value = re.sub(r"(? None: + self.settings = settings or get_settings() + self.policy_guard = policy_guard or PolicyGuard() + self.event_bus = event_bus + self.answer_client = answer_client or AnthropicPlannerClient(settings=self.settings) + + async def run(self, *, db: AsyncSession, job: BackgroundJob) -> AgentExecutionResult: + ask: AskProtocol | None = None + if self.event_bus is not None: + ask = AskProtocol(db=db, event_bus=self.event_bus, job_id=int(job.job_id)) + await ask.start() + try: + return await self._run(db=db, job=job, ask=ask) + finally: + if ask is not None: + await ask.aclose() + + async def _run( + self, + *, + db: AsyncSession, + job: BackgroundJob, + ask: AskProtocol | None, + ) -> AgentExecutionResult: + if job.requested_by is None: + raise ApiError(status_code=400, code=400, message="Agent job is missing requestedBy") + request = ExecuteAgentRequest.model_validate(dict(job.payload or {})) + plan_job_id = _parse_job_id(request.plan_job_id) + plan_repo = AgentPlanRepository(db) + work_sessions = AgentWorkSessionRepository(db) + plan = await plan_repo.get_for_execute_binding( + job_id=plan_job_id, + user_id=int(job.requested_by), + plan_hash=request.plan_hash, + ) + if plan is None: + raise ApiError(status_code=409, code=409, message="Plan hash mismatch") + + actions = [ + AgentProposedAction.model_validate(item) + for item in (plan.proposed_actions_json or []) + ] + high_risk_confirmed = bool(request.approval.high_risk_confirmed) + router = ToolRouter(db=db, user_id=int(job.requested_by)) + action_logs = AgentActionLogRepository(db) + step_outputs: dict[int, dict[str, Any]] = {} + applied = 0 + warnings: list[str] = [] + await work_sessions.create_for_job( + job_id=int(job.job_id), + user_id=int(job.requested_by), + checkpoint_json={"planJobId": str(plan_job_id), "planHash": request.plan_hash}, + ) + await db.commit() + + inbox_repo = AgentInboxMessageRepository(db) if self.event_bus is not None else None + paused = False + for action in actions: + await db.refresh(job) + if job.cancel_requested_at is not None: + raise AgentJobCanceled() + if inbox_repo is not None: + paused, skip_current = await self._handle_step_boundary_controls( + db=db, + job=job, + inbox_repo=inbox_repo, + action=action, + warnings=warnings, + paused=paused, + ) + if skip_current: + continue + + decision = await self.policy_guard.evaluate_tool_call( + tool_name=action.tool, + high_risk_confirmed=high_risk_confirmed, + ) + if not decision.allowed: + raise ApiError( + status_code=409, + code=409, + message="High-risk action requires confirmation", + data={"reasons": decision.reasons, "step": action.step, "tool": action.tool}, + ) + + started = datetime.now(UTC) + try: + resolved_input = _resolve_references(action.input, step_outputs) + except Exception as exc: + duration_ms = int((datetime.now(UTC) - started).total_seconds() * 1000) + await action_logs.append_step( + job_id=int(job.job_id), + step_no=action.step, + tool_name=action.tool, + inputs_json=action.input, + status="failed", + started_at=started, + ) + await action_logs.finish_step( + job_id=int(job.job_id), + step_no=action.step, + outputs_json={}, + status="failed", + duration_ms=duration_ms, + error_message=f"{type(exc).__name__}: {exc}"[:2000], + ) + await db.commit() + await self._publish_tool( + "tool.failed", + job_id=int(job.job_id), + step=action.step, + tool=action.tool, + payload={"errorMessage": f"{type(exc).__name__}: {exc}"[:2000]}, + ) + raise + + await action_logs.append_step( + job_id=int(job.job_id), + step_no=action.step, + tool_name=action.tool, + inputs_json=resolved_input, + status="running", + started_at=started, + ) + await db.commit() + await self._publish_tool( + "tool.started", + job_id=int(job.job_id), + step=action.step, + tool=action.tool, + payload={"input": resolved_input}, + emitted_at=started, + ) + + try: + output = await router.dispatch( + ToolCall(tool_name=action.tool, arguments=resolved_input) + ) + except Exception as exc: + await db.rollback() + duration_ms = int((datetime.now(UTC) - started).total_seconds() * 1000) + await action_logs.finish_step( + job_id=int(job.job_id), + step_no=action.step, + outputs_json={}, + status="failed", + duration_ms=duration_ms, + error_message=f"{type(exc).__name__}: {exc}"[:2000], + ) + await db.commit() + await self._publish_tool( + "tool.failed", + job_id=int(job.job_id), + step=action.step, + tool=action.tool, + payload={"errorMessage": f"{type(exc).__name__}: {exc}"[:2000]}, + ) + raise + + safe_output = jsonable_encoder(output) + duration_ms = int((datetime.now(UTC) - started).total_seconds() * 1000) + await action_logs.finish_step( + job_id=int(job.job_id), + step_no=action.step, + outputs_json=safe_output, + status="succeeded", + duration_ms=duration_ms, + ) + await db.commit() + await self._publish_tool( + "tool.succeeded", + job_id=int(job.job_id), + step=action.step, + tool=action.tool, + payload={"output": safe_output, "durationMs": duration_ms}, + ) + step_outputs[action.step] = safe_output + applied += 1 + + skipped = max(0, len(actions) - applied) + if skipped: + warnings.append(f"{skipped} action(s) were skipped.") + await work_sessions.close_session(job_id=int(job.job_id), status="closed") + await db.commit() + answer = await _build_execution_answer( + task_input=str(getattr(plan, "input_text", "") or ""), + actions=actions, + step_outputs=step_outputs, + answer_client=self.answer_client, + ) + return AgentExecutionResult( + plan_job_id=str(plan_job_id), + execute_job_id=str(job.job_id), + summary=f"Execution completed with {applied} applied action(s).", + answer=answer, + applied_actions=applied, + skipped_actions=skipped, + warnings=warnings, + finished_at=datetime.now(UTC), + ) + + async def _handle_step_boundary_controls( + self, + *, + db: AsyncSession, + job: BackgroundJob, + inbox_repo: AgentInboxMessageRepository, + action: AgentProposedAction, + warnings: list[str], + paused: bool, + ) -> tuple[bool, bool]: + while True: + skip_current = False + pending = await inbox_repo.list_pending_controls(job_id=int(job.job_id)) + for ctrl in pending: + kind = AgentInboxKind(ctrl.kind) + if kind == AgentInboxKind.CONTROL_CANCEL: + await inbox_repo.mark_dropped(inbox_message_id=int(ctrl.inbox_message_id)) + job.cancel_requested_at = datetime.now(UTC) + await db.commit() + raise AgentJobCanceled() + if kind == AgentInboxKind.CONTROL_PAUSE: + paused = True + await inbox_repo.mark_dropped(inbox_message_id=int(ctrl.inbox_message_id)) + await self._publish_state("agent.paused", job_id=int(job.job_id)) + elif kind == AgentInboxKind.CONTROL_RESUME: + paused = False + await inbox_repo.mark_dropped(inbox_message_id=int(ctrl.inbox_message_id)) + await self._publish_state("agent.resumed", job_id=int(job.job_id)) + elif kind == AgentInboxKind.CONTROL_SKIP: + await inbox_repo.mark_dropped(inbox_message_id=int(ctrl.inbox_message_id)) + warnings.append(f"Step {action.step} skipped by user") + skip_current = True + else: + await inbox_repo.mark_dropped(inbox_message_id=int(ctrl.inbox_message_id)) + await db.commit() + if skip_current: + return paused, True + if not paused: + return paused, False + await asyncio.sleep(0.1) + + async def _publish_state(self, event_type: str, *, job_id: int) -> None: + if self.event_bus is None: + return + try: + await self.event_bus.publish( + AgentEventEnvelope( + job_id=job_id, + event_type=event_type, + payload={}, + emitted_at=datetime.now(UTC), + ) + ) + except Exception: + logger.exception( + "Failed to publish state event jobId=%s eventType=%s", + job_id, + event_type, + ) + + async def _publish_tool( + self, + event_type: str, + *, + job_id: int, + step: int, + tool: str, + payload: dict[str, Any], + emitted_at: datetime | None = None, + ) -> None: + if self.event_bus is None: + return + try: + await self.event_bus.publish( + AgentEventEnvelope( + job_id=job_id, + event_type=event_type, + payload={"step": int(step), "tool": str(tool), **payload}, + emitted_at=emitted_at or datetime.now(UTC), + ) + ) + except Exception: + logger.exception( + "Failed to publish tool event jobId=%s eventType=%s step=%s tool=%s", + job_id, + event_type, + step, + tool, + ) + + +def _parse_job_id(raw: str) -> int: + try: + value = int(raw) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid planJobId") from exc + if value <= 0: + raise ApiError(status_code=400, code=400, message="Invalid planJobId") + return value + + +def _resolve_references( + value: Any, + step_outputs: dict[int, dict[str, Any]], + *, + field_name: str | None = None, + field_path: str = "input", +) -> Any: + if isinstance(value, str): + reference = parse_step_reference(value) + if reference is not None: + step, path = reference + current: Any = step_outputs.get(step) + for part in path: + if isinstance(current, dict): + current = current.get(part) + else: + current = None + if current is None: + raise ApiError( + status_code=409, + code=409, + message=f"Unable to resolve tool reference: {value}", + ) + return current + if is_symbolic_id_placeholder(value=value, field_name=field_name): + raise ApiError( + status_code=409, + code=409, + message=( + f"Invalid tool input at '{field_path}': unresolved placeholder '{value}'. " + "Use '$stepN.field' references." + ), + ) + return value + if isinstance(value, list): + return [ + _resolve_references( + item, + step_outputs, + field_name=field_name, + field_path=f"{field_path}[{index}]", + ) + for index, item in enumerate(value) + ] + if isinstance(value, dict): + return { + key: _resolve_references( + item, + step_outputs, + field_name=key, + field_path=f"{field_path}.{key}", + ) + for key, item in value.items() + } + return value + + +async def _build_execution_answer( + *, + task_input: str = "", + actions: list[AgentProposedAction], + step_outputs: dict[int, dict[str, Any]], + answer_client: AnswerClient, +) -> str | None: + if not actions: + return None + user_prompt = _answer_user_prompt( + task_input=task_input, + actions=actions, + step_outputs=step_outputs, + ) + text = await answer_client.create_answer( + system_prompt=_answer_system_prompt(), + user_prompt=user_prompt, + max_tokens=640, + reasoning_effort="low", + ) + answer = _normalize_answer(text) + if answer is None: + raise ApiError(status_code=502, code=502, message="Agent answer model returned empty response") + return answer + + +def _answer_system_prompt() -> str: + return ( + "You are FileFlash execution answer generator. " + "Only describe results that are present in tool outputs. " + "Do not invent filenames, counts, or paths. " + "Keep the response concise and user-facing in the same language as the user input." + ) + + +def _answer_user_prompt( + *, + task_input: str, + actions: list[AgentProposedAction], + step_outputs: dict[int, dict[str, Any]], +) -> str: + payload_actions: list[dict[str, Any]] = [] + for action in sorted(actions, key=lambda item: item.step): + payload_actions.append( + { + "step": action.step, + "tool": action.tool, + "sideEffect": action.side_effect, + "input": action.input, + "output": _compact_output(step_outputs.get(action.step)), + } + ) + payload = { + "task": task_input, + "actions": payload_actions, + "responseGuidance": { + "includeNamesWhenAvailable": True, + "mentionTruncationWhenProvided": True, + "ifAmbiguous": "state candidate count and ask for clarification", + }, + } + return json.dumps(payload, ensure_ascii=False, sort_keys=True) + + +def _compact_output(value: dict[str, Any] | None) -> dict[str, Any] | None: + if not isinstance(value, dict): + return None + text = json.dumps(value, ensure_ascii=False, separators=(",", ":")) + if len(text) <= 12_000: + return value + compact = dict(value) + compact["truncated"] = True + compact["truncatedFields"] = sorted(compact.keys())[:16] + compact.pop("items", None) + compact.pop("sampleItems", None) + return compact + + +def _normalize_answer(text: str) -> str | None: + candidate = str(text or "").strip() + if not candidate: + return None + candidate = " ".join(candidate.split()) + if len(candidate) > 1200: + candidate = candidate[:1200].rstrip() + "…" + return candidate diff --git a/app/src/fileflash/agents/runtime/llm.py b/app/src/fileflash/agents/runtime/llm.py new file mode 100644 index 0000000..12a3ddc --- /dev/null +++ b/app/src/fileflash/agents/runtime/llm.py @@ -0,0 +1,721 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Awaitable, Callable +from typing import Any, Protocol + +import anthropic +from anthropic import AsyncAnthropic + +from ...core.errors import ApiError +from ...core.settings import Settings + +ToolExecutor = Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] +logger = logging.getLogger(__name__) + + +class PlannerClient(Protocol): + async def create_plan( + self, + *, + system_prompt: str, + user_prompt: str, + max_tokens: int, + reasoning_effort: str = "adaptive", + tools: list[dict[str, Any]] | None = None, + tool_executor: ToolExecutor | None = None, + max_tool_roundtrips: int = 4, + ) -> dict[str, Any]: ... + + +class AnswerClient(Protocol): + async def create_answer( + self, + *, + system_prompt: str, + user_prompt: str, + max_tokens: int, + reasoning_effort: str = "adaptive", + ) -> str: ... + + +class AnthropicPlannerClient: + def __init__(self, *, settings: Settings, client: AsyncAnthropic | None = None) -> None: + self.settings = settings + self._client = client + + async def create_plan( + self, + *, + system_prompt: str, + user_prompt: str, + max_tokens: int, + reasoning_effort: str = "adaptive", + tools: list[dict[str, Any]] | None = None, + tool_executor: ToolExecutor | None = None, + max_tool_roundtrips: int = 4, + ) -> dict[str, Any]: + api_key = (self.settings.agent_llm_api_key or "").strip() + if not api_key: + raise ApiError(status_code=503, code=503, message="Agent LLM API key is not configured") + plan_token_cap = _safe_plan_token_cap(self.settings) + + request_kwargs: dict[str, Any] = { + "model": self.settings.agent_llm_model, + "max_tokens": min(max_tokens, plan_token_cap), + "system": system_prompt, + "messages": [{"role": "user", "content": user_prompt}], + "timeout": 60.0, + } + tool_name_map = _tool_name_map(tools or []) + if tools: + request_kwargs["tools"] = _anthropic_tools_payload(tools) + request_kwargs["tool_choice"] = {"type": "auto"} + request_kwargs.update(_reasoning_params(reasoning_effort)) + + try: + parsed, usage = await self._request_and_parse_plan( + api_key=api_key, + request_kwargs=request_kwargs, + tool_name_map=tool_name_map, + tool_executor=tool_executor, + max_tool_roundtrips=max_tool_roundtrips, + ) + except ApiError as first_error: + if not _is_retryable_output_error(first_error): + raise + logger.warning( + "Planner LLM retrying attempt=%s reason=%s degraded=%s jsonOnly=%s", + 2, + first_error.message, + True, + False, + ) + degraded_kwargs = _degraded_plan_request_kwargs(request_kwargs) + try: + parsed, usage = await self._request_and_parse_plan( + api_key=api_key, + request_kwargs=degraded_kwargs, + tool_name_map=tool_name_map, + tool_executor=tool_executor, + max_tool_roundtrips=max_tool_roundtrips, + ) + except ApiError as second_error: + if not _is_retryable_output_error(second_error): + raise + logger.warning( + "Planner LLM retrying attempt=%s reason=%s degraded=%s jsonOnly=%s", + 3, + second_error.message, + True, + True, + ) + strict_kwargs = _strict_json_retry_kwargs( + degraded_kwargs, + max_tokens_cap=plan_token_cap, + ) + parsed, usage = await self._request_and_parse_plan( + api_key=api_key, + request_kwargs=strict_kwargs, + tool_name_map=tool_name_map, + tool_executor=tool_executor, + max_tool_roundtrips=max_tool_roundtrips, + ) + if isinstance(usage, dict): + parsed["_usage"] = usage + return parsed + + async def create_answer( + self, + *, + system_prompt: str, + user_prompt: str, + max_tokens: int, + reasoning_effort: str = "adaptive", + ) -> str: + api_key = (self.settings.agent_llm_api_key or "").strip() + if not api_key: + raise ApiError(status_code=503, code=503, message="Agent LLM API key is not configured") + request_kwargs: dict[str, Any] = { + "model": self.settings.agent_llm_model, + "max_tokens": min(max_tokens, 1024), + "system": system_prompt, + "messages": [{"role": "user", "content": user_prompt}], + "timeout": 60.0, + } + request_kwargs.update(_reasoning_params(reasoning_effort)) + message = await self._request_plan(api_key=api_key, request_kwargs=request_kwargs) + try: + return _extract_text(message) + except ApiError as exc: + if not _is_retryable_output_error(exc): + raise + degraded_kwargs = dict(request_kwargs) + degraded_kwargs.pop("thinking", None) + degraded_kwargs.pop("output_config", None) + message = await self._request_plan(api_key=api_key, request_kwargs=degraded_kwargs) + return _extract_text(message) + + async def _request_plan(self, *, api_key: str, request_kwargs: dict[str, Any]) -> Any: + try: + return await self._get_client(api_key).messages.create(**request_kwargs) + except anthropic.APIStatusError as exc: + raise ApiError( + status_code=503, + code=503, + message="Agent LLM request failed", + data={"status": exc.status_code, "details": _response_details(exc)}, + ) from exc + except ( + anthropic.APIConnectionError, + anthropic.APITimeoutError, + anthropic.APIError, + ) as exc: + raise ApiError( + status_code=503, + code=503, + message=f"Agent LLM request failed: {type(exc).__name__}", + ) from exc + + def _get_client(self, api_key: str) -> AsyncAnthropic: + if self._client is None: + base_url = (self.settings.agent_llm_base_url or "").strip() or None + self._client = AsyncAnthropic( + api_key=api_key, + base_url=base_url, + timeout=60.0, + max_retries=0, + ) + return self._client + + async def _parse_plan_response( + self, + *, + api_key: str, + request_kwargs: dict[str, Any], + message: Any, + tool_name_map: dict[str, str], + tool_executor: ToolExecutor | None, + max_tool_roundtrips: int, + ) -> tuple[dict[str, Any], dict[str, Any] | None]: + if tool_executor is None: + return _parse_plan_message(message, tool_name_map=tool_name_map) + tool_calls = _extract_tool_use_calls(message=message, tool_name_map=tool_name_map) + if not tool_calls: + return _parse_plan_message(message, tool_name_map=tool_name_map) + return await self._run_tool_loop( + api_key=api_key, + request_kwargs=request_kwargs, + initial_message=message, + tool_name_map=tool_name_map, + tool_executor=tool_executor, + max_tool_roundtrips=max_tool_roundtrips, + ) + + async def _request_and_parse_plan( + self, + *, + api_key: str, + request_kwargs: dict[str, Any], + tool_name_map: dict[str, str], + tool_executor: ToolExecutor | None, + max_tool_roundtrips: int, + ) -> tuple[dict[str, Any], dict[str, Any] | None]: + message = await self._request_plan(api_key=api_key, request_kwargs=request_kwargs) + return await self._parse_plan_response( + api_key=api_key, + request_kwargs=request_kwargs, + message=message, + tool_name_map=tool_name_map, + tool_executor=tool_executor, + max_tool_roundtrips=max_tool_roundtrips, + ) + + async def _run_tool_loop( + self, + *, + api_key: str, + request_kwargs: dict[str, Any], + initial_message: Any, + tool_name_map: dict[str, str], + tool_executor: ToolExecutor, + max_tool_roundtrips: int, + ) -> tuple[dict[str, Any], dict[str, Any] | None]: + max_rounds = max(1, min(int(max_tool_roundtrips or 0), 12)) + base_messages = request_kwargs.get("messages") + if not isinstance(base_messages, list): + raise ApiError(status_code=502, code=502, message="Agent LLM returned an invalid response") + messages: list[dict[str, Any]] = list(base_messages) + usage_total: dict[str, int] = {} + current_message = initial_message + + for _ in range(max_rounds): + usage_total = _merge_usage_totals(usage_total, _usage_payload(current_message)) + tool_calls = _extract_tool_use_calls(message=current_message, tool_name_map=tool_name_map) + if not tool_calls: + parsed, _ = _parse_plan_message(current_message, tool_name_map=tool_name_map) + return parsed, usage_total or None + + assistant_content = _content_block_mappings(current_message) + if assistant_content: + messages.append({"role": "assistant", "content": assistant_content}) + tool_results: list[dict[str, Any]] = [] + for call in tool_calls: + tool_output = await tool_executor(call["tool"], call["input"]) + tool_result: dict[str, Any] = { + "type": "tool_result", + "tool_use_id": call["id"], + "content": _tool_result_content(tool_output), + } + if _is_tool_error_payload(tool_output): + tool_result["is_error"] = True + tool_results.append(tool_result) + messages.append({"role": "user", "content": tool_results}) + loop_kwargs = dict(request_kwargs) + loop_kwargs["messages"] = messages + current_message = await self._request_plan(api_key=api_key, request_kwargs=loop_kwargs) + + raise ApiError( + status_code=502, + code=502, + message="Agent LLM exceeded planning tool rounds", + ) + + +def _extract_text(message: Any) -> str: + chunks = getattr(message, "content", None) + if isinstance(chunks, str): + text = chunks.strip() + if text: + return text + raise ApiError(status_code=502, code=502, message="Agent LLM returned an empty response") + if not isinstance(chunks, list): + raise ApiError(status_code=502, code=502, message="Agent LLM returned an invalid response") + parts: list[str] = [] + for chunk in chunks: + parts.extend(_extract_text_parts(chunk)) + text = "\n".join(part for part in parts if part).strip() + if not text: + raise ApiError(status_code=502, code=502, message="Agent LLM returned an empty response") + return text + + +def _content_block_mappings(message: Any) -> list[dict[str, Any]]: + chunks = getattr(message, "content", None) + if isinstance(chunks, str): + return [{"type": "text", "text": chunks}] + if not isinstance(chunks, list): + return [] + blocks: list[dict[str, Any]] = [] + for chunk in chunks: + if isinstance(chunk, dict): + blocks.append(chunk) + continue + if hasattr(chunk, "model_dump"): + dumped = chunk.model_dump() + if isinstance(dumped, dict): + blocks.append(dumped) + continue + blocks.append( + { + "type": getattr(chunk, "type", None), + "text": getattr(chunk, "text", None), + "name": getattr(chunk, "name", None), + "input": getattr(chunk, "input", None), + "id": getattr(chunk, "id", None), + } + ) + return blocks + + +def _extract_tool_use_payload( + message: Any, + *, + tool_name_map: dict[str, str], +) -> tuple[list[dict[str, Any]], str | None]: + actions: list[dict[str, Any]] = [] + text_parts: list[str] = [] + for block in _content_block_mappings(message): + block_type = str(block.get("type") or "") + if block_type == "tool_use": + provider_name = str(block.get("name") or "").strip() + action_input = _coerce_mapping(block.get("input")) + actions.append( + { + "step": len(actions) + 1, + "tool": tool_name_map.get(provider_name, provider_name), + "input": action_input, + } + ) + continue + text_parts.extend(_extract_text_parts_from_mapping(block)) + summary = "\n".join(part for part in text_parts if part).strip() + return actions, summary or None + + +def _coerce_mapping(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if hasattr(value, "model_dump"): + dumped = value.model_dump() + if isinstance(dumped, dict): + return dumped + return {} + + +def _extract_text_parts(chunk: Any) -> list[str]: + if chunk is None: + return [] + if isinstance(chunk, str): + candidate = chunk.strip() + return [candidate] if candidate else [] + if isinstance(chunk, dict): + return _extract_text_parts_from_mapping(chunk) + if hasattr(chunk, "model_dump"): + dumped = chunk.model_dump() + if isinstance(dumped, dict): + return _extract_text_parts_from_mapping(dumped) + return _extract_text_parts_from_mapping( + { + "type": getattr(chunk, "type", None), + "text": getattr(chunk, "text", None), + "output_text": getattr(chunk, "output_text", None), + "content": getattr(chunk, "content", None), + } + ) + + +def _extract_text_parts_from_mapping(payload: dict[str, Any]) -> list[str]: + out: list[str] = [] + if payload.get("type") == "text": + out.extend(_flatten_text_value(payload.get("text"))) + for key in ("text", "output_text", "content"): + out.extend(_flatten_text_value(payload.get(key))) + deduped: list[str] = [] + seen: set[str] = set() + for item in out: + if not item or item in seen: + continue + deduped.append(item) + seen.add(item) + return deduped + + +def _flatten_text_value(value: Any) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + candidate = value.strip() + return [candidate] if candidate else [] + if isinstance(value, list): + out: list[str] = [] + for item in value: + out.extend(_flatten_text_value(item)) + return out + if isinstance(value, dict): + if value.get("type") == "text": + return _flatten_text_value(value.get("text")) + if "text" in value: + return _flatten_text_value(value.get("text")) + return [] + return [] + + +def _usage_payload(message: Any) -> dict[str, Any] | None: + usage = getattr(message, "usage", None) + if usage is None: + return None + if isinstance(usage, dict): + return usage + if hasattr(usage, "model_dump"): + dumped = usage.model_dump() + return dumped if isinstance(dumped, dict) else None + payload: dict[str, Any] = {} + usage_token_fields = ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + for key in usage_token_fields: + value = getattr(usage, key, None) + if value is not None: + payload[key] = value + return payload or None + + +def _merge_usage_totals(base: dict[str, int], extra: dict[str, Any] | None) -> dict[str, int]: + merged = dict(base) + if not isinstance(extra, dict): + return merged + for key in ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ): + value = extra.get(key) + if value is None: + continue + try: + parsed = int(value) + except (TypeError, ValueError): + continue + merged[key] = int(merged.get(key) or 0) + parsed + return merged + + +def _tool_result_content(payload: dict[str, Any]) -> list[dict[str, str]]: + text = json.dumps(payload, ensure_ascii=False) + if len(text) > 12_000: + text = text[:12_000] + "…" + return [{"type": "text", "text": text}] + + +def _is_tool_error_payload(payload: Any) -> bool: + if not isinstance(payload, dict): + return False + return bool(payload.get("_toolError") is True or payload.get("isError") is True) + + +def _reasoning_params(reasoning_effort: str) -> dict[str, Any]: + effort = (reasoning_effort or "adaptive").strip().lower() + if effort == "adaptive": + return {"thinking": {"type": "adaptive"}} + if effort not in {"low", "medium", "high", "xhigh", "max"}: + effort = "high" + return { + "thinking": {"type": "enabled"}, + "output_config": {"effort": effort}, + } + + +def _response_details(error: anthropic.APIStatusError) -> str: + response = getattr(error, "response", None) + text = getattr(response, "text", "") if response is not None else "" + return str(text or "")[:800] + + +def _parse_plan_message( + message: Any, + *, + tool_name_map: dict[str, str] | None = None, +) -> tuple[dict[str, Any], dict[str, Any] | None]: + tool_actions, summary = _extract_tool_use_payload( + message, + tool_name_map=tool_name_map or {}, + ) + usage = _usage_payload(message) + if tool_actions: + return { + "summary": summary or f"Prepared {len(tool_actions)} file action(s).", + "proposedActions": tool_actions, + }, usage + text = _extract_text(message) + parsed = _parse_json_text(text) + return parsed, usage + + +def _extract_tool_use_calls( + message: Any, + *, + tool_name_map: dict[str, str], +) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + for block in _content_block_mappings(message): + if str(block.get("type") or "") != "tool_use": + continue + provider_name = str(block.get("name") or "").strip() + if not provider_name: + continue + tool_use_id = str(block.get("id") or "").strip() + if not tool_use_id: + tool_use_id = f"tool_use_{len(calls) + 1}" + calls.append( + { + "id": tool_use_id, + "tool": tool_name_map.get(provider_name, provider_name), + "input": _coerce_mapping(block.get("input")), + } + ) + return calls + + +def _tool_name_map(tools: list[dict[str, Any]]) -> dict[str, str]: + mapping: dict[str, str] = {} + for tool in tools: + provider_name = str(tool.get("name") or "").strip() + if not provider_name: + continue + internal_name = str( + tool.get("internalName") + or tool.get("internal_name") + or tool.get("tool") + or provider_name + ).strip() + mapping[provider_name] = internal_name or provider_name + return mapping + + +def _anthropic_tools_payload(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + payload: list[dict[str, Any]] = [] + for tool in tools: + name = str(tool.get("name") or "").strip() + if not name: + continue + payload.append( + { + "name": name, + "description": str(tool.get("description") or ""), + "input_schema": dict(tool.get("input_schema") or {"type": "object"}), + } + ) + return payload + + +def _is_retryable_output_error(error: ApiError) -> bool: + if error.status_code != 502: + return False + return error.message in { + "Agent LLM returned an invalid response", + "Agent LLM returned an empty response", + "Agent LLM did not return valid JSON", + "Agent LLM JSON must be an object", + } + + +def _safe_plan_token_cap(settings: Settings) -> int: + raw_value = getattr(settings, "agent_llm_plan_max_tokens", 8192) + try: + parsed = int(raw_value) + except (TypeError, ValueError): + return 8192 + return max(1, parsed) + + +def _degraded_plan_request_kwargs(request_kwargs: dict[str, Any]) -> dict[str, Any]: + degraded_kwargs = dict(request_kwargs) + degraded_kwargs.pop("thinking", None) + degraded_kwargs.pop("output_config", None) + return degraded_kwargs + + +def _strict_json_retry_kwargs( + request_kwargs: dict[str, Any], + *, + max_tokens_cap: int, +) -> dict[str, Any]: + strict_kwargs = dict(request_kwargs) + strict_kwargs["max_tokens"] = max(1, int(max_tokens_cap)) + messages = request_kwargs.get("messages") + strict_kwargs["messages"] = _append_json_only_retry_instruction(messages) + return strict_kwargs + + +def _append_json_only_retry_instruction(messages: Any) -> list[dict[str, Any]]: + instruction = ( + "Return ONLY one valid JSON object that matches outputSchema. " + "Do not include markdown fences, prose, or extra text." + ) + if not isinstance(messages, list): + return [{"role": "user", "content": instruction}] + cloned: list[dict[str, Any]] = [] + for item in messages: + if isinstance(item, dict): + cloned.append(dict(item)) + for idx in range(len(cloned) - 1, -1, -1): + if cloned[idx].get("role") != "user": + continue + content = cloned[idx].get("content") + if isinstance(content, str): + merged = content.rstrip() + if merged: + merged = f"{merged}\n\n{instruction}" + else: + merged = instruction + cloned[idx]["content"] = merged + return cloned + cloned.append({"role": "user", "content": instruction}) + return cloned + + +def _parse_json_text(text: str) -> dict[str, Any]: + candidate = _strip_code_fences(text) + try: + return _decode_json_object(candidate) + except ApiError: + raise + except json.JSONDecodeError as decode_error: + extracted = _extract_balanced_json_object(candidate) + if extracted is not None: + try: + return _decode_json_object(extracted) + except ApiError: + raise + except json.JSONDecodeError: + pass + raise ApiError( + status_code=502, + code=502, + message="Agent LLM did not return valid JSON", + ) from decode_error + + +def _decode_json_object(candidate: str) -> dict[str, Any]: + parsed = json.loads(candidate) + if not isinstance(parsed, dict): + raise ApiError(status_code=502, code=502, message="Agent LLM JSON must be an object") + return parsed + + +def _strip_code_fences(text: str) -> str: + candidate = text.strip() + if not candidate.startswith("```"): + return candidate + lines = candidate.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].startswith("```"): + lines = lines[:-1] + return "\n".join(lines).strip() + + +def _extract_balanced_json_object(text: str) -> str | None: + start = -1 + depth = 0 + in_string = False + escaped = False + for idx, ch in enumerate(text): + if start < 0: + if ch == "{": + start = idx + depth = 1 + in_string = False + escaped = False + continue + if in_string: + if escaped: + escaped = False + elif ch == "\\": + escaped = True + elif ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + continue + if ch == "{": + depth += 1 + continue + if ch == "}": + depth -= 1 + if depth == 0: + return text[start : idx + 1] + if depth < 0: + start = -1 + depth = 0 + return None + + +__all__ = ["AnswerClient", "AnthropicPlannerClient", "PlannerClient", "ToolExecutor"] diff --git a/app/src/fileflash/agents/runtime/plan_runner.py b/app/src/fileflash/agents/runtime/plan_runner.py index 46a5918..33e8483 100644 --- a/app/src/fileflash/agents/runtime/plan_runner.py +++ b/app/src/fileflash/agents/runtime/plan_runner.py @@ -1,6 +1,1086 @@ from __future__ import annotations +import hashlib +import inspect +import json +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import and_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...core.mime import resolve_file_mime_type +from ...core.settings import Settings, get_settings +from ...models import AgentPlan, AgentSkill, BackgroundJob, File, Folder +from ...models.enums import AgentExecutionPolicy as DbAgentExecutionPolicy +from ...models.enums import FileStatus, FolderStatus, FolderType +from ...repositories import AgentSkillRepository +from ...repositories.agent.contracts import AgentSkillCatalogEntry +from ...schemas.agent import ( + AgentChosenSkill, + AgentCostEstimate, + AgentPlanningEvidence, + AgentPlanResult, + AgentProposedAction, + PlanAgentRequest, +) +from ..harness.ask import AskProtocol +from ..harness.event_bus import AgentEventBus +from ..harness.policy import classify_tool_side_effect, normalize_action_risk +from ..harness.router import ToolCall, ToolRouter +from ..harness.tool_registry import REGISTRY +from .llm import AnthropicPlannerClient, PlannerClient +from .reference_rules import is_symbolic_id_placeholder, parse_step_reference + class PlanRunner: - async def run(self, *args, **kwargs): - raise NotImplementedError("PlanRunner is scaffolded only in this stage") + def __init__( + self, + *, + settings: Settings | None = None, + planner_client: PlannerClient | None = None, + event_bus: AgentEventBus | None = None, + ) -> None: + self.settings = settings or get_settings() + self.planner_client = planner_client or AnthropicPlannerClient(settings=self.settings) + self.event_bus = event_bus + + async def run(self, *, db: AsyncSession, job: BackgroundJob) -> AgentPlanResult: + ask: AskProtocol | None = None + if self.event_bus is not None: + ask = AskProtocol(db=db, event_bus=self.event_bus, job_id=int(job.job_id)) + await ask.start() + try: + return await self._run(db=db, job=job, ask=ask) + finally: + if ask is not None: + await ask.aclose() + + async def _run( + self, + *, + db: AsyncSession, + job: BackgroundJob, + ask: AskProtocol | None, + ) -> AgentPlanResult: + if job.requested_by is None: + raise ApiError(status_code=400, code=400, message="Agent job is missing requestedBy") + + request = PlanAgentRequest.model_validate(dict(job.payload or {})) + user_id = int(job.requested_by) + skill = await _choose_skill( + db, + user_id=user_id, + task_input=request.input, + prefer_skill_id=request.hints.prefer_skill_id, + ) + metadata = await _collect_context_metadata(db, user_id=user_id, request=request) + allowed_tools = _skill_tool_whitelist(skill) + allowed_tool_set = set(allowed_tools) + exploration_tools = tuple( + tool_name for tool_name in allowed_tools if REGISTRY.get(tool_name).side_effect == "read" + ) + exploration_tool_set = set(exploration_tools) + planner_router = ToolRouter(db=db, user_id=user_id) + tool_call_budget = min(self.settings.agent_job_max_tool_calls, 32) + planned_tool_calls = 0 + planning_evidence: list[AgentPlanningEvidence] = [] + + async def _planning_tool_executor(tool_name: str, args: dict[str, Any]) -> dict[str, Any]: + nonlocal planned_tool_calls + planned_tool_calls += 1 + if planned_tool_calls > tool_call_budget: + raise ApiError( + status_code=400, + code=400, + message="Planner exceeded exploratory tool-call budget", + ) + if tool_name not in allowed_tool_set: + blocked = _blocked_planning_tool_result( + tool_name=tool_name, + reason="Planner attempted a tool that is not allowed by the selected skill.", + ) + if len(planning_evidence) < 12: + planning_evidence.append( + AgentPlanningEvidence( + step=planned_tool_calls, + tool=tool_name, + input=_evidence_mapping(args), + output_preview=_evidence_preview(blocked), + ) + ) + return blocked + if tool_name not in exploration_tool_set: + blocked = _blocked_planning_tool_result( + tool_name=tool_name, + reason=( + "Planner exploratory tool call must be read-only. " + "Move write operations into final proposedActions." + ), + ) + if len(planning_evidence) < 12: + planning_evidence.append( + AgentPlanningEvidence( + step=planned_tool_calls, + tool=tool_name, + input=_evidence_mapping(args), + output_preview=_evidence_preview(blocked), + ) + ) + return blocked + output = await planner_router.dispatch(ToolCall(tool_name=tool_name, arguments=args)) + if len(planning_evidence) < 12: + planning_evidence.append( + AgentPlanningEvidence( + step=planned_tool_calls, + tool=tool_name, + input=_evidence_mapping(args), + output_preview=_evidence_preview(output), + ) + ) + return output + + llm_payload = await self.planner_client.create_plan( + system_prompt=_system_prompt(), + user_prompt=_user_prompt( + request=request, + skill=skill, + allowed_tools=allowed_tools, + exploration_tools=exploration_tools, + metadata=metadata, + ), + max_tokens=request.hints.budget_tokens, + reasoning_effort=request.hints.reasoning_effort, + tools=REGISTRY.anthropic_tools_for(exploration_tools), + tool_executor=_planning_tool_executor, + max_tool_roundtrips=6, + ) + + effective_max_steps: int | None + if self.settings.is_development_env: + effective_max_steps = None + else: + effective_max_steps = min( + request.hints.max_steps, + self.settings.agent_job_max_tool_calls, + ) + + actions = _normalize_actions( + llm_payload=llm_payload, + allowed_tools=allowed_tools, + max_steps=effective_max_steps, + ) + chosen_skill = _chosen_skill(skill) + llm_summary = str( + llm_payload.get("summary") or f"Prepared {len(actions)} file action(s)." + ).strip() + if not llm_summary: + llm_summary = f"Prepared {len(actions)} file action(s)." + summary = llm_summary + if _has_write_actions(actions): + summary = await _grounded_write_summary( + db, + user_id=user_id, + actions=actions, + fallback_summary=llm_summary, + ) + + requires_confirmation = ( + request.execution_policy != "autopilot" + or any(action.requires_confirmation for action in actions) + ) + cost_estimate = _cost_estimate(llm_payload=llm_payload, actions=actions, metadata=metadata) + plan_hash = _plan_hash( + chosen_skill=chosen_skill, + actions=actions, + summary=summary, + ) + result = AgentPlanResult( + plan_job_id=str(job.job_id), + plan_hash=plan_hash, + chosen_skill=chosen_skill, + proposed_actions=actions, + summary=summary, + requires_confirmation=requires_confirmation, + cost_estimate=cost_estimate, + planning_evidence=planning_evidence or None, + ) + await _upsert_agent_plan( + db, + job=job, + request=request, + result=result, + ) + try: + await db.commit() + except Exception: + await db.rollback() + raise + return result + + async def _ask( + self, + *, + ask: AskProtocol | None, + prompt: str, + schema: dict[str, Any], + ) -> Any | None: + if ask is None: + return None + return await ask.ask( + prompt=prompt, + schema=schema, + timeout_sec=float(self.settings.agent_inbox_ask_timeout_sec), + ) + + +async def _choose_skill( + db: AsyncSession, + *, + user_id: int, + task_input: str, + prefer_skill_id: str | None, +) -> AgentSkill | AgentSkillCatalogEntry | None: + repo = AgentSkillRepository(db) + if prefer_skill_id: + skill = await repo.get_by_key(skill_key=prefer_skill_id, user_id=user_id) + if skill is None: + raise ApiError(status_code=404, code=404, message="Preferred skill not found") + return skill + + candidates = await repo.list_visible(user_id=user_id, limit=50) + if not candidates: + return None + + normalized_input = task_input.lower() + best: tuple[int, AgentSkillCatalogEntry] | None = None + for candidate in candidates: + haystack = ( + f"{candidate.skill_key} {candidate.name} {candidate.description} " + f"{candidate.triggers_text or ''} {candidate.search_text}" + ).lower() + score = 0 + for token in _tokens(normalized_input): + if token in haystack: + score += 2 if token in {"organize", "整理", "classify", "分类"} else 1 + if "整理" in normalized_input and "organize" in haystack: + score += 4 + if best is None or score > best[0]: + best = (score, candidate) + + if best is not None and best[0] > 0: + return best[1] + return candidates[0] + + +def _tokens(text: str) -> list[str]: + return [token.strip(" ,.;:!?，。；：！？") for token in text.split() if token.strip()] + + +def _skill_key(skill: AgentSkill | AgentSkillCatalogEntry | None) -> str | None: + if skill is None: + return None + return str(skill.skill_key) + + +def _skill_name(skill: AgentSkill | AgentSkillCatalogEntry | None) -> str | None: + if skill is None: + return None + return str(skill.name) + + +def _skill_tool_whitelist(skill: AgentSkill | AgentSkillCatalogEntry | None) -> tuple[str, ...]: + raw: Any = None + if isinstance(skill, AgentSkill): + raw = skill.tool_whitelist_json + elif skill is not None: + raw = skill.tool_whitelist_json + if isinstance(raw, list) and raw: + tools = tuple(str(item) for item in raw if str(item).strip()) + unknown = REGISTRY.unknown_names(tools) + if unknown: + raise ApiError( + status_code=422, + code=422, + message="Unknown agent tool in selected skill", + data={"unknownTools": sorted(unknown)}, + ) + return tools + return REGISTRY.all_names() + + +def _chosen_skill(skill: AgentSkill | AgentSkillCatalogEntry | None) -> AgentChosenSkill | None: + key = _skill_key(skill) + name = _skill_name(skill) + if not key or not name: + return None + return AgentChosenSkill(id=key, name=name) + + +async def _collect_context_metadata( + db: AsyncSession, + *, + user_id: int, + request: PlanAgentRequest, +) -> dict[str, Any]: + context = request.context + selected_file_ids = _parse_ids(context.selected_file_ids, "fileId") + selected_folder_ids = _parse_ids(context.selected_folder_ids, "folderId") + + files: list[File] = [] + folders: list[Folder] = [] + if selected_file_ids: + files = list( + await db.scalars( + select(File).where( + and_( + File.owner_id == user_id, + File.file_id.in_(selected_file_ids), + File.status == FileStatus.ACTIVE, + File.is_latest.is_(True), + ) + ) + ) + ) + if selected_folder_ids: + folders = list( + await db.scalars( + select(Folder).where( + and_( + Folder.owner_id == user_id, + Folder.folder_id.in_(selected_folder_ids), + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + ) + + scope = "selected" if selected_file_ids or selected_folder_ids else "currentFolder" + folder_id = await _resolve_folder_id(db, user_id=user_id, folder_id=context.root_folder_id) + if scope == "currentFolder": + folders = list( + await db.scalars( + select(Folder) + .where( + and_( + Folder.owner_id == user_id, + Folder.parent_folder_id == folder_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + .order_by(Folder.folder_name.asc()) + .limit(200) + ) + ) + files = list( + await db.scalars( + select(File) + .where( + and_( + File.owner_id == user_id, + File.folder_id == folder_id, + File.status == FileStatus.ACTIVE, + File.is_latest.is_(True), + ) + ) + .order_by(File.file_name.asc()) + .limit(200) + ) + ) + + return { + "scope": scope, + "currentPath": context.current_path, + "rootFolderId": str(folder_id), + "files": [_file_metadata(row) for row in files], + "folders": [_folder_metadata(row) for row in folders], + } + + +async def _resolve_folder_id(db: AsyncSession, *, user_id: int, folder_id: str | None) -> int: + if not folder_id or folder_id == "root": + root = await db.scalar( + select(Folder).where( + and_( + Folder.owner_id == user_id, + Folder.parent_folder_id.is_(None), + Folder.folder_type == FolderType.ROOT, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + if root is None: + raise ApiError(status_code=404, code=404, message="Root folder not found") + return int(root.folder_id) + try: + parsed = int(folder_id) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid rootFolderId") from exc + exists = await db.scalar( + select(Folder.folder_id).where( + and_( + Folder.folder_id == parsed, + Folder.owner_id == user_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + if exists is None: + raise ApiError(status_code=404, code=404, message="Folder not found") + return parsed + + +def _parse_ids(raw_ids: list[str], field_name: str) -> list[int]: + parsed: list[int] = [] + seen: set[int] = set() + for raw in raw_ids: + if raw == "root": + continue + try: + value = int(raw) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message=f"Invalid {field_name}") from exc + if value <= 0 or value in seen: + continue + parsed.append(value) + seen.add(value) + return parsed + + +def _file_metadata(row: File) -> dict[str, Any]: + return { + "itemType": "file", + "id": str(row.file_id), + "name": row.file_name, + "size": int(row.file_size or 0), + "mimeType": resolve_file_mime_type( + mime_type=row.mime_type, + file_ext=row.file_ext, + file_name=row.file_name, + ), + "folderId": str(row.folder_id), + "createdAt": row.created_at.isoformat() if row.created_at else None, + "updatedAt": row.updated_at.isoformat() if row.updated_at else None, + } + + +def _folder_metadata(row: Folder) -> dict[str, Any]: + return { + "itemType": "folder", + "id": str(row.folder_id), + "name": row.folder_name, + "size": int(row.cached_size or 0), + "parentFolderId": str(row.parent_folder_id) if row.parent_folder_id else None, + "createdAt": row.created_at.isoformat() if row.created_at else None, + "updatedAt": row.updated_at.isoformat() if row.updated_at else None, + } + + +def _system_prompt() -> str: + return ( + "You are FileFlash Agent Planner. Build plans from tool-grounded facts, not assumptions. " + "If you need facts, first call read-only tools; then output one final JSON object that matches outputSchema. " + "Do not read or infer file contents. Deletions are high risk and must be explicit. " + "Cross-step dependencies must use '$stepN.field' references only and never symbolic placeholders " + "like 'newFolderId'." + ) + + +def _user_prompt( + *, + request: PlanAgentRequest, + skill: AgentSkill | AgentSkillCatalogEntry | None, + allowed_tools: tuple[str, ...], + exploration_tools: tuple[str, ...], + metadata: dict[str, Any], +) -> str: + payload = { + "task": request.input, + "executionPolicy": request.execution_policy, + "reasoningEffort": request.hints.reasoning_effort, + "dataPolicy": request.data_policy.model_dump(by_alias=True), + "skill": _skill_payload(skill), + "allowedTools": list(allowed_tools), + "explorationTools": list(exploration_tools), + "toolSchemas": _tool_schemas(allowed_tools), + "toolUseMode": ( + "Use exploratory tool_use only with read-only tools listed in explorationTools. " + "Never call write tools with tool_use; put write operations only in final proposedActions " + "and use '$stepN.field' references for cross-step ids." + ), + "plannerDefaults": { + "organizeRequest": { + "scope": "root recursive unless user constrained the scope", + "folderNaming": "reuse existing folders first; create english category folders if missing", + "writePlanPolicy": "generate executable write actions by default", + } + }, + "referenceContract": { + "syntax": "$stepN.field", + "rules": [ + "If an action needs data from a previous action, use only '$stepN.field'.", + "N must refer to an existing previous step number.", + "Never invent symbolic placeholders such as 'newFolderId'.", + ], + "requiredExample": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "13", "targetFolderId": "$step1.folderId"}, + }, + ], + }, + "fileMetadata": metadata, + "outputSchema": { + "summary": "string", + "proposedActions": [ + { + "step": "integer starting at 1", + "tool": "one of allowedTools", + "input": "object", + "sideEffect": "read or write", + } + ], + }, + } + return json.dumps(payload, ensure_ascii=False, sort_keys=True) + + +def _skill_payload(skill: AgentSkill | AgentSkillCatalogEntry | None) -> dict[str, Any] | None: + if skill is None: + return None + if isinstance(skill, AgentSkill): + return { + "skillKey": skill.skill_key, + "name": skill.name, + "description": skill.description, + "triggersText": skill.triggers_text, + "planTemplate": skill.plan_template_json or {}, + } + return { + "skillKey": skill.skill_key, + "name": skill.name, + "description": skill.description, + "triggersText": skill.triggers_text, + "planTemplate": skill.plan_template_json or {}, + } + + +def _tool_schemas(allowed_tools: tuple[str, ...]) -> list[dict[str, Any]]: + return REGISTRY.schemas_for(allowed_tools) + + +def _blocked_planning_tool_result(*, tool_name: str, reason: str) -> dict[str, Any]: + return { + "_plannerBlocked": True, + "_toolError": True, + "errorCode": "planner.exploration_tool_blocked", + "tool": tool_name, + "message": reason, + "guidance": ( + "Do not execute write tools during exploration. " + "Place write actions in final proposedActions and reference prior outputs via '$stepN.field'." + ), + } + + +def _normalize_actions( + *, + llm_payload: dict[str, Any], + allowed_tools: tuple[str, ...], + max_steps: int | None, +) -> list[AgentProposedAction]: + raw_actions = llm_payload.get("proposedActions", llm_payload.get("proposed_actions")) + if raw_actions is None: + raw_actions = llm_payload.get("actions") + if not isinstance(raw_actions, list): + raise ApiError(status_code=502, code=502, message="Agent plan JSON missing proposedActions") + if max_steps is not None and len(raw_actions) > max_steps: + raise ApiError(status_code=400, code=400, message="Agent plan exceeds maxSteps") + + allowed = set(allowed_tools) + normalized: list[AgentProposedAction] = [] + seen_steps: set[int] = set() + for index, raw_action in enumerate(raw_actions, start=1): + if not isinstance(raw_action, dict): + raise ApiError(status_code=502, code=502, message="Agent action must be an object") + tool = str(raw_action.get("tool") or raw_action.get("toolName") or "").strip() + if tool not in allowed: + raise ApiError( + status_code=400, + code=400, + message=f"Tool is not allowed by selected skill: {tool}", + ) + action_input = raw_action.get("input", raw_action.get("arguments", {})) + if not isinstance(action_input, dict): + raise ApiError( + status_code=502, + code=502, + message="Agent action input must be an object", + ) + try: + step = int(raw_action.get("step") or index) + except (TypeError, ValueError) as exc: + raise ApiError( + status_code=502, + code=502, + message="Agent action step is invalid", + ) from exc + if step < 1 or step in seen_steps: + raise ApiError(status_code=502, code=502, message="Agent action step is invalid") + seen_steps.add(step) + action = AgentProposedAction( + step=step, + tool=tool, + input=action_input, + side_effect=raw_action.get("sideEffect") or classify_tool_side_effect(tool), # type: ignore[arg-type] + risk_level=raw_action.get("riskLevel") or "low", # type: ignore[arg-type] + requires_confirmation=bool(raw_action.get("requiresConfirmation") or False), + confirmation_reason=raw_action.get("confirmationReason"), + ) + normalized.append(normalize_action_risk(action)) + sorted_actions = sorted(normalized, key=lambda action: action.step) + _validate_action_inputs(sorted_actions) + return sorted_actions + + +def _validate_action_inputs(actions: list[AgentProposedAction]) -> None: + steps = {action.step for action in actions} + for action in actions: + _validate_action_input_value( + value=action.input, + action_step=action.step, + known_steps=steps, + field_name=None, + field_path="input", + ) + + +def _validate_action_input_value( + *, + value: Any, + action_step: int, + known_steps: set[int], + field_name: str | None, + field_path: str, +) -> None: + if isinstance(value, str): + reference = parse_step_reference(value) + if reference is not None: + ref_step, _ = reference + if ref_step not in known_steps: + raise ApiError( + status_code=400, + code=400, + message=( + f"Invalid plan action at step {action_step}: field '{field_path}' references " + f"missing step {ref_step} via '{value}'." + ), + ) + if ref_step >= action_step: + raise ApiError( + status_code=400, + code=400, + message=( + f"Invalid plan action at step {action_step}: field '{field_path}' references " + f"future step {ref_step} via '{value}'. Use only previous-step references." + ), + ) + return + if is_symbolic_id_placeholder(value=value, field_name=field_name): + raise ApiError( + status_code=400, + code=400, + message=( + f"Invalid plan action at step {action_step}: field '{field_path}' uses " + f"unresolved placeholder '{value}'. Use '$stepN.field' references." + ), + ) + return + if isinstance(value, list): + for index, item in enumerate(value): + _validate_action_input_value( + value=item, + action_step=action_step, + known_steps=known_steps, + field_name=field_name, + field_path=f"{field_path}[{index}]", + ) + return + if isinstance(value, dict): + for key, item in value.items(): + _validate_action_input_value( + value=item, + action_step=action_step, + known_steps=known_steps, + field_name=key, + field_path=f"{field_path}.{key}", + ) + + +def _has_write_actions(actions: list[AgentProposedAction]) -> bool: + return any(action.side_effect == "write" for action in actions) + + +async def _grounded_write_summary( + db: AsyncSession, + *, + user_id: int, + actions: list[AgentProposedAction], + fallback_summary: str, +) -> str: + write_actions = [action for action in actions if action.side_effect == "write"] + if not write_actions: + return fallback_summary + + created_folder_names: list[str] = [] + created_folder_by_step: dict[int, str] = {} + move_file_actions: list[AgentProposedAction] = [] + move_folder_actions: list[AgentProposedAction] = [] + file_ids: set[int] = set() + folder_ids: set[int] = set() + + for action in write_actions: + if action.tool == "drive.createFolder": + folder_name = str(action.input.get("name") or action.input.get("folderName") or "").strip() + if folder_name: + created_folder_names.append(folder_name) + created_folder_by_step[action.step] = folder_name + continue + if action.tool == "drive.moveFile": + move_file_actions.append(action) + file_id = _coerce_positive_int(action.input.get("fileId")) + if file_id is not None: + file_ids.add(file_id) + target_folder_id = action.input.get("targetFolderId") + if isinstance(target_folder_id, str): + parsed_folder_id = _coerce_positive_int(target_folder_id) + if parsed_folder_id is not None: + folder_ids.add(parsed_folder_id) + continue + if action.tool == "drive.moveFolder": + move_folder_actions.append(action) + source_folder_id = _coerce_positive_int(action.input.get("folderId")) + if source_folder_id is not None: + folder_ids.add(source_folder_id) + target_parent_id = action.input.get("targetParentId", action.input.get("targetFolderId")) + if isinstance(target_parent_id, str): + parsed_parent_id = _coerce_positive_int(target_parent_id) + if parsed_parent_id is not None: + folder_ids.add(parsed_parent_id) + + file_name_map = await _safe_fetch_file_names(db, user_id=user_id, file_ids=file_ids) + folder_name_map = await _safe_fetch_folder_names(db, user_id=user_id, folder_ids=folder_ids) + + moved_file_names: list[str] = [] + destination_folder_names: list[str] = [] + for action in move_file_actions: + file_id = _coerce_positive_int(action.input.get("fileId")) + if file_id is not None: + file_name = file_name_map.get(file_id) + if file_name: + moved_file_names.append(file_name) + destination = _resolve_destination_folder_name( + action.input.get("targetFolderId"), + created_folder_by_step=created_folder_by_step, + folder_name_map=folder_name_map, + ) + if destination: + destination_folder_names.append(destination) + + moved_folder_count = len(move_folder_actions) + moved_file_count = len(move_file_actions) + clauses: list[str] = [] + + if created_folder_names: + unique_created = _unique_preserve_order(created_folder_names) + if len(unique_created) == 1: + clauses.append(f"创建“{unique_created[0]}”文件夹") + else: + clauses.append(f"创建 {len(unique_created)} 个文件夹") + + if moved_file_count > 0: + clauses.append( + f"将{_format_moved_file_subject(moved_file_names, moved_file_count)}移动到" + f"{_format_destination_folder(destination_folder_names)}" + ) + + if moved_folder_count > 0: + clauses.append(f"移动 {moved_folder_count} 个文件夹") + + if not clauses: + return fallback_summary + return "，并".join(clauses) + "。" + + +async def _safe_fetch_file_names( + db: AsyncSession, + *, + user_id: int, + file_ids: set[int], +) -> dict[int, str]: + if not file_ids: + return {} + try: + result = await db.execute( + select(File.file_id, File.file_name).where( + and_( + File.owner_id == user_id, + File.file_id.in_(sorted(file_ids)), + File.status == FileStatus.ACTIVE, + File.is_latest.is_(True), + ) + ) + ) + rows = result.all() if hasattr(result, "all") else [] + if inspect.isawaitable(rows): + rows = await rows + except Exception: + return {} + out: dict[int, str] = {} + if not isinstance(rows, list): + return out + for row in rows: + try: + file_id = row[0] + file_name = row[1] + except Exception: + continue + parsed = _coerce_positive_int(file_id) + name = str(file_name or "").strip() + if parsed is None or not name: + continue + out[parsed] = name + return out + + +async def _safe_fetch_folder_names( + db: AsyncSession, + *, + user_id: int, + folder_ids: set[int], +) -> dict[int, str]: + if not folder_ids: + return {} + try: + result = await db.execute( + select(Folder.folder_id, Folder.folder_name).where( + and_( + Folder.owner_id == user_id, + Folder.folder_id.in_(sorted(folder_ids)), + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + rows = result.all() if hasattr(result, "all") else [] + if inspect.isawaitable(rows): + rows = await rows + except Exception: + return {} + out: dict[int, str] = {} + if not isinstance(rows, list): + return out + for row in rows: + try: + folder_id = row[0] + folder_name = row[1] + except Exception: + continue + parsed = _coerce_positive_int(folder_id) + name = str(folder_name or "").strip() + if parsed is None or not name: + continue + out[parsed] = name + return out + + +def _resolve_destination_folder_name( + raw_value: Any, + *, + created_folder_by_step: dict[int, str], + folder_name_map: dict[int, str], +) -> str | None: + if not isinstance(raw_value, str): + return None + value = raw_value.strip() + if not value: + return None + reference = parse_step_reference(value) + if reference is not None: + step, path = reference + if path and path[0].lower() in {"folderid", "id"}: + return created_folder_by_step.get(step) + return None + parsed = _coerce_positive_int(value) + if parsed is None: + return None + return folder_name_map.get(parsed) + + +def _format_moved_file_subject(file_names: list[str], total_count: int) -> str: + unique_names = _unique_preserve_order( + [name.strip() for name in file_names if isinstance(name, str) and name.strip()] + ) + if not unique_names: + return f"{total_count} 个文件" + if len(unique_names) == 1 and total_count == 1: + return f"“{unique_names[0]}”" + preview = unique_names[:3] + quoted = "、".join(f"“{name}”" for name in preview) + if total_count > len(preview): + return f"{quoted}等 {total_count} 个文件" + if total_count > len(unique_names): + return f"{quoted}共 {total_count} 个文件" + return f"{quoted}共 {total_count} 个文件" + + +def _format_destination_folder(folder_names: list[str]) -> str: + unique_names = _unique_preserve_order( + [name.strip() for name in folder_names if isinstance(name, str) and name.strip()] + ) + if not unique_names: + return "目标文件夹" + if len(unique_names) == 1: + return f"“{unique_names[0]}”文件夹" + return "多个目标文件夹" + + +def _unique_preserve_order(items: list[str]) -> list[str]: + out: list[str] = [] + seen: set[str] = set() + for item in items: + if item in seen: + continue + seen.add(item) + out.append(item) + return out + + +def _coerce_positive_int(value: Any) -> int | None: + text = str(value or "").strip() + if not text.isdigit(): + return None + parsed = int(text) + if parsed <= 0: + return None + return parsed + + +def _evidence_mapping(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + preview = _evidence_value_preview(value, depth=0) + if isinstance(preview, dict): + return preview + return {} + + +def _evidence_preview(value: Any) -> dict[str, Any]: + preview = _evidence_value_preview(value, depth=0) + if isinstance(preview, dict): + return preview + return {"value": preview} + + +def _evidence_value_preview(value: Any, *, depth: int) -> Any: + if depth >= 3: + if isinstance(value, str): + return value[:120] + ("…" if len(value) > 120 else "") + if isinstance(value, (int, float, bool)) or value is None: + return value + return str(value)[:120] + + if isinstance(value, dict): + out: dict[str, Any] = {} + items = list(value.items()) + for index, (key, item) in enumerate(items): + if index >= 12: + out["_truncatedKeys"] = len(items) - 12 + break + out[str(key)] = _evidence_value_preview(item, depth=depth + 1) + return out + + if isinstance(value, list): + preview_items = [_evidence_value_preview(item, depth=depth + 1) for item in value[:6]] + if len(value) > 6: + preview_items.append(f"...({len(value) - 6} more)") + return preview_items + + if isinstance(value, str): + return value[:200] + ("…" if len(value) > 200 else "") + if isinstance(value, (int, float, bool)) or value is None: + return value + return str(value)[:200] + + +def _cost_estimate( + *, + llm_payload: dict[str, Any], + actions: list[AgentProposedAction], + metadata: dict[str, Any], +) -> AgentCostEstimate: + usage = llm_payload.get("_usage") + if isinstance(usage, dict): + input_tokens = int(usage.get("input_tokens") or 0) + output_tokens = int(usage.get("output_tokens") or 0) + tokens = input_tokens + output_tokens + else: + tokens = max(1000, len(json.dumps(metadata, ensure_ascii=False)) // 3 + len(actions) * 180) + return AgentCostEstimate( + tokens=tokens, + tool_calls=len(actions), + duration_sec_estimate=max(1, len(actions) * 3), + ) + + +def _plan_hash( + *, + chosen_skill: AgentChosenSkill | None, + actions: list[AgentProposedAction], + summary: str, +) -> str: + payload = { + "chosenSkill": chosen_skill.model_dump(by_alias=True) if chosen_skill else None, + "proposedActions": [action.model_dump(by_alias=True) for action in actions], + "summary": summary, + } + raw = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + return "sha256:" + hashlib.sha256(raw.encode("utf-8")).hexdigest() + + +async def _upsert_agent_plan( + db: AsyncSession, + *, + job: BackgroundJob, + request: PlanAgentRequest, + result: AgentPlanResult, +) -> None: + existing = await db.scalar(select(AgentPlan).where(AgentPlan.job_id == job.job_id)) + values = { + "job_id": int(job.job_id), + "user_id": int(job.requested_by or 0), + "input_text": request.input, + "context_json": request.context.model_dump(by_alias=True), + "execution_policy": DbAgentExecutionPolicy(request.execution_policy), + "data_policy_json": request.data_policy.model_dump(by_alias=True), + "chosen_skill_id": result.chosen_skill.id if result.chosen_skill else None, + "proposed_actions_json": [ + action.model_dump(by_alias=True) for action in result.proposed_actions + ], + "plan_hash": result.plan_hash, + "summary": result.summary, + "cost_estimate_json": result.cost_estimate.model_dump(by_alias=True), + "created_at": datetime.now(UTC), + } + if existing is None: + db.add(AgentPlan(**values)) + else: + for key, value in values.items(): + if key not in {"job_id", "created_at"}: + setattr(existing, key, value) + await db.flush() diff --git a/app/src/fileflash/agents/runtime/reference_rules.py b/app/src/fileflash/agents/runtime/reference_rules.py new file mode 100644 index 0000000..cafb309 --- /dev/null +++ b/app/src/fileflash/agents/runtime/reference_rules.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import re + +_STEP_REFERENCE = re.compile(r"^\$step(?P\d+)\.(?P[A-Za-z0-9_.-]+)$") +_SYMBOLIC_ID = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_REFERENCE_FIELDS = frozenset( + { + "id", + "fileid", + "folderid", + "parentfolderid", + "targetfolderid", + "targetparentid", + } +) + + +def parse_step_reference(value: str) -> tuple[int, tuple[str, ...]] | None: + match = _STEP_REFERENCE.match(value.strip()) + if not match: + return None + step = int(match.group("step")) + path = tuple(part for part in match.group("path").split(".") if part) + if not path: + return None + return step, path + + +def is_symbolic_id_placeholder(*, value: str, field_name: str | None) -> bool: + if not field_name or field_name.strip().lower() not in _REFERENCE_FIELDS: + return False + candidate = value.strip() + if not candidate: + return False + if candidate.lower() == "root": + return False + if candidate.isdigit(): + return False + if parse_step_reference(candidate) is not None: + return False + return _SYMBOLIC_ID.match(candidate) is not None + + +__all__ = ["is_symbolic_id_placeholder", "parse_step_reference"] diff --git a/app/src/fileflash/agents/tools/__init__.py b/app/src/fileflash/agents/tools/__init__.py new file mode 100644 index 0000000..7ea9a9d --- /dev/null +++ b/app/src/fileflash/agents/tools/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from . import drive as drive + +__all__ = ["drive"] diff --git a/app/src/fileflash/agents/tools/drive.py b/app/src/fileflash/agents/tools/drive.py new file mode 100644 index 0000000..f74d369 --- /dev/null +++ b/app/src/fileflash/agents/tools/drive.py @@ -0,0 +1,878 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import and_, select + +from ...core.errors import ApiError +from ...core.mime import resolve_file_mime_type +from ...models import File, Folder +from ...models.enums import FileStatus, FolderStatus, FolderType +from ...models.tables_storage import StorageObject +from ...schemas.file import ( + CreateFolderRequest, + GetFolderContentsQuery, + MoveFileRequest, + MoveFolderRequest, + RenameFileRequest, + RenameFolderRequest, +) +from ..harness.tool_registry import REGISTRY, ToolContext, ToolSpec + +_CATEGORIES = ("video", "audio", "image", "document", "archive", "other") +_COUNT_FILE_NAME_LIMIT = 12 + + +async def _list_folder(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = _first_value(args, "folderId", "parentFolderId") or "root" + query = GetFolderContentsQuery( + folder_id=str(folder_id), + page=_int_arg(args.get("page"), default=1, minimum=1), + per_page=_int_arg(args.get("perPage"), default=200, minimum=1, maximum=200), + search=_optional_text(args.get("search")), + ) + if str(folder_id) == "root": + result = await ctx.folder_service.get_root_contents(user_id=ctx.user_id, query=query) + else: + result = await ctx.folder_service.get_folder_contents(user_id=ctx.user_id, query=query) + return result.model_dump(by_alias=True, mode="json") + + +async def _count_files(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = str(_first_value(args, "folderId", "parentFolderId") or "root") + recursive = _bool_arg(args.get("recursive"), default=True) + category = _normalize_category(args.get("category")) + search = str(args.get("search") or "").strip() + root_folder_id = await _resolve_folder_id(ctx, folder_id=folder_id) + folder_ids = await _folder_scope_ids(ctx, root_folder_id=root_folder_id, recursive=recursive) + + statement = _active_files_query(ctx, folder_ids=folder_ids) + if search: + statement = statement.where(File.file_name.ilike(f"%{search}%")) + statement = statement.order_by(File.file_name.asc()) + + rows = list(await ctx.db.scalars(statement)) + by_mime_type: dict[str, int] = {} + sample_items: list[dict[str, Any]] = [] + item_names: list[str] = [] + names_truncated = False + total_items = 0 + for row in rows: + resolved_mime = _resolved_mime(row) + if category is not None and _category_for_file(row) != category: + continue + + total_items += 1 + by_mime_type[resolved_mime] = by_mime_type.get(resolved_mime, 0) + 1 + file_name = str(row.file_name or "").strip() + if file_name: + if len(item_names) < _COUNT_FILE_NAME_LIMIT: + item_names.append(file_name) + else: + names_truncated = True + if len(sample_items) < 5: + sample_items.append(await _file_payload(ctx, row, include_path=False)) + + return { + "totalItems": total_items, + "category": category, + "recursive": recursive, + "folderId": str(root_folder_id), + "search": search or None, + "byMimeType": dict(sorted(by_mime_type.items())), + "itemNames": item_names, + "itemNamesTruncated": names_truncated, + "sampleItems": sample_items, + } + + +async def _create_folder(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + name = _required_text(args, "name", "folderName") + parent_id = _first_value(args, "parentFolderId", "targetParentId", "folderId") or "root" + result = await ctx.folder_service.create_folder( + user_id=ctx.user_id, + payload=CreateFolderRequest(folder_name=name, parent_folder_id=str(parent_id)), + ) + data = result.model_dump(by_alias=True, mode="json") + data.setdefault("folderId", data.get("id")) + return data + + +async def _move_file(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + file_id = _required_text(args, "fileId", "id") + target_folder_id = _required_text(args, "targetFolderId", "targetParentId") + result = await ctx.file_service.move_file( + user_id=ctx.user_id, + file_id=file_id, + payload=MoveFileRequest( + target_folder_id=target_folder_id, + share_handling=str(args.get("shareHandling") or "keep"), + ), + ) + return result.model_dump(by_alias=True, mode="json") + + +async def _move_folder(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = _required_text(args, "folderId", "id") + target_parent_id = _required_text(args, "targetParentId", "targetFolderId") + result = await ctx.folder_service.move_folder( + user_id=ctx.user_id, + folder_id=folder_id, + payload=MoveFolderRequest( + target_parent_id=target_parent_id, + share_handling=str(args.get("shareHandling") or "keep"), + ), + ) + return result.model_dump(by_alias=True, mode="json") + + +async def _rename_file(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + file_id = _required_text(args, "fileId", "id") + file_name = _required_text(args, "fileName", "name") + result = await ctx.file_service.rename_file( + user_id=ctx.user_id, + file_id=file_id, + payload=RenameFileRequest(file_name=file_name), + ) + return result.model_dump(by_alias=True, mode="json") + + +async def _rename_folder(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = _required_text(args, "folderId", "id") + folder_name = _required_text(args, "folderName", "name") + result = await ctx.folder_service.rename_folder( + user_id=ctx.user_id, + folder_id=folder_id, + payload=RenameFolderRequest(folder_name=folder_name), + ) + return result.model_dump(by_alias=True, mode="json") + + +async def _delete_file(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + file_id = _required_text(args, "fileId", "id") + result = await ctx.file_service.delete_file(user_id=ctx.user_id, file_id=file_id) + return result.model_dump(by_alias=True, mode="json") + + +async def _delete_folder(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = _required_text(args, "folderId", "id") + result = await ctx.folder_service.delete_folder(user_id=ctx.user_id, folder_id=folder_id) + return result.model_dump(by_alias=True, mode="json") + + +async def _search_files(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + query = _optional_text(args.get("query")) or _optional_text(args.get("search")) or "" + folder_id = str(args.get("folderId") or "root") + recursive = _bool_arg(args.get("recursive"), default=True) + category = _normalize_category(args.get("category")) + mime_prefix = _optional_text(args.get("mimePrefix")) + modified_after = _parse_datetime_arg(args.get("modifiedAfter")) + limit = _int_arg(args.get("limit"), default=50, minimum=1, maximum=200) + + root_folder_id = await _resolve_folder_id(ctx, folder_id=folder_id) + folder_ids = await _folder_scope_ids(ctx, root_folder_id=root_folder_id, recursive=recursive) + statement = _active_files_query(ctx, folder_ids=folder_ids) + if query: + statement = statement.where(File.file_name.ilike(f"%{query}%")) + if modified_after is not None: + statement = statement.where(File.updated_at >= modified_after) + statement = statement.order_by(File.file_name.asc()) + + items: list[dict[str, Any]] = [] + for row in list(await ctx.db.scalars(statement)): + mime_type = _resolved_mime(row) + if mime_prefix and not mime_type.lower().startswith(mime_prefix.lower()): + continue + if category is not None and _category_for_file(row) != category: + continue + items.append(await _file_payload(ctx, row, include_path=True)) + if len(items) >= limit: + break + + return { + "items": items, + "totalItems": len(items), + "query": query or None, + "folderId": str(root_folder_id), + "recursive": recursive, + "category": category, + "mimePrefix": mime_prefix, + "modifiedAfter": modified_after.isoformat() if modified_after else None, + } + + +async def _get_file_info(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + file_id = _parse_positive_int(_required_text(args, "fileId", "id"), "fileId") + row = await ctx.db.scalar( + select(File).where( + and_( + File.file_id == file_id, + File.owner_id == ctx.user_id, + File.status == FileStatus.ACTIVE, + File.is_latest.is_(True), + ) + ) + ) + if row is None: + raise ApiError(status_code=404, code=404, message="File not found") + storage = await ctx.db.get(StorageObject, int(row.storage_object_id)) + payload = await _file_payload(ctx, row, include_path=True) + payload.update( + { + "objectHash": str(storage.object_hash) if storage and storage.object_hash else None, + "hashAlgorithm": str(storage.hash_algorithm) if storage else None, + "storageObjectId": str(row.storage_object_id), + "category": _category_for_file(row), + } + ) + return payload + + +async def _list_recent(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + limit = _int_arg(args.get("limit"), default=20, minimum=1, maximum=50) + since = _parse_datetime_arg(args.get("since")) + statement = _active_files_query(ctx, folder_ids=None) + if since is not None: + statement = statement.where(File.updated_at >= since) + statement = statement.order_by(File.updated_at.desc(), File.file_id.desc()).limit(limit) + rows = list(await ctx.db.scalars(statement)) + return { + "items": [await _file_payload(ctx, row, include_path=True) for row in rows], + "totalItems": len(rows), + "limit": limit, + "since": since.isoformat() if since else None, + } + + +async def _stats_by_category(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = str(args.get("folderId") or "root") + recursive = _bool_arg(args.get("recursive"), default=True) + root_folder_id = await _resolve_folder_id(ctx, folder_id=folder_id) + folder_ids = await _folder_scope_ids(ctx, root_folder_id=root_folder_id, recursive=recursive) + rows = list(await ctx.db.scalars(_active_files_query(ctx, folder_ids=folder_ids))) + + categories = { + category: {"count": 0, "totalSize": 0} + for category in _CATEGORIES + } + total_size = 0 + for row in rows: + category = _category_for_file(row) + size = int(row.file_size or 0) + categories[category]["count"] += 1 + categories[category]["totalSize"] += size + total_size += size + + return { + "folderId": str(root_folder_id), + "recursive": recursive, + "totalItems": len(rows), + "totalSize": total_size, + "categories": categories, + "video": categories["video"]["count"], + "audio": categories["audio"]["count"], + "image": categories["image"]["count"], + "document": categories["document"]["count"], + "archive": categories["archive"]["count"], + "other": categories["other"]["count"], + } + + +async def _find_duplicates(ctx: ToolContext, args: dict[str, Any]) -> dict[str, Any]: + folder_id = str(args.get("folderId") or "root") + recursive = _bool_arg(args.get("recursive"), default=True) + by = str(args.get("by") or "hash").strip() or "hash" + if by not in {"hash", "nameSize"}: + raise ApiError(status_code=400, code=400, message="Invalid duplicate mode") + + root_folder_id = await _resolve_folder_id(ctx, folder_id=folder_id) + folder_ids = await _folder_scope_ids(ctx, root_folder_id=root_folder_id, recursive=recursive) + rows = ( + await ctx.db.execute( + _active_files_query(ctx, folder_ids=folder_ids) + .join(StorageObject, StorageObject.object_id == File.storage_object_id) + .add_columns(StorageObject.object_hash, StorageObject.hash_algorithm) + .order_by(File.file_name.asc()) + ) + ).all() + + groups: dict[str, dict[str, Any]] = {} + for row in rows: + file_row: File = row[0] + object_hash = row[1] + hash_algorithm = row[2] + if by == "hash": + if not object_hash: + continue + key = f"{hash_algorithm}:{object_hash}:{int(file_row.file_size or 0)}" + else: + key = f"{file_row.file_name.lower()}:{int(file_row.file_size or 0)}" + group = groups.setdefault( + key, + { + "key": key, + "by": by, + "hash": str(object_hash) if object_hash else None, + "hashAlgorithm": str(hash_algorithm) if hash_algorithm else None, + "size": int(file_row.file_size or 0), + "files": [], + }, + ) + group["files"].append(await _file_payload(ctx, file_row, include_path=True)) + + duplicate_groups = [group for group in groups.values() if len(group["files"]) > 1] + return { + "folderId": str(root_folder_id), + "recursive": recursive, + "by": by, + "groups": duplicate_groups, + "totalGroups": len(duplicate_groups), + "totalFiles": sum(len(group["files"]) for group in duplicate_groups), + } + + +def _active_files_query(ctx: ToolContext, *, folder_ids: list[int] | None): + statement = select(File).where( + and_( + File.owner_id == ctx.user_id, + File.status == FileStatus.ACTIVE, + File.is_latest.is_(True), + ) + ) + if folder_ids is not None: + statement = statement.where(File.folder_id.in_(folder_ids)) + return statement + + +async def _file_payload( + ctx: ToolContext, + row: File, + *, + include_path: bool, +) -> dict[str, Any]: + payload = { + "id": str(row.file_id), + "fileId": str(row.file_id), + "name": str(row.file_name), + "size": int(row.file_size or 0), + "mimeType": _resolved_mime(row), + "folderId": str(row.folder_id), + "createdAt": row.created_at.isoformat() if row.created_at else None, + "updatedAt": row.updated_at.isoformat() if row.updated_at else None, + } + if include_path: + folder_path = await _folder_path(ctx, folder_id=int(row.folder_id)) + payload["path"] = f"{folder_path}/{row.file_name}" if folder_path else str(row.file_name) + return payload + + +async def _resolve_folder_id(ctx: ToolContext, *, folder_id: str) -> int: + if not folder_id or folder_id == "root": + root_id = await ctx.db.scalar( + select(Folder.folder_id).where( + and_( + Folder.owner_id == ctx.user_id, + Folder.parent_folder_id.is_(None), + Folder.folder_type == FolderType.ROOT, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + if root_id is None: + raise ApiError(status_code=404, code=404, message="Root folder not found") + return int(root_id) + parsed = _parse_positive_int(folder_id, "folderId") + exists = await ctx.db.scalar( + select(Folder.folder_id).where( + and_( + Folder.folder_id == parsed, + Folder.owner_id == ctx.user_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + if exists is None: + raise ApiError(status_code=404, code=404, message="Folder not found") + return parsed + + +async def _folder_scope_ids( + ctx: ToolContext, + *, + root_folder_id: int, + recursive: bool, +) -> list[int]: + if not recursive: + return [root_folder_id] + return await _active_descendant_folder_ids(ctx, root_folder_id=root_folder_id) + + +async def _active_descendant_folder_ids(ctx: ToolContext, *, root_folder_id: int) -> list[int]: + descendants = ( + select(Folder.folder_id) + .where( + and_( + Folder.folder_id == root_folder_id, + Folder.owner_id == ctx.user_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + .cte(name="agent_tool_descendants", recursive=True) + ) + descendants = descendants.union_all( + select(Folder.folder_id).where( + and_( + Folder.parent_folder_id == descendants.c.folder_id, + Folder.owner_id == ctx.user_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + folder_ids = list(await ctx.db.scalars(select(descendants.c.folder_id))) + return [int(folder_id) for folder_id in folder_ids] + + +async def _folder_path(ctx: ToolContext, *, folder_id: int) -> str: + parts: list[str] = [] + current_id: int | None = folder_id + while current_id is not None: + folder = await ctx.db.scalar( + select(Folder).where( + and_( + Folder.folder_id == current_id, + Folder.owner_id == ctx.user_id, + Folder.status == FolderStatus.ACTIVE, + ) + ) + ) + if folder is None: + break + parts.append(str(folder.folder_name)) + current_id = int(folder.parent_folder_id) if folder.parent_folder_id is not None else None + parts.reverse() + return "/" + "/".join(parts) if parts else "" + + +def _first_value(args: dict[str, Any], *keys: str) -> Any: + for key in keys: + value = args.get(key) + if value not in (None, ""): + return value + return None + + +def _required_text(args: dict[str, Any], *keys: str) -> str: + value = _first_value(args, *keys) + if value is None: + raise ApiError(status_code=400, code=400, message=f"Missing required tool input: {keys[0]}") + text = str(value).strip() + if not text: + raise ApiError(status_code=400, code=400, message=f"Missing required tool input: {keys[0]}") + return text + + +def _optional_text(value: Any) -> str | None: + text = str(value or "").strip() + return text or None + + +def _bool_arg(value: Any, *, default: bool) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + text = str(value).strip().lower() + if text in {"1", "true", "yes", "y"}: + return True + if text in {"0", "false", "no", "n"}: + return False + return default + + +def _int_arg( + value: Any, + *, + default: int, + minimum: int, + maximum: int | None = None, +) -> int: + try: + parsed = int(value if value is not None else default) + except (TypeError, ValueError): + parsed = default + parsed = max(minimum, parsed) + if maximum is not None: + parsed = min(maximum, parsed) + return parsed + + +def _parse_positive_int(value: str, field_name: str) -> int: + try: + parsed = int(value) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message=f"Invalid {field_name}") from exc + if parsed <= 0: + raise ApiError(status_code=400, code=400, message=f"Invalid {field_name}") + return parsed + + +def _parse_datetime_arg(value: Any) -> datetime | None: + text = str(value or "").strip() + if not text: + return None + try: + parsed = datetime.fromisoformat(text.replace("Z", "+00:00")) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid datetime") from exc + if parsed.tzinfo is None: + return parsed.replace(tzinfo=UTC) + return parsed + + +def _normalize_category(value: Any) -> str | None: + text = str(value or "").strip().lower() + aliases = { + "movies": "video", + "movie": "video", + "film": "video", + "films": "video", + "videos": "video", + "anime": "video", + "animation": "video", + "视频": "video", + "影片": "video", + "电影": "video", + "动漫": "video", + "番剧": "video", + "documents": "document", + "docs": "document", + "images": "image", + "pictures": "image", + "archives": "archive", + "compressed": "archive", + } + text = aliases.get(text, text) + if text in _CATEGORIES: + return text + return None + + +def _resolved_mime(row: File) -> str: + return resolve_file_mime_type( + mime_type=row.mime_type, + file_ext=row.file_ext, + file_name=row.file_name, + ) + + +def _category_for_file(row: File) -> str: + mime = _resolved_mime(row).lower() + ext = _normalized_extension(row.file_ext) or _filename_extension(row.file_name) + if mime.startswith("video/") or ext in {"mp4", "mov", "avi", "mkv", "webm", "m4v"}: + return "video" + if mime.startswith("audio/") or ext in {"mp3", "wav", "flac", "m4a", "aac", "ogg"}: + return "audio" + if mime.startswith("image/") or ext in {"jpg", "jpeg", "png", "gif", "webp", "svg", "bmp"}: + return "image" + if mime in {"application/pdf"} or ext in {"pdf", "doc", "docx", "xls", "xlsx", "ppt", "pptx", "txt", "md"}: + return "document" + if ext in {"zip", "rar", "7z", "tar", "gz", "bz2", "xz"}: + return "archive" + return "other" + + +def _count_files_answer(output: dict[str, Any]) -> str: + total_items = int(output.get("totalItems") or 0) + category = str(output.get("category") or "").strip().lower() + qualifier = _search_qualifier(output) + if category == "video": + return f"你上传了 {total_items} 部{qualifier}电影（按视频文件统计）。" + if category == "audio": + return f"你上传了 {total_items} 个{qualifier}音频文件。" + if category == "image": + return f"你上传了 {total_items} 张{qualifier}图片。" + if category == "document": + return f"你上传了 {total_items} 个{qualifier}文档。" + if category == "archive": + return f"你上传了 {total_items} 个{qualifier}压缩包。" + return f"你上传了 {total_items} 个{qualifier}文件。" + + +def _search_qualifier(output: dict[str, Any]) -> str: + search = str(output.get("search") or "").strip() + if not search: + return "" + return f"名称包含“{search}”的" + + +def _normalized_extension(value: str | None) -> str: + return str(value or "").strip().lower().lstrip(".") + + +def _filename_extension(value: str | None) -> str: + name = str(value or "").strip().lower() + if "." not in name: + return "" + return name.rsplit(".", 1)[-1] + + +def _schema(properties: dict[str, Any], required: list[str] | None = None) -> dict[str, Any]: + return { + "type": "object", + "properties": properties, + "required": required or [], + } + + +_FOLDER_ID = {"type": "string", "description": "Folder id, or root for the user's root folder."} +_FILE_ID = {"type": "string", "description": "File id owned by the current user."} +_CATEGORY = {"type": "string", "enum": list(_CATEGORIES)} +_SHARE_HANDLING = {"type": "string", "enum": ["keep", "revoke"], "default": "keep"} + + +REGISTRY.register( + ToolSpec( + name="drive.listFolder", + description="List direct files and folders inside a folder.", + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "page": {"type": "integer", "minimum": 1, "default": 1}, + "perPage": {"type": "integer", "minimum": 1, "maximum": 200, "default": 200}, + "search": {"type": "string"}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_list_folder, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.countFiles", + description=( + "Count files under a folder. Supports recursive counts, broad file categories, " + "and filename contains search." + ), + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "recursive": {"type": "boolean", "default": True}, + "category": _CATEGORY, + "search": {"type": "string"}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_count_files, + answer_formatter=_count_files_answer, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.createFolder", + description="Create a folder under parentFolderId with name.", + input_schema=_schema( + { + "parentFolderId": _FOLDER_ID, + "name": {"type": "string", "minLength": 1, "maxLength": 255}, + }, + required=["name"], + ), + side_effect="write", + risk_level="medium", + requires_confirmation=False, + handler=_create_folder, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.moveFile", + description="Move fileId into targetFolderId.", + input_schema=_schema( + { + "fileId": _FILE_ID, + "targetFolderId": _FOLDER_ID, + "shareHandling": _SHARE_HANDLING, + }, + required=["fileId", "targetFolderId"], + ), + side_effect="write", + risk_level="medium", + requires_confirmation=False, + handler=_move_file, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.moveFolder", + description="Move folderId into targetParentId.", + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "targetParentId": _FOLDER_ID, + "shareHandling": _SHARE_HANDLING, + }, + required=["folderId", "targetParentId"], + ), + side_effect="write", + risk_level="medium", + requires_confirmation=False, + handler=_move_folder, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.renameFile", + description="Rename fileId to fileName.", + input_schema=_schema( + { + "fileId": _FILE_ID, + "fileName": {"type": "string", "minLength": 1, "maxLength": 255}, + }, + required=["fileId", "fileName"], + ), + side_effect="write", + risk_level="medium", + requires_confirmation=False, + handler=_rename_file, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.renameFolder", + description="Rename folderId to folderName.", + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "folderName": {"type": "string", "minLength": 1, "maxLength": 255}, + }, + required=["folderId", "folderName"], + ), + side_effect="write", + risk_level="medium", + requires_confirmation=False, + handler=_rename_folder, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.deleteFile", + description="Soft-delete fileId into the recycle bin. This is high risk.", + input_schema=_schema({"fileId": _FILE_ID}, required=["fileId"]), + side_effect="write", + risk_level="high", + requires_confirmation=True, + handler=_delete_file, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.deleteFolder", + description="Soft-delete folderId into the recycle bin. This is high risk.", + input_schema=_schema({"folderId": _FOLDER_ID}, required=["folderId"]), + side_effect="write", + risk_level="high", + requires_confirmation=True, + handler=_delete_folder, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.searchFiles", + description="Search active files by filename, folder scope, category, MIME prefix, and update time.", + input_schema=_schema( + { + "query": {"type": "string"}, + "folderId": _FOLDER_ID, + "recursive": {"type": "boolean", "default": True}, + "category": _CATEGORY, + "mimePrefix": {"type": "string", "description": "MIME type prefix such as video/."}, + "modifiedAfter": {"type": "string", "format": "date-time"}, + "limit": {"type": "integer", "minimum": 1, "maximum": 200, "default": 50}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_search_files, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.getFileInfo", + description="Return detailed metadata for one active file.", + input_schema=_schema({"fileId": _FILE_ID}, required=["fileId"]), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_get_file_info, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.listRecent", + description="List recently updated active files.", + input_schema=_schema( + { + "limit": {"type": "integer", "minimum": 1, "maximum": 50, "default": 20}, + "since": {"type": "string", "format": "date-time"}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_list_recent, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.statsByCategory", + description="Compute file counts and sizes by broad category under a folder.", + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "recursive": {"type": "boolean", "default": True}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_stats_by_category, + ) +) + +REGISTRY.register( + ToolSpec( + name="drive.findDuplicates", + description="Find duplicate active files by content hash or by name plus size.", + input_schema=_schema( + { + "folderId": _FOLDER_ID, + "recursive": {"type": "boolean", "default": True}, + "by": {"type": "string", "enum": ["hash", "nameSize"], "default": "hash"}, + } + ), + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_find_duplicates, + ) +) + +__all__ = [] diff --git a/app/src/fileflash/agents/worker.py b/app/src/fileflash/agents/worker.py new file mode 100644 index 0000000..a122e31 --- /dev/null +++ b/app/src/fileflash/agents/worker.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +import asyncio +import logging +import uuid +from datetime import UTC, datetime +from typing import Any + +from fastapi.encoders import jsonable_encoder +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from ..core.errors import ApiError +from ..core.settings import get_settings +from ..db.session import SessionLocal +from ..db.transaction import apply_local_lock_timeout +from ..models import BackgroundJob +from ..services.job_queue import RedisStreamJobQueue +from ..workers.contracts import WorkerJobMessage +from .harness.event_bus import AgentEventBus, AgentEventEnvelope, build_agent_event_bus +from .runtime import AgentJobCanceled, ExecuteRunner, PlanRunner + +logger = logging.getLogger(__name__) + + +class AgentWorkerConsumer: + def __init__( + self, + *, + queue: RedisStreamJobQueue, + session_factory: async_sessionmaker[AsyncSession] = SessionLocal, + event_bus: AgentEventBus | None = None, + ) -> None: + self._settings = get_settings() + self._queue = queue + self._session_factory = session_factory + self._event_bus = event_bus or build_agent_event_bus(settings=self._settings) + + async def run(self) -> None: + logger.info( + "Agent worker started queue=%s group=%s concurrency=%s", + self._settings.agent_queue_stream, + self._settings.agent_queue_group, + self._settings.agent_worker_concurrency, + ) + async with asyncio.TaskGroup() as group: + for slot in range(max(1, self._settings.agent_worker_concurrency)): + group.create_task(self._run_slot(slot)) + + async def _run_slot(self, slot: int) -> None: + while True: + queued = await self._queue.consume_one(block_ms=self._settings.agent_queue_block_ms) + if queued is None: + continue + queue_message_id, message = queued + try: + await self._process_message(slot=slot, message=message) + finally: + await self._queue.ack(queue_message_id) + + async def _process_message(self, *, slot: int, message: WorkerJobMessage) -> None: + job = await self._mark_running(message) + if job is None: + return + + started = datetime.now(UTC) + try: + result, phase = await asyncio.wait_for( + self._run_job(job=job), + timeout=self._settings.agent_job_timeout_sec, + ) + except AgentJobCanceled: + await self._mark_canceled(job_id=message.job_id) + logger.info( + "Agent slot=%s canceled jobId=%s taskType=%s", + slot, + message.job_id, + message.task_type, + ) + return + except Exception as exc: + await self._mark_failed(job_id=message.job_id, error=exc) + logger.warning( + "Agent slot=%s failed jobId=%s taskType=%s error=%s", + slot, + message.job_id, + message.task_type, + exc, + ) + return + + await self._mark_succeeded(job_id=message.job_id, result=result, phase=phase) + duration_ms = int((datetime.now(UTC) - started).total_seconds() * 1000) + logger.info( + "Agent slot=%s succeeded jobId=%s taskType=%s durationMs=%s traceId=%s", + slot, + message.job_id, + message.task_type, + duration_ms, + message.trace_id, + ) + + async def _run_job(self, *, job: BackgroundJob) -> tuple[dict[str, Any], str]: + async with self._session_factory() as db: + fresh_job = await db.get(BackgroundJob, int(job.job_id)) + if fresh_job is None: + raise ApiError(status_code=404, code=404, message="Job not found") + if fresh_job.task_type == "agent.plan": + result = await PlanRunner( + settings=self._settings, + event_bus=self._event_bus, + ).run(db=db, job=fresh_job) + phase = "awaiting_confirm" if result.requires_confirmation else "completed" + return result.model_dump(by_alias=True, mode="json"), phase + if fresh_job.task_type == "agent.execute": + result = await ExecuteRunner(event_bus=self._event_bus).run(db=db, job=fresh_job) + return result.model_dump(by_alias=True, mode="json"), "completed" + raise ApiError( + status_code=400, + code=400, + message=f"Unsupported agent task: {fresh_job.task_type}", + ) + + async def _mark_running(self, message: WorkerJobMessage) -> BackgroundJob | None: + async with self._session_factory() as db: + async with db.begin(): + await apply_local_lock_timeout(db) + job = await db.scalar( + select(BackgroundJob) + .where(BackgroundJob.job_id == message.job_id) + .with_for_update() + ) + if job is None or job.status in {"succeeded", "failed", "canceled"}: + return None + now = datetime.now(UTC) + job.status = "running" + job.started_at = job.started_at or now + job.updated_at = now + job.error_message = None + job.agent_phase = "planning" if job.task_type == "agent.plan" else "executing" + return job + + async def _mark_succeeded(self, *, job_id: int, result: dict[str, Any], phase: str) -> None: + safe_result = jsonable_encoder(result) + should_publish = False + async with self._session_factory() as db: + async with db.begin(): + await apply_local_lock_timeout(db) + job = await db.scalar( + select(BackgroundJob).where(BackgroundJob.job_id == job_id).with_for_update() + ) + if job is None: + return + if job.status == "canceled" or job.cancel_requested_at is not None: + return + now = datetime.now(UTC) + job.status = "succeeded" + job.result = safe_result + job.error_message = None + job.agent_phase = phase + job.finished_at = now + job.updated_at = now + should_publish = True + if should_publish: + await self._publish_terminal( + job_id=job_id, + event_type="job.succeeded", + payload={"status": "succeeded", "agentPhase": phase, "data": {"result": safe_result}}, + ) + + async def _mark_failed(self, *, job_id: int, error: Exception) -> None: + message = _error_message(error) + should_publish = False + async with self._session_factory() as db: + async with db.begin(): + await apply_local_lock_timeout(db) + job = await db.scalar( + select(BackgroundJob).where(BackgroundJob.job_id == job_id).with_for_update() + ) + if job is None: + return + if job.status == "canceled" or job.cancel_requested_at is not None: + return + now = datetime.now(UTC) + job.status = "failed" + job.agent_phase = "failed" + job.error_message = message[:2000] + job.finished_at = now + job.updated_at = now + should_publish = True + if should_publish: + await self._publish_terminal( + job_id=job_id, + event_type="job.failed", + payload={ + "status": "failed", + "agentPhase": "failed", + "message": message[:2000], + "data": {"errorMessage": message[:2000]}, + }, + ) + + async def _mark_canceled(self, *, job_id: int) -> None: + should_publish = False + async with self._session_factory() as db: + async with db.begin(): + await apply_local_lock_timeout(db) + job = await db.scalar( + select(BackgroundJob).where(BackgroundJob.job_id == job_id).with_for_update() + ) + if job is None: + return + now = datetime.now(UTC) + job.status = "canceled" + job.agent_phase = "canceled" + job.cancel_requested_at = job.cancel_requested_at or now + job.finished_at = now + job.updated_at = now + should_publish = True + if should_publish: + await self._publish_terminal( + job_id=job_id, + event_type="job.canceled", + payload={"status": "canceled", "agentPhase": "canceled"}, + ) + + async def _publish_terminal( + self, + *, + job_id: int, + event_type: str, + payload: dict[str, Any] | None = None, + ) -> None: + try: + await self._event_bus.publish( + AgentEventEnvelope( + job_id=job_id, + event_type=event_type, + payload=payload or {}, + emitted_at=datetime.now(UTC), + ) + ) + except Exception: + logger.exception( + "Failed to publish terminal event jobId=%s eventType=%s", + job_id, + event_type, + ) + + +def _error_message(error: Exception) -> str: + if isinstance(error, ApiError): + return f"ApiError[{error.status_code}/{error.code}]: {error.message}" + return f"{type(error).__name__}: {error}" + + +async def run_agent_worker() -> None: + settings = get_settings() + queue = RedisStreamJobQueue( + redis_url=settings.redis_url, + stream_key=settings.agent_queue_stream, + group_name=settings.agent_queue_group, + consumer_name=f"agent-{uuid.uuid4().hex[:8]}", + ) + consumer = AgentWorkerConsumer(queue=queue) + try: + await consumer.run() + finally: + await queue.aclose() + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s [%(name)s] %(message)s", + ) + try: + asyncio.run(run_agent_worker()) + except KeyboardInterrupt: + logger.info("Agent worker stopped by keyboard interrupt") + + +if __name__ == "__main__": + main() diff --git a/app/src/fileflash/core/deps.py b/app/src/fileflash/core/deps.py index c489641..5ac1f45 100644 --- a/app/src/fileflash/core/deps.py +++ b/app/src/fileflash/core/deps.py @@ -5,9 +5,10 @@ from jwt import InvalidTokenError from sqlalchemy.ext.asyncio import AsyncSession +from ..agents.harness.event_bus import AgentEventBus, build_agent_event_bus from ..db.deps import get_db -from ..models.tables_identity import User from ..models.enums import UserRole +from ..models.tables_identity import User from ..repositories import ( AgentActionLogRepository, AgentMcpRepository, @@ -17,10 +18,27 @@ AgentSkillRepository, AgentWorkSessionRepository, ) +from ..s3 import MinioObjectStorageClient +from ..services.admin.files import AdminFilesService +from ..services.admin.logs import AdminLogsService +from ..services.admin.moderation import AdminModerationService +from ..services.admin.notifications import AdminNotificationsService +from ..services.admin.storage import AdminStorageService +from ..services.admin.system import AdminSystemService +from ..services.admin.users import AdminUsersService +from ..services.agent import ( + ExecuteService, + McpService, + MemoryService, + PlanService, + SessionService, + SettingsService, + SkillService, +) from ..services.archive import ArchiveService -from ..services.agent import ExecuteService, McpService, MemoryService, PlanService, SessionService, SettingsService, SkillService from ..services.auth import AuthService from ..services.background_jobs import BackgroundJobService +from ..services.download_rate_limit import DownloadRateLimitService from ..services.email_delivery import VerificationEmailDeliveryService from ..services.file import FileService from ..services.folder import FolderService @@ -30,7 +48,6 @@ from ..services.registration_email_domain_rule import RegistrationEmailDomainRuleService from ..services.share import ShareService from ..services.upload import UploadService -from ..s3 import MinioObjectStorageClient from .errors import ApiError from .security import decode_access_token from .settings import Settings, get_settings @@ -48,6 +65,7 @@ redis_url=_settings.redis_url, stream_key=_settings.agent_queue_stream, ) +_agent_event_bus_singleton: AgentEventBus | None = None def get_rate_limiter() -> RedisRateLimiter: @@ -70,6 +88,13 @@ def get_agent_job_queue_publisher() -> RedisStreamJobQueue: return _agent_job_queue_publisher +def get_agent_event_bus() -> AgentEventBus: + global _agent_event_bus_singleton + if _agent_event_bus_singleton is None: + _agent_event_bus_singleton = build_agent_event_bus(settings=_settings) + return _agent_event_bus_singleton + + def get_background_job_service( queue_publisher: RedisStreamJobQueue = Depends(get_job_queue_publisher), ) -> BackgroundJobService: @@ -86,6 +111,14 @@ def get_settings_dep() -> Settings: return _settings +def get_download_rate_limit_service( + db: AsyncSession = Depends(get_db), + settings: Settings = Depends(get_settings_dep), + rate_limiter: RedisRateLimiter = Depends(get_rate_limiter), +) -> DownloadRateLimitService: + return DownloadRateLimitService(db=db, settings=settings, rate_limiter=rate_limiter) + + def get_client_ip(request: Request) -> str: forwarded_for = request.headers.get("x-forwarded-for") if forwarded_for: @@ -120,6 +153,52 @@ def get_registration_email_domain_rule_service( return RegistrationEmailDomainRuleService(db=db) +def get_admin_users_service( + db: AsyncSession = Depends(get_db), +) -> AdminUsersService: + return AdminUsersService(db=db) + + +def get_admin_storage_service( + db: AsyncSession = Depends(get_db), + rate_limiter: RedisRateLimiter = Depends(get_rate_limiter), +) -> AdminStorageService: + return AdminStorageService(db=db, redis=getattr(rate_limiter, "_redis", None)) + + +def get_admin_files_service( + db: AsyncSession = Depends(get_db), + event_publisher: InProcessAuthEventPublisher = Depends(get_event_publisher), + storage: MinioObjectStorageClient = Depends(get_object_storage), +) -> AdminFilesService: + return AdminFilesService(db=db, publisher=event_publisher, storage=storage) + + +def get_admin_moderation_service( + db: AsyncSession = Depends(get_db), +) -> AdminModerationService: + return AdminModerationService(db=db) + + +def get_admin_logs_service( + db: AsyncSession = Depends(get_db), +) -> AdminLogsService: + return AdminLogsService(db=db) + + +def get_admin_notifications_service( + db: AsyncSession = Depends(get_db), +) -> AdminNotificationsService: + return AdminNotificationsService(db=db) + + +def get_admin_system_service( + db: AsyncSession = Depends(get_db), + settings: Settings = Depends(get_settings_dep), +) -> AdminSystemService: + return AdminSystemService(db=db, settings=settings) + + def get_upload_service( db: AsyncSession = Depends(get_db), settings: Settings = Depends(get_settings_dep), @@ -198,6 +277,7 @@ def get_agent_work_session_repository(db: AsyncSession = Depends(get_db)) -> Age def get_agent_plan_service( db: AsyncSession = Depends(get_db), + settings: Settings = Depends(get_settings_dep), jobs: BackgroundJobService = Depends(get_agent_background_job_service), plans: AgentPlanRepository = Depends(get_agent_plan_repository), settings_repo: AgentSettingsRepository = Depends(get_agent_settings_repository), @@ -205,21 +285,24 @@ def get_agent_plan_service( ) -> PlanService: return PlanService( db=db, + settings=settings, jobs=jobs, plans=plans, - settings=settings_repo, + settings_repo=settings_repo, work_sessions=work_sessions, ) def get_agent_execute_service( db: AsyncSession = Depends(get_db), + settings: Settings = Depends(get_settings_dep), jobs: BackgroundJobService = Depends(get_agent_background_job_service), plans: AgentPlanRepository = Depends(get_agent_plan_repository), work_sessions: AgentWorkSessionRepository = Depends(get_agent_work_session_repository), ) -> ExecuteService: return ExecuteService( db=db, + settings=settings, jobs=jobs, plans=plans, work_sessions=work_sessions, diff --git a/app/src/fileflash/core/security.py b/app/src/fileflash/core/security.py index 63f3dda..c932376 100644 --- a/app/src/fileflash/core/security.py +++ b/app/src/fileflash/core/security.py @@ -70,3 +70,61 @@ def decode_share_access_token(token: str, settings: Settings) -> dict[str, Any]: if token_type != "share": raise jwt.InvalidTokenError("Invalid token type") return payload + + +def create_file_preview_token( + *, + user_id: int, + file_id: int, + settings: Settings, + expires_at: datetime, +) -> str: + now = datetime.now(UTC) + payload: dict[str, Any] = { + "sub": str(user_id), + "typ": "file_preview", + "scope": "file.preview", + "fileId": str(file_id), + "iat": int(now.timestamp()), + "exp": int(expires_at.timestamp()), + "jti": str(uuid.uuid4()), + } + return jwt.encode(payload, settings.jwt_secret_key, algorithm=settings.jwt_algorithm) + + +def decode_file_preview_token(token: str, settings: Settings) -> dict[str, Any]: + payload = jwt.decode(token, settings.jwt_secret_key, algorithms=[settings.jwt_algorithm]) + token_type = payload.get("typ") + scope = payload.get("scope") + if token_type != "file_preview" or scope != "file.preview": + raise jwt.InvalidTokenError("Invalid token type") + return payload + + +def create_admin_file_preview_token( + *, + admin_user_id: int, + file_id: int, + settings: Settings, + expires_at: datetime, +) -> str: + now = datetime.now(UTC) + payload: dict[str, Any] = { + "sub": str(admin_user_id), + "typ": "admin_file_preview", + "scope": "admin.file.preview", + "fileId": str(file_id), + "iat": int(now.timestamp()), + "exp": int(expires_at.timestamp()), + "jti": str(uuid.uuid4()), + } + return jwt.encode(payload, settings.jwt_secret_key, algorithm=settings.jwt_algorithm) + + +def decode_admin_file_preview_token(token: str, settings: Settings) -> dict[str, Any]: + payload = jwt.decode(token, settings.jwt_secret_key, algorithms=[settings.jwt_algorithm]) + token_type = payload.get("typ") + scope = payload.get("scope") + if token_type != "admin_file_preview" or scope != "admin.file.preview": + raise jwt.InvalidTokenError("Invalid token type") + return payload diff --git a/app/src/fileflash/core/settings.py b/app/src/fileflash/core/settings.py index bc7e066..71dde62 100644 --- a/app/src/fileflash/core/settings.py +++ b/app/src/fileflash/core/settings.py @@ -28,6 +28,10 @@ class Settings(BaseSettings): api_v1_prefix: str = "/api/v1" app_env: str = Field(default="production", alias="APP_ENV") + default_admin_username: str | None = Field(default=None, alias="DEFAULT_ADMIN_USERNAME") + default_admin_email: str | None = Field(default=None, alias="DEFAULT_ADMIN_EMAIL") + default_admin_password: str | None = Field(default=None, alias="DEFAULT_ADMIN_PASSWORD") + database_url: str | None = Field(default=None, alias="DATABASE_URL") ff_db_uri: str | None = Field(default=None, alias="FF_DB_URI") @@ -39,6 +43,10 @@ class Settings(BaseSettings): jwt_algorithm: str = "HS256" access_token_expire_minutes: int = 60 * 24 * 3 refresh_token_expire_days: int = 7 + file_preview_url_ttl_seconds: int = Field( + default=4 * 60 * 60, + alias="FILE_PREVIEW_URL_TTL_SECONDS", + ) refresh_cookie_name: str = "refreshToken" refresh_cookie_secure: bool = False @@ -48,6 +56,22 @@ class Settings(BaseSettings): cors_origins: list[str] = Field(default_factory=lambda: ["http://localhost:5173", "http://localhost:8080"]) redis_url: str | None = Field(default=None, alias="REDIS_URL") + agent_inbox_ask_timeout_sec: int = Field( + default=1800, + alias="AGENT_INBOX_ASK_TIMEOUT_SEC", + ) + agent_event_channel_prefix: str = Field( + default="agent:job", + alias="AGENT_EVENT_CHANNEL_PREFIX", + ) + agent_inbox_channel_prefix: str = Field( + default="agent:inbox", + alias="AGENT_INBOX_CHANNEL_PREFIX", + ) + agent_event_bus_buffer_size: int = Field( + default=64, + alias="AGENT_EVENT_BUS_BUFFER_SIZE", + ) rabbitmq_url: str | None = Field(default=None, alias="RABBITMQ_URL") email_verify_base_url: str = Field(default="", alias="EMAIL_VERIFY_BASE_URL") @@ -68,28 +92,44 @@ class Settings(BaseSettings): object_storage_secure: bool = Field(default=False, alias="OBJECT_STORAGE_SECURE") object_storage_region: str | None = Field(default=None, alias="OBJECT_STORAGE_REGION") - upload_chunk_size_default: int = Field(default=5 * 1024 * 1024, alias="UPLOAD_CHUNK_SIZE_DEFAULT") + upload_chunk_size_default: int = Field( + default=5 * 1024 * 1024, + alias="UPLOAD_CHUNK_SIZE_DEFAULT", + ) upload_chunk_size_min: int = Field(default=1 * 1024 * 1024, alias="UPLOAD_CHUNK_SIZE_MIN") upload_chunk_size_max: int = Field(default=16 * 1024 * 1024, alias="UPLOAD_CHUNK_SIZE_MAX") - upload_single_file_size_max: int = Field(default=5 * 1024 * 1024 * 1024, alias="UPLOAD_SINGLE_FILE_SIZE_MAX") + upload_single_file_size_max: int = Field( + default=5 * 1024 * 1024 * 1024, + alias="UPLOAD_SINGLE_FILE_SIZE_MAX", + ) + upload_verify_merged_object_hash: bool = Field( + default=False, + alias="UPLOAD_VERIFY_MERGED_OBJECT_HASH", + ) starred_items_limit: int = Field(default=20, alias="STARRED_ITEMS_LIMIT") upload_session_ttl_hours: int = Field(default=24, alias="UPLOAD_SESSION_TTL_HOURS") upload_temp_prefix: str = Field(default="tmp", alias="UPLOAD_TEMP_PREFIX") upload_object_prefix: str = Field(default="objects", alias="UPLOAD_OBJECT_PREFIX") - max_failed_login_attempts: int = 5 - account_lock_minutes: int = 15 + max_failed_login_attempts: int = 8 + account_lock_minutes: int = 5 email_verification_expire_minutes: int = 60 password_reset_expire_minutes: int = 30 - register_rate_limit: int = 5 + register_rate_limit: int = 12 register_rate_window_seconds: int = 600 - login_rate_limit: int = 10 + login_rate_limit: int = 30 login_rate_window_seconds: int = 300 forgot_password_rate_limit: int = 5 forgot_password_rate_window_seconds: int = 600 resend_verification_rate_limit: int = 5 resend_verification_rate_window_seconds: int = 600 + download_rate_window_seconds: int = Field(default=600, alias="DOWNLOAD_RATE_WINDOW_SECONDS") + download_rate_limit_requests: int = Field(default=120, alias="DOWNLOAD_RATE_LIMIT_REQUESTS") + download_rate_limit_bytes: int = Field( + default=2 * 1024 * 1024 * 1024, + alias="DOWNLOAD_RATE_LIMIT_BYTES", + ) worker_poll_interval_seconds: float = Field( default=2.0, @@ -124,9 +164,10 @@ class Settings(BaseSettings): agent_user_concurrent_limit: int = Field(default=2, alias="AGENT_USER_CONCURRENT_LIMIT") agent_staging_ttl_sec: int = Field(default=86400, alias="AGENT_STAGING_TTL_SEC") agent_sse_enabled: bool = Field(default=False, alias="AGENT_SSE_ENABLED") - agent_llm_provider: str = Field(default="anthropic", alias="AGENT_LLM_PROVIDER") agent_llm_model: str = Field(default="claude-sonnet-4-6", alias="AGENT_LLM_MODEL") + agent_llm_base_url: str | None = Field(default=None, alias="AGENT_LLM_BASE_URL") agent_llm_api_key: str | None = Field(default=None, alias="AGENT_LLM_API_KEY") + agent_llm_plan_max_tokens: int = Field(default=8192, alias="AGENT_LLM_PLAN_MAX_TOKENS") agent_mcp_endpoints_raw: str = Field(default="[]", alias="AGENT_MCP_ENDPOINTS") ffmpeg_binary: str = Field(default="ffmpeg", alias="FFMPEG_BINARY") @@ -184,6 +225,25 @@ def security_configuration_issues(self) -> tuple[str, ...]: token_hash_secret = (self.token_hash_secret or "").strip() if token_hash_secret and len(token_hash_secret.encode("utf-8")) < self.MIN_SECRET_LENGTH: issues.append(f"TOKEN_HASH_SECRET must be at least {self.MIN_SECRET_LENGTH} bytes") + issues.extend(self.default_admin_configuration_issues) + return tuple(issues) + + @property + def default_admin_configuration_issues(self) -> tuple[str, ...]: + if not self.is_production_env: + return () + + issues: list[str] = [] + if not (self.default_admin_username or "").strip(): + issues.append("DEFAULT_ADMIN_USERNAME is required in production") + if not (self.default_admin_email or "").strip(): + issues.append("DEFAULT_ADMIN_EMAIL is required in production") + + password = (self.default_admin_password or "").strip() + if not password: + issues.append("DEFAULT_ADMIN_PASSWORD is required in production") + elif len(password.encode("utf-8")) < self.MIN_SECRET_LENGTH: + issues.append(f"DEFAULT_ADMIN_PASSWORD must be at least {self.MIN_SECRET_LENGTH} bytes") return tuple(issues) def assert_runtime_security(self) -> None: diff --git a/app/src/fileflash/exceptions/__init__.py b/app/src/fileflash/exceptions/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/app/src/fileflash/models/__init__.py b/app/src/fileflash/models/__init__.py index 575fa97..21cd3a5 100644 --- a/app/src/fileflash/models/__init__.py +++ b/app/src/fileflash/models/__init__.py @@ -2,6 +2,7 @@ from .tables import ( Acl, AgentActionLog, + AgentInboxMessage, AgentMcpServer, AgentMemory, AgentPlan, @@ -40,6 +41,7 @@ __all__ = [ "Acl", "AgentActionLog", + "AgentInboxMessage", "AgentMcpServer", "AgentMemory", "AgentPlan", diff --git a/app/src/fileflash/models/enums.py b/app/src/fileflash/models/enums.py index 0ce1e32..7228c78 100644 --- a/app/src/fileflash/models/enums.py +++ b/app/src/fileflash/models/enums.py @@ -147,6 +147,29 @@ class AgentMcpVisibility(BaseStrEnum): PRIVATE = "private" +class AgentInboxRole(BaseStrEnum): + AGENT = "agent" + USER = "user" + + +class AgentInboxKind(BaseStrEnum): + ASK = "ask" + REPLY = "reply" + CONTROL_PAUSE = "control.pause" + CONTROL_RESUME = "control.resume" + CONTROL_APPROVE = "control.approve" + CONTROL_DENY = "control.deny" + CONTROL_SKIP = "control.skip" + CONTROL_CANCEL = "control.cancel" + + +class AgentInboxStatus(BaseStrEnum): + WAITING = "waiting" + ANSWERED = "answered" + TIMED_OUT = "timed_out" + DROPPED = "dropped" + + __all__ = [ "BaseStrEnum", "UploadStatus", @@ -172,5 +195,8 @@ class AgentMcpVisibility(BaseStrEnum): "AgentMemoryKind", "AgentSkillVisibility", "AgentMcpVisibility", + "AgentInboxRole", + "AgentInboxKind", + "AgentInboxStatus", ] diff --git a/app/src/fileflash/models/tables.py b/app/src/fileflash/models/tables.py index 0040f7b..5f4930c 100644 --- a/app/src/fileflash/models/tables.py +++ b/app/src/fileflash/models/tables.py @@ -11,6 +11,7 @@ ) from .tables_agent import ( AgentActionLog, + AgentInboxMessage, AgentMcpServer, AgentMemory, AgentPlan, @@ -50,6 +51,7 @@ __all__ = [ "Acl", "AgentActionLog", + "AgentInboxMessage", "AgentMcpServer", "AgentMemory", "AgentPlan", diff --git a/app/src/fileflash/models/tables_agent.py b/app/src/fileflash/models/tables_agent.py index 33b4e79..3ad9591 100644 --- a/app/src/fileflash/models/tables_agent.py +++ b/app/src/fileflash/models/tables_agent.py @@ -21,6 +21,9 @@ from .base import Base from .enums import ( AgentExecutionPolicy, + AgentInboxKind, + AgentInboxRole, + AgentInboxStatus, AgentMcpVisibility, AgentMemoryKind, AgentMemoryScope, @@ -345,8 +348,55 @@ class AgentWorkSession(Base): closed_at: Mapped[datetime | None] = mapped_column(DateTime) +class AgentInboxMessage(Base): + __tablename__ = "agent_inbox_message" + __table_args__ = ( + Index("idx_agent_inbox_message_job_created", "job_id", "created_at"), + Index( + "idx_agent_inbox_message_job_status", + "job_id", + "status", + postgresql_where=text("status IS NOT NULL"), + ), + ) + + inbox_message_id: Mapped[int] = mapped_column(BigInteger, Identity(), primary_key=True) + job_id: Mapped[int] = mapped_column( + BigInteger, + ForeignKey("background_job.job_id", ondelete="CASCADE"), + nullable=False, + ) + role: Mapped[AgentInboxRole] = mapped_column( + pg_enum(AgentInboxRole, "agent_inbox_role_enum"), + nullable=False, + ) + kind: Mapped[AgentInboxKind] = mapped_column( + pg_enum(AgentInboxKind, "agent_inbox_kind_enum"), + nullable=False, + ) + payload_json: Mapped[dict[str, Any]] = mapped_column( + JSONB, + nullable=False, + server_default=text("'{}'::jsonb"), + ) + reply_to_id: Mapped[int | None] = mapped_column( + BigInteger, + ForeignKey("agent_inbox_message.inbox_message_id", ondelete="SET NULL"), + ) + status: Mapped[AgentInboxStatus | None] = mapped_column( + pg_enum(AgentInboxStatus, "agent_inbox_status_enum"), + ) + created_at: Mapped[datetime] = mapped_column( + DateTime, + nullable=False, + server_default=text("CURRENT_TIMESTAMP"), + ) + answered_at: Mapped[datetime | None] = mapped_column(DateTime) + + __all__ = [ "AgentActionLog", + "AgentInboxMessage", "AgentMcpServer", "AgentMemory", "AgentPlan", diff --git a/app/src/fileflash/repositories/__init__.py b/app/src/fileflash/repositories/__init__.py index 84b1407..c063ea1 100644 --- a/app/src/fileflash/repositories/__init__.py +++ b/app/src/fileflash/repositories/__init__.py @@ -1,5 +1,6 @@ from .agent import ( AgentActionLogRepository, + AgentInboxMessageRepository, AgentMcpCatalogEntry, AgentMcpRepository, AgentMemoryActiveEntry, @@ -13,6 +14,7 @@ __all__ = [ "AgentActionLogRepository", + "AgentInboxMessageRepository", "AgentMcpCatalogEntry", "AgentMcpRepository", "AgentMemoryActiveEntry", diff --git a/app/src/fileflash/repositories/agent/__init__.py b/app/src/fileflash/repositories/agent/__init__.py index ac9837d..c551d06 100644 --- a/app/src/fileflash/repositories/agent/__init__.py +++ b/app/src/fileflash/repositories/agent/__init__.py @@ -1,5 +1,6 @@ from .action_log import AgentActionLogRepository from .contracts import AgentMcpCatalogEntry, AgentMemoryActiveEntry, AgentSkillCatalogEntry +from .inbox import AgentInboxMessageRepository from .mcp import AgentMcpRepository from .memory import AgentMemoryRepository from .plan import AgentPlanRepository @@ -9,6 +10,7 @@ __all__ = [ "AgentActionLogRepository", + "AgentInboxMessageRepository", "AgentMcpCatalogEntry", "AgentMcpRepository", "AgentMemoryActiveEntry", diff --git a/app/src/fileflash/repositories/agent/action_log.py b/app/src/fileflash/repositories/agent/action_log.py index 6bcf707..1cd3c08 100644 --- a/app/src/fileflash/repositories/agent/action_log.py +++ b/app/src/fileflash/repositories/agent/action_log.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from typing import Any +from fastapi.encoders import jsonable_encoder from sqlalchemy import select, text from sqlalchemy.ext.asyncio import AsyncSession @@ -27,7 +28,7 @@ async def append_step( job_id=job_id, step_no=step_no, tool_name=tool_name, - inputs_json=inputs_json or {}, + inputs_json=jsonable_encoder(inputs_json or {}), status=status, started_at=started_at or datetime.now(UTC), ) @@ -56,7 +57,7 @@ async def finish_step( if entry is None: return None - entry.outputs_json = outputs_json or {} + entry.outputs_json = jsonable_encoder(outputs_json or {}) entry.status = status entry.duration_ms = duration_ms entry.error_message = error_message diff --git a/app/src/fileflash/repositories/agent/inbox.py b/app/src/fileflash/repositories/agent/inbox.py new file mode 100644 index 0000000..d6996b2 --- /dev/null +++ b/app/src/fileflash/repositories/agent/inbox.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import and_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import AgentInboxMessage +from ...models.enums import AgentInboxKind, AgentInboxRole, AgentInboxStatus + +_CONTROL_KINDS = frozenset( + { + AgentInboxKind.CONTROL_PAUSE, + AgentInboxKind.CONTROL_RESUME, + AgentInboxKind.CONTROL_APPROVE, + AgentInboxKind.CONTROL_DENY, + AgentInboxKind.CONTROL_SKIP, + AgentInboxKind.CONTROL_CANCEL, + } +) + + +class AgentInboxMessageRepository: + def __init__(self, db: AsyncSession) -> None: + self._db = db + + async def create_ask( + self, + *, + job_id: int, + payload: dict[str, Any], + ) -> AgentInboxMessage: + msg = AgentInboxMessage( + job_id=job_id, + role=AgentInboxRole.AGENT, + kind=AgentInboxKind.ASK, + payload_json=payload, + status=AgentInboxStatus.WAITING, + created_at=datetime.now(UTC), + ) + self._db.add(msg) + await self._db.flush() + return msg + + async def record_user_message( + self, + *, + job_id: int, + kind: AgentInboxKind, + payload: dict[str, Any], + reply_to_id: int | None = None, + ) -> AgentInboxMessage: + msg = AgentInboxMessage( + job_id=job_id, + role=AgentInboxRole.USER, + kind=kind, + payload_json=payload, + reply_to_id=reply_to_id, + status=None, + created_at=datetime.now(UTC), + ) + self._db.add(msg) + await self._db.flush() + return msg + + async def mark_answered( + self, + *, + inbox_message_id: int, + answered_at: datetime, + ) -> AgentInboxMessage: + msg = await self._db.get(AgentInboxMessage, inbox_message_id) + if msg is None: + raise ValueError(f"AgentInboxMessage {inbox_message_id} not found") + msg.status = AgentInboxStatus.ANSWERED + msg.answered_at = answered_at + await self._db.flush() + return msg + + async def mark_timed_out( + self, + *, + inbox_message_id: int, + answered_at: datetime, + ) -> AgentInboxMessage: + msg = await self._db.get(AgentInboxMessage, inbox_message_id) + if msg is None: + raise ValueError(f"AgentInboxMessage {inbox_message_id} not found") + msg.status = AgentInboxStatus.TIMED_OUT + msg.answered_at = answered_at + await self._db.flush() + return msg + + async def mark_dropped(self, *, inbox_message_id: int) -> None: + msg = await self._db.get(AgentInboxMessage, inbox_message_id) + if msg is None: + return + if msg.kind in _CONTROL_KINDS: + msg.status = AgentInboxStatus.DROPPED + msg.answered_at = datetime.now(UTC) + await self._db.flush() + + async def get_ask(self, *, inbox_message_id: int) -> AgentInboxMessage | None: + msg = await self._db.get(AgentInboxMessage, inbox_message_id) + if msg is None or msg.kind != AgentInboxKind.ASK: + return None + return msg + + async def get_reply_for(self, *, ask_id: int) -> AgentInboxMessage | None: + return await self._db.scalar( + select(AgentInboxMessage).where( + and_( + AgentInboxMessage.reply_to_id == ask_id, + AgentInboxMessage.kind == AgentInboxKind.REPLY, + ) + ) + ) + + async def list_pending_controls(self, *, job_id: int) -> list[AgentInboxMessage]: + rows = await self._db.scalars( + select(AgentInboxMessage) + .where( + and_( + AgentInboxMessage.job_id == job_id, + AgentInboxMessage.role == AgentInboxRole.USER, + AgentInboxMessage.kind.in_(list(_CONTROL_KINDS)), + AgentInboxMessage.status.is_(None), + ) + ) + .order_by(AgentInboxMessage.created_at.asc()) + ) + return list(rows) + + +__all__ = ["AgentInboxMessageRepository"] diff --git a/app/src/fileflash/repositories/agent/mcp.py b/app/src/fileflash/repositories/agent/mcp.py index cb3dc86..fba358a 100644 --- a/app/src/fileflash/repositories/agent/mcp.py +++ b/app/src/fileflash/repositories/agent/mcp.py @@ -35,7 +35,7 @@ async def list_visible(self, *, user_id: int | None, enabled_only: bool = False) FROM v_agent_mcp_catalog WHERE ( visibility = 'system' - OR (:user_id IS NOT NULL AND owner_user_id = :user_id) + OR owner_user_id = CAST(:user_id AS BIGINT) ) AND (:enabled_only = FALSE OR enabled = TRUE) ORDER BY CASE WHEN visibility = 'system' THEN 0 ELSE 1 END, created_at DESC diff --git a/app/src/fileflash/repositories/agent/skill.py b/app/src/fileflash/repositories/agent/skill.py index 11bc593..e7ec0f3 100644 --- a/app/src/fileflash/repositories/agent/skill.py +++ b/app/src/fileflash/repositories/agent/skill.py @@ -34,7 +34,7 @@ async def list_visible(self, *, user_id: int | None, limit: int = 50) -> list[Ag search_text FROM v_agent_skill_catalog WHERE visibility = 'global' - OR (:user_id IS NOT NULL AND owner_user_id = :user_id) + OR owner_user_id = CAST(:user_id AS BIGINT) ORDER BY CASE WHEN visibility = 'global' THEN 0 ELSE 1 END, created_at DESC LIMIT :limit """ @@ -61,7 +61,7 @@ async def search_visible(self, *, user_id: int | None, query_text: str, limit: i updated_at, search_text FROM v_agent_skill_catalog - WHERE (visibility = 'global' OR (:user_id IS NOT NULL AND owner_user_id = :user_id)) + WHERE (visibility = 'global' OR owner_user_id = CAST(:user_id AS BIGINT)) AND ( :query_text = '' OR search_text ILIKE '%' || :query_text || '%' @@ -117,8 +117,8 @@ async def list_catalog_paginated( where_clause = """ ( (:visibility = 'global' AND visibility = 'global') - OR (:visibility = 'private' AND :user_id IS NOT NULL AND visibility = 'private' AND owner_user_id = :user_id) - OR (:visibility = 'all' AND (visibility = 'global' OR (:user_id IS NOT NULL AND owner_user_id = :user_id))) + OR (:visibility = 'private' AND visibility = 'private' AND owner_user_id = CAST(:user_id AS BIGINT)) + OR (:visibility = 'all' AND (visibility = 'global' OR owner_user_id = CAST(:user_id AS BIGINT))) ) AND ( :query_text = '' diff --git a/app/src/fileflash/routers/__init__.py b/app/src/fileflash/routers/__init__.py index 2f14ff0..d73b994 100644 --- a/app/src/fileflash/routers/__init__.py +++ b/app/src/fileflash/routers/__init__.py @@ -1,6 +1,13 @@ from fastapi import APIRouter from .auth import router as auth_router +from .admin_users import router as admin_users_router +from .admin_storage import router as admin_storage_router +from .admin_files import router as admin_files_router +from .admin_moderation import router as admin_moderation_router +from .admin_logs import router as admin_logs_router +from .admin_notifications import router as admin_notifications_router +from .admin_system import router as admin_system_router from .admin_registration_email_domain_rules import router as admin_registration_email_domain_rules_router from .files import router as files_router from .folders import router as folders_router @@ -10,10 +17,18 @@ from .shares import router as shares_router from .storage import router as storage_router from .uploads import router as uploads_router +from .agent import router as agent_router from .agent_skills import router as agent_skills_router api_router = APIRouter() api_router.include_router(auth_router) +api_router.include_router(admin_users_router) +api_router.include_router(admin_storage_router) +api_router.include_router(admin_files_router) +api_router.include_router(admin_moderation_router) +api_router.include_router(admin_logs_router) +api_router.include_router(admin_notifications_router) +api_router.include_router(admin_system_router) api_router.include_router(admin_registration_email_domain_rules_router) api_router.include_router(files_router) api_router.include_router(folders_router) @@ -23,6 +38,7 @@ api_router.include_router(shares_router) api_router.include_router(storage_router) api_router.include_router(uploads_router) +api_router.include_router(agent_router) api_router.include_router(agent_skills_router) __all__ = ["api_router"] diff --git a/app/src/fileflash/routers/admin_files.py b/app/src/fileflash/routers/admin_files.py new file mode 100644 index 0000000..024a680 --- /dev/null +++ b/app/src/fileflash/routers/admin_files.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from urllib.parse import urlencode + +from fastapi import APIRouter, Depends, Header, Query, Request +from fastapi.responses import StreamingResponse +from jwt import InvalidTokenError + +from ..core.deps import get_admin_files_service, get_settings_dep, require_admin +from ..core.errors import ApiError, api_success +from ..core.security import create_admin_file_preview_token, decode_admin_file_preview_token +from ..core.settings import Settings +from ..models.tables_identity import User +from ..schemas.admin.files import ListAdminFilesQuery +from ..schemas.file import FilePreviewUrlResponse +from ..services.admin.files import AdminFilesService + +router = APIRouter(prefix="/admin/files", tags=["admin"]) + + +@router.get("") +async def list_admin_files( + query: ListAdminFilesQuery = Depends(), + _: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + data = await service.list_files(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Files fetched") + + +@router.get("/{file_id}") +async def get_admin_file_detail( + file_id: int, + _: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + result = await service.get_file_detail(file_id=file_id) + return api_success(data=result.model_dump(by_alias=True), message="File audit detail fetched") + + +@router.get("/{file_id}/preview") +async def preview_admin_file( + file_id: int, + range_header: str | None = Header(default=None, alias="Range"), + _: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + result = await service.get_preview_stream(file_id=file_id, range_header=range_header) + return StreamingResponse( + result.stream, + media_type=result.content_type, + headers=result.headers, + status_code=result.status_code, + ) + + +@router.post("/{file_id}/preview-url") +async def create_admin_file_preview_url( + file_id: int, + request: Request, + admin: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), + settings: Settings = Depends(get_settings_dep), +): + await service.get_file_detail(file_id=file_id) + expires_at = datetime.now(UTC) + timedelta(seconds=settings.file_preview_url_ttl_seconds) + token = create_admin_file_preview_token( + admin_user_id=int(admin.user_id), + file_id=file_id, + settings=settings, + expires_at=expires_at, + ) + stream_url = str(request.url_for("admin_preview_file_stream", file_id=str(file_id))) + result = FilePreviewUrlResponse( + url=f"{stream_url}?{urlencode({'token': token})}", + expires_at=expires_at, + ) + return api_success(data=result.model_dump(by_alias=True), message="Preview URL created") + + +@router.get("/{file_id}/preview-stream", name="admin_preview_file_stream") +async def preview_admin_file_stream( + file_id: int, + token: str = Query(..., min_length=1), + range_header: str | None = Header(default=None, alias="Range"), + service: AdminFilesService = Depends(get_admin_files_service), + settings: Settings = Depends(get_settings_dep), +): + try: + payload = decode_admin_file_preview_token(token, settings) + token_file_id = str(payload["fileId"]) + except (InvalidTokenError, KeyError, ValueError): + raise ApiError(status_code=401, code=401, message="Invalid or expired preview token") from None + + if token_file_id != str(file_id): + raise ApiError(status_code=403, code=403, message="Preview token does not match file") + + result = await service.get_preview_stream(file_id=file_id, range_header=range_header) + return StreamingResponse( + result.stream, + media_type=result.content_type, + headers=result.headers, + status_code=result.status_code, + ) + + +@router.post("/{file_id}/rescan") +async def rescan_admin_file( + file_id: int, + admin: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + result = await service.request_rescan(file_id=file_id, requested_by=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Rescan requested") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_logs.py b/app/src/fileflash/routers/admin_logs.py new file mode 100644 index 0000000..70274ae --- /dev/null +++ b/app/src/fileflash/routers/admin_logs.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_logs_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..schemas.admin.logs import ListAdminLogsQuery +from ..services.admin.logs import AdminLogsService + +router = APIRouter(prefix="/admin/logs", tags=["admin"]) + + +@router.get("") +async def list_admin_logs( + query: ListAdminLogsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminLogsService = Depends(get_admin_logs_service), +): + data = await service.list_logs(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Logs fetched") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_moderation.py b/app/src/fileflash/routers/admin_moderation.py new file mode 100644 index 0000000..aa61604 --- /dev/null +++ b/app/src/fileflash/routers/admin_moderation.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_moderation_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..schemas.admin.moderation import ListViolationsQuery +from ..services.admin.moderation import AdminModerationService + +router = APIRouter(prefix="/admin/violations", tags=["admin"]) + + +@router.get("") +async def list_violations( + query: ListViolationsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminModerationService = Depends(get_admin_moderation_service), +): + data = await service.list_violations(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Violations fetched") + + +@router.post("/{case_id}/resolve") +async def resolve_violation( + case_id: int, + admin: User = Depends(require_admin), + service: AdminModerationService = Depends(get_admin_moderation_service), +): + result = await service.resolve_case(case_id=case_id, handled_by=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Violation resolved") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_notifications.py b/app/src/fileflash/routers/admin_notifications.py new file mode 100644 index 0000000..019d616 --- /dev/null +++ b/app/src/fileflash/routers/admin_notifications.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_notifications_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..schemas.admin.notifications import BroadcastRequest, ListAdminNotificationsQuery +from ..services.admin.notifications import AdminNotificationsService + +router = APIRouter(prefix="/admin/notifications", tags=["admin"]) + + +@router.get("") +async def list_admin_notifications( + query: ListAdminNotificationsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + data = await service.list_notifications(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Notifications fetched") + + +@router.post("/broadcast") +async def broadcast_notification( + payload: BroadcastRequest, + admin: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + result = await service.broadcast(payload=payload, sender_id=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Broadcast sent") + + +@router.delete("/{notification_id}") +async def archive_admin_notification( + notification_id: int, + _: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + await service.archive(notification_id=notification_id) + return api_success( + data={"notificationId": str(notification_id), "status": "archived"}, + message="Notification archived", + ) + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_storage.py b/app/src/fileflash/routers/admin_storage.py new file mode 100644 index 0000000..9b7d6ef --- /dev/null +++ b/app/src/fileflash/routers/admin_storage.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_storage_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..schemas.admin.storage import ListStorageUsersQuery, UpdateQuotaRequest, UsageTrendQuery +from ..services.admin.storage import AdminStorageService + +router = APIRouter(prefix="/admin/storage", tags=["admin"]) + + +@router.get("/summary") +async def get_admin_storage_summary( + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + data = await service.summary() + return api_success(data=data.model_dump(by_alias=True), message="Storage summary fetched") + + +@router.get("/users") +async def list_admin_storage_users( + query: ListStorageUsersQuery = Depends(), + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + data = await service.list_storage_users(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Storage users fetched") + + +@router.patch("/users/{user_id}/quota") +async def update_admin_user_quota( + user_id: int, + payload: UpdateQuotaRequest, + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + result = await service.update_quota(user_id=user_id, new_limit=payload.storage_limit) + return api_success(data=result.model_dump(by_alias=True), message="Quota updated") + + +@router.get("/usage-trend") +async def get_storage_usage_trend( + query: UsageTrendQuery = Depends(), + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + result = await service.usage_trend(query=query) + return api_success(data=result.model_dump(by_alias=True), message="Usage trend fetched") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_system.py b/app/src/fileflash/routers/admin_system.py new file mode 100644 index 0000000..f8ef13a --- /dev/null +++ b/app/src/fileflash/routers/admin_system.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_system_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..services.admin.system import AdminSystemService + +router = APIRouter(prefix="/admin/system", tags=["admin"]) + + +@router.get("/health") +async def get_system_health( + _: User = Depends(require_admin), + service: AdminSystemService = Depends(get_admin_system_service), +): + data = await service.health() + return api_success(data=data.model_dump(by_alias=True), message="System health fetched") + + +@router.get("/rate-limit") +async def get_rate_limit_status( + _: User = Depends(require_admin), + service: AdminSystemService = Depends(get_admin_system_service), +): + data = await service.rate_limit_status() + return api_success(data=data.model_dump(by_alias=True), message="Rate limit fetched") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/admin_users.py b/app/src/fileflash/routers/admin_users.py new file mode 100644 index 0000000..0eb8619 --- /dev/null +++ b/app/src/fileflash/routers/admin_users.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from fastapi import APIRouter, Depends + +from ..core.deps import get_admin_users_service, require_admin +from ..core.errors import api_success +from ..models.tables_identity import User +from ..schemas.admin.users import ListAdminUsersQuery, UpdateUserStatusRequest +from ..services.admin.users import AdminUsersService + +router = APIRouter(prefix="/admin/users", tags=["admin"]) + + +def get_list_admin_users_query(query: ListAdminUsersQuery = Depends()) -> ListAdminUsersQuery: + try: + query.resolve_usage_window() + except ValueError as exc: + from ..core.errors import ApiError + + raise ApiError(status_code=400, code=400, message=str(exc)) from exc + return query + + +@router.get("") +async def list_admin_users( + query: ListAdminUsersQuery = Depends(get_list_admin_users_query), + _: User = Depends(require_admin), + service: AdminUsersService = Depends(get_admin_users_service), +): + data = await service.list_users(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Users fetched") + + +@router.patch("/{user_id}/status") +async def update_admin_user_status( + user_id: int, + payload: UpdateUserStatusRequest, + _: User = Depends(require_admin), + service: AdminUsersService = Depends(get_admin_users_service), +): + result = await service.set_status(user_id=user_id, external_status=payload.status) + return api_success(data=result.model_dump(by_alias=True), message="Status updated") + + +__all__ = ["router"] diff --git a/app/src/fileflash/routers/agent.py b/app/src/fileflash/routers/agent.py index 46216dc..4cc869f 100644 --- a/app/src/fileflash/routers/agent.py +++ b/app/src/fileflash/routers/agent.py @@ -1,9 +1,385 @@ from __future__ import annotations -from fastapi import APIRouter +import json +from typing import Annotated, Any, get_args + +from fastapi import APIRouter, Depends +from fastapi.responses import StreamingResponse +from sqlalchemy import and_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ..agents.harness.event_bus import AgentEventBus, AgentEventEnvelope +from ..agents.harness.inbox import AgentInbox +from ..core.deps import ( + get_agent_event_bus, + get_agent_execute_service, + get_agent_plan_service, + get_current_user, +) +from ..core.errors import ApiError, api_success +from ..db.deps import get_db +from ..models import AgentActionLog, BackgroundJob +from ..models.enums import AgentInboxKind +from ..models.tables_identity import User +from ..schemas.agent import ( + AgentInboxMessageRequest, + AgentInboxMessageResponse, + AgentJobEvent, + AgentJobEventType, + ExecuteAgentRequest, + PlanAgentRequest, +) +from ..services.agent import ExecuteService, PlanService -# The agent router is intentionally scaffold-only for this stage. -# Do not include it in api_router until runtime/services are implemented. router = APIRouter(prefix="/agent", tags=["agent"]) + +@router.post("/plan") +async def plan_agent_task( + payload: PlanAgentRequest, + current_user: Annotated[User, Depends(get_current_user)], + plan_service: Annotated[PlanService, Depends(get_agent_plan_service)], +): + data = await plan_service.enqueue_plan(user_id=current_user.user_id, payload=payload) + return api_success( + data=data.model_dump(by_alias=True), + message="Plan job created", + ) + + +@router.post("/execute") +async def execute_agent_plan( + payload: ExecuteAgentRequest, + current_user: Annotated[User, Depends(get_current_user)], + execute_service: Annotated[ExecuteService, Depends(get_agent_execute_service)], +): + data = await execute_service.enqueue_execute(user_id=current_user.user_id, payload=payload) + return api_success( + data=data.model_dump(by_alias=True), + message="Execute job created", + ) + + +@router.get("/jobs/{job_id}/events") +async def stream_agent_job_events( + job_id: str, + current_user: Annotated[User, Depends(get_current_user)], + db: Annotated[AsyncSession, Depends(get_db)], + event_bus: Annotated[AgentEventBus, Depends(get_agent_event_bus)], +): + parsed_job_id = _parse_job_id(job_id) + initial_events, initial_terminal = await _agent_job_events_for_job( + db=db, + job_id=parsed_job_id, + user_id=int(current_user.user_id), + ) + + async def event_stream(): + seen: set[str] = set() + for event in initial_events: + seen.add(event.id) + yield _format_sse_event(event) + if initial_terminal: + return + async with event_bus.subscribe(job_id=parsed_job_id) as stream: + while True: + try: + envelope = await stream.next(timeout=30.0) + except TimeoutError: + yield ": keep-alive\n\n" + continue + event = _envelope_to_job_event(envelope) + if event is None: + continue + if event.id in seen: + continue + seen.add(event.id) + yield _format_sse_event(event) + if event.type in {"job.succeeded", "job.failed", "job.canceled"}: + break + + return StreamingResponse( + event_stream(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + +@router.post("/jobs/{job_id}/messages") +async def post_agent_job_message( + job_id: str, + payload: AgentInboxMessageRequest, + current_user: Annotated[User, Depends(get_current_user)], + db: Annotated[AsyncSession, Depends(get_db)], + event_bus: Annotated[AgentEventBus, Depends(get_agent_event_bus)], +): + parsed_job_id = _parse_job_id(job_id) + job = await db.scalar( + select(BackgroundJob) + .where( + and_( + BackgroundJob.job_id == parsed_job_id, + BackgroundJob.requested_by == current_user.user_id, + BackgroundJob.task_type.in_(["agent.plan", "agent.execute"]), + ) + ) + ) + if job is None: + raise ApiError(status_code=404, code=404, message="Job not found") + + kind = AgentInboxKind(payload.kind) + reply_to_id: int | None = None + if payload.reply_to is not None: + try: + reply_to_id = int(payload.reply_to) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid replyTo") from exc + + inbox = AgentInbox(db=db, event_bus=event_bus) + try: + msg = await inbox.handle( + job_id=parsed_job_id, + kind=kind, + payload=_inbox_payload_from_request(payload), + reply_to_id=reply_to_id, + ) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message=str(exc)) from exc + await db.commit() + + data = AgentInboxMessageResponse( + inbox_message_id=str(msg.inbox_message_id), + kind=payload.kind, + accepted_at=msg.created_at, + ) + return api_success(data=data.model_dump(by_alias=True), message="Message accepted") + + +def _inbox_payload_from_request(req: AgentInboxMessageRequest) -> dict[str, Any]: + body: dict[str, Any] = {} + if req.value is not None: + body["value"] = req.value + if req.metadata: + body["metadata"] = req.metadata + return body + + +def _parse_job_id(raw: str) -> int: + try: + parsed_job_id = int(raw) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid jobId") from exc + if parsed_job_id <= 0: + raise ApiError(status_code=400, code=400, message="Invalid jobId") + return parsed_job_id + + +async def _agent_job_events_for_job( + *, + db: AsyncSession, + job_id: int, + user_id: int, +) -> tuple[list[AgentJobEvent], bool]: + job = await db.scalar( + select(BackgroundJob).where( + and_( + BackgroundJob.job_id == job_id, + BackgroundJob.requested_by == user_id, + BackgroundJob.task_type.in_(["agent.plan", "agent.execute"]), + ) + ) + ) + if job is None: + raise ApiError(status_code=404, code=404, message="Job not found") + + terminal = str(job.status) in {"succeeded", "failed", "canceled"} + events: list[AgentJobEvent] = [] + if job.task_type == "agent.plan" and job.status == "succeeded" and job.result: + events.append(_plan_ready_event(job)) + events.append(_job_status_event(job)) + elif job.task_type != "agent.execute" or not terminal: + events.append(_job_status_event(job)) + + if job.task_type == "agent.execute": + action_logs = list( + await db.scalars( + select(AgentActionLog) + .where(AgentActionLog.job_id == job_id) + .order_by(AgentActionLog.step_no.asc()) + ) + ) + for action_log in action_logs: + events.extend(_tool_events(job=job, action_log=action_log)) + if terminal: + events.append(_job_status_event(job)) + + return events, terminal + + +def _job_status_event(job: BackgroundJob) -> AgentJobEvent: + status = str(job.status) + event_type = { + "pending": "job.queued", + "running": "job.running", + "succeeded": "job.succeeded", + "failed": "job.failed", + "canceled": "job.canceled", + }.get(status, "job.running") + timestamp = job.updated_at or job.created_at + return AgentJobEvent( + id=f"{job.job_id}:job:{status}:{timestamp.isoformat()}", + job_id=str(job.job_id), + task_type=str(job.task_type), + type=event_type, # type: ignore[arg-type] + status=status, + agent_phase=job.agent_phase, + message=_job_status_message(job), + data=_job_status_data(job), + timestamp=timestamp, + ) + + +def _plan_ready_event(job: BackgroundJob) -> AgentJobEvent: + timestamp = job.finished_at or job.updated_at or job.created_at + return AgentJobEvent( + id=f"{job.job_id}:plan-ready", + job_id=str(job.job_id), + task_type=str(job.task_type), + type="plan.ready", + status=str(job.status), + agent_phase=job.agent_phase, + message="计划已生成。", + data={"result": dict(job.result or {})}, + timestamp=timestamp, + ) + + +def _tool_events(*, job: BackgroundJob, action_log: AgentActionLog) -> list[AgentJobEvent]: + events = [ + AgentJobEvent( + id=f"{job.job_id}:tool:{action_log.action_log_id}:started", + job_id=str(job.job_id), + task_type=str(job.task_type), + type="tool.started", + status=str(job.status), + agent_phase=job.agent_phase, + message=_tool_started_message(action_log), + data=_tool_event_data(action_log, include_output=False), + timestamp=action_log.started_at, + ) + ] + if action_log.status in {"succeeded", "failed"} and action_log.finished_at is not None: + events.append( + AgentJobEvent( + id=f"{job.job_id}:tool:{action_log.action_log_id}:{action_log.status}", + job_id=str(job.job_id), + task_type=str(job.task_type), + type="tool.succeeded" if action_log.status == "succeeded" else "tool.failed", + status=str(job.status), + agent_phase=job.agent_phase, + message=_tool_finished_message(action_log), + data=_tool_event_data(action_log, include_output=True), + timestamp=action_log.finished_at, + ) + ) + return events + + +def _job_status_message(job: BackgroundJob) -> str: + status = str(job.status) + if status == "pending": + return "任务已排队。" + if status == "running": + return "正在规划任务。" if job.task_type == "agent.plan" else "正在执行计划。" + if status == "succeeded": + result = dict(job.result or {}) + answer = result.get("answer") + if isinstance(answer, str) and answer.strip(): + return "答案已生成。" + return "任务已完成。" + if status == "failed": + return str(job.error_message or "任务失败。") + if status == "canceled": + return "任务已取消。" + return "任务状态已更新。" + + +def _job_status_data(job: BackgroundJob) -> dict[str, object]: + data: dict[str, object] = {} + if job.status in {"succeeded", "failed", "canceled"}: + data["result"] = dict(job.result or {}) + if job.error_message: + data["errorMessage"] = job.error_message + return data + + +def _tool_started_message(action_log: AgentActionLog) -> str: + if action_log.tool_name == "drive.countFiles": + inputs = dict(action_log.inputs_json or {}) + search = str(inputs.get("search") or "").strip() + category = str(inputs.get("category") or "").strip() + target = "视频文件" if category == "video" else "文件" + if search: + return f"正在读取名称包含“{search}”的{target}数量。" + return f"正在读取{target}数量。" + return f"正在调用 {action_log.tool_name}。" + + +def _tool_finished_message(action_log: AgentActionLog) -> str: + if action_log.status == "failed": + return str(action_log.error_message or f"{action_log.tool_name} 调用失败。") + if action_log.tool_name == "drive.countFiles": + outputs = dict(action_log.outputs_json or {}) + total_items = int(outputs.get("totalItems") or 0) + return f"读取完成，匹配 {total_items} 个文件。" + return f"{action_log.tool_name} 已完成。" + + +def _tool_event_data(action_log: AgentActionLog, *, include_output: bool) -> dict[str, object]: + data: dict[str, object] = { + "step": int(action_log.step_no), + "tool": str(action_log.tool_name), + "input": dict(action_log.inputs_json or {}), + } + if include_output: + data["output"] = dict(action_log.outputs_json or {}) + if action_log.duration_ms is not None: + data["durationMs"] = int(action_log.duration_ms) + if action_log.error_message: + data["errorMessage"] = action_log.error_message + return data + + +def _envelope_to_job_event(env: AgentEventEnvelope) -> AgentJobEvent | None: + if env.event_type.startswith("agent.inbox."): + return None + if env.event_type not in get_args(AgentJobEventType): + return None + payload = dict(env.payload or {}) + data = payload.get("data") + return AgentJobEvent( + id=env.event_id or f"{env.job_id}:{env.event_type}:{env.emitted_at.isoformat()}", + job_id=str(env.job_id), + task_type=str(payload.get("taskType") or "agent.execute"), + type=env.event_type, # type: ignore[arg-type] + status=str(payload.get("status") or "running"), + agent_phase=payload.get("agentPhase"), + message=str(payload.get("message") or ""), + data=dict(data) if isinstance(data, dict) else payload, + timestamp=env.emitted_at, + ) + + +def _format_sse_event(event: AgentJobEvent) -> str: + payload = event.model_dump(by_alias=True, mode="json") + return ( + f"id: {event.id}\n" + f"event: {event.type}\n" + f"data: {json.dumps(payload, ensure_ascii=False, separators=(',', ':'))}\n\n" + ) + + __all__ = ["router"] diff --git a/app/src/fileflash/routers/files.py b/app/src/fileflash/routers/files.py index b5fe019..8938b4d 100644 --- a/app/src/fileflash/routers/files.py +++ b/app/src/fileflash/routers/files.py @@ -1,18 +1,30 @@ from __future__ import annotations import os +from datetime import UTC, datetime, timedelta +from urllib.parse import urlencode -from fastapi import APIRouter, Depends, Header, Query +from fastapi import APIRouter, Depends, Header, Query, Request from fastapi.responses import FileResponse, StreamingResponse +from jwt import InvalidTokenError from starlette.background import BackgroundTask -from ..core.deps import get_archive_service, get_current_user, get_file_service +from ..core.deps import ( + get_archive_service, + get_current_user, + get_download_rate_limit_service, + get_file_service, + get_settings_dep, +) from ..core.errors import ApiError, api_success +from ..core.security import create_file_preview_token, decode_file_preview_token +from ..core.settings import Settings from ..models.tables_identity import User from ..schemas.archive import ArchiveExtractRequest from ..schemas.file import ( BatchDownloadRequest, BatchFilesRequest, + FilePreviewUrlResponse, GetFilesQuery, MoveFileRequest, RenameFileRequest, @@ -20,11 +32,19 @@ ) from ..schemas.job import to_background_job_response from ..services.archive import ArchiveService +from ..services.download_rate_limit import DownloadRateLimitService from ..services.file import FileService router = APIRouter(prefix="/files", tags=["files"]) +def _content_length(headers: dict[str, str]) -> int: + try: + return max(0, int(headers.get("Content-Length") or 0)) + except ValueError: + return 0 + + @router.get("") async def list_files( folder_id: str | None = Query(None, alias="folderId"), @@ -106,11 +126,14 @@ async def batch_download_files( payload: BatchDownloadRequest, current_user: User = Depends(get_current_user), file_service: FileService = Depends(get_file_service), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), ): - archive_path, archive_name = await file_service.create_batch_download_archive( + plan = await file_service.create_batch_download_plan( user_id=current_user.user_id, payload=payload, ) + await download_limiter.enforce_user(user=current_user, bytes_count=plan.estimated_bytes) + archive_path, archive_name = await file_service.create_batch_download_archive_from_plan(plan=plan) return FileResponse( archive_path, media_type="application/zip", @@ -140,12 +163,14 @@ async def download_file( range_header: str | None = Header(default=None, alias="Range"), current_user: User = Depends(get_current_user), file_service: FileService = Depends(get_file_service), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), ): result = await file_service.get_download_stream( user_id=current_user.user_id, file_id=file_id, range_header=range_header, ) + await download_limiter.enforce_user(user=current_user, bytes_count=_content_length(result.headers)) return StreamingResponse( result.stream, media_type=result.content_type, @@ -160,12 +185,76 @@ async def preview_file( range_header: str | None = Header(default=None, alias="Range"), current_user: User = Depends(get_current_user), file_service: FileService = Depends(get_file_service), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), ): result = await file_service.get_preview_stream( user_id=current_user.user_id, file_id=file_id, range_header=range_header, ) + await download_limiter.enforce_user(user=current_user, bytes_count=_content_length(result.headers)) + return StreamingResponse( + result.stream, + media_type=result.content_type, + headers=result.headers, + status_code=result.status_code, + ) + + +@router.post("/{file_id}/preview-url") +async def create_file_preview_url( + file_id: str, + request: Request, + current_user: User = Depends(get_current_user), + file_service: FileService = Depends(get_file_service), + settings: Settings = Depends(get_settings_dep), +): + try: + fid = int(file_id) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid fileId") from exc + + await file_service.get_file(user_id=current_user.user_id, file_id=fid) + expires_at = datetime.now(UTC) + timedelta(seconds=settings.file_preview_url_ttl_seconds) + token = create_file_preview_token( + user_id=int(current_user.user_id), + file_id=fid, + settings=settings, + expires_at=expires_at, + ) + stream_url = str(request.url_for("preview_file_stream", file_id=str(fid))) + result = FilePreviewUrlResponse( + url=f"{stream_url}?{urlencode({'token': token})}", + expires_at=expires_at, + ) + return api_success(data=result.model_dump(by_alias=True)) + + +@router.get("/{file_id}/preview-stream", name="preview_file_stream") +async def preview_file_stream( + file_id: str, + token: str = Query(..., min_length=1), + range_header: str | None = Header(default=None, alias="Range"), + file_service: FileService = Depends(get_file_service), + settings: Settings = Depends(get_settings_dep), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), +): + try: + payload = decode_file_preview_token(token, settings) + user_id = int(payload["sub"]) + token_file_id = str(payload["fileId"]) + except (InvalidTokenError, KeyError, ValueError): + raise ApiError(status_code=401, code=401, message="Invalid or expired preview token") from None + + if token_file_id != str(file_id): + raise ApiError(status_code=403, code=403, message="Preview token does not match file") + + result = await file_service.get_preview_stream( + user_id=user_id, + file_id=file_id, + range_header=range_header, + ) + await download_limiter.enforce_user_id(user_id=user_id, bytes_count=_content_length(result.headers)) return StreamingResponse( result.stream, media_type=result.content_type, diff --git a/app/src/fileflash/routers/shares.py b/app/src/fileflash/routers/shares.py index 4341854..4ac18c0 100644 --- a/app/src/fileflash/routers/shares.py +++ b/app/src/fileflash/routers/shares.py @@ -3,7 +3,14 @@ from fastapi import APIRouter, Depends, Header from fastapi.responses import StreamingResponse -from ..core.deps import get_client_ip, get_share_service, get_user_agent, get_current_user, require_verified_user +from ..core.deps import ( + get_client_ip, + get_current_user, + get_download_rate_limit_service, + get_share_service, + get_user_agent, + require_verified_user, +) from ..core.errors import api_success from ..core.http_headers import build_content_disposition from ..models.tables_identity import User @@ -14,6 +21,7 @@ SaveShareRequest, UpdateShareSettingsRequest, ) +from ..services.download_rate_limit import DownloadRateLimitService from ..services.share import ShareService router = APIRouter(prefix="/shares", tags=["shares"]) @@ -85,6 +93,19 @@ def _sanitize_stream_headers( return sanitized +def _content_length(headers: dict[str, str] | None) -> int: + if not headers: + return 0 + for key, value in headers.items(): + if key.lower() != "content-length": + continue + try: + return max(0, int(value)) + except ValueError: + return 0 + return 0 + + @router.post("") async def create_share( payload: CreateShareRequest, @@ -195,6 +216,7 @@ async def download_shared_file( client_ip: str = Depends(get_client_ip), user_agent: str | None = Depends(get_user_agent), share_service: ShareService = Depends(get_share_service), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), ): token = _extract_bearer_token(authorization) if not token: @@ -210,6 +232,10 @@ async def download_shared_file( range_header=range_header, ip_address=client_ip, user_agent=user_agent, + rate_limit_check=lambda bytes_count: download_limiter.enforce_share_ip( + client_ip=client_ip, + bytes_count=bytes_count, + ), ) else: raw = await share_service.get_shared_file_stream( @@ -220,6 +246,8 @@ async def download_shared_file( user_agent=user_agent, ) stream, filename, content_type, status_code, headers = _extract_share_stream(tuple(raw)) + if not hasattr(share_service, "get_shared_file_download_stream_response"): + await download_limiter.enforce_share_ip(client_ip=client_ip, bytes_count=_content_length(headers)) response_headers = _sanitize_stream_headers(headers=headers, filename=filename, disposition="attachment") return StreamingResponse(stream, media_type=content_type, headers=response_headers, status_code=status_code) @@ -232,6 +260,7 @@ async def preview_shared_file( client_ip: str = Depends(get_client_ip), user_agent: str | None = Depends(get_user_agent), share_service: ShareService = Depends(get_share_service), + download_limiter: DownloadRateLimitService = Depends(get_download_rate_limit_service), ): token = _extract_bearer_token(authorization) if not token: @@ -247,6 +276,10 @@ async def preview_shared_file( range_header=range_header, ip_address=client_ip, user_agent=user_agent, + rate_limit_check=lambda bytes_count: download_limiter.enforce_share_ip( + client_ip=client_ip, + bytes_count=bytes_count, + ), ) else: raw = await share_service.get_shared_file_stream( @@ -257,6 +290,8 @@ async def preview_shared_file( user_agent=user_agent, ) stream, filename, content_type, status_code, headers = _extract_share_stream(tuple(raw)) + if not hasattr(share_service, "get_shared_file_download_stream_response"): + await download_limiter.enforce_share_ip(client_ip=client_ip, bytes_count=_content_length(headers)) response_headers = _sanitize_stream_headers(headers=headers, filename=filename, disposition="inline") return StreamingResponse(stream, media_type=content_type, headers=response_headers, status_code=status_code) diff --git a/app/src/fileflash/routers/uploads.py b/app/src/fileflash/routers/uploads.py index 9996c87..10a6abb 100644 --- a/app/src/fileflash/routers/uploads.py +++ b/app/src/fileflash/routers/uploads.py @@ -22,6 +22,18 @@ async def preflight_upload( return api_success(data=response.model_dump(by_alias=True), message="Ready for upload") +@router.get("/recoverable") +async def recoverable_uploads( + current_user: User = Depends(get_current_user), + upload_service: UploadService = Depends(get_upload_service), +): + sessions = await upload_service.list_recoverable_sessions(user_id=current_user.user_id) + return api_success( + data=[session.model_dump(by_alias=True) for session in sessions], + message="Recoverable upload sessions fetched", + ) + + @router.post("/{upload_id}/chunk") async def upload_chunk( upload_id: str, @@ -59,3 +71,19 @@ async def merge_chunks( code=201, status_code=201, ) + + +@router.post("/{upload_id}/cancel") +async def cancel_upload( + upload_id: str, + current_user: User = Depends(get_current_user), + upload_service: UploadService = Depends(get_upload_service), +): + response = await upload_service.cancel_upload_session( + user_id=current_user.user_id, + upload_id=upload_id, + ) + return api_success( + data=response.model_dump(by_alias=True), + message="Upload session canceled", + ) diff --git a/app/src/fileflash/schemas/__init__.py b/app/src/fileflash/schemas/__init__.py index 39e1ac4..1e577d9 100644 --- a/app/src/fileflash/schemas/__init__.py +++ b/app/src/fileflash/schemas/__init__.py @@ -1,12 +1,32 @@ -from .auth import ( - ForgotPasswordRequest, - ForgotPasswordResponse, - LoginRequest, - RegisterResponseData, - RegisterRequest, - ResetPasswordRequest, - TokenResponse, - VerifyEmailRequest, +from .admin.files import ( + AdminFileAuditDetail, + AdminFileAuditItem, + AdminFileAuditOwner, + AdminFileLatestScan, + ListAdminFilesQuery, + RescanResponse, +) +from .agent import ( + AgentApproval, + AgentChosenSkill, + AgentCostEstimate, + AgentDataPolicy, + AgentExecutionResult, + AgentHints, + AgentInboxMessageKind, + AgentInboxMessageRequest, + AgentInboxMessageResponse, + AgentJobEvent, + AgentPlanContext, + AgentPlanningEvidence, + AgentPlanResult, + AgentProposedAction, + AgentReasoningEffort, + CancelAgentResponse, + ExecuteAgentRequest, + ExecuteAgentResponse, + PlanAgentRequest, + PlanAgentResponse, ) from .agent_skill import ( AgentSkillItem, @@ -18,13 +38,22 @@ ListAgentSkillsQuery, UpdateAgentSkillRequest, ) +from .auth import ( + ForgotPasswordRequest, + ForgotPasswordResponse, + LoginRequest, + RegisterRequest, + RegisterResponseData, + ResetPasswordRequest, + TokenResponse, + VerifyEmailRequest, +) from .common import ApiResponse, CamelModel, PageQuery, PaginatedData, PaginationMeta from .file import ( - AdminFileAuditItem, BatchDownloadRequest, - BatchMoveItemResult, BatchFilesRequest, BatchFilesResponse, + BatchMoveItemResult, ContentItem, CopyFileRequest, CopyFileResponse, @@ -73,12 +102,6 @@ PermissionItem, UpdatePermissionRequest, ) -from .registration_email_domain_rule import ( - CreateRegistrationEmailDomainRuleRequest, - ListRegistrationEmailDomainRulesQuery, - RegistrationEmailDomainRuleItem, - UpdateRegistrationEmailDomainRuleRequest, -) from .recycle import ( ClearRecycleBinResponse, GetRecycleBinQuery, @@ -87,11 +110,17 @@ RestoreRecycleItemRequest, RestoreRecycleItemResponse, ) +from .registration_email_domain_rule import ( + CreateRegistrationEmailDomainRuleRequest, + ListRegistrationEmailDomainRulesQuery, + RegistrationEmailDomainRuleItem, + UpdateRegistrationEmailDomainRuleRequest, +) from .share import ( + AcceptSharedItemResponse, AccessShareRequest, AccessShareResponseData, AccessUrls, - AcceptSharedItemResponse, CreateShareRequest, DeleteShareResponse, GetSharedItemsQuery, @@ -148,9 +177,27 @@ "AccessUrls", "AcceptSharedItemResponse", "ActivityItem", + "AgentApproval", + "AgentChosenSkill", + "AgentCostEstimate", + "AgentDataPolicy", + "AgentExecutionResult", + "AgentHints", + "AgentInboxMessageKind", + "AgentInboxMessageRequest", + "AgentInboxMessageResponse", + "AgentJobEvent", + "AgentPlanContext", + "AgentPlanningEvidence", + "AgentPlanResult", + "AgentProposedAction", + "AgentReasoningEffort", "AddGroupMemberRequest", "AddGroupMemberResponse", + "AdminFileAuditDetail", "AdminFileAuditItem", + "AdminFileAuditOwner", + "AdminFileLatestScan", "ApiResponse", "BatchFilesRequest", "BatchFilesResponse", @@ -159,6 +206,7 @@ "BroadcastNotificationRequest", "BreakdownDetail", "CamelModel", + "CancelAgentResponse", "ChangePasswordRequest", "ClearRecycleBinResponse", "ContentItem", @@ -174,6 +222,8 @@ "DeleteNotificationResponse", "DeletePermissionResponse", "DeleteShareResponse", + "ExecuteAgentRequest", + "ExecuteAgentResponse", "FileDetails", "FileItem", "FilterSummary", @@ -200,6 +250,7 @@ "LoginRequest", "LogItem", "LogsList", + "ListAdminFilesQuery", "MarkAllAsReadResponse", "MarkAsReadResponse", "MergeChunksRequest", @@ -214,6 +265,8 @@ "PaginatedData", "PaginationMeta", "PermissionItem", + "PlanAgentRequest", + "PlanAgentResponse", "PermanentDeleteResponse", "CreateRegistrationEmailDomainRuleRequest", "RateLimitRule", @@ -226,6 +279,7 @@ "RenameFileRequest", "RenameFolderRequest", "RescanAdminFileResponse", + "RescanResponse", "ResolveViolationResponse", "ResetPasswordRequest", "RestoreRecycleItemRequest", diff --git a/app/src/fileflash/schemas/admin/__init__.py b/app/src/fileflash/schemas/admin/__init__.py new file mode 100644 index 0000000..33e52be --- /dev/null +++ b/app/src/fileflash/schemas/admin/__init__.py @@ -0,0 +1 @@ +"""Admin-only schemas. All requests/responses share CamelModel from schemas.common.""" diff --git a/app/src/fileflash/schemas/admin/files.py b/app/src/fileflash/schemas/admin/files.py new file mode 100644 index 0000000..d96dd74 --- /dev/null +++ b/app/src/fileflash/schemas/admin/files.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from ..common import CamelModel, PageQuery + +VirusStatus = Literal["clean", "pending", "flagged"] + + +class AdminFileLatestScan(CamelModel): + scan_type: str + scan_result: str + virus_status: VirusStatus + scanned_at: datetime + details: dict[str, Any] | None = None + + +class AdminFileAuditItem(CamelModel): + id: str + object_id: str + name: str + size: int + mime_type: str + hash: str + virus_status: VirusStatus + is_shared: bool + owner_name: str + upload_count: int + owner_count: int + scanned_at: datetime | None = None + updated_at: datetime + created_at: datetime + + +class AdminFileAuditOwner(CamelModel): + user_id: str + username: str + email: str + file_count: int + first_uploaded_at: datetime + last_uploaded_at: datetime + + +class AdminFileAuditDetail(AdminFileAuditItem): + object_hash: str | None = None + hash_algorithm: str + storage_status: str + latest_scan: AdminFileLatestScan | None = None + owners: list[AdminFileAuditOwner] + + +class ListAdminFilesQuery(PageQuery): + search: str | None = None + virus_status: VirusStatus | None = None + owner_id: str | None = None + mime_type: str | None = None + sort: Literal["name", "size", "createdAt", "updatedAt"] = "updatedAt" + order: Literal["asc", "desc"] = "desc" + + +class RescanResponse(CamelModel): + file_id: str + virus_status: VirusStatus + scanned_at: datetime + + +__all__ = [ + "AdminFileAuditDetail", + "AdminFileAuditItem", + "AdminFileAuditOwner", + "AdminFileLatestScan", + "ListAdminFilesQuery", + "RescanResponse", + "VirusStatus", +] diff --git a/app/src/fileflash/schemas/admin/logs.py b/app/src/fileflash/schemas/admin/logs.py new file mode 100644 index 0000000..8139d46 --- /dev/null +++ b/app/src/fileflash/schemas/admin/logs.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + + +class LogItem(CamelModel): + id: str + user_id: str | None + operation: str + operation_name: str + target_type: str | None + target_id: str | None + result: str + ip_address: str | None + user_agent: str | None + performed_at: datetime + details: str | None = None + metadata: dict[str, Any] = Field(default_factory=dict) + + +class ListAdminLogsQuery(PageQuery): + user_id: str | None = None + operation: str | None = None + result: Literal["success", "failure"] | None = None + from_at: datetime | None = None + to_at: datetime | None = None + + +class AdminLogsResponse(CamelModel): + logs: list[LogItem] + pagination: dict[str, Any] + + +__all__ = ["AdminLogsResponse", "ListAdminLogsQuery", "LogItem"] diff --git a/app/src/fileflash/schemas/admin/moderation.py b/app/src/fileflash/schemas/admin/moderation.py new file mode 100644 index 0000000..e6384a6 --- /dev/null +++ b/app/src/fileflash/schemas/admin/moderation.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from ..common import CamelModel, PageQuery + +ViolationLevel = Literal["low", "medium", "high"] +ViolationStatus = Literal["pending", "under_review", "resolved"] + + +class ViolationItem(CamelModel): + id: str + file_id: str | None + file_name: str | None + type: str + level: ViolationLevel + reported_at: datetime + status: ViolationStatus + + +class ListViolationsQuery(PageQuery): + status: ViolationStatus | None = None + + +class ResolveViolationResponse(CamelModel): + violation_id: str + resolved_at: datetime + + +__all__ = [ + "ListViolationsQuery", + "ResolveViolationResponse", + "ViolationItem", + "ViolationLevel", + "ViolationStatus", +] diff --git a/app/src/fileflash/schemas/admin/notifications.py b/app/src/fileflash/schemas/admin/notifications.py new file mode 100644 index 0000000..8ee7ecf --- /dev/null +++ b/app/src/fileflash/schemas/admin/notifications.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + + +class AdminNotificationItem(CamelModel): + id: str + message: str + title: str | None + type: str + status: str + is_read: bool + created_at: datetime + updated_at: datetime + recipient_count: int | None = None + + +class ListAdminNotificationsQuery(PageQuery): + status: str | None = None + type: str | None = None + + +class BroadcastRequest(CamelModel): + title: str | None = Field(default=None, max_length=255) + message: str = Field(min_length=1, max_length=2000) + type: Literal["system", "announcement"] = "system" + + +class BroadcastResponse(CamelModel): + broadcast_id: str + recipient_count: int + sent_at: datetime + + +__all__ = [ + "AdminNotificationItem", + "BroadcastRequest", + "BroadcastResponse", + "ListAdminNotificationsQuery", +] diff --git a/app/src/fileflash/schemas/admin/storage.py b/app/src/fileflash/schemas/admin/storage.py new file mode 100644 index 0000000..13e718e --- /dev/null +++ b/app/src/fileflash/schemas/admin/storage.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + + +class AdminStorageSummary(CamelModel): + storage_used: int + storage_limit: int + storage_percentage: float + file_count: int + user_count: int + updated_at: datetime + + +class AdminStorageUserItem(CamelModel): + user_id: str + username: str + email: str + storage_limit: int + storage_used: int + usage_percentage: float + updated_at: datetime + + +class ListStorageUsersQuery(PageQuery): + sort: Literal["storageUsed", "usagePercentage", "username"] = "storageUsed" + order: Literal["asc", "desc"] = "desc" + + +class UpdateQuotaRequest(CamelModel): + storage_limit: int = Field(ge=0) + + +class UpdateQuotaResponse(CamelModel): + user_id: str + storage_limit: int + storage_used: int + usage_percentage: float + updated_at: datetime + + +class UsageTrendQuery(CamelModel): + days: Literal[7, 14, 30] = 7 + + +class UsageTrendPoint(CamelModel): + date: str + used: int + + +class UsageTrendResponse(CamelModel): + trends: list[UsageTrendPoint] + is_estimated: bool = False + + +__all__ = [ + "AdminStorageSummary", + "AdminStorageUserItem", + "ListStorageUsersQuery", + "UpdateQuotaRequest", + "UpdateQuotaResponse", + "UsageTrendQuery", + "UsageTrendPoint", + "UsageTrendResponse", +] diff --git a/app/src/fileflash/schemas/admin/system.py b/app/src/fileflash/schemas/admin/system.py new file mode 100644 index 0000000..bc09025 --- /dev/null +++ b/app/src/fileflash/schemas/admin/system.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from datetime import datetime + +from ..common import CamelModel + + +class SystemHealth(CamelModel): + platform_targets: list[str] + max_concurrent_uploads: int + active_upload_sessions: int + virus_scan_enabled: bool + thumbnail_generation_enabled: bool + registration_mail_enabled: bool + hash_computation_enabled: bool + last_updated_at: datetime + + +class RateLimitRule(CamelModel): + rule_id: str + scope: str + window_seconds: int + limit: int + current_usage: int + blocked_requests: int + + +class RateLimitStatus(CamelModel): + rules: list[RateLimitRule] + evaluated_at: datetime + + +__all__ = ["RateLimitRule", "RateLimitStatus", "SystemHealth"] diff --git a/app/src/fileflash/schemas/admin/users.py b/app/src/fileflash/schemas/admin/users.py new file mode 100644 index 0000000..9eab779 --- /dev/null +++ b/app/src/fileflash/schemas/admin/users.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + +ExternalUserStatus = Literal["active", "suspended", "pending_verification"] + +DEFAULT_USAGE_WINDOW = timedelta(days=7) +MAX_USAGE_WINDOW = timedelta(days=90) + + +def _normalize_datetime(value: datetime) -> datetime: + if value.tzinfo is None: + return value.replace(tzinfo=UTC) + return value.astimezone(UTC) + + +class AdminUserUsageStats(CamelModel): + traffic_bytes: int = Field(ge=0) + agent_tokens: int = Field(ge=0) + + +class AdminUserItem(CamelModel): + user_id: str + username: str + email: str + role: str + status: ExternalUserStatus + email_verified: bool + email_verified_at: datetime | None = None + storage_limit: int + storage_used: int + usage_percentage: float + last_login_at: datetime | None = None + last_active_at: datetime | None = None + created_at: datetime + usage_stats: AdminUserUsageStats + + +class ListAdminUsersQuery(PageQuery): + search: str | None = None + status: Literal["active", "suspended"] | None = None + role: Literal["USER", "ADMIN"] | None = None + sort: Literal["username", "createdAt", "storageUsed"] = "createdAt" + order: Literal["asc", "desc"] = "desc" + usage_from: datetime | None = None + usage_to: datetime | None = None + + def resolve_usage_window(self, *, now: datetime | None = None) -> tuple[datetime, datetime]: + resolved_now = _normalize_datetime(now or datetime.now(UTC)) + if self.usage_from is None and self.usage_to is None: + return resolved_now - DEFAULT_USAGE_WINDOW, resolved_now + if self.usage_from is None or self.usage_to is None: + raise ValueError("usageFrom and usageTo must be provided together") + + usage_from = _normalize_datetime(self.usage_from) + usage_to = _normalize_datetime(self.usage_to) + if usage_from > usage_to: + raise ValueError("usageFrom must be earlier than or equal to usageTo") + if usage_to - usage_from > MAX_USAGE_WINDOW: + raise ValueError("usage window must not exceed 90 days") + return usage_from, usage_to + + +class UpdateUserStatusRequest(CamelModel): + status: Literal["active", "suspended"] + + +class UpdateUserStatusResponse(CamelModel): + user_id: str + status: ExternalUserStatus + updated_at: datetime + + +__all__ = [ + "AdminUserItem", + "AdminUserUsageStats", + "ListAdminUsersQuery", + "UpdateUserStatusRequest", + "UpdateUserStatusResponse", +] diff --git a/app/src/fileflash/schemas/agent.py b/app/src/fileflash/schemas/agent.py new file mode 100644 index 0000000..47c273d --- /dev/null +++ b/app/src/fileflash/schemas/agent.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from pydantic import Field + +from .common import CamelModel + +AgentExecutionPolicy = Literal["planOnly", "confirm", "autopilot"] +AgentActionSideEffect = Literal["read", "write"] +AgentRiskLevel = Literal["low", "medium", "high"] +AgentReasoningEffort = Literal["adaptive", "low", "medium", "high", "xhigh", "max"] +AgentJobPhase = Literal[ + "planning", + "awaiting_confirm", + "executing", + "awaiting_commit", + "completed", + "failed", + "canceled", +] +AgentJobEventType = Literal[ + "job.queued", + "job.running", + "plan.ready", + "tool.started", + "tool.succeeded", + "tool.failed", + "tool.partial", + "agent.thinking", + "agent.progress", + "agent.ask", + "agent.paused", + "agent.resumed", + "job.succeeded", + "job.failed", + "job.canceled", +] + + +class AgentDataPolicy(CamelModel): + allow_file_content: bool = False + max_read_bytes: int = Field(default=1_048_576, ge=0) + allowed_mime_types: list[str] = Field(default_factory=lambda: ["*/*"]) + + +class AgentHints(CamelModel): + prefer_skill_id: str | None = None + max_steps: int = Field(default=12, ge=1) + budget_tokens: int = Field(default=8_000, ge=1) + reasoning_effort: AgentReasoningEffort = "adaptive" + + +class AgentPlanContext(CamelModel): + root_folder_id: str = "root" + selected_file_ids: list[str] = Field(default_factory=list) + selected_folder_ids: list[str] = Field(default_factory=list) + current_path: str = "/My Files" + + +class PlanAgentRequest(CamelModel): + input: str = Field(min_length=1, max_length=4_000) + context: AgentPlanContext + execution_policy: AgentExecutionPolicy = "confirm" + data_policy: AgentDataPolicy = Field(default_factory=AgentDataPolicy) + hints: AgentHints = Field(default_factory=AgentHints) + + +class PlanAgentResponse(CamelModel): + job_id: str + status: str + task_type: Literal["agent.plan"] = "agent.plan" + + +class AgentProposedAction(CamelModel): + step: int = Field(ge=1) + tool: str = Field(min_length=1, max_length=120) + input: dict[str, Any] = Field(default_factory=dict) + side_effect: AgentActionSideEffect + risk_level: AgentRiskLevel = "low" + requires_confirmation: bool = False + confirmation_reason: str | None = None + + +class AgentCostEstimate(CamelModel): + tokens: int = Field(ge=0) + tool_calls: int = Field(ge=0) + duration_sec_estimate: int = Field(ge=0) + + +class AgentChosenSkill(CamelModel): + id: str + name: str + + +class AgentPlanningEvidence(CamelModel): + step: int = Field(ge=1) + tool: str = Field(min_length=1, max_length=120) + input: dict[str, Any] = Field(default_factory=dict) + output_preview: dict[str, Any] = Field(default_factory=dict) + + +class AgentPlanResult(CamelModel): + plan_job_id: str + plan_hash: str + chosen_skill: AgentChosenSkill | None = None + proposed_actions: list[AgentProposedAction] + summary: str + requires_confirmation: bool + cost_estimate: AgentCostEstimate + planning_evidence: list[AgentPlanningEvidence] | None = None + + +class AgentApproval(CamelModel): + confirmed_by: str + confirmed_at: datetime + high_risk_confirmed: bool = False + high_risk_confirmed_at: datetime | None = None + + +class ExecuteAgentRequest(CamelModel): + plan_job_id: str + plan_hash: str + approval: AgentApproval + + +class ExecuteAgentResponse(CamelModel): + job_id: str + status: str + task_type: Literal["agent.execute"] = "agent.execute" + + +class CancelAgentResponse(CamelModel): + job_id: str + status: str + canceled_at: datetime + + +class AgentExecutionResult(CamelModel): + plan_job_id: str + execute_job_id: str + summary: str + answer: str | None = None + applied_actions: int = Field(ge=0) + skipped_actions: int = Field(ge=0) + warnings: list[str] = Field(default_factory=list) + finished_at: datetime + + +class AgentJobEvent(CamelModel): + id: str + job_id: str + task_type: str + type: AgentJobEventType + status: str + agent_phase: str | None = None + message: str + data: dict[str, Any] = Field(default_factory=dict) + timestamp: datetime + + +AgentInboxMessageKind = Literal[ + "reply", + "control.pause", + "control.resume", + "control.approve", + "control.deny", + "control.skip", + "control.cancel", +] + + +class AgentInboxMessageRequest(CamelModel): + kind: AgentInboxMessageKind + reply_to: str | None = None + value: Any = None + metadata: dict[str, Any] = Field(default_factory=dict) + + +class AgentInboxMessageResponse(CamelModel): + inbox_message_id: str + kind: AgentInboxMessageKind + accepted_at: datetime + + +__all__ = [ + "AgentActionSideEffect", + "AgentApproval", + "AgentChosenSkill", + "AgentCostEstimate", + "AgentDataPolicy", + "AgentExecutionPolicy", + "AgentExecutionResult", + "AgentHints", + "AgentInboxMessageKind", + "AgentInboxMessageRequest", + "AgentInboxMessageResponse", + "AgentJobPhase", + "AgentJobEvent", + "AgentJobEventType", + "AgentPlanContext", + "AgentPlanningEvidence", + "AgentPlanResult", + "AgentProposedAction", + "AgentReasoningEffort", + "AgentRiskLevel", + "CancelAgentResponse", + "ExecuteAgentRequest", + "ExecuteAgentResponse", + "PlanAgentRequest", + "PlanAgentResponse", +] diff --git a/app/src/fileflash/schemas/file.py b/app/src/fileflash/schemas/file.py index d9dd1ab..3827366 100644 --- a/app/src/fileflash/schemas/file.py +++ b/app/src/fileflash/schemas/file.py @@ -128,6 +128,20 @@ class UploadPreflightResponse(CamelModel): uploaded_chunk_indexes: list[int] | None = None +class RecoverableUploadSession(CamelModel): + upload_id: str + file_name: str = Field(min_length=1, max_length=255) + file_size: int = Field(ge=0) + uploaded_bytes: int = Field(ge=0) + chunk_size: int = Field(gt=0) + file_hash: str = Field(min_length=8, max_length=128) + mime_type: str = Field(min_length=1, max_length=255) + parent_id: str + updated_at: datetime + expired_at: datetime | None = None + status: Literal["init", "uploading"] + + class MergeChunksRequest(CamelModel): file_hash: str = Field(min_length=8, max_length=128) file_name: str = Field(min_length=1, max_length=255) @@ -147,6 +161,11 @@ class MergeChunksResponse(CamelModel): download_url: str +class UploadCancelResponse(CamelModel): + upload_id: str + canceled_at: datetime + + class FileDetails(FileItem): status: bool @@ -158,6 +177,11 @@ class MediaOptimization(CamelModel): updated_at: datetime +class FilePreviewUrlResponse(CamelModel): + url: str + expires_at: datetime + + class RenameFileRequest(CamelModel): file_name: str = Field(min_length=1, max_length=255) diff --git a/app/src/fileflash/scripts/init_dev_accounts.py b/app/src/fileflash/scripts/init_dev_accounts.py index 80c5f6d..da67db6 100644 --- a/app/src/fileflash/scripts/init_dev_accounts.py +++ b/app/src/fileflash/scripts/init_dev_accounts.py @@ -11,11 +11,11 @@ def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description="Initialize development test accounts for FileFlash.") + parser = argparse.ArgumentParser(description="Initialize seeded accounts for FileFlash.") parser.add_argument( "--reset-password", action="store_true", - help="Force reset passwords to defaults (admin/admin123, demo/demo123).", + help="Force reset seeded account passwords from the active environment configuration.", ) return parser @@ -23,8 +23,8 @@ def build_parser() -> argparse.ArgumentParser: async def run(reset_password: bool) -> int: settings = get_settings() if settings.is_production_env: - logger.warning( - "Manual dev-account initialization is running under APP_ENV=%s. This is not executed automatically in production.", + logger.info( + "Manual account initialization is using DEFAULT_ADMIN_* for APP_ENV=%s.", settings.app_env, ) diff --git a/app/src/fileflash/scripts/run_with_workers.py b/app/src/fileflash/scripts/run_with_workers.py index 7f2db6d..a6b19a7 100644 --- a/app/src/fileflash/scripts/run_with_workers.py +++ b/app/src/fileflash/scripts/run_with_workers.py @@ -48,6 +48,17 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Start API only (without file workers).", ) + parser.add_argument( + "--no-agent-worker", + action="store_true", + help="Do not start agent workers even when AGENT_ENABLED=true.", + ) + parser.add_argument( + "--agent-worker-count", + type=int, + default=max(1, int(getattr(settings, "agent_worker_concurrency", 1))), + help="Number of agent worker consumer processes when AGENT_ENABLED=true.", + ) return parser @@ -164,6 +175,20 @@ def main() -> int: processes.append(worker_proc) print(f"[run-with-workers] started {worker_name}: {_format_cmd(worker_cmd)}") + settings = get_settings() + should_start_agent_workers = ( + not args.no_worker + and getattr(settings, "agent_enabled", False) + and not args.no_agent_worker + ) + if should_start_agent_workers: + for index in range(max(1, int(args.agent_worker_count))): + worker_name = f"agent-worker-{index + 1}" + worker_cmd = [python, "-m", "fileflash.agents.worker"] + worker_proc = _spawn_process(worker_name, worker_cmd, cwd) + processes.append(worker_proc) + print(f"[run-with-workers] started {worker_name}: {_format_cmd(worker_cmd)}") + while True: for managed in processes: exit_code = managed.process.poll() diff --git a/app/src/fileflash/services/__init__.py b/app/src/fileflash/services/__init__.py index e523bba..551c022 100644 --- a/app/src/fileflash/services/__init__.py +++ b/app/src/fileflash/services/__init__.py @@ -2,6 +2,7 @@ from .agent import ExecuteService, McpService, MemoryService, PlanService, SessionService, SettingsService, SkillService from .auth import AuthService from .background_jobs import BackgroundJobService +from .download_rate_limit import DownloadRateLimitService from .email_delivery import VerificationEmailDeliveryService from .file import FileService from .folder import FolderService @@ -17,6 +18,7 @@ "AuthService", "ArchiveService", "BackgroundJobService", + "DownloadRateLimitService", "VerificationEmailDeliveryService", "ExecuteService", "FileService", diff --git a/app/src/fileflash/services/admin/__init__.py b/app/src/fileflash/services/admin/__init__.py new file mode 100644 index 0000000..0edfeeb --- /dev/null +++ b/app/src/fileflash/services/admin/__init__.py @@ -0,0 +1 @@ +"""Admin-only services. Each module owns one /admin/* surface.""" diff --git a/app/src/fileflash/services/admin/_status.py b/app/src/fileflash/services/admin/_status.py new file mode 100644 index 0000000..d238975 --- /dev/null +++ b/app/src/fileflash/services/admin/_status.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from ...core.errors import ApiError +from ...models.enums import UserStatus + +_EXTERNAL_TO_INTERNAL = { + "active": UserStatus.ACTIVE, + "suspended": UserStatus.DISABLED, +} + +_INTERNAL_TO_EXTERNAL = { + UserStatus.ACTIVE: "active", + UserStatus.DISABLED: "suspended", + UserStatus.LOCKED: "suspended", + UserStatus.PENDING_VERIFICATION: "pending_verification", +} + + +def external_to_internal(value: str) -> UserStatus: + try: + return _EXTERNAL_TO_INTERNAL[value] + except KeyError as exc: + raise ApiError( + status_code=422, + code=422, + message=f"Invalid user status: {value!r}", + ) from exc + + +def internal_to_external(value: UserStatus) -> str: + return _INTERNAL_TO_EXTERNAL.get(value, value.value) + + +__all__ = ["external_to_internal", "internal_to_external"] diff --git a/app/src/fileflash/services/admin/files.py b/app/src/fileflash/services/admin/files.py new file mode 100644 index 0000000..a818da5 --- /dev/null +++ b/app/src/fileflash/services/admin/files.py @@ -0,0 +1,563 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any, Protocol + +from sqlalchemy import and_, func, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...core.http_headers import build_content_disposition +from ...core.mime import DEFAULT_MIME_TYPE, resolve_file_mime_type +from ...models.enums import FileStatus, ScanResult, ShareStatus, UploadStatus +from ...models.tables_access_share import Share +from ...models.tables_audit_security import ObjectScanResult +from ...models.tables_identity import User +from ...models.tables_storage import File, FileMediaMetadata, StorageObject +from ...s3.minio_client import MinioObjectStorageClient +from ...schemas.admin.files import ( + AdminFileAuditDetail, + AdminFileAuditItem, + AdminFileAuditOwner, + AdminFileLatestScan, + ListAdminFilesQuery, + RescanResponse, + VirusStatus, +) +from ...schemas.common import PaginatedData, PaginationMeta + + +class EventPublisherProtocol(Protocol): + async def publish(self, event_name: str, payload: dict[str, Any]) -> None: ... + + +@dataclass(slots=True) +class AdminFileStreamResult: + stream: AsyncIterator[bytes] + filename: str + content_type: str + status_code: int + headers: dict[str, str] + + +@dataclass(slots=True) +class ResolvedAdminStreamObject: + storage_object: StorageObject + content_type_override: str | None = None + + +_VIRUS_STATUS_MAP: dict[ScanResult, VirusStatus] = { + ScanResult.CLEAN: "clean", + ScanResult.PENDING: "pending", + ScanResult.INFECTED: "flagged", + ScanResult.BLOCKED: "flagged", + ScanResult.FAILED: "pending", +} + +TRANSCODE_READY_STATUS = "ready" + + +class AdminFilesService: + def __init__( + self, + db: AsyncSession, + publisher: EventPublisherProtocol, + storage: MinioObjectStorageClient | None = None, + ) -> None: + self.db = db + self.publisher = publisher + self.storage = storage + + async def list_files(self, *, query: ListAdminFilesQuery) -> PaginatedData[AdminFileAuditItem]: + latest_scan = self._latest_scan_subquery() + object_stats = self._object_stats_subquery() + share_stats = self._share_stats_subquery() + + statement = ( + select( + File, + StorageObject, + User, + ObjectScanResult, + object_stats.c.upload_count, + object_stats.c.owner_count, + share_stats.c.share_count, + ) + .join(StorageObject, File.storage_object_id == StorageObject.object_id) + .join(User, File.owner_id == User.user_id) + .join(object_stats, object_stats.c.object_id == StorageObject.object_id, isouter=True) + .join(share_stats, share_stats.c.object_id == StorageObject.object_id, isouter=True) + .join(latest_scan, latest_scan.c.object_id == StorageObject.object_id, isouter=True) + .join( + ObjectScanResult, + and_( + ObjectScanResult.object_id == latest_scan.c.object_id, + ObjectScanResult.scanned_at == latest_scan.c.scanned_at, + ), + isouter=True, + ) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + ) + + if query.search: + kw = f"%{query.search.strip().lower()}%" + statement = statement.where(func.lower(File.file_name).like(kw)) + if query.owner_id: + statement = statement.where(File.owner_id == int(query.owner_id)) + if query.mime_type: + statement = statement.where(File.mime_type == query.mime_type) + if query.virus_status: + wanted = [raw for raw, mapped in _VIRUS_STATUS_MAP.items() if mapped == query.virus_status] + statement = statement.where(ObjectScanResult.result.in_(wanted)) + + total = int(await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0) + + sort_column = { + "name": File.file_name, + "size": File.file_size, + "createdAt": File.created_at, + "updatedAt": File.updated_at, + }[query.sort] + statement = statement.order_by(sort_column.desc() if query.order == "desc" else sort_column.asc()) + + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = (await self.db.execute(statement.offset(offset).limit(query.per_page))).all() + items = [ + self._to_item( + file_row, + object_row, + owner_row, + scan_row, + upload_count=upload_count, + owner_count=owner_count, + share_count=share_count, + ) + for ( + file_row, + object_row, + owner_row, + scan_row, + upload_count, + owner_count, + share_count, + ) in rows + ] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def get_file_detail(self, *, file_id: int) -> AdminFileAuditDetail: + file_row, object_row, owner_row, scan_row = await self._get_active_file_context(file_id=file_id) + owners = await self._load_object_owners(object_id=int(object_row.object_id)) + upload_count = sum(owner.file_count for owner in owners) + owner_count = len(owners) + is_shared = await self._object_is_shared(object_id=int(object_row.object_id)) + + item = self._to_item( + file_row, + object_row, + owner_row, + scan_row, + upload_count=upload_count, + owner_count=owner_count, + share_count=1 if is_shared else 0, + ) + return AdminFileAuditDetail( + **item.model_dump(), + object_hash=object_row.object_hash, + hash_algorithm=object_row.hash_algorithm, + storage_status=self._enum_value(object_row.upload_status), + latest_scan=self._to_latest_scan(scan_row), + owners=owners, + ) + + async def get_preview_stream( + self, + *, + file_id: int, + range_header: str | None, + ) -> AdminFileStreamResult: + if self.storage is None: + raise ApiError(status_code=503, code=503, message="Object storage is unavailable") + + file_row, object_row, _owner_row, _scan_row = await self._get_active_file_context(file_id=file_id) + resolved_object = await self._resolve_stream_storage_object( + file_row=file_row, + source_object=object_row, + prefer_optimized=True, + ) + storage_object = resolved_object.storage_object if resolved_object is not None else None + if storage_object is None or not self._is_upload_status_active(storage_object.upload_status): + raise ApiError(status_code=404, code=404, message="File content not found") + + object_size = int(storage_object.object_size or file_row.file_size or 0) + if object_size <= 0: + raise ApiError(status_code=404, code=404, message="File content not found") + + content_type = resolve_file_mime_type( + mime_type=( + resolved_object.content_type_override + if resolved_object is not None and resolved_object.content_type_override + else file_row.mime_type or storage_object.content_type + ), + file_ext=file_row.file_ext, + file_name=file_row.file_name, + default=DEFAULT_MIME_TYPE, + ) + headers = { + "Accept-Ranges": "bytes", + "Content-Disposition": build_content_disposition(file_row.file_name, disposition="inline"), + } + + byte_range = self._parse_range_header(range_header=range_header, file_size=object_size) + if byte_range is None: + headers["Content-Length"] = str(object_size) + stream = self.storage.iter_object( + bucket_name=storage_object.bucket_name, + object_key=storage_object.object_key, + ) + return AdminFileStreamResult( + stream=stream, + filename=file_row.file_name, + content_type=content_type, + status_code=200, + headers=headers, + ) + + start, end = byte_range + headers["Content-Length"] = str(end - start + 1) + headers["Content-Range"] = f"bytes {start}-{end}/{object_size}" + stream = self.storage.iter_object_range( + bucket_name=storage_object.bucket_name, + object_key=storage_object.object_key, + start=start, + end=end, + ) + return AdminFileStreamResult( + stream=stream, + filename=file_row.file_name, + content_type=content_type, + status_code=206, + headers=headers, + ) + + async def request_rescan(self, *, file_id: int, requested_by: int) -> RescanResponse: + file_row = await self.db.get(File, file_id) + if file_row is None or file_row.deleted_at is not None or file_row.status != FileStatus.ACTIVE: + raise ApiError(status_code=404, code=404, message="File not found") + + now = datetime.now(UTC) + self.db.add( + ObjectScanResult( + object_id=int(file_row.storage_object_id), + scan_type="virus", + result=ScanResult.PENDING, + details={"requestedBy": requested_by}, + scanned_at=now, + created_at=now, + ) + ) + await self.db.commit() + + await self.publisher.publish( + "files.rescan_requested", + { + "fileId": str(file_id), + "objectId": str(file_row.storage_object_id), + "requestedBy": requested_by, + }, + ) + return RescanResponse(file_id=str(file_id), virus_status="pending", scanned_at=now) + + async def _get_active_file_context( + self, + *, + file_id: int, + ) -> tuple[File, StorageObject, User, ObjectScanResult | None]: + latest_scan = self._latest_scan_subquery() + row = ( + await self.db.execute( + select(File, StorageObject, User, ObjectScanResult) + .join(StorageObject, File.storage_object_id == StorageObject.object_id) + .join(User, File.owner_id == User.user_id) + .join(latest_scan, latest_scan.c.object_id == StorageObject.object_id, isouter=True) + .join( + ObjectScanResult, + and_( + ObjectScanResult.object_id == latest_scan.c.object_id, + ObjectScanResult.scanned_at == latest_scan.c.scanned_at, + ), + isouter=True, + ) + .where(File.file_id == file_id) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + ) + ).first() + if row is None: + raise ApiError(status_code=404, code=404, message="File not found") + file_row, object_row, owner_row, scan_row = row + return file_row, object_row, owner_row, scan_row + + async def _load_object_owners(self, *, object_id: int) -> list[AdminFileAuditOwner]: + first_uploaded = func.min(File.created_at) + last_uploaded = func.max(File.created_at) + rows = ( + await self.db.execute( + select( + User.user_id, + User.username, + User.email, + func.count(File.file_id), + first_uploaded, + last_uploaded, + ) + .join(User, User.user_id == File.owner_id) + .where(File.storage_object_id == object_id) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + .group_by(User.user_id, User.username, User.email) + .order_by(last_uploaded.desc()) + ) + ).all() + return [ + AdminFileAuditOwner( + user_id=str(user_id), + username=username, + email=email, + file_count=int(file_count or 0), + first_uploaded_at=first_uploaded_at, + last_uploaded_at=last_uploaded_at, + ) + for user_id, username, email, file_count, first_uploaded_at, last_uploaded_at in rows + ] + + async def _object_is_shared(self, *, object_id: int) -> bool: + now = datetime.now(UTC) + share_id = await self.db.scalar( + select(Share.share_id) + .join(File, File.file_id == Share.file_id) + .where(File.storage_object_id == object_id) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + .where(Share.status == ShareStatus.ACTIVE) + .where(or_(Share.expire_time.is_(None), Share.expire_time > now)) + .limit(1) + ) + return share_id is not None + + async def _resolve_stream_storage_object( + self, + *, + file_row: File, + source_object: StorageObject, + prefer_optimized: bool, + ) -> ResolvedAdminStreamObject | None: + if not prefer_optimized: + return ResolvedAdminStreamObject(storage_object=source_object) + + metadata_row = await self.db.scalar( + select(FileMediaMetadata) + .where(FileMediaMetadata.source_object_id == int(file_row.storage_object_id)) + .limit(1) + ) + if not isinstance(metadata_row, FileMediaMetadata): + return ResolvedAdminStreamObject(storage_object=source_object) + transcode = (metadata_row.extra_metadata or {}).get("transcode") + if not isinstance(transcode, dict): + return ResolvedAdminStreamObject(storage_object=source_object) + if str(transcode.get("status") or "").strip().lower() != TRANSCODE_READY_STATUS: + return ResolvedAdminStreamObject(storage_object=source_object) + + bucket_name = str(transcode.get("optimizedBucketName") or "").strip() + object_key = str(transcode.get("optimizedObjectKey") or "").strip() + if not bucket_name or not object_key: + return ResolvedAdminStreamObject(storage_object=source_object) + + optimized_mime_type = str(transcode.get("optimizedMimeType") or "").strip() or None + optimized_object = await self.db.scalar( + select(StorageObject) + .where( + and_( + StorageObject.bucket_name == bucket_name, + StorageObject.object_key == object_key, + StorageObject.upload_status == UploadStatus.ACTIVE, + ) + ) + .limit(1) + ) + if isinstance(optimized_object, StorageObject): + return ResolvedAdminStreamObject( + storage_object=optimized_object, + content_type_override=optimized_mime_type or optimized_object.content_type, + ) + + if self.storage is None: + return ResolvedAdminStreamObject(storage_object=source_object) + exists = await self.storage.object_exists(bucket_name=bucket_name, object_key=object_key) + if not exists: + return ResolvedAdminStreamObject(storage_object=source_object) + + stat = await self.storage.stat_object(bucket_name=bucket_name, object_key=object_key) + created = StorageObject( + bucket_name=bucket_name, + object_key=object_key, + object_size=int(stat.size), + etag=stat.etag, + version_id=stat.version_id, + content_type=stat.content_type, + upload_status=UploadStatus.ACTIVE, + ) + self.db.add(created) + await self.db.flush() + return ResolvedAdminStreamObject( + storage_object=created, + content_type_override=optimized_mime_type or created.content_type, + ) + + @staticmethod + def _latest_scan_subquery(): + return ( + select(ObjectScanResult.object_id, func.max(ObjectScanResult.scanned_at).label("scanned_at")) + .group_by(ObjectScanResult.object_id) + .subquery() + ) + + @staticmethod + def _object_stats_subquery(): + return ( + select( + File.storage_object_id.label("object_id"), + func.count(File.file_id).label("upload_count"), + func.count(func.distinct(File.owner_id)).label("owner_count"), + ) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + .group_by(File.storage_object_id) + .subquery() + ) + + @staticmethod + def _share_stats_subquery(): + now = datetime.now(UTC) + return ( + select( + File.storage_object_id.label("object_id"), + func.count(Share.share_id).label("share_count"), + ) + .join(Share, Share.file_id == File.file_id) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + .where(Share.status == ShareStatus.ACTIVE) + .where(or_(Share.expire_time.is_(None), Share.expire_time > now)) + .group_by(File.storage_object_id) + .subquery() + ) + + @staticmethod + def _to_item( + file_row: File, + object_row: StorageObject, + owner_row: User, + scan_row: ObjectScanResult | None, + *, + upload_count: int | None, + owner_count: int | None, + share_count: int | None, + ) -> AdminFileAuditItem: + return AdminFileAuditItem( + id=str(file_row.file_id), + object_id=str(object_row.object_id), + name=file_row.file_name, + size=int(file_row.file_size), + mime_type=file_row.mime_type or object_row.content_type or DEFAULT_MIME_TYPE, + hash=(object_row.object_hash or "")[:16], + virus_status=_VIRUS_STATUS_MAP.get(scan_row.result, "pending") if scan_row else "pending", + is_shared=bool(share_count or 0), + owner_name=owner_row.username, + upload_count=int(upload_count or 1), + owner_count=int(owner_count or 1), + scanned_at=scan_row.scanned_at if scan_row else None, + updated_at=file_row.updated_at, + created_at=file_row.created_at, + ) + + @staticmethod + def _to_latest_scan(scan_row: ObjectScanResult | None) -> AdminFileLatestScan | None: + if scan_row is None: + return None + result = AdminFilesService._enum_value(scan_row.result) + return AdminFileLatestScan( + scan_type=scan_row.scan_type, + scan_result=result, + virus_status=_VIRUS_STATUS_MAP.get(scan_row.result, "pending"), + scanned_at=scan_row.scanned_at, + details=scan_row.details, + ) + + @staticmethod + def _enum_value(value: object) -> str: + return str(getattr(value, "value", value)) + + @staticmethod + def _is_upload_status_active(value: object) -> bool: + return value == UploadStatus.ACTIVE or AdminFilesService._enum_value(value) == UploadStatus.ACTIVE.value + + @staticmethod + def _parse_range_header(range_header: str | None, file_size: int) -> tuple[int, int] | None: + if not range_header: + return None + + value = range_header.strip() + if not value.lower().startswith("bytes="): + raise ApiError(status_code=416, code=416, message="Invalid Range header") + + spec = value[6:].strip() + if "," in spec: + raise ApiError(status_code=416, code=416, message="Multiple ranges are not supported") + + if spec.startswith("-"): + suffix_part = spec[1:].strip() + if not suffix_part.isdigit(): + raise ApiError(status_code=416, code=416, message="Invalid Range header") + suffix = int(suffix_part) + if suffix <= 0: + raise ApiError(status_code=416, code=416, message="Invalid Range header") + start = max(file_size - suffix, 0) + end = file_size - 1 + return start, end + + if "-" not in spec: + raise ApiError(status_code=416, code=416, message="Invalid Range header") + + start_part, end_part = spec.split("-", 1) + if not start_part.strip().isdigit(): + raise ApiError(status_code=416, code=416, message="Invalid Range header") + start = int(start_part.strip()) + end = file_size - 1 + if end_part.strip(): + if not end_part.strip().isdigit(): + raise ApiError(status_code=416, code=416, message="Invalid Range header") + end = int(end_part.strip()) + + if start < 0 or start >= file_size or end < start: + raise ApiError(status_code=416, code=416, message="Requested range is not satisfiable") + + if end >= file_size: + end = file_size - 1 + return start, end + + +__all__ = ["AdminFilesService", "AdminFileStreamResult"] diff --git a/app/src/fileflash/services/admin/logs.py b/app/src/fileflash/services/admin/logs.py new file mode 100644 index 0000000..93010e5 --- /dev/null +++ b/app/src/fileflash/services/admin/logs.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +from sqlalchemy import and_, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models.tables_audit_security import Log +from ...schemas.admin.logs import AdminLogsResponse, ListAdminLogsQuery, LogItem + + +class AdminLogsService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_logs(self, *, query: ListAdminLogsQuery) -> AdminLogsResponse: + statement = select(Log) + conditions = [] + if query.user_id: + conditions.append(Log.user_id == int(query.user_id)) + if query.operation: + conditions.append(Log.operation == query.operation) + if query.result: + conditions.append(Log.result == query.result) + if query.from_at: + conditions.append(Log.performed_at >= query.from_at) + if query.to_at: + conditions.append(Log.performed_at <= query.to_at) + if conditions: + statement = statement.where(and_(*conditions)) + statement = statement.order_by(Log.performed_at.desc(), Log.id.desc()) + + total = int(await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + + logs = [ + LogItem( + id=str(row.id), + user_id=str(row.user_id) if row.user_id else None, + operation=row.operation, + operation_name=row.operation, + target_type=row.target_type, + target_id=str(row.target_id) if row.target_id else None, + result=row.result, + ip_address=row.ip_address, + user_agent=row.user_agent, + performed_at=row.performed_at or datetime.now(UTC), + details=row.details, + metadata=row.metadata_payload or {}, + ) + for row in rows + ] + return AdminLogsResponse( + logs=logs, + pagination={ + "totalItems": total, + "totalPages": total_pages, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": query.page > 1, + "hasNext": query.page < total_pages, + }, + ) + + +__all__ = ["AdminLogsService"] diff --git a/app/src/fileflash/services/admin/moderation.py b/app/src/fileflash/services/admin/moderation.py new file mode 100644 index 0000000..b153fd1 --- /dev/null +++ b/app/src/fileflash/services/admin/moderation.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from decimal import Decimal + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.tables_audit_security import ModerationCase +from ...models.tables_storage import File +from ...schemas.admin.moderation import ( + ListViolationsQuery, + ResolveViolationResponse, + ViolationItem, + ViolationLevel, +) +from ...schemas.common import PaginatedData, PaginationMeta + +_OPEN_STATES = ("pending", "under_review") + + +class AdminModerationService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_violations(self, *, query: ListViolationsQuery) -> PaginatedData[ViolationItem]: + statement = select(ModerationCase, File).join(File, File.file_id == ModerationCase.file_id, isouter=True) + if query.status: + statement = statement.where(ModerationCase.status == query.status) + statement = statement.order_by(ModerationCase.created_at.desc()) + + total = int(await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = (await self.db.execute(statement.offset(offset).limit(query.per_page))).all() + items = [self._to_item(case_row, file_row) for case_row, file_row in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def resolve_case(self, *, case_id: int, handled_by: int) -> ResolveViolationResponse: + case_row = await self.db.get(ModerationCase, case_id, with_for_update=True) + if case_row is None: + raise ApiError(status_code=404, code=404, message="Violation case not found") + if case_row.status not in _OPEN_STATES: + raise ApiError(status_code=409, code=409, message="Case already resolved") + + now = datetime.now(UTC) + case_row.status = "resolved" + case_row.resolution = "admin_clear" + case_row.handled_by = handled_by + case_row.handled_at = now + case_row.updated_at = now + await self.db.commit() + return ResolveViolationResponse(violation_id=str(case_id), resolved_at=now) + + @staticmethod + def _to_item(case_row: ModerationCase, file_row: File | None) -> ViolationItem: + return ViolationItem( + id=str(case_row.case_id), + file_id=str(case_row.file_id) if case_row.file_id else None, + file_name=file_row.file_name if file_row else None, + type=case_row.reason_type, + level=_level_from_confidence(case_row.confidence), + reported_at=case_row.created_at, + status=case_row.status, # type: ignore[arg-type] + ) + + +def _level_from_confidence(confidence: Decimal | None) -> ViolationLevel: + if confidence is None: + return "low" + value = float(confidence) + if value > 0.8: + return "high" + if value > 0.5: + return "medium" + return "low" + + +__all__ = ["AdminModerationService"] diff --git a/app/src/fileflash/services/admin/notifications.py b/app/src/fileflash/services/admin/notifications.py new file mode 100644 index 0000000..2325140 --- /dev/null +++ b/app/src/fileflash/services/admin/notifications.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from typing import AsyncIterator + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.enums import UserStatus +from ...models.tables_audit_security import Notification +from ...models.tables_identity import User +from ...schemas.admin.notifications import ( + AdminNotificationItem, + BroadcastRequest, + BroadcastResponse, + ListAdminNotificationsQuery, +) +from ...schemas.common import PaginatedData, PaginationMeta + +MAX_BROADCAST_RECIPIENTS = 50_000 +_CHUNK_SIZE = 500 + + +class AdminNotificationsService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_notifications( + self, + *, + query: ListAdminNotificationsQuery, + ) -> PaginatedData[AdminNotificationItem]: + statement = select(Notification) + if query.status: + statement = statement.where(Notification.status == query.status) + if query.type: + statement = statement.where(Notification.notification_type == query.type) + statement = statement.order_by(Notification.created_at.desc(), Notification.id.desc()) + + total = int(await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + items = [self._to_item(row) for row in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def broadcast(self, *, payload: BroadcastRequest, sender_id: int) -> BroadcastResponse: + message = (payload.message or "").strip() + if not message: + raise ApiError(status_code=422, code=422, message="Broadcast message cannot be empty") + + recipient_count = int( + await self.db.scalar( + select(func.count(User.user_id)).where( + User.status == UserStatus.ACTIVE, + User.deleted_at.is_(None), + ) + ) + or 0 + ) + if recipient_count > MAX_BROADCAST_RECIPIENTS: + raise ApiError( + status_code=422, + code=422, + message=f"Recipient count {recipient_count} exceeds limit {MAX_BROADCAST_RECIPIENTS}", + ) + + broadcast_id = str(uuid.uuid4()) + now = datetime.now(UTC) + delivered = 0 + async for chunk in self._iter_active_user_ids(): + rows = [ + Notification( + user_id=user_id, + title=payload.title, + notification_type=payload.type, + channel="in_app", + message=message, + payload={"broadcastId": broadcast_id}, + sender_user_id=sender_id, + status="sent", + sent_at=now, + is_read=False, + created_at=now, + updated_at=now, + ) + for user_id in chunk + ] + if rows: + self.db.add_all(rows) + await self.db.commit() + delivered += len(rows) + + return BroadcastResponse( + broadcast_id=broadcast_id, + recipient_count=delivered, + sent_at=now, + ) + + async def archive(self, *, notification_id: int) -> None: + row = await self.db.get(Notification, notification_id) + if row is None: + raise ApiError(status_code=404, code=404, message="Notification not found") + row.status = "archived" + row.updated_at = datetime.now(UTC) + await self.db.commit() + + async def _iter_active_user_ids(self) -> AsyncIterator[list[int]]: + stream = await self.db.scalars( + select(User.user_id).where(User.status == UserStatus.ACTIVE, User.deleted_at.is_(None)) + ) + buffer: list[int] = [] + for user_id in stream: + buffer.append(int(user_id)) + if len(buffer) >= _CHUNK_SIZE: + yield buffer + buffer = [] + if buffer: + yield buffer + + @staticmethod + def _to_item(row: Notification) -> AdminNotificationItem: + return AdminNotificationItem( + id=str(row.id), + message=row.message, + title=row.title, + type=row.notification_type, + status=row.status, + is_read=bool(row.is_read), + created_at=row.created_at or datetime.now(UTC), + updated_at=row.updated_at, + recipient_count=None, + ) + + +__all__ = ["AdminNotificationsService", "MAX_BROADCAST_RECIPIENTS"] diff --git a/app/src/fileflash/services/admin/storage.py b/app/src/fileflash/services/admin/storage.py new file mode 100644 index 0000000..95d4e75 --- /dev/null +++ b/app/src/fileflash/services/admin/storage.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from datetime import UTC, date, datetime, timedelta + +from redis.asyncio import Redis +from sqlalchemy import Float, and_, case, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.tables_audit_security import Log +from ...models.tables_identity import User +from ...models.tables_storage import File +from ...schemas.admin.storage import ( + AdminStorageSummary, + AdminStorageUserItem, + ListStorageUsersQuery, + UpdateQuotaResponse, + UsageTrendPoint, + UsageTrendQuery, + UsageTrendResponse, +) +from ...schemas.common import PaginatedData, PaginationMeta + +_STORAGE_EVENT_OPS = ("file.created", "file.deleted", "file.restored") +_TREND_CACHE_TTL = 300 + + +class AdminStorageService: + def __init__(self, db: AsyncSession, redis: Redis | None) -> None: + self.db = db + self.redis = redis + + async def summary(self) -> AdminStorageSummary: + used_sum = await self.db.scalar( + select(func.coalesce(func.sum(User.storage_used), 0)).where(User.deleted_at.is_(None)) + ) + limit_sum = await self.db.scalar( + select(func.coalesce(func.sum(User.storage_limit), 0)).where(User.deleted_at.is_(None)) + ) + file_count = await self.db.scalar( + select(func.count(File.file_id)).where(File.deleted_at.is_(None)) + ) + user_count = await self.db.scalar( + select(func.count(User.user_id)).where(User.deleted_at.is_(None)) + ) + + used = int(used_sum or 0) + limit = int(limit_sum or 0) + return AdminStorageSummary( + storage_used=used, + storage_limit=limit, + storage_percentage=round((used / limit) * 100, 2) if limit else 0.0, + file_count=int(file_count or 0), + user_count=int(user_count or 0), + updated_at=datetime.now(UTC), + ) + + async def list_storage_users(self, *, query: ListStorageUsersQuery) -> PaginatedData[AdminStorageUserItem]: + usage_ratio = case( + (User.storage_limit > 0, User.storage_used.cast(Float) / User.storage_limit.cast(Float)), + else_=0.0, + ) + sort_column = { + "storageUsed": User.storage_used, + "usagePercentage": usage_ratio, + "username": User.username, + }[query.sort] + + statement = ( + select(User) + .where(User.deleted_at.is_(None)) + .order_by(sort_column.desc() if query.order == "desc" else sort_column.asc()) + ) + + total = await self.db.scalar(select(func.count()).select_from(statement.subquery())) + total_items = int(total or 0) + total_pages = max(1, -(-total_items // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + items = [self._to_item(row) for row in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total_items, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def update_quota(self, *, user_id: int, new_limit: int) -> UpdateQuotaResponse: + target = await self.db.get(User, user_id, with_for_update=True) + if target is None or target.deleted_at is not None: + raise ApiError(status_code=404, code=404, message="User not found") + if new_limit < int(target.storage_used): + raise ApiError(status_code=409, code=409, message="New quota cannot be below current usage") + + target.storage_limit = int(new_limit) + target.updated_at = datetime.now(UTC) + await self.db.commit() + await self.db.refresh(target) + return UpdateQuotaResponse( + user_id=str(target.user_id), + storage_limit=int(target.storage_limit), + storage_used=int(target.storage_used), + usage_percentage=round((int(target.storage_used) / max(int(target.storage_limit), 1)) * 100, 2), + updated_at=target.updated_at, + ) + + async def usage_trend(self, *, query: UsageTrendQuery) -> UsageTrendResponse: + cached = await self._cache_get(query.days) + if cached is not None: + return cached + + current_total = int( + await self.db.scalar( + select(func.coalesce(func.sum(User.storage_used), 0)).where(User.deleted_at.is_(None)) + ) + or 0 + ) + cutoff = datetime.now(UTC) - timedelta(days=query.days) + rows = await self.db.execute( + select(Log.operation, Log.metadata_payload, Log.performed_at) + .where(Log.operation.in_(_STORAGE_EVENT_OPS)) + .where(Log.performed_at >= cutoff) + ) + events = rows.all() + + deltas: dict[date, int] = {} + for operation, metadata_payload, performed_at in events: + if performed_at is None: + continue + size = int((metadata_payload or {}).get("size") or 0) + sign = 1 if operation in {"file.created", "file.restored"} else -1 + day = performed_at.astimezone(UTC).date() + deltas[day] = deltas.get(day, 0) + sign * size + + today = datetime.now(UTC).date() + points: list[UsageTrendPoint] = [] + running = current_total + for offset in range(query.days): + day = today - timedelta(days=offset) + points.append(UsageTrendPoint(date=day.isoformat(), used=max(running, 0))) + running -= deltas.get(day, 0) + points.reverse() + + result = UsageTrendResponse(trends=points, is_estimated=not events) + await self._cache_set(query.days, result) + return result + + @staticmethod + def _to_item(row: User) -> AdminStorageUserItem: + limit = max(int(row.storage_limit), 1) + return AdminStorageUserItem( + user_id=str(row.user_id), + username=row.username, + email=row.email, + storage_limit=int(row.storage_limit), + storage_used=int(row.storage_used), + usage_percentage=round((int(row.storage_used) / limit) * 100, 2), + updated_at=row.updated_at, + ) + + async def _cache_get(self, days: int) -> UsageTrendResponse | None: + if self.redis is None: + return None + try: + raw = await self.redis.get(f"admin:storage:trend:{days}") + except Exception: + return None + if not raw: + return None + return UsageTrendResponse.model_validate_json(raw) + + async def _cache_set(self, days: int, payload: UsageTrendResponse) -> None: + if self.redis is None: + return + try: + await self.redis.setex( + f"admin:storage:trend:{days}", + _TREND_CACHE_TTL, + payload.model_dump_json(by_alias=True), + ) + except Exception: + return + + +__all__ = ["AdminStorageService"] diff --git a/app/src/fileflash/services/admin/system.py b/app/src/fileflash/services/admin/system.py new file mode 100644 index 0000000..1cbbffe --- /dev/null +++ b/app/src/fileflash/services/admin/system.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.settings import Settings +from ...models.enums import UploadTaskStatus +from ...models.tables_storage import UploadTask +from ...schemas.admin.system import RateLimitRule, RateLimitStatus, SystemHealth + + +class AdminSystemService: + def __init__(self, db: AsyncSession, settings: Settings) -> None: + self.db = db + self.settings = settings + + async def health(self) -> SystemHealth: + active_uploads = int( + await self.db.scalar( + select(func.count(UploadTask.task_id)).where( + UploadTask.status.in_([UploadTaskStatus.INIT, UploadTaskStatus.UPLOADING]) + ) + ) + or 0 + ) + + platform_targets: list[str] = [] + if self.settings.object_storage_bucket: + platform_targets.append(f"s3://{self.settings.object_storage_bucket}") + if self.settings.redis_url: + platform_targets.append("redis") + + return SystemHealth( + platform_targets=platform_targets, + max_concurrent_uploads=getattr(self.settings, "max_concurrent_uploads", 4), + active_upload_sessions=active_uploads, + virus_scan_enabled=bool(getattr(self.settings, "virus_scan_enabled", False)), + thumbnail_generation_enabled=bool(getattr(self.settings, "thumbnail_generation_enabled", True)), + registration_mail_enabled=bool(self.settings.mail_server and self.settings.mail_from), + hash_computation_enabled=bool(self.settings.upload_verify_merged_object_hash), + last_updated_at=datetime.now(UTC), + ) + + async def rate_limit_status(self) -> RateLimitStatus: + rules = [ + RateLimitRule( + rule_id="login", + scope="auth.login", + window_seconds=self.settings.login_rate_window_seconds, + limit=self.settings.login_rate_limit, + current_usage=0, + blocked_requests=0, + ), + RateLimitRule( + rule_id="register", + scope="auth.register", + window_seconds=self.settings.register_rate_window_seconds, + limit=self.settings.register_rate_limit, + current_usage=0, + blocked_requests=0, + ), + RateLimitRule( + rule_id="forgot_password", + scope="auth.forgot_password", + window_seconds=self.settings.forgot_password_rate_window_seconds, + limit=self.settings.forgot_password_rate_limit, + current_usage=0, + blocked_requests=0, + ), + RateLimitRule( + rule_id="resend_verification", + scope="auth.resend_verification", + window_seconds=self.settings.resend_verification_rate_window_seconds, + limit=self.settings.resend_verification_rate_limit, + current_usage=0, + blocked_requests=0, + ), + ] + return RateLimitStatus(rules=rules, evaluated_at=datetime.now(UTC)) + + +__all__ = ["AdminSystemService"] diff --git a/app/src/fileflash/services/admin/users.py b/app/src/fileflash/services/admin/users.py new file mode 100644 index 0000000..7c86a9a --- /dev/null +++ b/app/src/fileflash/services/admin/users.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +from sqlalchemy import and_, func, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.enums import UploadTaskStatus, UserRole, UserStatus +from ...models.tables_identity import User, UserSession +from ...models.tables_storage import UploadTask +from ...models.tables_worker import BackgroundJob +from ...schemas.admin.users import ( + AdminUserItem, + AdminUserUsageStats, + ListAdminUsersQuery, + UpdateUserStatusResponse, +) +from ...schemas.common import PaginatedData, PaginationMeta +from ._status import external_to_internal, internal_to_external + + +class AdminUsersService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_users(self, *, query: ListAdminUsersQuery) -> PaginatedData[AdminUserItem]: + statement = select(User).where(User.deleted_at.is_(None)) + if query.search: + keyword = f"%{query.search.strip().lower()}%" + statement = statement.where( + func.lower(User.username).like(keyword) | func.lower(User.email).like(keyword) + ) + if query.status: + statement = statement.where(User.status == external_to_internal(query.status)) + if query.role: + statement = statement.where(User.role == UserRole(query.role)) + + sort_column = { + "username": User.username, + "createdAt": User.created_at, + "storageUsed": User.storage_used, + }[query.sort] + statement = statement.order_by(sort_column.desc() if query.order == "desc" else sort_column.asc()) + + total = await self.db.scalar(select(func.count()).select_from(statement.subquery())) + total_items = int(total or 0) + total_pages = max(1, -(-total_items // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + + user_ids = [int(row.user_id) for row in rows] + last_seen_map = await self._collect_last_seen(user_ids) + usage_from, usage_to = self._resolve_usage_window(query) + usage_map = await self._collect_usage_stats( + user_ids=user_ids, + usage_from=usage_from, + usage_to=usage_to, + ) + items = [ + self._to_item( + row, + last_seen_map.get(int(row.user_id)), + usage_map.get( + int(row.user_id), + AdminUserUsageStats(traffic_bytes=0, agent_tokens=0), + ), + ) + for row in rows + ] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total_items, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + @staticmethod + def _resolve_usage_window(query: ListAdminUsersQuery) -> tuple[datetime, datetime]: + try: + return query.resolve_usage_window() + except ValueError as exc: + raise ApiError(status_code=400, code=400, message=str(exc)) from exc + + async def _collect_usage_stats( + self, + *, + user_ids: list[int], + usage_from: datetime, + usage_to: datetime, + ) -> dict[int, AdminUserUsageStats]: + if not user_ids: + return {} + + traffic_rows = await self.db.execute( + select(UploadTask.user_id, func.coalesce(func.sum(UploadTask.total_size), 0)) + .where( + and_( + UploadTask.user_id.in_(user_ids), + UploadTask.status == UploadTaskStatus.COMPLETED, + UploadTask.completed_at.is_not(None), + UploadTask.completed_at >= usage_from, + UploadTask.completed_at <= usage_to, + ) + ) + .group_by(UploadTask.user_id) + ) + stats: dict[int, AdminUserUsageStats] = { + int(user_id): AdminUserUsageStats(traffic_bytes=int(total or 0), agent_tokens=0) + for user_id, total in traffic_rows.all() + } + + token_expr = BackgroundJob.result["costEstimate"]["tokens"].as_integer() + agent_rows = await self.db.execute( + select( + BackgroundJob.requested_by, + func.coalesce(func.sum(func.coalesce(token_expr, 0)), 0), + ) + .where( + and_( + BackgroundJob.requested_by.in_(user_ids), + BackgroundJob.task_type == "agent.plan", + BackgroundJob.status == "succeeded", + BackgroundJob.finished_at.is_not(None), + BackgroundJob.finished_at >= usage_from, + BackgroundJob.finished_at <= usage_to, + ) + ) + .group_by(BackgroundJob.requested_by) + ) + for user_id, total in agent_rows.all(): + if user_id is None: + continue + key = int(user_id) + current = stats.get(key, AdminUserUsageStats(traffic_bytes=0, agent_tokens=0)) + stats[key] = AdminUserUsageStats( + traffic_bytes=current.traffic_bytes, + agent_tokens=int(total or 0), + ) + + return stats + + async def set_status(self, *, user_id: int, external_status: str) -> UpdateUserStatusResponse: + target = await self.db.get(User, user_id) + if target is None or target.deleted_at is not None: + raise ApiError(status_code=404, code=404, message="User not found") + + new_internal = external_to_internal(external_status) + if ( + new_internal == UserStatus.DISABLED + and target.role == UserRole.ADMIN + and target.status == UserStatus.ACTIVE + ): + remaining = await self.db.scalar( + select(func.count(User.user_id)).where( + and_( + User.role == UserRole.ADMIN, + User.status == UserStatus.ACTIVE, + User.user_id != user_id, + User.deleted_at.is_(None), + ) + ) + ) + if int(remaining or 0) == 0: + raise ApiError( + status_code=409, + code=409, + message="Cannot suspend the last active admin", + ) + + target.status = new_internal + target.updated_at = datetime.now(UTC) + if new_internal == UserStatus.DISABLED: + now = datetime.now(UTC) + await self.db.execute( + update(UserSession) + .where(and_(UserSession.user_id == user_id, UserSession.revoked_at.is_(None))) + .values(revoked_at=now, last_seen_at=now) + ) + + await self.db.commit() + await self.db.refresh(target) + return UpdateUserStatusResponse( + user_id=str(target.user_id), + status=internal_to_external(target.status), + updated_at=target.updated_at, + ) + + async def _collect_last_seen(self, user_ids: list[int]) -> dict[int, datetime]: + if not user_ids: + return {} + rows = await self.db.execute( + select(UserSession.user_id, func.max(UserSession.last_seen_at)) + .where(and_(UserSession.user_id.in_(user_ids), UserSession.revoked_at.is_(None))) + .group_by(UserSession.user_id) + ) + return {int(user_id): seen for user_id, seen in rows.all()} + + @staticmethod + def _to_item( + row: User, + last_active_at: datetime | None, + usage_stats: AdminUserUsageStats, + ) -> AdminUserItem: + limit = max(int(row.storage_limit), 1) + return AdminUserItem( + user_id=str(row.user_id), + username=row.username, + email=row.email, + role=row.role.value if hasattr(row.role, "value") else str(row.role), + status=internal_to_external(row.status), + email_verified=bool(row.email_verified), + email_verified_at=row.email_verified_at, + storage_limit=int(row.storage_limit), + storage_used=int(row.storage_used), + usage_percentage=round((int(row.storage_used) / limit) * 100, 2), + last_login_at=row.last_login_at, + last_active_at=last_active_at, + created_at=row.created_at, + usage_stats=usage_stats, + ) + + +__all__ = ["AdminUsersService"] diff --git a/app/src/fileflash/services/agent/execute_service.py b/app/src/fileflash/services/agent/execute_service.py index b2031c4..e3f7c41 100644 --- a/app/src/fileflash/services/agent/execute_service.py +++ b/app/src/fileflash/services/agent/execute_service.py @@ -1,8 +1,14 @@ from __future__ import annotations +from sqlalchemy import and_, select from sqlalchemy.ext.asyncio import AsyncSession +from ...agents.harness.policy import classify_tool_risk +from ...core.errors import ApiError +from ...core.settings import Settings +from ...models import BackgroundJob from ...repositories import AgentPlanRepository, AgentWorkSessionRepository +from ...schemas.agent import AgentProposedAction, ExecuteAgentRequest, ExecuteAgentResponse from ..background_jobs import BackgroundJobService @@ -11,14 +17,113 @@ def __init__( self, *, db: AsyncSession, + settings: Settings, jobs: BackgroundJobService, plans: AgentPlanRepository, work_sessions: AgentWorkSessionRepository, ) -> None: self.db = db + self.settings = settings self.jobs = jobs self.plans = plans self.work_sessions = work_sessions - async def enqueue_execute(self, *args, **kwargs): - raise NotImplementedError("Agent execute service is scaffolded only in this stage") + async def enqueue_execute( + self, + *, + user_id: int, + payload: ExecuteAgentRequest, + ) -> ExecuteAgentResponse: + if not self.settings.agent_enabled: + raise ApiError(status_code=503, code=503, message="Agent runtime is disabled") + + plan_job_id = _parse_job_id(payload.plan_job_id) + plan_job = await self.db.scalar( + select(BackgroundJob).where( + and_( + BackgroundJob.job_id == plan_job_id, + BackgroundJob.requested_by == user_id, + BackgroundJob.task_type == "agent.plan", + ) + ) + ) + if plan_job is None: + raise ApiError(status_code=404, code=404, message="Plan job not found") + if plan_job.status != "succeeded": + raise ApiError(status_code=409, code=409, message="Plan job is not ready for execution") + + plan = await self.plans.get_for_execute_binding( + job_id=plan_job_id, + user_id=user_id, + plan_hash=payload.plan_hash, + ) + if plan is None: + raise ApiError(status_code=409, code=409, message="planHash mismatch") + + high_risk_actions = _high_risk_actions(plan.proposed_actions_json or []) + if high_risk_actions and not payload.approval.high_risk_confirmed: + raise ApiError( + status_code=409, + code=409, + message="High-risk action requires confirmation", + data={"highRiskActions": high_risk_actions}, + ) + + idempotency_key = f"agent.execute:{plan_job_id}" + existing_execute = await self.db.scalar( + select(BackgroundJob).where( + and_( + BackgroundJob.task_type == "agent.execute", + BackgroundJob.idempotency_key == idempotency_key, + ) + ) + ) + if existing_execute is not None: + raise ApiError( + status_code=409, + code=409, + message="Plan has already been executed", + data={ + "jobId": str(existing_execute.job_id), + "status": str(existing_execute.status), + }, + ) + + job = await self.jobs.enqueue( + self.db, + task_type="agent.execute", + payload=payload.model_dump(by_alias=True, mode="json"), + idempotency_key=idempotency_key, + requested_by=user_id, + max_attempts=1, + priority=100, + agent_phase="executing", + ) + return ExecuteAgentResponse( + job_id=str(job.job_id), + status=str(job.status), + task_type="agent.execute", + ) + + +def _parse_job_id(raw: str) -> int: + try: + value = int(raw) + except ValueError as exc: + raise ApiError(status_code=400, code=400, message="Invalid planJobId") from exc + if value <= 0: + raise ApiError(status_code=400, code=400, message="Invalid planJobId") + return value + + +def _high_risk_actions(raw_actions: object) -> list[dict[str, object]]: + if not isinstance(raw_actions, list): + return [] + risky: list[dict[str, object]] = [] + for item in raw_actions: + if not isinstance(item, dict): + continue + action = AgentProposedAction.model_validate(item) + if action.risk_level == "high" or classify_tool_risk(action.tool) == "high": + risky.append(action.model_dump(by_alias=True)) + return risky diff --git a/app/src/fileflash/services/agent/plan_service.py b/app/src/fileflash/services/agent/plan_service.py index 94f14d0..f5fc3a5 100644 --- a/app/src/fileflash/services/agent/plan_service.py +++ b/app/src/fileflash/services/agent/plan_service.py @@ -1,8 +1,15 @@ from __future__ import annotations +from datetime import UTC, datetime, timedelta + +from sqlalchemy import and_, func, select from sqlalchemy.ext.asyncio import AsyncSession +from ...core.errors import ApiError +from ...core.settings import Settings +from ...models import BackgroundJob from ...repositories import AgentPlanRepository, AgentSettingsRepository, AgentWorkSessionRepository +from ...schemas.agent import PlanAgentRequest, PlanAgentResponse from ..background_jobs import BackgroundJobService @@ -11,16 +18,72 @@ def __init__( self, *, db: AsyncSession, + settings: Settings, jobs: BackgroundJobService, plans: AgentPlanRepository, - settings: AgentSettingsRepository, + settings_repo: AgentSettingsRepository, work_sessions: AgentWorkSessionRepository, ) -> None: self.db = db + self.settings = settings self.jobs = jobs self.plans = plans - self.settings = settings + self.settings_repo = settings_repo self.work_sessions = work_sessions - async def enqueue_plan(self, *args, **kwargs): - raise NotImplementedError("Agent plan service is scaffolded only in this stage") + async def enqueue_plan(self, *, user_id: int, payload: PlanAgentRequest) -> PlanAgentResponse: + if not self.settings.agent_enabled: + raise ApiError(status_code=503, code=503, message="Agent runtime is disabled") + if payload.hints.budget_tokens > self.settings.agent_job_max_tokens: + raise ApiError( + status_code=400, + code=400, + message="Agent token budget exceeds server limit", + ) + if ( + not self.settings.is_development_env + and payload.hints.max_steps > self.settings.agent_job_max_tool_calls + ): + raise ApiError(status_code=400, code=400, message="Agent maxSteps exceeds server limit") + + await self._enforce_limits(user_id=user_id) + job = await self.jobs.enqueue( + self.db, + task_type="agent.plan", + payload=payload.model_dump(by_alias=True), + requested_by=user_id, + max_attempts=1, + priority=100, + agent_phase="planning", + ) + return PlanAgentResponse( + job_id=str(job.job_id), + status=str(job.status), + task_type="agent.plan", + ) + + async def _enforce_limits(self, *, user_id: int) -> None: + concurrent = await self.db.scalar( + select(func.count(BackgroundJob.job_id)).where( + and_( + BackgroundJob.requested_by == user_id, + BackgroundJob.task_type.in_(["agent.plan", "agent.execute"]), + BackgroundJob.status.in_(["pending", "running", "retrying"]), + ) + ) + ) + if int(concurrent or 0) >= self.settings.agent_user_concurrent_limit: + raise ApiError(status_code=429, code=429, message="Agent concurrent job limit exceeded") + + since = datetime.now(UTC) - timedelta(days=1) + daily = await self.db.scalar( + select(func.count(BackgroundJob.job_id)).where( + and_( + BackgroundJob.requested_by == user_id, + BackgroundJob.task_type == "agent.plan", + BackgroundJob.created_at >= since, + ) + ) + ) + if int(daily or 0) >= self.settings.agent_user_daily_limit: + raise ApiError(status_code=429, code=429, message="Agent daily job limit exceeded") diff --git a/app/src/fileflash/services/agent/skill_service.py b/app/src/fileflash/services/agent/skill_service.py index cc06582..634c2d9 100644 --- a/app/src/fileflash/services/agent/skill_service.py +++ b/app/src/fileflash/services/agent/skill_service.py @@ -23,6 +23,7 @@ UpdateAgentSkillRequest, ) from ...schemas.common import PaginatedData, PaginationMeta +from ...agents.harness.tool_registry import REGISTRY class SkillService: @@ -66,6 +67,19 @@ def _coerce_tool_whitelist(raw: Any) -> list[str]: return [str(item) for item in raw if isinstance(item, (str, int, float))] return [] + @staticmethod + def _validate_tool_whitelist(raw: list[str]) -> list[str]: + tools = [str(item).strip() for item in raw if str(item).strip()] + unknown = REGISTRY.unknown_names(tools) + if unknown: + raise ApiError( + status_code=422, + code=422, + message="Unknown agent tool in toolWhitelist", + data={"unknownTools": sorted(unknown)}, + ) + return tools + @classmethod def _to_item(cls, entity: AgentSkill) -> AgentSkillItem: visibility_value = ( @@ -140,13 +154,14 @@ async def get_skill(self, *, user_id: int, skill_key: str) -> AgentSkillItem: async def create_custom_skill(self, *, user_id: int, payload: CreateAgentSkillRequest) -> AgentSkillItem: skill_key = await self._generate_unique_user_skill_key(user_id=user_id, name=payload.name) + tool_whitelist = self._validate_tool_whitelist(payload.tool_whitelist) entity = await self.skills.create( values={ "skill_key": skill_key, "name": payload.name, "description": payload.description, "triggers_text": payload.triggers_text, - "tool_whitelist_json": payload.tool_whitelist, + "tool_whitelist_json": tool_whitelist, "plan_template_json": payload.plan_template, "inputs_schema_json": payload.inputs_schema, "outputs_schema_json": payload.outputs_schema, @@ -178,7 +193,9 @@ async def update_custom_skill(self, *, user_id: int, skill_key: str, payload: Up if "triggers_text" in fields_set: values["triggers_text"] = payload.triggers_text if "tool_whitelist" in fields_set: - values["tool_whitelist_json"] = payload.tool_whitelist or [] + values["tool_whitelist_json"] = self._validate_tool_whitelist( + payload.tool_whitelist or [] + ) if "plan_template" in fields_set: values["plan_template_json"] = payload.plan_template or {} if "inputs_schema" in fields_set: @@ -232,12 +249,13 @@ async def import_global_skills(self, *, payload: ImportAgentSkillsRequest) -> Im results: list[ImportAgentSkillResult] = [] for item in payload.items: + tool_whitelist = self._validate_tool_whitelist(item.tool_whitelist) existing = existing_by_key.get(item.skill_key) values = { "name": item.name, "description": item.description, "triggers_text": item.triggers_text, - "tool_whitelist_json": item.tool_whitelist, + "tool_whitelist_json": tool_whitelist, "plan_template_json": item.plan_template, "inputs_schema_json": item.inputs_schema, "outputs_schema_json": item.outputs_schema, diff --git a/app/src/fileflash/services/background_jobs.py b/app/src/fileflash/services/background_jobs.py index 4ad5147..7d45fec 100644 --- a/app/src/fileflash/services/background_jobs.py +++ b/app/src/fileflash/services/background_jobs.py @@ -4,6 +4,7 @@ from datetime import UTC, datetime from typing import Any +from fastapi.encoders import jsonable_encoder from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession @@ -29,12 +30,14 @@ async def enqueue( requested_by: int | None = None, max_attempts: int = 5, priority: int = 100, + agent_phase: str | None = None, ) -> BackgroundJob: now = datetime.now(UTC) + normalized_payload = jsonable_encoder(payload) job = BackgroundJob( task_type=task_type, status="pending", - payload=payload, + payload=normalized_payload, result={}, error_message=None, attempt=0, @@ -42,6 +45,7 @@ async def enqueue( scheduled_at=now, trace_id=str(uuid.uuid4()), idempotency_key=idempotency_key, + agent_phase=agent_phase, requested_by=requested_by, priority=priority, ) diff --git a/app/src/fileflash/services/dev_seed.py b/app/src/fileflash/services/dev_seed.py index 0845280..6c66fd6 100644 --- a/app/src/fileflash/services/dev_seed.py +++ b/app/src/fileflash/services/dev_seed.py @@ -59,19 +59,23 @@ async def initialize_dev_accounts( reset_password: bool = False, auto_run: bool = False, ) -> bool: - if auto_run and not settings.is_development_env: - if settings.is_production_env: - logger.info("Skip dev account auto-initialization in production environment: %s", settings.app_env) - else: - logger.info("Skip dev account auto-initialization for non-dev APP_ENV=%s", settings.app_env) + if auto_run and not settings.is_development_env and not settings.is_production_env: + logger.info("Skip account auto-initialization for non-dev APP_ENV=%s", settings.app_env) return False + if settings.is_production_env: + settings.assert_runtime_security() + + accounts = _seed_accounts_for_settings(settings) + async with SessionLocal() as db: - seeder = DevAccountSeeder(db=db) + seeder = DevAccountSeeder(db=db, accounts=accounts) summary = await seeder.seed(reset_password=reset_password) logger.info( - "Dev accounts initialized: createdUsers=%s updatedUsers=%s resetPasswordUsers=%s createdPreferences=%s createdRoots=%s", + "%s initialized: createdUsers=%s updatedUsers=%s resetPasswordUsers=%s " + "createdPreferences=%s createdRoots=%s", + "Production default admin" if settings.is_production_env else "Dev accounts", summary.created_users, summary.updated_users, summary.reset_password_users, @@ -81,15 +85,44 @@ async def initialize_dev_accounts( return True +def _seed_accounts_for_settings(settings: Settings) -> tuple[DevSeedAccount, ...]: + if not settings.is_production_env: + return DEV_SEED_ACCOUNTS + + username = (settings.default_admin_username or "").strip() + email = (settings.default_admin_email or "").strip() + password = (settings.default_admin_password or "").strip() + if not username or not email or not password: + raise ValueError( + "DEFAULT_ADMIN_USERNAME, DEFAULT_ADMIN_EMAIL, and DEFAULT_ADMIN_PASSWORD are required" + ) + + return ( + DevSeedAccount( + username=username, + email=email, + password=password, + role=UserRole.ADMIN, + language=UiLanguage.ZH_CN, + ), + ) + + class DevAccountSeeder: - def __init__(self, *, db: AsyncSession) -> None: + def __init__( + self, + *, + db: AsyncSession, + accounts: tuple[DevSeedAccount, ...] = DEV_SEED_ACCOUNTS, + ) -> None: self.db = db + self.accounts = accounts async def seed(self, *, reset_password: bool = False) -> DevSeedSummary: summary = DevSeedSummary() now = datetime.now(UTC) - for spec in DEV_SEED_ACCOUNTS: + for spec in self.accounts: user = await self._find_existing_user(spec=spec) created = user is None @@ -129,7 +162,12 @@ async def seed(self, *, reset_password: bool = False) -> DevSeedSummary: if created: user.password_changed_at = now - await self._ensure_preference(user=user, language=spec.language, now=now, summary=summary) + await self._ensure_preference( + user=user, + language=spec.language, + now=now, + summary=summary, + ) await self._ensure_root_folder(user=user, now=now, summary=summary) await self.db.commit() @@ -156,7 +194,9 @@ async def _ensure_preference( now: datetime, summary: DevSeedSummary, ) -> None: - preference = await self.db.scalar(select(UserPreference).where(UserPreference.user_id == user.user_id)) + preference = await self.db.scalar( + select(UserPreference).where(UserPreference.user_id == user.user_id) + ) if preference is None: self.db.add( UserPreference( diff --git a/app/src/fileflash/services/download_rate_limit.py b/app/src/fileflash/services/download_rate_limit.py new file mode 100644 index 0000000..7f1870b --- /dev/null +++ b/app/src/fileflash/services/download_rate_limit.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.errors import ApiError +from ..core.settings import Settings +from ..models.enums import UserRole +from ..models.tables_identity import User +from .rate_limiter import RedisRateLimiter + + +class DownloadRateLimitService: + def __init__( + self, + *, + db: AsyncSession, + settings: Settings, + rate_limiter: RedisRateLimiter, + ) -> None: + self.db = db + self.settings = settings + self.rate_limiter = rate_limiter + + async def enforce_user(self, *, user: User, bytes_count: int) -> None: + if user.role == UserRole.ADMIN: + return + await self._enforce(scope=f"user:{int(user.user_id)}", bytes_count=bytes_count) + + async def enforce_user_id(self, *, user_id: int, bytes_count: int) -> None: + user = await self.db.get(User, user_id) + if user is not None and user.role == UserRole.ADMIN: + return + await self._enforce(scope=f"user:{int(user_id)}", bytes_count=bytes_count) + + async def enforce_share_ip(self, *, client_ip: str, bytes_count: int) -> None: + await self._enforce(scope=f"share-ip:{client_ip}", bytes_count=bytes_count) + + async def _enforce(self, *, scope: str, bytes_count: int) -> None: + window_seconds = max(1, int(self.settings.download_rate_window_seconds)) + request_limit = max(1, int(self.settings.download_rate_limit_requests)) + byte_limit = max(1, int(self.settings.download_rate_limit_bytes)) + normalized_bytes = max(0, int(bytes_count)) + + request_allowed = await self.rate_limiter.allow( + key=f"download-rate:{scope}:requests", + limit=request_limit, + window_seconds=window_seconds, + ) + if not request_allowed: + raise ApiError(status_code=429, code=429, message="Download rate limit exceeded") + + bytes_allowed = await self.rate_limiter.allow_weighted( + key=f"download-rate:{scope}:bytes", + limit=byte_limit, + window_seconds=window_seconds, + weight=normalized_bytes, + ) + if not bytes_allowed: + raise ApiError(status_code=429, code=429, message="Download bandwidth limit exceeded") + + +__all__ = ["DownloadRateLimitService"] diff --git a/app/src/fileflash/services/file.py b/app/src/fileflash/services/file.py index 3d910da..7d18814 100644 --- a/app/src/fileflash/services/file.py +++ b/app/src/fileflash/services/file.py @@ -74,6 +74,18 @@ class DownloadStreamResult: headers: dict[str, str] +@dataclass(slots=True) +class BatchDownloadPlan: + files: list[tuple[File, StorageObject, str]] + estimated_bytes: int + + +@dataclass(slots=True) +class ResolvedStreamObject: + storage_object: StorageObject + content_type_override: str | None = None + + class FileService: def __init__( self, @@ -284,10 +296,11 @@ async def _get_file_stream( raise ApiError(status_code=503, code=503, message="Object storage is unavailable") file_row = await self._get_active_file(user_id=user_id, file_id=file_id) - storage_object = await self._resolve_stream_storage_object( + resolved_object = await self._resolve_stream_storage_object( file_row=file_row, prefer_optimized=(content_disposition == "inline"), ) + storage_object = resolved_object.storage_object if resolved_object is not None else None if storage_object is None or storage_object.upload_status != UploadStatus.ACTIVE: raise ApiError(status_code=404, code=404, message="File content not found") @@ -296,7 +309,11 @@ async def _get_file_stream( raise ApiError(status_code=404, code=404, message="File content not found") content_type = resolve_file_mime_type( - mime_type=file_row.mime_type or storage_object.content_type, + mime_type=( + resolved_object.content_type_override + if resolved_object is not None and resolved_object.content_type_override + else file_row.mime_type or storage_object.content_type + ), file_ext=file_row.file_ext, file_name=file_row.file_name, default=DEFAULT_MIME_TYPE, @@ -347,6 +364,15 @@ async def create_batch_download_archive( user_id: int, payload: BatchDownloadRequest, ) -> tuple[str, str]: + plan = await self.create_batch_download_plan(user_id=user_id, payload=payload) + return await self.create_batch_download_archive_from_plan(plan=plan) + + async def create_batch_download_plan( + self, + *, + user_id: int, + payload: BatchDownloadRequest, + ) -> BatchDownloadPlan: if self.storage is None: raise ApiError(status_code=503, code=503, message="Object storage is unavailable") @@ -416,6 +442,31 @@ async def create_batch_download_archive( if not files_with_storage: raise ApiError(status_code=404, code=404, message="No downloadable files found") + files = [ + ( + file_row, + storage_object, + self._safe_zip_path(file_paths.get(int(file_row.file_id), file_row.file_name)), + ) + for file_row, storage_object in files_with_storage + ] + estimated_bytes = sum( + int(storage_object.object_size or file_row.file_size or 0) + for file_row, storage_object, _zip_path in files + ) + return BatchDownloadPlan(files=files, estimated_bytes=max(0, estimated_bytes)) + + async def create_batch_download_archive_from_plan( + self, + *, + plan: BatchDownloadPlan, + ) -> tuple[str, str]: + if self.storage is None: + raise ApiError(status_code=503, code=503, message="Object storage is unavailable") + + if not plan.files: + raise ApiError(status_code=404, code=404, message="No downloadable files found") + archive_name = f"fileflash-download-{datetime.now(UTC).strftime('%Y%m%d-%H%M%S')}.zip" tmp = tempfile.NamedTemporaryFile(prefix="fileflash-download-", suffix=".zip", delete=False) tmp_path = tmp.name @@ -423,8 +474,7 @@ async def create_batch_download_archive( try: with zipfile.ZipFile(tmp_path, mode="w", compression=zipfile.ZIP_DEFLATED, allowZip64=True) as archive: - for file_row, storage_object in files_with_storage: - zip_path = self._safe_zip_path(file_paths.get(int(file_row.file_id), file_row.file_name)) + for _file_row, storage_object, zip_path in plan.files: with archive.open(zip_path, mode="w") as entry: async for chunk in self.storage.iter_object( bucket_name=storage_object.bucket_name, @@ -1832,12 +1882,12 @@ async def _resolve_stream_storage_object( *, file_row: File, prefer_optimized: bool, - ) -> StorageObject | None: + ) -> ResolvedStreamObject | None: source_object = await self.db.get(StorageObject, int(file_row.storage_object_id)) if source_object is None: return None if not prefer_optimized: - return source_object + return ResolvedStreamObject(storage_object=source_object) metadata_row = await self.db.scalar( select(FileMediaMetadata) @@ -1845,17 +1895,19 @@ async def _resolve_stream_storage_object( .limit(1) ) if not isinstance(metadata_row, FileMediaMetadata): - return source_object + return ResolvedStreamObject(storage_object=source_object) transcode = (metadata_row.extra_metadata or {}).get("transcode") if not isinstance(transcode, dict): - return source_object + return ResolvedStreamObject(storage_object=source_object) if str(transcode.get("status") or "").strip().lower() != TRANSCODE_READY_STATUS: - return source_object + return ResolvedStreamObject(storage_object=source_object) bucket_name = str(transcode.get("optimizedBucketName") or "").strip() object_key = str(transcode.get("optimizedObjectKey") or "").strip() if not bucket_name or not object_key: - return source_object + return ResolvedStreamObject(storage_object=source_object) + + optimized_mime_type = str(transcode.get("optimizedMimeType") or "").strip() or None optimized_object = await self.db.scalar( select(StorageObject) @@ -1869,13 +1921,16 @@ async def _resolve_stream_storage_object( .limit(1) ) if isinstance(optimized_object, StorageObject): - return optimized_object + return ResolvedStreamObject( + storage_object=optimized_object, + content_type_override=optimized_mime_type or optimized_object.content_type, + ) if self.storage is None: - return source_object + return ResolvedStreamObject(storage_object=source_object) exists = await self.storage.object_exists(bucket_name=bucket_name, object_key=object_key) if not exists: - return source_object + return ResolvedStreamObject(storage_object=source_object) stat = await self.storage.stat_object(bucket_name=bucket_name, object_key=object_key) created = StorageObject( @@ -1889,7 +1944,10 @@ async def _resolve_stream_storage_object( ) self.db.add(created) await self.db.flush() - return created + return ResolvedStreamObject( + storage_object=created, + content_type_override=optimized_mime_type or created.content_type, + ) @staticmethod def _parse_datetime(raw: object) -> datetime | None: diff --git a/app/src/fileflash/services/rate_limiter.py b/app/src/fileflash/services/rate_limiter.py index 9ced38d..c872fe2 100644 --- a/app/src/fileflash/services/rate_limiter.py +++ b/app/src/fileflash/services/rate_limiter.py @@ -6,6 +6,7 @@ from redis.exceptions import RedisError logger = logging.getLogger(__name__) +_LOOP_RUNTIME_ERROR_MARKERS = ("different loop", "attached to a different loop", "event loop is closed") class RedisRateLimiter: @@ -21,21 +22,57 @@ async def _client(self) -> Redis | None: return self._redis async def allow(self, key: str, limit: int, window_seconds: int) -> bool: + return await self.allow_weighted(key=key, limit=limit, window_seconds=window_seconds, weight=1) + + async def allow_weighted(self, key: str, limit: int, window_seconds: int, weight: int) -> bool: client = await self._client() if client is None: return True + normalized_weight = max(0, int(weight)) try: - current = await client.incr(key) - if current == 1: + current = await client.incrby(key, normalized_weight) + if current == normalized_weight: await client.expire(key, window_seconds) return current <= limit except RedisError: logger.exception("Redis unavailable, rate limiter degraded for key=%s", key) return True + except RuntimeError as exc: + if not _is_loop_runtime_error(exc): + raise + logger.exception("Redis loop error, rate limiter degraded for key=%s", key) + await self._drop_cached_client() + return True async def close(self) -> None: - if self._redis is not None: - await self._redis.aclose() - self._redis = None + client = self._redis + self._redis = None + if client is None: + return + try: + await client.aclose() + except RuntimeError as exc: + if _is_loop_runtime_error(exc): + logger.warning("Ignore Redis close loop error: %s", exc) + return + raise + + async def _drop_cached_client(self) -> None: + client = self._redis + self._redis = None + if client is None: + return + try: + await client.aclose() + except RuntimeError as exc: + if _is_loop_runtime_error(exc): + logger.warning("Ignore Redis reset loop error: %s", exc) + return + raise + + +def _is_loop_runtime_error(exc: RuntimeError) -> bool: + message = str(exc).lower() + return any(marker in message for marker in _LOOP_RUNTIME_ERROR_MARKERS) diff --git a/app/src/fileflash/services/share.py b/app/src/fileflash/services/share.py index 4d0ce35..a2b0755 100644 --- a/app/src/fileflash/services/share.py +++ b/app/src/fileflash/services/share.py @@ -2,7 +2,7 @@ import logging import secrets -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from datetime import UTC, datetime from math import ceil from pathlib import Path @@ -422,6 +422,7 @@ async def get_shared_file_download_stream_response( range_header: str | None, ip_address: str, user_agent: str | None, + rate_limit_check: Callable[[int], Awaitable[None]] | None = None, ) -> tuple[AsyncIterator[bytes], str, str, int, dict[str, str]]: async def _operation() -> tuple[AsyncIterator[bytes], str, str, int, dict[str, str]]: await apply_local_lock_timeout(self.db) @@ -449,6 +450,11 @@ async def _operation() -> tuple[AsyncIterator[bytes], str, str, int, dict[str, s if object_size <= 0: raise ApiError(status_code=404, code=404, message="Shared file content not found") + byte_range = self._parse_range_header(range_header=range_header, file_size=object_size) + bytes_to_send = object_size if byte_range is None else byte_range[1] - byte_range[0] + 1 + if rate_limit_check is not None: + await rate_limit_check(bytes_to_send) + if action == "download": await self.db.execute( update(Share) @@ -479,7 +485,6 @@ async def _operation() -> tuple[AsyncIterator[bytes], str, str, int, dict[str, s ), } - byte_range = self._parse_range_header(range_header=range_header, file_size=object_size) if byte_range is None: headers["Content-Length"] = str(object_size) return ( diff --git a/app/src/fileflash/services/upload.py b/app/src/fileflash/services/upload.py index 0b221d8..491e6d8 100644 --- a/app/src/fileflash/services/upload.py +++ b/app/src/fileflash/services/upload.py @@ -33,7 +33,14 @@ from ..models.tables_storage import File, FileMediaMetadata, Folder, StorageObject, UploadTask, UploadTaskPart from ..models.tables_worker import BackgroundJob from ..s3.minio_client import MinioObjectStorageClient, ObjectStorageError -from ..schemas.file import MergeChunksRequest, MergeChunksResponse, UploadPreflightRequest, UploadPreflightResponse +from ..schemas.file import ( + MergeChunksRequest, + MergeChunksResponse, + RecoverableUploadSession, + UploadCancelResponse, + UploadPreflightRequest, + UploadPreflightResponse, +) from .background_jobs import BackgroundJobService logger = logging.getLogger(__name__) @@ -149,6 +156,65 @@ async def _operation() -> UploadPreflightResponse: raise to_retryable_concurrency_error(exc) from exc raise + async def list_recoverable_sessions(self, *, user_id: int) -> list[RecoverableUploadSession]: + now = datetime.now(UTC) + rows = list( + await self.db.scalars( + select(UploadTask) + .where( + and_( + UploadTask.user_id == user_id, + UploadTask.upload_id.is_not(None), + UploadTask.status.in_((UploadTaskStatus.INIT, UploadTaskStatus.UPLOADING)), + or_(UploadTask.expired_at.is_(None), UploadTask.expired_at > now), + ) + ) + .order_by(UploadTask.updated_at.desc(), UploadTask.task_id.desc()) + ) + ) + + sessions: list[RecoverableUploadSession] = [] + for row in rows: + if row.expired_at and row.expired_at <= now: + continue + if row.status not in (UploadTaskStatus.INIT, UploadTaskStatus.UPLOADING): + continue + upload_id = (row.upload_id or "").strip() + file_name = (row.file_name or "").strip() + file_hash = self._normalize_task_hash(row.object_hash) + if not upload_id or not file_name or not file_hash: + continue + + file_size = max(0, int(row.total_size or 0)) + uploaded_bytes = max(0, min(file_size, int(row.uploaded_bytes or 0))) + chunk_size = int(row.chunk_size or self._resolved_chunk_size()) + parent_id = str(row.folder_id) if row.folder_id is not None else "root" + status = "init" if row.status == UploadTaskStatus.INIT else "uploading" + resolved_mime_type = resolve_file_mime_type( + mime_type=row.mime_type, + file_ext=self._extract_ext(file_name), + file_name=file_name, + default=DEFAULT_MIME_TYPE, + ) + + sessions.append( + RecoverableUploadSession( + upload_id=upload_id, + file_name=file_name, + file_size=file_size, + uploaded_bytes=uploaded_bytes, + chunk_size=max(1, chunk_size), + file_hash=file_hash, + mime_type=resolved_mime_type, + parent_id=parent_id, + updated_at=row.updated_at or now, + expired_at=row.expired_at, + status=status, + ) + ) + + return sessions + async def upload_chunk(self, *, user_id: int, upload_id: str, chunk_index: int, chunk_bytes: bytes) -> None: if chunk_index < 0: raise ApiError(status_code=400, code=400, message="chunkIndex must be >= 0") @@ -445,16 +511,17 @@ async def _operation() -> MergeChunksResponse: await self.db.commit() raise ApiError(status_code=422, code=422, message="Composed file size mismatch") - actual_hash = await self.storage.compute_object_hash( - object_key=task.object_key, - algorithm=hash_algorithm, - ) - if actual_hash != object_hash: - await self.storage.remove_object(object_key=task.object_key) - task.status = UploadTaskStatus.FAILED - task.last_error = "Composed file hash mismatch" - await self.db.commit() - raise ApiError(status_code=422, code=422, message="Composed file hash mismatch") + if self.settings.upload_verify_merged_object_hash: + actual_hash = await self.storage.compute_object_hash( + object_key=task.object_key, + algorithm=hash_algorithm, + ) + if actual_hash != object_hash: + await self.storage.remove_object(object_key=task.object_key) + task.status = UploadTaskStatus.FAILED + task.last_error = "Composed file hash mismatch" + await self.db.commit() + raise ApiError(status_code=422, code=422, message="Composed file hash mismatch") storage_object = await self._find_storage_object( object_hash=object_hash, @@ -546,6 +613,34 @@ async def _operation() -> MergeChunksResponse: raise to_retryable_concurrency_error(exc) from exc raise + async def cancel_upload_session(self, *, user_id: int, upload_id: str) -> UploadCancelResponse: + task = await self._get_task_for_update(user_id=user_id, upload_id=upload_id) + if task is None: + raise ApiError(status_code=404, code=404, message="Upload session not found") + if task.status == UploadTaskStatus.COMPLETED: + raise ApiError(status_code=409, code=409, message="Upload session already completed") + if task.status == UploadTaskStatus.FAILED: + raise ApiError(status_code=409, code=409, message="Upload session is not cancelable") + if task.status == UploadTaskStatus.ABORTED: + return UploadCancelResponse( + upload_id=upload_id, + canceled_at=task.updated_at or datetime.now(UTC), + ) + + canceled_at = datetime.now(UTC) + cleanup_keys = await self._collect_task_cleanup_keys(task=task) + task.status = UploadTaskStatus.ABORTED + task.last_error = "Upload canceled by client" + await self.db.commit() + + if cleanup_keys: + try: + await self.storage.remove_objects(object_keys=cleanup_keys) + except Exception: # noqa: BLE001 + logger.exception("Failed to cleanup canceled upload temp objects uploadId=%s", upload_id) + + return UploadCancelResponse(upload_id=upload_id, canceled_at=canceled_at) + async def _cleanup_expired_tasks(self, *, user_id: int) -> None: now = datetime.now(UTC) expired_tasks = list( @@ -567,10 +662,7 @@ async def _cleanup_expired_tasks(self, *, user_id: int) -> None: for task in expired_tasks: task.status = UploadTaskStatus.ABORTED task.last_error = "Upload session expired" - if task.object_key: - cleanup_keys.append(task.object_key) - part_indexes = await self._list_uploaded_indexes(task_id=task.task_id) - cleanup_keys.extend(self._build_part_object_key(task=task, chunk_index=index) for index in part_indexes) + cleanup_keys.extend(await self._collect_task_cleanup_keys(task=task)) await self.db.commit() if cleanup_keys: @@ -758,6 +850,15 @@ async def _list_uploaded_indexes(self, *, task_id: int) -> list[int]: ) ) + async def _collect_task_cleanup_keys(self, *, task: UploadTask) -> list[str]: + keys: list[str] = [] + if task.object_key: + keys.append(task.object_key) + part_indexes = await self._list_uploaded_indexes(task_id=task.task_id) + keys.extend(self._build_part_object_key(task=task, chunk_index=index) for index in part_indexes) + # Preserve order but deduplicate possible repeated entries. + return list(dict.fromkeys(keys)) + async def _abort_task(self, *, task: UploadTask, reason: str) -> None: task.status = UploadTaskStatus.ABORTED task.last_error = reason diff --git a/app/src/fileflash/tasks/transcode.py b/app/src/fileflash/tasks/transcode.py index 31632a9..adbfa7f 100644 --- a/app/src/fileflash/tasks/transcode.py +++ b/app/src/fileflash/tasks/transcode.py @@ -13,6 +13,11 @@ from ..s3.minio_client import MinioObjectStorageClient TRANSCODE_PROFILE_VERSION = "mp4-v1" +_FFMPEG_DETERMINISTIC_ERROR_PATTERNS = ( + "unsupported channel layout", + "error while opening encoder", + "could not open encoder before eof", +) @dataclass(slots=True) @@ -182,6 +187,8 @@ def build_ffmpeg_command( audio_codec, "-b:a", f"{audio_bitrate}k", + "-ac", + "2", "-map", "0:v:0", "-map", @@ -196,6 +203,8 @@ def build_ffmpeg_command( "aac", "-b:a", f"{audio_bitrate}k", + "-ac", + "2", "-movflags", "+faststart", "-map", @@ -279,10 +288,33 @@ def _run_command(command: list[str], *, timeout_seconds: int) -> subprocess.Comp raise RuntimeError(f"Binary not found for command: {command[0]}") from exc if result.returncode != 0: stderr = (result.stderr or "").strip() - raise RuntimeError(f"Command failed ({result.returncode}): {' '.join(command)} | {stderr}") + error_message = f"Command failed ({result.returncode}): {' '.join(command)} | {stderr}" + if _is_deterministic_ffmpeg_failure(command=command, stderr=stderr): + raise ValueError(error_message) + raise RuntimeError(error_message) return result +def _is_deterministic_ffmpeg_failure(*, command: list[str], stderr: str) -> bool: + if not command: + return False + + binary_name = Path(command[0]).name.lower() + if "ffmpeg" not in binary_name: + return False + + stderr_text = stderr.lower() + if any(pattern in stderr_text for pattern in _FFMPEG_DETERMINISTIC_ERROR_PATTERNS): + return True + + if "invalid argument" in stderr_text and ( + "enc:aac" in stderr_text or "aac @" in stderr_text or "channel layout" in stderr_text + ): + return True + + return False + + def _run_async(awaitable: Any) -> Any: return asyncio.run(awaitable) diff --git a/app/tests/test_admin_files_routes.py b/app/tests/test_admin_files_routes.py new file mode 100644 index 0000000..3b91678 --- /dev/null +++ b/app/tests/test_admin_files_routes.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_files_service, get_settings_dep, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.core.security import create_admin_file_preview_token +from fileflash.core.settings import Settings +from fileflash.routers.admin_files import router as admin_router +from fileflash.schemas.admin.files import RescanResponse +from fileflash.services.admin.files import AdminFileStreamResult + + +class StubService: + async def list_files(self, *, query): # noqa: ANN001 + return SimpleNamespace( + model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, + "totalPages": 1, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": False, + "hasNext": False, + }, + } + ) + + async def get_file_detail(self, *, file_id): # noqa: ANN001 + return SimpleNamespace( + model_dump=lambda **_: { + "id": str(file_id), + "objectId": "42", + "name": "demo.txt", + "size": 12, + "mimeType": "text/plain", + "hash": "abc123", + "virusStatus": "clean", + "isShared": False, + "ownerName": "owner", + "uploadCount": 2, + "ownerCount": 1, + "scannedAt": "2026-01-01T00:00:00Z", + "updatedAt": "2026-01-01T00:00:00Z", + "createdAt": "2026-01-01T00:00:00Z", + "objectHash": "abc123full", + "hashAlgorithm": "sha256", + "storageStatus": "active", + "latestScan": { + "scanType": "virus", + "scanResult": "clean", + "virusStatus": "clean", + "scannedAt": "2026-01-01T00:00:00Z", + "details": {}, + }, + "owners": [ + { + "userId": "1", + "username": "owner", + "email": "owner@example.com", + "fileCount": 2, + "firstUploadedAt": "2026-01-01T00:00:00Z", + "lastUploadedAt": "2026-01-01T00:00:00Z", + } + ], + } + ) + + async def get_preview_stream(self, *, file_id, range_header): # noqa: ANN001 + _ = file_id + + async def _stream(content: bytes) -> AsyncIterator[bytes]: + yield content + + headers = { + "Accept-Ranges": "bytes", + "Content-Disposition": 'inline; filename="demo.txt"', + } + if range_header: + headers["Content-Length"] = "4" + headers["Content-Range"] = "bytes 0-3/12" + return AdminFileStreamResult( + stream=_stream(b"prev"), + filename="demo.txt", + content_type="text/plain", + status_code=206, + headers=headers, + ) + + headers["Content-Length"] = "12" + return AdminFileStreamResult( + stream=_stream(b"preview-bytes"), + filename="demo.txt", + content_type="text/plain", + status_code=200, + headers=headers, + ) + + async def request_rescan(self, *, file_id, requested_by): # noqa: ANN001 + _ = requested_by + return RescanResponse( + file_id=str(file_id), + virus_status="pending", + scanned_at=datetime.now(UTC), + ) + + +def _client(*, admin: bool = True, settings: Settings | None = None) -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(admin_router, prefix="/api/v1") + app.dependency_overrides[get_admin_files_service] = lambda: StubService() + + def _admin_override(): + if not admin: + raise ApiError(status_code=403, code=403, message="Admin access required") + return SimpleNamespace(user_id=99) + + resolved_settings = settings or Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + + app.dependency_overrides[require_admin] = _admin_override + app.dependency_overrides[get_settings_dep] = lambda: resolved_settings + return TestClient(app) + + +def test_list_files_returns_empty() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/files") + assert resp.status_code == 200 + + +def test_rescan_returns_pending() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/files/7/rescan") + assert resp.status_code == 200 + assert resp.json()["data"]["virusStatus"] == "pending" + + +def test_get_file_detail_returns_owner_and_upload_counts() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/files/7") + assert resp.status_code == 200 + data = resp.json()["data"] + assert data["uploadCount"] == 2 + assert data["ownerCount"] == 1 + assert data["owners"][0]["email"] == "owner@example.com" + + +def test_preview_route_returns_inline_stream() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/files/7/preview", headers={"Range": "bytes=0-3"}) + assert resp.status_code == 206 + assert resp.headers["content-range"] == "bytes 0-3/12" + assert resp.content == b"prev" + + +def test_create_preview_url_returns_admin_stream_url() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/files/7/preview-url") + assert resp.status_code == 200 + data = resp.json()["data"] + assert data["url"].startswith("http://testserver/api/v1/admin/files/7/preview-stream?token=") + assert data["expiresAt"] + + +def test_admin_preview_requires_admin() -> None: + with _client(admin=False) as c: + resp = c.get("/api/v1/admin/files/7/preview") + assert resp.status_code == 403 + assert resp.json()["message"] == "Admin access required" + + +def test_preview_stream_supports_range_with_valid_token() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_admin_file_preview_token( + admin_user_id=99, + file_id=7, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + + with _client(settings=settings) as c: + resp = c.get(f"/api/v1/admin/files/7/preview-stream?token={token}", headers={"Range": "bytes=0-3"}) + + assert resp.status_code == 206 + assert resp.headers["content-range"] == "bytes 0-3/12" + assert resp.content == b"prev" + + +def test_preview_stream_rejects_mismatched_token_file_id() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_admin_file_preview_token( + admin_user_id=99, + file_id=8, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + + with _client(settings=settings) as c: + resp = c.get(f"/api/v1/admin/files/7/preview-stream?token={token}") + + assert resp.status_code == 403 + assert resp.json()["message"] == "Preview token does not match file" + + +def test_preview_stream_rejects_expired_token() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_admin_file_preview_token( + admin_user_id=99, + file_id=7, + settings=settings, + expires_at=datetime.now(UTC) - timedelta(minutes=10), + ) + + with _client(settings=settings) as c: + resp = c.get(f"/api/v1/admin/files/7/preview-stream?token={token}") + + assert resp.status_code == 401 + assert resp.json()["message"] == "Invalid or expired preview token" diff --git a/app/tests/test_admin_files_service.py b/app/tests/test_admin_files_service.py new file mode 100644 index 0000000..8889275 --- /dev/null +++ b/app/tests/test_admin_files_service.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import FileStatus, ScanResult +from fileflash.schemas.admin.files import AdminFileAuditOwner, ListAdminFilesQuery +from fileflash.services.admin.files import AdminFilesService + + +class DummyPublisher: + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + + async def publish(self, event_name: str, payload: dict) -> None: # noqa: ANN001 + self.calls.append((event_name, payload)) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.execute = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + + +@pytest.mark.asyncio +async def test_list_returns_paginated_empty_when_no_files() -> None: + session = DummySession() + session.scalar.return_value = 0 + session.execute.return_value = Mock(all=lambda: []) + service = AdminFilesService(db=session, publisher=DummyPublisher()) # type: ignore[arg-type] + + result = await service.list_files(query=ListAdminFilesQuery()) + + assert result.items == [] + assert result.pagination.total_items == 0 + + +@pytest.mark.asyncio +async def test_rescan_missing_file_returns_404() -> None: + session = DummySession() + session.get.return_value = None + service = AdminFilesService(db=session, publisher=DummyPublisher()) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.request_rescan(file_id=1, requested_by=99) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_rescan_inserts_scan_record_and_publishes_event() -> None: + session = DummySession() + file_row = Mock( + file_id=1, + storage_object_id=2, + deleted_at=None, + status=FileStatus.ACTIVE, + ) + session.get.return_value = file_row + publisher = DummyPublisher() + service = AdminFilesService(db=session, publisher=publisher) # type: ignore[arg-type] + + result = await service.request_rescan(file_id=1, requested_by=99) + + assert result.virus_status == "pending" + assert publisher.calls and publisher.calls[0][0] == "files.rescan_requested" + session.commit.assert_awaited() + + +@pytest.mark.asyncio +async def test_detail_aggregates_active_object_owner_counts(monkeypatch: pytest.MonkeyPatch) -> None: + session = DummySession() + service = AdminFilesService(db=session, publisher=DummyPublisher()) # type: ignore[arg-type] + now = datetime.now(UTC) + file_row = Mock( + file_id=10, + file_name="report.pdf", + file_size=128, + mime_type="application/pdf", + file_ext="pdf", + updated_at=now, + created_at=now, + ) + object_row = Mock( + object_id=20, + object_hash="a" * 64, + hash_algorithm="sha256", + content_type="application/pdf", + upload_status="active", + ) + owner_row = Mock(username="alice") + scan_row = Mock( + result=ScanResult.CLEAN, + scan_type="virus", + scanned_at=now, + details={"engine": "unit"}, + ) + owners = [ + AdminFileAuditOwner( + user_id="1", + username="alice", + email="alice@example.com", + file_count=2, + first_uploaded_at=now, + last_uploaded_at=now, + ), + AdminFileAuditOwner( + user_id="2", + username="bob", + email="bob@example.com", + file_count=1, + first_uploaded_at=now, + last_uploaded_at=now, + ), + ] + + monkeypatch.setattr( + service, + "_get_active_file_context", + AsyncMock(return_value=(file_row, object_row, owner_row, scan_row)), + ) + monkeypatch.setattr(service, "_load_object_owners", AsyncMock(return_value=owners)) + monkeypatch.setattr(service, "_object_is_shared", AsyncMock(return_value=True)) + + detail = await service.get_file_detail(file_id=10) + + assert detail.object_id == "20" + assert detail.upload_count == 3 + assert detail.owner_count == 2 + assert detail.is_shared is True + assert detail.latest_scan and detail.latest_scan.virus_status == "clean" diff --git a/app/tests/test_admin_logs_routes.py b/app/tests/test_admin_logs_routes.py new file mode 100644 index 0000000..4d82e73 --- /dev/null +++ b/app/tests/test_admin_logs_routes.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_logs_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_logs import router as logs_router +from fileflash.schemas.admin.logs import AdminLogsResponse + + +class StubService: + async def list_logs(self, *, query): # noqa: ANN001 + return AdminLogsResponse( + logs=[], + pagination={ + "totalItems": 0, + "totalPages": 1, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": False, + "hasNext": False, + }, + ) + + +def test_admin_can_list_logs() -> None: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(logs_router, prefix="/api/v1") + app.dependency_overrides[get_admin_logs_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + with TestClient(app) as c: + resp = c.get("/api/v1/admin/logs") + assert resp.status_code == 200 + assert resp.json()["data"]["logs"] == [] diff --git a/app/tests/test_admin_moderation_routes.py b/app/tests/test_admin_moderation_routes.py new file mode 100644 index 0000000..04112b4 --- /dev/null +++ b/app/tests/test_admin_moderation_routes.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_moderation_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_moderation import router as mod_router +from fileflash.schemas.admin.moderation import ResolveViolationResponse + + +class StubService: + async def list_violations(self, *, query): # noqa: ANN001 + return SimpleNamespace( + model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, + "totalPages": 1, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": False, + "hasNext": False, + }, + } + ) + + async def resolve_case(self, *, case_id, handled_by): # noqa: ANN001 + _ = handled_by + return ResolveViolationResponse(violation_id=str(case_id), resolved_at=datetime.now(UTC)) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(mod_router, prefix="/api/v1") + app.dependency_overrides[get_admin_moderation_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_list_violations() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/violations") + assert resp.status_code == 200 + + +def test_resolve_violation() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/violations/3/resolve") + assert resp.status_code == 200 + assert resp.json()["data"]["violationId"] == "3" diff --git a/app/tests/test_admin_moderation_service.py b/app/tests/test_admin_moderation_service.py new file mode 100644 index 0000000..a0007f6 --- /dev/null +++ b/app/tests/test_admin_moderation_service.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.schemas.admin.moderation import ListViolationsQuery +from fileflash.services.admin.moderation import AdminModerationService + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_resolve_missing_case_returns_404() -> None: + session = DummySession() + session.get.return_value = None + service = AdminModerationService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.resolve_case(case_id=1, handled_by=2) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_resolve_already_resolved_returns_409() -> None: + session = DummySession() + case = Mock(case_id=1, status="resolved") + session.get.return_value = case + service = AdminModerationService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.resolve_case(case_id=1, handled_by=2) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_resolve_pending_case_sets_resolved() -> None: + session = DummySession() + case = Mock(case_id=1, status="pending", handled_by=None, handled_at=None, resolution=None, file_id=10) + session.get.return_value = case + service = AdminModerationService(db=session) # type: ignore[arg-type] + + result = await service.resolve_case(case_id=1, handled_by=2) + + assert result.violation_id == "1" + assert case.status == "resolved" + assert case.resolution == "admin_clear" + assert case.handled_by == 2 + session.commit.assert_awaited() + + +@pytest.mark.asyncio +async def test_list_returns_empty_paginated() -> None: + session = DummySession() + session.scalar.return_value = 0 + session.execute.return_value = Mock(all=lambda: []) + service = AdminModerationService(db=session) # type: ignore[arg-type] + + result = await service.list_violations(query=ListViolationsQuery()) + assert result.items == [] + assert result.pagination.total_items == 0 diff --git a/app/tests/test_admin_notifications_routes.py b/app/tests/test_admin_notifications_routes.py new file mode 100644 index 0000000..c708e2b --- /dev/null +++ b/app/tests/test_admin_notifications_routes.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_notifications_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_notifications import router as n_router +from fileflash.schemas.admin.notifications import BroadcastResponse + + +class StubService: + async def list_notifications(self, *, query): # noqa: ANN001 + return SimpleNamespace( + model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, + "totalPages": 1, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": False, + "hasNext": False, + }, + } + ) + + async def broadcast(self, *, payload, sender_id): # noqa: ANN001 + _ = payload, sender_id + return BroadcastResponse( + broadcast_id="b1", + recipient_count=2, + sent_at=datetime.now(UTC), + ) + + async def archive(self, *, notification_id): # noqa: ANN001 + _ = notification_id + return None + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(n_router, prefix="/api/v1") + app.dependency_overrides[get_admin_notifications_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_list() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/notifications") + assert resp.status_code == 200 + + +def test_broadcast() -> None: + with _client() as c: + resp = c.post( + "/api/v1/admin/notifications/broadcast", + json={"message": "hello", "type": "system"}, + ) + assert resp.status_code == 200 + assert resp.json()["data"]["broadcastId"] == "b1" + + +def test_archive() -> None: + with _client() as c: + resp = c.delete("/api/v1/admin/notifications/9") + assert resp.status_code == 200 diff --git a/app/tests/test_admin_notifications_service.py b/app/tests/test_admin_notifications_service.py new file mode 100644 index 0000000..d8236ca --- /dev/null +++ b/app/tests/test_admin_notifications_service.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.schemas.admin.notifications import BroadcastRequest +from fileflash.services.admin.notifications import ( + MAX_BROADCAST_RECIPIENTS, + AdminNotificationsService, +) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.add_all = Mock() + self.commit = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_broadcast_rejects_empty_message() -> None: + session = DummySession() + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + with pytest.raises(ApiError): + await service.broadcast( + payload=BroadcastRequest.model_construct(title=None, message="", type="system"), + sender_id=1, + ) + + +@pytest.mark.asyncio +async def test_broadcast_too_many_recipients_returns_422() -> None: + session = DummySession() + session.scalar.return_value = MAX_BROADCAST_RECIPIENTS + 1 + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + with pytest.raises(ApiError) as exc: + await service.broadcast(payload=BroadcastRequest(message="hi", type="system"), sender_id=1) + assert exc.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_broadcast_writes_one_row_per_recipient() -> None: + session = DummySession() + session.scalar.return_value = 3 + session.scalars.return_value = iter([10, 11, 12]) + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + + result = await service.broadcast(payload=BroadcastRequest(message="ping", type="system"), sender_id=1) + + assert result.recipient_count == 3 + assert session.add_all.called diff --git a/app/tests/test_admin_storage_routes.py b/app/tests/test_admin_storage_routes.py new file mode 100644 index 0000000..69c3396 --- /dev/null +++ b/app/tests/test_admin_storage_routes.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_storage_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_storage import router as storage_router +from fileflash.schemas.admin.storage import ( + AdminStorageSummary, + UpdateQuotaResponse, + UsageTrendPoint, + UsageTrendResponse, +) + + +class StubService: + async def summary(self): # noqa: ANN201 + return AdminStorageSummary( + storage_used=1000, + storage_limit=10000, + storage_percentage=10.0, + file_count=3, + user_count=2, + updated_at=datetime.now(UTC), + ) + + async def list_storage_users(self, *, query): # noqa: ANN001, ANN201 + return SimpleNamespace( + model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, + "totalPages": 1, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": False, + "hasNext": False, + }, + } + ) + + async def update_quota(self, *, user_id, new_limit): # noqa: ANN001, ANN201 + return UpdateQuotaResponse( + user_id=str(user_id), + storage_limit=new_limit, + storage_used=0, + usage_percentage=0.0, + updated_at=datetime.now(UTC), + ) + + async def usage_trend(self, *, query): # noqa: ANN001, ANN201 + _ = query + return UsageTrendResponse( + trends=[UsageTrendPoint(date="2026-05-24", used=1)], + is_estimated=False, + ) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(storage_router, prefix="/api/v1") + app.dependency_overrides[get_admin_storage_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_summary_returns_camel_case() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/storage/summary") + assert resp.status_code == 200 + assert "storageUsed" in resp.json()["data"] + + +def test_update_quota_passes_storage_limit() -> None: + with _client() as c: + resp = c.patch("/api/v1/admin/storage/users/7/quota", json={"storageLimit": 1024}) + assert resp.status_code == 200 + assert resp.json()["data"]["storageLimit"] == 1024 + + +def test_usage_trend_default_days() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/storage/usage-trend") + assert resp.status_code == 200 + assert len(resp.json()["data"]["trends"]) == 1 diff --git a/app/tests/test_admin_storage_service.py b/app/tests/test_admin_storage_service.py new file mode 100644 index 0000000..fb4233b --- /dev/null +++ b/app/tests/test_admin_storage_service.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import UserRole, UserStatus +from fileflash.models.tables_identity import User +from fileflash.schemas.admin.storage import UsageTrendQuery +from fileflash.services.admin.storage import AdminStorageService + + +def _user(**kwargs) -> User: + base = dict( + user_id=1, + username="bob", + email="b@x.com", + password_hash="x", + role=UserRole.USER, + status=UserStatus.ACTIVE, + email_verified=True, + email_verified_at=datetime.now(UTC), + storage_limit=10 * 1024 * 1024 * 1024, + storage_used=2 * 1024 * 1024 * 1024, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + base.update(kwargs) + return User(**base) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.refresh = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_summary_aggregates_users_and_files() -> None: + session = DummySession() + session.scalar.side_effect = [ + 2 * 1024 * 1024 * 1024, + 10 * 1024 * 1024 * 1024, + 42, + 7, + ] + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.summary() + + assert result.storage_used == 2 * 1024 * 1024 * 1024 + assert result.user_count == 7 + assert result.file_count == 42 + assert round(result.storage_percentage, 2) == 20.0 + + +@pytest.mark.asyncio +async def test_update_quota_rejects_below_usage() -> None: + session = DummySession() + target = _user(storage_used=5 * 1024 * 1024 * 1024) + session.get.return_value = target + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.update_quota(user_id=1, new_limit=1 * 1024 * 1024 * 1024) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_update_quota_success_updates_and_commits() -> None: + session = DummySession() + target = _user(storage_used=1 * 1024 * 1024 * 1024) + session.get.return_value = target + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.update_quota(user_id=1, new_limit=20 * 1024 * 1024 * 1024) + + assert result.storage_limit == 20 * 1024 * 1024 * 1024 + assert target.storage_limit == 20 * 1024 * 1024 * 1024 + session.commit.assert_awaited() + + +@pytest.mark.asyncio +async def test_usage_trend_returns_n_points() -> None: + session = DummySession() + session.scalar.return_value = 2 * 1024 * 1024 * 1024 + session.execute.return_value = Mock(all=lambda: []) + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.usage_trend(query=UsageTrendQuery(days=7)) + + assert len(result.trends) == 7 + assert result.is_estimated is True diff --git a/app/tests/test_admin_system_routes.py b/app/tests/test_admin_system_routes.py new file mode 100644 index 0000000..f758ba6 --- /dev/null +++ b/app/tests/test_admin_system_routes.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_system_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_system import router as sys_router +from fileflash.schemas.admin.system import RateLimitRule, RateLimitStatus, SystemHealth + + +class StubService: + async def health(self): # noqa: ANN201 + return SystemHealth( + platform_targets=["s3://bkt"], + max_concurrent_uploads=4, + active_upload_sessions=0, + virus_scan_enabled=False, + thumbnail_generation_enabled=True, + registration_mail_enabled=False, + hash_computation_enabled=True, + last_updated_at=datetime.now(UTC), + ) + + async def rate_limit_status(self): # noqa: ANN201 + return RateLimitStatus( + rules=[ + RateLimitRule( + rule_id="login", + scope="auth.login", + window_seconds=60, + limit=5, + current_usage=0, + blocked_requests=0, + ) + ], + evaluated_at=datetime.now(UTC), + ) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(sys_router, prefix="/api/v1") + app.dependency_overrides[get_admin_system_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_health() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/system/health") + assert resp.status_code == 200 + assert resp.json()["data"]["activeUploadSessions"] == 0 + + +def test_rate_limit() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/system/rate-limit") + assert resp.status_code == 200 + assert resp.json()["data"]["rules"][0]["scope"] == "auth.login" diff --git a/app/tests/test_admin_system_service.py b/app/tests/test_admin_system_service.py new file mode 100644 index 0000000..2c94b41 --- /dev/null +++ b/app/tests/test_admin_system_service.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import pytest + +from fileflash.core.settings import Settings +from fileflash.services.admin.system import AdminSystemService + + +class DummySession: + async def scalar(self, _query: object) -> int: + return 0 + + +def make_settings(**overrides: object) -> Settings: + payload = { + "FF_DB_URI": "postgresql://root:pwd@localhost:5432/fileflash", + "JWT_SECRET_KEY": "unit-test-secret-key-1234567890abcd", + } + payload.update(overrides) + return Settings(**payload) + + +@pytest.mark.asyncio +async def test_health_hash_computation_enabled_follows_settings() -> None: + disabled_service = AdminSystemService( + db=DummySession(), + settings=make_settings(UPLOAD_VERIFY_MERGED_OBJECT_HASH=False), + ) + disabled_health = await disabled_service.health() + assert disabled_health.hash_computation_enabled is False + + enabled_service = AdminSystemService( + db=DummySession(), + settings=make_settings(UPLOAD_VERIFY_MERGED_OBJECT_HASH=True), + ) + enabled_health = await enabled_service.health() + assert enabled_health.hash_computation_enabled is True + + +@pytest.mark.asyncio +async def test_rate_limit_status_uses_auth_default_limits() -> None: + service = AdminSystemService( + db=DummySession(), + settings=make_settings(), + ) + + status = await service.rate_limit_status() + rules_by_scope = {rule.scope: rule for rule in status.rules} + + assert rules_by_scope["auth.login"].limit == 30 + assert rules_by_scope["auth.login"].window_seconds == 300 + assert rules_by_scope["auth.register"].limit == 12 + assert rules_by_scope["auth.register"].window_seconds == 600 diff --git a/app/tests/test_admin_user_status_mapping.py b/app/tests/test_admin_user_status_mapping.py new file mode 100644 index 0000000..351aeb3 --- /dev/null +++ b/app/tests/test_admin_user_status_mapping.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import UserStatus +from fileflash.services.admin._status import external_to_internal, internal_to_external + + +def test_active_maps_both_ways() -> None: + assert external_to_internal("active") == UserStatus.ACTIVE + assert internal_to_external(UserStatus.ACTIVE) == "active" + + +def test_suspended_maps_to_disabled() -> None: + assert external_to_internal("suspended") == UserStatus.DISABLED + assert internal_to_external(UserStatus.DISABLED) == "suspended" + + +def test_locked_externally_appears_as_suspended() -> None: + assert internal_to_external(UserStatus.LOCKED) == "suspended" + + +def test_pending_verification_passthrough() -> None: + assert internal_to_external(UserStatus.PENDING_VERIFICATION) == "pending_verification" + + +def test_invalid_external_status_raises() -> None: + with pytest.raises(ApiError): + external_to_internal("garbage") diff --git a/app/tests/test_admin_users_routes.py b/app/tests/test_admin_users_routes.py new file mode 100644 index 0000000..bff1702 --- /dev/null +++ b/app/tests/test_admin_users_routes.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_admin_users_service, require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_users import router as admin_users_router +from fileflash.schemas.admin.users import AdminUserItem, AdminUserUsageStats, UpdateUserStatusResponse +from fileflash.schemas.common import PaginatedData, PaginationMeta + + +class StubService: + async def list_users(self, *, query): # noqa: ANN001 + item = AdminUserItem( + user_id="1", + username="alice", + email="a@x.com", + role="USER", + status="active", + email_verified=True, + email_verified_at=None, + storage_limit=1024, + storage_used=0, + usage_percentage=0.0, + last_login_at=None, + last_active_at=None, + created_at=datetime.now(UTC), + usage_stats=AdminUserUsageStats(traffic_bytes=1024, agent_tokens=42), + ) + return PaginatedData( + items=[item], + pagination=PaginationMeta( + total_items=1, + total_pages=1, + per_page=query.per_page, + current_page=query.page, + has_prev=False, + has_next=False, + ), + ) + + async def set_status(self, *, user_id, external_status): # noqa: ANN001 + return UpdateUserStatusResponse( + user_id=str(user_id), + status=external_status, + updated_at=datetime.now(UTC), + ) + + +def _client(admin: bool) -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(admin_users_router, prefix="/api/v1") + app.dependency_overrides[get_admin_users_service] = lambda: StubService() + if admin: + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + else: + + async def _deny(): + raise ApiError(status_code=403, code=403, message="forbidden") + + app.dependency_overrides[require_admin] = _deny + return TestClient(app) + + +def test_admin_can_list_users() -> None: + with _client(admin=True) as c: + resp = c.get("/api/v1/admin/users") + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["data"]["items"][0]["username"] == "alice" + assert body["data"]["items"][0]["usageStats"] == {"trafficBytes": 1024, "agentTokens": 42} + + +def test_non_admin_gets_403() -> None: + with _client(admin=False) as c: + resp = c.get("/api/v1/admin/users") + assert resp.status_code == 403 + + +def test_usage_window_requires_both_bounds() -> None: + with _client(admin=True) as c: + resp = c.get("/api/v1/admin/users?usageFrom=2026-01-01T00:00:00Z") + assert resp.status_code == 400 + + +def test_usage_window_rejects_reversed_bounds() -> None: + with _client(admin=True) as c: + resp = c.get( + "/api/v1/admin/users" + "?usageFrom=2026-02-01T00:00:00Z" + "&usageTo=2026-01-01T00:00:00Z" + ) + assert resp.status_code == 400 + + +def test_usage_window_rejects_more_than_90_days() -> None: + with _client(admin=True) as c: + resp = c.get( + "/api/v1/admin/users" + "?usageFrom=2026-01-01T00:00:00Z" + "&usageTo=2026-04-02T00:00:00Z" + ) + assert resp.status_code == 400 + + +def test_admin_can_patch_status() -> None: + with _client(admin=True) as c: + resp = c.patch("/api/v1/admin/users/42/status", json={"status": "suspended"}) + assert resp.status_code == 200 + assert resp.json()["data"]["status"] == "suspended" diff --git a/app/tests/test_admin_users_service.py b/app/tests/test_admin_users_service.py new file mode 100644 index 0000000..1370023 --- /dev/null +++ b/app/tests/test_admin_users_service.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import UserRole, UserStatus +from fileflash.models.tables_identity import User +from fileflash.schemas.admin.users import ListAdminUsersQuery +from fileflash.services.admin.users import AdminUsersService + + +def _user(**kwargs) -> User: + base = dict( + user_id=1, + username="alice", + email="alice@example.com", + password_hash="x", + role=UserRole.USER, + status=UserStatus.ACTIVE, + email_verified=True, + email_verified_at=datetime.now(UTC), + storage_limit=10 * 1024 * 1024 * 1024, + storage_used=1024, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + base.update(kwargs) + return User(**base) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.refresh = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +class ResultRows: + def __init__(self, rows) -> None: # noqa: ANN001 + self._rows = rows + + def all(self): # noqa: ANN201 + return self._rows + + +@pytest.mark.asyncio +async def test_list_users_returns_paginated_items() -> None: + session = DummySession() + session.scalar.return_value = 1 + session.scalars.return_value = [_user()] + session.execute.return_value = Mock(all=lambda: []) + service = AdminUsersService(db=session) # type: ignore[arg-type] + + result = await service.list_users(query=ListAdminUsersQuery()) + + assert result.pagination.total_items == 1 + assert result.items[0].username == "alice" + assert result.items[0].status == "active" + assert result.items[0].usage_stats.traffic_bytes == 0 + assert result.items[0].usage_stats.agent_tokens == 0 + + +def test_list_users_query_default_usage_window() -> None: + now = datetime(2026, 5, 26, 12, 0, tzinfo=UTC) + usage_from, usage_to = ListAdminUsersQuery().resolve_usage_window(now=now) + + assert usage_to == now + assert (usage_to - usage_from).days == 7 + + +@pytest.mark.asyncio +async def test_collect_usage_stats_aggregates_traffic_and_tokens() -> None: + session = DummySession() + session.execute = AsyncMock( + side_effect=[ + ResultRows([(1, 2048), (2, 4096)]), + ResultRows([(1, 1500), (3, None)]), + ] + ) + service = AdminUsersService(db=session) # type: ignore[arg-type] + + stats = await service._collect_usage_stats( + user_ids=[1, 2, 3], + usage_from=datetime(2026, 5, 1, tzinfo=UTC), + usage_to=datetime(2026, 5, 26, tzinfo=UTC), + ) + + assert stats[1].traffic_bytes == 2048 + assert stats[1].agent_tokens == 1500 + assert stats[2].traffic_bytes == 4096 + assert stats[2].agent_tokens == 0 + assert stats[3].traffic_bytes == 0 + assert stats[3].agent_tokens == 0 + + +@pytest.mark.asyncio +async def test_set_status_user_not_found() -> None: + session = DummySession() + session.get.return_value = None + service = AdminUsersService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.set_status(user_id=999, external_status="suspended") + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_set_status_suspend_last_admin_blocked() -> None: + session = DummySession() + admin = _user(user_id=2, role=UserRole.ADMIN, status=UserStatus.ACTIVE) + session.get.return_value = admin + session.scalar.return_value = 0 + service = AdminUsersService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.set_status(user_id=2, external_status="suspended") + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_set_status_suspends_user_and_revokes_sessions() -> None: + session = DummySession() + target = _user(user_id=3) + session.get.return_value = target + session.scalar.return_value = 5 + service = AdminUsersService(db=session) # type: ignore[arg-type] + + result = await service.set_status(user_id=3, external_status="suspended") + + assert result.status == "suspended" + assert target.status == UserStatus.DISABLED + session.commit.assert_awaited() diff --git a/app/tests/test_agent_a_end_to_end.py b/app/tests/test_agent_a_end_to_end.py new file mode 100644 index 0000000..377fcbc --- /dev/null +++ b/app/tests/test_agent_a_end_to_end.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from test_agent_inbox_repository import InboxSession + +from fileflash.agents.harness.event_bus import InMemoryAgentEventBus +from fileflash.agents.harness.inbox import AgentInbox +from fileflash.agents.runtime import execute_runner as execute_module +from fileflash.agents.runtime.execute_runner import AgentJobCanceled, ExecuteRunner +from fileflash.models import BackgroundJob +from fileflash.models.enums import AgentInboxKind + + +class RuntimeInboxSession(InboxSession): + async def refresh(self, _job: BackgroundJob) -> None: + return None + + async def rollback(self) -> None: + return None + + +def _execute_job() -> BackgroundJob: + now = datetime.now(UTC) + return BackgroundJob( + job_id=800, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "500", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": now.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=now, + created_at=now, + updated_at=now, + ) + + +@pytest.mark.asyncio +async def test_user_pause_then_cancel_via_inbox(monkeypatch: pytest.MonkeyPatch): + action = { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace(proposed_actions_json=[action]) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + + session = RuntimeInboxSession() + bus = InMemoryAgentEventBus() + inbox = AgentInbox(db=session, event_bus=bus) # type: ignore[arg-type] + job = _execute_job() + seen_events: list[str] = [] + + await inbox.handle(job_id=int(job.job_id), kind=AgentInboxKind.CONTROL_PAUSE, payload={}) + await session.commit() + + async def cancel_when_paused() -> None: + async with bus.subscribe(job_id=int(job.job_id)) as stream: + while True: + event = await stream.next(timeout=2.0) + seen_events.append(event.event_type) + if event.event_type == "agent.paused": + await inbox.handle( + job_id=int(job.job_id), + kind=AgentInboxKind.CONTROL_CANCEL, + payload={}, + ) + await session.commit() + return + + listener = asyncio.create_task(cancel_when_paused()) + await asyncio.sleep(0) + with pytest.raises(AgentJobCanceled): + await ExecuteRunner(event_bus=bus).run(db=session, job=job) # type: ignore[arg-type] + await listener + + assert "agent.paused" in seen_events + assert job.cancel_requested_at is not None diff --git a/app/tests/test_agent_ask_protocol.py b/app/tests/test_agent_ask_protocol.py new file mode 100644 index 0000000..d935075 --- /dev/null +++ b/app/tests/test_agent_ask_protocol.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import asyncio + +import pytest +from test_agent_inbox_repository import InboxSession + +from fileflash.agents.harness.ask import AskProtocol, AskTimedOut +from fileflash.agents.harness.event_bus import InMemoryAgentEventBus +from fileflash.agents.harness.inbox import AgentInbox +from fileflash.models.enums import AgentInboxKind, AgentInboxStatus + + +@pytest.mark.asyncio +async def test_ask_returns_when_reply_arrives() -> None: + session = InboxSession() + bus = InMemoryAgentEventBus() + inbox = AgentInbox(db=session, event_bus=bus) # type: ignore[arg-type] + protocol = AskProtocol(db=session, event_bus=bus, job_id=1) # type: ignore[arg-type] + + await protocol.start() + try: + async def reply_later() -> None: + for _ in range(20): + asks = [msg for msg in session.messages if msg.kind == AgentInboxKind.ASK] + if asks: + ask = asks[-1] + await inbox.handle( + job_id=1, + kind=AgentInboxKind.REPLY, + payload={"value": "A"}, + reply_to_id=int(ask.inbox_message_id), + ) + await session.commit() + return + await asyncio.sleep(0.01) + raise AssertionError("ask message was not created") + + replier = asyncio.create_task(reply_later()) + result = await protocol.ask( + prompt="choose", + schema={"choice": ["A", "B"]}, + timeout_sec=2.0, + ) + await replier + finally: + await protocol.aclose() + + assert result == "A" + + +@pytest.mark.asyncio +async def test_ask_times_out() -> None: + session = InboxSession() + bus = InMemoryAgentEventBus() + protocol = AskProtocol(db=session, event_bus=bus, job_id=1) # type: ignore[arg-type] + + await protocol.start() + try: + with pytest.raises(AskTimedOut): + await protocol.ask(prompt="?", schema={}, timeout_sec=0.1) + finally: + await protocol.aclose() + + asks = [msg for msg in session.messages if msg.kind == AgentInboxKind.ASK] + assert asks + assert asks[-1].status == AgentInboxStatus.TIMED_OUT diff --git a/app/tests/test_agent_event_bus.py b/app/tests/test_agent_event_bus.py new file mode 100644 index 0000000..6214640 --- /dev/null +++ b/app/tests/test_agent_event_bus.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +from fileflash.agents.harness.event_bus import AgentEventEnvelope, InMemoryAgentEventBus + + +@pytest.mark.asyncio +async def test_subscriber_receives_published_event() -> None: + bus = InMemoryAgentEventBus() + envelope = AgentEventEnvelope( + job_id=42, + event_type="agent.ask", + payload={"prompt": "choose"}, + emitted_at=datetime.now(UTC), + ) + + async with bus.subscribe(job_id=42) as stream: + await bus.publish(envelope) + received = await stream.next(timeout=1.0) + + assert received == envelope + + +@pytest.mark.asyncio +async def test_only_subscribers_of_same_job_receive() -> None: + bus = InMemoryAgentEventBus() + own = AgentEventEnvelope( + job_id=1, + event_type="job.running", + payload={}, + emitted_at=datetime.now(UTC), + ) + other = AgentEventEnvelope( + job_id=2, + event_type="job.running", + payload={}, + emitted_at=datetime.now(UTC), + ) + + async with bus.subscribe(job_id=1) as stream: + await bus.publish(other) + await bus.publish(own) + first = await stream.next(timeout=1.0) + + assert first == own + + +@pytest.mark.asyncio +async def test_empty_subscriber_times_out() -> None: + bus = InMemoryAgentEventBus() + async with bus.subscribe(job_id=7) as stream: + with pytest.raises(TimeoutError): + await stream.next(timeout=0.1) + + +def test_event_envelope_json_serializes_nested_datetime_payload() -> None: + now = datetime.now(UTC).replace(microsecond=0) + envelope = AgentEventEnvelope( + job_id=9, + event_type="job.succeeded", + payload={ + "data": { + "result": { + "finishedAt": now, + "steps": [{"completedAt": now}], + } + } + }, + emitted_at=now, + event_id="evt-1", + ) + + decoded = AgentEventEnvelope.from_json(envelope.to_json()) + + assert decoded.job_id == envelope.job_id + assert decoded.event_type == envelope.event_type + assert decoded.event_id == "evt-1" + assert decoded.payload["data"]["result"]["finishedAt"] == now.isoformat() + assert decoded.payload["data"]["result"]["steps"][0]["completedAt"] == now.isoformat() diff --git a/app/tests/test_agent_inbox.py b/app/tests/test_agent_inbox.py new file mode 100644 index 0000000..7520ffa --- /dev/null +++ b/app/tests/test_agent_inbox.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import pytest +from test_agent_inbox_repository import InboxSession + +from fileflash.agents.harness.event_bus import InMemoryAgentEventBus +from fileflash.agents.harness.inbox import AgentInbox +from fileflash.models import AgentInboxMessage +from fileflash.models.enums import AgentInboxKind +from fileflash.repositories import AgentInboxMessageRepository + + +@pytest.mark.asyncio +async def test_handle_reply_persists_and_publishes() -> None: + session = InboxSession() + repo = AgentInboxMessageRepository(session) # type: ignore[arg-type] + ask = await repo.create_ask(job_id=1, payload={"prompt": "?"}) + await session.commit() + + bus = InMemoryAgentEventBus() + inbox = AgentInbox(db=session, event_bus=bus) # type: ignore[arg-type] + + async with bus.subscribe(job_id=1) as stream: + msg = await inbox.handle( + job_id=1, + kind=AgentInboxKind.REPLY, + payload={"value": "yes"}, + reply_to_id=int(ask.inbox_message_id), + ) + await session.commit() + event = await stream.next(timeout=1.0) + + assert isinstance(msg, AgentInboxMessage) + assert msg.kind == AgentInboxKind.REPLY + assert event.event_type == "agent.inbox.reply" + assert event.payload["replyTo"] == str(ask.inbox_message_id) + assert event.payload["value"] == "yes" + + +@pytest.mark.asyncio +async def test_reply_with_unknown_ask_raises() -> None: + session = InboxSession() + bus = InMemoryAgentEventBus() + inbox = AgentInbox(db=session, event_bus=bus) # type: ignore[arg-type] + + with pytest.raises(ValueError): + await inbox.handle( + job_id=1, + kind=AgentInboxKind.REPLY, + payload={"value": "yes"}, + reply_to_id=999999, + ) diff --git a/app/tests/test_agent_inbox_model.py b/app/tests/test_agent_inbox_model.py new file mode 100644 index 0000000..b9fae17 --- /dev/null +++ b/app/tests/test_agent_inbox_model.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +from fileflash.models import AgentInboxMessage +from fileflash.models.enums import AgentInboxKind, AgentInboxRole, AgentInboxStatus + + +def test_agent_inbox_message_model_fields() -> None: + msg = AgentInboxMessage( + job_id=1, + role=AgentInboxRole.AGENT, + kind=AgentInboxKind.ASK, + payload_json={"prompt": "which one?", "schema": {}}, + status=AgentInboxStatus.WAITING, + created_at=datetime.now(UTC), + ) + + assert AgentInboxMessage.__tablename__ == "agent_inbox_message" + assert msg.kind == AgentInboxKind.ASK + assert msg.status == AgentInboxStatus.WAITING + assert msg.payload_json["prompt"] == "which one?" diff --git a/app/tests/test_agent_inbox_repository.py b/app/tests/test_agent_inbox_repository.py new file mode 100644 index 0000000..7f8f62d --- /dev/null +++ b/app/tests/test_agent_inbox_repository.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest + +from fileflash.models import AgentInboxMessage +from fileflash.models.enums import AgentInboxKind, AgentInboxRole, AgentInboxStatus +from fileflash.repositories import AgentInboxMessageRepository + + +class InboxSession: + def __init__(self) -> None: + self.messages: list[AgentInboxMessage] = [] + self._next_id = 1 + self.commits = 0 + + def add(self, msg: AgentInboxMessage) -> None: + msg.inbox_message_id = self._next_id + self._next_id += 1 + self.messages.append(msg) + + async def flush(self) -> None: + return None + + async def commit(self) -> None: + self.commits += 1 + + async def get(self, _model, inbox_message_id: int): # noqa: ANN001 + for msg in self.messages: + if msg.inbox_message_id == inbox_message_id: + return msg + return None + + async def scalar(self, _query): # noqa: ANN001 + return None + + async def scalars(self, _query): # noqa: ANN001 + controls = { + AgentInboxKind.CONTROL_PAUSE, + AgentInboxKind.CONTROL_RESUME, + AgentInboxKind.CONTROL_APPROVE, + AgentInboxKind.CONTROL_DENY, + AgentInboxKind.CONTROL_SKIP, + AgentInboxKind.CONTROL_CANCEL, + } + return [ + msg + for msg in self.messages + if msg.role == AgentInboxRole.USER + and msg.kind in controls + and msg.status is None + ] + + +@pytest.mark.asyncio +async def test_create_ask_then_record_reply() -> None: + session = InboxSession() + repo = AgentInboxMessageRepository(session) # type: ignore[arg-type] + + ask = await repo.create_ask( + job_id=1, + payload={"prompt": "choose", "schema": {"choice": ["A", "B"]}}, + ) + await session.commit() + assert ask.status == AgentInboxStatus.WAITING + assert ask.role == AgentInboxRole.AGENT + assert ask.kind == AgentInboxKind.ASK + + reply = await repo.record_user_message( + job_id=1, + kind=AgentInboxKind.REPLY, + payload={"value": "A"}, + reply_to_id=int(ask.inbox_message_id), + ) + await session.commit() + assert reply.role == AgentInboxRole.USER + assert reply.reply_to_id == ask.inbox_message_id + + answered = await repo.mark_answered( + inbox_message_id=int(ask.inbox_message_id), + answered_at=datetime.now(UTC), + ) + await session.commit() + assert answered.status == AgentInboxStatus.ANSWERED + assert answered.answered_at is not None + + +@pytest.mark.asyncio +async def test_pending_controls_excludes_consumed() -> None: + session = InboxSession() + repo = AgentInboxMessageRepository(session) # type: ignore[arg-type] + pause = await repo.record_user_message( + job_id=1, + kind=AgentInboxKind.CONTROL_PAUSE, + payload={}, + ) + await session.commit() + + pending = await repo.list_pending_controls(job_id=1) + assert [msg.inbox_message_id for msg in pending] == [pause.inbox_message_id] + + await repo.mark_dropped(inbox_message_id=int(pause.inbox_message_id)) + await session.commit() + pending_after = await repo.list_pending_controls(job_id=1) + assert pending_after == [] diff --git a/app/tests/test_agent_inbox_schema.py b/app/tests/test_agent_inbox_schema.py new file mode 100644 index 0000000..f6f7cd2 --- /dev/null +++ b/app/tests/test_agent_inbox_schema.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from fileflash.schemas.agent import AgentInboxMessageRequest + + +def test_reply_with_value_validates() -> None: + msg = AgentInboxMessageRequest.model_validate( + {"kind": "reply", "replyTo": "42", "value": "yes"} + ) + + assert msg.kind == "reply" + assert msg.reply_to == "42" + assert msg.value == "yes" + + +def test_unknown_kind_rejected() -> None: + with pytest.raises(ValidationError): + AgentInboxMessageRequest.model_validate({"kind": "control.explode"}) diff --git a/app/tests/test_agent_plan_execute_runtime.py b/app/tests/test_agent_plan_execute_runtime.py new file mode 100644 index 0000000..245bf4b --- /dev/null +++ b/app/tests/test_agent_plan_execute_runtime.py @@ -0,0 +1,2793 @@ +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from fileflash.agents.harness.event_bus import AgentEventEnvelope +from fileflash.agents.harness.policy import PolicyGuard, classify_tool_risk +from fileflash.agents.harness.router import ToolCall, ToolRouter +from fileflash.agents.runtime import execute_runner as execute_module +from fileflash.agents.runtime import plan_runner as plan_module +from fileflash.agents.runtime.execute_runner import AgentJobCanceled, ExecuteRunner +from fileflash.agents.runtime.llm import AnthropicPlannerClient +from fileflash.agents.runtime.plan_runner import PlanRunner +from fileflash.core.errors import ApiError +from fileflash.models import BackgroundJob +from fileflash.models.enums import AgentInboxKind +from fileflash.repositories import ( + AgentPlanRepository, + AgentSettingsRepository, + AgentWorkSessionRepository, +) +from fileflash.schemas.agent import ExecuteAgentRequest, PlanAgentRequest +from fileflash.services.agent.execute_service import ExecuteService +from fileflash.services.agent.plan_service import PlanService + + +class DummyDb: + def __init__(self) -> None: + self.scalar = AsyncMock() + self.execute = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.add = AsyncMock() + self.flush = AsyncMock() + self.commit = AsyncMock() + self.rollback = AsyncMock() + self.refresh = AsyncMock() + + +class FakeJobs: + def __init__(self) -> None: + self.kwargs = {} + + async def enqueue(self, db, **kwargs): # noqa: ANN001 + self.kwargs = kwargs + return BackgroundJob( + job_id=123, + task_type=kwargs["task_type"], + status="pending", + payload=kwargs["payload"], + result={}, + requested_by=kwargs["requested_by"], + max_attempts=kwargs["max_attempts"], + priority=kwargs["priority"], + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + +def settings(**overrides): + base = { + "agent_enabled": True, + "agent_job_max_tokens": 50_000, + "agent_job_max_tool_calls": 100, + "is_development_env": False, + "agent_user_concurrent_limit": 2, + "agent_user_daily_limit": 50, + "agent_llm_base_url": None, + "agent_llm_plan_max_tokens": 8192, + } + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.mark.asyncio +async def test_plan_enqueue_returns_frontend_shape_and_sets_phase(): + db = DummyDb() + db.scalar = AsyncMock(side_effect=[0, 0]) + jobs = FakeJobs() + service = PlanService( + db=db, + settings=settings(), + jobs=jobs, # type: ignore[arg-type] + plans=AgentPlanRepository(db), # type: ignore[arg-type] + settings_repo=AgentSettingsRepository(db), # type: ignore[arg-type] + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = PlanAgentRequest.model_validate( + { + "input": "整理当前文件夹", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + "dataPolicy": { + "allowFileContent": False, + "maxReadBytes": 1024, + "allowedMimeTypes": ["*/*"], + }, + "hints": { + "preferSkillId": None, + "maxSteps": 12, + "budgetTokens": 8000, + "reasoningEffort": "adaptive", + }, + } + ) + + result = await service.enqueue_plan(user_id=7, payload=payload) + + assert result.job_id == "123" + assert result.task_type == "agent.plan" + assert jobs.kwargs["agent_phase"] == "planning" + assert jobs.kwargs["payload"]["executionPolicy"] == "confirm" + assert jobs.kwargs["payload"]["hints"]["reasoningEffort"] == "adaptive" + + +@pytest.mark.asyncio +async def test_plan_enqueue_rejects_max_steps_above_server_limit_in_non_dev(): + db = DummyDb() + jobs = FakeJobs() + service = PlanService( + db=db, + settings=settings(), + jobs=jobs, # type: ignore[arg-type] + plans=AgentPlanRepository(db), # type: ignore[arg-type] + settings_repo=AgentSettingsRepository(db), # type: ignore[arg-type] + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = PlanAgentRequest.model_validate( + { + "input": "整理当前文件夹", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + "hints": { + "preferSkillId": None, + "maxSteps": 101, + "budgetTokens": 8000, + "reasoningEffort": "adaptive", + }, + } + ) + + with pytest.raises(ApiError) as exc: + await service.enqueue_plan(user_id=7, payload=payload) + + assert exc.value.status_code == 400 + assert exc.value.message == "Agent maxSteps exceeds server limit" + + +@pytest.mark.asyncio +async def test_plan_enqueue_allows_max_steps_above_server_limit_in_dev(): + db = DummyDb() + db.scalar = AsyncMock(side_effect=[0, 0]) + jobs = FakeJobs() + service = PlanService( + db=db, + settings=settings(is_development_env=True), + jobs=jobs, # type: ignore[arg-type] + plans=AgentPlanRepository(db), # type: ignore[arg-type] + settings_repo=AgentSettingsRepository(db), # type: ignore[arg-type] + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = PlanAgentRequest.model_validate( + { + "input": "整理当前文件夹", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + "hints": { + "preferSkillId": None, + "maxSteps": 5000, + "budgetTokens": 8000, + "reasoningEffort": "adaptive", + }, + } + ) + + result = await service.enqueue_plan(user_id=7, payload=payload) + + assert result.job_id == "123" + assert jobs.kwargs["payload"]["hints"]["maxSteps"] == 5000 + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_uses_sdk_and_parses_text_blocks(): + class FakeMessages: + def __init__(self) -> None: + self.kwargs = {} + + async def create(self, **kwargs): # noqa: ANN003 + self.kwargs = kwargs + return SimpleNamespace( + content=[ + SimpleNamespace( + type="text", + text='{"summary":"ok","proposedActions":[]}', + ) + ], + usage=SimpleNamespace(input_tokens=3, output_tokens=4), + ) + + fake_messages = FakeMessages() + fake_client = SimpleNamespace(messages=fake_messages) + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=fake_client, # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=9000, + reasoning_effort="adaptive", + ) + + assert fake_messages.kwargs["model"] == "claude-test" + assert fake_messages.kwargs["max_tokens"] == 8192 + assert fake_messages.kwargs["system"] == "system" + assert fake_messages.kwargs["messages"] == [{"role": "user", "content": "user"}] + assert fake_messages.kwargs["thinking"] == {"type": "adaptive"} + assert "output_config" not in fake_messages.kwargs + assert result["summary"] == "ok" + assert result["_usage"] == {"input_tokens": 3, "output_tokens": 4} + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_maps_reasoning_effort_to_output_config(): + class FakeMessages: + def __init__(self) -> None: + self.kwargs = {} + + async def create(self, **kwargs): # noqa: ANN003 + self.kwargs = kwargs + return SimpleNamespace( + content=[SimpleNamespace(type="text", text='{"proposedActions":[]}')], + usage={}, + ) + + fake_messages = FakeMessages() + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=800, + reasoning_effort="xhigh", + ) + + assert fake_messages.kwargs["max_tokens"] == 800 + assert fake_messages.kwargs["thinking"] == {"type": "enabled"} + assert fake_messages.kwargs["output_config"] == {"effort": "xhigh"} + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_retries_without_reasoning_on_empty_output(): + class FakeMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def create(self, **kwargs): # noqa: ANN003 + self.calls.append(dict(kwargs)) + if len(self.calls) == 1: + return SimpleNamespace( + content=[SimpleNamespace(type="thinking", thinking="...")], + usage={}, + ) + return SimpleNamespace( + content=[SimpleNamespace(type="text", text='{"summary":"fallback","proposedActions":[]}')], + usage={}, + ) + + fake_messages = FakeMessages() + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1024, + reasoning_effort="high", + ) + + assert len(fake_messages.calls) == 2 + assert "thinking" in fake_messages.calls[0] + assert "output_config" in fake_messages.calls[0] + assert "thinking" not in fake_messages.calls[1] + assert "output_config" not in fake_messages.calls[1] + assert result["summary"] == "fallback" + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_retries_with_json_only_prompt_on_invalid_json(): + class FakeMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def create(self, **kwargs): # noqa: ANN003 + self.calls.append(dict(kwargs)) + if len(self.calls) == 1: + return SimpleNamespace( + content=[SimpleNamespace(type="text", text="I think we should move files.")], + usage={}, + ) + if len(self.calls) == 2: + return SimpleNamespace( + content=[SimpleNamespace(type="text", text="summary: move files")], + usage={}, + ) + return SimpleNamespace( + content=[SimpleNamespace(type="text", text='{"summary":"strict","proposedActions":[]}')], + usage={}, + ) + + fake_messages = FakeMessages() + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1024, + reasoning_effort="high", + ) + + assert len(fake_messages.calls) == 3 + assert "thinking" in fake_messages.calls[0] + assert "output_config" in fake_messages.calls[0] + assert "thinking" not in fake_messages.calls[1] + assert "output_config" not in fake_messages.calls[1] + assert "thinking" not in fake_messages.calls[2] + assert "output_config" not in fake_messages.calls[2] + assert fake_messages.calls[1]["max_tokens"] == 1024 + assert fake_messages.calls[2]["max_tokens"] == 8192 + third_messages = fake_messages.calls[2]["messages"] + assert isinstance(third_messages, list) + assert "Return ONLY one valid JSON object" in third_messages[0]["content"] + assert result["summary"] == "strict" + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_parses_json_from_wrapped_text(): + class FakeMessages: + async def create(self, **kwargs): # noqa: ANN003 + return SimpleNamespace( + content=[ + SimpleNamespace( + type="text", + text='Here is the result:\n{"summary":"ok","proposedActions":[]}\nThanks!', + ) + ], + usage={}, + ) + + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=FakeMessages()), # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1000, + ) + + assert result["summary"] == "ok" + assert result["proposedActions"] == [] + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_raises_after_three_invalid_json_responses(): + class FakeMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def create(self, **kwargs): # noqa: ANN003 + self.calls.append(dict(kwargs)) + return SimpleNamespace( + content=[SimpleNamespace(type="text", text="not valid json")], + usage={}, + ) + + fake_messages = FakeMessages() + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + with pytest.raises(ApiError) as exc: + await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1024, + reasoning_effort="high", + ) + + assert exc.value.status_code == 502 + assert exc.value.message == "Agent LLM did not return valid JSON" + assert len(fake_messages.calls) == 3 + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_uses_tools_and_parses_tool_use_blocks(): + class FakeMessages: + def __init__(self) -> None: + self.kwargs = {} + + async def create(self, **kwargs): # noqa: ANN003 + self.kwargs = kwargs + return SimpleNamespace( + content=[ + SimpleNamespace(type="text", text="Count matching videos."), + SimpleNamespace( + type="tool_use", + id="toolu_1", + name="drive_count_files", + input={"folderId": "root", "category": "video"}, + ), + ], + usage={"input_tokens": 5, "output_tokens": 6}, + ) + + fake_messages = FakeMessages() + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1024, + tools=[ + { + "name": "drive_count_files", + "description": "Count files.", + "input_schema": {"type": "object"}, + "internalName": "drive.countFiles", + } + ], + ) + + assert fake_messages.kwargs["tool_choice"] == {"type": "auto"} + assert fake_messages.kwargs["tools"] == [ + { + "name": "drive_count_files", + "description": "Count files.", + "input_schema": {"type": "object"}, + } + ] + assert result["summary"] == "Count matching videos." + assert result["proposedActions"] == [ + { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "category": "video"}, + } + ] + assert result["_usage"] == {"input_tokens": 5, "output_tokens": 6} + + +@pytest.mark.asyncio +async def test_anthropic_planner_client_executes_tool_loop_before_final_plan(): + class FakeMessages: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + async def create(self, **kwargs): # noqa: ANN003 + self.calls.append(dict(kwargs)) + if len(self.calls) == 1: + return SimpleNamespace( + content=[ + SimpleNamespace(type="text", text="Need to inspect candidates first."), + SimpleNamespace( + type="tool_use", + id="toolu_1", + name="drive_search_files", + input={"folderId": "root", "query": "银翼杀手", "category": "video"}, + ), + ], + usage={"input_tokens": 10, "output_tokens": 5}, + ) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="text", + text=( + '{"summary":"move matched file","proposedActions":[' + '{"step":1,"tool":"drive.moveFile","input":{"fileId":"11","targetFolderId":"21"}}' + "]}" + ), + ) + ], + usage={"input_tokens": 8, "output_tokens": 9}, + ) + + fake_messages = FakeMessages() + tool_executor = AsyncMock(return_value={"items": [{"id": "11", "name": "银翼杀手2049.mp4"}], "totalItems": 1}) + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + ), + client=SimpleNamespace(messages=fake_messages), # type: ignore[arg-type] + ) + + result = await client.create_plan( + system_prompt="system", + user_prompt="user", + max_tokens=1024, + tools=[ + { + "name": "drive_search_files", + "description": "Search files.", + "input_schema": {"type": "object"}, + "internalName": "drive.searchFiles", + }, + { + "name": "drive_move_file", + "description": "Move file.", + "input_schema": {"type": "object"}, + "internalName": "drive.moveFile", + }, + ], + tool_executor=tool_executor, + max_tool_roundtrips=4, + ) + + tool_executor.assert_awaited_once_with( + "drive.searchFiles", + {"folderId": "root", "query": "银翼杀手", "category": "video"}, + ) + assert len(fake_messages.calls) == 2 + second_messages = fake_messages.calls[1]["messages"] + assert isinstance(second_messages, list) + tool_result_blocks = second_messages[-1]["content"] + assert isinstance(tool_result_blocks, list) + assert tool_result_blocks[0]["type"] == "tool_result" + assert result["summary"] == "move matched file" + assert result["proposedActions"][0]["tool"] == "drive.moveFile" + assert result["_usage"] == {"input_tokens": 18, "output_tokens": 14} + + +def test_anthropic_planner_client_uses_configured_base_url(monkeypatch: pytest.MonkeyPatch): + captured: dict[str, object] = {} + + def fake_async_anthropic(**kwargs): # noqa: ANN003 + captured.update(kwargs) + return SimpleNamespace(messages=SimpleNamespace()) + + monkeypatch.setattr("fileflash.agents.runtime.llm.AsyncAnthropic", fake_async_anthropic) + client = AnthropicPlannerClient( + settings=settings( + agent_llm_api_key="test-key", + agent_llm_model="claude-test", + agent_llm_base_url="https://api.deepseek.com/anthropic", + ), + ) + + assert client._get_client("test-key").messages is not None + assert captured["api_key"] == "test-key" + assert captured["base_url"] == "https://api.deepseek.com/anthropic" + + +@pytest.mark.asyncio +async def test_execute_rejects_high_risk_plan_without_confirmation(): + db = DummyDb() + db.scalar.return_value = BackgroundJob( + job_id=99, + task_type="agent.plan", + status="succeeded", + payload={}, + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + plans = AgentPlanRepository(db) # type: ignore[arg-type] + plans.get_for_execute_binding = AsyncMock( + return_value=SimpleNamespace( + proposed_actions_json=[ + { + "step": 1, + "tool": "drive.deleteFile", + "input": {"fileId": "1"}, + "sideEffect": "write", + "riskLevel": "high", + "requiresConfirmation": True, + } + ] + ) + ) + service = ExecuteService( + db=db, + settings=settings(), + jobs=FakeJobs(), # type: ignore[arg-type] + plans=plans, + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = ExecuteAgentRequest.model_validate( + { + "planJobId": "99", + "planHash": "sha256:test", + "approval": {"confirmedBy": "7", "confirmedAt": datetime.now(UTC).isoformat()}, + } + ) + + with pytest.raises(ApiError) as exc: + await service.enqueue_execute(user_id=7, payload=payload) + + assert exc.value.status_code == 409 + assert exc.value.data["highRiskActions"][0]["tool"] == "drive.deleteFile" + + +@pytest.mark.asyncio +async def test_execute_enqueue_serializes_approval_datetime_as_json_string(): + db = DummyDb() + plan_job = BackgroundJob( + job_id=99, + task_type="agent.plan", + status="succeeded", + payload={}, + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db.scalar = AsyncMock(side_effect=[plan_job, None]) + plans = AgentPlanRepository(db) # type: ignore[arg-type] + plans.get_for_execute_binding = AsyncMock(return_value=SimpleNamespace(proposed_actions_json=[])) + jobs = FakeJobs() + service = ExecuteService( + db=db, + settings=settings(), + jobs=jobs, # type: ignore[arg-type] + plans=plans, + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = ExecuteAgentRequest.model_validate( + { + "planJobId": "99", + "planHash": "sha256:test", + "approval": {"confirmedBy": "7", "confirmedAt": "2026-05-25T10:00:00Z"}, + } + ) + + await service.enqueue_execute(user_id=7, payload=payload) + + approval_payload = jobs.kwargs["payload"]["approval"] + assert isinstance(approval_payload["confirmedAt"], str) + assert approval_payload["confirmedAt"] == "2026-05-25T10:00:00Z" + assert jobs.kwargs["idempotency_key"] == "agent.execute:99" + + +@pytest.mark.asyncio +async def test_execute_rejects_repeat_when_existing_execute_job_exists(): + db = DummyDb() + plan_job = BackgroundJob( + job_id=99, + task_type="agent.plan", + status="succeeded", + payload={}, + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + existing_execute = BackgroundJob( + job_id=200, + task_type="agent.execute", + status="running", + payload={}, + result={}, + requested_by=7, + idempotency_key="agent.execute:99", + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db.scalar = AsyncMock(side_effect=[plan_job, existing_execute]) + plans = AgentPlanRepository(db) # type: ignore[arg-type] + plans.get_for_execute_binding = AsyncMock(return_value=SimpleNamespace(proposed_actions_json=[])) + jobs = FakeJobs() + service = ExecuteService( + db=db, + settings=settings(), + jobs=jobs, # type: ignore[arg-type] + plans=plans, + work_sessions=AgentWorkSessionRepository(db), # type: ignore[arg-type] + ) + payload = ExecuteAgentRequest.model_validate( + { + "planJobId": "99", + "planHash": "sha256:test", + "approval": {"confirmedBy": "7", "confirmedAt": "2026-05-25T10:00:00Z"}, + } + ) + + with pytest.raises(ApiError) as exc: + await service.enqueue_execute(user_id=7, payload=payload) + + assert exc.value.status_code == 409 + assert "already been executed" in exc.value.message.lower() or "already" in exc.value.message.lower() + assert exc.value.data["jobId"] == "200" + assert jobs.kwargs == {} + + +@pytest.mark.asyncio +async def test_plan_runner_generates_stable_hash(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "Move files into folders.", + "proposedActions": [ + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "1", "targetFolderId": "2"}, + }, + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Docs"}, + } + ], + } + ) + client = SimpleNamespace(create_plan=planner) + request = PlanAgentRequest.model_validate( + { + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "autopilot", + "dataPolicy": { + "allowFileContent": False, + "maxReadBytes": 1024, + "allowedMimeTypes": ["*/*"], + }, + "hints": { + "preferSkillId": None, + "maxSteps": 12, + "budgetTokens": 8000, + "reasoningEffort": "high", + }, + } + ) + job = BackgroundJob( + job_id=321, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + runner = PlanRunner(settings=settings(), planner_client=client) # type: ignore[arg-type] + + first = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + second = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert first.plan_hash == second.plan_hash + assert first.requires_confirmation is False + assert [action.step for action in first.proposed_actions] == [1, 2] + assert "reasoningEffort" in planner.await_args.kwargs["user_prompt"] + + +@pytest.mark.asyncio +async def test_plan_runner_ignores_requested_max_steps_in_dev(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "Move files into folders.", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Docs"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "1", "targetFolderId": "2"}, + }, + ], + } + ) + request = PlanAgentRequest.model_validate( + { + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "autopilot", + "hints": { + "preferSkillId": None, + "maxSteps": 1, + "budgetTokens": 8000, + "reasoningEffort": "adaptive", + }, + } + ) + job = BackgroundJob( + job_id=3221, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + runner = PlanRunner( + settings=settings(is_development_env=True), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 2 + assert [action.step for action in result.proposed_actions] == [1, 2] + + +@pytest.mark.asyncio +async def test_plan_runner_commits_after_upsert(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock(return_value={"summary": "ok", "proposedActions": []}) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + } + ) + job = BackgroundJob( + job_id=321, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db = DummyDb() + + await runner.run(db=db, job=job) # type: ignore[arg-type] + + db.commit.assert_awaited_once() + db.rollback.assert_not_called() + + +@pytest.mark.asyncio +async def test_plan_runner_rolls_back_when_commit_fails(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock(return_value={"summary": "ok", "proposedActions": []}) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + } + ) + job = BackgroundJob( + job_id=322, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db = DummyDb() + db.commit.side_effect = RuntimeError("commit failed") + + with pytest.raises(RuntimeError, match="commit failed"): + await runner.run(db=db, job=job) # type: ignore[arg-type] + + db.commit.assert_awaited_once() + db.rollback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_plan_runner_propagates_llm_output_errors_without_fallback( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock( + side_effect=ApiError( + status_code=502, + code=502, + message="Agent LLM returned an empty response", + ) + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "autopilot", + } + ) + job = BackgroundJob( + job_id=333, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db = DummyDb() + + with pytest.raises(ApiError) as exc: + await runner.run(db=db, job=job) # type: ignore[arg-type] + assert exc.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_plan_runner_uses_planner_returned_count_action_for_movie_question( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock( + return_value={ + "summary": "count videos", + "proposedActions": [ + { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + } + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "我上传了多少部电影？", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=334, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 1 + assert result.proposed_actions[0].tool == "drive.countFiles" + assert result.proposed_actions[0].input["category"] == "video" + assert result.proposed_actions[0].side_effect == "read" + + +@pytest.mark.asyncio +async def test_plan_runner_uses_planner_returned_count_action_for_anime_question( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock( + return_value={ + "summary": "count anime videos", + "proposedActions": [ + { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + } + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "我上传了多少动漫？", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=336, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 1 + action = result.proposed_actions[0] + assert action.tool == "drive.countFiles" + assert action.input["category"] == "video" + assert action.input.get("search") in (None, "") + + +@pytest.mark.asyncio +async def test_plan_runner_delegates_count_question_with_search_term_to_planner( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + planner = AsyncMock( + return_value={ + "summary": "count matching movies", + "proposedActions": [ + { + "step": 1, + "tool": "drive.countFiles", + "input": { + "folderId": "root", + "recursive": True, + "category": "video", + "search": "银翼杀手", + }, + } + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "我上传了几部银翼杀手？", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=335, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + planner.assert_awaited_once() + assert len(result.proposed_actions) == 1 + action = result.proposed_actions[0] + assert action.tool == "drive.countFiles" + assert action.input["category"] == "video" + assert action.input["search"] == "银翼杀手" + assert action.side_effect == "read" + + +@pytest.mark.asyncio +async def test_plan_runner_blocks_write_tool_in_exploratory_loop_and_continues( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + dispatch_mock = AsyncMock(return_value={"ok": True}) + monkeypatch.setattr( + plan_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace(dispatch=dispatch_mock), + ) + + async def fake_create_plan(**kwargs): # noqa: ANN003 + tool_executor = kwargs["tool_executor"] + blocked = await tool_executor("drive.createFolder", {"parentFolderId": "root", "name": "Movies"}) + assert blocked["_plannerBlocked"] is True + assert blocked["_toolError"] is True + assert blocked["tool"] == "drive.createFolder" + return { + "summary": "create then move", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "1", "targetFolderId": "$step1.folderId"}, + }, + ], + } + + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=fake_create_plan), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "整理文件", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=340, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + assert len(result.proposed_actions) == 2 + assert result.proposed_actions[0].tool == "drive.createFolder" + assert result.proposed_actions[1].input["targetFolderId"] == "$step1.folderId" + dispatch_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_plan_runner_passes_only_read_tools_to_planner_tool_use( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + captured_tools: list[dict[str, Any]] = [] + + async def fake_create_plan(**kwargs): # noqa: ANN003 + nonlocal captured_tools + tools = kwargs.get("tools") + if isinstance(tools, list): + captured_tools = list(tools) + return {"summary": "ok", "proposedActions": []} + + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=fake_create_plan), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "整理当前文件夹", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=341, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert captured_tools + internal_names = {str(item.get("internalName") or "") for item in captured_tools} + assert "drive.createFolder" not in internal_names + assert "drive.moveFile" not in internal_names + for internal_name in internal_names: + assert plan_module.REGISTRY.get(internal_name).side_effect == "read" + + +@pytest.mark.asyncio +async def test_plan_runner_uses_planner_returned_move_action_when_unique_matches( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "move one matching file", + "proposedActions": [ + { + "step": 1, + "tool": "drive.moveFile", + "input": {"fileId": "11", "targetFolderId": "21", "shareHandling": "keep"}, + } + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "把银翼杀手电影放到银翼杀手文件夹下", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=337, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 1 + action = result.proposed_actions[0] + assert action.tool == "drive.moveFile" + assert action.input["fileId"] == "11" + assert action.input["targetFolderId"] == "21" + assert action.side_effect == "write" + + +@pytest.mark.asyncio +async def test_plan_runner_uses_planner_returned_create_then_move_when_target_missing( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "create missing folder then move", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "银翼杀手"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "11", "targetFolderId": "$step1.folderId"}, + }, + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "把银翼杀手电影放到银翼杀手文件夹下", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=338, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 2 + assert result.proposed_actions[0].tool == "drive.createFolder" + assert result.proposed_actions[0].input["name"] == "银翼杀手" + assert result.proposed_actions[1].tool == "drive.moveFile" + assert result.proposed_actions[1].input["fileId"] == "11" + assert result.proposed_actions[1].input["targetFolderId"] == "$step1.folderId" + + +@pytest.mark.asyncio +async def test_plan_runner_rewrites_write_summary_with_grounded_facts( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "创建银翼杀手文件夹，然后将 V字仇杀队 文件夹中的2部银翼杀手电影移入该文件夹。", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "银翼杀手"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "19", "targetFolderId": "$step1.folderId"}, + }, + { + "step": 3, + "tool": "drive.moveFile", + "input": {"fileId": "20", "targetFolderId": "$step1.folderId"}, + }, + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "把银翼杀手两部，移到银翼杀手文件夹里", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=348, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + db = DummyDb() + + class _Rows: + def __init__(self, rows: list[tuple[Any, Any]]) -> None: + self._rows = rows + + def all(self) -> list[tuple[Any, Any]]: + return list(self._rows) + + db.execute = AsyncMock(return_value=_Rows([(19, "银翼杀手1982.mp4"), (20, "银翼杀手2049.mp4")])) + + result = await runner.run(db=db, job=job) # type: ignore[arg-type] + + assert "V字仇杀队" not in result.summary + assert "创建“银翼杀手”文件夹" in result.summary + assert "2 个文件" in result.summary + + +@pytest.mark.asyncio +async def test_plan_runner_records_planning_evidence_from_read_tools( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + + class _FakeToolRouter: + async def dispatch(self, _call): # noqa: ANN001 + return { + "items": [ + {"fileId": "19", "name": "银翼杀手1982.mp4", "path": "/My Files/videos/银翼杀手1982.mp4"}, + {"fileId": "20", "name": "银翼杀手2049.mp4", "path": "/My Files/videos/银翼杀手2049.mp4"}, + {"fileId": "21", "name": "Blade Runner Trailer.mp4", "path": "/My Files/videos/Blade Runner Trailer.mp4"}, + {"fileId": "22", "name": "x1.mp4", "path": "/My Files/videos/x1.mp4"}, + {"fileId": "23", "name": "x2.mp4", "path": "/My Files/videos/x2.mp4"}, + {"fileId": "24", "name": "x3.mp4", "path": "/My Files/videos/x3.mp4"}, + {"fileId": "25", "name": "x4.mp4", "path": "/My Files/videos/x4.mp4"}, + ], + "totalItems": 7, + "query": "银翼杀手", + "folderId": "1", + "recursive": True, + "category": "video", + } + + monkeypatch.setattr(plan_module, "ToolRouter", lambda **kwargs: _FakeToolRouter()) + + async def _planner_with_read_tool(**kwargs): # noqa: ANN003 + tool_executor = kwargs["tool_executor"] + await tool_executor( + "drive.searchFiles", + {"folderId": "root", "query": "银翼杀手", "category": "video"}, + ) + return { + "summary": "search first", + "proposedActions": [ + { + "step": 1, + "tool": "drive.searchFiles", + "input": {"folderId": "root", "query": "银翼杀手", "category": "video"}, + } + ], + } + + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=_planner_with_read_tool), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "找出银翼杀手视频文件", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=349, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert result.planning_evidence is not None + assert len(result.planning_evidence) == 1 + evidence = result.planning_evidence[0] + assert evidence.step == 1 + assert evidence.tool == "drive.searchFiles" + assert evidence.input["query"] == "银翼杀手" + assert evidence.output_preview["totalItems"] == 7 + assert isinstance(evidence.output_preview.get("items"), list) + assert "..." in str(evidence.output_preview["items"][-1]) + + +@pytest.mark.asyncio +async def test_plan_runner_uses_planner_returned_read_only_candidates_when_ambiguous( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(plan_module, "_choose_skill", AsyncMock(return_value=None)) + monkeypatch.setattr( + plan_module, + "_collect_context_metadata", + AsyncMock(return_value={"scope": "currentFolder", "rootFolderId": "root", "files": [], "folders": []}), + ) + monkeypatch.setattr(plan_module, "_upsert_agent_plan", AsyncMock(return_value=None)) + planner = AsyncMock( + return_value={ + "summary": "ambiguous, return candidates", + "proposedActions": [ + { + "step": 1, + "tool": "drive.searchFiles", + "input": {"folderId": "root", "query": "银翼杀手", "category": "video"}, + } + ], + } + ) + runner = PlanRunner( + settings=settings(), + planner_client=SimpleNamespace(create_plan=planner), # type: ignore[arg-type] + ) + request = PlanAgentRequest.model_validate( + { + "input": "把银翼杀手电影放到银翼杀手文件夹下", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + } + ) + job = BackgroundJob( + job_id=339, + task_type="agent.plan", + status="running", + payload=request.model_dump(by_alias=True), + result={}, + requested_by=7, + scheduled_at=datetime.now(UTC), + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + result = await runner.run(db=DummyDb(), job=job) # type: ignore[arg-type] + + assert len(result.proposed_actions) == 1 + action = result.proposed_actions[0] + assert action.tool == "drive.searchFiles" + assert action.side_effect == "read" + assert "ambiguous" in result.summary + + +def test_normalize_actions_rejects_symbolic_placeholder_target_folder(): + with pytest.raises(ApiError) as exc: + plan_module._normalize_actions( + llm_payload={ + "summary": "organize movies", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "13", "targetFolderId": "newFolderId"}, + }, + ], + }, + allowed_tools=("drive.createFolder", "drive.moveFile"), + max_steps=10, + ) + + assert exc.value.status_code == 400 + assert "step 2" in exc.value.message + assert "targetFolderId" in exc.value.message + assert "newFolderId" in exc.value.message + + +def test_normalize_actions_allows_unbounded_steps_when_max_steps_none(): + actions = plan_module._normalize_actions( + llm_payload={ + "summary": "organize movies", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "13", "targetFolderId": "$step1.folderId"}, + }, + ], + }, + allowed_tools=("drive.createFolder", "drive.moveFile"), + max_steps=None, + ) + + assert len(actions) == 2 + + +def test_normalize_actions_accepts_previous_step_reference(): + actions = plan_module._normalize_actions( + llm_payload={ + "summary": "organize movies", + "proposedActions": [ + { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + { + "step": 2, + "tool": "drive.moveFile", + "input": {"fileId": "13", "targetFolderId": "$step1.folderId"}, + }, + ], + }, + allowed_tools=("drive.createFolder", "drive.moveFile"), + max_steps=10, + ) + + assert len(actions) == 2 + assert actions[1].input["targetFolderId"] == "$step1.folderId" + + +def test_normalize_actions_rejects_future_step_reference(): + with pytest.raises(ApiError) as exc: + plan_module._normalize_actions( + llm_payload={ + "summary": "organize movies", + "proposedActions": [ + { + "step": 3, + "tool": "drive.moveFile", + "input": {"fileId": "13", "targetFolderId": "$step4.folderId"}, + }, + { + "step": 4, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + }, + ], + }, + allowed_tools=("drive.createFolder", "drive.moveFile"), + max_steps=10, + ) + + assert exc.value.status_code == 400 + assert "future step 4" in exc.value.message + assert "$step4.folderId" in exc.value.message + + +def test_execute_reference_resolution_rejects_symbolic_placeholder(): + with pytest.raises(ApiError) as exc: + execute_module._resolve_references( + {"targetFolderId": "newFolderId"}, + step_outputs={}, + ) + + assert exc.value.status_code == 409 + assert "targetFolderId" in exc.value.message + assert "$stepN.field" in exc.value.message + + +@pytest.mark.asyncio +async def test_policy_guard_blocks_delete_without_confirmation(): + decision = await PolicyGuard().evaluate_tool_call( + tool_name="drive.deleteFile", + high_risk_confirmed=False, + ) + assert decision.allowed is False + assert classify_tool_risk("drive.deleteFolder") == "high" + + +@pytest.mark.asyncio +async def test_tool_router_dispatches_move_file(): + router = ToolRouter(db=DummyDb(), user_id=7) # type: ignore[arg-type] + router.file_service.move_file = AsyncMock( + return_value=SimpleNamespace( + model_dump=lambda **kwargs: {"fileId": "1", "targetFolderId": "2"} + ) + ) + + result = await router.dispatch( + ToolCall( + tool_name="drive.moveFile", + arguments={"fileId": "1", "targetFolderId": "2"}, + ) + ) + + assert result == {"fileId": "1", "targetFolderId": "2"} + router.file_service.move_file.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_tool_router_count_files_counts_recursive_videos(): + db = DummyDb() + db.scalar = AsyncMock(return_value=1) + db.scalars = AsyncMock( + side_effect=[ + [1, 2], + [ + SimpleNamespace( + file_id=10, + file_name="movie.mp4", + file_size=100, + mime_type="application/octet-stream", + file_ext="mp4", + folder_id=1, + created_at=None, + updated_at=None, + ), + SimpleNamespace( + file_id=11, + file_name="clip.mkv", + file_size=200, + mime_type="video/x-matroska", + file_ext="mkv", + folder_id=2, + created_at=None, + updated_at=None, + ), + SimpleNamespace( + file_id=12, + file_name="notes.txt", + file_size=10, + mime_type="text/plain", + file_ext="txt", + folder_id=1, + created_at=None, + updated_at=None, + ), + ], + ] + ) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall( + tool_name="drive.countFiles", + arguments={"folderId": "root", "recursive": True, "category": "video"}, + ) + ) + + assert result["totalItems"] == 2 + assert result["category"] == "video" + assert result["recursive"] is True + assert result["byMimeType"] == {"video/mp4": 1, "video/x-matroska": 1} + assert [item["name"] for item in result["sampleItems"]] == ["movie.mp4", "clip.mkv"] + executed_statement = str(db.scalars.await_args_list[-1].args[0]) + assert "file.status" in executed_statement + assert "file.is_latest" in executed_statement + + +@pytest.mark.asyncio +async def test_tool_router_count_files_filters_by_search_term(): + db = DummyDb() + db.scalar = AsyncMock(return_value=1) + db.scalars = AsyncMock( + side_effect=[ + [1], + [ + SimpleNamespace( + file_id=10, + file_name="银翼杀手.mp4", + file_size=100, + mime_type="video/mp4", + file_ext="mp4", + folder_id=1, + created_at=None, + updated_at=None, + ), + ], + ] + ) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall( + tool_name="drive.countFiles", + arguments={ + "folderId": "root", + "recursive": True, + "category": "video", + "search": "银翼杀手", + }, + ) + ) + + assert result["totalItems"] == 1 + assert result["search"] == "银翼杀手" + executed_statement = str(db.scalars.await_args_list[-1].args[0]) + assert "file_name" in executed_statement + + +@pytest.mark.asyncio +async def test_execute_runner_normalizes_tool_output_before_action_log(monkeypatch: pytest.MonkeyPatch): + started = datetime.now(UTC) + output_time = datetime.now(UTC) + job = BackgroundJob( + job_id=600, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "500", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.createFolder", + "input": {"parentFolderId": "root", "name": "Movies"}, + "sideEffect": "write", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + mock_plan_repo = SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace( + proposed_actions_json=[action], + ) + ) + ) + monkeypatch.setattr(execute_module, "AgentPlanRepository", lambda _db: mock_plan_repo) + + mock_work_sessions = SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ) + monkeypatch.setattr(execute_module, "AgentWorkSessionRepository", lambda _db: mock_work_sessions) + + captured_outputs: list[dict[str, object]] = [] + mock_action_logs = SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock( + side_effect=lambda **kwargs: captured_outputs.append(dict(kwargs)) or None + ), + ) + monkeypatch.setattr(execute_module, "AgentActionLogRepository", lambda _db: mock_action_logs) + + mock_router = SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "id": "9", + "createdAt": output_time, + "updatedAt": output_time, + } + ) + ) + monkeypatch.setattr(execute_module, "ToolRouter", lambda **kwargs: mock_router) + + result = await ExecuteRunner( + answer_client=SimpleNamespace(create_answer=AsyncMock(return_value="ok")) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.applied_actions == 1 + assert captured_outputs + success_call = next(item for item in captured_outputs if item.get("status") == "succeeded") + outputs_json = success_call["outputs_json"] + assert isinstance(outputs_json, dict) + assert isinstance(outputs_json["createdAt"], str) + assert outputs_json["createdAt"] == output_time.isoformat() + + +@pytest.mark.asyncio +async def test_execute_runner_propagates_answer_model_errors(monkeypatch: pytest.MonkeyPatch): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=610, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "510", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock(return_value=SimpleNamespace(proposed_actions_json=[action])) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "totalItems": 3, + "category": "video", + "recursive": True, + "folderId": "1", + "byMimeType": {"video/mp4": 3}, + } + ) + ), + ) + + with pytest.raises(ApiError) as exc: + await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock( + side_effect=ApiError(status_code=503, code=503, message="Agent LLM API key is not configured") + ) + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert exc.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_execute_runner_returns_count_files_answer(monkeypatch: pytest.MonkeyPatch): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=601, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "501", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + mock_plan_repo = SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace( + proposed_actions_json=[action], + ) + ) + ) + monkeypatch.setattr(execute_module, "AgentPlanRepository", lambda _db: mock_plan_repo) + + mock_work_sessions = SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ) + monkeypatch.setattr(execute_module, "AgentWorkSessionRepository", lambda _db: mock_work_sessions) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "totalItems": 3, + "category": "video", + "recursive": True, + "folderId": "1", + "byMimeType": {"video/mp4": 3}, + "sampleItems": [], + } + ) + ), + ) + + result = await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock(return_value="你上传了 3 部电影（按视频文件统计）。") + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.answer == "你上传了 3 部电影（按视频文件统计）。" + assert result.applied_actions == 1 + + +@pytest.mark.asyncio +async def test_execute_runner_returns_count_files_answer_with_search_term( + monkeypatch: pytest.MonkeyPatch, +): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=602, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "502", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.countFiles", + "input": { + "folderId": "root", + "recursive": True, + "category": "video", + "search": "银翼杀手", + }, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace(proposed_actions_json=[action]) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "totalItems": 2, + "category": "video", + "recursive": True, + "folderId": "1", + "search": "银翼杀手", + "byMimeType": {"video/mp4": 2}, + "sampleItems": [], + } + ) + ), + ) + + result = await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock(return_value="你上传了 2 部名称包含“银翼杀手”的电影（按视频文件统计）。") + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.answer == "你上传了 2 部名称包含“银翼杀手”的电影（按视频文件统计）。" + assert "只读操作" not in (result.answer or "") + + +@pytest.mark.asyncio +async def test_execute_runner_returns_count_files_answer_with_names_when_asked( + monkeypatch: pytest.MonkeyPatch, +): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=604, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "504", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "archive"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace( + input_text="我上传了多少压缩包，叫什么名字", + proposed_actions_json=[action], + ) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "totalItems": 2, + "category": "archive", + "recursive": True, + "folderId": "1", + "itemNames": ["photos.zip", "backup.7z"], + "itemNamesTruncated": False, + "byMimeType": {"application/zip": 1, "application/x-7z-compressed": 1}, + "sampleItems": [], + } + ) + ), + ) + + result = await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock(return_value="你上传了 2 个压缩包，名字是：photos.zip、backup.7z。") + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.answer == "你上传了 2 个压缩包，名字是：photos.zip、backup.7z。" + + +@pytest.mark.asyncio +async def test_execute_runner_lists_archive_names_for_read_only_archive_question( + monkeypatch: pytest.MonkeyPatch, +): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=603, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "503", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.listFolder", + "input": {"folderId": "root"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace( + input_text="我上传了哪些压缩包", + proposed_actions_json=[action], + ) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "items": [ + { + "itemType": "file", + "id": "1", + "name": "photos.zip", + "size": 10, + "mimeType": "application/zip", + }, + { + "itemType": "file", + "id": "2", + "name": "movie.mp4", + "size": 20, + "mimeType": "video/mp4", + }, + { + "itemType": "file", + "id": "3", + "name": "backup.7z", + "size": 30, + "mimeType": "application/octet-stream", + }, + ], + "pagination": { + "totalItems": 3, + "totalPages": 1, + "perPage": 200, + "currentPage": 1, + "hasPrev": False, + "hasNext": False, + }, + } + ) + ), + ) + + result = await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock(return_value="当前文件夹中的压缩包有 2 个：photos.zip、backup.7z。") + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.answer == "当前文件夹中的压缩包有 2 个：photos.zip、backup.7z。" + + +@pytest.mark.asyncio +async def test_execute_runner_returns_search_files_candidate_answer( + monkeypatch: pytest.MonkeyPatch, +): + started = datetime.now(UTC) + job = BackgroundJob( + job_id=605, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "505", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + action = { + "step": 1, + "tool": "drive.searchFiles", + "input": {"folderId": "root", "query": "银翼杀手", "category": "video"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + db = DummyDb() + db.refresh = AsyncMock() + + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace( + input_text="把银翼杀手电影放到银翼杀手文件夹下", + proposed_actions_json=[action], + ) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "items": [ + {"id": "11", "name": "银翼杀手1982.mp4", "mimeType": "video/mp4"}, + {"id": "12", "name": "银翼杀手2049.mp4", "mimeType": "video/mp4"}, + ], + "totalItems": 2, + "query": "银翼杀手", + "folderId": "1", + "recursive": False, + "category": "video", + } + ) + ), + ) + + result = await ExecuteRunner( + answer_client=SimpleNamespace( + create_answer=AsyncMock( + return_value="找到 2 个名称包含“银翼杀手”的文件：银翼杀手1982.mp4、银翼杀手2049.mp4。" + ) + ) # type: ignore[arg-type] + ).run(db=db, job=job) # type: ignore[arg-type] + + assert result.answer == "找到 2 个名称包含“银翼杀手”的文件：银翼杀手1982.mp4、银翼杀手2049.mp4。" + + +class _NeverStream: + async def next(self, *, timeout=None): # noqa: ANN001 + await asyncio.Future() + + async def aclose(self) -> None: + return None + + +class _CaptureBus: + def __init__(self) -> None: + self.events: list[AgentEventEnvelope] = [] + + async def publish(self, envelope: AgentEventEnvelope) -> None: + self.events.append(envelope) + + @asynccontextmanager + async def subscribe(self, *, job_id: int): # noqa: ARG002 + yield _NeverStream() + + +def _execute_job_for_controls() -> BackgroundJob: + started = datetime.now(UTC) + return BackgroundJob( + job_id=700, + task_type="agent.execute", + status="running", + payload={ + "planJobId": "500", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": started.isoformat(), + "highRiskConfirmed": False, + }, + }, + result={}, + requested_by=7, + scheduled_at=started, + created_at=started, + updated_at=started, + ) + + +def _patch_execute_dependencies( + monkeypatch: pytest.MonkeyPatch, + *, + controls: list[list[SimpleNamespace]], + dropped: list[int], +) -> None: + action = { + "step": 1, + "tool": "drive.countFiles", + "input": {"folderId": "root", "recursive": True, "category": "video"}, + "sideEffect": "read", + "riskLevel": "low", + "requiresConfirmation": False, + } + + monkeypatch.setattr( + execute_module, + "AgentPlanRepository", + lambda _db: SimpleNamespace( + get_for_execute_binding=AsyncMock( + return_value=SimpleNamespace(proposed_actions_json=[action]) + ) + ), + ) + monkeypatch.setattr( + execute_module, + "AgentWorkSessionRepository", + lambda _db: SimpleNamespace( + create_for_job=AsyncMock(return_value=None), + close_session=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "AgentActionLogRepository", + lambda _db: SimpleNamespace( + append_step=AsyncMock(return_value=None), + finish_step=AsyncMock(return_value=None), + ), + ) + monkeypatch.setattr( + execute_module, + "ToolRouter", + lambda **kwargs: SimpleNamespace( + dispatch=AsyncMock( + return_value={ + "totalItems": 1, + "category": "video", + "recursive": True, + "folderId": "1", + "byMimeType": {"video/mp4": 1}, + "sampleItems": [], + } + ) + ), + ) + + class FakeInboxRepository: + def __init__(self, _db) -> None: # noqa: ANN001 + return None + + async def list_pending_controls(self, *, job_id: int): # noqa: ARG002 + if controls: + return controls.pop(0) + return [] + + async def mark_dropped(self, *, inbox_message_id: int) -> None: + dropped.append(inbox_message_id) + + monkeypatch.setattr(execute_module, "AgentInboxMessageRepository", FakeInboxRepository) + + +@pytest.mark.asyncio +async def test_execute_runner_pauses_then_resumes_at_step_boundary( + monkeypatch: pytest.MonkeyPatch, +): + controls = [ + [SimpleNamespace(inbox_message_id=1, kind=AgentInboxKind.CONTROL_PAUSE)], + [SimpleNamespace(inbox_message_id=2, kind=AgentInboxKind.CONTROL_RESUME)], + ] + dropped: list[int] = [] + _patch_execute_dependencies(monkeypatch, controls=controls, dropped=dropped) + db = DummyDb() + db.refresh = AsyncMock() + bus = _CaptureBus() + + result = await ExecuteRunner( + event_bus=bus, + answer_client=SimpleNamespace(create_answer=AsyncMock(return_value="ok")), # type: ignore[arg-type] + ).run(db=db, job=_execute_job_for_controls()) # type: ignore[arg-type] + + assert result.applied_actions == 1 + assert dropped == [1, 2] + event_types = [event.event_type for event in bus.events] + assert "agent.paused" in event_types + assert "agent.resumed" in event_types + assert "tool.started" in event_types + assert "tool.succeeded" in event_types + + +@pytest.mark.asyncio +async def test_execute_runner_canceled_via_inbox_at_step_boundary( + monkeypatch: pytest.MonkeyPatch, +): + controls = [[SimpleNamespace(inbox_message_id=1, kind=AgentInboxKind.CONTROL_CANCEL)]] + dropped: list[int] = [] + _patch_execute_dependencies(monkeypatch, controls=controls, dropped=dropped) + db = DummyDb() + db.refresh = AsyncMock() + + with pytest.raises(AgentJobCanceled): + await ExecuteRunner( + event_bus=_CaptureBus(), + answer_client=SimpleNamespace(create_answer=AsyncMock(return_value="ok")), # type: ignore[arg-type] + ).run( # type: ignore[arg-type] + db=db, + job=_execute_job_for_controls(), + ) + + assert dropped == [1] + + +@pytest.mark.asyncio +async def test_execute_runner_publish_state_ignores_event_bus_failures(): + bus = SimpleNamespace(publish=AsyncMock(side_effect=RuntimeError("boom"))) + runner = ExecuteRunner(event_bus=bus) # type: ignore[arg-type] + + await runner._publish_state("agent.paused", job_id=1) + + assert bus.publish.await_count == 1 + + +@pytest.mark.asyncio +async def test_execute_runner_publish_tool_ignores_event_bus_failures(): + bus = SimpleNamespace(publish=AsyncMock(side_effect=RuntimeError("boom"))) + runner = ExecuteRunner(event_bus=bus) # type: ignore[arg-type] + + await runner._publish_tool( + "tool.started", + job_id=1, + step=1, + tool="drive.listFolder", + payload={"input": {"folderId": "root"}}, + ) + + assert bus.publish.await_count == 1 diff --git a/app/tests/test_agent_repositories.py b/app/tests/test_agent_repositories.py index 129d9c6..f440bda 100644 --- a/app/tests/test_agent_repositories.py +++ b/app/tests/test_agent_repositories.py @@ -28,6 +28,14 @@ def mappings(self): return list(self._rows) +class FakeScalarResult: + def __init__(self, value): + self._value = value + + def scalar(self): + return self._value + + class DummySession: def __init__(self) -> None: self.add = Mock() @@ -134,6 +142,86 @@ async def test_agent_skill_list_visible_maps_catalog_rows(): repo = AgentSkillRepository(session) items = await repo.list_visible(user_id=7) assert [item.skill_key for item in items] == ["builtin:organizeByType", "user:cleanup"] + statement = session.execute.await_args.args[0] + assert "owner_user_id = CAST(:user_id AS BIGINT)" in str(statement) + + +@pytest.mark.asyncio +async def test_agent_skill_list_visible_accepts_null_user_id(): + session = DummySession() + session.execute.return_value = FakeMappingResult([]) + + repo = AgentSkillRepository(session) + items = await repo.list_visible(user_id=None) + + assert items == [] + params = session.execute.await_args.args[1] + assert params["user_id"] is None + + +@pytest.mark.asyncio +async def test_agent_skill_search_visible_uses_bigint_cast_with_empty_query(): + session = DummySession() + session.execute.return_value = FakeMappingResult([]) + + repo = AgentSkillRepository(session) + items = await repo.search_visible(user_id=1, query_text="", limit=20) + + assert items == [] + statement = session.execute.await_args.args[0] + assert "owner_user_id = CAST(:user_id AS BIGINT)" in str(statement) + params = session.execute.await_args.args[1] + assert params["query_text"] == "" + + +@pytest.mark.asyncio +async def test_agent_skill_list_catalog_paginated_uses_bigint_cast_with_null_user(): + session = DummySession() + session.execute = AsyncMock( + side_effect=[ + FakeScalarResult(1), + FakeMappingResult( + [ + { + "skill_id": 1, + "skill_key": "builtin:organizeByType", + "name": "organizeByType", + "description": "Organize files by type", + "triggers_text": "organize, classify", + "tool_whitelist_json": ["drive.listFolder"], + "plan_template_json": {}, + "inputs_schema_json": {}, + "outputs_schema_json": {}, + "visibility": "global", + "owner_user_id": None, + "created_at": datetime.now(UTC), + "updated_at": datetime.now(UTC), + "search_text": "organize files by type", + } + ] + ), + ] + ) + + repo = AgentSkillRepository(session) + items, total_items = await repo.list_catalog_paginated( + user_id=None, + visibility="all", + query_text="", + page=1, + per_page=20, + ) + + assert total_items == 1 + assert [item.skill_key for item in items] == ["builtin:organizeByType"] + count_statement = session.execute.await_args_list[0].args[0] + list_statement = session.execute.await_args_list[1].args[0] + assert "owner_user_id = CAST(:user_id AS BIGINT)" in str(count_statement) + assert "owner_user_id = CAST(:user_id AS BIGINT)" in str(list_statement) + params = session.execute.await_args_list[1].args[1] + assert params["query_text"] == "" + assert params["offset"] == 0 + assert params["limit"] == 20 @pytest.mark.asyncio @@ -179,6 +267,21 @@ async def test_agent_mcp_list_visible_maps_catalog_rows(): repo = AgentMcpRepository(session) items = await repo.list_visible(user_id=7) assert [item.name for item in items] == ["web-search", "private-python"] + statement = session.execute.await_args.args[0] + assert "owner_user_id = CAST(:user_id AS BIGINT)" in str(statement) + + +@pytest.mark.asyncio +async def test_agent_mcp_list_visible_accepts_null_user_id(): + session = DummySession() + session.execute.return_value = FakeMappingResult([]) + + repo = AgentMcpRepository(session) + items = await repo.list_visible(user_id=None) + + assert items == [] + params = session.execute.await_args.args[1] + assert params["user_id"] is None @pytest.mark.asyncio @@ -256,6 +359,59 @@ async def test_agent_action_log_finish_refreshes_work_session_metrics(): session.execute.assert_awaited() +@pytest.mark.asyncio +async def test_agent_action_log_append_step_normalizes_datetime_inputs(): + session = DummySession() + session.scalar = AsyncMock(return_value=88) + captured: list[SimpleNamespace] = [] + + def add(entry: SimpleNamespace) -> None: + captured.append(entry) + + session.add = Mock(side_effect=add) + repo = AgentActionLogRepository(session) + now = datetime.now(UTC) + + await repo.append_step( + job_id=101, + step_no=1, + tool_name="drive.createFolder", + inputs_json={"createdAt": now}, + status="running", + ) + + assert captured + saved = captured[0] + assert isinstance(saved.inputs_json["createdAt"], str) + assert saved.inputs_json["createdAt"] == now.isoformat() + + +@pytest.mark.asyncio +async def test_agent_action_log_finish_step_normalizes_datetime_outputs(): + session = DummySession() + entry = SimpleNamespace( + outputs_json={}, + status="running", + duration_ms=None, + error_message=None, + finished_at=None, + ) + session.scalar = AsyncMock(side_effect=[entry, 88]) + repo = AgentActionLogRepository(session) + now = datetime.now(UTC) + + await repo.finish_step( + job_id=101, + step_no=2, + outputs_json={"updatedAt": now}, + status="succeeded", + duration_ms=5, + ) + + assert isinstance(entry.outputs_json["updatedAt"], str) + assert entry.outputs_json["updatedAt"] == now.isoformat() + + @pytest.mark.asyncio async def test_agent_work_session_refresh_metrics_executes_db_function(): session = DummySession() diff --git a/app/tests/test_agent_routes.py b/app/tests/test_agent_routes.py new file mode 100644 index 0000000..68d225e --- /dev/null +++ b/app/tests/test_agent_routes.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +from contextlib import asynccontextmanager +from datetime import UTC, datetime + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.agents.harness.event_bus import AgentEventEnvelope, InMemoryAgentEventBus +from fileflash.core.deps import ( + get_agent_event_bus, + get_agent_execute_service, + get_agent_plan_service, + get_current_user, +) +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.db.deps import get_db +from fileflash.models import AgentActionLog, AgentInboxMessage, BackgroundJob +from fileflash.models.enums import AgentInboxRole +from fileflash.models.tables_identity import User +from fileflash.routers.agent import router +from fileflash.schemas.agent import ExecuteAgentResponse, PlanAgentResponse + + +class StubPlanService: + async def enqueue_plan(self, *, user_id, payload): # noqa: ANN001 + return PlanAgentResponse(job_id="10", status="pending", task_type="agent.plan") + + +class StubExecuteService: + async def enqueue_execute(self, *, user_id, payload): # noqa: ANN001 + return ExecuteAgentResponse(job_id="11", status="pending", task_type="agent.execute") + + +class StubDb: + def __init__(self) -> None: + now = datetime.now(UTC) + self.job = BackgroundJob( + job_id=12, + task_type="agent.execute", + status="pending", + payload={}, + result={}, + requested_by=7, + scheduled_at=now, + created_at=now, + updated_at=now, + ) + self.messages: list[AgentInboxMessage] = [] + self._next_inbox_id = 1 + + async def scalar(self, _query): # noqa: ANN001 + return self.job + + async def scalars(self, _query): # noqa: ANN001 + return [] + + def add(self, msg: AgentInboxMessage) -> None: + msg.inbox_message_id = self._next_inbox_id + self._next_inbox_id += 1 + self.messages.append(msg) + + async def flush(self) -> None: + return None + + async def commit(self) -> None: + return None + + async def refresh(self, _job: BackgroundJob) -> None: + return None + + async def get(self, _model, _id: int): # noqa: ANN001 + return None + + +class RunningJobDb(StubDb): + def __init__(self) -> None: + super().__init__() + self.job.status = "running" + + +class EventsDb(StubDb): + def __init__(self) -> None: + super().__init__() + now = datetime.now(UTC) + self.job.status = "succeeded" + self.job.result = { + "planJobId": "10", + "executeJobId": "12", + "summary": "done", + "answer": "你上传了 2 部名称包含“银翼杀手”的电影（按视频文件统计）。", + "appliedActions": 1, + "skippedActions": 0, + "warnings": [], + "finishedAt": now.isoformat(), + } + self.job.finished_at = now + self.job.updated_at = now + self.action_log = AgentActionLog( + action_log_id=1, + job_id=12, + step_no=1, + tool_name="drive.countFiles", + inputs_json={"folderId": "root", "category": "video", "search": "银翼杀手"}, + outputs_json={"totalItems": 2, "category": "video", "search": "银翼杀手"}, + status="succeeded", + duration_ms=12, + started_at=now, + finished_at=now, + ) + + async def scalars(self, _query): # noqa: ANN001 + return [self.action_log] + + +def _user() -> User: + return User(user_id=7, username="u7", email="u7@example.com", password_hash="x") + + +def _client() -> TestClient: + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) + app.dependency_overrides[get_current_user] = _user + app.dependency_overrides[get_agent_plan_service] = lambda: StubPlanService() + app.dependency_overrides[get_agent_execute_service] = lambda: StubExecuteService() + app.dependency_overrides[get_db] = lambda: StubDb() + return TestClient(app) + + +def _client_with_running_job() -> TestClient: + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) + app.dependency_overrides[get_current_user] = _user + app.dependency_overrides[get_agent_plan_service] = lambda: StubPlanService() + app.dependency_overrides[get_agent_execute_service] = lambda: StubExecuteService() + app.dependency_overrides[get_db] = lambda: RunningJobDb() + return TestClient(app) + + +def _client_with_events() -> TestClient: + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) + app.dependency_overrides[get_current_user] = _user + app.dependency_overrides[get_agent_plan_service] = lambda: StubPlanService() + app.dependency_overrides[get_agent_execute_service] = lambda: StubExecuteService() + app.dependency_overrides[get_db] = lambda: EventsDb() + return TestClient(app) + + +def test_plan_route_returns_response_shell(): + response = _client().post( + "/api/v1/agent/plan", + json={ + "input": "organize", + "context": { + "rootFolderId": "root", + "selectedFileIds": [], + "selectedFolderIds": [], + "currentPath": "/My Files", + }, + "executionPolicy": "confirm", + "dataPolicy": { + "allowFileContent": False, + "maxReadBytes": 1024, + "allowedMimeTypes": ["*/*"], + }, + "hints": {"preferSkillId": None, "maxSteps": 12, "budgetTokens": 8000}, + }, + ) + + assert response.status_code == 200 + body = response.json() + assert body["success"] is True + assert body["data"]["jobId"] == "10" + assert body["data"]["taskType"] == "agent.plan" + + +def test_execute_route_returns_response_shell(): + response = _client().post( + "/api/v1/agent/execute", + json={ + "planJobId": "10", + "planHash": "sha256:test", + "approval": { + "confirmedBy": "7", + "confirmedAt": datetime.now(UTC).isoformat(), + "highRiskConfirmed": True, + "highRiskConfirmedAt": datetime.now(UTC).isoformat(), + }, + }, + ) + + assert response.status_code == 200 + body = response.json() + assert body["success"] is True + assert body["data"]["jobId"] == "11" + assert body["data"]["taskType"] == "agent.execute" + + +def test_post_message_control_pause_returns_response_shell(): + bus = InMemoryAgentEventBus() + db = StubDb() + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) + app.dependency_overrides[get_current_user] = _user + app.dependency_overrides[get_db] = lambda: db + app.dependency_overrides[get_agent_event_bus] = lambda: bus + client = TestClient(app) + + response = client.post("/api/v1/agent/jobs/12/messages", json={"kind": "control.pause"}) + + assert response.status_code == 200 + body = response.json() + assert body["success"] is True + assert body["data"]["kind"] == "control.pause" + assert body["data"]["inboxMessageId"] == "1" + assert db.messages[0].role == AgentInboxRole.USER + + +def test_job_events_route_streams_tool_and_final_answer_events(): + response = _client_with_events().get("/api/v1/agent/jobs/12/events") + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + body = response.text + assert "event: tool.started" in body + assert "event: tool.succeeded" in body + assert "event: job.succeeded" in body + assert "正在读取名称包含" in body + assert "银翼杀手" in body + assert "answer" in body + + +def test_job_events_route_streams_event_bus_events_after_initial_replay(): + now = datetime.now(UTC) + events = [ + AgentEventEnvelope( + job_id=12, + event_type="agent.progress", + payload={"step": 1, "total": 3, "message": "halfway"}, + emitted_at=now, + ), + AgentEventEnvelope( + job_id=12, + event_type="job.succeeded", + payload={"status": "succeeded"}, + emitted_at=now, + ), + ] + + class StaticStream: + async def next(self, *, timeout=None): # noqa: ANN001 + if not events: + raise TimeoutError + return events.pop(0) + + async def aclose(self) -> None: + return None + + class StaticBus: + async def publish(self, envelope): # noqa: ANN001 + return None + + @asynccontextmanager + async def subscribe(self, *, job_id: int): # noqa: ARG002 + yield StaticStream() + + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) + app.dependency_overrides[get_current_user] = _user + app.dependency_overrides[get_db] = lambda: RunningJobDb() + app.dependency_overrides[get_agent_event_bus] = lambda: StaticBus() + client = TestClient(app) + + response = client.get("/api/v1/agent/jobs/12/events") + + assert response.status_code == 200 + assert "event: agent.progress" in response.text + assert "event: job.succeeded" in response.text diff --git a/app/tests/test_agent_skill_service.py b/app/tests/test_agent_skill_service.py index 79d8ffe..bc4d13f 100644 --- a/app/tests/test_agent_skill_service.py +++ b/app/tests/test_agent_skill_service.py @@ -76,6 +76,29 @@ async def test_update_custom_skill_requires_owner_private(): session.commit.assert_not_awaited() +@pytest.mark.asyncio +async def test_create_custom_skill_rejects_unknown_tool(): + session = DummySession() + session.scalar.return_value = None + + repo = AgentSkillRepository(session) + service = SkillService(db=session, skills=repo) + + with pytest.raises(ApiError) as exc: + await service.create_custom_skill( + user_id=7, + payload=CreateAgentSkillRequest( + name="Unsafe", + description="bad tool", + tool_whitelist=["drive.listFolder", "files.list"], + ), + ) + + assert exc.value.status_code == 422 + assert exc.value.data == {"unknownTools": ["files.list"]} + session.commit.assert_not_awaited() + + @pytest.mark.asyncio async def test_delete_custom_skill_requires_owner_private(): session = DummySession() @@ -91,6 +114,33 @@ async def test_delete_custom_skill_requires_owner_private(): session.commit.assert_not_awaited() +@pytest.mark.asyncio +async def test_import_rejects_unknown_tool(): + session = DummySession() + session.scalars.return_value = [] + + repo = AgentSkillRepository(session) + service = SkillService(db=session, skills=repo) + + payload = ImportAgentSkillsRequest( + items=[ + ImportAgentSkillItem( + skill_key="builtin:bad", + name="bad", + description="bad", + tool_whitelist=["drive.missing"], + ) + ], + ) + + with pytest.raises(ApiError) as exc: + await service.import_global_skills(payload=payload) + + assert exc.value.status_code == 422 + assert exc.value.data == {"unknownTools": ["drive.missing"]} + session.commit.assert_not_awaited() + + @pytest.mark.asyncio async def test_import_insert_only_conflict_raises_409(): session = DummySession() diff --git a/app/tests/test_agent_tools.py b/app/tests/test_agent_tools.py new file mode 100644 index 0000000..36287ec --- /dev/null +++ b/app/tests/test_agent_tools.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from fileflash.agents.harness.router import ToolCall, ToolRouter +from fileflash.agents.harness.tool_registry import REGISTRY, ToolRegistry, ToolSpec +from fileflash.models.enums import FileStatus + + +async def _noop_handler(_ctx, _args): # noqa: ANN001 + return {"ok": True} + + +def test_tool_registry_registers_and_maps_provider_names(): + registry = ToolRegistry() + registry.register( + ToolSpec( + name="drive.testTool", + description="test", + input_schema={"type": "object"}, + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_noop_handler, + ) + ) + + assert registry.all_names() == ("drive.testTool",) + assert registry.get("drive.testTool").anthropic_name == "drive_test_tool" + assert registry.get_by_provider_name("drive_test_tool").name == "drive.testTool" + assert registry.anthropic_tools_for(["drive.testTool"])[0]["internalName"] == "drive.testTool" + + +def test_tool_registry_rejects_duplicate_names(): + registry = ToolRegistry() + spec = ToolSpec( + name="drive.testTool", + description="test", + input_schema={"type": "object"}, + side_effect="read", + risk_level="low", + requires_confirmation=False, + handler=_noop_handler, + ) + registry.register(spec) + + with pytest.raises(ValueError): + registry.register(spec) + + +def test_builtin_registry_contains_new_query_tools(): + names = set(REGISTRY.all_names()) + + assert { + "drive.searchFiles", + "drive.getFileInfo", + "drive.listRecent", + "drive.statsByCategory", + "drive.findDuplicates", + }.issubset(names) + assert REGISTRY.get("drive.deleteFile").risk_level == "high" + + +class DummyDb: + def __init__(self) -> None: + self.scalar = AsyncMock(return_value=1) + self.scalars = AsyncMock(return_value=[]) + self.execute = AsyncMock() + self.get = AsyncMock() + + +@pytest.mark.asyncio +async def test_tool_router_dispatches_new_search_files_tool(): + db = DummyDb() + db.scalar = AsyncMock( + side_effect=[ + 1, + SimpleNamespace(folder_name="My Files", parent_folder_id=None), + ] + ) + db.scalars = AsyncMock( + side_effect=[ + [1], + [ + SimpleNamespace( + file_id=10, + file_name="movie.mp4", + file_size=100, + mime_type="video/mp4", + file_ext="mp4", + folder_id=1, + storage_object_id=20, + status=FileStatus.ACTIVE, + is_latest=True, + created_at=None, + updated_at=None, + ) + ], + ] + ) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall( + tool_name="drive.searchFiles", + arguments={"folderId": "root", "query": "movie", "category": "video"}, + ) + ) + + assert result["totalItems"] == 1 + assert result["items"][0]["name"] == "movie.mp4" + + +@pytest.mark.asyncio +async def test_tool_router_dispatches_stats_by_category_tool(): + db = DummyDb() + db.scalars = AsyncMock( + side_effect=[ + [1], + [ + SimpleNamespace( + file_id=10, + file_name="movie.mp4", + file_size=100, + mime_type="video/mp4", + file_ext="mp4", + folder_id=1, + storage_object_id=20, + created_at=None, + updated_at=None, + ), + SimpleNamespace( + file_id=11, + file_name="notes.txt", + file_size=10, + mime_type="text/plain", + file_ext="txt", + folder_id=1, + storage_object_id=21, + created_at=None, + updated_at=None, + ), + ], + ] + ) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall(tool_name="drive.statsByCategory", arguments={"folderId": "root"}) + ) + + assert result["video"] == 1 + assert result["document"] == 1 + assert result["totalSize"] == 110 + + +@pytest.mark.asyncio +async def test_tool_router_count_files_accepts_anime_alias_and_returns_item_names(): + db = DummyDb() + db.scalar = AsyncMock(return_value=1) + db.scalars = AsyncMock( + side_effect=[ + [1], + [ + SimpleNamespace( + file_id=10, + file_name="银河动漫番剧.mp4", + file_size=100, + mime_type="video/mp4", + file_ext="mp4", + folder_id=1, + storage_object_id=20, + created_at=None, + updated_at=None, + ), + SimpleNamespace( + file_id=11, + file_name="notes.txt", + file_size=10, + mime_type="text/plain", + file_ext="txt", + folder_id=1, + storage_object_id=21, + created_at=None, + updated_at=None, + ), + ], + ] + ) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall( + tool_name="drive.countFiles", + arguments={"folderId": "root", "recursive": True, "category": "anime"}, + ) + ) + + assert result["category"] == "video" + assert result["totalItems"] == 1 + assert result["itemNames"] == ["银河动漫番剧.mp4"] + assert result["itemNamesTruncated"] is False + + +@pytest.mark.asyncio +async def test_tool_router_count_files_truncates_item_names_at_limit(): + db = DummyDb() + db.scalar = AsyncMock(return_value=1) + rows = [ + SimpleNamespace( + file_id=100 + index, + file_name=f"video-{index:02d}.mp4", + file_size=100, + mime_type="video/mp4", + file_ext="mp4", + folder_id=1, + storage_object_id=200 + index, + created_at=None, + updated_at=None, + ) + for index in range(13) + ] + db.scalars = AsyncMock(side_effect=[[1], rows]) + router = ToolRouter(db=db, user_id=7) # type: ignore[arg-type] + + result = await router.dispatch( + ToolCall( + tool_name="drive.countFiles", + arguments={"folderId": "root", "recursive": True, "category": "video"}, + ) + ) + + assert result["totalItems"] == 13 + assert len(result["itemNames"]) == 12 + assert result["itemNames"][0] == "video-00.mp4" + assert result["itemNames"][-1] == "video-11.mp4" + assert result["itemNamesTruncated"] is True diff --git a/app/tests/test_agent_worker.py b/app/tests/test_agent_worker.py new file mode 100644 index 0000000..1275bfe --- /dev/null +++ b/app/tests/test_agent_worker.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from fileflash.agents.harness.event_bus import AgentEventEnvelope +from fileflash.agents.worker import AgentWorkerConsumer +from fileflash.models import BackgroundJob + + +class _AsyncContextManager: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, exc_type, exc, tb): + return None + + +class DummySession: + def __init__(self, job: BackgroundJob): + self._job = job + + def begin(self): + return _AsyncContextManager(SimpleNamespace()) + + async def scalar(self, _query): # noqa: ANN001 + return self._job + + +class CaptureBus: + def __init__(self) -> None: + self.events: list[AgentEventEnvelope] = [] + + async def publish(self, envelope: AgentEventEnvelope) -> None: + self.events.append(envelope) + + +class FailingBus: + async def publish(self, envelope: AgentEventEnvelope) -> None: # noqa: ARG002 + raise RuntimeError("publish unavailable") + + +def _job(*, status: str, cancel_requested_at: datetime | None) -> BackgroundJob: + now = datetime.now(UTC) + return BackgroundJob( + job_id=65, + task_type="agent.plan", + status=status, + payload={}, + result={}, + requested_by=7, + cancel_requested_at=cancel_requested_at, + scheduled_at=now, + created_at=now, + updated_at=now, + ) + + +@pytest.mark.asyncio +async def test_mark_succeeded_does_not_override_canceled_job(monkeypatch: pytest.MonkeyPatch): + canceled_at = datetime.now(UTC) + job = _job(status="canceled", cancel_requested_at=canceled_at) + session = DummySession(job) + consumer = AgentWorkerConsumer( + queue=SimpleNamespace(), + session_factory=lambda: _AsyncContextManager(session), # type: ignore[arg-type] + ) + monkeypatch.setattr("fileflash.agents.worker.apply_local_lock_timeout", AsyncMock(return_value=None)) + + await consumer._mark_succeeded(job_id=65, result={"summary": "ok"}, phase="completed") + + assert job.status == "canceled" + assert job.cancel_requested_at == canceled_at + assert job.result == {} + + +@pytest.mark.asyncio +async def test_mark_failed_does_not_override_job_with_cancel_request(monkeypatch: pytest.MonkeyPatch): + canceled_at = datetime.now(UTC) + job = _job(status="running", cancel_requested_at=canceled_at) + session = DummySession(job) + consumer = AgentWorkerConsumer( + queue=SimpleNamespace(), + session_factory=lambda: _AsyncContextManager(session), # type: ignore[arg-type] + ) + monkeypatch.setattr("fileflash.agents.worker.apply_local_lock_timeout", AsyncMock(return_value=None)) + + await consumer._mark_failed(job_id=65, error=RuntimeError("boom")) + + assert job.status == "running" + assert job.cancel_requested_at == canceled_at + assert job.error_message is None + + +@pytest.mark.asyncio +async def test_mark_succeeded_publishes_terminal_event(monkeypatch: pytest.MonkeyPatch): + job = _job(status="running", cancel_requested_at=None) + session = DummySession(job) + bus = CaptureBus() + consumer = AgentWorkerConsumer( + queue=SimpleNamespace(), + session_factory=lambda: _AsyncContextManager(session), # type: ignore[arg-type] + event_bus=bus, + ) + monkeypatch.setattr("fileflash.agents.worker.apply_local_lock_timeout", AsyncMock(return_value=None)) + + finished_at = datetime.now(UTC).replace(microsecond=0) + await consumer._mark_succeeded( + job_id=65, + result={"summary": "ok", "finishedAt": finished_at}, + phase="completed", + ) + + assert [event.event_type for event in bus.events] == ["job.succeeded"] + assert bus.events[0].payload["status"] == "succeeded" + assert bus.events[0].payload["data"]["result"]["finishedAt"] == finished_at.isoformat() + assert job.result["finishedAt"] == finished_at.isoformat() + + +@pytest.mark.asyncio +async def test_mark_succeeded_ignores_publish_failures(monkeypatch: pytest.MonkeyPatch): + job = _job(status="running", cancel_requested_at=None) + session = DummySession(job) + consumer = AgentWorkerConsumer( + queue=SimpleNamespace(), + session_factory=lambda: _AsyncContextManager(session), # type: ignore[arg-type] + event_bus=FailingBus(), + ) + monkeypatch.setattr("fileflash.agents.worker.apply_local_lock_timeout", AsyncMock(return_value=None)) + + await consumer._mark_succeeded(job_id=65, result={"summary": "ok"}, phase="completed") + + assert job.status == "succeeded" + assert job.result["summary"] == "ok" diff --git a/app/tests/test_background_jobs_service.py b/app/tests/test_background_jobs_service.py index 05c64be..22dcda0 100644 --- a/app/tests/test_background_jobs_service.py +++ b/app/tests/test_background_jobs_service.py @@ -145,3 +145,27 @@ async def test_enqueue_transcode_job_uses_object_storage_payload(): assert payload["outputObjectKey"].endswith(".mp4") assert payload["fileId"] == 999 assert payload["requestedBy"] == 1 + + +@pytest.mark.asyncio +async def test_enqueue_normalizes_payload_for_jsonb_storage(): + session = DummySession() + queue = SimpleNamespace(publish=AsyncMock(return_value="1-0")) + service = BackgroundJobService(queue_publisher=queue) + confirmed_at = datetime(2026, 5, 25, 10, 0, tzinfo=UTC) + + job = await service.enqueue( + session, # type: ignore[arg-type] + task_type="agent.execute", + payload={ + "planJobId": "46", + "approval": {"confirmedAt": confirmed_at}, + }, + requested_by=7, + max_attempts=1, + priority=100, + agent_phase="executing", + ) + + assert isinstance(job.payload["approval"]["confirmedAt"], str) + assert job.payload["approval"]["confirmedAt"] == confirmed_at.isoformat() diff --git a/app/tests/test_dev_seed.py b/app/tests/test_dev_seed.py index 2ac47c2..6681920 100644 --- a/app/tests/test_dev_seed.py +++ b/app/tests/test_dev_seed.py @@ -5,8 +5,9 @@ import pytest from fileflash.core.settings import Settings +from fileflash.models.enums import UserRole from fileflash.models.tables_identity import User -from fileflash.services.dev_seed import DevAccountSeeder, initialize_dev_accounts +from fileflash.services.dev_seed import DevAccountSeeder, DevSeedSummary, initialize_dev_accounts class DummySeedSession: @@ -38,7 +39,10 @@ async def _find_existing_user(self, *, spec): # type: ignore[override] for obj in self.db.added: if not isinstance(obj, User): continue - if obj.username.lower() == spec.username.lower() or obj.email.lower() == spec.email.lower(): + if ( + obj.username.lower() == spec.username.lower() + or obj.email.lower() == spec.email.lower() + ): return obj return None @@ -59,7 +63,7 @@ def make_settings(**overrides: object) -> Settings: "JWT_SECRET_KEY": "unit-test-secret-key-1234567890abcd", } payload.update(overrides) - return Settings(**payload) + return Settings(_env_file=None, **payload) @pytest.mark.asyncio @@ -95,12 +99,14 @@ async def test_dev_account_seeder_is_idempotent_and_supports_password_reset(): @pytest.mark.asyncio -async def test_initialize_dev_accounts_skips_auto_run_in_production(monkeypatch: pytest.MonkeyPatch): +async def test_initialize_dev_accounts_skips_auto_run_outside_dev_and_prod( + monkeypatch: pytest.MonkeyPatch, +): guard = AsyncMock(side_effect=AssertionError("SessionLocal should not be called")) monkeypatch.setattr("fileflash.services.dev_seed.SessionLocal", guard) result = await initialize_dev_accounts( - settings=make_settings(APP_ENV="production"), + settings=make_settings(APP_ENV="staging"), auto_run=True, reset_password=False, ) @@ -108,3 +114,59 @@ async def test_initialize_dev_accounts_skips_auto_run_in_production(monkeypatch: assert result is False guard.assert_not_called() + +@pytest.mark.asyncio +async def test_initialize_dev_accounts_uses_env_admin_in_production( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + class SessionLocalStub: + calls = 0 + + def __call__(self): + self.calls += 1 + return self + + async def __aenter__(self): + return object() + + async def __aexit__(self, exc_type, exc, traceback): + return None + + class SeederStub: + def __init__(self, *, db: object, accounts: tuple[object, ...]) -> None: + captured["db"] = db + captured["accounts"] = accounts + + async def seed(self, *, reset_password: bool = False) -> DevSeedSummary: + captured["reset_password"] = reset_password + return DevSeedSummary(created_users=1) + + session_local = SessionLocalStub() + monkeypatch.setattr("fileflash.services.dev_seed.SessionLocal", session_local) + monkeypatch.setattr("fileflash.services.dev_seed.DevAccountSeeder", SeederStub) + + result = await initialize_dev_accounts( + settings=make_settings( + APP_ENV="production", + DEFAULT_ADMIN_USERNAME="root-admin", + DEFAULT_ADMIN_EMAIL="root-admin@example.com", + DEFAULT_ADMIN_PASSWORD="p" * 32, + ), + auto_run=True, + reset_password=False, + ) + + assert result is True + assert session_local.calls == 1 + accounts = captured["accounts"] + assert isinstance(accounts, tuple) + assert len(accounts) == 1 + account = accounts[0] + assert account.username == "root-admin" + assert account.email == "root-admin@example.com" + assert account.password == "p" * 32 + assert account.role == UserRole.ADMIN + assert captured["reset_password"] is False + diff --git a/app/tests/test_file_download_recycle_service.py b/app/tests/test_file_download_recycle_service.py index cf5ec2c..e0b6031 100644 --- a/app/tests/test_file_download_recycle_service.py +++ b/app/tests/test_file_download_recycle_service.py @@ -1,15 +1,29 @@ from __future__ import annotations +import tempfile from datetime import UTC, datetime +from types import SimpleNamespace from unittest.mock import AsyncMock import pytest - -from fileflash.core.errors import ApiError -from fileflash.models.enums import FileStatus, FolderStatus, FolderType, UploadStatus +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_current_user, get_download_rate_limit_service, get_file_service +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.models.enums import ( + FileStatus, + FolderStatus, + FolderType, + UploadStatus, + UserRole, + UserStatus, +) +from fileflash.models.tables_identity import User from fileflash.models.tables_storage import File, FileMediaMetadata, Folder, StorageObject -from fileflash.schemas.file import BatchFilesRequest -from fileflash.services.file import FileService +from fileflash.routers.files import router as files_router +from fileflash.schemas.file import BatchDownloadRequest, BatchFilesRequest +from fileflash.services.file import BatchDownloadPlan, DownloadStreamResult, FileService class DummyStorage: @@ -41,6 +55,14 @@ def __init__(self) -> None: self.delete = AsyncMock() +class ResultRows: + def __init__(self, rows) -> None: # noqa: ANN001 + self._rows = rows + + def all(self): # noqa: ANN201 + return self._rows + + def make_file_row(*, file_id: int = 1, file_name: str = "demo.txt", folder_id: int = 10) -> File: return File( file_id=file_id, @@ -73,6 +95,22 @@ def make_folder_row(*, folder_id: int = 10, folder_name: str = "Docs") -> Folder ) +def make_user(*, role: UserRole = UserRole.USER) -> User: + return User( + user_id=1, + username="alice", + email="alice@example.com", + password_hash="x", + role=role, + status=UserStatus.ACTIVE, + email_verified=True, + storage_limit=1024, + storage_used=0, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + + @pytest.mark.asyncio async def test_get_download_stream_supports_single_range(monkeypatch: pytest.MonkeyPatch): session = DummySession() @@ -291,9 +329,9 @@ async def test_get_preview_stream_prefers_transcoded_object_when_ready(monkeypat storage = DummyStorage() service = FileService(db=session, storage=storage) - file_row = make_file_row(file_id=11, file_name="preview.mp4") - file_row.file_ext = "mp4" - file_row.mime_type = "video/mp4" + file_row = make_file_row(file_id=11, file_name="preview.mkv") + file_row.file_ext = "mkv" + file_row.mime_type = "video/x-matroska" file_row.storage_object_id = 101 source_object = StorageObject( object_id=101, @@ -330,6 +368,7 @@ async def test_get_preview_stream_prefers_transcoded_object_when_ready(monkeypat result = await service.get_preview_stream(user_id=1, file_id="11", range_header=None) assert result.status_code == 200 + assert result.content_type == "video/mp4" assert result.headers["Content-Length"] == "128" @@ -374,6 +413,145 @@ async def test_get_preview_stream_falls_back_to_source_when_transcoded_missing(m assert result.headers["Content-Length"] == "512" +@pytest.mark.asyncio +async def test_batch_download_plan_estimates_source_file_size() -> None: + session = DummySession() + storage = DummyStorage() + service = FileService(db=session, storage=storage) + file_row = make_file_row(file_id=7, file_name="archive.bin") + file_row.storage_object_id = 99 + storage_object = StorageObject( + object_id=99, + bucket_name="fileflash", + object_key="objects/u1/archive", + object_size=512, + upload_status=UploadStatus.ACTIVE, + content_type="application/octet-stream", + ) + session.scalars = AsyncMock(return_value=[file_row]) + session.execute = AsyncMock(return_value=ResultRows([(file_row, storage_object)])) + + plan = await service.create_batch_download_plan( + user_id=1, + payload=BatchDownloadRequest(fileIds=["7"]), + ) + + assert plan.estimated_bytes == 512 + assert plan.files[0][2] == "archive.bin" + + +class StubDownloadLimiter: + def __init__(self, *, deny: bool = False) -> None: + self.deny = deny + self.calls: list[tuple[str, int]] = [] + + async def enforce_user(self, *, user: User, bytes_count: int) -> None: + self.calls.append((f"user:{user.user_id}", bytes_count)) + if self.deny and user.role != UserRole.ADMIN: + raise ApiError(status_code=429, code=429, message="Download rate limit exceeded") + + async def enforce_user_id(self, *, user_id: int, bytes_count: int) -> None: + self.calls.append((f"user:{user_id}", bytes_count)) + if self.deny: + raise ApiError(status_code=429, code=429, message="Download rate limit exceeded") + + +class StubFileRouteService: + async def get_download_stream( + self, + *, + user_id: int, # noqa: ARG002 + file_id: str, # noqa: ARG002 + range_header: str | None, + ) -> DownloadStreamResult: + async def _stream(): + yield b"0123456789" + + headers = {"Content-Length": "4" if range_header else "10", "Accept-Ranges": "bytes"} + if range_header: + headers["Content-Range"] = "bytes 0-3/10" + return DownloadStreamResult( + stream=_stream(), + filename="demo.txt", + content_type="text/plain", + status_code=206 if range_header else 200, + headers=headers, + ) + + async def get_preview_stream(self, **kwargs) -> DownloadStreamResult: # noqa: ANN003 + return await self.get_download_stream(**kwargs) + + async def create_batch_download_plan( + self, + *, + user_id: int, # noqa: ARG002 + payload: BatchDownloadRequest, # noqa: ARG002 + ) -> BatchDownloadPlan: + return SimpleNamespace(estimated_bytes=10, files=[object()]) # type: ignore[return-value] + + async def create_batch_download_archive_from_plan(self, *, plan: BatchDownloadPlan): # noqa: ANN201, ARG002 + tmp = tempfile.NamedTemporaryFile(prefix="fileflash-test-", suffix=".zip", delete=False) + tmp.write(b"zip") + tmp.close() + return tmp.name, "test.zip" + + +def _files_client(*, role: UserRole, limiter: StubDownloadLimiter) -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(files_router, prefix="/api/v1") + app.dependency_overrides[get_current_user] = lambda: make_user(role=role) + app.dependency_overrides[get_file_service] = lambda: StubFileRouteService() + app.dependency_overrides[get_download_rate_limit_service] = lambda: limiter + return TestClient(app) + + +def test_download_route_returns_429_when_limiter_rejects_user() -> None: + limiter = StubDownloadLimiter(deny=True) + with _files_client(role=UserRole.USER, limiter=limiter) as client: + response = client.get("/api/v1/files/1/download") + + assert response.status_code == 429 + assert limiter.calls == [("user:1", 10)] + + +def test_download_route_preserves_range_response_when_allowed() -> None: + limiter = StubDownloadLimiter() + with _files_client(role=UserRole.USER, limiter=limiter) as client: + response = client.get("/api/v1/files/1/download", headers={"Range": "bytes=0-3"}) + + assert response.status_code == 206 + assert response.headers["content-range"] == "bytes 0-3/10" + assert limiter.calls == [("user:1", 4)] + + +def test_admin_download_route_is_not_rejected_by_user_limiter() -> None: + limiter = StubDownloadLimiter(deny=True) + with _files_client(role=UserRole.ADMIN, limiter=limiter) as client: + response = client.get("/api/v1/files/1/download") + + assert response.status_code == 200 + assert limiter.calls == [("user:1", 10)] + + +def test_batch_download_route_returns_429_before_archive_when_limited() -> None: + limiter = StubDownloadLimiter(deny=True) + with _files_client(role=UserRole.USER, limiter=limiter) as client: + response = client.post("/api/v1/files/batch-download", json={"fileIds": ["1"]}) + + assert response.status_code == 429 + assert limiter.calls == [("user:1", 10)] + + +def test_admin_batch_download_route_is_not_rejected_by_user_limiter() -> None: + limiter = StubDownloadLimiter(deny=True) + with _files_client(role=UserRole.ADMIN, limiter=limiter) as client: + response = client.post("/api/v1/files/batch-download", json={"fileIds": ["1"]}) + + assert response.status_code == 200 + assert response.content == b"zip" + + @pytest.mark.asyncio async def test_delete_file_marks_record_deleted(monkeypatch: pytest.MonkeyPatch): session = DummySession() diff --git a/app/tests/test_file_folder_patch_routes.py b/app/tests/test_file_folder_patch_routes.py index b289d06..1c38aac 100644 --- a/app/tests/test_file_folder_patch_routes.py +++ b/app/tests/test_file_folder_patch_routes.py @@ -1,12 +1,15 @@ from __future__ import annotations from collections.abc import AsyncIterator -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from fastapi import FastAPI from fastapi.testclient import TestClient -from fileflash.core.deps import get_current_user, get_file_service, get_folder_service +from fileflash.core.deps import get_current_user, get_file_service, get_folder_service, get_settings_dep +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.core.security import create_file_preview_token +from fileflash.core.settings import Settings from fileflash.models.tables_identity import User from fileflash.routers.files import router as files_router from fileflash.routers.folders import router as folders_router @@ -47,6 +50,9 @@ def _make_folder_item(*, is_starred: bool) -> FolderItem: class StubFileService: + async def get_file(self, *, user_id: int, file_id: int) -> FileDetails: # noqa: ARG002 + return _make_file_details(name="demo.txt", is_starred=False) + async def rename_file(self, *, user_id: int, file_id: str, payload: RenameFileRequest) -> FileDetails: # noqa: ARG002 return _make_file_details(name=payload.file_name, is_starred=False) @@ -93,17 +99,26 @@ async def toggle_folder_star(self, *, user_id: int, folder_id: str, is_starred: return _make_folder_item(is_starred=is_starred) -def _build_client() -> TestClient: +def _build_client(*, authenticated: bool = True, settings: Settings | None = None) -> TestClient: app = FastAPI() app.include_router(files_router, prefix="/api/v1") app.include_router(folders_router, prefix="/api/v1") + app.add_exception_handler(ApiError, api_error_handler) async def _current_user_override() -> User: + if not authenticated: + raise ApiError(status_code=401, code=401, message="Missing authorization token") return User(user_id=1, username="owner", email="owner@example.com", password_hash="hash") + resolved_settings = settings or Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + app.dependency_overrides[get_current_user] = _current_user_override app.dependency_overrides[get_file_service] = lambda: StubFileService() app.dependency_overrides[get_folder_service] = lambda: StubFolderService() + app.dependency_overrides[get_settings_dep] = lambda: resolved_settings return TestClient(app) @@ -142,3 +157,88 @@ def test_get_file_preview_route_returns_stream() -> None: assert response.headers["content-range"] == "bytes 0-3/12" assert response.headers["content-type"].startswith("text/plain") assert response.content == b"prev" + + +def test_create_file_preview_url_route_returns_signed_stream_url() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + with _build_client(settings=settings) as client: + response = client.post("/api/v1/files/1/preview-url") + + assert response.status_code == 200 + payload = response.json() + assert payload["success"] is True + assert payload["data"]["url"].startswith("http://testserver/api/v1/files/1/preview-stream?token=") + assert payload["data"]["expiresAt"] + + +def test_create_file_preview_url_route_requires_authentication() -> None: + with _build_client(authenticated=False) as client: + response = client.post("/api/v1/files/1/preview-url") + + assert response.status_code == 401 + assert response.json()["message"] == "Missing authorization token" + + +def test_get_file_preview_stream_route_supports_range_with_valid_token() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_file_preview_token( + user_id=1, + file_id=1, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + + with _build_client(settings=settings) as client: + response = client.get( + f"/api/v1/files/1/preview-stream?token={token}", + headers={"Range": "bytes=0-3"}, + ) + + assert response.status_code == 206 + assert response.headers["content-range"] == "bytes 0-3/12" + assert response.headers["content-length"] == "4" + assert response.content == b"prev" + + +def test_get_file_preview_stream_route_rejects_mismatched_token_file_id() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_file_preview_token( + user_id=1, + file_id=2, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + + with _build_client(settings=settings) as client: + response = client.get(f"/api/v1/files/1/preview-stream?token={token}") + + assert response.status_code == 403 + assert response.json()["message"] == "Preview token does not match file" + + +def test_get_file_preview_stream_route_rejects_expired_token() -> None: + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_file_preview_token( + user_id=1, + file_id=1, + settings=settings, + expires_at=datetime.now(UTC) - timedelta(minutes=10), + ) + + with _build_client(settings=settings) as client: + response = client.get(f"/api/v1/files/1/preview-stream?token={token}") + + assert response.status_code == 401 + assert response.json()["message"] == "Invalid or expired preview token" diff --git a/app/tests/test_rate_limiter.py b/app/tests/test_rate_limiter.py new file mode 100644 index 0000000..6720f6f --- /dev/null +++ b/app/tests/test_rate_limiter.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest +from redis.exceptions import RedisError + +from fileflash.core.settings import Settings +from fileflash.models.enums import UserRole, UserStatus +from fileflash.models.tables_identity import User +from fileflash.services.download_rate_limit import DownloadRateLimitService +from fileflash.services.rate_limiter import RedisRateLimiter + + +class FakeRedis: + def __init__( + self, + *, + fail: bool = False, + runtime_error: RuntimeError | None = None, + close_runtime_error: RuntimeError | None = None, + ) -> None: + self.fail = fail + self.runtime_error = runtime_error + self.close_runtime_error = close_runtime_error + self.values: dict[str, int] = {} + self.expired: list[tuple[str, int]] = [] + + async def incrby(self, key: str, amount: int) -> int: + if self.runtime_error is not None: + raise self.runtime_error + if self.fail: + raise RedisError("down") + self.values[key] = self.values.get(key, 0) + amount + return self.values[key] + + async def expire(self, key: str, window_seconds: int) -> None: + self.expired.append((key, window_seconds)) + + async def aclose(self) -> None: + if self.close_runtime_error is not None: + raise self.close_runtime_error + + +@pytest.mark.asyncio +async def test_allow_weighted_uses_incrby_and_sets_ttl() -> None: + limiter = RedisRateLimiter("redis://example") + fake = FakeRedis() + limiter._redis = fake # type: ignore[assignment] + + allowed = await limiter.allow_weighted("k", limit=10, window_seconds=60, weight=4) + + assert allowed is True + assert fake.values["k"] == 4 + assert fake.expired == [("k", 60)] + + +@pytest.mark.asyncio +async def test_allow_weighted_rejects_over_limit() -> None: + limiter = RedisRateLimiter("redis://example") + fake = FakeRedis() + limiter._redis = fake # type: ignore[assignment] + + assert await limiter.allow_weighted("k", limit=5, window_seconds=60, weight=4) is True + assert await limiter.allow_weighted("k", limit=5, window_seconds=60, weight=2) is False + + +@pytest.mark.asyncio +async def test_allow_weighted_degrades_open_when_redis_fails() -> None: + limiter = RedisRateLimiter("redis://example") + limiter._redis = FakeRedis(fail=True) # type: ignore[assignment] + + assert await limiter.allow_weighted("k", limit=1, window_seconds=60, weight=10) is True + + +@pytest.mark.asyncio +async def test_allow_weighted_degrades_open_and_resets_client_on_loop_runtime_error() -> None: + limiter = RedisRateLimiter("redis://example") + limiter._redis = FakeRedis( + runtime_error=RuntimeError("Future pending attached to a different loop") + ) # type: ignore[assignment] + + assert await limiter.allow_weighted("k", limit=1, window_seconds=60, weight=10) is True + assert limiter._redis is None + + +@pytest.mark.asyncio +async def test_allow_weighted_reraises_non_loop_runtime_error() -> None: + limiter = RedisRateLimiter("redis://example") + limiter._redis = FakeRedis(runtime_error=RuntimeError("unexpected runtime failure")) # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="unexpected runtime failure"): + await limiter.allow_weighted("k", limit=1, window_seconds=60, weight=10) + + +@pytest.mark.asyncio +async def test_close_ignores_loop_runtime_error() -> None: + limiter = RedisRateLimiter("redis://example") + limiter._redis = FakeRedis(close_runtime_error=RuntimeError("Event loop is closed")) # type: ignore[assignment] + + await limiter.close() + + assert limiter._redis is None + + +class FakeRateLimiter: + def __init__(self) -> None: + self.allow = AsyncMock(return_value=False) + self.allow_weighted = AsyncMock(return_value=False) + + +@pytest.mark.asyncio +async def test_download_rate_limiter_skips_admin_user_id() -> None: + admin = User( + user_id=1, + username="admin", + email="admin@example.com", + password_hash="x", + role=UserRole.ADMIN, + status=UserStatus.ACTIVE, + email_verified=True, + storage_limit=1024, + storage_used=0, + ) + db = type("Db", (), {"get": AsyncMock(return_value=admin)})() + rate_limiter = FakeRateLimiter() + service = DownloadRateLimitService( + db=db, # type: ignore[arg-type] + settings=Settings(DATABASE_URL="sqlite+aiosqlite:///:memory:"), + rate_limiter=rate_limiter, # type: ignore[arg-type] + ) + + await service.enforce_user_id(user_id=1, bytes_count=100) + + rate_limiter.allow.assert_not_awaited() + rate_limiter.allow_weighted.assert_not_awaited() diff --git a/app/tests/test_security.py b/app/tests/test_security.py index fe29f57..6fb4a73 100644 --- a/app/tests/test_security.py +++ b/app/tests/test_security.py @@ -1,9 +1,15 @@ from __future__ import annotations +from datetime import UTC, datetime, timedelta + from fileflash.core.security import ( + create_admin_file_preview_token, create_access_token, + create_file_preview_token, create_refresh_token, + decode_admin_file_preview_token, decode_access_token, + decode_file_preview_token, get_password_hash, hash_token, verify_password, @@ -31,6 +37,42 @@ def test_access_token_round_trip(): assert "exp" in payload +def test_file_preview_token_round_trip(): + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_file_preview_token( + user_id=42, + file_id=99, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + payload = decode_file_preview_token(token=token, settings=settings) + assert payload["sub"] == "42" + assert payload["fileId"] == "99" + assert payload["scope"] == "file.preview" + assert payload["typ"] == "file_preview" + + +def test_admin_file_preview_token_round_trip(): + settings = Settings( + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + FF_DB_URI="postgresql://u:p@localhost:5432/db", + ) + token = create_admin_file_preview_token( + admin_user_id=7, + file_id=99, + settings=settings, + expires_at=datetime.now(UTC) + timedelta(minutes=10), + ) + payload = decode_admin_file_preview_token(token=token, settings=settings) + assert payload["sub"] == "7" + assert payload["fileId"] == "99" + assert payload["scope"] == "admin.file.preview" + assert payload["typ"] == "admin_file_preview" + + def test_refresh_token_hash_is_deterministic(): settings = Settings( JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", diff --git a/app/tests/test_settings.py b/app/tests/test_settings.py index de26b55..52356ea 100644 --- a/app/tests/test_settings.py +++ b/app/tests/test_settings.py @@ -5,73 +5,86 @@ from fileflash.core.settings import Settings +def make_settings(**overrides: object) -> Settings: + payload = {"FF_DB_URI": "postgresql://root:pwd@localhost:5432/fileflash"} + payload.update(overrides) + return Settings(_env_file=None, **payload) + + def test_async_database_url_conversion(): - settings = Settings(FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash") + settings = make_settings() assert settings.async_database_url == "postgresql+asyncpg://root:pwd@localhost:5432/fileflash" def test_upload_related_settings_defaults(): - settings = Settings(FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash") + settings = make_settings() assert settings.object_storage_bucket == "fileflash" assert settings.upload_chunk_size_default == 5 * 1024 * 1024 + assert settings.upload_single_file_size_max == 5 * 1024 * 1024 * 1024 + assert settings.upload_verify_merged_object_hash is False assert settings.upload_session_ttl_seconds == 24 * 3600 assert settings.starred_items_limit == 20 assert settings.worker_process_count == 1 def test_agent_related_settings_defaults(): - settings = Settings(FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash") + settings = make_settings() assert settings.agent_queue_stream == "fileflash:agents" assert settings.agent_job_timeout_sec == 600 assert settings.agent_tool_timeout_sec == 30 + assert settings.agent_llm_plan_max_tokens == 8192 assert settings.agent_mcp_endpoints == () +def test_auth_risk_control_defaults(): + settings = make_settings() + assert settings.register_rate_limit == 12 + assert settings.login_rate_limit == 30 + assert settings.max_failed_login_attempts == 8 + assert settings.account_lock_minutes == 5 + + def test_app_env_detection(): - dev = Settings(FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", APP_ENV="development") + dev = make_settings(APP_ENV="development") assert dev.is_development_env is True assert dev.is_production_env is False - prod = Settings(FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", APP_ENV="prod") + prod = make_settings(APP_ENV="prod") assert prod.is_development_env is False assert prod.is_production_env is True def test_worker_process_count_from_env(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( WORKER_PROCESS_COUNT="3", ) assert settings.worker_process_count == 3 def test_starred_items_limit_from_env(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( STARRED_ITEMS_LIMIT="12", ) assert settings.starred_items_limit == 12 def test_verify_base_url_defaults_to_localhost_in_development(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( APP_ENV="development", + EMAIL_VERIFY_BASE_URL="", ) assert settings.normalized_email_verify_base_url == "http://localhost:8080" def test_verify_base_url_adds_http_when_scheme_missing(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( EMAIL_VERIFY_BASE_URL="localhost:3000", ) assert settings.normalized_email_verify_base_url == "http://localhost:3000" def test_mail_configuration_issues_includes_missing_required_fields(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( APP_ENV="development", MAIL_FROM="", MAIL_SERVER="", @@ -85,8 +98,7 @@ def test_mail_configuration_issues_includes_missing_required_fields(): def test_mail_configuration_rejects_both_tls_modes_enabled(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( EMAIL_VERIFY_BASE_URL="http://localhost:5173", MAIL_FROM="demo@example.com", MAIL_SERVER="smtp.example.com", @@ -96,12 +108,15 @@ def test_mail_configuration_rejects_both_tls_modes_enabled(): MAIL_STARTTLS=True, MAIL_SSL_TLS=True, ) - assert "MAIL_SSL_TLS and MAIL_STARTTLS cannot both be true" in settings.mail_configuration_issues + assert ( + "MAIL_SSL_TLS and MAIL_STARTTLS cannot both be true" + in settings.mail_configuration_issues + ) def test_assert_runtime_security_raises_when_jwt_secret_too_short(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( + APP_ENV="development", JWT_SECRET_KEY="short-key", ) with pytest.raises(ValueError, match="JWT_SECRET_KEY must be at least 32 bytes"): @@ -109,11 +124,46 @@ def test_assert_runtime_security_raises_when_jwt_secret_too_short(): def test_assert_runtime_security_raises_when_token_hash_secret_too_short(): - settings = Settings( - FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + settings = make_settings( + APP_ENV="development", JWT_SECRET_KEY="x" * 32, TOKEN_HASH_SECRET="short-key", ) with pytest.raises(ValueError, match="TOKEN_HASH_SECRET must be at least 32 bytes"): settings.assert_runtime_security() + +def test_assert_runtime_security_requires_default_admin_env_in_production(): + settings = make_settings( + APP_ENV="production", + JWT_SECRET_KEY="x" * 32, + ) + + with pytest.raises(ValueError, match="DEFAULT_ADMIN_USERNAME is required in production"): + settings.assert_runtime_security() + + +def test_assert_runtime_security_rejects_short_default_admin_password_in_production(): + settings = make_settings( + APP_ENV="production", + JWT_SECRET_KEY="x" * 32, + DEFAULT_ADMIN_USERNAME="admin", + DEFAULT_ADMIN_EMAIL="admin@example.com", + DEFAULT_ADMIN_PASSWORD="short-password", + ) + + with pytest.raises(ValueError, match="DEFAULT_ADMIN_PASSWORD must be at least 32 bytes"): + settings.assert_runtime_security() + + +def test_assert_runtime_security_accepts_default_admin_env_in_production(): + settings = make_settings( + APP_ENV="production", + JWT_SECRET_KEY="x" * 32, + DEFAULT_ADMIN_USERNAME="admin", + DEFAULT_ADMIN_EMAIL="admin@example.com", + DEFAULT_ADMIN_PASSWORD="p" * 32, + ) + + settings.assert_runtime_security() + diff --git a/app/tests/test_share_routes.py b/app/tests/test_share_routes.py index 48cb008..e6d5b6f 100644 --- a/app/tests/test_share_routes.py +++ b/app/tests/test_share_routes.py @@ -5,7 +5,13 @@ from fastapi import FastAPI from fastapi.testclient import TestClient -from fileflash.core.deps import get_client_ip, get_share_service, get_user_agent +from fileflash.core.deps import ( + get_client_ip, + get_download_rate_limit_service, + get_share_service, + get_user_agent, +) +from fileflash.core.errors import ApiError, api_error_handler from fileflash.routers.shares import router as shares_router @@ -19,6 +25,7 @@ async def get_shared_file_download_stream_response( range_header: str | None, # noqa: ARG002 ip_address: str, # noqa: ARG002 user_agent: str | None, # noqa: ARG002 + rate_limit_check=None, # noqa: ANN001 ) -> tuple[AsyncIterator[bytes], str, str, int, dict[str, str]]: async def _stream() -> AsyncIterator[bytes]: yield b"data" @@ -32,15 +39,30 @@ async def _stream() -> AsyncIterator[bytes]: "Accept-Ranges": "bytes", "Content-Length": "4", } + if rate_limit_check is not None: + await rate_limit_check(4) return _stream(), "测试文档.pdf", "application/pdf", 200, headers -def _build_client() -> TestClient: +class StubDownloadLimiter: + def __init__(self, *, deny: bool = False) -> None: + self.deny = deny + self.calls: list[tuple[str, int]] = [] + + async def enforce_share_ip(self, *, client_ip: str, bytes_count: int) -> None: + self.calls.append((client_ip, bytes_count)) + if self.deny: + raise ApiError(status_code=429, code=429, message="Download rate limit exceeded") + + +def _build_client(limiter: StubDownloadLimiter | None = None) -> TestClient: app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) app.include_router(shares_router, prefix="/api/v1") app.dependency_overrides[get_share_service] = lambda: StubShareService() app.dependency_overrides[get_client_ip] = lambda: "127.0.0.1" app.dependency_overrides[get_user_agent] = lambda: "pytest" + app.dependency_overrides[get_download_rate_limit_service] = lambda: limiter or StubDownloadLimiter() return TestClient(app) @@ -71,3 +93,27 @@ def test_shared_preview_handles_unicode_filename_header() -> None: assert 'filename*=UTF-8\'\'' in header header.encode("latin-1") assert response.content == b"data" + + +def test_shared_download_returns_429_when_ip_limited() -> None: + limiter = StubDownloadLimiter(deny=True) + with _build_client(limiter) as client: + response = client.get( + "/api/v1/shares/ABCD/download", + headers={"Authorization": "Bearer test-share-token"}, + ) + + assert response.status_code == 429 + assert limiter.calls == [("127.0.0.1", 4)] + + +def test_shared_preview_returns_429_when_ip_limited() -> None: + limiter = StubDownloadLimiter(deny=True) + with _build_client(limiter) as client: + response = client.get( + "/api/v1/shares/ABCD/preview", + headers={"Authorization": "Bearer test-share-token"}, + ) + + assert response.status_code == 429 + assert limiter.calls == [("127.0.0.1", 4)] diff --git a/app/tests/test_startup_fail_fast.py b/app/tests/test_startup_fail_fast.py index a2ea383..b6c23e8 100644 --- a/app/tests/test_startup_fail_fast.py +++ b/app/tests/test_startup_fail_fast.py @@ -5,15 +5,30 @@ import pytest +from fileflash.core.settings import Settings from fileflash.main import lifespan from fileflash.s3.minio_client import ObjectStorageAuthError +def make_settings(**overrides: object) -> Settings: + payload = { + "FF_DB_URI": "postgresql://root:pwd@localhost:5432/fileflash", + "APP_ENV": "production", + "JWT_SECRET_KEY": "unit-test-secret-key-1234567890abcd", + "DEFAULT_ADMIN_USERNAME": "admin", + "DEFAULT_ADMIN_EMAIL": "admin@example.com", + "DEFAULT_ADMIN_PASSWORD": "p" * 32, + } + payload.update(overrides) + return Settings(_env_file=None, **payload) + + @pytest.mark.asyncio async def test_lifespan_fails_fast_when_database_check_fails(monkeypatch: pytest.MonkeyPatch): verify = AsyncMock(side_effect=RuntimeError("database unavailable")) verify_schema = AsyncMock() seed = AsyncMock() + monkeypatch.setattr("fileflash.main.settings", make_settings()) monkeypatch.setattr("fileflash.main.verify_database_connection", verify) monkeypatch.setattr("fileflash.main.verify_schema_compatibility", verify_schema) monkeypatch.setattr("fileflash.main.initialize_dev_accounts", seed) @@ -32,6 +47,7 @@ async def test_lifespan_fails_fast_when_schema_check_fails(monkeypatch: pytest.M verify = AsyncMock() verify_schema = AsyncMock(side_effect=RuntimeError("schema incompatible")) seed = AsyncMock() + monkeypatch.setattr("fileflash.main.settings", make_settings()) monkeypatch.setattr("fileflash.main.verify_database_connection", verify) monkeypatch.setattr("fileflash.main.verify_schema_compatibility", verify_schema) monkeypatch.setattr("fileflash.main.initialize_dev_accounts", seed) @@ -50,7 +66,10 @@ async def test_lifespan_fails_fast_when_object_storage_check_fails(monkeypatch: verify = AsyncMock() verify_schema = AsyncMock() seed = AsyncMock() - storage = SimpleNamespace(ensure_bucket=AsyncMock(side_effect=ObjectStorageAuthError("bad credentials"))) + storage = SimpleNamespace( + ensure_bucket=AsyncMock(side_effect=ObjectStorageAuthError("bad credentials")) + ) + monkeypatch.setattr("fileflash.main.settings", make_settings()) monkeypatch.setattr("fileflash.main.verify_database_connection", verify) monkeypatch.setattr("fileflash.main.verify_schema_compatibility", verify_schema) monkeypatch.setattr("fileflash.main.get_object_storage", lambda: storage) @@ -64,3 +83,54 @@ async def test_lifespan_fails_fast_when_object_storage_check_fails(monkeypatch: verify_schema.assert_awaited_once() storage.ensure_bucket.assert_awaited_once() seed.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lifespan_fails_fast_when_prod_admin_env_missing(monkeypatch: pytest.MonkeyPatch): + verify = AsyncMock() + verify_schema = AsyncMock() + seed = AsyncMock() + monkeypatch.setattr( + "fileflash.main.settings", + Settings( + _env_file=None, + FF_DB_URI="postgresql://root:pwd@localhost:5432/fileflash", + APP_ENV="production", + JWT_SECRET_KEY="unit-test-secret-key-1234567890abcd", + ), + ) + monkeypatch.setattr("fileflash.main.verify_database_connection", verify) + monkeypatch.setattr("fileflash.main.verify_schema_compatibility", verify_schema) + monkeypatch.setattr("fileflash.main.initialize_dev_accounts", seed) + + with pytest.raises(ValueError, match="DEFAULT_ADMIN_USERNAME is required in production"): + async with lifespan(object()): + pass + + verify.assert_not_awaited() + verify_schema.assert_not_awaited() + seed.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lifespan_fails_fast_when_prod_admin_password_too_short( + monkeypatch: pytest.MonkeyPatch, +): + verify = AsyncMock() + verify_schema = AsyncMock() + seed = AsyncMock() + monkeypatch.setattr( + "fileflash.main.settings", + make_settings(DEFAULT_ADMIN_PASSWORD="short-password"), + ) + monkeypatch.setattr("fileflash.main.verify_database_connection", verify) + monkeypatch.setattr("fileflash.main.verify_schema_compatibility", verify_schema) + monkeypatch.setattr("fileflash.main.initialize_dev_accounts", seed) + + with pytest.raises(ValueError, match="DEFAULT_ADMIN_PASSWORD must be at least 32 bytes"): + async with lifespan(object()): + pass + + verify.assert_not_awaited() + verify_schema.assert_not_awaited() + seed.assert_not_awaited() diff --git a/app/tests/test_tasks_transcode.py b/app/tests/test_tasks_transcode.py index 63ad82a..3998257 100644 --- a/app/tests/test_tasks_transcode.py +++ b/app/tests/test_tasks_transcode.py @@ -5,7 +5,9 @@ from pathlib import Path from types import SimpleNamespace -from fileflash.tasks.transcode import build_ffmpeg_command, run_media_transcode +import pytest + +from fileflash.tasks.transcode import _run_command, build_ffmpeg_command, run_media_transcode def test_build_ffmpeg_command_for_video_contains_profile_flags(): @@ -24,9 +26,53 @@ def test_build_ffmpeg_command_for_video_contains_profile_flags(): assert "+faststart" in command assert "-pix_fmt" in command assert "yuv420p" in command + assert "-ac" in command + assert command[command.index("-ac") + 1] == "2" assert command[-1] == "output.mp4" +def test_build_ffmpeg_command_for_audio_forces_stereo(): + command = build_ffmpeg_command( + input_path=Path("input.wav"), + output_path=Path("output.m4a"), + media_type="audio", + ffmpeg_binary="ffmpeg", + payload={}, + ) + + assert command[:4] == ["ffmpeg", "-y", "-i", "input.wav"] + assert "-vn" in command + assert "-c:a" in command + assert "-ac" in command + assert command[command.index("-ac") + 1] == "2" + assert command[-1] == "output.m4a" + + +def test_run_command_raises_value_error_for_deterministic_ffmpeg_encoder_failure(monkeypatch): + def fake_run( + command: list[str], + *, + check: bool, + capture_output: bool, + text: bool, + timeout: int, + ) -> subprocess.CompletedProcess[str]: + _ = (check, capture_output, text, timeout) + return subprocess.CompletedProcess( + command, + returncode=1, + stdout="", + stderr='[aac @ 0000020b7eff3d00] Unsupported channel layout "6 channels"\n' + "Error while opening encoder - maybe incorrect parameters\n" + "[af#0:1 @ 0000020b003fde00] Task finished with error code: -22 (Invalid argument)", + ) + + monkeypatch.setattr("fileflash.tasks.transcode.subprocess.run", fake_run) + + with pytest.raises(ValueError, match="Command failed"): + _run_command(["ffmpeg", "-i", "input.mp4", "output.mp4"], timeout_seconds=30) + + def test_run_media_transcode_storage_mode_with_mocked_subprocess(monkeypatch, tmp_path: Path): source_probe = { "streams": [{"codec_type": "audio", "codec_name": "pcm_s16le", "sample_rate": "44100"}], diff --git a/app/tests/test_upload_routes.py b/app/tests/test_upload_routes.py new file mode 100644 index 0000000..a5b9855 --- /dev/null +++ b/app/tests/test_upload_routes.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import get_current_user, get_upload_service +from fileflash.models.tables_identity import User +from fileflash.routers.uploads import router as uploads_router +from fileflash.schemas.file import RecoverableUploadSession, UploadCancelResponse + + +class StubUploadService: + async def cancel_upload_session(self, *, user_id: int, upload_id: str) -> UploadCancelResponse: # noqa: ARG002 + return UploadCancelResponse( + upload_id=upload_id, + canceled_at=datetime(2026, 5, 24, 9, 0, 0, tzinfo=UTC), + ) + + async def list_recoverable_sessions(self, *, user_id: int) -> list[RecoverableUploadSession]: # noqa: ARG002 + return [ + RecoverableUploadSession( + upload_id="upload-1", + file_name="demo.mp4", + file_size=1024, + uploaded_bytes=512, + chunk_size=256, + file_hash="a" * 64, + mime_type="video/mp4", + parent_id="root", + updated_at=datetime(2026, 5, 24, 8, 0, 0, tzinfo=UTC), + expired_at=datetime(2026, 5, 25, 8, 0, 0, tzinfo=UTC), + status="uploading", + ) + ] + + +def _build_client() -> TestClient: + app = FastAPI() + app.include_router(uploads_router, prefix="/api/v1") + + async def _current_user_override() -> User: + return User(user_id=1, username="owner", email="owner@example.com", password_hash="hash") + + app.dependency_overrides[get_current_user] = _current_user_override + app.dependency_overrides[get_upload_service] = lambda: StubUploadService() + return TestClient(app) + + +def test_post_upload_cancel_route_returns_success_shell() -> None: + with _build_client() as client: + response = client.post("/api/v1/uploads/upload-123/cancel") + + assert response.status_code == 200 + payload = response.json() + assert payload["success"] is True + assert payload["message"] == "Upload session canceled" + assert payload["data"]["uploadId"] == "upload-123" + assert str(payload["data"]["canceledAt"]).startswith("2026-05-24T09:00:00") + + +def test_get_recoverable_uploads_route_returns_success_shell() -> None: + with _build_client() as client: + response = client.get("/api/v1/uploads/recoverable") + + assert response.status_code == 200 + payload = response.json() + assert payload["success"] is True + assert payload["message"] == "Recoverable upload sessions fetched" + assert isinstance(payload["data"], list) + assert payload["data"][0]["uploadId"] == "upload-1" + assert payload["data"][0]["uploadedBytes"] == 512 diff --git a/app/tests/test_upload_service.py b/app/tests/test_upload_service.py index 298d019..58d3382 100644 --- a/app/tests/test_upload_service.py +++ b/app/tests/test_upload_service.py @@ -168,6 +168,162 @@ async def test_preflight_returns_503_when_storage_is_unavailable(monkeypatch: py cleanup_mock.assert_not_awaited() +@pytest.mark.asyncio +async def test_preflight_allows_20gb_file_size_boundary(monkeypatch: pytest.MonkeyPatch): + max_file_size = 20 * 1024 * 1024 * 1024 + settings = make_settings(UPLOAD_SINGLE_FILE_SIZE_MAX=str(max_file_size)) + session = DummySession() + service, _storage = make_service(session, settings=settings) + + monkeypatch.setattr(service, "_cleanup_expired_tasks", AsyncMock()) + monkeypatch.setattr(service, "_resolve_folder_id", AsyncMock(return_value=7)) + monkeypatch.setattr(service, "_find_storage_object", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_find_active_task", AsyncMock(return_value=None)) + + response = await service.preflight( + user_id=2, + payload=UploadPreflightRequest( + fileHash="b" * 64, + fileName="video.mp4", + fileSize=max_file_size, + mimeType="video/mp4", + parentId="7", + ), + ) + + assert response.status == "UPLOADING" + assert response.upload_id is not None + + +@pytest.mark.asyncio +async def test_preflight_rejects_file_size_over_20gb(): + max_file_size = 20 * 1024 * 1024 * 1024 + settings = make_settings(UPLOAD_SINGLE_FILE_SIZE_MAX=str(max_file_size)) + session = DummySession() + service, storage = make_service(session, settings=settings) + + with pytest.raises(ApiError) as exc: + await service.preflight( + user_id=2, + payload=UploadPreflightRequest( + fileHash="b" * 64, + fileName="video.mp4", + fileSize=max_file_size + 1, + mimeType="video/mp4", + parentId="7", + ), + ) + + assert exc.value.status_code == 413 + assert exc.value.code == 413 + storage.ensure_bucket.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_list_recoverable_sessions_returns_only_active_non_expired_tasks(): + session = DummySession() + service, _storage = make_service(session) + now = datetime.now(UTC) + active_uploading = UploadTask( + task_id=100, + user_id=1, + folder_id=7, + file_name="resume.bin", + mime_type="application/octet-stream", + bucket_name="fileflash", + object_key="objects/u1/resume", + object_hash="a" * 64, + total_size=100, + uploaded_bytes=140, + chunk_size=None, + upload_id="upload-resume-1", + status=UploadTaskStatus.UPLOADING, + expired_at=now + timedelta(hours=1), + updated_at=now, + ) + active_init = UploadTask( + task_id=101, + user_id=1, + folder_id=None, + file_name="init.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/init", + object_hash="b" * 64, + total_size=20, + uploaded_bytes=5, + chunk_size=10, + upload_id="upload-resume-2", + status=UploadTaskStatus.INIT, + expired_at=now + timedelta(hours=2), + updated_at=now, + ) + expired_task = UploadTask( + task_id=102, + user_id=1, + folder_id=9, + file_name="expired.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/expired", + object_hash="c" * 64, + total_size=30, + chunk_size=10, + upload_id="upload-expired", + status=UploadTaskStatus.UPLOADING, + expired_at=now - timedelta(minutes=1), + updated_at=now, + ) + failed_task = UploadTask( + task_id=103, + user_id=1, + folder_id=9, + file_name="failed.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/failed", + object_hash="d" * 64, + total_size=30, + chunk_size=10, + upload_id="upload-failed", + status=UploadTaskStatus.FAILED, + expired_at=now + timedelta(hours=1), + updated_at=now, + ) + missing_hash_task = UploadTask( + task_id=104, + user_id=1, + folder_id=9, + file_name="missing-hash.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/nohash", + object_hash=None, + total_size=30, + chunk_size=10, + upload_id="upload-nohash", + status=UploadTaskStatus.UPLOADING, + expired_at=now + timedelta(hours=1), + updated_at=now, + ) + session.scalars_queue = [[active_uploading, active_init, expired_task, failed_task, missing_hash_task]] + + sessions = await service.list_recoverable_sessions(user_id=1) + + assert len(sessions) == 2 + first = sessions[0] + assert first.upload_id == "upload-resume-1" + assert first.file_size == 100 + assert first.uploaded_bytes == 100 + assert first.chunk_size == 5 * 1024 * 1024 + assert first.parent_id == "7" + assert first.status == "uploading" + second = sessions[1] + assert second.upload_id == "upload-resume-2" + assert second.parent_id == "root" + assert second.status == "init" + + @pytest.mark.asyncio async def test_upload_chunk_rejects_out_of_range_index(monkeypatch: pytest.MonkeyPatch): session = DummySession() @@ -201,6 +357,108 @@ async def test_upload_chunk_rejects_out_of_range_index(monkeypatch: pytest.Monke assert "out of range" in exc.value.message +@pytest.mark.asyncio +async def test_cancel_upload_session_marks_aborted_and_cleans_objects(monkeypatch: pytest.MonkeyPatch): + session = DummySession() + service, storage = make_service(session) + task = UploadTask( + task_id=9, + user_id=1, + folder_id=1, + file_name="clip.mp4", + mime_type="video/mp4", + bucket_name="fileflash", + object_key="objects/u1/demo", + object_hash="a" * 64, + total_size=1024, + chunk_size=512, + upload_id="upload-cancel-active", + status=UploadTaskStatus.UPLOADING, + expired_at=datetime.now(UTC) + timedelta(hours=1), + ) + + monkeypatch.setattr(service, "_get_task_for_update", AsyncMock(return_value=task)) + monkeypatch.setattr( + service, + "_collect_task_cleanup_keys", + AsyncMock(return_value=["objects/u1/demo", "temp/u1/upload-cancel-active/part-00000000"]), + ) + + response = await service.cancel_upload_session(user_id=1, upload_id="upload-cancel-active") + + assert response.upload_id == "upload-cancel-active" + assert task.status == UploadTaskStatus.ABORTED + assert task.last_error == "Upload canceled by client" + assert session.commits == 1 + storage.remove_objects.assert_awaited_once_with( + object_keys=["objects/u1/demo", "temp/u1/upload-cancel-active/part-00000000"], + ) + + +@pytest.mark.asyncio +async def test_cancel_upload_session_is_idempotent_for_aborted_task(monkeypatch: pytest.MonkeyPatch): + session = DummySession() + service, storage = make_service(session) + canceled_at = datetime.now(UTC) + task = UploadTask( + task_id=10, + user_id=1, + folder_id=1, + file_name="already-aborted.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/aborted", + object_hash="b" * 64, + total_size=8, + chunk_size=4, + upload_id="upload-cancel-aborted", + status=UploadTaskStatus.ABORTED, + expired_at=datetime.now(UTC) + timedelta(hours=1), + updated_at=canceled_at, + ) + monkeypatch.setattr(service, "_get_task_for_update", AsyncMock(return_value=task)) + collect_mock = AsyncMock(return_value=["objects/u1/aborted"]) + monkeypatch.setattr(service, "_collect_task_cleanup_keys", collect_mock) + + response = await service.cancel_upload_session(user_id=1, upload_id="upload-cancel-aborted") + + assert response.upload_id == "upload-cancel-aborted" + assert response.canceled_at == canceled_at + assert session.commits == 0 + collect_mock.assert_not_awaited() + storage.remove_objects.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_cancel_upload_session_rejects_completed_task(monkeypatch: pytest.MonkeyPatch): + session = DummySession() + service, storage = make_service(session) + task = UploadTask( + task_id=11, + user_id=1, + folder_id=1, + file_name="done.txt", + mime_type="text/plain", + bucket_name="fileflash", + object_key="objects/u1/done", + object_hash="c" * 64, + total_size=8, + chunk_size=4, + upload_id="upload-cancel-completed", + status=UploadTaskStatus.COMPLETED, + expired_at=datetime.now(UTC) + timedelta(hours=1), + ) + monkeypatch.setattr(service, "_get_task_for_update", AsyncMock(return_value=task)) + + with pytest.raises(ApiError) as exc: + await service.cancel_upload_session(user_id=1, upload_id="upload-cancel-completed") + + assert exc.value.status_code == 409 + assert "already completed" in exc.value.message + assert session.commits == 0 + storage.remove_objects.assert_not_awaited() + + @pytest.mark.asyncio async def test_merge_returns_conflict_without_strategy(monkeypatch: pytest.MonkeyPatch): session = DummySession() @@ -251,9 +509,56 @@ async def test_merge_returns_conflict_without_strategy(monkeypatch: pytest.Monke @pytest.mark.asyncio -async def test_merge_marks_failed_on_hash_mismatch(monkeypatch: pytest.MonkeyPatch): +async def test_merge_skips_full_hash_verification_in_fast_mode(monkeypatch: pytest.MonkeyPatch): session = DummySession() service, storage = make_service(session) + task = UploadTask( + task_id=19, + user_id=2, + folder_id=2, + file_name="asset.bin", + mime_type="application/octet-stream", + bucket_name="fileflash", + object_key="objects/u2/new", + object_hash="a" * 64, + total_size=4, + chunk_size=2, + upload_id="upload-fast-mode", + status=UploadTaskStatus.UPLOADING, + expired_at=datetime.now(UTC) + timedelta(hours=1), + ) + parts = [ + UploadTaskPart(task_id=19, part_number=0, part_size=2, status=UploadPartStatus.UPLOADED), + UploadTaskPart(task_id=19, part_number=1, part_size=2, status=UploadPartStatus.UPLOADED), + ] + session.scalars_queue = [parts] + + monkeypatch.setattr(service, "_get_task_for_update", AsyncMock(return_value=task)) + monkeypatch.setattr(service, "_resolve_folder_id", AsyncMock(return_value=2)) + monkeypatch.setattr(service, "_find_conflict_file", AsyncMock(return_value=None)) + monkeypatch.setattr(service, "_find_storage_object", AsyncMock(return_value=None)) + + response = await service.merge_chunks( + user_id=2, + upload_id="upload-fast-mode", + payload=MergeChunksRequest( + fileHash="a" * 64, + fileName="asset.bin", + mimeType="application/octet-stream", + parentId="2", + conflictStrategy="rename", + ), + ) + + assert response.file_name == "asset.bin" + storage.compute_object_hash.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_merge_marks_failed_on_hash_mismatch(monkeypatch: pytest.MonkeyPatch): + settings = make_settings(UPLOAD_VERIFY_MERGED_OBJECT_HASH=True) + session = DummySession() + service, storage = make_service(session, settings=settings) task = UploadTask( task_id=20, user_id=2, diff --git a/app/uv.lock b/app/uv.lock index f1aa17f..4913fb8 100644 --- a/app/uv.lock +++ b/app/uv.lock @@ -34,6 +34,25 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anthropic" +version = "0.104.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "docstring-parser" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/22/c7/7a655b948916f777354648ce979f68b94d5b8dbdb5f61fed1f37fad9378c/anthropic-0.104.1.tar.gz", hash = "sha256:17362b6c45f527afcc9b0fdf62011ffd359726ab2ebcb1978ea0cc41bd8d8d40", size = 850081, upload-time = "2026-05-22T15:36:57.432Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b8/12/d9ab42790494d7c428391a46cd28492395566a6a8ccb138d681978594455/anthropic-0.104.1-py3-none-any.whl", hash = "sha256:35c8cb456f5a4405aafe1f10f03f6fcc54fa51fa8ec01d655cc4b437d120e9b7", size = 832996, upload-time = "2026-05-22T15:36:59.519Z" }, +] + [[package]] name = "anyio" version = "4.12.1" @@ -463,6 +482,15 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1b/82/ca4893968aeb2709aacfb57a30dec6fa2ab25b10fa9f064b8882ce33f599/cryptography-46.0.6-cp38-abi3-win_amd64.whl", hash = "sha256:79e865c642cfc5c0b3eb12af83c35c5aeff4fa5c672dc28c43721c2c9fdd2f0f", size = 3471160, upload-time = "2026-03-25T23:34:37.191Z" }, ] +[[package]] +name = "distro" +version = "1.9.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722, upload-time = "2023-12-24T09:54:32.31Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" }, +] + [[package]] name = "dnspython" version = "2.8.0" @@ -472,6 +500,15 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, ] +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + [[package]] name = "email-validator" version = "2.3.0" @@ -527,6 +564,7 @@ name = "fileflash" version = "0.1.0" source = { editable = "." } dependencies = [ + { name = "anthropic" }, { name = "asyncpg" }, { name = "fastapi" }, { name = "fastapi-mail" }, @@ -554,6 +592,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "anthropic", specifier = ">=0.104.1" }, { name = "asyncpg", specifier = ">=0.30.0" }, { name = "fastapi", specifier = ">=0.135.1" }, { name = "fastapi-mail", specifier = ">=1.6.2" }, @@ -745,6 +784,78 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "jiter" +version = "0.15.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/66/b5/55f06bb281d92fb3cc86d14e1def2bd908bb77693183e7cb1f5a3c388b0c/jiter-0.15.0.tar.gz", hash = "sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76", size = 166640, upload-time = "2026-05-19T10:09:48.361Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/44/53/4f6bddbcde3c71e56d0aa1337ec95950f3d27dd4153e25aadf0feac71751/jiter-0.15.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d", size = 308793, upload-time = "2026-05-19T10:07:35.25Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/01/84/c01099b59a285a1ebba64ae93f62bfa036675340fd1b0045ae65890a0442/jiter-0.15.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8c9004af7c8d67cce7f1aae1026fb55607f4aa600710d08ede3a3ce4aeefe7e0", size = 309570, upload-time = "2026-05-19T10:07:36.919Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/58/64/8fb7f9d45bb98190355454cd04dad8d8f27223d6bd52f83af07f637168a6/jiter-0.15.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c210f8b35dc6f30aafd4b4365ca89b9d1189f21ab49b8e68fa6322a847aef138", size = 336783, upload-time = "2026-05-19T10:07:38.694Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c3/b6/f5739011d009b3a30f6a53c5240979030ba29ae46a8c67e3a15759f7c37d/jiter-0.15.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5f30bae8bc1c2d613e28e5af3e8cceb09b742f1c8a8a5f839fb67afaffc03b61", size = 363555, upload-time = "2026-05-19T10:07:40.832Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/12/98a9d9f766665e8a3b6252454e17cb0c464606a28cf2fa09399b003345fa/jiter-0.15.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c60e71b6d10cfc284c9bf36bd885e8d44c46f688ce50aa91b5edd90181dea687", size = 452255, upload-time = "2026-05-19T10:07:42.62Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/d5/60f972840f79c5e7544fce567c56f1e4e50468f996baba3e78d823dd62a6/jiter-0.15.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ab068bce62a45aa3e7367eceaffb5dde60b7eb853be8dece45132e3d0ff4879", size = 373559, upload-time = "2026-05-19T10:07:44.201Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ee/cf/d46ef1234ba335aabc2f013210db8e0821a22f5e644a2e9449df199ecc23/jiter-0.15.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa248c9eb220197d363f688818dac2fd4b2f0cd7d843ca7105d652034823427d", size = 346055, upload-time = "2026-05-19T10:07:46.005Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/63/4d2749d8d54d230bad9b3a6b0d00cc28c6ff6b2fdffc26a8ccf76cc5a974/jiter-0.15.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2a77aadd57cac1682e4401a72724d2796d89a4ba129b1a5812aa94ee480826eb", size = 351406, upload-time = "2026-05-19T10:07:47.855Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d9/b9/9965b990035d8773328e0a8c8b457a87bf2b19f6c4126d9d99296be5d16a/jiter-0.15.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2ae901f3a55bfafdde31d289590fa25e3245735a2b1e8c7cc15871710a002871", size = 389357, upload-time = "2026-05-19T10:07:49.665Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2d/55/9ddf903deda1413e87fed792f416b7123daee5b8efbad6a202a7421c36a5/jiter-0.15.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:f0b271b462769543716f92d3a4f90527df6ef5ed05ee95ec4137f513e21e1b77", size = 517263, upload-time = "2026-05-19T10:07:51.537Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/76/a0c40ad064d3a20a4fde231e35d56e9a01ce82164278180e82d5daf85469/jiter-0.15.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:2fb6a5d26af81fc0f00f9360a891e05cf755e149bba391c4d563adc54812973d", size = 548646, upload-time = "2026-05-19T10:07:53.196Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/23/4f/eca9b954942916ba2f453891b8593ab444cd872396fe66a3936616f236f3/jiter-0.15.0-cp312-cp312-win32.whl", hash = "sha256:c2f6bb8b5216ab9e7873bc08b5d7bef2b8abbb578a3069bf1cd14a45d71d771d", size = 206427, upload-time = "2026-05-19T10:07:55.307Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/95/bf/8ead82a87495149542748e828d153fd232a512a22c83b02c4815c1a9c7d8/jiter-0.15.0-cp312-cp312-win_amd64.whl", hash = "sha256:40b2c7e92c44a84d748d21706c68dc6ff8161d80b59c99d774721a0d2317d7c7", size = 197300, upload-time = "2026-05-19T10:07:56.651Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f4/e4/9b8a78fb2d894471bc344e37f1949bdd784bd914d031dba0ba3a40c71dd7/jiter-0.15.0-cp312-cp312-win_arm64.whl", hash = "sha256:cc0bc345cf2df9d1c00ac443f50d543c1ccfa8b0422cb85b1ab70d681c0b255b", size = 192702, upload-time = "2026-05-19T10:07:58.307Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/f4/f708c900ecee41b2025ef8413d5351e5649eb2125c506f6720cc69b06f5c/jiter-0.15.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1c11465f97e2abf45a014b83b730222f8f1c5335e802c7055a67d50de6f1f4e3", size = 307829, upload-time = "2026-05-19T10:07:59.704Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/86/59/db537c0949e83668c38481d426b9f2fd5ab758c4ee53a811dd0a510626a0/jiter-0.15.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d1e7b1776f0797956c509e123d0952d10d293a9492dea9f288ab9570ec01d1a5", size = 308445, upload-time = "2026-05-19T10:08:01.184Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/37/38/ea0e13b18c30ef951da0d47d39e7fa9edb82a93a62990ffbd7cea9b622d4/jiter-0.15.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:351a341c2105aa430b7047e30f1bf7975f6313b00165d3fc07be2edaf741f279", size = 336181, upload-time = "2026-05-19T10:08:02.688Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/58/fc/2303901b16c4ba05865588990a420c0b4156270b44379c20931544a1d962/jiter-0.15.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4ab395feec8d249ec4044e228e98a7033f043426a265df439dc3698823f0a4e4", size = 362985, upload-time = "2026-05-19T10:08:04.394Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5b/6f/11bace093c52e7d4d26c8e606ccd7ae8c972189622469ec0d9e28161e28b/jiter-0.15.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a2a438005b6f22d0273413484d6094d7c2c5d10ec1b3a3bf128e0d1d3ba53258", size = 453292, upload-time = "2026-05-19T10:08:05.967Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/22/db/987f2f086ca4d7a6582eb4ccd513f9b26b42d9e4243a087609a3137a8fc7/jiter-0.15.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f18f85e4218d1b40f000f42a92239a7a61a902cd42c65e6c360dbd17dcb20894", size = 373501, upload-time = "2026-05-19T10:08:07.857Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8f/7c/89fbcabb2739b7a5b8dc959a1b6c5761f6484f5fed3486854b3c789bb1de/jiter-0.15.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d1aa62e277fc1cbd80e6deacae6f4d983b41b3d7728e0645c5d741a6149bba45", size = 344683, upload-time = "2026-05-19T10:08:09.431Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/30/6f/6cca7692e7dddfec6d8d76c54dc97f2af2a41df4ac0674b999df1f09a5f3/jiter-0.15.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:6550fa135c7deb8ead6af49ed7ff648532ea8334a1447fe34a36315ef79c5c29", size = 350892, upload-time = "2026-05-19T10:08:11.352Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/39/14/0338d6190cb8e6d22e677ab1d4eabd4117f67cca70c54cd04b82ff64e068/jiter-0.15.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:066f8f33f18b2419cd8213b2436fa7fbc9c499f315971cfa3ce1f9820c001b1b", size = 388723, upload-time = "2026-05-19T10:08:12.912Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/90/31/cc19f4a1bdb6afb09ce6a2f2615aa8d44d994eba0d8e6105ed1af920e736/jiter-0.15.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:75e8a04e91432dde9f1838373cf93d23726c79d3e908d319acf0e796f85592e7", size = 516648, upload-time = "2026-05-19T10:08:14.808Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/49/9f/833c541512cd091b63c10c0381973dfe11bc7a503a818c16384417e0c81e/jiter-0.15.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:a97261f1fccb8e50ecd2890a96e46efdc3f57c80a197324c6777827231eca712", size = 547382, upload-time = "2026-05-19T10:08:16.927Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d2/11/e7b70e91f90bc4477e8eee9e8a5f7cf3cb41b4525d6394dc98a714eb8f7f/jiter-0.15.0-cp313-cp313-win32.whl", hash = "sha256:c77496cb10bd7549690fbbab3e5ec05857b83e49276f4a9423a766ddd2afcd4c", size = 205845, upload-time = "2026-05-19T10:08:18.401Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4b/23/5c20d9ad6f02c493e4023e5d2d09e1c1f15fe2753c9102c544aff068a88e/jiter-0.15.0-cp313-cp313-win_amd64.whl", hash = "sha256:b15741f501469009ae0ae90b7147958a664a7dede40aa7ff174a8a4645f546d0", size = 196842, upload-time = "2026-05-19T10:08:20.131Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6b/11/1eb400ef248e8c925fd883fbe325daf5e42cd1b0d308539dd332bd4f7ffc/jiter-0.15.0-cp313-cp313-win_arm64.whl", hash = "sha256:5d6a60072b44c3c2b797a7ddcbcbbf2b34ea3cfd4721580fbfd2a09d9d9b84ba", size = 192212, upload-time = "2026-05-19T10:08:21.807Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8a/60/2fd8d7c79da8acf9b7b277c7616847773779356b92acfc9bb158452174da/jiter-0.15.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ef1fd24d9413f6209e00d3d5a453e67acfe004a25cc6c8e8484faed4311ab9e8", size = 315065, upload-time = "2026-05-19T10:08:23.218Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/46/f4/008fb7d65e8ac2abf00811651a661e025c4ba80bbc6f378450384ddd3aed/jiter-0.15.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:144f8e72cb53dab146347b91cceac01f5481237f2b93b4a339a1ee8f8878b67c", size = 339444, upload-time = "2026-05-19T10:08:24.701Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/00/55/90b0c7b9c6896c0f2a591dd36d36b71d22e09674bfef178fa03ba3f81499/jiter-0.15.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:553fcac2ef2cb990877f9fc0833b8b629a3e6a5670b6b5fd58219b41a653ddc4", size = 347779, upload-time = "2026-05-19T10:08:26.408Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/51/6b/69666cec5000fd57734c118437394516c749ae8dbeea9fb66d6fef9c4775/jiter-0.15.0-cp313-cp313t-win_amd64.whl", hash = "sha256:774f93f65031856bf14ad9f59bdcab8b8cad501e5ceabd51ba3525f76937a25b", size = 200395, upload-time = "2026-05-19T10:08:28.055Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/39/04/a6aa62cd27e8149b0d28df5561f10f6cceaf7935a9ccf3f1c5a05f9a0cd8/jiter-0.15.0-cp313-cp313t-win_arm64.whl", hash = "sha256:f1e1754960f38ec40613a07e5e372df67acb3b890fb383b6fb3de3e49ddbf3c7", size = 190516, upload-time = "2026-05-19T10:08:29.35Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/eb/d2/079f350ebf7859d081de30aa890f9e3be68516f754f3ba32366ffff4dcee/jiter-0.15.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:ac0d9ddea4350974be7a221fc25895f251a8fee748c889bdced2141c0fec1a49", size = 308884, upload-time = "2026-05-19T10:08:31.667Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/04/4e/a2c30a7f69b48c03b20935d647479106fe932f6e63f75faf53937197e05d/jiter-0.15.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:01a8222cf05ab1128e239421156c207949808acaaea2bdfd33130ae666786e86", size = 310028, upload-time = "2026-05-19T10:08:33.304Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/40/90/2e7cdfd3cf8ca967be38c48f5cf474d79f089efaf559a40f15984a77ae69/jiter-0.15.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:182226cbc930c9fab81bc2e41a4da672f89539906dadb05e75670ac07b94f71f", size = 337485, upload-time = "2026-05-19T10:08:35.259Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9b/11/15a1aa28b120b8ee5b4f1fb894c125046225f09847738bd64233d3b84883/jiter-0.15.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:71683c38c825452999b5717fcae07ea708e8c93003e808be4319c1b02e3d176e", size = 364223, upload-time = "2026-05-19T10:08:36.694Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b7/25/f442e8af5f3d0dcf47b39e83a0efd9ee45ea946aa6d04625dc3181eae3b6/jiter-0.15.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:30f2218e6a9e5c18bc10fe6d41ac189c442c88eacf11bad9f28ef95a9bef00e6", size = 456387, upload-time = "2026-05-19T10:08:38.143Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/f4/37f2d2c9f64f49af7da652ed7532bb5a2372e588e6927c3fdd76f911db65/jiter-0.15.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5157de9f76eb4bc5ea74a1219366a25f945ad305641d74e04f59c54087091aa9", size = 374461, upload-time = "2026-05-19T10:08:39.869Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/28/edcfbbbf0cb15436f36664a8908a0df47ab9006298d4cd937dc08ea932d6/jiter-0.15.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90c5db5527c221249a876160663ab891ace358c17f7b9c93ec1478b7f0550e5c", size = 345924, upload-time = "2026-05-19T10:08:41.668Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/47/13/89fba6398dab7f202b7278c4b4aac122399d2c0183971c4a57a3b7088df5/jiter-0.15.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:3e4540b8e74e4268811ac05db226a6a128ff572e7e0ce3f1163b693cadb184cd", size = 352283, upload-time = "2026-05-19T10:08:43.091Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1b/da/0f6af8cef2c565a1ab44d970f268c43ccaa72707386ea6388e6fe2b6cd26/jiter-0.15.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:62ebd14e47e9aed9df4472afcb2663668ce4d74891cd54f86bf6e44029d6dc89", size = 389985, upload-time = "2026-05-19T10:08:44.915Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a1/ec/b9cb7d6d29e24ee14910266157d2a279d7a8f60ee0df7fa840882976ba64/jiter-0.15.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0be6f5ad41a809f303f416d17cec92a7a725902fb9b4f3de3d19362ac0ef8554", size = 517695, upload-time = "2026-05-19T10:08:46.486Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/64/5e/6d1bda880723aae0ad86b4b763f044362448efe31e3e819635d41cb03451/jiter-0.15.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:813dfbb17d65328bf86e5f0905dd277ba2265d3ca20556e86c0c7035b7182e5a", size = 548868, upload-time = "2026-05-19T10:08:48.026Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0c/72/7de501cf38dcacaf35098796f3a50e0f2e338baba18a58946c618544b809/jiter-0.15.0-cp314-cp314-win32.whl", hash = "sha256:50e51156192722a9c58db112837d3f8ef96fb3c5ecc14e95f409134b08b158ec", size = 206380, upload-time = "2026-05-19T10:08:49.738Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/a9/e19addf4b0c1bdce52c6da12351e6bc42c340c45e7c09e2158e46d293ccc/jiter-0.15.0-cp314-cp314-win_amd64.whl", hash = "sha256:30ce1a5d16b5641dc935d50ef775af6a0871e3d14ab05d6fc54dff371b78e558", size = 197687, upload-time = "2026-05-19T10:08:51.088Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f2/c9/776b1db01db25fc6c1d58d1979a37b0a9fe787e5f5b1d062d2eaacb77923/jiter-0.15.0-cp314-cp314-win_arm64.whl", hash = "sha256:510c8b3c17a0ed9ac69850c0438dada3c9b82d9c4d589fcb62002a5a9cf3a866", size = 192571, upload-time = "2026-05-19T10:08:52.451Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a0/f6/45bb4670bacf300fd2c7abadbfb3af376e5f1b6ae75fd9bc069891d15870/jiter-0.15.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7553333dd0930c104a5a0db8df72bf7219fe663d731383b576bb6ed6351c984d", size = 317151, upload-time = "2026-05-19T10:08:53.867Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/68/ed635ad5acd7b73e454283083bbb7c8205ad10e88b0d9d7d793b09fe8226/jiter-0.15.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2143ab06181d2b029eedcb6af3cebe95f11bbac62441781860f98ee9330a6a6", size = 341243, upload-time = "2026-05-19T10:08:55.383Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5d/db/3ff4176b817b8ea33879e71e13d8bc2b0d481a7ed3fe9e080f333d415c16/jiter-0.15.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6eac374c5c975709b69c10f09afd199df74150172156ad10c8d4fd785b7da995", size = 363629, upload-time = "2026-05-19T10:08:56.928Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ab/24/5f8270e0ba9c883582f96f722f8a0b58015c7ce1f8c6d4571cf394e99b6b/jiter-0.15.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b3b3b775e33d3bfaec9899edc526ae97b0da0bf9d071a46124ba419149a414f8", size = 456198, upload-time = "2026-05-19T10:08:58.618Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/45/5b/76fc02b0b5c54c3d18c60653156e2f76fde1816f9b4722db68d6ee2c897e/jiter-0.15.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:eda3071db3346334beae1360b46da4606da57bf3528c167b3c38533afaf9f2c5", size = 373710, upload-time = "2026-05-19T10:09:00.151Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/52/4310821b0ea9277994d3e1f49fc6a4b34e4800caebacb2c0af81da59a454/jiter-0.15.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c6694a173ecabc12eb60efbc0b474464ead1951ff65cd8b1e72100715c64512b", size = 349901, upload-time = "2026-05-19T10:09:01.621Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/93/fe/67648c35b3594fba8854ac64cc8a826d8bcd18324bbdb53d77697c60b6ef/jiter-0.15.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:a254e10b593624d230c365b6d616b22ca0ad65e63a16e6631c2b3466022e6ba8", size = 352438, upload-time = "2026-05-19T10:09:03.216Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cb/28/0a1879d07ad6b3e025a2750027363452ced93c2d16d1c9d4b153ffd51c91/jiter-0.15.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d8d2955167274e15d79a7a020afdd9b39c990eb80b2d89fca695d92dcfdd38ec", size = 388152, upload-time = "2026-05-19T10:09:04.741Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c1/78/46c6f6b56ba85c90021f4afd72ed42f691f8f84daacb5fe27277070e3858/jiter-0.15.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:acf4ee4d1fc55917239fe72972fb292dd773055d05eb040d36f4326e02cc2c0e", size = 517707, upload-time = "2026-05-19T10:09:06.231Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ca/cb/720662d4c88fcad606e826fef5424365527ba43ce4868a479aed8f8c507e/jiter-0.15.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:e7196e56f1cd69af1dbb07dff02dcfb260a50b45a82d409d92a06fedb32473b5", size = 548241, upload-time = "2026-05-19T10:09:08.093Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/60/e3/935b8034fd143f21125c87d51404a9e0e1449186a494405721ff5d1d695e/jiter-0.15.0-cp314-cp314t-win32.whl", hash = "sha256:7f6163c0f10b055245f814dcc59f4818da60dfe72f3e72ab89fc24b6bd5e9c52", size = 207950, upload-time = "2026-05-19T10:09:09.616Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/93/59/984fd9ece895953dad3e0880a650e766f5a2da2c5514f0eafdaaabbeb5f9/jiter-0.15.0-cp314-cp314t-win_amd64.whl", hash = "sha256:980c256edb05b78a111b99c4de3b1d32e31634b867fd1fc2cf726e7b7bba9854", size = 200055, upload-time = "2026-05-19T10:09:11.367Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0e/a4/cf8d779feb133a27a2e3bc833bccb9e13aa332cdf820497ebf72c10ce8c3/jiter-0.15.0-cp314-cp314t-win_arm64.whl", hash = "sha256:66b1880df2d01e206e8339769d1c7c1753bcb653efd6289e203f6f24ebada0c0", size = 191244, upload-time = "2026-05-19T10:09:12.74Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/73/38/505941b2b092fd5bbbd60a52a880db1173f1690ae6751bed3af1c9ddcb4e/jiter-0.15.0-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:631f13a3d04e97d4e083993b10f4b99530e3a10d953e2eb5e196b7dc7f812ce0", size = 303769, upload-time = "2026-05-19T10:09:42.203Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e7/95/a06692b29e77473f286e1ec1f426d3ca44d7b5843be8ad21d7a5f3fcdcc0/jiter-0.15.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:b6c0ffae686c39bf3737be60793783267628783ea42545632c10b291105aee45", size = 305128, upload-time = "2026-05-19T10:09:43.657Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/23/85/7270d7ad41d6061a25b950c6bf91d638bd9aacb113200a8c8d57a055fd67/jiter-0.15.0-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1d54fb5b31dea401a41af3f8a7d2512e9b6a6a005491e6166c7e4ffab9639a9c", size = 340459, upload-time = "2026-05-19T10:09:45.452Z" }, + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c8/8d/302cb2057b7513327b4d575cff6b1d066ee6431a5357fc3f8867cd684406/jiter-0.15.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:54d5d6090cdc1b7c9e780dfb04949a990adb1e301a2fc0bbcee7de4638d33f9a", size = 344469, upload-time = "2026-05-19T10:09:46.864Z" }, +] + [[package]] name = "markupsafe" version = "3.0.3" @@ -1352,6 +1463,15 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7f/d0/578c47dd68152ddddddf31cd7fc67dc30b7cdf639a86275fda821b0d9d98/ruff-0.15.6-py3-none-win_arm64.whl", hash = "sha256:c34de3dd0b0ba203be50ae70f5910b17188556630e2178fd7d79fc030eb0d837", size = 11060497, upload-time = "2026-03-12T23:05:25.968Z" }, ] +[[package]] +name = "sniffio" +version = "1.3.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, +] + [[package]] name = "sortedcontainers" version = "2.4.0" diff --git a/deploy.sh b/deploy.sh new file mode 100755 index 0000000..b0f7091 --- /dev/null +++ b/deploy.sh @@ -0,0 +1,113 @@ +#!/bin/bash +# FileFlash 一键部署脚本 +# 用法: 将整个项目上传到服务器后，在项目根目录执行 bash deploy.sh +# bash deploy.sh --fresh # 强制全量重建镜像（首次部署或大版本更新） + +set -e + +echo "========================================" +echo " FileFlash 生产环境部署" +echo "========================================" + +RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m' + +# ---------- 1. 检查环境 ---------- +echo -e "\n${YELLOW}[1/5] 检查环境...${NC}" + +if ! command -v docker &>/dev/null; then + echo -e "${RED}未安装 Docker，请先安装${NC}" + exit 1 +fi + +if ! docker compose version &>/dev/null; then + echo -e "${RED}需要 Docker Compose v2+${NC}" + exit 1 +fi + +# ---------- 2. 生成密钥 ---------- +echo -e "\n${YELLOW}[2/5] 生成安全密钥...${NC}" + +JWT_SECRET=$(openssl rand -hex 32 2>/dev/null || python3 -c "import secrets;print(secrets.token_hex(32))" 2>/dev/null) +ADMIN_PASSWORD=$(openssl rand -hex 32 2>/dev/null || python3 -c "import secrets;print(secrets.token_hex(32))" 2>/dev/null) +DB_PASSWORD=$(openssl rand -hex 16 2>/dev/null || python3 -c "import secrets;print(secrets.token_hex(16))" 2>/dev/null) +MINIO_PASSWORD=$(openssl rand -hex 16 2>/dev/null || python3 -c "import secrets;print(secrets.token_hex(16))" 2>/dev/null) + +if [ -z "$JWT_SECRET" ] || [ -z "$ADMIN_PASSWORD" ] || [ -z "$DB_PASSWORD" ] || [ -z "$MINIO_PASSWORD" ]; then + echo -e "${RED}无法生成密钥，请手动修改 docker/.env.production${NC}" + exit 1 +fi + +# 替换 .env.production +sed -i "s|^JWT_SECRET_KEY=.*|JWT_SECRET_KEY=$JWT_SECRET|" docker/.env.production +sed -i "s|^DEFAULT_ADMIN_PASSWORD=.*|DEFAULT_ADMIN_PASSWORD=$ADMIN_PASSWORD|" docker/.env.production +sed -i "s|^OBJECT_STORAGE_SECRET_KEY=.*|OBJECT_STORAGE_SECRET_KEY=$MINIO_PASSWORD|" docker/.env.production +sed -i "s|:psgl-ff-db@|:$DB_PASSWORD@|" docker/.env.production + +# 替换 docker-compose.prod.yml +sed -i "s|POSTGRES_PASSWORD: psgl-ff-db|POSTGRES_PASSWORD: $DB_PASSWORD|" docker/docker-compose.prod.yml +sed -i "s|MINIO_ROOT_PASSWORD: minio-admin|MINIO_ROOT_PASSWORD: $MINIO_PASSWORD|" docker/docker-compose.prod.yml +sed -i "s|-password=psgl-ff-db|-password=$DB_PASSWORD|" docker/docker-compose.prod.yml + +echo -e "${GREEN}JWT_SECRET_KEY 已生成${NC}" +echo -e "${GREEN}DEFAULT_ADMIN_PASSWORD 已生成${NC}" +echo -e "${GREEN}数据库密码 已生成${NC}" +echo -e "${GREEN}MinIO 密码 已生成${NC}" + +# ---------- 3. 构建镜像 ---------- +echo -e "\n${YELLOW}[3/5] 构建 Docker 镜像...${NC}" + +BUILD_FLAGS="" +if [ "$1" = "--fresh" ]; then + BUILD_FLAGS="--no-cache" + echo -e "${YELLOW} 首次部署模式: 强制全量构建${NC}" +else + echo " 增量构建 (使用层缓存，加 --fresh 可强制全量)" +fi + +docker compose -f docker/docker-compose.prod.yml build $BUILD_FLAGS + +echo -e "${GREEN}镜像构建完成${NC}" + +# ---------- 4. 启动服务 ---------- +echo -e "\n${YELLOW}[4/5] 启动服务...${NC}" + +docker compose -f docker/docker-compose.prod.yml up -d + +# ---------- 5. 等待就绪 ---------- +echo -e "\n${YELLOW}[5/5] 等待服务就绪...${NC}" + +echo -n "等待服务启动" +for i in $(seq 1 60); do + if curl -s -o /dev/null http://localhost/health 2>/dev/null; then + echo "" + echo -e "${GREEN}服务已就绪${NC}" + break + fi + echo -n "." + sleep 3 +done + +echo "" +echo "========================================" +echo -e "${GREEN} 部署完成!${NC}" +echo "========================================" +echo "" +echo "管理员账号:" +echo " 用户名: administrator" +echo " 密码: $ADMIN_PASSWORD" +echo "" +echo "数据库密码: $DB_PASSWORD" +echo "MinIO 密码: $MINIO_PASSWORD" +echo "" +echo "访问 http://$(hostname -I | awk '{print $1}')" +echo "" +echo "常用命令:" +echo " docker compose -f docker/docker-compose.prod.yml logs -f # 查看日志" +echo " docker compose -f docker/docker-compose.prod.yml down # 停止" +echo " docker compose -f docker/docker-compose.prod.yml restart # 重启" +echo " docker compose -f docker/docker-compose.prod.yml ps # 状态" +echo "" +echo "如需开启 AI Agent:" +echo " 1. 修改 docker/.env.production 中 AGENT_ENABLED=true" +echo " 2. 配置 AGENT_LLM_BASE_URL 和 AGENT_LLM_API_KEY" +echo " 3. docker compose -f docker/docker-compose.prod.yml --profile agent up -d" diff --git a/docker/Dockerfile.backend b/docker/Dockerfile.backend new file mode 100644 index 0000000..5a82bdf --- /dev/null +++ b/docker/Dockerfile.backend @@ -0,0 +1,29 @@ +FROM docker.m.daocloud.io/library/python:3.12-slim + +WORKDIR /app + +# 换成阿里云 debian 镜像源 +RUN sed -i 's/deb.debian.org/mirrors.aliyun.com/g' /etc/apt/sources.list.d/debian.sources + +# 再安装 ffmpeg +RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg && rm -rf /var/lib/apt/lists/* + +# pip 走阿里云 +RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ +RUN pip install --no-cache-dir uv + +# uv 走阿里云 +ENV UV_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ + +# Layer 1: 仅依赖（除非 pyproject.toml / uv.lock 变动，否则缓存命中） +COPY app/pyproject.toml app/uv.lock ./ +RUN uv sync --no-dev --frozen --no-install-project + +# Layer 2: 源码 + 安装项目本身 +COPY app/src ./src +COPY app/README.md ./README.md +RUN uv sync --no-dev --frozen + +EXPOSE 8080 + +CMD ["uv", "run", "uvicorn", "fileflash.main:app", "--host", "0.0.0.0", "--port", "8080"] diff --git a/docker/Dockerfile.frontend b/docker/Dockerfile.frontend new file mode 100644 index 0000000..4eec7cf --- /dev/null +++ b/docker/Dockerfile.frontend @@ -0,0 +1,35 @@ +FROM docker.m.daocloud.io/library/node:22-alpine AS builder + +WORKDIR /build + +RUN npm config set registry https://registry.npmmirror.com +ENV BUN_CONFIG_REGISTRY=https://registry.npmmirror.com + +RUN npm install -g bun + +COPY web/package.json web/bun.lock* ./ + +# Docker 里只需要打包 Vue，不需要 Electron 二进制 +ENV ELECTRON_SKIP_BINARY_DOWNLOAD=1 + +RUN bun install --frozen-lockfile +COPY web ./ + +ARG VITE_BASE_URL=/api/v1 +ENV VITE_BASE_URL=$VITE_BASE_URL + +ARG VITE_ENABLE_MOCKS=false +ENV VITE_ENABLE_MOCKS=$VITE_ENABLE_MOCKS + +# 跳过 vue-tsc 类型检查（CI 已做），仅生产构建 +RUN bun x vite build + +# Stage 2: Serve static files with nginx +FROM docker.m.daocloud.io/library/nginx:1.27-alpine + +COPY --from=builder /build/dist /usr/share/nginx/html +COPY docker/nginx.conf /etc/nginx/conf.d/default.conf + +EXPOSE 80 + +CMD ["nginx", "-g", "daemon off;"] diff --git a/docker/docker-compose.prod.yml b/docker/docker-compose.prod.yml new file mode 100644 index 0000000..7d86c01 --- /dev/null +++ b/docker/docker-compose.prod.yml @@ -0,0 +1,159 @@ +services: + # ========== PostgreSQL + Flyway 迁移 ========== + postgres: + image: docker.m.daocloud.io/library/postgres:16 + container_name: ff-postgres + restart: always + environment: + POSTGRES_USER: admin + POSTGRES_PASSWORD: 539b537cec0417e79748d68b6fa67bce + POSTGRES_DB: fileflash + TZ: Asia/Shanghai + volumes: + - postgres_data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U admin -d fileflash"] + interval: 10s + timeout: 5s + retries: 5 + networks: + - fileflash + + flyway: + image: docker.m.daocloud.io/flyway/flyway:latest + container_name: ff-flyway + command: > + -url=jdbc:postgresql://postgres:5432/fileflash + -user=admin + -password=539b537cec0417e79748d68b6fa67bce + -connectRetries=15 + migrate + volumes: + - ./flyway/migrations:/flyway/sql + depends_on: + postgres: + condition: service_healthy + networks: + - fileflash + + # ========== Redis ========== + redis: + image: docker.m.daocloud.io/library/redis:7-alpine + container_name: ff-redis + restart: always + command: redis-server --appendonly yes + volumes: + - redis_data:/data + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 10s + timeout: 5s + retries: 5 + networks: + - fileflash + + # ========== MinIO 对象存储 ========== + minio: + image: quay.io/minio/minio:latest + container_name: ff-minio + restart: always + environment: + MINIO_ROOT_USER: admin + MINIO_ROOT_PASSWORD: b6d4e8ed2ca164ced8e7b2f0ed252c97 + volumes: + - minio_data:/data + command: server /data + healthcheck: + test: ["CMD", "curl" ,"-f","http://localhost:9000/minio/health/live"] + interval: 30s + timeout: 20s + retries: 3 + networks: + - fileflash + + # ========== FastAPI 后端 ========== + backend: + build: + context: .. + dockerfile: docker/Dockerfile.backend + container_name: ff-backend + restart: always + env_file: + - .env.production + depends_on: + flyway: + condition: service_completed_successfully + redis: + condition: service_healthy + minio: + condition: service_healthy + networks: + - fileflash + + # ========== Worker 后台任务(转码/扫描等) ========== + worker: + build: + context: .. + dockerfile: docker/Dockerfile.backend + container_name: ff-worker + restart: always + env_file: + - .env.production + command: ["uv", "run", "python", "-m", "fileflash.workers.consumer"] + depends_on: + flyway: + condition: service_completed_successfully + redis: + condition: service_healthy + minio: + condition: service_healthy + networks: + - fileflash + + # ========== Agent Worker (开启 AGENT_ENABLED=true 时需要) ========== + agent-worker: + build: + context: .. + dockerfile: docker/Dockerfile.backend + container_name: ff-agent-worker + restart: always + env_file: + - .env.production + command: ["uv", "run", "python", "-m", "fileflash.agents.worker"] + depends_on: + flyway: + condition: service_completed_successfully + redis: + condition: service_healthy + minio: + condition: service_healthy + networks: + - fileflash + profiles: + - agent + + # ========== Vue 前端 + nginx ========== + frontend: + build: + context: .. + dockerfile: docker/Dockerfile.frontend + args: + VITE_BASE_URL: /api/v1 + VITE_ENABLE_MOCKS: 'false' + container_name: ff-frontend + restart: always + ports: + - "80:80" + depends_on: + - backend + networks: + - fileflash + +volumes: + postgres_data: + redis_data: + minio_data: + +networks: + fileflash: + driver: bridge diff --git a/docker/flyway/migrations/V13__agent_builtin_skills.sql b/docker/flyway/migrations/V13__agent_builtin_skills.sql new file mode 100644 index 0000000..da8af9a --- /dev/null +++ b/docker/flyway/migrations/V13__agent_builtin_skills.sql @@ -0,0 +1,47 @@ +INSERT INTO agent_skill ( + skill_key, + name, + description, + triggers_text, + tool_whitelist_json, + plan_template_json, + inputs_schema_json, + outputs_schema_json, + visibility, + owner_user_id +) +VALUES ( + 'builtin:organizeByType', + 'Organize By Type', + 'Analyze the current folder metadata and propose folder creation, moves, and safe renames to organize files by type.', + 'organize files, organize by type, classify files, 整理文件, 文件分类', + '[ + "drive.listFolder", + "drive.createFolder", + "drive.moveFile", + "drive.moveFolder", + "drive.renameFile", + "drive.renameFolder", + "drive.deleteFile", + "drive.deleteFolder" + ]'::jsonb, + '{ + "strategy": "Group direct children by broad file type unless the user asks for a different organization rule.", + "scope": "selected items when present; otherwise current folder direct children only", + "deletePolicy": "delete actions are high risk and require explicit user confirmation" + }'::jsonb, + '{"type":"object"}'::jsonb, + '{"type":"object"}'::jsonb, + 'global', + NULL +) +ON CONFLICT (skill_key) DO UPDATE SET + name = EXCLUDED.name, + description = EXCLUDED.description, + triggers_text = EXCLUDED.triggers_text, + tool_whitelist_json = EXCLUDED.tool_whitelist_json, + plan_template_json = EXCLUDED.plan_template_json, + inputs_schema_json = EXCLUDED.inputs_schema_json, + outputs_schema_json = EXCLUDED.outputs_schema_json, + visibility = EXCLUDED.visibility, + owner_user_id = EXCLUDED.owner_user_id; diff --git a/docker/flyway/migrations/V14__agent_inbox.sql b/docker/flyway/migrations/V14__agent_inbox.sql new file mode 100644 index 0000000..16ab9e0 --- /dev/null +++ b/docker/flyway/migrations/V14__agent_inbox.sql @@ -0,0 +1,49 @@ +-- ========================= +-- Domain: agent inbox +-- ========================= + +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'agent_inbox_role_enum') THEN + CREATE TYPE agent_inbox_role_enum AS ENUM ('agent', 'user'); + END IF; + IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'agent_inbox_kind_enum') THEN + CREATE TYPE agent_inbox_kind_enum AS ENUM ( + 'ask', + 'reply', + 'control.pause', + 'control.resume', + 'control.approve', + 'control.deny', + 'control.skip', + 'control.cancel' + ); + END IF; + IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'agent_inbox_status_enum') THEN + CREATE TYPE agent_inbox_status_enum AS ENUM ('waiting', 'answered', 'timed_out', 'dropped'); + END IF; +END +$$; + +CREATE TABLE IF NOT EXISTS agent_inbox_message ( + inbox_message_id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, + job_id BIGINT NOT NULL, + role agent_inbox_role_enum NOT NULL, + kind agent_inbox_kind_enum NOT NULL, + payload_json JSONB NOT NULL DEFAULT '{}'::jsonb, + reply_to_id BIGINT NULL, + status agent_inbox_status_enum NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + answered_at TIMESTAMP NULL, + CONSTRAINT fk_agent_inbox_message_job + FOREIGN KEY (job_id) REFERENCES background_job(job_id) ON DELETE CASCADE, + CONSTRAINT fk_agent_inbox_message_reply_to + FOREIGN KEY (reply_to_id) REFERENCES agent_inbox_message(inbox_message_id) ON DELETE SET NULL +); + +CREATE INDEX IF NOT EXISTS idx_agent_inbox_message_job_created + ON agent_inbox_message (job_id, created_at); + +CREATE INDEX IF NOT EXISTS idx_agent_inbox_message_job_status + ON agent_inbox_message (job_id, status) + WHERE status IS NOT NULL; diff --git a/docker/flyway/migrations/V15__agent_tool_registry_skills.sql b/docker/flyway/migrations/V15__agent_tool_registry_skills.sql new file mode 100644 index 0000000..f1a9079 --- /dev/null +++ b/docker/flyway/migrations/V15__agent_tool_registry_skills.sql @@ -0,0 +1,19 @@ +UPDATE agent_skill +SET tool_whitelist_json = '[ + "drive.listFolder", + "drive.countFiles", + "drive.searchFiles", + "drive.getFileInfo", + "drive.listRecent", + "drive.statsByCategory", + "drive.findDuplicates", + "drive.createFolder", + "drive.moveFile", + "drive.moveFolder", + "drive.renameFile", + "drive.renameFolder", + "drive.deleteFile", + "drive.deleteFolder" +]'::jsonb, +updated_at = CURRENT_TIMESTAMP +WHERE skill_key = 'builtin:organizeByType'; diff --git a/docker/nginx.conf b/docker/nginx.conf new file mode 100644 index 0000000..2cbd2d5 --- /dev/null +++ b/docker/nginx.conf @@ -0,0 +1,39 @@ +server { + listen 80; + server_name _; + + # Gzip compression + gzip on; + gzip_types text/plain text/css application/json application/javascript text/xml application/xml; + + # Frontend SPA + root /usr/share/nginx/html; + index index.html; + + location / { + try_files $uri $uri/ /index.html; + } + + # Proxy API requests to backend + location /api/ { + proxy_pass http://backend:8080; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + + # Upload support + client_max_body_size 5G; + proxy_read_timeout 900s; + proxy_send_timeout 900s; + + # Disable buffering for streaming (video/audio) + proxy_buffering off; + proxy_request_buffering off; + } + + # Health check + location /health { + proxy_pass http://backend:8080/health; + } +} diff --git a/docs/superpowers/plans/2026-05-24-admin-console-backend.md b/docs/superpowers/plans/2026-05-24-admin-console-backend.md new file mode 100644 index 0000000..ebdb421 --- /dev/null +++ b/docs/superpowers/plans/2026-05-24-admin-console-backend.md @@ -0,0 +1,3174 @@ +# Admin Console Backend Implementation Plan (Plan A) + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 在 `app/src/fileflash/` 下补齐 7 个 `/admin/*` 路由（users / storage / files / moderation / logs / notifications / system），全部接真实 DB；对外形状与现有前端 `web/src/api/*.ts` 期望一致。 + +**Architecture:** 沿用现有 `routers/admin_registration_email_domain_rules.py` 的模式：`routers/admin_<域>.py` 仅做参数编排 + `api_success` 壳；业务逻辑落在 `services/admin/<域>.py`；Pydantic schema 在 `schemas/admin/<域>.py`（`CamelModel` + `by_alias`）。所有路由 `Depends(require_admin)`。 + +**Tech Stack:** FastAPI + SQLAlchemy AsyncSession + Pydantic v2 (CamelModel) + Redis (rate-limit) + pytest/AsyncMock。 + +**Reference Spec:** `docs/superpowers/specs/2026-05-24-admin-console-design.md` + +**前置知识（执行前必读）** + +1. `core/errors.py`：错误抛 `ApiError(status_code, code, message)`；`code` 是数字（=HTTP 状态码），不是字符串符号。message 用英文短语 + 可选中文。 +2. `core/deps.py`：`require_admin` / `get_current_user` / `get_db` 已经存在；新增 service 都按 `get_admin_*_service` 命名加 dep factory。 +3. `schemas/common.py`：`CamelModel` 强制 camelCase by alias、`extra="forbid"`、`str_strip_whitespace=True`；`PaginatedData[T]` / `PageQuery(page, per_page)` 直接复用。 +4. `agents.md §5`：**用户状态外部语义是 `active | suspended`**，但 `enums.UserStatus` 内部只有 `ACTIVE / PENDING_VERIFICATION / LOCKED / DISABLED`，**没有 SUSPENDED**。映射： + - 外部 `'active'` ↔ 内部 `UserStatus.ACTIVE` + - 外部 `'suspended'` ↔ 内部 `UserStatus.DISABLED`（admin 主动停用；`LOCKED` 留给自动锁定） +5. `RegistrationEmailDomainRule` 路由 / 服务 / 测试**保持原样**，不动。 +6. 测试目录：`app/tests/test_admin_<域>_routes.py`（路由层 + stub service）和 `test_admin_<域>_service.py`（service 层 + DummySession 或 in-memory 假实现）。 +7. 执行命令：在 `app/` 目录下 `uv run pytest tests/test_admin_<域>*.py -v`；类型检查不存在，构建用 `uv run python -c "from fileflash.main import app; print(app.title)"`。 + +**对 spec 的两处简化（已与执行者明确）** + +- spec §3.8 Rate Limit：初版**不扫 Redis**。后端只维护一份"已注册限流规则"的静态枚举，`current_usage / blocked_requests` 暂返回 0。原因：现有 `RedisRateLimiter` 仅暴露 `allow()`，无规则注册中心；如需运行时统计，需先扩展 limiter，留作后续。 +- spec §5.3 `notification.broadcast_completed` 事件：本期**不发布**。`Notification` 行写入 + `Log` 写入已足以观测，事件名空挂以待后续 worker 接入。 + +--- + +## File Structure + +**新建** + +``` +app/src/fileflash/ + schemas/admin/__init__.py + schemas/admin/users.py + schemas/admin/storage.py + schemas/admin/files.py + schemas/admin/moderation.py + schemas/admin/logs.py + schemas/admin/notifications.py + schemas/admin/system.py + services/admin/__init__.py + services/admin/users.py + services/admin/storage.py + services/admin/files.py + services/admin/moderation.py + services/admin/logs.py + services/admin/notifications.py + services/admin/system.py + routers/admin_users.py + routers/admin_storage.py + routers/admin_files.py + routers/admin_moderation.py + routers/admin_logs.py + routers/admin_notifications.py + routers/admin_system.py + +app/tests/ + test_admin_users_routes.py + test_admin_users_service.py + test_admin_storage_routes.py + test_admin_storage_service.py + test_admin_files_routes.py + test_admin_files_service.py + test_admin_moderation_routes.py + test_admin_moderation_service.py + test_admin_logs_routes.py + test_admin_notifications_routes.py + test_admin_notifications_service.py + test_admin_system_routes.py +``` + +**修改** + +``` +app/src/fileflash/ + core/deps.py ← 加 7 个 get_admin_*_service factory + routers/__init__.py ← include 7 个新 router + services/messaging.py ← 加新事件常量（可选） + services/rate_limiter.py ← 加 record_blocked() 累计（5.x 任务） +``` + +--- + +## Task 0: 基础设施（admin 包 + 状态映射） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/__init__.py` +- Create: `app/src/fileflash/services/admin/__init__.py` +- Create: `app/src/fileflash/services/admin/_status.py` +- Test: `app/tests/test_admin_user_status_mapping.py` + +- [ ] **Step 0.1: 建空包（schemas/admin, services/admin）** + +```bash +# 创建两个空 __init__.py +``` + +`app/src/fileflash/schemas/admin/__init__.py`： + +```python +"""Admin-only schemas. All requests/responses share CamelModel from schemas.common.""" +``` + +`app/src/fileflash/services/admin/__init__.py`： + +```python +"""Admin-only services. Each module owns one /admin/* surface.""" +``` + +- [ ] **Step 0.2: 写失败测试 — 状态映射** + +`app/tests/test_admin_user_status_mapping.py`： + +```python +from __future__ import annotations + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import UserStatus +from fileflash.services.admin._status import external_to_internal, internal_to_external + + +def test_active_maps_both_ways() -> None: + assert external_to_internal("active") == UserStatus.ACTIVE + assert internal_to_external(UserStatus.ACTIVE) == "active" + + +def test_suspended_maps_to_disabled() -> None: + assert external_to_internal("suspended") == UserStatus.DISABLED + assert internal_to_external(UserStatus.DISABLED) == "suspended" + + +def test_locked_externally_appears_as_suspended() -> None: + # 自动锁定外部仍展示为 suspended（用户不需要区分原因） + assert internal_to_external(UserStatus.LOCKED) == "suspended" + + +def test_pending_verification_passthrough() -> None: + assert internal_to_external(UserStatus.PENDING_VERIFICATION) == "pending_verification" + + +def test_invalid_external_status_raises() -> None: + with pytest.raises(ApiError): + external_to_internal("garbage") +``` + +- [ ] **Step 0.3: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_user_status_mapping.py -v` +Expected: ImportError on `fileflash.services.admin._status` + +- [ ] **Step 0.4: 实现 _status.py** + +`app/src/fileflash/services/admin/_status.py`： + +```python +from __future__ import annotations + +from ...core.errors import ApiError +from ...models.enums import UserStatus + +_EXTERNAL_TO_INTERNAL = { + "active": UserStatus.ACTIVE, + "suspended": UserStatus.DISABLED, +} + +_INTERNAL_TO_EXTERNAL = { + UserStatus.ACTIVE: "active", + UserStatus.DISABLED: "suspended", + UserStatus.LOCKED: "suspended", + UserStatus.PENDING_VERIFICATION: "pending_verification", +} + + +def external_to_internal(value: str) -> UserStatus: + try: + return _EXTERNAL_TO_INTERNAL[value] + except KeyError as exc: + raise ApiError( + status_code=422, + code=422, + message=f"Invalid user status: {value!r}", + ) from exc + + +def internal_to_external(value: UserStatus) -> str: + return _INTERNAL_TO_EXTERNAL.get(value, value.value) + + +__all__ = ["external_to_internal", "internal_to_external"] +``` + +- [ ] **Step 0.5: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_user_status_mapping.py -v` +Expected: 5 passed + +- [ ] **Step 0.6: Commit** + +```bash +cd app && git add src/fileflash/schemas/admin/__init__.py src/fileflash/services/admin/__init__.py src/fileflash/services/admin/_status.py +cd .. && git add app/tests/test_admin_user_status_mapping.py +git commit -m "feat(admin): scaffold admin packages and user status mapping" +``` + +--- + +## Task 1: Admin Users — schemas + service + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/users.py` +- Create: `app/src/fileflash/services/admin/users.py` +- Test: `app/tests/test_admin_users_service.py` + +- [ ] **Step 1.1: 写 schemas** + +`app/src/fileflash/schemas/admin/users.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + +ExternalUserStatus = Literal["active", "suspended", "pending_verification"] + + +class AdminUserItem(CamelModel): + user_id: str + username: str + email: str + role: str + status: ExternalUserStatus + email_verified: bool + email_verified_at: datetime | None = None + storage_limit: int + storage_used: int + usage_percentage: float + last_login_at: datetime | None = None + last_active_at: datetime | None = None + created_at: datetime + + +class ListAdminUsersQuery(PageQuery): + search: str | None = None + status: Literal["active", "suspended"] | None = None + role: Literal["USER", "ADMIN"] | None = None + sort: Literal["username", "createdAt", "storageUsed"] = "createdAt" + order: Literal["asc", "desc"] = "desc" + + +class UpdateUserStatusRequest(CamelModel): + status: Literal["active", "suspended"] + + +class UpdateUserStatusResponse(CamelModel): + user_id: str + status: ExternalUserStatus + updated_at: datetime + + +__all__ = [ + "AdminUserItem", + "ListAdminUsersQuery", + "UpdateUserStatusRequest", + "UpdateUserStatusResponse", +] +``` + +- [ ] **Step 1.2: 写失败测试 — list 默认分页** + +`app/tests/test_admin_users_service.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.enums import UserRole, UserStatus +from fileflash.models.tables_identity import User +from fileflash.schemas.admin.users import ListAdminUsersQuery +from fileflash.services.admin.users import AdminUsersService + + +def _user(**kwargs) -> User: + base = dict( + user_id=1, + username="alice", + email="alice@example.com", + password_hash="x", + role=UserRole.USER, + status=UserStatus.ACTIVE, + email_verified=True, + email_verified_at=datetime.now(UTC), + storage_limit=10 * 1024 * 1024 * 1024, + storage_used=1024, + created_at=datetime.now(UTC), + updated_at=datetime.now(UTC), + ) + base.update(kwargs) + return User(**base) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.refresh = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_list_users_returns_paginated_items() -> None: + session = DummySession() + session.scalar.return_value = 1 + session.scalars.return_value = [_user()] + session.execute.return_value = Mock(all=lambda: []) # no sessions joined + service = AdminUsersService(db=session) # type: ignore[arg-type] + + result = await service.list_users(query=ListAdminUsersQuery()) + + assert result.pagination.total_items == 1 + assert result.items[0].username == "alice" + assert result.items[0].status == "active" +``` + +- [ ] **Step 1.3: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_users_service.py -v` +Expected: ImportError on `fileflash.services.admin.users` + +- [ ] **Step 1.4: 实现 AdminUsersService 的 list_users** + +`app/src/fileflash/services/admin/users.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import and_, func, or_, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.enums import UserRole, UserStatus +from ...models.tables_identity import User, UserSession +from ...schemas.admin.users import ( + AdminUserItem, + ListAdminUsersQuery, + UpdateUserStatusResponse, +) +from ...schemas.common import PaginatedData, PaginationMeta +from ._status import external_to_internal, internal_to_external + + +class AdminUsersService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_users( + self, + *, + query: ListAdminUsersQuery, + ) -> PaginatedData[AdminUserItem]: + statement = select(User).where(User.deleted_at.is_(None)) + + if query.search: + keyword = f"%{query.search.strip().lower()}%" + statement = statement.where( + or_( + func.lower(User.username).like(keyword), + func.lower(User.email).like(keyword), + ) + ) + if query.status: + statement = statement.where(User.status == external_to_internal(query.status)) + if query.role: + statement = statement.where(User.role == UserRole(query.role)) + + sort_column = { + "username": User.username, + "createdAt": User.created_at, + "storageUsed": User.storage_used, + }[query.sort] + statement = statement.order_by( + sort_column.desc() if query.order == "desc" else sort_column.asc() + ) + + total = await self.db.scalar(select(func.count()).select_from(statement.subquery())) + total_items = int(total or 0) + total_pages = max(1, -(-total_items // query.per_page)) + offset = (query.page - 1) * query.per_page + + rows = list( + await self.db.scalars(statement.offset(offset).limit(query.per_page)) + ) + + last_seen_map = await self._collect_last_seen([row.user_id for row in rows]) + items = [self._to_item(row, last_seen_map.get(row.user_id)) for row in rows] + + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total_items, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def _collect_last_seen(self, user_ids: list[int]) -> dict[int, datetime]: + if not user_ids: + return {} + rows = await self.db.execute( + select(UserSession.user_id, func.max(UserSession.last_seen_at)) + .where( + and_( + UserSession.user_id.in_(user_ids), + UserSession.revoked_at.is_(None), + ) + ) + .group_by(UserSession.user_id) + ) + return {user_id: seen for user_id, seen in rows.all()} + + @staticmethod + def _to_item(row: User, last_active_at: datetime | None) -> AdminUserItem: + limit = max(row.storage_limit, 1) + return AdminUserItem( + user_id=str(row.user_id), + username=row.username, + email=row.email, + role=row.role.value if hasattr(row.role, "value") else str(row.role), + status=internal_to_external(row.status), # type: ignore[arg-type] + email_verified=row.email_verified, + email_verified_at=row.email_verified_at, + storage_limit=row.storage_limit, + storage_used=row.storage_used, + usage_percentage=round((row.storage_used / limit) * 100, 2), + last_login_at=row.last_login_at, + last_active_at=last_active_at, + created_at=row.created_at, + ) + + +__all__ = ["AdminUsersService"] +``` + +- [ ] **Step 1.5: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_users_service.py -v` +Expected: 1 passed + +- [ ] **Step 1.6: 写失败测试 — set_status（成功 + 防误锁 + 404）** + +追加到 `app/tests/test_admin_users_service.py`： + +```python +@pytest.mark.asyncio +async def test_set_status_user_not_found() -> None: + session = DummySession() + session.get.return_value = None + service = AdminUsersService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.set_status(user_id=999, external_status="suspended") + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_set_status_suspend_last_admin_blocked() -> None: + session = DummySession() + admin = _user(user_id=2, role=UserRole.ADMIN, status=UserStatus.ACTIVE) + session.get.return_value = admin + # 当前活跃 admin 仅 1 名 => 不允许 suspend + session.scalar.return_value = 1 + service = AdminUsersService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.set_status(user_id=2, external_status="suspended") + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_set_status_suspends_user_and_revokes_sessions() -> None: + session = DummySession() + target = _user(user_id=3) + session.get.return_value = target + session.scalar.return_value = 5 # active admins enough + service = AdminUsersService(db=session) # type: ignore[arg-type] + + result = await service.set_status(user_id=3, external_status="suspended") + + assert result.status == "suspended" + assert target.status == UserStatus.DISABLED + session.commit.assert_awaited() +``` + +- [ ] **Step 1.7: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_users_service.py -v` +Expected: AttributeError on `service.set_status` + +- [ ] **Step 1.8: 实现 set_status** + +追加到 `app/src/fileflash/services/admin/users.py`（在类内）： + +```python + async def set_status( + self, + *, + user_id: int, + external_status: str, + ) -> UpdateUserStatusResponse: + target = await self.db.get(User, user_id) + if target is None or target.deleted_at is not None: + raise ApiError(status_code=404, code=404, message="User not found") + + new_internal = external_to_internal(external_status) + + if ( + new_internal == UserStatus.DISABLED + and target.role == UserRole.ADMIN + and target.status == UserStatus.ACTIVE + ): + remaining = await self.db.scalar( + select(func.count(User.user_id)).where( + and_( + User.role == UserRole.ADMIN, + User.status == UserStatus.ACTIVE, + User.user_id != user_id, + User.deleted_at.is_(None), + ) + ) + ) + if int(remaining or 0) == 0: + raise ApiError( + status_code=409, + code=409, + message="Cannot suspend the last active admin", + ) + + target.status = new_internal + target.updated_at = datetime.now(UTC) + + if new_internal == UserStatus.DISABLED: + await self.db.execute( + UserSession.__table__.update() + .where( + and_( + UserSession.user_id == user_id, + UserSession.revoked_at.is_(None), + ) + ) + .values(revoked_at=datetime.now(UTC)) + ) + + await self.db.commit() + await self.db.refresh(target) + + return UpdateUserStatusResponse( + user_id=str(target.user_id), + status=internal_to_external(target.status), # type: ignore[arg-type] + updated_at=target.updated_at, + ) +``` + +- [ ] **Step 1.9: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_users_service.py -v` +Expected: 4 passed + +- [ ] **Step 1.10: Commit** + +```bash +git add app/src/fileflash/schemas/admin/users.py app/src/fileflash/services/admin/users.py app/tests/test_admin_users_service.py +git commit -m "feat(admin): users service with list and set_status (last-admin guard)" +``` + +--- + +## Task 2: Admin Users — router + deps + 注册 + +**Files:** + +- Create: `app/src/fileflash/routers/admin_users.py` +- Modify: `app/src/fileflash/core/deps.py` +- Modify: `app/src/fileflash/routers/__init__.py` +- Test: `app/tests/test_admin_users_routes.py` + +- [ ] **Step 2.1: 写失败测试 — 403 给非 admin / 200 给 admin** + +`app/tests/test_admin_users_routes.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_users import ( + router as admin_users_router, + get_admin_users_service, +) +from fileflash.schemas.admin.users import ( + AdminUserItem, + UpdateUserStatusResponse, +) +from fileflash.schemas.common import PaginatedData, PaginationMeta + + +class StubService: + async def list_users(self, *, query): # noqa: ANN001 + item = AdminUserItem( + user_id="1", username="alice", email="a@x.com", role="USER", + status="active", email_verified=True, email_verified_at=None, + storage_limit=1024, storage_used=0, usage_percentage=0.0, + last_login_at=None, last_active_at=None, created_at=datetime.now(UTC), + ) + return PaginatedData( + items=[item], + pagination=PaginationMeta( + total_items=1, total_pages=1, + per_page=query.per_page, current_page=query.page, + has_prev=False, has_next=False, + ), + ) + + async def set_status(self, *, user_id, external_status): # noqa: ANN001 + return UpdateUserStatusResponse( + user_id=str(user_id), status=external_status, + updated_at=datetime.now(UTC), + ) + + +def _client(admin: bool) -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(admin_users_router, prefix="/api/v1") + app.dependency_overrides[get_admin_users_service] = lambda: StubService() + if admin: + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + else: + async def _deny(): + raise ApiError(status_code=403, code=403, message="forbidden") + app.dependency_overrides[require_admin] = _deny + return TestClient(app) + + +def test_admin_can_list_users() -> None: + with _client(admin=True) as c: + resp = c.get("/api/v1/admin/users") + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["data"]["items"][0]["username"] == "alice" + + +def test_non_admin_gets_403() -> None: + with _client(admin=False) as c: + resp = c.get("/api/v1/admin/users") + assert resp.status_code == 403 + + +def test_admin_can_patch_status() -> None: + with _client(admin=True) as c: + resp = c.patch("/api/v1/admin/users/42/status", json={"status": "suspended"}) + assert resp.status_code == 200 + assert resp.json()["data"]["status"] == "suspended" +``` + +- [ ] **Step 2.2: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_users_routes.py -v` +Expected: ImportError on `fileflash.routers.admin_users` + +- [ ] **Step 2.3: 实现 router + 局部 dep factory** + +`app/src/fileflash/routers/admin_users.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.users import ( + ListAdminUsersQuery, + UpdateUserStatusRequest, +) +from ..services.admin.users import AdminUsersService + +router = APIRouter(prefix="/admin/users", tags=["admin"]) + + +def get_admin_users_service(db: AsyncSession = Depends(get_db)) -> AdminUsersService: + return AdminUsersService(db=db) + + +@router.get("") +async def list_admin_users( + query: ListAdminUsersQuery = Depends(), + _: User = Depends(require_admin), + service: AdminUsersService = Depends(get_admin_users_service), +): + data = await service.list_users(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Users fetched") + + +@router.patch("/{user_id}/status") +async def update_admin_user_status( + user_id: int, + payload: UpdateUserStatusRequest, + _: User = Depends(require_admin), + service: AdminUsersService = Depends(get_admin_users_service), +): + result = await service.set_status(user_id=user_id, external_status=payload.status) + return api_success(data=result.model_dump(by_alias=True), message="Status updated") + + +__all__ = ["router", "get_admin_users_service"] +``` + +- [ ] **Step 2.4: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_users_routes.py -v` +Expected: 3 passed + +- [ ] **Step 2.5: 注册 router 到 api_router** + +Modify `app/src/fileflash/routers/__init__.py`：在 imports 区追加 + +```python +from .admin_users import router as admin_users_router +``` + +在 `include_router` 区追加 + +```python +api_router.include_router(admin_users_router) +``` + +- [ ] **Step 2.6: 冒烟启动** + +Run: `cd app && uv run python -c "from fileflash.main import app; print([r.path for r in app.routes if '/admin/users' in r.path])"` +Expected: prints `['/api/v1/admin/users', '/api/v1/admin/users/{user_id}/status']` + +- [ ] **Step 2.7: Commit** + +```bash +git add app/src/fileflash/routers/admin_users.py app/src/fileflash/routers/__init__.py app/tests/test_admin_users_routes.py +git commit -m "feat(admin): /admin/users list + /admin/users/{id}/status routes" +``` + +--- + +## Task 3: Admin Storage（summary / users / quota / usage-trend） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/storage.py` +- Create: `app/src/fileflash/services/admin/storage.py` +- Create: `app/src/fileflash/routers/admin_storage.py` +- Modify: `app/src/fileflash/routers/__init__.py` +- Test: `app/tests/test_admin_storage_service.py`, `app/tests/test_admin_storage_routes.py` + +- [ ] **Step 3.1: 写 schemas** + +`app/src/fileflash/schemas/admin/storage.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + + +class AdminStorageSummary(CamelModel): + storage_used: int + storage_limit: int + storage_percentage: float + file_count: int + user_count: int + updated_at: datetime + + +class AdminStorageUserItem(CamelModel): + user_id: str + username: str + email: str + storage_limit: int + storage_used: int + usage_percentage: float + updated_at: datetime + + +class ListStorageUsersQuery(PageQuery): + sort: Literal["storageUsed", "usagePercentage", "username"] = "storageUsed" + order: Literal["asc", "desc"] = "desc" + + +class UpdateQuotaRequest(CamelModel): + storage_limit: int = Field(ge=0) + + +class UpdateQuotaResponse(CamelModel): + user_id: str + storage_limit: int + storage_used: int + usage_percentage: float + updated_at: datetime + + +class UsageTrendQuery(CamelModel): + days: Literal[7, 14, 30] = 7 + + +class UsageTrendPoint(CamelModel): + date: str + used: int + + +class UsageTrendResponse(CamelModel): + trends: list[UsageTrendPoint] + is_estimated: bool = False + + +__all__ = [ + "AdminStorageSummary", + "AdminStorageUserItem", + "ListStorageUsersQuery", + "UpdateQuotaRequest", + "UpdateQuotaResponse", + "UsageTrendQuery", + "UsageTrendPoint", + "UsageTrendResponse", +] +``` + +- [ ] **Step 3.2: 写失败测试 — summary / quota / usage-trend** + +`app/tests/test_admin_storage_service.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.models.tables_identity import User +from fileflash.models.enums import UserRole, UserStatus +from fileflash.schemas.admin.storage import ListStorageUsersQuery, UsageTrendQuery +from fileflash.services.admin.storage import AdminStorageService + + +def _user(**kwargs) -> User: + base = dict( + user_id=1, username="bob", email="b@x.com", password_hash="x", + role=UserRole.USER, status=UserStatus.ACTIVE, + email_verified=True, email_verified_at=datetime.now(UTC), + storage_limit=10 * 1024 * 1024 * 1024, storage_used=2 * 1024 * 1024 * 1024, + created_at=datetime.now(UTC), updated_at=datetime.now(UTC), + ) + base.update(kwargs) + return User(**base) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.refresh = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_summary_aggregates_users_and_files() -> None: + session = DummySession() + # 三次 scalar 调用：sum(storage_used), sum(storage_limit), count(files), count(users) + session.scalar.side_effect = [ + 2 * 1024 * 1024 * 1024, # storage_used sum + 10 * 1024 * 1024 * 1024, # storage_limit sum + 42, # file count + 7, # user count + ] + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.summary() + + assert result.storage_used == 2 * 1024 * 1024 * 1024 + assert result.user_count == 7 + assert result.file_count == 42 + assert round(result.storage_percentage, 2) == 20.0 + + +@pytest.mark.asyncio +async def test_update_quota_rejects_below_usage() -> None: + session = DummySession() + target = _user(storage_used=5 * 1024 * 1024 * 1024) + session.get.return_value = target + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.update_quota(user_id=1, new_limit=1 * 1024 * 1024 * 1024) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_update_quota_success_updates_and_commits() -> None: + session = DummySession() + target = _user(storage_used=1 * 1024 * 1024 * 1024) + session.get.return_value = target + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.update_quota(user_id=1, new_limit=20 * 1024 * 1024 * 1024) + + assert result.storage_limit == 20 * 1024 * 1024 * 1024 + assert target.storage_limit == 20 * 1024 * 1024 * 1024 + session.commit.assert_awaited() + + +@pytest.mark.asyncio +async def test_usage_trend_returns_n_points() -> None: + session = DummySession() + session.scalar.return_value = 2 * 1024 * 1024 * 1024 # T_now + session.execute.return_value = Mock(all=lambda: []) # no events + service = AdminStorageService(db=session, redis=None) # type: ignore[arg-type] + + result = await service.usage_trend(query=UsageTrendQuery(days=7)) + + assert len(result.trends) == 7 + assert result.is_estimated is True # no events => 估算 +``` + +- [ ] **Step 3.3: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_storage_service.py -v` +Expected: ImportError + +- [ ] **Step 3.4: 实现 AdminStorageService** + +`app/src/fileflash/services/admin/storage.py`： + +```python +from __future__ import annotations + +import json +from datetime import UTC, date, datetime, timedelta +from typing import Any + +from redis.asyncio import Redis +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.tables_audit_security import Log +from ...models.tables_identity import User +from ...models.tables_storage import File +from ...schemas.admin.storage import ( + AdminStorageSummary, + AdminStorageUserItem, + ListStorageUsersQuery, + UpdateQuotaResponse, + UsageTrendPoint, + UsageTrendQuery, + UsageTrendResponse, +) +from ...schemas.common import PaginatedData, PaginationMeta + +_STORAGE_EVENT_OPS = ("file.created", "file.deleted", "file.restored") +_TREND_CACHE_TTL = 300 + + +class AdminStorageService: + def __init__(self, db: AsyncSession, redis: Redis | None) -> None: + self.db = db + self.redis = redis + + async def summary(self) -> AdminStorageSummary: + used_sum = await self.db.scalar(select(func.coalesce(func.sum(User.storage_used), 0))) + limit_sum = await self.db.scalar(select(func.coalesce(func.sum(User.storage_limit), 0))) + file_count = await self.db.scalar(select(func.count(File.file_id))) + user_count = await self.db.scalar( + select(func.count(User.user_id)).where(User.deleted_at.is_(None)) + ) + + used = int(used_sum or 0) + limit = int(limit_sum or 0) + return AdminStorageSummary( + storage_used=used, + storage_limit=limit, + storage_percentage=round((used / limit) * 100, 2) if limit else 0.0, + file_count=int(file_count or 0), + user_count=int(user_count or 0), + updated_at=datetime.now(UTC), + ) + + async def list_storage_users( + self, + *, + query: ListStorageUsersQuery, + ) -> PaginatedData[AdminStorageUserItem]: + statement = select(User).where(User.deleted_at.is_(None)) + + sort_column = { + "storageUsed": User.storage_used, + "username": User.username, + "usagePercentage": User.storage_used, # 近似按 storage_used 排序 + }[query.sort] + statement = statement.order_by( + sort_column.desc() if query.order == "desc" else sort_column.asc() + ) + + total = await self.db.scalar(select(func.count()).select_from(statement.subquery())) + total_items = int(total or 0) + total_pages = max(1, -(-total_items // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + + items = [self._to_item(row) for row in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total_items, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def update_quota(self, *, user_id: int, new_limit: int) -> UpdateQuotaResponse: + target = await self.db.get(User, user_id, with_for_update=True) + if target is None or target.deleted_at is not None: + raise ApiError(status_code=404, code=404, message="User not found") + + if new_limit < target.storage_used: + raise ApiError( + status_code=409, + code=409, + message="New quota cannot be below current usage", + ) + + target.storage_limit = new_limit + target.updated_at = datetime.now(UTC) + await self.db.commit() + await self.db.refresh(target) + return UpdateQuotaResponse( + user_id=str(target.user_id), + storage_limit=target.storage_limit, + storage_used=target.storage_used, + usage_percentage=round((target.storage_used / max(target.storage_limit, 1)) * 100, 2), + updated_at=target.updated_at, + ) + + async def usage_trend(self, *, query: UsageTrendQuery) -> UsageTrendResponse: + cached = await self._cache_get(query.days) + if cached is not None: + return cached + + t_now = int( + await self.db.scalar(select(func.coalesce(func.sum(User.storage_used), 0))) or 0 + ) + + cutoff = datetime.now(UTC) - timedelta(days=query.days) + rows = await self.db.execute( + select(Log.operation, Log.metadata_payload, Log.performed_at) + .where(Log.operation.in_(_STORAGE_EVENT_OPS)) + .where(Log.performed_at >= cutoff) + ) + events = rows.all() + + deltas: dict[date, int] = {} + for op, metadata, performed_at in events: + if not performed_at: + continue + size = int((metadata or {}).get("size") or 0) + sign = 1 if op == "file.created" or op == "file.restored" else -1 + day = performed_at.astimezone(UTC).date() + deltas[day] = deltas.get(day, 0) + sign * size + + today = datetime.now(UTC).date() + points: list[UsageTrendPoint] = [] + running = t_now + for offset in range(query.days): + d = today - timedelta(days=offset) + points.append(UsageTrendPoint(date=d.isoformat(), used=max(running, 0))) + running -= deltas.get(d, 0) + points.reverse() + + response = UsageTrendResponse( + trends=points, + is_estimated=not events, + ) + await self._cache_set(query.days, response) + return response + + @staticmethod + def _to_item(row: User) -> AdminStorageUserItem: + limit = max(row.storage_limit, 1) + return AdminStorageUserItem( + user_id=str(row.user_id), + username=row.username, + email=row.email, + storage_limit=row.storage_limit, + storage_used=row.storage_used, + usage_percentage=round((row.storage_used / limit) * 100, 2), + updated_at=row.updated_at, + ) + + async def _cache_get(self, days: int) -> UsageTrendResponse | None: + if self.redis is None: + return None + try: + raw = await self.redis.get(f"admin:storage:trend:{days}") + except Exception: + return None + if not raw: + return None + return UsageTrendResponse.model_validate_json(raw) + + async def _cache_set(self, days: int, payload: UsageTrendResponse) -> None: + if self.redis is None: + return + try: + await self.redis.setex( + f"admin:storage:trend:{days}", + _TREND_CACHE_TTL, + payload.model_dump_json(by_alias=True), + ) + except Exception: + return + + +__all__ = ["AdminStorageService"] +``` + +- [ ] **Step 3.5: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_storage_service.py -v` +Expected: 4 passed + +- [ ] **Step 3.6: 实现 router + 注册** + +`app/src/fileflash/routers/admin_storage.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import get_rate_limiter, require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.storage import ( + ListStorageUsersQuery, + UpdateQuotaRequest, + UsageTrendQuery, +) +from ..services.admin.storage import AdminStorageService +from ..services.rate_limiter import RedisRateLimiter + +router = APIRouter(prefix="/admin/storage", tags=["admin"]) + + +def get_admin_storage_service( + db: AsyncSession = Depends(get_db), + rate_limiter: RedisRateLimiter = Depends(get_rate_limiter), +) -> AdminStorageService: + # 复用现有 RedisRateLimiter 持有的连接，仅当存在时启用缓存 + return AdminStorageService(db=db, redis=getattr(rate_limiter, "_redis", None)) + + +@router.get("/summary") +async def get_admin_storage_summary( + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + data = await service.summary() + return api_success(data=data.model_dump(by_alias=True), message="Storage summary fetched") + + +@router.get("/users") +async def list_admin_storage_users( + query: ListStorageUsersQuery = Depends(), + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + data = await service.list_storage_users(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Storage users fetched") + + +@router.patch("/users/{user_id}/quota") +async def update_admin_user_quota( + user_id: int, + payload: UpdateQuotaRequest, + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + result = await service.update_quota(user_id=user_id, new_limit=payload.storage_limit) + return api_success(data=result.model_dump(by_alias=True), message="Quota updated") + + +@router.get("/usage-trend") +async def get_storage_usage_trend( + query: UsageTrendQuery = Depends(), + _: User = Depends(require_admin), + service: AdminStorageService = Depends(get_admin_storage_service), +): + result = await service.usage_trend(query=query) + return api_success(data=result.model_dump(by_alias=True), message="Usage trend fetched") + + +__all__ = ["router", "get_admin_storage_service"] +``` + +- [ ] **Step 3.7: 写路由测试** + +`app/tests/test_admin_storage_routes.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_storage import router as storage_router, get_admin_storage_service +from fileflash.schemas.admin.storage import ( + AdminStorageSummary, + UpdateQuotaResponse, + UsageTrendPoint, + UsageTrendResponse, +) + + +class StubService: + async def summary(self): + return AdminStorageSummary( + storage_used=1000, storage_limit=10000, storage_percentage=10.0, + file_count=3, user_count=2, updated_at=datetime.now(UTC), + ) + + async def list_storage_users(self, *, query): # noqa: ANN001 + return SimpleNamespace(model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, "totalPages": 1, + "perPage": query.per_page, "currentPage": query.page, + "hasPrev": False, "hasNext": False, + }, + }) + + async def update_quota(self, *, user_id, new_limit): # noqa: ANN001 + return UpdateQuotaResponse( + user_id=str(user_id), storage_limit=new_limit, + storage_used=0, usage_percentage=0.0, updated_at=datetime.now(UTC), + ) + + async def usage_trend(self, *, query): # noqa: ANN001 + return UsageTrendResponse( + trends=[UsageTrendPoint(date="2026-05-24", used=1)], + is_estimated=False, + ) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(storage_router, prefix="/api/v1") + app.dependency_overrides[get_admin_storage_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_summary_returns_camel_case() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/storage/summary") + assert resp.status_code == 200 + assert "storageUsed" in resp.json()["data"] + + +def test_update_quota_passes_storage_limit() -> None: + with _client() as c: + resp = c.patch("/api/v1/admin/storage/users/7/quota", json={"storageLimit": 1024}) + assert resp.status_code == 200 + assert resp.json()["data"]["storageLimit"] == 1024 + + +def test_usage_trend_default_days() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/storage/usage-trend") + assert resp.status_code == 200 + assert len(resp.json()["data"]["trends"]) == 1 +``` + +- [ ] **Step 3.8: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_storage_routes.py -v` +Expected: 3 passed + +- [ ] **Step 3.9: 注册 router** + +Modify `app/src/fileflash/routers/__init__.py`： + +```python +from .admin_storage import router as admin_storage_router +# ... +api_router.include_router(admin_storage_router) +``` + +- [ ] **Step 3.10: Commit** + +```bash +git add app/src/fileflash/schemas/admin/storage.py \ + app/src/fileflash/services/admin/storage.py \ + app/src/fileflash/routers/admin_storage.py \ + app/src/fileflash/routers/__init__.py \ + app/tests/test_admin_storage_service.py \ + app/tests/test_admin_storage_routes.py +git commit -m "feat(admin): /admin/storage summary, users, quota, usage-trend" +``` + +--- + +## Task 4: Admin Files（list + rescan） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/files.py` +- Create: `app/src/fileflash/services/admin/files.py` +- Create: `app/src/fileflash/routers/admin_files.py` +- Test: `app/tests/test_admin_files_service.py`, `app/tests/test_admin_files_routes.py` +- Modify: `app/src/fileflash/routers/__init__.py` + +- [ ] **Step 4.1: 写 schemas** + +`app/src/fileflash/schemas/admin/files.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from ..common import CamelModel, PageQuery + + +VirusStatus = Literal["clean", "pending", "flagged"] + + +class AdminFileAuditItem(CamelModel): + id: str + name: str + size: int + mime_type: str + hash: str + virus_status: VirusStatus + is_shared: bool + owner_name: str + updated_at: datetime + created_at: datetime + + +class ListAdminFilesQuery(PageQuery): + search: str | None = None + virus_status: VirusStatus | None = None + owner_id: str | None = None + mime_type: str | None = None + sort: Literal["name", "size", "createdAt", "updatedAt"] = "updatedAt" + order: Literal["asc", "desc"] = "desc" + + +class RescanResponse(CamelModel): + file_id: str + virus_status: VirusStatus + scanned_at: datetime + + +__all__ = ["AdminFileAuditItem", "ListAdminFilesQuery", "RescanResponse", "VirusStatus"] +``` + +- [ ] **Step 4.2: 写失败测试 — list + rescan** + +`app/tests/test_admin_files_service.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.schemas.admin.files import ListAdminFilesQuery +from fileflash.services.admin.files import AdminFilesService + + +class DummyPublisher: + def __init__(self) -> None: + self.calls: list[tuple[str, dict]] = [] + + async def publish(self, event_name: str, payload: dict) -> None: # noqa: ANN001 + self.calls.append((event_name, payload)) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.execute = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + + +@pytest.mark.asyncio +async def test_list_returns_paginated_empty_when_no_files() -> None: + session = DummySession() + session.scalar.return_value = 0 + session.execute.return_value = Mock(all=lambda: []) + service = AdminFilesService(db=session, publisher=DummyPublisher()) # type: ignore[arg-type] + + result = await service.list_files(query=ListAdminFilesQuery()) + + assert result.items == [] + assert result.pagination.total_items == 0 + + +@pytest.mark.asyncio +async def test_rescan_missing_file_returns_404() -> None: + session = DummySession() + session.get.return_value = None + service = AdminFilesService(db=session, publisher=DummyPublisher()) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.request_rescan(file_id=1, requested_by=99) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_rescan_inserts_scan_record_and_publishes_event() -> None: + session = DummySession() + file_row = Mock(file_id=1, object_id=2) + session.get.return_value = file_row + publisher = DummyPublisher() + service = AdminFilesService(db=session, publisher=publisher) # type: ignore[arg-type] + + result = await service.request_rescan(file_id=1, requested_by=99) + + assert result.virus_status == "pending" + assert publisher.calls and publisher.calls[0][0] == "files.rescan_requested" + session.commit.assert_awaited() +``` + +- [ ] **Step 4.3: 跑测试验证失败** + +Run: `cd app && uv run pytest tests/test_admin_files_service.py -v` +Expected: ImportError + +- [ ] **Step 4.4: 实现 AdminFilesService** + +`app/src/fileflash/services/admin/files.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any, Protocol + +from sqlalchemy import and_, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.enums import FileStatus, ScanResult +from ...models.tables_audit_security import ObjectScanResult +from ...models.tables_identity import User +from ...models.tables_storage import File, StorageObject +from ...schemas.admin.files import ( + AdminFileAuditItem, + ListAdminFilesQuery, + RescanResponse, + VirusStatus, +) +from ...schemas.common import PaginatedData, PaginationMeta + + +class EventPublisherProtocol(Protocol): + async def publish(self, event_name: str, payload: dict[str, Any]) -> None: ... + + +_VIRUS_STATUS_MAP: dict[ScanResult, VirusStatus] = { + ScanResult.CLEAN: "clean", + ScanResult.PENDING: "pending", + ScanResult.INFECTED: "flagged", + ScanResult.BLOCKED: "flagged", + ScanResult.FAILED: "pending", +} + + +class AdminFilesService: + def __init__(self, db: AsyncSession, publisher: EventPublisherProtocol) -> None: + self.db = db + self.publisher = publisher + + async def list_files( + self, + *, + query: ListAdminFilesQuery, + ) -> PaginatedData[AdminFileAuditItem]: + # 最近一次扫描结果 sub-query + latest_scan = ( + select( + ObjectScanResult.object_id, + func.max(ObjectScanResult.scanned_at).label("scanned_at"), + ) + .group_by(ObjectScanResult.object_id) + .subquery() + ) + + statement = ( + select(File, StorageObject, User, ObjectScanResult) + .join(StorageObject, File.storage_object_id == StorageObject.object_id) + .join(User, File.owner_id == User.user_id) + .join( + latest_scan, + latest_scan.c.object_id == StorageObject.object_id, + isouter=True, + ) + .join( + ObjectScanResult, + and_( + ObjectScanResult.object_id == latest_scan.c.object_id, + ObjectScanResult.scanned_at == latest_scan.c.scanned_at, + ), + isouter=True, + ) + .where(File.status == FileStatus.ACTIVE) + .where(File.deleted_at.is_(None)) + ) + + if query.search: + kw = f"%{query.search.strip().lower()}%" + statement = statement.where(func.lower(File.file_name).like(kw)) + if query.owner_id: + statement = statement.where(File.owner_id == int(query.owner_id)) + if query.mime_type: + statement = statement.where(File.mime_type == query.mime_type) + if query.virus_status: + wanted = [k for k, v in _VIRUS_STATUS_MAP.items() if v == query.virus_status] + statement = statement.where(ObjectScanResult.result.in_(wanted)) + + sort_column = { + "name": File.file_name, + "size": File.file_size, + "createdAt": File.created_at, + "updatedAt": File.updated_at, + }[query.sort] + statement = statement.order_by( + sort_column.desc() if query.order == "desc" else sort_column.asc() + ) + + count_stmt = select(func.count()).select_from(statement.subquery()) + total = int(await self.db.scalar(count_stmt) or 0) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + + rows = (await self.db.execute(statement.offset(offset).limit(query.per_page))).all() + items = [self._to_item(file_, obj, owner, scan) for file_, obj, owner, scan in rows] + + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def request_rescan(self, *, file_id: int, requested_by: int) -> RescanResponse: + file_row = await self.db.get(File, file_id) + if file_row is None or file_row.deleted_at is not None or file_row.status != FileStatus.ACTIVE: + raise ApiError(status_code=404, code=404, message="File not found") + + now = datetime.now(UTC) + scan = ObjectScanResult( + object_id=file_row.storage_object_id, + scan_type="virus", + result=ScanResult.PENDING, + details={"requestedBy": requested_by}, + scanned_at=now, + created_at=now, + ) + self.db.add(scan) + await self.db.commit() + + await self.publisher.publish( + "files.rescan_requested", + { + "fileId": str(file_id), + "objectId": str(file_row.storage_object_id), + "requestedBy": requested_by, + }, + ) + + return RescanResponse( + file_id=str(file_id), + virus_status="pending", + scanned_at=now, + ) + + @staticmethod + def _to_item( + file_: File, + obj: StorageObject, + owner: User, + scan: ObjectScanResult | None, + ) -> AdminFileAuditItem: + return AdminFileAuditItem( + id=str(file_.file_id), + name=file_.file_name, + size=file_.file_size, + mime_type=file_.mime_type or obj.content_type or "application/octet-stream", + hash=(obj.object_hash or "")[:16], + virus_status=_VIRUS_STATUS_MAP.get(scan.result, "pending") if scan else "pending", + is_shared=False, + owner_name=owner.username, + updated_at=file_.updated_at, + created_at=file_.created_at, + ) + + +__all__ = ["AdminFilesService"] +``` + +> **校对依据：** `File` 字段 `file_name / file_size / mime_type / storage_object_id / owner_id / status / deleted_at`；`StorageObject` 字段 `object_hash / content_type / object_size`；`ScanResult` 枚举 `PENDING / CLEAN / INFECTED / BLOCKED / FAILED`。已与 `app/src/fileflash/models/tables_storage.py` 和 `enums.py` 核对一致。 + +- [ ] **Step 4.5: 写 rescan service 测试中 file_row 的 mock 字段对齐** + +更新 Step 4.2 中 `test_rescan_inserts_scan_record_and_publishes_event`： + +```python +@pytest.mark.asyncio +async def test_rescan_inserts_scan_record_and_publishes_event() -> None: + from fileflash.models.enums import FileStatus + session = DummySession() + file_row = Mock( + file_id=1, storage_object_id=2, + deleted_at=None, status=FileStatus.ACTIVE, + ) + session.get.return_value = file_row + publisher = DummyPublisher() + service = AdminFilesService(db=session, publisher=publisher) # type: ignore[arg-type] + + result = await service.request_rescan(file_id=1, requested_by=99) + + assert result.virus_status == "pending" + assert publisher.calls and publisher.calls[0][0] == "files.rescan_requested" + session.commit.assert_awaited() +``` + +- [ ] **Step 4.6: 跑测试验证通过** + +Run: `cd app && uv run pytest tests/test_admin_files_service.py -v` +Expected: 3 passed + +- [ ] **Step 4.7: 实现 router + 测试 + 注册** + +`app/src/fileflash/routers/admin_files.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import get_event_publisher, require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.files import ListAdminFilesQuery +from ..services.admin.files import AdminFilesService +from ..services.messaging import InProcessAuthEventPublisher + +router = APIRouter(prefix="/admin/files", tags=["admin"]) + + +def get_admin_files_service( + db: AsyncSession = Depends(get_db), + publisher: InProcessAuthEventPublisher = Depends(get_event_publisher), +) -> AdminFilesService: + return AdminFilesService(db=db, publisher=publisher) + + +@router.get("") +async def list_admin_files( + query: ListAdminFilesQuery = Depends(), + _: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + data = await service.list_files(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Files fetched") + + +@router.post("/{file_id}/rescan") +async def rescan_admin_file( + file_id: int, + admin: User = Depends(require_admin), + service: AdminFilesService = Depends(get_admin_files_service), +): + result = await service.request_rescan(file_id=file_id, requested_by=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Rescan requested") + + +__all__ = ["router", "get_admin_files_service"] +``` + +`app/tests/test_admin_files_routes.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_files import router as admin_router, get_admin_files_service +from fileflash.schemas.admin.files import RescanResponse + + +class StubService: + async def list_files(self, *, query): # noqa: ANN001 + return SimpleNamespace(model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, "totalPages": 1, + "perPage": query.per_page, "currentPage": query.page, + "hasPrev": False, "hasNext": False, + }, + }) + + async def request_rescan(self, *, file_id, requested_by): # noqa: ANN001 + return RescanResponse( + file_id=str(file_id), virus_status="pending", + scanned_at=datetime.now(UTC), + ) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(admin_router, prefix="/api/v1") + app.dependency_overrides[get_admin_files_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=99) + return TestClient(app) + + +def test_list_files_returns_empty() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/files") + assert resp.status_code == 200 + + +def test_rescan_returns_pending() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/files/7/rescan") + assert resp.status_code == 200 + assert resp.json()["data"]["virusStatus"] == "pending" +``` + +- [ ] **Step 4.8: 跑测试 + 注册 router + commit** + +Run: `cd app && uv run pytest tests/test_admin_files_routes.py -v` +Expected: 2 passed + +Modify `app/src/fileflash/routers/__init__.py` 加入 `admin_files_router` 的 import 和 include。 + +```bash +git add app/src/fileflash/schemas/admin/files.py app/src/fileflash/services/admin/files.py app/src/fileflash/routers/admin_files.py app/src/fileflash/routers/__init__.py app/tests/test_admin_files_*.py +git commit -m "feat(admin): /admin/files list + /admin/files/{id}/rescan with event" +``` + +--- + +## Task 5: Admin Moderation（violations list + resolve） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/moderation.py` +- Create: `app/src/fileflash/services/admin/moderation.py` +- Create: `app/src/fileflash/routers/admin_moderation.py` +- Test: `app/tests/test_admin_moderation_service.py`, `app/tests/test_admin_moderation_routes.py` + +- [ ] **Step 5.1: 写 schemas** + +`app/src/fileflash/schemas/admin/moderation.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from ..common import CamelModel, PageQuery + + +ViolationLevel = Literal["low", "medium", "high"] +ViolationStatus = Literal["pending", "under_review", "resolved"] + + +class ViolationItem(CamelModel): + id: str + file_id: str | None + file_name: str | None + type: str + level: ViolationLevel + reported_at: datetime + status: ViolationStatus + + +class ListViolationsQuery(PageQuery): + status: ViolationStatus | None = None + + +class ResolveViolationResponse(CamelModel): + violation_id: str + resolved_at: datetime + + +__all__ = [ + "ListViolationsQuery", + "ResolveViolationResponse", + "ViolationItem", + "ViolationLevel", + "ViolationStatus", +] +``` + +- [ ] **Step 5.2: 写失败测试** + +`app/tests/test_admin_moderation_service.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from decimal import Decimal +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.schemas.admin.moderation import ListViolationsQuery +from fileflash.services.admin.moderation import AdminModerationService + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.commit = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.get = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_resolve_missing_case_returns_404() -> None: + session = DummySession() + session.get.return_value = None + service = AdminModerationService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.resolve_case(case_id=1, handled_by=2) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_resolve_already_resolved_returns_409() -> None: + session = DummySession() + case = Mock(case_id=1, status="resolved") + session.get.return_value = case + service = AdminModerationService(db=session) # type: ignore[arg-type] + + with pytest.raises(ApiError) as exc: + await service.resolve_case(case_id=1, handled_by=2) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_resolve_pending_case_sets_resolved() -> None: + session = DummySession() + case = Mock(case_id=1, status="pending", handled_by=None, handled_at=None, + resolution=None, file_id=10) + session.get.return_value = case + service = AdminModerationService(db=session) # type: ignore[arg-type] + + result = await service.resolve_case(case_id=1, handled_by=2) + + assert result.violation_id == "1" + assert case.status == "resolved" + assert case.resolution == "admin_clear" + assert case.handled_by == 2 + session.commit.assert_awaited() + + +@pytest.mark.asyncio +async def test_list_returns_empty_paginated() -> None: + session = DummySession() + session.scalar.return_value = 0 + session.execute.return_value = Mock(all=lambda: []) + service = AdminModerationService(db=session) # type: ignore[arg-type] + + result = await service.list_violations(query=ListViolationsQuery()) + assert result.items == [] + assert result.pagination.total_items == 0 +``` + +- [ ] **Step 5.3: 实现 service** + +`app/src/fileflash/services/admin/moderation.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from decimal import Decimal + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.tables_audit_security import ModerationCase +from ...models.tables_storage import File +from ...schemas.admin.moderation import ( + ListViolationsQuery, + ResolveViolationResponse, + ViolationItem, + ViolationLevel, +) +from ...schemas.common import PaginatedData, PaginationMeta + +_OPEN_STATES = ("pending", "under_review") + + +class AdminModerationService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_violations( + self, + *, + query: ListViolationsQuery, + ) -> PaginatedData[ViolationItem]: + statement = ( + select(ModerationCase, File) + .join(File, File.file_id == ModerationCase.file_id, isouter=True) + ) + if query.status: + statement = statement.where(ModerationCase.status == query.status) + statement = statement.order_by(ModerationCase.created_at.desc()) + + total = int( + await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0 + ) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = (await self.db.execute(statement.offset(offset).limit(query.per_page))).all() + + items = [self._to_item(case, file_) for case, file_ in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def resolve_case(self, *, case_id: int, handled_by: int) -> ResolveViolationResponse: + case = await self.db.get(ModerationCase, case_id, with_for_update=True) + if case is None: + raise ApiError(status_code=404, code=404, message="Violation case not found") + if case.status not in _OPEN_STATES: + raise ApiError(status_code=409, code=409, message="Case already resolved") + + now = datetime.now(UTC) + case.status = "resolved" + case.resolution = "admin_clear" + case.handled_by = handled_by + case.handled_at = now + case.updated_at = now + await self.db.commit() + + return ResolveViolationResponse(violation_id=str(case_id), resolved_at=now) + + @staticmethod + def _to_item(case: ModerationCase, file_: File | None) -> ViolationItem: + return ViolationItem( + id=str(case.case_id), + file_id=str(case.file_id) if case.file_id else None, + file_name=file_.name if file_ else None, + type=case.reason_type, + level=_level_from_confidence(case.confidence), + reported_at=case.created_at, + status=case.status, # type: ignore[arg-type] + ) + + +def _level_from_confidence(confidence: Decimal | None) -> ViolationLevel: + if confidence is None: + return "low" + value = float(confidence) + if value > 0.8: + return "high" + if value > 0.5: + return "medium" + return "low" + + +__all__ = ["AdminModerationService"] +``` + +- [ ] **Step 5.4: 跑 service 测试** + +Run: `cd app && uv run pytest tests/test_admin_moderation_service.py -v` +Expected: 4 passed + +- [ ] **Step 5.5: 实现 router + 测试** + +`app/src/fileflash/routers/admin_moderation.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.moderation import ListViolationsQuery +from ..services.admin.moderation import AdminModerationService + +router = APIRouter(prefix="/admin/violations", tags=["admin"]) + + +def get_admin_moderation_service( + db: AsyncSession = Depends(get_db), +) -> AdminModerationService: + return AdminModerationService(db=db) + + +@router.get("") +async def list_violations( + query: ListViolationsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminModerationService = Depends(get_admin_moderation_service), +): + data = await service.list_violations(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Violations fetched") + + +@router.post("/{case_id}/resolve") +async def resolve_violation( + case_id: int, + admin: User = Depends(require_admin), + service: AdminModerationService = Depends(get_admin_moderation_service), +): + result = await service.resolve_case(case_id=case_id, handled_by=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Violation resolved") + + +__all__ = ["router", "get_admin_moderation_service"] +``` + +`app/tests/test_admin_moderation_routes.py`（参照 §2.1 模式，2 测例：list 通过 / resolve 通过）： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_moderation import ( + router as mod_router, + get_admin_moderation_service, +) +from fileflash.schemas.admin.moderation import ResolveViolationResponse + + +class StubService: + async def list_violations(self, *, query): # noqa: ANN001 + return SimpleNamespace(model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, "totalPages": 1, + "perPage": query.per_page, "currentPage": query.page, + "hasPrev": False, "hasNext": False, + }, + }) + + async def resolve_case(self, *, case_id, handled_by): # noqa: ANN001 + return ResolveViolationResponse(violation_id=str(case_id), resolved_at=datetime.now(UTC)) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(mod_router, prefix="/api/v1") + app.dependency_overrides[get_admin_moderation_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_list_violations() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/violations") + assert resp.status_code == 200 + + +def test_resolve_violation() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/violations/3/resolve") + assert resp.status_code == 200 + assert resp.json()["data"]["violationId"] == "3" +``` + +- [ ] **Step 5.6: 跑测试 + 注册 + commit** + +Run: `cd app && uv run pytest tests/test_admin_moderation_routes.py -v` + +Modify `routers/__init__.py` 加入 `admin_moderation_router` import + include。 + +```bash +git add app/src/fileflash/schemas/admin/moderation.py app/src/fileflash/services/admin/moderation.py app/src/fileflash/routers/admin_moderation.py app/src/fileflash/routers/__init__.py app/tests/test_admin_moderation_*.py +git commit -m "feat(admin): /admin/violations list + resolve via ModerationCase" +``` + +--- + +## Task 6: Admin Logs + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/logs.py` +- Create: `app/src/fileflash/services/admin/logs.py` +- Create: `app/src/fileflash/routers/admin_logs.py` +- Test: `app/tests/test_admin_logs_routes.py` + +- [ ] **Step 6.1: schemas** + +`app/src/fileflash/schemas/admin/logs.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from ..common import CamelModel, PageQuery + + +class LogItem(CamelModel): + id: str + user_id: str | None + operation: str + operation_name: str + target_type: str | None + target_id: str | None + result: str + ip_address: str | None + user_agent: str | None + performed_at: datetime + details: str | None = None + metadata: dict[str, Any] = {} + + +class ListAdminLogsQuery(PageQuery): + user_id: str | None = None + operation: str | None = None + result: Literal["success", "failure"] | None = None + from_at: datetime | None = None + to_at: datetime | None = None + + +class AdminLogsResponse(CamelModel): + logs: list[LogItem] + pagination: dict[str, Any] + + +__all__ = ["AdminLogsResponse", "ListAdminLogsQuery", "LogItem"] +``` + +- [ ] **Step 6.2: 实现 service（轻量；只 list）** + +`app/src/fileflash/services/admin/logs.py`： + +```python +from __future__ import annotations + +from sqlalchemy import and_, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models.tables_audit_security import Log +from ...schemas.admin.logs import AdminLogsResponse, ListAdminLogsQuery, LogItem + + +class AdminLogsService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_logs(self, *, query: ListAdminLogsQuery) -> AdminLogsResponse: + statement = select(Log) + conditions = [] + if query.user_id: + conditions.append(Log.user_id == int(query.user_id)) + if query.operation: + conditions.append(Log.operation == query.operation) + if query.result: + conditions.append(Log.result == query.result) + if query.from_at: + conditions.append(Log.performed_at >= query.from_at) + if query.to_at: + conditions.append(Log.performed_at <= query.to_at) + if conditions: + statement = statement.where(and_(*conditions)) + statement = statement.order_by(Log.performed_at.desc()) + + total = int( + await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0 + ) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + + logs = [ + LogItem( + id=str(row.id), + user_id=str(row.user_id) if row.user_id else None, + operation=row.operation, + operation_name=row.operation, + target_type=row.target_type, + target_id=str(row.target_id) if row.target_id else None, + result=row.result, + ip_address=row.ip_address, + user_agent=row.user_agent, + performed_at=row.performed_at, + details=row.details, + metadata=row.metadata_payload or {}, + ) + for row in rows + ] + + return AdminLogsResponse( + logs=logs, + pagination={ + "totalItems": total, + "totalPages": total_pages, + "perPage": query.per_page, + "currentPage": query.page, + "hasPrev": query.page > 1, + "hasNext": query.page < total_pages, + }, + ) + + +__all__ = ["AdminLogsService"] +``` + +- [ ] **Step 6.3: 实现 router** + +`app/src/fileflash/routers/admin_logs.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.logs import ListAdminLogsQuery +from ..services.admin.logs import AdminLogsService + +router = APIRouter(prefix="/admin/logs", tags=["admin"]) + + +def get_admin_logs_service(db: AsyncSession = Depends(get_db)) -> AdminLogsService: + return AdminLogsService(db=db) + + +@router.get("") +async def list_admin_logs( + query: ListAdminLogsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminLogsService = Depends(get_admin_logs_service), +): + data = await service.list_logs(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Logs fetched") + + +__all__ = ["router", "get_admin_logs_service"] +``` + +- [ ] **Step 6.4: 测试** + +`app/tests/test_admin_logs_routes.py`： + +```python +from __future__ import annotations + +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_logs import router as logs_router, get_admin_logs_service +from fileflash.schemas.admin.logs import AdminLogsResponse + + +class StubService: + async def list_logs(self, *, query): # noqa: ANN001 + return AdminLogsResponse( + logs=[], + pagination={ + "totalItems": 0, "totalPages": 1, + "perPage": query.per_page, "currentPage": query.page, + "hasPrev": False, "hasNext": False, + }, + ) + + +def test_admin_can_list_logs() -> None: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(logs_router, prefix="/api/v1") + app.dependency_overrides[get_admin_logs_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + with TestClient(app) as c: + resp = c.get("/api/v1/admin/logs") + assert resp.status_code == 200 + assert resp.json()["data"]["logs"] == [] +``` + +- [ ] **Step 6.5: 跑测试 + 注册 + commit** + +Run: `cd app && uv run pytest tests/test_admin_logs_routes.py -v` + +Modify `routers/__init__.py` 加入 `admin_logs_router` import + include。 + +```bash +git add app/src/fileflash/schemas/admin/logs.py app/src/fileflash/services/admin/logs.py app/src/fileflash/routers/admin_logs.py app/src/fileflash/routers/__init__.py app/tests/test_admin_logs_routes.py +git commit -m "feat(admin): /admin/logs list with filters" +``` + +--- + +## Task 7: Admin Notifications（list / broadcast / read / archive） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/notifications.py` +- Create: `app/src/fileflash/services/admin/notifications.py` +- Create: `app/src/fileflash/routers/admin_notifications.py` +- Test: `app/tests/test_admin_notifications_service.py`, `app/tests/test_admin_notifications_routes.py` + +- [ ] **Step 7.1: schemas** + +`app/src/fileflash/schemas/admin/notifications.py`： + +```python +from __future__ import annotations + +from datetime import datetime +from typing import Any, Literal + +from pydantic import Field + +from ..common import CamelModel, PageQuery + + +class AdminNotificationItem(CamelModel): + id: str + message: str + title: str | None + type: str + status: str + is_read: bool + created_at: datetime + updated_at: datetime + recipient_count: int | None = None + + +class ListAdminNotificationsQuery(PageQuery): + status: str | None = None + type: str | None = None + + +class BroadcastRequest(CamelModel): + title: str | None = Field(default=None, max_length=255) + message: str = Field(min_length=1, max_length=2000) + type: Literal["system", "announcement"] = "system" + + +class BroadcastResponse(CamelModel): + broadcast_id: str + recipient_count: int + sent_at: datetime + + +__all__ = [ + "AdminNotificationItem", + "BroadcastRequest", + "BroadcastResponse", + "ListAdminNotificationsQuery", +] +``` + +- [ ] **Step 7.2: 写失败测试** + +`app/tests/test_admin_notifications_service.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import AsyncMock, Mock + +import pytest + +from fileflash.core.errors import ApiError +from fileflash.schemas.admin.notifications import BroadcastRequest +from fileflash.services.admin.notifications import ( + AdminNotificationsService, + MAX_BROADCAST_RECIPIENTS, +) + + +class DummySession: + def __init__(self) -> None: + self.add = Mock() + self.add_all = Mock() + self.commit = AsyncMock() + self.scalar = AsyncMock() + self.scalars = AsyncMock() + self.execute = AsyncMock() + + +@pytest.mark.asyncio +async def test_broadcast_rejects_empty_message() -> None: + # Pydantic 会先拦截；这里测 service 层显式拒绝（防御） + session = DummySession() + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + with pytest.raises(ApiError): + await service.broadcast( + payload=BroadcastRequest.model_construct(title=None, message="", type="system"), + sender_id=1, + ) + + +@pytest.mark.asyncio +async def test_broadcast_too_many_recipients_returns_422() -> None: + session = DummySession() + session.scalar.return_value = MAX_BROADCAST_RECIPIENTS + 1 + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + with pytest.raises(ApiError) as exc: + await service.broadcast( + payload=BroadcastRequest(message="hi", type="system"), + sender_id=1, + ) + assert exc.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_broadcast_writes_one_row_per_recipient() -> None: + session = DummySession() + session.scalar.return_value = 3 + session.scalars.return_value = iter([10, 11, 12]) + service = AdminNotificationsService(db=session) # type: ignore[arg-type] + + result = await service.broadcast( + payload=BroadcastRequest(message="ping", type="system"), + sender_id=1, + ) + + assert result.recipient_count == 3 + assert session.add_all.called +``` + +- [ ] **Step 7.3: 实现 service** + +`app/src/fileflash/services/admin/notifications.py`： + +```python +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from typing import AsyncIterator + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.errors import ApiError +from ...models.enums import UserStatus +from ...models.tables_audit_security import Notification +from ...models.tables_identity import User +from ...schemas.admin.notifications import ( + AdminNotificationItem, + BroadcastRequest, + BroadcastResponse, + ListAdminNotificationsQuery, +) +from ...schemas.common import PaginatedData, PaginationMeta + +MAX_BROADCAST_RECIPIENTS = 50_000 +_CHUNK_SIZE = 500 + + +class AdminNotificationsService: + def __init__(self, db: AsyncSession) -> None: + self.db = db + + async def list_notifications( + self, + *, + query: ListAdminNotificationsQuery, + ) -> PaginatedData[AdminNotificationItem]: + statement = select(Notification) + if query.status: + statement = statement.where(Notification.status == query.status) + if query.type: + statement = statement.where(Notification.notification_type == query.type) + statement = statement.order_by(Notification.created_at.desc()) + + total = int( + await self.db.scalar(select(func.count()).select_from(statement.subquery())) or 0 + ) + total_pages = max(1, -(-total // query.per_page)) + offset = (query.page - 1) * query.per_page + rows = list(await self.db.scalars(statement.offset(offset).limit(query.per_page))) + + items = [self._to_item(row) for row in rows] + return PaginatedData( + items=items, + pagination=PaginationMeta( + total_items=total, + total_pages=total_pages, + per_page=query.per_page, + current_page=query.page, + has_prev=query.page > 1, + has_next=query.page < total_pages, + ), + ) + + async def broadcast( + self, + *, + payload: BroadcastRequest, + sender_id: int, + ) -> BroadcastResponse: + message = (payload.message or "").strip() + if not message: + raise ApiError(status_code=422, code=422, message="Broadcast message cannot be empty") + + recipient_count = int( + await self.db.scalar( + select(func.count(User.user_id)).where( + User.status == UserStatus.ACTIVE, + User.deleted_at.is_(None), + ) + ) or 0 + ) + if recipient_count > MAX_BROADCAST_RECIPIENTS: + raise ApiError( + status_code=422, + code=422, + message=f"Recipient count {recipient_count} exceeds limit {MAX_BROADCAST_RECIPIENTS}", + ) + + broadcast_id = str(uuid.uuid4()) + now = datetime.now(UTC) + delivered = 0 + + async for chunk in self._iter_active_user_ids(): + rows = [ + Notification( + user_id=uid, + title=payload.title, + notification_type=payload.type, + channel="in_app", + message=message, + payload={"broadcastId": broadcast_id}, + sender_user_id=sender_id, + status="sent", + sent_at=now, + is_read=False, + created_at=now, + updated_at=now, + ) + for uid in chunk + ] + if rows: + self.db.add_all(rows) + await self.db.commit() + delivered += len(rows) + + return BroadcastResponse( + broadcast_id=broadcast_id, + recipient_count=delivered, + sent_at=now, + ) + + async def _iter_active_user_ids(self) -> AsyncIterator[list[int]]: + stream = await self.db.scalars( + select(User.user_id).where( + User.status == UserStatus.ACTIVE, + User.deleted_at.is_(None), + ) + ) + buffer: list[int] = [] + for user_id in stream: + buffer.append(user_id) + if len(buffer) >= _CHUNK_SIZE: + yield buffer + buffer = [] + if buffer: + yield buffer + + async def archive(self, *, notification_id: int) -> None: + row = await self.db.get(Notification, notification_id) + if row is None: + raise ApiError(status_code=404, code=404, message="Notification not found") + row.status = "archived" + row.updated_at = datetime.now(UTC) + await self.db.commit() + + @staticmethod + def _to_item(row: Notification) -> AdminNotificationItem: + broadcast_id = (row.payload or {}).get("broadcastId") if row.payload else None + return AdminNotificationItem( + id=str(row.id), + message=row.message, + title=row.title, + type=row.notification_type, + status=row.status, + is_read=bool(row.is_read), + created_at=row.created_at or datetime.now(UTC), + updated_at=row.updated_at, + recipient_count=None if broadcast_id is None else None, + ) + + +__all__ = ["AdminNotificationsService", "MAX_BROADCAST_RECIPIENTS"] +``` + +- [ ] **Step 7.4: 跑测试** + +Run: `cd app && uv run pytest tests/test_admin_notifications_service.py -v` +Expected: 3 passed + +- [ ] **Step 7.5: 实现 router + 路由测试** + +`app/src/fileflash/routers/admin_notifications.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import require_admin +from ..core.errors import api_success +from ..db.deps import get_db +from ..models.tables_identity import User +from ..schemas.admin.notifications import ( + BroadcastRequest, + ListAdminNotificationsQuery, +) +from ..services.admin.notifications import AdminNotificationsService + +router = APIRouter(prefix="/admin/notifications", tags=["admin"]) + + +def get_admin_notifications_service( + db: AsyncSession = Depends(get_db), +) -> AdminNotificationsService: + return AdminNotificationsService(db=db) + + +@router.get("") +async def list_admin_notifications( + query: ListAdminNotificationsQuery = Depends(), + _: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + data = await service.list_notifications(query=query) + return api_success(data=data.model_dump(by_alias=True), message="Notifications fetched") + + +@router.post("/broadcast") +async def broadcast_notification( + payload: BroadcastRequest, + admin: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + result = await service.broadcast(payload=payload, sender_id=admin.user_id) + return api_success(data=result.model_dump(by_alias=True), message="Broadcast sent") + + +@router.delete("/{notification_id}") +async def archive_admin_notification( + notification_id: int, + _: User = Depends(require_admin), + service: AdminNotificationsService = Depends(get_admin_notifications_service), +): + await service.archive(notification_id=notification_id) + return api_success( + data={"notificationId": str(notification_id), "status": "archived"}, + message="Notification archived", + ) + + +__all__ = ["router", "get_admin_notifications_service"] +``` + +`app/tests/test_admin_notifications_routes.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_notifications import ( + router as n_router, + get_admin_notifications_service, +) +from fileflash.schemas.admin.notifications import BroadcastResponse + + +class StubService: + async def list_notifications(self, *, query): # noqa: ANN001 + return SimpleNamespace(model_dump=lambda **_: { + "items": [], + "pagination": { + "totalItems": 0, "totalPages": 1, + "perPage": query.per_page, "currentPage": query.page, + "hasPrev": False, "hasNext": False, + }, + }) + + async def broadcast(self, *, payload, sender_id): # noqa: ANN001 + return BroadcastResponse( + broadcast_id="b1", recipient_count=2, sent_at=datetime.now(UTC), + ) + + async def archive(self, *, notification_id): # noqa: ANN001 + return None + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(n_router, prefix="/api/v1") + app.dependency_overrides[get_admin_notifications_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_list() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/notifications") + assert resp.status_code == 200 + + +def test_broadcast() -> None: + with _client() as c: + resp = c.post("/api/v1/admin/notifications/broadcast", + json={"message": "hello", "type": "system"}) + assert resp.status_code == 200 + assert resp.json()["data"]["broadcastId"] == "b1" + + +def test_archive() -> None: + with _client() as c: + resp = c.delete("/api/v1/admin/notifications/9") + assert resp.status_code == 200 +``` + +- [ ] **Step 7.6: 跑测试 + 注册 + commit** + +```bash +cd app && uv run pytest tests/test_admin_notifications_routes.py -v +``` + +Modify `routers/__init__.py` 加 `admin_notifications_router` import + include。 + +```bash +git add app/src/fileflash/schemas/admin/notifications.py app/src/fileflash/services/admin/notifications.py app/src/fileflash/routers/admin_notifications.py app/src/fileflash/routers/__init__.py app/tests/test_admin_notifications_*.py +git commit -m "feat(admin): /admin/notifications list, broadcast, archive" +``` + +--- + +## Task 8: Admin System（health + rate-limit） + +**Files:** + +- Create: `app/src/fileflash/schemas/admin/system.py` +- Create: `app/src/fileflash/services/admin/system.py` +- Create: `app/src/fileflash/routers/admin_system.py` +- Test: `app/tests/test_admin_system_routes.py` + +- [ ] **Step 8.1: schemas** + +`app/src/fileflash/schemas/admin/system.py`： + +```python +from __future__ import annotations + +from datetime import datetime + +from ..common import CamelModel + + +class SystemHealth(CamelModel): + platform_targets: list[str] + max_concurrent_uploads: int + active_upload_sessions: int + virus_scan_enabled: bool + thumbnail_generation_enabled: bool + registration_mail_enabled: bool + hash_computation_enabled: bool + last_updated_at: datetime + + +class RateLimitRule(CamelModel): + rule_id: str + scope: str + window_seconds: int + limit: int + current_usage: int + blocked_requests: int + + +class RateLimitStatus(CamelModel): + rules: list[RateLimitRule] + evaluated_at: datetime + + +__all__ = ["RateLimitRule", "RateLimitStatus", "SystemHealth"] +``` + +- [ ] **Step 8.2: 实现 service（轻量；health 来自 settings，rate-limit 静态枚举）** + +`app/src/fileflash/services/admin/system.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...core.settings import Settings +from ...models.enums import UploadTaskStatus +from ...models.tables_storage import UploadTask +from ...schemas.admin.system import RateLimitRule, RateLimitStatus, SystemHealth + +# 静态规则定义；如 auth router 有 N 个 RateLimiter.allow 调用，把它们对应的 scope 列在这。 +_RATE_LIMIT_RULES = [ + {"rule_id": "login", "scope": "auth.login", "window_seconds": 60, "limit": 5}, + {"rule_id": "register", "scope": "auth.register", "window_seconds": 600, "limit": 3}, + {"rule_id": "forgot_password", "scope": "auth.forgot_password", "window_seconds": 600, "limit": 3}, + {"rule_id": "resend_verification", "scope": "auth.resend_verification", "window_seconds": 600, "limit": 3}, +] + + +class AdminSystemService: + def __init__(self, db: AsyncSession, settings: Settings) -> None: + self.db = db + self.settings = settings + + async def health(self) -> SystemHealth: + active_uploads = int( + await self.db.scalar( + select(func.count(UploadTask.task_id)).where( + UploadTask.status.in_([UploadTaskStatus.INIT, UploadTaskStatus.UPLOADING]) + ) + ) or 0 + ) + + platform_targets: list[str] = [] + if self.settings.object_storage_bucket: + platform_targets.append(f"s3://{self.settings.object_storage_bucket}") + if self.settings.redis_url: + platform_targets.append("redis") + + return SystemHealth( + platform_targets=platform_targets, + max_concurrent_uploads=getattr(self.settings, "max_concurrent_uploads", 4), + active_upload_sessions=active_uploads, + virus_scan_enabled=bool(getattr(self.settings, "virus_scan_enabled", False)), + thumbnail_generation_enabled=bool(getattr(self.settings, "thumbnail_generation_enabled", True)), + registration_mail_enabled=bool(self.settings.mail_server and self.settings.mail_from), + hash_computation_enabled=True, + last_updated_at=datetime.now(UTC), + ) + + async def rate_limit_status(self) -> RateLimitStatus: + rules = [ + RateLimitRule( + rule_id=rule["rule_id"], + scope=rule["scope"], + window_seconds=rule["window_seconds"], + limit=rule["limit"], + current_usage=0, # 不扫 Redis；运行时统计后续扩展 + blocked_requests=0, + ) + for rule in _RATE_LIMIT_RULES + ] + return RateLimitStatus(rules=rules, evaluated_at=datetime.now(UTC)) + + +__all__ = ["AdminSystemService"] +``` + +- [ ] **Step 8.3: 实现 router** + +`app/src/fileflash/routers/admin_system.py`： + +```python +from __future__ import annotations + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from ..core.deps import get_settings_dep, require_admin +from ..core.errors import api_success +from ..core.settings import Settings +from ..db.deps import get_db +from ..models.tables_identity import User +from ..services.admin.system import AdminSystemService + +router = APIRouter(prefix="/admin/system", tags=["admin"]) + + +def get_admin_system_service( + db: AsyncSession = Depends(get_db), + settings: Settings = Depends(get_settings_dep), +) -> AdminSystemService: + return AdminSystemService(db=db, settings=settings) + + +@router.get("/health") +async def get_system_health( + _: User = Depends(require_admin), + service: AdminSystemService = Depends(get_admin_system_service), +): + data = await service.health() + return api_success(data=data.model_dump(by_alias=True), message="System health fetched") + + +@router.get("/rate-limit") +async def get_rate_limit_status( + _: User = Depends(require_admin), + service: AdminSystemService = Depends(get_admin_system_service), +): + data = await service.rate_limit_status() + return api_success(data=data.model_dump(by_alias=True), message="Rate limit fetched") + + +__all__ = ["router", "get_admin_system_service"] +``` + +- [ ] **Step 8.4: 测试** + +`app/tests/test_admin_system_routes.py`： + +```python +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fileflash.core.deps import require_admin +from fileflash.core.errors import ApiError, api_error_handler +from fileflash.routers.admin_system import router as sys_router, get_admin_system_service +from fileflash.schemas.admin.system import ( + RateLimitRule, RateLimitStatus, SystemHealth, +) + + +class StubService: + async def health(self): + return SystemHealth( + platform_targets=["s3://bkt"], max_concurrent_uploads=4, + active_upload_sessions=0, virus_scan_enabled=False, + thumbnail_generation_enabled=True, registration_mail_enabled=False, + hash_computation_enabled=True, last_updated_at=datetime.now(UTC), + ) + + async def rate_limit_status(self): + return RateLimitStatus( + rules=[ + RateLimitRule( + rule_id="login", scope="auth.login", + window_seconds=60, limit=5, current_usage=0, blocked_requests=0, + ) + ], + evaluated_at=datetime.now(UTC), + ) + + +def _client() -> TestClient: + app = FastAPI() + app.add_exception_handler(ApiError, api_error_handler) + app.include_router(sys_router, prefix="/api/v1") + app.dependency_overrides[get_admin_system_service] = lambda: StubService() + app.dependency_overrides[require_admin] = lambda: SimpleNamespace(user_id=1) + return TestClient(app) + + +def test_health() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/system/health") + assert resp.status_code == 200 + assert resp.json()["data"]["activeUploadSessions"] == 0 + + +def test_rate_limit() -> None: + with _client() as c: + resp = c.get("/api/v1/admin/system/rate-limit") + assert resp.status_code == 200 + assert resp.json()["data"]["rules"][0]["scope"] == "auth.login" +``` + +- [ ] **Step 8.5: 跑测试 + 注册 + commit** + +```bash +cd app && uv run pytest tests/test_admin_system_routes.py -v +``` + +Modify `routers/__init__.py` 加入 `admin_system_router` import + include。 + +```bash +git add app/src/fileflash/schemas/admin/system.py app/src/fileflash/services/admin/system.py app/src/fileflash/routers/admin_system.py app/src/fileflash/routers/__init__.py app/tests/test_admin_system_routes.py +git commit -m "feat(admin): /admin/system health and rate-limit" +``` + +--- + +## Task 9: 端到端冒烟 + 全量回归 + +- [ ] **Step 9.1: 全量启动冒烟** + +Run: `cd app && uv run python -c "from fileflash.main import app; print(sorted(r.path for r in app.routes if '/admin/' in r.path))"` +Expected: prints 14+ paths including `/api/v1/admin/users`, `/api/v1/admin/storage/summary`, `/api/v1/admin/files`, `/api/v1/admin/violations`, `/api/v1/admin/logs`, `/api/v1/admin/notifications`, `/api/v1/admin/system/health` 等。 + +- [ ] **Step 9.2: 跑全量 admin 测试** + +Run: `cd app && uv run pytest tests/test_admin_ -v` +Expected: 所有 admin_* 测试通过（含 Task 0-8 累计 ≥ 25 个测例）。 + +- [ ] **Step 9.3: 跑全量回归** + +Run: `cd app && uv run pytest -q` +Expected: 不引入新 failure；现有测试全绿。 + +- [ ] **Step 9.4: 编写 changelog / README 摘要（可选）** + +如果项目维护 README/CHANGELOG，添加一段： + +> 新增 Admin Console 后端 API（/admin/users, /admin/storage/*, /admin/files, /admin/violations, /admin/logs, /admin/notifications, /admin/system/*）。前端请参照 Plan B 完成 Console 重构。 + +如果项目无此惯例，跳过本步。 + +- [ ] **Step 9.5: Commit final** + +```bash +git status +# 若 Step 9.4 有改动则提交： +# git add README.md && git commit -m "docs: note admin console API surface" +``` + +--- + +## Plan A 完工标准 + +- 所有 `/admin/*` 路由：`/users`, `/storage/*`, `/files*`, `/violations*`, `/logs`, `/notifications*`, `/system/*` 可经过真实 DB 工作，鉴权由 `require_admin` 控制。 +- 14+ 新增测例全部通过。 +- 全量 `uv run pytest -q` 仍绿。 +- 现有 `/admin/registration-email-domain-rules` 零改动、零回归。 +- 前端在 mock 关闭、指向真实后端时（设 `VITE_USE_MOCK=false`），Console 各页（Plan B 完成后）可正常工作。 diff --git a/docs/superpowers/plans/2026-05-24-admin-console-frontend.md b/docs/superpowers/plans/2026-05-24-admin-console-frontend.md new file mode 100644 index 0000000..f236394 --- /dev/null +++ b/docs/superpowers/plans/2026-05-24-admin-console-frontend.md @@ -0,0 +1,1996 @@ +# Admin Console Frontend Implementation Plan (Plan B) + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 把 `web/src/pages/dashboard/Dashboard.vue` 单页式管理面板重构为 `web/src/pages/console/` 多页式 Console（9 个子页 + 共享侧栏与组件），并补齐与 Plan A 后端契约对齐的 api 函数与 mock。 + +**Architecture:** 复用 `MainLayout` 顶层框架，在其下嵌一层 `ConsoleLayout`（含 `ConsoleSidebar` + ``）；9 个子页 lazy import；共享展示组件不持有数据，由子页 `onMounted` 拉数据后注入 props。视觉沿用 [[frontend_aesthetic]] 工业风。 + +**Tech Stack:** Vue 3 ` + + + + +``` + +- [ ] **Step 1.2: 写 ConsoleLayout** + +`web/src/pages/console/ConsoleLayout.vue`： + +```vue + + + + + +``` + +- [ ] **Step 1.3: 写入口 index.ts（每个子目录都需要类似一个）** + +`web/src/pages/console/index.ts`： + +```typescript +export { default as ConsoleLayout } from './ConsoleLayout.vue'; +``` + +(每个子页的 `/index.ts` 同样默认导出页面组件，例：) + +`web/src/pages/console/overview/index.ts`（占位，Task 3 才实现）： + +```typescript +import OverviewPage from './OverviewPage.vue'; +export default OverviewPage; +``` + +> 在本任务里，先为 9 个子页都各建一个空 `.vue`（含最小模板 ``）+ `index.ts`，让路由可以挂上不报错。代码示例： + +```vue + + + + +``` + +为 9 个子目录都生成同样形状的 placeholder：`overview, users, storage, content, moderation, system, logs, notifications, rules`。 + +- [ ] **Step 1.4: 修改 router/routes.ts** + +`web/src/router/routes.ts` —— 在 `children` 数组中把 `dashboard` 路由替换为 console 嵌套： + +```typescript +// 删除原 dashboard 子路由块；新增： +{ + path: 'console', + component: () => import('../pages/console/ConsoleLayout.vue'), + meta: { navId: 'console', requiresAdmin: true }, + children: [ + { path: '', redirect: '/console/overview' }, + { path: 'overview', name: 'ConsoleOverview', component: () => import('../pages/console/overview/index.ts') }, + { path: 'users', name: 'ConsoleUsers', component: () => import('../pages/console/users/index.ts') }, + { path: 'storage', name: 'ConsoleStorage', component: () => import('../pages/console/storage/index.ts') }, + { path: 'content', name: 'ConsoleContent', component: () => import('../pages/console/content/index.ts') }, + { path: 'moderation', name: 'ConsoleModeration', component: () => import('../pages/console/moderation/index.ts') }, + { path: 'system', name: 'ConsoleSystem', component: () => import('../pages/console/system/index.ts') }, + { path: 'logs', name: 'ConsoleLogs', component: () => import('../pages/console/logs/index.ts') }, + { path: 'notifications', name: 'ConsoleNotifications', component: () => import('../pages/console/notifications/index.ts') }, + { path: 'rules', name: 'ConsoleRules', component: () => import('../pages/console/rules/index.ts') }, + ], +}, +{ path: '/dashboard', redirect: '/console/overview' }, +``` + +旧 `dashboard` route 移除（注意 `/` 父路由还是 MainLayout，console 是它的 child）。 + +- [ ] **Step 1.5: 类型检查 + 启动 dev server 烟测** + +```bash +cd web && bun run check +cd web && bun run dev # 浏览器访问 http://localhost:5173/console +``` + +预期：以 admin 用户登录后访问 `/console`，重定向至 `/console/overview`，能看到 sidebar + "overview (TODO)" 占位。 + +- [ ] **Step 1.6: Commit** + +```bash +git add web/src/pages/console web/src/router/routes.ts +git commit -m "feat(web): scaffold Console layout, sidebar, and 9 subpage routes" +``` + +--- + +## Task 2: 共享组件（components/console/） + +**Files:** + +- Create: `web/src/components/console/KpiCard.vue`, `StatusBadge.vue`, `FilterBar.vue`, `AdminTable.vue`, `TrendChart.vue`, `BroadcastComposer.vue`, `QuotaEditor.vue`, `index.ts` + +- [ ] **Step 2.1: KpiCard.vue** + +```vue + + + + + +``` + +- [ ] **Step 2.2: StatusBadge.vue** + +```vue + + + + + +``` + +- [ ] **Step 2.3: FilterBar.vue** + +```vue + + + + + +``` + +- [ ] **Step 2.4: AdminTable.vue** + +```vue + + + + + +``` + +- [ ] **Step 2.5: TrendChart.vue** + +```vue + + + + + +``` + +- [ ] **Step 2.6: BroadcastComposer.vue** + +```vue + + +