mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-10-16 06:58:51 +08:00
fix(dashboard): bound upload memory and add resumable chunked uploads (#10123)
* refactor(dashboard): consolidate upload saves into streaming helper
Five copies of the same upload-save logic existed: three file wrappers
(UploadFileAdapter, RequestUploadFile, PluginUploadFile) plus one
_save_upload() in each of BackupService and SkillsService. The backup
variant fell back to a full-body read() for raw FastAPI UploadFile,
buffering entire uploads in memory.
Introduce save_upload_stream() as the single streaming implementation
(1 MiB blocks, optional max_bytes with in-loop abort and partial-file
cleanup). The wrappers keep their public signatures and delegate to it;
skills/chat/config save paths call save() directly with max_bytes,
replacing their post-save size checks.
* fix(dashboard): bound backup upload size and harden chunked flow
Backup routes passed raw FastAPI UploadFile into the service layer,
hitting the full-body read() fallback: a single authenticated request
with an oversized chunk or whole-file upload could OOM-kill the
process (verified with a 512 MiB memory-limited container, exit 137).
- Wrap uploads in UploadFileAdapter so all saves stream to disk
- Enforce CHUNK_SIZE (1 MiB) per chunk and 128 MiB for whole-file
uploads, aborting mid-stream with partial-file cleanup
- Cap chunked backup totals at 8 GiB and point larger backups at
manual restore via the data directory
- Verify the merged backup size against the declared total
* fix(dashboard): enforce MAX_CONTENT_LENGTH with per-route overrides
The 128 MiB MAX_CONTENT_LENGTH value was written into the compat-layer
config dict but nothing read it, so the FastAPI/hypercorn path had no
request body limit at all.
Add a middleware that rejects oversized bodies with 413 before auth or
body parsing, using MAX_CONTENT_LENGTH as the default with per-route
overrides for legitimate large uploads (attachments 512 MiB, plugin
config files 500 MiB, knowledge bases 512 MiB) and a tight 2 MiB cap
on the backup chunk endpoint. Requests without a Content-Length header
pass through and remain bounded by per-endpoint max_bytes at save time.
* refactor(dashboard): extract generic chunked upload session service
Pull the chunk session machinery out of BackupService into a reusable
ChunkedUploadService owning session lifecycle, chunk storage, streaming
assembly and expiry cleanup, in preparation for webchat attachment
chunked uploads.
Sessions are now bound to an owner (the authenticated username) and a
purpose; chunk/complete/abort calls from another owner fail with the
same "not found or expired" error as unknown sessions, so upload IDs
cannot be probed or consumed cross-user. BackupService keeps business
concerns (zip validation, size caps, manifest marking) and delegates
the rest; API shapes and frontend flows are unchanged.
* feat(dashboard): add resumable upload status endpoint
Expose per-session progress (received chunk indexes, totals, expiry)
so clients can resume an interrupted chunked upload by sending only
the missing chunks instead of starting over.
Abuse considerations baked in: the lookup inherits owner binding and
answers "not found or expired" for both unknown sessions and sessions
owned by someone else (no probing), the response carries no filename
(minimal disclosure), and the query is strictly read-only — it never
touches last_activity, so polling it cannot anchor a session and its
on-disk chunks alive past expiry.
* fix(dashboard): bound file signature memory with chunked hashing
getFileSignature() read the entire file into browser memory via
arrayBuffer() before computing SHA-256, spiking renderer memory by the
full attachment size (up to the 512 MiB upload cap).
The signature is only used for in-session staging dedup and never
leaves the browser, so it does not need to be the canonical whole-file
SHA-256. Hash 8 MiB blocks individually and digest the concatenated
block digests instead; the prefix changes to sha256m: to mark the
algorithm change (old signatures vanish with the page session, so no
compatibility concern).
* feat(dashboard): resumable chunked upload UI for backup import
Extract the chunk scheduling logic from BackupDialog into a reusable
useChunkedUpload composable (5-way concurrency, per-chunk silent
retries, progress accounting) and add resume support on top of the
new upload status endpoint: a failed upload now offers a "Resume
Upload" action that queries received chunk indexes and sends only the
missing ones. Expired sessions transparently re-initialize and start
over; a failed complete() drops the poisoned session by design.
Wires the new /backups/upload/status endpoint through the OpenAPI
spec and regenerated client, and adds the resumeUpload label to all
four locales.
* feat(dashboard): chunked upload endpoints for webchat attachments
Mirror the backup chunked-upload flow for chat attachments so large
files (>128 MiB request-body default, up to the 512 MiB attachment
cap) can be uploaded in slices: POST /api/v1/files/upload/{init,
chunk,complete,abort,status}, owner-bound sessions, and a per-route
body limit on the chunk endpoint.
ChatService reuses the generic ChunkedUploadService with a new meta
passthrough carrying the client content_type through the session;
completion merges outside the chunk dir and feeds the result through
save_uploaded_file() unchanged via a move-based LocalUploadFile
adapter, keeping MIME sniffing, type classification and attachment
insertion in exactly one place.
Also repairs two pre-existing test stubs that predated the
max_bytes-aware save() contract and were already failing on this
branch.
* feat(dashboard): chunked upload with resume for webchat attachments
Route staged files >= 32 MiB through the new chunked upload endpoints
via useChunkedUpload; smaller files keep the single POST path so paste
and quick picks stay low-latency. The content_type travels through the
init payload so server-side attachment classification is unchanged.
A failed chunked upload is kept as an error chip in the input area
with a retry action that resumes from the last received chunk (and
transparently restarts when the session expired) plus a discard action
that aborts the session. Failed entries hold their uploader in a
shallowRef so Vue's deep ref unwrapping cannot strip the composable's
internal refs, and their signatures stay reserved to block duplicate
staging while a failed entry exists.
* docs: regenerate openapi artifacts for file upload endpoints
* fix(dashboard): harden chunk save and request body limits
Address review findings on the upload-hardening branch:
- save_chunk() now writes to a uniquely named temp file, verifies the
save() byte count against the expected chunk length, and publishes
via atomic rename. A failed or concurrent same-index retry can no
longer truncate and delete an already-received chunk (previously
left received_chunks claiming data that assemble() then failed to
find), and short writes are rejected instead of published.
- Requests with multipart bodies but no Content-Length are rejected
with 411: form parsing spools large bodies to disk before any
per-file size check can run. This is a documented stopgap — other
lengthless bodies still pass and are bounded only at save time;
closing that fully requires counting bytes as they arrive.
- Whole-file upload overrides gain 1 MiB of multipart framing slack
so a file exactly at the size cap is no longer rejected with 413.
The body-limit decision is extracted into a pure function for direct
unit tests; the assemble-time size recheck is now covered by an
out-of-band corruption test since chunk-level verification rejects
short writes earlier.
* fix(dashboard): drain upload pool and pin session on resume
A failed chunk made runPool() reject immediately while the remaining
in-flight requests kept writing; a resume could then race them on the
same chunk, stale retries read the mutable uploadId (potentially
writing into a re-initialized session), and late completions kept
counting progress (150% after resume).
The pool now stops scheduling on first failure but drains issued
requests via allSettled before surfacing the error, each request pins
the uploadId captured at run start, and a run-generation counter
(bumped on start/resume/cancel) blocks stale completions from the
progress accounting.
* fix(dashboard): close 411 bypass and temp-file residue on publish
Address the second round of review findings:
- The 411 media-type check used a raw startswith(), but the form
parser strips surrounding whitespace, so a leading-space
Content-Type was let through without a Content-Length while
Starlette still parsed the body as multipart. Detection now uses
the same parse_options_header the parser uses.
- A failing os.replace() during chunk publish left the uniquely
named temp file behind; it is now covered by the same cleanup
discipline as the save path.
Backup direct-upload routes deliberately get no framing slack: they
are legacy paths with no remaining callers, and the pre-existing
exact-limit 413 edge is answered by chunked upload.
* fix(dashboard): abort sessions created after a cancel
Cancelling while initUpload was still in flight cleared local state
before an upload id existed; when the delayed init returned, the run
flipped from idle to error and the freshly created server session was
never aborted. handleFailure now lets cancel win over any late
resolution and aborts the session id captured at that point.
* fix(dashboard): check cancel flag after completeUpload resolves
Cancelling while completeUpload was in flight let the late success
response flip the uploader from idle to done and hand the result to
the caller, which proceeded to the precheck. completeSession() now
re-checks the cancel flag after the response arrives and routes the
late success through the cancelled path instead.
* fix(dashboard): gate the 411 rule to body-capable methods
The 411 stopgap rejected GET /api/v1/logs/live: the SSE log client
sends a bogus "Content-Type: multipart/form-data" header on a GET
(copy-paste relic), and a GET with no Content-Length matched the
multipart-without-length rule. Body-less methods never trigger form
parsing, so no disk-spool risk exists for them — the 411 now applies
only to POST/PUT/PATCH. The junk header on the SSE client is removed
as well; the other two SSE consumers never had it.
* feat(dashboard): show in-progress chips for chunked attachments
Chunked uploads gave no feedback while in flight: the staged chip
only appeared after completion, leaving multi-minute uploads silent.
Track in-flight uploads alongside the failed ones and render progress
chips (filename, percent, thin progress bar) with a cancel action in
the input area; on success the chip hands off to the staged
attachment, on failure to the resumable error chip.
* fix(dashboard): type active upload chips upfront and align progress bar
The in-progress chip used a generic cloud icon even though the file
type is known locally before the upload starts; it now derives
icon/color/label from the filename via attachmentPresentation, same
as the staged chip. The progress bar's left edge is aligned with the
filename text instead of the card padding.
* fix(dashboard): keep active-chip progress bar inside the card
v-progress-linear ships width:100%, which wins over the right offset
in absolute positioning; the bar ran 52px past the card's right edge
and was clipped by overflow:hidden. Give it width:auto so the left/
right offsets determine the width.
* fix(dashboard): rework active-chip progress track and cancel button
Positioning a Vuetify component root directly proved fragile — the
progress bar drifted to the card's top edge. Wrap it in a plain track
div that owns the absolute positioning and let the component fill it.
Replace the corner tonal badge cancel with an inline text-variant
icon button after the percent label, and drop the padding reserved
for the badge.
* fix(dashboard): enforce session expiry and atomic save semantics
Address the third review round (sourcery-ai findings):
- get_session() now actually enforces the idle expiry it advertises
in its "not found or expired" error; expired sessions stay
abortable for cleanup, and backup upload_init starts the janitor
just like chat's does, so chunk dirs cannot linger indefinitely.
- cleanup_session() only unregisters a session after rmtree()
succeeds, letting the janitor retry on transient filesystem errors
instead of permanently leaking the chunk directory.
- save_upload_stream() writes to a uniquely named temp sibling and
atomically renames on success, so a failed or cancelled overwrite
can no longer delete a pre-existing destination (e.g. re-uploading
over an existing attachment), and all callers inherit the behavior.
The write()-return-value finding is a non-issue: BufferedWriter
writes fully or raises, and the chunk path already verifies byte
counts.
* fix(dashboard): guard failed-upload retry against double clicks
retryFailedUpload() ran resume() without any in-flight marker, so a
second click started a concurrent resume on the same session; the
loser hit "not found or expired" and left a failed chip next to the
successfully staged attachment. The entry now moves to activeUploads
while resuming, which blocks the second click and shows live
progress, and returns to failedUploads only if the resume errors.
This commit is contained in:
Ruochen Pan
+13
-12
@@ -12,6 +12,8 @@ from starlette.datastructures import Headers
|
||||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
from starlette.responses import StreamingResponse
|
||||
|
||||
from astrbot.core.utils.upload import save_upload_stream
|
||||
|
||||
ValueT = TypeVar("ValueT")
|
||||
DefaultT = TypeVar("DefaultT")
|
||||
ConvertedT = TypeVar("ConvertedT")
|
||||
@@ -109,23 +111,22 @@ class PluginUploadFile:
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
async def save(self, destination: str | Path) -> None:
|
||||
async def save(
|
||||
self, destination: str | Path, *, max_bytes: int | None = None
|
||||
) -> int:
|
||||
"""Save the uploaded file to disk.
|
||||
|
||||
Args:
|
||||
destination: Destination file path.
|
||||
max_bytes: Optional hard limit; the partial file is removed and
|
||||
UploadTooLargeError is raised when the upload exceeds it.
|
||||
|
||||
Returns:
|
||||
Number of bytes written.
|
||||
"""
|
||||
path = Path(destination)
|
||||
try:
|
||||
await self._upload_file.seek(0)
|
||||
except Exception:
|
||||
pass
|
||||
with path.open("wb") as output:
|
||||
while True:
|
||||
chunk = await self._upload_file.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
output.write(chunk)
|
||||
return await save_upload_stream(
|
||||
self._upload_file, destination, max_bytes=max_bytes
|
||||
)
|
||||
|
||||
async def read(self, size: int = -1) -> bytes:
|
||||
"""Read bytes from the uploaded file.
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Shared streaming save for uploaded files.
|
||||
|
||||
Internal helper used by the dashboard/plugin upload wrappers so that every
|
||||
upload endpoint writes request files to disk with bounded memory. Not part of
|
||||
the plugin-facing API surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
COPY_BLOCK_SIZE = 1024 * 1024
|
||||
|
||||
|
||||
class UploadTooLargeError(ValueError):
|
||||
"""Raised when an upload exceeds the allowed size while saving."""
|
||||
|
||||
def __init__(self, max_bytes: int) -> None:
|
||||
super().__init__(f"Uploaded file exceeds the {max_bytes} bytes limit")
|
||||
self.max_bytes = max_bytes
|
||||
|
||||
|
||||
async def save_upload_stream(
|
||||
upload: Any,
|
||||
dest: str | Path,
|
||||
*,
|
||||
max_bytes: int | None = None,
|
||||
) -> int:
|
||||
"""Stream an uploaded file to disk in 1 MiB blocks.
|
||||
|
||||
The payload lands in a uniquely named sibling temp file and is
|
||||
atomically renamed to ``dest`` only after the complete stream succeeds,
|
||||
so a failed or cancelled save never destroys a pre-existing
|
||||
destination (e.g. re-uploading over an existing attachment).
|
||||
|
||||
Args:
|
||||
upload: Upload object exposing async ``read(size)``; ``seek(0)`` is
|
||||
attempted best-effort before copying.
|
||||
dest: Destination file path.
|
||||
max_bytes: Optional hard limit. When the upload exceeds it, the
|
||||
partial temp file is removed and UploadTooLargeError is raised.
|
||||
|
||||
Returns:
|
||||
Number of bytes written.
|
||||
|
||||
Raises:
|
||||
UploadTooLargeError: The upload exceeded ``max_bytes``.
|
||||
"""
|
||||
path = Path(dest)
|
||||
temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
await upload.seek(0)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
written = 0
|
||||
try:
|
||||
with temp_path.open("wb") as output:
|
||||
while True:
|
||||
block = await upload.read(COPY_BLOCK_SIZE)
|
||||
if not block:
|
||||
break
|
||||
written += len(block)
|
||||
if max_bytes is not None and written > max_bytes:
|
||||
raise UploadTooLargeError(max_bytes)
|
||||
output.write(block)
|
||||
await asyncio.to_thread(os.replace, temp_path, path)
|
||||
except BaseException:
|
||||
# Only ever remove our own temp file, never the destination.
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
return written
|
||||
@@ -18,6 +18,7 @@ from astrbot.dashboard.services.backup_service import (
|
||||
)
|
||||
|
||||
from .auth import AuthContext, require_dashboard_user, require_scope
|
||||
from .multipart import UploadFileAdapter
|
||||
|
||||
router = APIRouter(tags=["Backups"])
|
||||
legacy_router = APIRouter(
|
||||
@@ -151,7 +152,10 @@ async def upload_backup(
|
||||
_auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(lambda: service.upload_backup(file), prefix="上传备份文件失败")
|
||||
return await _run(
|
||||
lambda: service.upload_backup(UploadFileAdapter(file)),
|
||||
prefix="上传备份文件失败",
|
||||
)
|
||||
|
||||
|
||||
@legacy_router.post("/upload")
|
||||
@@ -160,17 +164,20 @@ async def upload_dashboard_backup(
|
||||
_username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(lambda: service.upload_backup(file), prefix="上传备份文件失败")
|
||||
return await _run(
|
||||
lambda: service.upload_backup(UploadFileAdapter(file)),
|
||||
prefix="上传备份文件失败",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/backups/upload/init")
|
||||
async def init_backup_upload(
|
||||
payload: BackupUploadInitRequest,
|
||||
_auth: AuthContext = Depends(require_system_scope),
|
||||
auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_init(_model_dict(payload)),
|
||||
lambda: service.upload_init(_model_dict(payload), owner=auth.username),
|
||||
prefix="初始化分片上传失败",
|
||||
)
|
||||
|
||||
@@ -178,11 +185,11 @@ async def init_backup_upload(
|
||||
@legacy_router.post("/upload/init")
|
||||
async def init_dashboard_backup_upload(
|
||||
payload: BackupUploadInitRequest,
|
||||
_username: str = Depends(require_dashboard_user),
|
||||
username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_init(_model_dict(payload)),
|
||||
lambda: service.upload_init(_model_dict(payload), owner=username),
|
||||
prefix="初始化分片上传失败",
|
||||
)
|
||||
|
||||
@@ -192,14 +199,15 @@ async def upload_backup_chunk(
|
||||
upload_id: str = Form(...),
|
||||
chunk_index: str = Form(...),
|
||||
chunk: UploadFile = File(...),
|
||||
_auth: AuthContext = Depends(require_system_scope),
|
||||
auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str=chunk_index,
|
||||
chunk_file=chunk,
|
||||
chunk_file=UploadFileAdapter(chunk),
|
||||
owner=auth.username,
|
||||
),
|
||||
prefix="上传分片失败",
|
||||
)
|
||||
@@ -210,14 +218,15 @@ async def upload_dashboard_backup_chunk(
|
||||
upload_id: str = Form(...),
|
||||
chunk_index: str = Form(...),
|
||||
chunk: UploadFile = File(...),
|
||||
_username: str = Depends(require_dashboard_user),
|
||||
username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str=chunk_index,
|
||||
chunk_file=chunk,
|
||||
chunk_file=UploadFileAdapter(chunk),
|
||||
owner=username,
|
||||
),
|
||||
prefix="上传分片失败",
|
||||
)
|
||||
@@ -226,11 +235,11 @@ async def upload_dashboard_backup_chunk(
|
||||
@router.post("/backups/upload/complete")
|
||||
async def complete_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
_auth: AuthContext = Depends(require_system_scope),
|
||||
auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_complete(_model_dict(payload)),
|
||||
lambda: service.upload_complete(_model_dict(payload), owner=auth.username),
|
||||
prefix="完成分片上传失败",
|
||||
)
|
||||
|
||||
@@ -238,11 +247,11 @@ async def complete_backup_upload(
|
||||
@legacy_router.post("/upload/complete")
|
||||
async def complete_dashboard_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
_username: str = Depends(require_dashboard_user),
|
||||
username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_complete(_model_dict(payload)),
|
||||
lambda: service.upload_complete(_model_dict(payload), owner=username),
|
||||
prefix="完成分片上传失败",
|
||||
)
|
||||
|
||||
@@ -250,11 +259,11 @@ async def complete_dashboard_backup_upload(
|
||||
@router.post("/backups/upload/abort")
|
||||
async def abort_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
_auth: AuthContext = Depends(require_system_scope),
|
||||
auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_abort(_model_dict(payload)),
|
||||
lambda: service.upload_abort(_model_dict(payload), owner=auth.username),
|
||||
prefix="取消上传失败",
|
||||
)
|
||||
|
||||
@@ -262,15 +271,39 @@ async def abort_backup_upload(
|
||||
@legacy_router.post("/upload/abort")
|
||||
async def abort_dashboard_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
_username: str = Depends(require_dashboard_user),
|
||||
username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_abort(_model_dict(payload)),
|
||||
lambda: service.upload_abort(_model_dict(payload), owner=username),
|
||||
prefix="取消上传失败",
|
||||
)
|
||||
|
||||
|
||||
@router.post("/backups/upload/status")
|
||||
async def status_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
auth: AuthContext = Depends(require_system_scope),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_status(_model_dict(payload), owner=auth.username),
|
||||
prefix="查询上传状态失败",
|
||||
)
|
||||
|
||||
|
||||
@legacy_router.post("/upload/status")
|
||||
async def status_dashboard_backup_upload(
|
||||
payload: BackupUploadSessionRequest,
|
||||
username: str = Depends(require_dashboard_user),
|
||||
service: BackupService = Depends(get_service),
|
||||
):
|
||||
return await _run(
|
||||
lambda: service.upload_status(_model_dict(payload), owner=username),
|
||||
prefix="查询上传状态失败",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/backups/tasks/{task_id}")
|
||||
async def get_backup_progress(
|
||||
task_id: str,
|
||||
|
||||
@@ -1,10 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, Query, Request, UploadFile
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
Depends,
|
||||
File,
|
||||
Form,
|
||||
HTTPException,
|
||||
Query,
|
||||
Request,
|
||||
UploadFile,
|
||||
)
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from astrbot.dashboard.async_utils import run_maybe_async
|
||||
from astrbot.dashboard.responses import error, ok
|
||||
from astrbot.dashboard.schemas import ChatUploadInitRequest, ChatUploadSessionRequest
|
||||
from astrbot.dashboard.services.chat_service import ChatService, ChatServiceError
|
||||
from astrbot.dashboard.services.file_service import FileService, FileServiceError
|
||||
|
||||
@@ -49,15 +59,19 @@ async def _run_file(operation, *, error_message: str = "File access error"):
|
||||
return error(error_message)
|
||||
|
||||
|
||||
async def _upload_file(file: UploadFile, service: ChatService):
|
||||
result = await _run_file(
|
||||
lambda: service.save_uploaded_file(UploadFileAdapter(file))
|
||||
)
|
||||
async def _run_chat_upload(operation):
|
||||
result = await _run_file(operation)
|
||||
if isinstance(result, dict) and result.get("status") == "error":
|
||||
return result
|
||||
return ok(result)
|
||||
|
||||
|
||||
async def _upload_file(file: UploadFile, service: ChatService):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.save_uploaded_file(UploadFileAdapter(file))
|
||||
)
|
||||
|
||||
|
||||
@router.get("/files/tokens/{file_token}")
|
||||
async def get_token_file(
|
||||
file_token: str,
|
||||
@@ -75,6 +89,76 @@ async def upload_file(
|
||||
return await _upload_file(file, service)
|
||||
|
||||
|
||||
@router.post("/files/upload/init")
|
||||
async def init_file_upload(
|
||||
payload: ChatUploadInitRequest,
|
||||
auth: AuthContext = Depends(require_file_scope),
|
||||
service: ChatService = Depends(get_chat_service),
|
||||
):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.upload_init(
|
||||
payload.model_dump(exclude_none=True), owner=auth.username
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/files/upload/chunk")
|
||||
async def upload_file_chunk(
|
||||
upload_id: str = Form(...),
|
||||
chunk_index: str = Form(...),
|
||||
chunk: UploadFile = File(...),
|
||||
auth: AuthContext = Depends(require_file_scope),
|
||||
service: ChatService = Depends(get_chat_service),
|
||||
):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str=chunk_index,
|
||||
chunk_file=UploadFileAdapter(chunk),
|
||||
owner=auth.username,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/files/upload/complete")
|
||||
async def complete_file_upload(
|
||||
payload: ChatUploadSessionRequest,
|
||||
auth: AuthContext = Depends(require_file_scope),
|
||||
service: ChatService = Depends(get_chat_service),
|
||||
):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.upload_complete(
|
||||
payload.model_dump(exclude_none=True), owner=auth.username
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/files/upload/abort")
|
||||
async def abort_file_upload(
|
||||
payload: ChatUploadSessionRequest,
|
||||
auth: AuthContext = Depends(require_file_scope),
|
||||
service: ChatService = Depends(get_chat_service),
|
||||
):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.upload_abort(
|
||||
payload.model_dump(exclude_none=True), owner=auth.username
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/files/upload/status")
|
||||
async def status_file_upload(
|
||||
payload: ChatUploadSessionRequest,
|
||||
auth: AuthContext = Depends(require_file_scope),
|
||||
service: ChatService = Depends(get_chat_service),
|
||||
):
|
||||
return await _run_chat_upload(
|
||||
lambda: service.upload_status(
|
||||
payload.model_dump(exclude_none=True), owner=auth.username
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/files/content")
|
||||
async def get_file_by_name(
|
||||
filename: str | None = Query(default=None),
|
||||
|
||||
@@ -7,6 +7,8 @@ from typing import Any
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
|
||||
from astrbot.core.utils.upload import save_upload_stream
|
||||
|
||||
|
||||
class UploadFileAdapter:
|
||||
def __init__(self, upload_file: StarletteUploadFile) -> None:
|
||||
@@ -23,18 +25,12 @@ class UploadFileAdapter:
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
async def save(self, destination: str | Path) -> None:
|
||||
path = Path(destination)
|
||||
try:
|
||||
await self._upload_file.seek(0)
|
||||
except Exception:
|
||||
pass
|
||||
with path.open("wb") as output:
|
||||
while True:
|
||||
chunk = await self._upload_file.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
output.write(chunk)
|
||||
async def save(
|
||||
self, destination: str | Path, *, max_bytes: int | None = None
|
||||
) -> int:
|
||||
return await save_upload_stream(
|
||||
self._upload_file, destination, max_bytes=max_bytes
|
||||
)
|
||||
|
||||
|
||||
class MultiDict:
|
||||
|
||||
@@ -15,6 +15,8 @@ from fastapi.responses import FileResponse, JSONResponse, Response
|
||||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
from starlette.responses import StreamingResponse
|
||||
|
||||
from astrbot.core.utils.upload import save_upload_stream
|
||||
|
||||
_request_var: contextvars.ContextVar[DashboardRequest] = contextvars.ContextVar(
|
||||
"dashboard_request"
|
||||
)
|
||||
@@ -99,18 +101,12 @@ class RequestUploadFile:
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
async def save(self, destination: str | Path) -> None:
|
||||
path = Path(destination)
|
||||
try:
|
||||
await self._upload_file.seek(0)
|
||||
except Exception:
|
||||
pass
|
||||
with path.open("wb") as output:
|
||||
while True:
|
||||
chunk = await self._upload_file.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
output.write(chunk)
|
||||
async def save(
|
||||
self, destination: str | Path, *, max_bytes: int | None = None
|
||||
) -> int:
|
||||
return await save_upload_stream(
|
||||
self._upload_file, destination, max_bytes=max_bytes
|
||||
)
|
||||
|
||||
def __getattr__(self, key: str):
|
||||
return getattr(self._upload_file, key)
|
||||
|
||||
@@ -70,6 +70,16 @@ class BackupUploadSessionRequest(OpenModel):
|
||||
upload_id: str | None = None
|
||||
|
||||
|
||||
class ChatUploadInitRequest(OpenModel):
|
||||
filename: str | None = None
|
||||
total_size: int | None = None
|
||||
content_type: str | None = None
|
||||
|
||||
|
||||
class ChatUploadSessionRequest(OpenModel):
|
||||
upload_id: str | None = None
|
||||
|
||||
|
||||
class BackupImportRequest(OpenModel):
|
||||
confirmed: bool | None = None
|
||||
|
||||
|
||||
@@ -27,15 +27,105 @@ from astrbot.dashboard.asgi_runtime import (
|
||||
FastAPIAppAdapter,
|
||||
)
|
||||
from astrbot.dashboard.responses import error
|
||||
from astrbot.dashboard.services.backup_service import CHUNK_SIZE
|
||||
from astrbot.dashboard.services.chat_service import MAX_UPLOAD_FILE_SIZE_BYTES
|
||||
from astrbot.dashboard.services.config_service import MAX_FILE_BYTES
|
||||
|
||||
from .api.app import create_dashboard_asgi_app
|
||||
from .plugin_page_auth import PluginPageAuth
|
||||
from .services.auth_service import DASHBOARD_JWT_COOKIE_NAME
|
||||
|
||||
try: # Mirror starlette.requests so media type detection matches the parser.
|
||||
from python_multipart.multipart import parse_options_header
|
||||
except ImportError: # pragma: no cover
|
||||
try:
|
||||
from multipart.multipart import parse_options_header
|
||||
except ImportError:
|
||||
parse_options_header = None
|
||||
|
||||
if os.name == "nt":
|
||||
# Windows 的 mimetypes 会把 .svg 映射成非标准的 image/svg,这里强制覆盖为标准类型
|
||||
mimetypes.add_type("image/svg+xml", ".svg", strict=True)
|
||||
|
||||
# Multipart framing (boundaries, part headers) rides on top of the file
|
||||
# payload, so whole-file upload routes get slack beyond the file size limit;
|
||||
# otherwise a file exactly at the limit would be rejected with 413.
|
||||
_MULTIPART_OVERHEAD_BYTES = 1024 * 1024
|
||||
|
||||
# Per-route request body limits overriding the default MAX_CONTENT_LENGTH.
|
||||
# More specific prefixes must come first. Routes not listed here fall back
|
||||
# to the default; requests without a Content-Length header pass through and
|
||||
# are bounded by the per-endpoint max_bytes checks at save time.
|
||||
_BODY_LIMIT_OVERRIDES: tuple[tuple[str, int], ...] = (
|
||||
("/api/v1/backups/upload/chunk", CHUNK_SIZE * 2),
|
||||
("/api/backup/upload/chunk", CHUNK_SIZE * 2),
|
||||
("/api/v1/files/upload/chunk", CHUNK_SIZE * 2),
|
||||
("/api/v1/files", MAX_UPLOAD_FILE_SIZE_BYTES + _MULTIPART_OVERHEAD_BYTES),
|
||||
(
|
||||
"/api/chat/post_file",
|
||||
MAX_UPLOAD_FILE_SIZE_BYTES + _MULTIPART_OVERHEAD_BYTES,
|
||||
),
|
||||
("/api/v1/plugins/config-files", MAX_FILE_BYTES + _MULTIPART_OVERHEAD_BYTES),
|
||||
(
|
||||
"/api/v1/knowledge-bases/",
|
||||
MAX_UPLOAD_FILE_SIZE_BYTES + _MULTIPART_OVERHEAD_BYTES,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# Methods with request-body semantics; the 411 stopgap only applies to
|
||||
# these, since form parsing (and its disk spooling) cannot trigger
|
||||
# for body-less methods like GET.
|
||||
_BODY_METHODS = frozenset({"POST", "PUT", "PATCH"})
|
||||
|
||||
|
||||
def _check_body_limit(
|
||||
path: str,
|
||||
content_length: int | None,
|
||||
content_type: str,
|
||||
*,
|
||||
method: str = "POST",
|
||||
default_limit: int,
|
||||
) -> tuple[int, str] | None:
|
||||
"""Decide whether an /api request body must be rejected up front.
|
||||
|
||||
Returns:
|
||||
A (status_code, message) rejection, or None to pass through.
|
||||
|
||||
Note:
|
||||
The 411 rule is a stopgap scoped to multipart uploads on methods
|
||||
with body semantics: their form parsing spools large bodies to
|
||||
disk before any per-file size check can run, so they must declare
|
||||
a length that can be bounded before parsing. Body-less methods
|
||||
(GET etc.) never trigger form parsing and pass even with a bogus
|
||||
multipart Content-Type header. Other lengthless bodies are still
|
||||
bounded only at save time; closing that gap fully requires
|
||||
counting bytes as they arrive, which is out of scope here.
|
||||
"""
|
||||
if not path.startswith("/api"):
|
||||
return None
|
||||
if content_length is None:
|
||||
if method not in _BODY_METHODS:
|
||||
return None
|
||||
# Identify the media type by the same rule the form parser uses:
|
||||
# parse_options_header strips surrounding whitespace, so a leading-
|
||||
# space Content-Type that startswith() would miss is still caught.
|
||||
media_type = (
|
||||
parse_options_header(content_type)[0] if parse_options_header else b""
|
||||
)
|
||||
if media_type == b"multipart/form-data":
|
||||
return 411, "Content-Length header is required for uploads"
|
||||
return None
|
||||
limit = default_limit
|
||||
for prefix, route_limit in _BODY_LIMIT_OVERRIDES:
|
||||
if path.startswith(prefix):
|
||||
limit = route_limit
|
||||
break
|
||||
if content_length > limit:
|
||||
return 413, f"Request body exceeds the {limit} bytes limit"
|
||||
return None
|
||||
|
||||
|
||||
_RATE_LIMITED_ENDPOINTS: frozenset = frozenset(
|
||||
{
|
||||
"/api/config/astrbot/update",
|
||||
@@ -210,6 +300,27 @@ class AstrBotDashboard:
|
||||
return auth_response
|
||||
return await call_next(request_)
|
||||
|
||||
@self.asgi_app.middleware("http")
|
||||
async def dashboard_body_limit_middleware(request_, call_next):
|
||||
# Registered after the auth middleware so it runs outermost and
|
||||
# can reject oversized bodies before any parsing happens.
|
||||
raw_length = request_.headers.get("content-length")
|
||||
try:
|
||||
content_length = int(raw_length) if raw_length else None
|
||||
except ValueError:
|
||||
content_length = None
|
||||
rejection = _check_body_limit(
|
||||
request_.url.path,
|
||||
content_length,
|
||||
request_.headers.get("content-type", ""),
|
||||
method=request_.method,
|
||||
default_limit=self.app.config["MAX_CONTENT_LENGTH"],
|
||||
)
|
||||
if rejection is not None:
|
||||
status_code, message = rejection
|
||||
return JSONResponse(error(message), status_code=status_code)
|
||||
return await call_next(request_)
|
||||
|
||||
self.shutdown_event = shutdown_event
|
||||
|
||||
async def auth_middleware(self, current_request: Request):
|
||||
|
||||
@@ -2,11 +2,8 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
import zipfile
|
||||
@@ -26,9 +23,18 @@ from astrbot.core.utils.astrbot_path import (
|
||||
get_astrbot_backups_path,
|
||||
get_astrbot_data_path,
|
||||
)
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
from astrbot.dashboard.services.chunked_upload_service import (
|
||||
ChunkedUploadError,
|
||||
ChunkedUploadService,
|
||||
)
|
||||
|
||||
CHUNK_SIZE = 1024 * 1024
|
||||
UPLOAD_EXPIRE_SECONDS = 3600
|
||||
# Hard caps against disk exhaustion: a backup is never legitimately larger
|
||||
# than this, and the whole-file endpoint is only for small backups (large
|
||||
# ones must use the chunked flow).
|
||||
MAX_BACKUP_TOTAL_BYTES = 8 * 1024 * 1024 * 1024
|
||||
MAX_DIRECT_UPLOAD_BYTES = 128 * 1024 * 1024
|
||||
|
||||
|
||||
class BackupServiceError(Exception):
|
||||
@@ -72,30 +78,12 @@ class BackupService:
|
||||
self.chunks_dir = os.path.join(self.backup_dir, ".chunks")
|
||||
self.backup_tasks: dict[str, dict] = {}
|
||||
self.backup_progress: dict[str, dict] = {}
|
||||
self.upload_sessions: dict[str, dict] = {}
|
||||
self._cleanup_task: asyncio.Task | None = None
|
||||
self.chunked_uploads = ChunkedUploadService(self.chunks_dir)
|
||||
|
||||
@staticmethod
|
||||
def _payload(data: object) -> dict[str, Any]:
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
@staticmethod
|
||||
async def _save_upload(file: Any, target_path: str) -> None:
|
||||
if hasattr(file, "save"):
|
||||
result = file.save(target_path)
|
||||
if hasattr(result, "__await__"):
|
||||
await result
|
||||
return
|
||||
|
||||
if hasattr(file, "read"):
|
||||
data = file.read()
|
||||
if hasattr(data, "__await__"):
|
||||
data = await data
|
||||
Path(target_path).write_bytes(data)
|
||||
return
|
||||
|
||||
raise BackupServiceError("无效的上传文件")
|
||||
|
||||
@staticmethod
|
||||
def _validate_backup_filename(filename: str | None, *, missing: str) -> str:
|
||||
if not filename:
|
||||
@@ -176,45 +164,10 @@ class BackupService:
|
||||
return _callback
|
||||
|
||||
def ensure_cleanup_task_started(self) -> None:
|
||||
if self._cleanup_task is None or self._cleanup_task.done():
|
||||
try:
|
||||
self._cleanup_task = asyncio.create_task(
|
||||
self._cleanup_expired_uploads()
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
async def _cleanup_expired_uploads(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(300)
|
||||
current_time = time.time()
|
||||
expired_sessions = []
|
||||
|
||||
for upload_id, session in self.upload_sessions.items():
|
||||
last_activity = session.get("last_activity", session["created_at"])
|
||||
if current_time - last_activity > UPLOAD_EXPIRE_SECONDS:
|
||||
expired_sessions.append(upload_id)
|
||||
|
||||
for upload_id in expired_sessions:
|
||||
await self.cleanup_upload_session(upload_id)
|
||||
logger.info(f"清理过期的上传会话: {upload_id}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.error(f"清理过期上传会话失败: {exc}")
|
||||
self.chunked_uploads.ensure_cleanup_task_started()
|
||||
|
||||
async def cleanup_upload_session(self, upload_id: str) -> None:
|
||||
if upload_id in self.upload_sessions:
|
||||
session = self.upload_sessions[upload_id]
|
||||
chunk_dir = session.get("chunk_dir")
|
||||
if chunk_dir and os.path.exists(chunk_dir):
|
||||
try:
|
||||
shutil.rmtree(chunk_dir)
|
||||
except Exception as exc:
|
||||
logger.warning(f"清理分片目录失败: {exc}")
|
||||
del self.upload_sessions[upload_id]
|
||||
await self.chunked_uploads.cleanup_session(upload_id)
|
||||
|
||||
def get_backup_manifest(self, zip_path: str) -> dict | None:
|
||||
try:
|
||||
@@ -315,7 +268,13 @@ class BackupService:
|
||||
|
||||
Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
|
||||
zip_path = os.path.join(self.backup_dir, unique_filename)
|
||||
await self._save_upload(file, zip_path)
|
||||
try:
|
||||
await file.save(zip_path, max_bytes=MAX_DIRECT_UPLOAD_BYTES)
|
||||
except UploadTooLargeError as exc:
|
||||
raise BackupServiceError(
|
||||
f"Backup file exceeds the size limit ({MAX_DIRECT_UPLOAD_BYTES // (1024**2)} MB); "
|
||||
"use chunked upload instead."
|
||||
) from exc
|
||||
|
||||
logger.info(
|
||||
f"上传的备份文件已保存: {unique_filename} (原始名称: {file.filename})"
|
||||
@@ -326,7 +285,7 @@ class BackupService:
|
||||
"size": os.path.getsize(zip_path),
|
||||
}
|
||||
|
||||
def upload_init(self, data: object) -> dict:
|
||||
def upload_init(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = self._payload(data)
|
||||
filename = payload.get("filename")
|
||||
total_size = payload.get("total_size", 0)
|
||||
@@ -337,35 +296,30 @@ class BackupService:
|
||||
raise BackupServiceError("请上传 ZIP 格式的备份文件")
|
||||
if total_size <= 0:
|
||||
raise BackupServiceError("无效的文件大小")
|
||||
if total_size > MAX_BACKUP_TOTAL_BYTES:
|
||||
raise BackupServiceError(
|
||||
f"Backup file exceeds the size limit ({MAX_BACKUP_TOTAL_BYTES // (1024**3)} GB). "
|
||||
"You can copy it into the backups folder of the data directory "
|
||||
"via FTP/SFTP and restore it from the backup list."
|
||||
)
|
||||
|
||||
total_chunks = math.ceil(total_size / CHUNK_SIZE)
|
||||
upload_id = str(uuid.uuid4())
|
||||
chunk_dir = os.path.join(self.chunks_dir, upload_id)
|
||||
Path(chunk_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
safe_filename = secure_filename(filename)
|
||||
unique_filename = generate_unique_filename(safe_filename)
|
||||
current_time = time.time()
|
||||
self.upload_sessions[upload_id] = {
|
||||
"filename": unique_filename,
|
||||
"original_filename": filename,
|
||||
"total_size": total_size,
|
||||
"total_chunks": total_chunks,
|
||||
"received_chunks": set(),
|
||||
"created_at": current_time,
|
||||
"last_activity": current_time,
|
||||
"chunk_dir": chunk_dir,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
f"初始化分片上传: upload_id={upload_id}, "
|
||||
f"filename={unique_filename}, total_chunks={total_chunks}"
|
||||
)
|
||||
unique_filename = generate_unique_filename(secure_filename(filename))
|
||||
self.chunked_uploads.ensure_cleanup_task_started()
|
||||
try:
|
||||
session = self.chunked_uploads.init_session(
|
||||
owner=owner,
|
||||
purpose="backup",
|
||||
filename=unique_filename,
|
||||
original_filename=filename,
|
||||
total_size=total_size,
|
||||
)
|
||||
except ChunkedUploadError as exc:
|
||||
raise BackupServiceError(str(exc)) from exc
|
||||
|
||||
return {
|
||||
"upload_id": upload_id,
|
||||
"chunk_size": CHUNK_SIZE,
|
||||
"total_chunks": total_chunks,
|
||||
"upload_id": session.id,
|
||||
"chunk_size": session.chunk_size,
|
||||
"total_chunks": session.total_chunks,
|
||||
"filename": unique_filename,
|
||||
}
|
||||
|
||||
@@ -375,6 +329,7 @@ class BackupService:
|
||||
upload_id: str | None,
|
||||
chunk_index_str: str | None,
|
||||
chunk_file: Any | None,
|
||||
owner: str = "",
|
||||
) -> dict:
|
||||
if not upload_id or chunk_index_str is None:
|
||||
raise BackupServiceError("缺少必要参数")
|
||||
@@ -386,29 +341,13 @@ class BackupService:
|
||||
|
||||
if not chunk_file:
|
||||
raise BackupServiceError("缺少分片数据")
|
||||
if upload_id not in self.upload_sessions:
|
||||
raise BackupServiceError("上传会话不存在或已过期")
|
||||
|
||||
session = self.upload_sessions[upload_id]
|
||||
if chunk_index < 0 or chunk_index >= session["total_chunks"]:
|
||||
raise BackupServiceError("分片索引超出范围")
|
||||
|
||||
chunk_path = os.path.join(session["chunk_dir"], f"{chunk_index}.part")
|
||||
await self._save_upload(chunk_file, chunk_path)
|
||||
session["received_chunks"].add(chunk_index)
|
||||
session["last_activity"] = time.time()
|
||||
|
||||
received_count = len(session["received_chunks"])
|
||||
total_chunks = session["total_chunks"]
|
||||
logger.debug(
|
||||
f"接收分片: upload_id={upload_id}, chunk={chunk_index + 1}/{total_chunks}"
|
||||
)
|
||||
|
||||
return {
|
||||
"received": received_count,
|
||||
"total": total_chunks,
|
||||
"chunk_index": chunk_index,
|
||||
}
|
||||
try:
|
||||
return await self.chunked_uploads.save_chunk(
|
||||
upload_id, chunk_index, chunk_file, owner=owner
|
||||
)
|
||||
except ChunkedUploadError as exc:
|
||||
raise BackupServiceError(str(exc)) from exc
|
||||
|
||||
def mark_backup_as_uploaded(self, zip_path: str) -> None:
|
||||
try:
|
||||
@@ -428,67 +367,62 @@ class BackupService:
|
||||
except Exception as exc:
|
||||
logger.warning(f"标记备份来源失败: {exc}")
|
||||
|
||||
async def upload_complete(self, data: object) -> dict:
|
||||
async def upload_complete(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = self._payload(data)
|
||||
upload_id = payload.get("upload_id")
|
||||
|
||||
if not upload_id:
|
||||
raise BackupServiceError("缺少 upload_id 参数")
|
||||
if upload_id not in self.upload_sessions:
|
||||
raise BackupServiceError("上传会话不存在或已过期")
|
||||
|
||||
session = self.upload_sessions[upload_id]
|
||||
received = session["received_chunks"]
|
||||
total = session["total_chunks"]
|
||||
|
||||
if len(received) != total:
|
||||
missing = set(range(total)) - received
|
||||
raise BackupServiceError(f"分片不完整，缺少: {sorted(missing)[:10]}...")
|
||||
|
||||
chunk_dir = session["chunk_dir"]
|
||||
filename = session["filename"]
|
||||
|
||||
Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
|
||||
output_path = os.path.join(self.backup_dir, filename)
|
||||
|
||||
try:
|
||||
with open(output_path, "wb") as outfile:
|
||||
for i in range(total):
|
||||
chunk_path = os.path.join(chunk_dir, f"{i}.part")
|
||||
with open(chunk_path, "rb") as chunk_file:
|
||||
while True:
|
||||
data_block = chunk_file.read(8192)
|
||||
if not data_block:
|
||||
break
|
||||
outfile.write(data_block)
|
||||
session = self.chunked_uploads.get_session(upload_id, owner=owner)
|
||||
Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
|
||||
output_path = os.path.join(self.backup_dir, session.filename)
|
||||
file_size = await self.chunked_uploads.assemble(
|
||||
upload_id, output_path, owner=owner
|
||||
)
|
||||
except ChunkedUploadError as exc:
|
||||
raise BackupServiceError(str(exc)) from exc
|
||||
|
||||
file_size = os.path.getsize(output_path)
|
||||
self.mark_backup_as_uploaded(output_path)
|
||||
logger.info(f"分片上传完成: {filename}, size={file_size}, chunks={total}")
|
||||
await self.cleanup_upload_session(upload_id)
|
||||
self.mark_backup_as_uploaded(output_path)
|
||||
logger.info(
|
||||
f"分片上传完成: {session.filename}, size={file_size}, "
|
||||
f"chunks={session.total_chunks}"
|
||||
)
|
||||
|
||||
return {
|
||||
"filename": filename,
|
||||
"original_filename": session["original_filename"],
|
||||
"size": file_size,
|
||||
}
|
||||
except Exception:
|
||||
if os.path.exists(output_path):
|
||||
os.remove(output_path)
|
||||
raise
|
||||
return {
|
||||
"filename": session.filename,
|
||||
"original_filename": session.original_filename,
|
||||
"size": file_size,
|
||||
}
|
||||
|
||||
async def upload_abort(self, data: object) -> tuple[dict | None, str | None]:
|
||||
async def upload_abort(
|
||||
self, data: object, *, owner: str = ""
|
||||
) -> tuple[dict | None, str | None]:
|
||||
payload = self._payload(data)
|
||||
upload_id = payload.get("upload_id")
|
||||
if not upload_id:
|
||||
raise BackupServiceError("缺少 upload_id 参数")
|
||||
|
||||
if upload_id in self.upload_sessions:
|
||||
await self.cleanup_upload_session(upload_id)
|
||||
logger.info(f"取消分片上传: {upload_id}")
|
||||
try:
|
||||
if await self.chunked_uploads.abort(upload_id, owner=owner):
|
||||
logger.info(f"取消分片上传: {upload_id}")
|
||||
except ChunkedUploadError as exc:
|
||||
raise BackupServiceError(str(exc)) from exc
|
||||
|
||||
return None, "上传已取消"
|
||||
|
||||
def upload_status(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = self._payload(data)
|
||||
upload_id = payload.get("upload_id")
|
||||
if not upload_id:
|
||||
raise BackupServiceError("缺少 upload_id 参数")
|
||||
|
||||
try:
|
||||
return self.chunked_uploads.session_status(upload_id, owner=owner)
|
||||
except ChunkedUploadError as exc:
|
||||
raise BackupServiceError(str(exc)) from exc
|
||||
|
||||
def check_backup(self, data: object) -> dict:
|
||||
payload = self._payload(data)
|
||||
filename = self._validate_backup_filename(
|
||||
|
||||
@@ -33,6 +33,11 @@ from astrbot.core.utils.media_utils import (
|
||||
MEDIA_MIME_EXTENSIONS,
|
||||
detect_image_mime_type_async,
|
||||
)
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
from astrbot.dashboard.services.chunked_upload_service import (
|
||||
ChunkedUploadError,
|
||||
ChunkedUploadService,
|
||||
)
|
||||
|
||||
SSE_HEARTBEAT = ": heartbeat\n\n"
|
||||
CHAT_RUN_SUBSCRIBER_QUEUE_SIZE = 256
|
||||
@@ -58,6 +63,30 @@ def sanitize_upload_filename(filename: str | None) -> str:
|
||||
return name
|
||||
|
||||
|
||||
class LocalUploadFile:
|
||||
"""Adapt a merged local file to the upload contract save_uploaded_file() consumes.
|
||||
|
||||
Mirrors UploadFileAdapter (api/multipart.py) but sources bytes from a
|
||||
local path; save() moves the file so the chunked-upload merge result
|
||||
lands in the attachments directory without a second full copy.
|
||||
"""
|
||||
|
||||
def __init__(self, path: Path, filename: str, content_type: str | None) -> None:
|
||||
self._path = path
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
self.headers: dict = {}
|
||||
self.content_length = path.stat().st_size
|
||||
|
||||
async def save(
|
||||
self, destination: str | Path, *, max_bytes: int | None = None
|
||||
) -> int:
|
||||
if max_bytes is not None and self.content_length > max_bytes:
|
||||
raise UploadTooLargeError(max_bytes)
|
||||
await asyncio.to_thread(os.replace, self._path, destination)
|
||||
return self.content_length
|
||||
|
||||
|
||||
def normalize_reasoning_message_parts(
|
||||
message_parts: list[dict] | None,
|
||||
reasoning: str = "",
|
||||
@@ -521,6 +550,10 @@ class ChatService:
|
||||
self.chat_runs: dict[str, ChatRunState] = {}
|
||||
self.chat_runs_by_session: dict[str, set[str]] = {}
|
||||
|
||||
self.chunked_uploads = ChunkedUploadService(
|
||||
os.path.join(get_astrbot_data_path(), "webchat", ".chunks")
|
||||
)
|
||||
|
||||
async def build_user_message_parts(self, message: str | list) -> list[dict]:
|
||||
return await build_webchat_message_parts(
|
||||
message,
|
||||
@@ -620,12 +653,12 @@ class ChatService:
|
||||
if not file_path.is_relative_to(attachments_dir):
|
||||
raise ChatServiceError("Invalid filename")
|
||||
|
||||
await file.save(str(file_path))
|
||||
if file_path.stat().st_size > MAX_UPLOAD_FILE_SIZE_BYTES:
|
||||
file_path.unlink(missing_ok=True)
|
||||
try:
|
||||
await file.save(str(file_path), max_bytes=MAX_UPLOAD_FILE_SIZE_BYTES)
|
||||
except UploadTooLargeError as exc:
|
||||
raise ChatServiceError(
|
||||
f"File too large (limit {MAX_UPLOAD_FILE_SIZE_MB} MB)"
|
||||
)
|
||||
) from exc
|
||||
if attach_type == "image":
|
||||
detected_mime_type = await detect_image_mime_type_async(
|
||||
file_path,
|
||||
@@ -663,6 +696,118 @@ class ChatService:
|
||||
raise ChatServiceError("Missing key: file")
|
||||
return await self.save_uploaded_file(files["file"])
|
||||
|
||||
def upload_init(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = data if isinstance(data, dict) else {}
|
||||
original_filename = str(payload.get("filename") or "")
|
||||
total_size = payload.get("total_size", 0)
|
||||
|
||||
if not original_filename:
|
||||
raise ChatServiceError("Missing key: filename")
|
||||
if not isinstance(total_size, int) or total_size <= 0:
|
||||
raise ChatServiceError("Invalid file size")
|
||||
if total_size > MAX_UPLOAD_FILE_SIZE_BYTES:
|
||||
raise ChatServiceError(
|
||||
f"File too large (limit {MAX_UPLOAD_FILE_SIZE_MB} MB)"
|
||||
)
|
||||
|
||||
self.chunked_uploads.ensure_cleanup_task_started()
|
||||
try:
|
||||
session = self.chunked_uploads.init_session(
|
||||
owner=owner,
|
||||
purpose="chat_attachment",
|
||||
filename=sanitize_upload_filename(original_filename),
|
||||
original_filename=original_filename,
|
||||
total_size=total_size,
|
||||
meta={
|
||||
"content_type": payload.get("content_type")
|
||||
or "application/octet-stream"
|
||||
},
|
||||
)
|
||||
except ChunkedUploadError as exc:
|
||||
raise ChatServiceError(str(exc)) from exc
|
||||
|
||||
return {
|
||||
"upload_id": session.id,
|
||||
"chunk_size": session.chunk_size,
|
||||
"total_chunks": session.total_chunks,
|
||||
}
|
||||
|
||||
async def upload_chunk(
|
||||
self,
|
||||
*,
|
||||
upload_id: str | None,
|
||||
chunk_index_str: str | None,
|
||||
chunk_file: Any | None,
|
||||
owner: str = "",
|
||||
) -> dict:
|
||||
if not upload_id or chunk_index_str is None or not chunk_file:
|
||||
raise ChatServiceError("Missing required parameters")
|
||||
|
||||
try:
|
||||
chunk_index = int(chunk_index_str)
|
||||
except ValueError as exc:
|
||||
raise ChatServiceError("Invalid chunk_index") from exc
|
||||
|
||||
try:
|
||||
return await self.chunked_uploads.save_chunk(
|
||||
upload_id, chunk_index, chunk_file, owner=owner
|
||||
)
|
||||
except ChunkedUploadError as exc:
|
||||
raise ChatServiceError(str(exc)) from exc
|
||||
|
||||
async def upload_complete(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = data if isinstance(data, dict) else {}
|
||||
upload_id = payload.get("upload_id")
|
||||
if not upload_id:
|
||||
raise ChatServiceError("Missing key: upload_id")
|
||||
|
||||
try:
|
||||
session = self.chunked_uploads.get_session(upload_id, owner=owner)
|
||||
# Merge outside the chunk dir: assemble() removes that dir on success.
|
||||
merged_path = self.chunked_uploads.chunks_root / f"{upload_id}.merged"
|
||||
await self.chunked_uploads.assemble(upload_id, merged_path, owner=owner)
|
||||
except ChunkedUploadError as exc:
|
||||
raise ChatServiceError(str(exc)) from exc
|
||||
|
||||
try:
|
||||
upload = LocalUploadFile(
|
||||
merged_path,
|
||||
session.filename,
|
||||
session.meta.get("content_type"),
|
||||
)
|
||||
result = await self.save_uploaded_file(upload)
|
||||
finally:
|
||||
merged_path.unlink(missing_ok=True)
|
||||
|
||||
logger.info(
|
||||
f"Chunked attachment upload completed: {session.filename}, "
|
||||
f"size={session.total_size}, chunks={session.total_chunks}"
|
||||
)
|
||||
return result
|
||||
|
||||
async def upload_abort(self, data: object, *, owner: str = "") -> None:
|
||||
payload = data if isinstance(data, dict) else {}
|
||||
upload_id = payload.get("upload_id")
|
||||
if not upload_id:
|
||||
return
|
||||
|
||||
try:
|
||||
if await self.chunked_uploads.abort(upload_id, owner=owner):
|
||||
logger.info(f"Aborted chunked attachment upload: {upload_id}")
|
||||
except ChunkedUploadError as exc:
|
||||
raise ChatServiceError(str(exc)) from exc
|
||||
|
||||
def upload_status(self, data: object, *, owner: str = "") -> dict:
|
||||
payload = data if isinstance(data, dict) else {}
|
||||
upload_id = payload.get("upload_id")
|
||||
if not upload_id:
|
||||
raise ChatServiceError("Missing key: upload_id")
|
||||
|
||||
try:
|
||||
return self.chunked_uploads.session_status(upload_id, owner=owner)
|
||||
except ChunkedUploadError as exc:
|
||||
raise ChatServiceError(str(exc)) from exc
|
||||
|
||||
async def delete_threads_by_ids(self, thread_ids: list[str], creator: str) -> None:
|
||||
for thread_id in thread_ids:
|
||||
unified_msg_origin = build_thread_unified_msg_origin(creator, thread_id)
|
||||
|
||||
@@ -0,0 +1,327 @@
|
||||
"""Generic chunked-upload session management.
|
||||
|
||||
Business services keep ownership of file validation, final naming and
|
||||
post-merge handling; this service owns the session lifecycle, chunk storage,
|
||||
streaming assembly and expiry cleanup. Every session is bound to an owner
|
||||
(the authenticated username) and a purpose (the consuming business), so a
|
||||
session created for one business cannot be consumed by another caller or
|
||||
another upload flow.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from astrbot.core import logger
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
|
||||
DEFAULT_CHUNK_SIZE = 1024 * 1024
|
||||
DEFAULT_EXPIRE_SECONDS = 3600
|
||||
CLEANUP_INTERVAL_SECONDS = 300
|
||||
|
||||
|
||||
class ChunkedUploadError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class UploadSession:
|
||||
id: str
|
||||
owner: str
|
||||
purpose: str
|
||||
filename: str
|
||||
original_filename: str
|
||||
total_size: int
|
||||
total_chunks: int
|
||||
chunk_size: int
|
||||
chunk_dir: Path
|
||||
meta: dict = field(default_factory=dict)
|
||||
received_chunks: set[int] = field(default_factory=set)
|
||||
created_at: float = 0.0
|
||||
last_activity: float = 0.0
|
||||
|
||||
|
||||
class ChunkedUploadService:
|
||||
"""Manage chunked upload sessions on disk with bounded memory."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunks_root: str | Path,
|
||||
*,
|
||||
chunk_size: int = DEFAULT_CHUNK_SIZE,
|
||||
expire_seconds: int = DEFAULT_EXPIRE_SECONDS,
|
||||
) -> None:
|
||||
self.chunks_root = Path(chunks_root)
|
||||
self.chunk_size = chunk_size
|
||||
self.expire_seconds = expire_seconds
|
||||
self.sessions: dict[str, UploadSession] = {}
|
||||
self._cleanup_task: asyncio.Task | None = None
|
||||
|
||||
def init_session(
|
||||
self,
|
||||
*,
|
||||
owner: str,
|
||||
purpose: str,
|
||||
filename: str,
|
||||
original_filename: str,
|
||||
total_size: int,
|
||||
meta: dict | None = None,
|
||||
) -> UploadSession:
|
||||
"""Create a session and its chunk directory.
|
||||
|
||||
Args:
|
||||
owner: Authenticated username the session is bound to.
|
||||
purpose: Consuming business identifier (e.g. "backup").
|
||||
filename: Sanitized final file name chosen by the business.
|
||||
original_filename: Name as reported by the client.
|
||||
total_size: Declared total size in bytes; the business validates
|
||||
any upper bound before calling, so only sanity is checked here.
|
||||
meta: Opaque business metadata carried with the session (e.g. the
|
||||
client-reported content type) and read back at assembly time.
|
||||
|
||||
Returns:
|
||||
The created session.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: The declared size is not positive.
|
||||
"""
|
||||
if total_size <= 0:
|
||||
raise ChunkedUploadError("Invalid file size")
|
||||
|
||||
upload_id = str(uuid.uuid4())
|
||||
chunk_dir = self.chunks_root / upload_id
|
||||
chunk_dir.mkdir(parents=True, exist_ok=True)
|
||||
now = time.time()
|
||||
session = UploadSession(
|
||||
id=upload_id,
|
||||
owner=owner,
|
||||
purpose=purpose,
|
||||
filename=filename,
|
||||
original_filename=original_filename,
|
||||
total_size=total_size,
|
||||
total_chunks=math.ceil(total_size / self.chunk_size),
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_dir=chunk_dir,
|
||||
meta=meta or {},
|
||||
created_at=now,
|
||||
last_activity=now,
|
||||
)
|
||||
self.sessions[upload_id] = session
|
||||
logger.info(
|
||||
f"Chunked upload session started: id={upload_id}, owner={owner}, "
|
||||
f"purpose={purpose}, file={filename}, chunks={session.total_chunks}"
|
||||
)
|
||||
return session
|
||||
|
||||
def get_session(self, upload_id: str, *, owner: str | None = None) -> UploadSession:
|
||||
"""Look up a session, optionally enforcing its owner.
|
||||
|
||||
Missing sessions, expired sessions and owner mismatches raise the
|
||||
same error so callers cannot probe which sessions exist. The janitor
|
||||
reaps expired chunk directories; this check keeps expired sessions
|
||||
from staying usable in the meantime.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: Session unknown, expired or owned by someone else.
|
||||
"""
|
||||
session = self.sessions.get(upload_id)
|
||||
if (
|
||||
session is None
|
||||
or (owner is not None and session.owner != owner)
|
||||
or time.time() - session.last_activity > self.expire_seconds
|
||||
):
|
||||
raise ChunkedUploadError("Upload session not found or expired")
|
||||
return session
|
||||
|
||||
async def save_chunk(
|
||||
self,
|
||||
upload_id: str,
|
||||
chunk_index: int,
|
||||
file: Any,
|
||||
*,
|
||||
owner: str | None = None,
|
||||
) -> dict:
|
||||
"""Persist one chunk; repeats and out-of-order arrivals are idempotent.
|
||||
|
||||
Args:
|
||||
file: Upload object exposing the adapter ``save()`` contract.
|
||||
|
||||
Returns:
|
||||
Progress dict with received/total/chunk_index.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: Unknown session, bad index or oversized chunk.
|
||||
"""
|
||||
session = self.get_session(upload_id, owner=owner)
|
||||
if chunk_index < 0 or chunk_index >= session.total_chunks:
|
||||
raise ChunkedUploadError("Chunk index out of range")
|
||||
|
||||
# Save to a uniquely named temp file and publish it atomically:
|
||||
# a failed or concurrent retry for the same index must not clobber
|
||||
# the chunk that was already received.
|
||||
chunk_path = session.chunk_dir / f"{chunk_index}.part"
|
||||
temp_path = session.chunk_dir / f"{chunk_index}.{uuid.uuid4().hex}.tmp"
|
||||
try:
|
||||
written = await file.save(temp_path, max_bytes=session.chunk_size)
|
||||
except BaseException as exc:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
if isinstance(exc, UploadTooLargeError):
|
||||
raise ChunkedUploadError("Chunk exceeds the size limit") from exc
|
||||
raise
|
||||
|
||||
# The save contract returns the byte count; a short or absent write
|
||||
# means the chunk is unusable and must not be published.
|
||||
expected = min(
|
||||
session.total_size - chunk_index * session.chunk_size,
|
||||
session.chunk_size,
|
||||
)
|
||||
if written != expected:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise ChunkedUploadError(
|
||||
f"Chunk size mismatch: got {written} bytes, expected {expected}"
|
||||
)
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(os.replace, temp_path, chunk_path)
|
||||
except BaseException:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
session.received_chunks.add(chunk_index)
|
||||
session.last_activity = time.time()
|
||||
|
||||
logger.debug(
|
||||
f"Chunk received: id={upload_id}, "
|
||||
f"chunk={chunk_index + 1}/{session.total_chunks}"
|
||||
)
|
||||
return {
|
||||
"received": len(session.received_chunks),
|
||||
"total": session.total_chunks,
|
||||
"chunk_index": chunk_index,
|
||||
}
|
||||
|
||||
def session_status(self, upload_id: str, *, owner: str | None = None) -> dict:
|
||||
"""Return resumable progress for a session.
|
||||
|
||||
Read-only on purpose: it must not refresh ``last_activity``, or
|
||||
polling this endpoint would anchor a session (and its on-disk
|
||||
chunks) alive forever, defeating the expiry mechanism.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: Session unknown, expired or owned by someone else.
|
||||
"""
|
||||
session = self.get_session(upload_id, owner=owner)
|
||||
remaining = self.expire_seconds - (time.time() - session.last_activity)
|
||||
return {
|
||||
"received_chunks": sorted(session.received_chunks),
|
||||
"total_chunks": session.total_chunks,
|
||||
"chunk_size": session.chunk_size,
|
||||
"expires_in": max(0, int(remaining)),
|
||||
}
|
||||
|
||||
async def assemble(
|
||||
self,
|
||||
upload_id: str,
|
||||
dest: str | Path,
|
||||
*,
|
||||
owner: str | None = None,
|
||||
) -> int:
|
||||
"""Merge all chunks into dest, verify the size and clean up.
|
||||
|
||||
Returns:
|
||||
Size in bytes of the merged file.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: Chunks missing or merged size mismatch.
|
||||
"""
|
||||
session = self.get_session(upload_id, owner=owner)
|
||||
if len(session.received_chunks) != session.total_chunks:
|
||||
missing = sorted(set(range(session.total_chunks)) - session.received_chunks)
|
||||
raise ChunkedUploadError(f"Chunks incomplete, missing: {missing[:10]}...")
|
||||
|
||||
dest_path = Path(dest)
|
||||
try:
|
||||
with dest_path.open("wb") as outfile:
|
||||
for i in range(session.total_chunks):
|
||||
with (session.chunk_dir / f"{i}.part").open("rb") as part:
|
||||
shutil.copyfileobj(part, outfile)
|
||||
size = dest_path.stat().st_size
|
||||
if size != session.total_size:
|
||||
raise ChunkedUploadError(
|
||||
f"Merged size ({size}) does not match the declared size "
|
||||
f"({session.total_size})"
|
||||
)
|
||||
except BaseException:
|
||||
dest_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
await self.cleanup_session(upload_id)
|
||||
return size
|
||||
|
||||
async def abort(self, upload_id: str, *, owner: str | None = None) -> bool:
|
||||
"""Abort and clean up a session. Unknown sessions abort silently.
|
||||
|
||||
Works on expired sessions too — cleanup is the point — but never on
|
||||
sessions owned by someone else.
|
||||
|
||||
Returns:
|
||||
True when a session existed and was removed.
|
||||
|
||||
Raises:
|
||||
ChunkedUploadError: The session exists but belongs to another owner.
|
||||
"""
|
||||
session = self.sessions.get(upload_id)
|
||||
if session is None:
|
||||
return False
|
||||
if owner is not None and session.owner != owner:
|
||||
raise ChunkedUploadError("Upload session not found or expired")
|
||||
await self.cleanup_session(upload_id)
|
||||
return True
|
||||
|
||||
async def cleanup_session(self, upload_id: str) -> None:
|
||||
session = self.sessions.get(upload_id)
|
||||
if session is None:
|
||||
return
|
||||
if session.chunk_dir.exists():
|
||||
try:
|
||||
shutil.rmtree(session.chunk_dir)
|
||||
except Exception as exc:
|
||||
logger.warning(f"Failed to remove chunk dir {session.chunk_dir}: {exc}")
|
||||
# Keep the session registered so the janitor can retry the
|
||||
# leftover directory instead of forgetting it permanently.
|
||||
return
|
||||
self.sessions.pop(upload_id, None)
|
||||
|
||||
def ensure_cleanup_task_started(self) -> None:
|
||||
if self._cleanup_task is None or self._cleanup_task.done():
|
||||
try:
|
||||
self._cleanup_task = asyncio.create_task(
|
||||
self._cleanup_expired_sessions()
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
async def _cleanup_expired_sessions(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(CLEANUP_INTERVAL_SECONDS)
|
||||
now = time.time()
|
||||
expired = [
|
||||
upload_id
|
||||
for upload_id, session in self.sessions.items()
|
||||
if now - session.last_activity > self.expire_seconds
|
||||
]
|
||||
for upload_id in expired:
|
||||
await self.cleanup_session(upload_id)
|
||||
logger.info(f"Cleaned up expired upload session: {upload_id}")
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to clean up expired upload sessions: {exc}")
|
||||
@@ -36,6 +36,7 @@ from astrbot.core.utils.totp import (
|
||||
set_pending_totp_secret,
|
||||
verify_configured_2fa_code,
|
||||
)
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
from astrbot.core.utils.webhook_utils import ensure_platform_webhook_config
|
||||
from astrbot.dashboard.async_utils import run_maybe_async
|
||||
from astrbot.dashboard.responses import ApiError
|
||||
@@ -1238,9 +1239,9 @@ class ConfigFileService:
|
||||
continue
|
||||
|
||||
save_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
await file.save(str(save_path))
|
||||
if save_path.is_file() and save_path.stat().st_size > MAX_FILE_BYTES:
|
||||
save_path.unlink()
|
||||
try:
|
||||
await file.save(str(save_path), max_bytes=MAX_FILE_BYTES)
|
||||
except UploadTooLargeError:
|
||||
errors.append(f"File too large: {filename}")
|
||||
continue
|
||||
uploaded.append(rel_path)
|
||||
|
||||
@@ -103,23 +103,6 @@ class SkillsService:
|
||||
"You are not permitted to do this operation in demo mode"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _save_upload(file: Any, target_path: str) -> None:
|
||||
if hasattr(file, "save"):
|
||||
maybe_awaitable = file.save(target_path)
|
||||
if hasattr(maybe_awaitable, "__await__"):
|
||||
await maybe_awaitable
|
||||
return
|
||||
|
||||
if hasattr(file, "read"):
|
||||
data = file.read()
|
||||
if hasattr(data, "__await__"):
|
||||
data = await data
|
||||
Path(target_path).write_bytes(data)
|
||||
return
|
||||
|
||||
raise SkillsServiceError("Invalid upload file")
|
||||
|
||||
def resolve_local_skill_dir(self, name: str) -> Path:
|
||||
skill_name = str(name or "").strip()
|
||||
if not skill_name:
|
||||
@@ -308,7 +291,7 @@ class SkillsService:
|
||||
temp_path = _next_available_temp_path(temp_dir, filename)
|
||||
|
||||
try:
|
||||
await self._save_upload(file, temp_path)
|
||||
await file.save(temp_path)
|
||||
try:
|
||||
skill_name = skill_mgr.install_skill_from_zip(
|
||||
temp_path,
|
||||
@@ -365,7 +348,7 @@ class SkillsService:
|
||||
continue
|
||||
|
||||
temp_path = _next_available_temp_path(temp_dir, filename)
|
||||
await self._save_upload(file, temp_path)
|
||||
await file.save(temp_path)
|
||||
|
||||
try:
|
||||
skill_name = skill_mgr.install_skill_from_zip(
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -67,6 +67,12 @@ export type BotRegistrationRequest = {
|
||||
|
||||
export type action = 'start' | 'poll';
|
||||
|
||||
export type ChatChunkUploadRequest = {
|
||||
upload_id: string;
|
||||
chunk_index: number;
|
||||
chunk: (Blob | File);
|
||||
};
|
||||
|
||||
/**
|
||||
* Per-request ChatUI feature flags. A value here takes priority over its legacy top-level field, followed by the documented default.
|
||||
*/
|
||||
@@ -190,6 +196,16 @@ export type ChatThreadMessageRequest = {
|
||||
flags?: ChatFlags;
|
||||
};
|
||||
|
||||
export type ChatUploadInitRequest = {
|
||||
filename: string;
|
||||
total_size: number;
|
||||
content_type?: string;
|
||||
};
|
||||
|
||||
export type ChatUploadSessionRequest = {
|
||||
upload_id: string;
|
||||
};
|
||||
|
||||
export type CommandPatchRequest = {
|
||||
enabled?: boolean;
|
||||
alias?: string;
|
||||
@@ -1666,6 +1682,46 @@ export type UploadFileResponse = (SuccessEnvelope);
|
||||
|
||||
export type UploadFileError = unknown;
|
||||
|
||||
export type InitFileUploadData = {
|
||||
body: ChatUploadInitRequest;
|
||||
};
|
||||
|
||||
export type InitFileUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type InitFileUploadError = unknown;
|
||||
|
||||
export type UploadFileChunkData = {
|
||||
body: ChatChunkUploadRequest;
|
||||
};
|
||||
|
||||
export type UploadFileChunkResponse = (SuccessEnvelope);
|
||||
|
||||
export type UploadFileChunkError = unknown;
|
||||
|
||||
export type CompleteFileUploadData = {
|
||||
body: ChatUploadSessionRequest;
|
||||
};
|
||||
|
||||
export type CompleteFileUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type CompleteFileUploadError = unknown;
|
||||
|
||||
export type AbortFileUploadData = {
|
||||
body: ChatUploadSessionRequest;
|
||||
};
|
||||
|
||||
export type AbortFileUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type AbortFileUploadError = unknown;
|
||||
|
||||
export type StatusFileUploadData = {
|
||||
body: ChatUploadSessionRequest;
|
||||
};
|
||||
|
||||
export type StatusFileUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type StatusFileUploadError = unknown;
|
||||
|
||||
export type UploadOpenApiFileData = {
|
||||
body: FileUploadRequest;
|
||||
};
|
||||
@@ -3448,6 +3504,14 @@ export type AbortBackupUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type AbortBackupUploadError = unknown;
|
||||
|
||||
export type StatusBackupUploadData = {
|
||||
body: BackupUploadSessionRequest;
|
||||
};
|
||||
|
||||
export type StatusBackupUploadResponse = (SuccessEnvelope);
|
||||
|
||||
export type StatusBackupUploadError = unknown;
|
||||
|
||||
export type GetBackupProgressData = {
|
||||
path: {
|
||||
task_id: string;
|
||||
|
||||
@@ -11,6 +11,7 @@ import {
|
||||
type BackupUploadSessionRequest,
|
||||
type BotConfigRequest,
|
||||
type BotRegistrationRequest,
|
||||
type ChatChunkUploadRequest,
|
||||
type ChatMessagePatchRequest,
|
||||
type ChatMessageRegenerateRequest,
|
||||
type ChatProjectRequest,
|
||||
@@ -19,6 +20,8 @@ import {
|
||||
type ChatSessionPatchRequest,
|
||||
type ChatThreadCreateRequest,
|
||||
type ChatThreadMessageRequest,
|
||||
type ChatUploadInitRequest,
|
||||
type ChatUploadSessionRequest,
|
||||
type CommandPatchRequest,
|
||||
type ConfigRouteUpsertRequest,
|
||||
type ConfigRoutesReplaceRequest,
|
||||
@@ -771,6 +774,9 @@ export const backupApi = {
|
||||
abortUpload(payload: BackupUploadSessionRequest) {
|
||||
return typed<OpenConfig>(openApiV1.abortBackupUpload({ body: payload }));
|
||||
},
|
||||
statusUpload(payload: BackupUploadSessionRequest) {
|
||||
return typed<any>(openApiV1.statusBackupUpload({ body: payload }));
|
||||
},
|
||||
check(filename: string) {
|
||||
return typed<any>(
|
||||
openApiV1.checkBackup({ path: { filename } }),
|
||||
@@ -981,6 +987,23 @@ export const fileApi = {
|
||||
openApiV1.uploadFile({ body: generatedFormData(formData) }),
|
||||
);
|
||||
},
|
||||
initUpload(payload: ChatUploadInitRequest) {
|
||||
return typed<any>(openApiV1.initFileUpload({ body: payload }));
|
||||
},
|
||||
uploadChunk(formData: FormData | ChatChunkUploadRequest) {
|
||||
return typed<any>(
|
||||
openApiV1.uploadFileChunk({ body: generatedFormData(formData) }),
|
||||
);
|
||||
},
|
||||
completeUpload(payload: ChatUploadSessionRequest) {
|
||||
return typed<any>(openApiV1.completeFileUpload({ body: payload }));
|
||||
},
|
||||
abortUpload(payload: ChatUploadSessionRequest) {
|
||||
return typed<any>(openApiV1.abortFileUpload({ body: payload }));
|
||||
},
|
||||
statusUpload(payload: ChatUploadSessionRequest) {
|
||||
return typed<any>(openApiV1.statusFileUpload({ body: payload }));
|
||||
},
|
||||
getByName(filename: string) {
|
||||
return openApiV1.getFileByName({
|
||||
query: { filename },
|
||||
|
||||
@@ -248,12 +248,17 @@
|
||||
:reply-to="chatInputReplyTarget"
|
||||
:send-shortcut="sendShortcut"
|
||||
:show-provider-selector="false"
|
||||
:failed-uploads="failedUploadViews"
|
||||
:active-uploads="activeUploadViews"
|
||||
:placeholder="tm('input.projectPlaceholder')"
|
||||
@send="sendCurrentMessage"
|
||||
@stop="stopCurrentSession"
|
||||
@remove-image="removeImage"
|
||||
@remove-audio="removeAudio"
|
||||
@remove-file="removeFile"
|
||||
@retry-failed-upload="retryFailedUpload"
|
||||
@discard-failed-upload="discardFailedUpload"
|
||||
@cancel-active-upload="cancelActiveUpload"
|
||||
@start-recording="startRecording"
|
||||
@stop-recording="stopRecording"
|
||||
@paste-image="handlePaste"
|
||||
@@ -388,6 +393,8 @@
|
||||
:reply-to="chatInputReplyTarget"
|
||||
:send-shortcut="sendShortcut"
|
||||
:show-provider-selector="false"
|
||||
:failed-uploads="failedUploadViews"
|
||||
:active-uploads="activeUploadViews"
|
||||
:placeholder="
|
||||
activeProject ? tm('input.projectPlaceholder') : undefined
|
||||
"
|
||||
@@ -396,6 +403,9 @@
|
||||
@remove-image="removeImage"
|
||||
@remove-audio="removeAudio"
|
||||
@remove-file="removeFile"
|
||||
@retry-failed-upload="retryFailedUpload"
|
||||
@discard-failed-upload="discardFailedUpload"
|
||||
@cancel-active-upload="cancelActiveUpload"
|
||||
@start-recording="startRecording"
|
||||
@stop-recording="stopRecording"
|
||||
@paste-image="handlePaste"
|
||||
@@ -592,12 +602,17 @@ const {
|
||||
stagedImagesUrl,
|
||||
stagedAudioUrl,
|
||||
stagedNonImageFiles,
|
||||
failedUploadViews,
|
||||
activeUploadViews,
|
||||
processAndUploadImage,
|
||||
processAndUploadFile,
|
||||
handlePaste,
|
||||
removeImage,
|
||||
removeAudio,
|
||||
removeFile,
|
||||
retryFailedUpload,
|
||||
discardFailedUpload,
|
||||
cancelActiveUpload,
|
||||
clearStaged,
|
||||
cleanupMediaCache,
|
||||
} = useMediaHandling();
|
||||
|
||||
@@ -82,6 +82,69 @@
|
||||
variant="tonal"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div
|
||||
v-for="(active, index) in activeUploads"
|
||||
:key="'active-' + index"
|
||||
class="attachment-card file-preview attachment-card--active"
|
||||
>
|
||||
<div
|
||||
class="attachment-icon"
|
||||
:style="{ '--attachment-color': activePresentation(active).color }"
|
||||
>
|
||||
<v-icon :icon="activePresentation(active).icon" size="24"></v-icon>
|
||||
<span class="attachment-ext">{{
|
||||
activePresentation(active).label
|
||||
}}</span>
|
||||
</div>
|
||||
<span class="attachment-name">{{ active.name }}</span>
|
||||
<span class="attachment-progress-text">{{ active.percent }}%</span>
|
||||
<v-btn
|
||||
@click="$emit('cancelActiveUpload', index)"
|
||||
class="cancel-active-btn"
|
||||
icon="mdi-close"
|
||||
size="x-small"
|
||||
color="grey-darken-1"
|
||||
variant="text"
|
||||
/>
|
||||
<div class="attachment-progress-track">
|
||||
<v-progress-linear
|
||||
:model-value="active.percent"
|
||||
color="primary"
|
||||
height="3"
|
||||
rounded
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div
|
||||
v-for="(failed, index) in failedUploads"
|
||||
:key="'failed-' + index"
|
||||
class="attachment-card file-preview attachment-card--failed"
|
||||
>
|
||||
<div class="attachment-icon attachment-icon--failed">
|
||||
<v-icon icon="mdi-alert-circle-outline" size="24"></v-icon>
|
||||
</div>
|
||||
<span class="attachment-name" :title="failed.error">{{
|
||||
failed.name
|
||||
}}</span>
|
||||
<v-btn
|
||||
@click="$emit('retryFailedUpload', index)"
|
||||
class="remove-attachment-btn retry-attachment-btn"
|
||||
icon="mdi-refresh"
|
||||
size="x-small"
|
||||
color="primary"
|
||||
variant="tonal"
|
||||
/>
|
||||
<v-btn
|
||||
@click="$emit('discardFailedUpload', index)"
|
||||
class="remove-attachment-btn"
|
||||
icon="mdi-close"
|
||||
size="x-small"
|
||||
color="error"
|
||||
variant="tonal"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</transition>
|
||||
|
||||
@@ -277,6 +340,7 @@ import StyledMenu from "@/components/shared/StyledMenu.vue";
|
||||
import CommandSuggestion from "./CommandSuggestion.vue";
|
||||
import { attachmentPresentation } from "./attachmentPresentation";
|
||||
import type { Session } from "@/composables/useSessions";
|
||||
import type { FailedUploadView, ActiveUploadView } from "@/composables/useMediaHandling";
|
||||
import type { SuggestionCommand } from "./CommandSuggestion.vue";
|
||||
|
||||
interface StagedFileInfo {
|
||||
@@ -304,6 +368,8 @@ interface Props {
|
||||
stagedImagesUrl: string[];
|
||||
stagedAudioUrl: string;
|
||||
stagedFiles?: StagedFileInfo[];
|
||||
failedUploads?: FailedUploadView[];
|
||||
activeUploads?: ActiveUploadView[];
|
||||
disabled: boolean;
|
||||
showSettings?: boolean;
|
||||
isRecording: boolean;
|
||||
@@ -323,6 +389,8 @@ const props = withDefaults(defineProps<Props>(), {
|
||||
currentSession: null,
|
||||
configId: null,
|
||||
stagedFiles: () => [],
|
||||
failedUploads: () => [],
|
||||
activeUploads: () => [],
|
||||
replyTo: null,
|
||||
sendShortcut: "shift_enter",
|
||||
showProviderSelector: true,
|
||||
@@ -338,6 +406,9 @@ const emit = defineEmits<{
|
||||
removeImage: [index: number];
|
||||
removeAudio: [];
|
||||
removeFile: [index: number];
|
||||
retryFailedUpload: [index: number];
|
||||
discardFailedUpload: [index: number];
|
||||
cancelActiveUpload: [index: number];
|
||||
startRecording: [];
|
||||
stopRecording: [];
|
||||
pasteImage: [event: ClipboardEvent];
|
||||
@@ -515,7 +586,9 @@ const hasStagedAttachments = computed(() => {
|
||||
return (
|
||||
props.stagedImagesUrl.length > 0 ||
|
||||
props.stagedAudioUrl ||
|
||||
(props.stagedFiles && props.stagedFiles.length > 0)
|
||||
(props.stagedFiles && props.stagedFiles.length > 0) ||
|
||||
(props.failedUploads && props.failedUploads.length > 0) ||
|
||||
(props.activeUploads && props.activeUploads.length > 0)
|
||||
);
|
||||
});
|
||||
|
||||
@@ -523,6 +596,12 @@ function filePresentation(file: StagedFileInfo) {
|
||||
return attachmentPresentation(file);
|
||||
}
|
||||
|
||||
// In-flight uploads are typed by their local filename already; no need to
|
||||
// wait for the server-side classification.
|
||||
function activePresentation(active: ActiveUploadView) {
|
||||
return attachmentPresentation({ original_name: active.name });
|
||||
}
|
||||
|
||||
// Ctrl+B 长按录音相关
|
||||
const ctrlKeyDown = ref(false);
|
||||
const ctrlKeyTimer = ref<number | null>(null);
|
||||
@@ -1258,6 +1337,50 @@ defineExpose({
|
||||
opacity: 1;
|
||||
}
|
||||
|
||||
.attachment-card--failed {
|
||||
--attachment-color: rgb(var(--v-theme-error));
|
||||
padding-right: 56px;
|
||||
}
|
||||
|
||||
.attachment-icon--failed {
|
||||
color: rgb(var(--v-theme-error));
|
||||
}
|
||||
|
||||
.retry-attachment-btn {
|
||||
right: 30px;
|
||||
}
|
||||
|
||||
.attachment-card--active {
|
||||
padding-right: 10px;
|
||||
}
|
||||
|
||||
.attachment-progress-text {
|
||||
flex-shrink: 0;
|
||||
font-size: 12px;
|
||||
font-weight: 600;
|
||||
color: rgb(var(--v-theme-primary));
|
||||
}
|
||||
|
||||
.cancel-active-btn {
|
||||
flex-shrink: 0;
|
||||
width: 22px !important;
|
||||
height: 22px !important;
|
||||
min-width: 22px !important;
|
||||
opacity: 0.7;
|
||||
}
|
||||
|
||||
.cancel-active-btn:hover {
|
||||
opacity: 1;
|
||||
}
|
||||
|
||||
.attachment-progress-track {
|
||||
position: absolute;
|
||||
/* Align with the filename text: card padding + icon width + gap. */
|
||||
left: 52px;
|
||||
right: 10px;
|
||||
bottom: 3px;
|
||||
}
|
||||
|
||||
.fade-in {
|
||||
animation: fadeIn 0.3s ease-in-out;
|
||||
}
|
||||
|
||||
@@ -234,7 +234,17 @@
|
||||
<v-alert type="error" variant="tonal" class="mb-4">
|
||||
{{ importError }}
|
||||
</v-alert>
|
||||
<v-btn color="primary" variant="tonal" @click="resetImport">
|
||||
<v-btn
|
||||
v-if="canResume"
|
||||
color="primary"
|
||||
variant="tonal"
|
||||
class="mr-2"
|
||||
@click="resumeAndCheck"
|
||||
>
|
||||
<v-icon class="mr-2">mdi-play</v-icon>
|
||||
{{ t('features.settings.backup.import.resumeUpload') }}
|
||||
</v-btn>
|
||||
<v-btn color="grey-darken-1" variant="text" @click="resetImport">
|
||||
{{ t('features.settings.backup.import.retry') }}
|
||||
</v-btn>
|
||||
</div>
|
||||
@@ -370,6 +380,7 @@
|
||||
<script setup>
|
||||
import { ref, computed, watch } from 'vue'
|
||||
import { backupApi } from '@/api/v1'
|
||||
import { useChunkedUpload } from '@/composables/useChunkedUpload'
|
||||
import { useI18n } from '@/i18n/composables'
|
||||
import { askForConfirmation, useConfirmDialog } from '@/utils/confirmDialog'
|
||||
import { restartAstrBot as restartAstrBotRuntime } from '@/utils/restartAstrBot'
|
||||
@@ -399,16 +410,20 @@ const importError = ref('')
|
||||
const uploadedFilename = ref('') // 已上传的文件名
|
||||
const checkResult = ref(null) // 预检查结果
|
||||
|
||||
// 分片上传状态
|
||||
const CONCURRENT_UPLOADS = 5 // 并发上传数
|
||||
const uploadId = ref('')
|
||||
const chunkSize = ref(0) // 分片大小（从后端获取）
|
||||
const uploadProgress = ref({
|
||||
uploaded: 0,
|
||||
total: 0,
|
||||
percent: 0,
|
||||
message: ''
|
||||
})
|
||||
// 分片上传状态（调度由 useChunkedUpload 管理）
|
||||
const uploader = useChunkedUpload(backupApi)
|
||||
const { canResume } = uploader
|
||||
const uploadMessageOverride = ref('') // 预检查阶段覆盖上传进度文案
|
||||
const uploadProgress = computed(() => ({
|
||||
uploaded: uploader.uploadedBytes.value,
|
||||
total: uploader.totalBytes.value,
|
||||
percent: uploader.percent.value,
|
||||
message: uploadMessageOverride.value || {
|
||||
init: t('features.settings.backup.import.uploadInit'),
|
||||
chunks: t('features.settings.backup.import.uploadingChunks'),
|
||||
complete: t('features.settings.backup.import.uploadComplete')
|
||||
}[uploader.phase.value]
|
||||
}))
|
||||
|
||||
// 备份列表
|
||||
const loadingList = ref(false)
|
||||
@@ -550,124 +565,50 @@ const resetExport = () => {
|
||||
exportError.value = ''
|
||||
}
|
||||
|
||||
/**
|
||||
* 并发上传分片
|
||||
*
|
||||
* 使用并发控制同时上传多个分片，提升上传速度。
|
||||
* 后端按分片索引命名文件（如 0.part, 1.part），合并时按顺序读取，
|
||||
* 因此分片到达顺序不影响最终结果。
|
||||
*/
|
||||
const uploadChunksInParallel = async (file, totalChunks, currentUploadId, currentChunkSize) => {
|
||||
// 跟踪已完成的字节数（使用原子操作避免并发问题）
|
||||
let completedBytes = 0
|
||||
const chunkSizes = []
|
||||
|
||||
// 预计算每个分片的大小（使用后端返回的 chunk_size）
|
||||
for (let i = 0; i < totalChunks; i++) {
|
||||
const start = i * currentChunkSize
|
||||
const end = Math.min(start + currentChunkSize, file.size)
|
||||
chunkSizes[i] = end - start
|
||||
}
|
||||
|
||||
// 上传单个分片的函数
|
||||
const uploadSingleChunk = async (chunkIndex) => {
|
||||
const start = chunkIndex * currentChunkSize
|
||||
const end = Math.min(start + currentChunkSize, file.size)
|
||||
const chunk = file.slice(start, end)
|
||||
|
||||
const response = await backupApi.uploadChunk({
|
||||
upload_id: currentUploadId,
|
||||
chunk_index: chunkIndex,
|
||||
chunk
|
||||
})
|
||||
|
||||
if (response.data.status !== 'ok') {
|
||||
throw new Error(response.data.message)
|
||||
}
|
||||
|
||||
// 更新进度（累加已完成字节）
|
||||
completedBytes += chunkSizes[chunkIndex]
|
||||
uploadProgress.value.uploaded = completedBytes
|
||||
uploadProgress.value.percent = Math.round((completedBytes / file.size) * 100)
|
||||
|
||||
return response
|
||||
}
|
||||
|
||||
// 创建分片索引队列
|
||||
const pendingChunks = Array.from({ length: totalChunks }, (_, i) => i)
|
||||
const activePromises = []
|
||||
|
||||
// 处理队列中的分片
|
||||
while (pendingChunks.length > 0 || activePromises.length > 0) {
|
||||
// 填充并发槽位
|
||||
while (pendingChunks.length > 0 && activePromises.length < CONCURRENT_UPLOADS) {
|
||||
const chunkIndex = pendingChunks.shift()
|
||||
const promise = uploadSingleChunk(chunkIndex).then(() => {
|
||||
// 完成后从活动列表移除
|
||||
const idx = activePromises.indexOf(promise)
|
||||
if (idx > -1) activePromises.splice(idx, 1)
|
||||
})
|
||||
activePromises.push(promise)
|
||||
}
|
||||
|
||||
// 等待至少一个完成
|
||||
if (activePromises.length > 0) {
|
||||
await Promise.race(activePromises)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 上传并检查
|
||||
const uploadAndCheck = async () => {
|
||||
if (!importFile.value) return
|
||||
|
||||
importStatus.value = 'uploading'
|
||||
const file = importFile.value
|
||||
uploadMessageOverride.value = ''
|
||||
|
||||
const result = await uploader.start(importFile.value)
|
||||
if (!result) {
|
||||
if (uploader.status.value === 'error') {
|
||||
importStatus.value = 'failed'
|
||||
importError.value = uploader.errorMessage.value
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
uploadedFilename.value = result.filename
|
||||
await checkUploadedBackup()
|
||||
}
|
||||
|
||||
// 断点续传：补传缺失分片后接着预检查
|
||||
const resumeAndCheck = async () => {
|
||||
importStatus.value = 'uploading'
|
||||
importError.value = ''
|
||||
uploadMessageOverride.value = ''
|
||||
|
||||
const result = await uploader.resume()
|
||||
if (!result) {
|
||||
if (uploader.status.value === 'error') {
|
||||
importStatus.value = 'failed'
|
||||
importError.value = uploader.errorMessage.value
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
uploadedFilename.value = result.filename
|
||||
await checkUploadedBackup()
|
||||
}
|
||||
|
||||
// 上传完成后的预检查
|
||||
const checkUploadedBackup = async () => {
|
||||
uploadMessageOverride.value = t('features.settings.backup.import.checking')
|
||||
|
||||
try {
|
||||
// 初始化上传进度
|
||||
uploadProgress.value = {
|
||||
uploaded: 0,
|
||||
total: file.size,
|
||||
percent: 0,
|
||||
message: t('features.settings.backup.import.uploadInit')
|
||||
}
|
||||
|
||||
// 步骤1: 初始化分片上传（后端计算并返回 chunk_size 和 total_chunks）
|
||||
const initResponse = await backupApi.initUpload({
|
||||
filename: file.name,
|
||||
total_size: file.size
|
||||
})
|
||||
|
||||
if (initResponse.data.status !== 'ok') {
|
||||
throw new Error(initResponse.data.message)
|
||||
}
|
||||
|
||||
uploadId.value = initResponse.data.data.upload_id
|
||||
chunkSize.value = initResponse.data.data.chunk_size
|
||||
const totalChunks = initResponse.data.data.total_chunks
|
||||
|
||||
// 步骤2: 并行分片上传（5个并发连接）
|
||||
uploadProgress.value.message = t('features.settings.backup.import.uploadingChunks')
|
||||
|
||||
await uploadChunksInParallel(file, totalChunks, uploadId.value, chunkSize.value)
|
||||
|
||||
// 步骤3: 完成上传
|
||||
uploadProgress.value.message = t('features.settings.backup.import.uploadComplete')
|
||||
|
||||
const completeResponse = await backupApi.completeUpload({
|
||||
upload_id: uploadId.value
|
||||
})
|
||||
|
||||
if (completeResponse.data.status !== 'ok') {
|
||||
throw new Error(completeResponse.data.message)
|
||||
}
|
||||
|
||||
uploadedFilename.value = completeResponse.data.data.filename
|
||||
|
||||
// 步骤4: 预检查
|
||||
uploadProgress.value.message = t('features.settings.backup.import.checking')
|
||||
|
||||
const checkResponse = await backupApi.check(uploadedFilename.value)
|
||||
|
||||
if (checkResponse.data.status !== 'ok') {
|
||||
@@ -675,7 +616,7 @@ const uploadAndCheck = async () => {
|
||||
}
|
||||
|
||||
checkResult.value = checkResponse.data.data
|
||||
|
||||
|
||||
// 检查是否有效
|
||||
if (!checkResult.value.valid) {
|
||||
importStatus.value = 'failed'
|
||||
@@ -687,17 +628,6 @@ const uploadAndCheck = async () => {
|
||||
importStatus.value = 'confirm'
|
||||
|
||||
} catch (error) {
|
||||
// 上传失败时尝试清理已上传的分片
|
||||
if (uploadId.value) {
|
||||
try {
|
||||
await backupApi.abortUpload({
|
||||
upload_id: uploadId.value
|
||||
})
|
||||
} catch (abortError) {
|
||||
console.error('Failed to abort upload:', abortError)
|
||||
}
|
||||
}
|
||||
|
||||
importStatus.value = 'failed'
|
||||
importError.value = error.response?.data?.message || error.message || 'Upload failed'
|
||||
}
|
||||
@@ -759,17 +689,9 @@ const pollImportProgress = async () => {
|
||||
|
||||
// 重置导入状态
|
||||
const resetImport = async () => {
|
||||
// 如果有进行中的上传，先取消
|
||||
if (uploadId.value && importStatus.value === 'uploading') {
|
||||
try {
|
||||
await backupApi.abortUpload({
|
||||
upload_id: uploadId.value
|
||||
})
|
||||
} catch (error) {
|
||||
console.error('Failed to abort upload:', error)
|
||||
}
|
||||
}
|
||||
|
||||
// 取消并清理可能存在的上传会话（包括失败后可续传的会话）
|
||||
await uploader.cancel()
|
||||
|
||||
importStatus.value = 'idle'
|
||||
importFile.value = null
|
||||
importTaskId.value = null
|
||||
@@ -777,9 +699,7 @@ const resetImport = async () => {
|
||||
importError.value = ''
|
||||
uploadedFilename.value = ''
|
||||
checkResult.value = null
|
||||
uploadId.value = ''
|
||||
chunkSize.value = 0
|
||||
uploadProgress.value = { uploaded: 0, total: 0, percent: 0, message: '' }
|
||||
uploadMessageOverride.value = ''
|
||||
}
|
||||
|
||||
// 下载备份（使用浏览器原生下载，可显示下载进度）
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
import { ref, computed } from 'vue';
|
||||
|
||||
/**
|
||||
* Business-agnostic chunked upload scheduler.
|
||||
*
|
||||
* Owns slicing, bounded concurrency, silent per-chunk retries and
|
||||
* session resume. Pages wire in the five endpoint calls of their
|
||||
* business and render from the exposed state only.
|
||||
*/
|
||||
|
||||
export interface ChunkedUploadApi {
|
||||
initUpload(payload: { filename: string; total_size: number }): Promise<any>;
|
||||
uploadChunk(payload: { upload_id: string; chunk_index: number; chunk: Blob }): Promise<any>;
|
||||
completeUpload(payload: { upload_id: string }): Promise<any>;
|
||||
abortUpload(payload: { upload_id: string }): Promise<any>;
|
||||
statusUpload(payload: { upload_id: string }): Promise<any>;
|
||||
}
|
||||
|
||||
export type ChunkedUploadStatus = 'idle' | 'uploading' | 'error' | 'done';
|
||||
export type ChunkedUploadPhase = 'init' | 'chunks' | 'complete';
|
||||
|
||||
const CONCURRENT_UPLOADS = 5;
|
||||
// Silent retries per chunk before the whole run is considered failed.
|
||||
const CHUNK_MAX_ATTEMPTS = 3;
|
||||
|
||||
export function useChunkedUpload(api: ChunkedUploadApi) {
|
||||
const status = ref<ChunkedUploadStatus>('idle');
|
||||
const phase = ref<ChunkedUploadPhase>('init');
|
||||
const uploadedBytes = ref(0);
|
||||
const totalBytes = ref(0);
|
||||
const errorMessage = ref('');
|
||||
|
||||
const percent = computed(() =>
|
||||
totalBytes.value > 0 ? Math.round((uploadedBytes.value / totalBytes.value) * 100) : 0,
|
||||
);
|
||||
// Resume is only meaningful after a failed run while the file is known.
|
||||
const canResume = computed(() => status.value === 'error');
|
||||
|
||||
// Session state kept across a failure so resume() can continue it.
|
||||
let file: File | null = null;
|
||||
let uploadId = '';
|
||||
let chunkSize = 0;
|
||||
let totalChunks = 0;
|
||||
let chunkSizes: number[] = [];
|
||||
let cancelled = false;
|
||||
// Bumped on every start/resume/cancel; stale completions from a
|
||||
// previous run must not count toward the current run's progress.
|
||||
let runGeneration = 0;
|
||||
|
||||
function envelopeData(response: any): any {
|
||||
if (response.data?.status !== 'ok') {
|
||||
throw new Error(response.data?.message || 'Upload failed');
|
||||
}
|
||||
return response.data.data;
|
||||
}
|
||||
|
||||
function toMessage(err: any): string {
|
||||
return err?.response?.data?.message || err?.message || 'Upload failed';
|
||||
}
|
||||
|
||||
function planChunks() {
|
||||
chunkSizes = [];
|
||||
for (let i = 0; i < totalChunks; i++) {
|
||||
const start = i * chunkSize;
|
||||
chunkSizes[i] = Math.min(start + chunkSize, file!.size) - start;
|
||||
}
|
||||
}
|
||||
|
||||
// sessionId is pinned per run: retries must keep writing to the
|
||||
// session they started in, even if a newer run re-initialized one.
|
||||
async function uploadOneChunk(chunkIndex: number, sessionId: string, generation: number) {
|
||||
const start = chunkIndex * chunkSize;
|
||||
const chunk = file!.slice(start, start + chunkSize);
|
||||
let lastError: any;
|
||||
for (let attempt = 0; attempt < CHUNK_MAX_ATTEMPTS; attempt++) {
|
||||
if (cancelled) throw new Error('cancelled');
|
||||
try {
|
||||
envelopeData(
|
||||
await api.uploadChunk({
|
||||
upload_id: sessionId,
|
||||
chunk_index: chunkIndex,
|
||||
chunk,
|
||||
}),
|
||||
);
|
||||
if (generation === runGeneration) {
|
||||
uploadedBytes.value += chunkSizes[chunkIndex];
|
||||
}
|
||||
return;
|
||||
} catch (err) {
|
||||
lastError = err;
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
async function runPool(indexes: number[], sessionId: string, generation: number) {
|
||||
const pending = [...indexes];
|
||||
const active: Promise<void>[] = [];
|
||||
let failure: any = null;
|
||||
while (!failure && !cancelled && (pending.length > 0 || active.length > 0)) {
|
||||
while (pending.length > 0 && active.length < CONCURRENT_UPLOADS) {
|
||||
const chunkIndex = pending.shift()!;
|
||||
const promise = uploadOneChunk(chunkIndex, sessionId, generation).finally(() => {
|
||||
const idx = active.indexOf(promise);
|
||||
if (idx > -1) active.splice(idx, 1);
|
||||
});
|
||||
active.push(promise);
|
||||
}
|
||||
if (active.length > 0) {
|
||||
try {
|
||||
await Promise.race(active);
|
||||
} catch (err) {
|
||||
// Stop scheduling; already-issued requests are drained
|
||||
// below so a later resume never races them on a chunk.
|
||||
failure = err;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (active.length > 0) await Promise.allSettled(active);
|
||||
if (cancelled) throw new Error('cancelled');
|
||||
if (failure) throw failure;
|
||||
}
|
||||
|
||||
async function initSession() {
|
||||
phase.value = 'init';
|
||||
const data = envelopeData(
|
||||
await api.initUpload({ filename: file!.name, total_size: file!.size }),
|
||||
);
|
||||
uploadId = data.upload_id;
|
||||
chunkSize = data.chunk_size;
|
||||
totalChunks = data.total_chunks;
|
||||
planChunks();
|
||||
uploadedBytes.value = 0;
|
||||
}
|
||||
|
||||
async function completeSession() {
|
||||
phase.value = 'complete';
|
||||
const result = envelopeData(await api.completeUpload({ upload_id: uploadId }));
|
||||
// The request was in flight when the user cancelled; a late success
|
||||
// must not flip the state back to done or hand a result to the caller.
|
||||
if (cancelled) throw new Error('cancelled');
|
||||
return result;
|
||||
}
|
||||
|
||||
async function start(f: File): Promise<any> {
|
||||
file = f;
|
||||
cancelled = false;
|
||||
status.value = 'uploading';
|
||||
errorMessage.value = '';
|
||||
totalBytes.value = f.size;
|
||||
try {
|
||||
await initSession();
|
||||
phase.value = 'chunks';
|
||||
await runPool(
|
||||
Array.from({ length: totalChunks }, (_, i) => i),
|
||||
uploadId,
|
||||
++runGeneration,
|
||||
);
|
||||
const result = await completeSession();
|
||||
status.value = 'done';
|
||||
return result;
|
||||
} catch (err) {
|
||||
return handleFailure(err);
|
||||
}
|
||||
}
|
||||
|
||||
async function resume(): Promise<any> {
|
||||
if (!file) return undefined;
|
||||
cancelled = false;
|
||||
status.value = 'uploading';
|
||||
errorMessage.value = '';
|
||||
try {
|
||||
let received: number[] = [];
|
||||
if (uploadId) {
|
||||
try {
|
||||
const data = await envelopeData(api.statusUpload({ upload_id: uploadId }));
|
||||
received = data.received_chunks;
|
||||
} catch {
|
||||
// Session expired or unknown: fall back to a fresh one.
|
||||
uploadId = '';
|
||||
}
|
||||
}
|
||||
if (!uploadId) {
|
||||
await initSession();
|
||||
} else {
|
||||
uploadedBytes.value = received.reduce((sum, i) => sum + chunkSizes[i], 0);
|
||||
}
|
||||
phase.value = 'chunks';
|
||||
const missing = Array.from({ length: totalChunks }, (_, i) => i).filter(
|
||||
i => !received.includes(i),
|
||||
);
|
||||
await runPool(missing, uploadId, ++runGeneration);
|
||||
const result = await completeSession();
|
||||
status.value = 'done';
|
||||
return result;
|
||||
} catch (err) {
|
||||
return handleFailure(err);
|
||||
}
|
||||
}
|
||||
|
||||
function handleFailure(err: any): undefined {
|
||||
// Cancel wins over any late resolution: a run cancelled while its
|
||||
// init request was still in flight lands here once init returns,
|
||||
// and its freshly created session must not leak server-side.
|
||||
if (cancelled || err?.message === 'cancelled') {
|
||||
const id = uploadId;
|
||||
reset();
|
||||
if (id) {
|
||||
void api.abortUpload({ upload_id: id }).catch(e =>
|
||||
console.error('Failed to abort upload:', e),
|
||||
);
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
// A failed complete() means the merged result was rejected; the
|
||||
// session cannot be trusted to succeed on retry, so drop it and let
|
||||
// resume() start a fresh one.
|
||||
if (phase.value === 'complete') uploadId = '';
|
||||
status.value = 'error';
|
||||
errorMessage.value = toMessage(err);
|
||||
return undefined;
|
||||
}
|
||||
|
||||
async function cancel() {
|
||||
cancelled = true;
|
||||
runGeneration++;
|
||||
const id = uploadId;
|
||||
reset();
|
||||
if (id) {
|
||||
try {
|
||||
await api.abortUpload({ upload_id: id });
|
||||
} catch (err) {
|
||||
console.error('Failed to abort upload:', err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function reset() {
|
||||
status.value = 'idle';
|
||||
phase.value = 'init';
|
||||
uploadedBytes.value = 0;
|
||||
totalBytes.value = 0;
|
||||
errorMessage.value = '';
|
||||
file = null;
|
||||
uploadId = '';
|
||||
chunkSize = 0;
|
||||
totalChunks = 0;
|
||||
chunkSizes = [];
|
||||
}
|
||||
|
||||
return {
|
||||
status,
|
||||
phase,
|
||||
percent,
|
||||
uploadedBytes,
|
||||
totalBytes,
|
||||
errorMessage,
|
||||
canResume,
|
||||
start,
|
||||
resume,
|
||||
cancel,
|
||||
reset,
|
||||
};
|
||||
}
|
||||
@@ -1,5 +1,10 @@
|
||||
import { ref, computed } from 'vue';
|
||||
import { ref, shallowRef, computed } from 'vue';
|
||||
import { fileApi } from '@/api/v1';
|
||||
import { useChunkedUpload } from '@/composables/useChunkedUpload';
|
||||
|
||||
// Files at or above this size use resumable chunked upload instead of a
|
||||
// single multipart POST.
|
||||
const CHUNKED_UPLOAD_THRESHOLD = 32 * 1024 * 1024;
|
||||
|
||||
export interface StagedFileInfo {
|
||||
attachment_id: string;
|
||||
@@ -10,19 +15,71 @@ export interface StagedFileInfo {
|
||||
signature?: string;
|
||||
}
|
||||
|
||||
export interface FailedUploadView {
|
||||
name: string;
|
||||
size: number;
|
||||
error: string;
|
||||
}
|
||||
|
||||
export interface ActiveUploadView {
|
||||
name: string;
|
||||
size: number;
|
||||
percent: number;
|
||||
}
|
||||
|
||||
interface UploadEntry {
|
||||
file: File;
|
||||
signature: string;
|
||||
uploader: ReturnType<typeof useChunkedUpload>;
|
||||
}
|
||||
|
||||
export function useMediaHandling() {
|
||||
const stagedFiles = ref<StagedFileInfo[]>([]);
|
||||
const mediaCache = ref<Record<string, string>>({});
|
||||
const pendingFileSignatures = new Set<string>();
|
||||
// shallowRef: entries hold uploader instances whose internal refs must
|
||||
// stay intact (a deep ref() would unwrap them), so updates reassign.
|
||||
const failedUploads = shallowRef<UploadEntry[]>([]);
|
||||
const activeUploads = shallowRef<UploadEntry[]>([]);
|
||||
|
||||
// Display-only projection of failed uploads for the input area.
|
||||
const failedUploadViews = computed<FailedUploadView[]>(() =>
|
||||
failedUploads.value.map(entry => ({
|
||||
name: entry.file.name,
|
||||
size: entry.file.size,
|
||||
error: entry.uploader.errorMessage.value
|
||||
}))
|
||||
);
|
||||
|
||||
// Display-only projection of in-flight uploads (progress chips).
|
||||
const activeUploadViews = computed<ActiveUploadView[]>(() =>
|
||||
activeUploads.value.map(entry => ({
|
||||
name: entry.file.name,
|
||||
size: entry.file.size,
|
||||
percent: entry.uploader.percent.value
|
||||
}))
|
||||
);
|
||||
|
||||
async function getFileSignature(file: File): Promise<string> {
|
||||
if (crypto?.subtle) {
|
||||
const buffer = await file.arrayBuffer();
|
||||
const digest = await crypto.subtle.digest('SHA-256', buffer);
|
||||
// Digest per block, then digest the concatenated block digests.
|
||||
// The signature is only used for in-session dedup and never
|
||||
// leaves the browser, so it does not need to be the canonical
|
||||
// SHA-256 of the whole file — only stable per file. This keeps
|
||||
// memory bounded at BLOCK bytes regardless of file size.
|
||||
const BLOCK = 8 * 1024 * 1024;
|
||||
const blockHashes: Uint8Array[] = [];
|
||||
for (let offset = 0; offset < file.size; offset += BLOCK) {
|
||||
const buf = await file.slice(offset, offset + BLOCK).arrayBuffer();
|
||||
blockHashes.push(new Uint8Array(await crypto.subtle.digest('SHA-256', buf)));
|
||||
}
|
||||
const combined = new Uint8Array(blockHashes.length * 32);
|
||||
blockHashes.forEach((h, i) => combined.set(h, i * 32));
|
||||
const digest = await crypto.subtle.digest('SHA-256', combined);
|
||||
const hash = Array.from(new Uint8Array(digest))
|
||||
.map(byte => byte.toString(16).padStart(2, '0'))
|
||||
.join('');
|
||||
return `sha256:${hash}`;
|
||||
return `sha256m:${hash}`;
|
||||
}
|
||||
|
||||
return `meta:${file.name}:${file.size}:${file.type}:${file.lastModified}`;
|
||||
@@ -31,7 +88,8 @@ export function useMediaHandling() {
|
||||
function isDuplicateFile(signature: string) {
|
||||
return (
|
||||
pendingFileSignatures.has(signature) ||
|
||||
stagedFiles.value.some(file => file.signature === signature)
|
||||
stagedFiles.value.some(file => file.signature === signature) ||
|
||||
failedUploads.value.some(entry => entry.signature === signature)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -52,28 +110,32 @@ export function useMediaHandling() {
|
||||
}
|
||||
}
|
||||
|
||||
function stageUploaded(file: File, data: any, signature: string): StagedFileInfo {
|
||||
const stagedFile = {
|
||||
attachment_id: data.attachment_id,
|
||||
filename: data.filename,
|
||||
original_name: file.name,
|
||||
url: URL.createObjectURL(file),
|
||||
type: data.type,
|
||||
signature
|
||||
};
|
||||
stagedFiles.value.push(stagedFile);
|
||||
return stagedFile;
|
||||
}
|
||||
|
||||
async function uploadStagedFile(file: File): Promise<StagedFileInfo | undefined> {
|
||||
const signature = await getFileSignature(file);
|
||||
if (isDuplicateFile(signature)) return undefined;
|
||||
|
||||
pendingFileSignatures.add(signature);
|
||||
const formData = new FormData();
|
||||
formData.append('file', file);
|
||||
|
||||
try {
|
||||
if (file.size >= CHUNKED_UPLOAD_THRESHOLD) {
|
||||
return await uploadChunkedStagedFile(file, signature);
|
||||
}
|
||||
const formData = new FormData();
|
||||
formData.append('file', file);
|
||||
const response = await fileApi.upload(formData);
|
||||
|
||||
const { attachment_id, filename, type } = response.data.data;
|
||||
const stagedFile = {
|
||||
attachment_id,
|
||||
filename,
|
||||
original_name: file.name,
|
||||
url: URL.createObjectURL(file),
|
||||
type,
|
||||
signature
|
||||
};
|
||||
stagedFiles.value.push(stagedFile);
|
||||
return stagedFile;
|
||||
return stageUploaded(file, response.data.data, signature);
|
||||
} catch (err) {
|
||||
console.error('Error uploading file:', err);
|
||||
return undefined;
|
||||
@@ -82,6 +144,64 @@ export function useMediaHandling() {
|
||||
}
|
||||
}
|
||||
|
||||
// Large files go through the resumable chunked upload endpoints; a
|
||||
// failure keeps the session entry so the user can resume from the
|
||||
// last received chunk instead of starting over.
|
||||
async function uploadChunkedStagedFile(file: File, signature: string): Promise<StagedFileInfo | undefined> {
|
||||
const uploader = useChunkedUpload({
|
||||
initUpload: ({ filename, total_size }) =>
|
||||
fileApi.initUpload({ filename, total_size, content_type: file.type }),
|
||||
uploadChunk: fileApi.uploadChunk,
|
||||
completeUpload: fileApi.completeUpload,
|
||||
abortUpload: fileApi.abortUpload,
|
||||
statusUpload: fileApi.statusUpload
|
||||
});
|
||||
const entry: UploadEntry = { file, signature, uploader };
|
||||
activeUploads.value = [...activeUploads.value, entry];
|
||||
const result = await uploader.start(file);
|
||||
activeUploads.value = activeUploads.value.filter(e => e !== entry);
|
||||
if (result) {
|
||||
return stageUploaded(file, result, signature);
|
||||
}
|
||||
if (uploader.status.value === 'error') {
|
||||
failedUploads.value = [...failedUploads.value, entry];
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function cancelActiveUpload(index: number) {
|
||||
const entry = activeUploads.value[index];
|
||||
if (!entry) return;
|
||||
// The start() promise settles as cancelled; the entry is removed by
|
||||
// the settle path in uploadChunkedStagedFile.
|
||||
void entry.uploader.cancel();
|
||||
}
|
||||
|
||||
async function retryFailedUpload(index: number): Promise<StagedFileInfo | undefined> {
|
||||
const entry = failedUploads.value[index];
|
||||
if (!entry) return undefined;
|
||||
// Move to active: blocks a second click from running resume()
|
||||
// concurrently on the same session and shows live progress.
|
||||
failedUploads.value = failedUploads.value.filter(e => e !== entry);
|
||||
activeUploads.value = [...activeUploads.value, entry];
|
||||
const result = await entry.uploader.resume();
|
||||
activeUploads.value = activeUploads.value.filter(e => e !== entry);
|
||||
if (!result) {
|
||||
if (entry.uploader.status.value === 'error') {
|
||||
failedUploads.value = [...failedUploads.value, entry];
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
return stageUploaded(entry.file, result, entry.signature);
|
||||
}
|
||||
|
||||
async function discardFailedUpload(index: number) {
|
||||
const entry = failedUploads.value[index];
|
||||
if (!entry) return;
|
||||
failedUploads.value = failedUploads.value.filter(e => e !== entry);
|
||||
await entry.uploader.cancel();
|
||||
}
|
||||
|
||||
async function processAndUploadImage(file: File) {
|
||||
return uploadStagedFile(file);
|
||||
}
|
||||
@@ -196,6 +316,8 @@ export function useMediaHandling() {
|
||||
stagedAudioUrl,
|
||||
stagedFiles,
|
||||
stagedNonImageFiles,
|
||||
failedUploadViews,
|
||||
activeUploadViews,
|
||||
getMediaFile,
|
||||
processAndUploadImage,
|
||||
processAndUploadFile,
|
||||
@@ -203,6 +325,9 @@ export function useMediaHandling() {
|
||||
removeImage,
|
||||
removeAudio,
|
||||
removeFile,
|
||||
retryFailedUpload,
|
||||
discardFailedUpload,
|
||||
cancelActiveUpload,
|
||||
clearStaged,
|
||||
cleanupMediaCache
|
||||
};
|
||||
|
||||
@@ -209,6 +209,7 @@
|
||||
"restartNow": "Restart Now",
|
||||
"failed": "Import Failed",
|
||||
"retry": "Retry",
|
||||
"resumeUpload": "Resume Upload",
|
||||
"version": {
|
||||
"backupVersion": "Backup Version",
|
||||
"currentVersion": "Current Version",
|
||||
|
||||
@@ -209,6 +209,7 @@
|
||||
"restartNow": "今すぐ再起動",
|
||||
"failed": "インポートに失敗しました",
|
||||
"retry": "再試行",
|
||||
"resumeUpload": "アップロードを再開",
|
||||
"version": {
|
||||
"backupVersion": "バックアップのバージョン",
|
||||
"currentVersion": "現在のバージョン",
|
||||
|
||||
@@ -209,6 +209,7 @@
|
||||
"restartNow": "Перезапустить сейчас",
|
||||
"failed": "Ошибка импорта",
|
||||
"retry": "Повторить",
|
||||
"resumeUpload": "Продолжить загрузку",
|
||||
"version": {
|
||||
"backupVersion": "Версия бэкапа",
|
||||
"currentVersion": "Текущая версия",
|
||||
|
||||
@@ -209,6 +209,7 @@
|
||||
"restartNow": "立即重启",
|
||||
"failed": "导入失败",
|
||||
"retry": "重试",
|
||||
"resumeUpload": "继续上传",
|
||||
"version": {
|
||||
"backupVersion": "备份版本",
|
||||
"currentVersion": "当前版本",
|
||||
|
||||
@@ -28,7 +28,6 @@ export const useCommonStore = defineStore("common", {
|
||||
// 注意：这里如果之前改过 Polyfill 的话，可能需要保持原样
|
||||
// 如果是用 fetch 的话，这里是支持 Authorization Header 的
|
||||
const headers = {
|
||||
"Content-Type": "multipart/form-data",
|
||||
Authorization: "Bearer " + localStorage.getItem("token"),
|
||||
};
|
||||
|
||||
|
||||
@@ -207,6 +207,11 @@ Upload and download chat attachments.
|
||||
| `POST` | `/api/v1/file` | — |
|
||||
| `POST` | `/api/v1/files` | — |
|
||||
| `GET` | `/api/v1/files/content` | — |
|
||||
| `POST` | `/api/v1/files/upload/abort` | — |
|
||||
| `POST` | `/api/v1/files/upload/chunk` | — |
|
||||
| `POST` | `/api/v1/files/upload/complete` | — |
|
||||
| `POST` | `/api/v1/files/upload/init` | — |
|
||||
| `POST` | `/api/v1/files/upload/status` | — |
|
||||
| `GET` | `/api/v1/files/{attachment_id}` | — |
|
||||
| `DELETE` | `/api/v1/files/{attachment_id}` | — |
|
||||
| `GET` | `/api/v1/files/{attachment_id}/content` | — |
|
||||
|
||||
@@ -2580,6 +2580,136 @@
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/files/upload/init": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Files"
|
||||
],
|
||||
"summary": "Initialize chunked file upload",
|
||||
"operationId": "initFileUpload",
|
||||
"x-astrbot-scope": "file",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ChatUploadInitRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"$ref": "#/components/responses/Ok"
|
||||
}
|
||||
},
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/files/upload/chunk": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Files"
|
||||
],
|
||||
"summary": "Upload a file chunk",
|
||||
"operationId": "uploadFileChunk",
|
||||
"x-astrbot-scope": "file",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"multipart/form-data": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ChatChunkUploadRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"$ref": "#/components/responses/Ok"
|
||||
}
|
||||
},
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/files/upload/complete": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Files"
|
||||
],
|
||||
"summary": "Complete chunked file upload",
|
||||
"operationId": "completeFileUpload",
|
||||
"x-astrbot-scope": "file",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ChatUploadSessionRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"$ref": "#/components/responses/Ok"
|
||||
}
|
||||
},
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/files/upload/abort": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Files"
|
||||
],
|
||||
"summary": "Abort chunked file upload",
|
||||
"operationId": "abortFileUpload",
|
||||
"x-astrbot-scope": "file",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ChatUploadSessionRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"$ref": "#/components/responses/Ok"
|
||||
}
|
||||
},
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/files/upload/status": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Files"
|
||||
],
|
||||
"summary": "Query chunked file upload status",
|
||||
"operationId": "statusFileUpload",
|
||||
"x-astrbot-scope": "file",
|
||||
"requestBody": {
|
||||
"required": true,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ChatUploadSessionRequest"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"$ref": "#/components/responses/Ok"
|
||||
}
|
||||
},
|
||||
"description": "**Required scope:** `file`"
|
||||
}
|
||||
},
|
||||
"/api/v1/file": {
|
||||
"post": {
|
||||
"tags": [
|
||||
@@ -8061,6 +8191,59 @@
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"ChatUploadInitRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"filename",
|
||||
"total_size"
|
||||
],
|
||||
"properties": {
|
||||
"filename": {
|
||||
"type": "string"
|
||||
},
|
||||
"total_size": {
|
||||
"type": "integer",
|
||||
"minimum": 1
|
||||
},
|
||||
"content_type": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"ChatUploadSessionRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"upload_id"
|
||||
],
|
||||
"properties": {
|
||||
"upload_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"ChatChunkUploadRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"upload_id",
|
||||
"chunk_index",
|
||||
"chunk"
|
||||
],
|
||||
"properties": {
|
||||
"upload_id": {
|
||||
"type": "string"
|
||||
},
|
||||
"chunk_index": {
|
||||
"type": "integer",
|
||||
"minimum": 0
|
||||
},
|
||||
"chunk": {
|
||||
"type": "string",
|
||||
"format": "binary"
|
||||
}
|
||||
}
|
||||
},
|
||||
"T2iTemplateRequest": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
|
||||
@@ -207,6 +207,11 @@ outline: deep
|
||||
| `POST` | `/api/v1/file` | — |
|
||||
| `POST` | `/api/v1/files` | — |
|
||||
| `GET` | `/api/v1/files/content` | — |
|
||||
| `POST` | `/api/v1/files/upload/abort` | — |
|
||||
| `POST` | `/api/v1/files/upload/chunk` | — |
|
||||
| `POST` | `/api/v1/files/upload/complete` | — |
|
||||
| `POST` | `/api/v1/files/upload/init` | — |
|
||||
| `POST` | `/api/v1/files/upload/status` | — |
|
||||
| `GET` | `/api/v1/files/{attachment_id}` | — |
|
||||
| `DELETE` | `/api/v1/files/{attachment_id}` | — |
|
||||
| `GET` | `/api/v1/files/{attachment_id}/content` | — |
|
||||
|
||||
@@ -1701,6 +1701,86 @@ paths:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/files/upload/init:
|
||||
post:
|
||||
tags: [Files]
|
||||
summary: Initialize chunked file upload
|
||||
operationId: initFileUpload
|
||||
x-astrbot-scope: file
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: "#/components/schemas/ChatUploadInitRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/files/upload/chunk:
|
||||
post:
|
||||
tags: [Files]
|
||||
summary: Upload a file chunk
|
||||
operationId: uploadFileChunk
|
||||
x-astrbot-scope: file
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
multipart/form-data:
|
||||
schema:
|
||||
$ref: "#/components/schemas/ChatChunkUploadRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/files/upload/complete:
|
||||
post:
|
||||
tags: [Files]
|
||||
summary: Complete chunked file upload
|
||||
operationId: completeFileUpload
|
||||
x-astrbot-scope: file
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: "#/components/schemas/ChatUploadSessionRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/files/upload/abort:
|
||||
post:
|
||||
tags: [Files]
|
||||
summary: Abort chunked file upload
|
||||
operationId: abortFileUpload
|
||||
x-astrbot-scope: file
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: "#/components/schemas/ChatUploadSessionRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/files/upload/status:
|
||||
post:
|
||||
tags: [Files]
|
||||
summary: Query chunked file upload status
|
||||
operationId: statusFileUpload
|
||||
x-astrbot-scope: file
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: "#/components/schemas/ChatUploadSessionRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/file:
|
||||
post:
|
||||
tags: [Open API]
|
||||
@@ -4645,6 +4725,22 @@ paths:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/backups/upload/status:
|
||||
post:
|
||||
tags: [Backups]
|
||||
summary: Query chunked backup upload status
|
||||
operationId: statusBackupUpload
|
||||
x-astrbot-scope: system
|
||||
requestBody:
|
||||
required: true
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: "#/components/schemas/BackupUploadSessionRequest"
|
||||
responses:
|
||||
"200":
|
||||
$ref: "#/components/responses/Ok"
|
||||
|
||||
/api/v1/backups/tasks/{task_id}:
|
||||
get:
|
||||
tags: [Backups]
|
||||
@@ -6388,6 +6484,40 @@ components:
|
||||
type: string
|
||||
format: binary
|
||||
|
||||
ChatUploadInitRequest:
|
||||
type: object
|
||||
required: [filename, total_size]
|
||||
properties:
|
||||
filename:
|
||||
type: string
|
||||
total_size:
|
||||
type: integer
|
||||
minimum: 1
|
||||
content_type:
|
||||
type: string
|
||||
additionalProperties: false
|
||||
|
||||
ChatUploadSessionRequest:
|
||||
type: object
|
||||
required: [upload_id]
|
||||
properties:
|
||||
upload_id:
|
||||
type: string
|
||||
additionalProperties: false
|
||||
|
||||
ChatChunkUploadRequest:
|
||||
type: object
|
||||
required: [upload_id, chunk_index, chunk]
|
||||
properties:
|
||||
upload_id:
|
||||
type: string
|
||||
chunk_index:
|
||||
type: integer
|
||||
minimum: 0
|
||||
chunk:
|
||||
type: string
|
||||
format: binary
|
||||
|
||||
BackupRenameRequest:
|
||||
type: object
|
||||
required: [new_name]
|
||||
|
||||
@@ -5,6 +5,7 @@ import os
|
||||
import re
|
||||
import zipfile
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
@@ -27,8 +28,13 @@ from astrbot.core.config.default import VERSION
|
||||
from astrbot.core.db.po import (
|
||||
ConversationV2,
|
||||
)
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
from astrbot.core.utils.version_comparator import VersionComparator
|
||||
from astrbot.dashboard.services.backup_service import (
|
||||
CHUNK_SIZE,
|
||||
MAX_BACKUP_TOTAL_BYTES,
|
||||
BackupService,
|
||||
BackupServiceError,
|
||||
generate_unique_filename,
|
||||
secure_filename,
|
||||
)
|
||||
@@ -1093,3 +1099,362 @@ class TestBackupIntegration:
|
||||
# 读取主数据库
|
||||
main_db = json.loads(zf.read("databases/main_db.json"))
|
||||
assert "platform_stats" in main_db
|
||||
|
||||
|
||||
class _StubUploadFile:
|
||||
"""模拟 adapter 风格 save() 契约的上传文件对象"""
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self._data = data
|
||||
|
||||
async def save(self, destination, *, max_bytes=None) -> int:
|
||||
if max_bytes is not None and len(self._data) > max_bytes:
|
||||
raise UploadTooLargeError(max_bytes)
|
||||
Path(destination).write_bytes(self._data)
|
||||
return len(self._data)
|
||||
|
||||
|
||||
class TestBackupUploadLimits:
|
||||
"""备份上传大小限制测试"""
|
||||
|
||||
@pytest.fixture
|
||||
def backup_service(self, tmp_path):
|
||||
"""创建使用临时目录的 BackupService"""
|
||||
service = BackupService(db=MagicMock(), core_lifecycle=MagicMock())
|
||||
service.backup_dir = str(tmp_path / "backups")
|
||||
service.chunks_dir = str(tmp_path / "backups" / ".chunks")
|
||||
service.chunked_uploads.chunks_root = Path(service.chunks_dir)
|
||||
return service
|
||||
|
||||
def test_upload_init_rejects_oversized_total(self, backup_service):
|
||||
"""声明总大小超过上限时拒绝初始化"""
|
||||
with pytest.raises(BackupServiceError, match="size limit"):
|
||||
backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": MAX_BACKUP_TOTAL_BYTES + 1},
|
||||
owner="tester",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_chunk_rejects_oversized_chunk(self, backup_service):
|
||||
"""超过 CHUNK_SIZE 的分片被拒绝且不产生残留文件"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": CHUNK_SIZE}, owner="tester"
|
||||
)
|
||||
big_chunk = _StubUploadFile(b"x" * (CHUNK_SIZE + 1))
|
||||
|
||||
with pytest.raises(BackupServiceError, match="Chunk exceeds the size limit"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="0",
|
||||
chunk_file=big_chunk,
|
||||
owner="tester",
|
||||
)
|
||||
|
||||
await backup_service.cleanup_upload_session(session["upload_id"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_chunk_accepts_small_chunk(self, backup_service):
|
||||
"""正常大小的分片可以上传"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="tester"
|
||||
)
|
||||
result = await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * 100),
|
||||
owner="tester",
|
||||
)
|
||||
assert result["received"] == 1
|
||||
assert result["total"] == 1
|
||||
|
||||
await backup_service.cleanup_upload_session(session["upload_id"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_complete_rejects_size_mismatch(self, backup_service):
|
||||
"""合并后大小与声明大小不一致时拒绝完成（分片带外损坏的兜底）"""
|
||||
declared = CHUNK_SIZE + 100
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": declared}, owner="tester"
|
||||
)
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * CHUNK_SIZE),
|
||||
owner="tester",
|
||||
)
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="1",
|
||||
chunk_file=_StubUploadFile(b"x" * 100),
|
||||
owner="tester",
|
||||
)
|
||||
# 传片时的字节数校验已拦截短写；此处绕过保存路径直接篡改磁盘上的
|
||||
# 分片（模拟带外损坏），验证合并时的大小复核仍然兜底
|
||||
chunk_path = (
|
||||
backup_service.chunked_uploads.chunks_root / session["upload_id"] / "1.part"
|
||||
)
|
||||
chunk_path.write_bytes(b"x" * 50)
|
||||
|
||||
with pytest.raises(BackupServiceError, match="does not match"):
|
||||
await backup_service.upload_complete(
|
||||
{"upload_id": session["upload_id"]}, owner="tester"
|
||||
)
|
||||
|
||||
await backup_service.cleanup_upload_session(session["upload_id"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_rejects_other_owner(self, backup_service):
|
||||
"""会话绑定创建者，其他用户无法传片、合并或取消"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * 100),
|
||||
owner="mallory",
|
||||
)
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
await backup_service.upload_complete(
|
||||
{"upload_id": session["upload_id"]}, owner="mallory"
|
||||
)
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
await backup_service.upload_abort(
|
||||
{"upload_id": session["upload_id"]}, owner="mallory"
|
||||
)
|
||||
|
||||
# 会话在攻击后仍然存活，真正的主人可以正常使用
|
||||
result = await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
assert result["received"] == 1
|
||||
|
||||
await backup_service.cleanup_upload_session(session["upload_id"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_status_reports_progress(self, backup_service):
|
||||
"""状态查询返回已收分片，支持乱序后的续传定位"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": CHUNK_SIZE * 2}, owner="tester"
|
||||
)
|
||||
# 故意乱序：只传第 1 片（索引 1）
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=session["upload_id"],
|
||||
chunk_index_str="1",
|
||||
chunk_file=_StubUploadFile(b"x" * CHUNK_SIZE),
|
||||
owner="tester",
|
||||
)
|
||||
|
||||
status = backup_service.upload_status(
|
||||
{"upload_id": session["upload_id"]}, owner="tester"
|
||||
)
|
||||
assert status["received_chunks"] == [1]
|
||||
assert status["total_chunks"] == 2
|
||||
assert status["chunk_size"] == CHUNK_SIZE
|
||||
assert 0 < status["expires_in"] <= 3600
|
||||
|
||||
await backup_service.cleanup_upload_session(session["upload_id"])
|
||||
|
||||
def test_upload_status_bound_to_owner(self, backup_service):
|
||||
"""状态查询同样绑定 owner，且不能探测他人会话"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
backup_service.upload_status(
|
||||
{"upload_id": session["upload_id"]}, owner="mallory"
|
||||
)
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
backup_service.upload_status({"upload_id": "no-such-id"}, owner="alice")
|
||||
|
||||
def test_upload_status_does_not_extend_lifetime(self, backup_service):
|
||||
"""查询状态不得刷新 last_activity，否则轮询会让会话永不过期"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
inner = backup_service.chunked_uploads.get_session(
|
||||
session["upload_id"], owner="alice"
|
||||
)
|
||||
inner.last_activity -= 100 # 模拟会话已经闲置了 100 秒
|
||||
|
||||
status = backup_service.upload_status(
|
||||
{"upload_id": session["upload_id"]}, owner="alice"
|
||||
)
|
||||
assert status["expires_in"] <= 3600 - 100 + 1
|
||||
assert (
|
||||
backup_service.chunked_uploads.get_session(
|
||||
session["upload_id"], owner="alice"
|
||||
).last_activity
|
||||
== inner.last_activity
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_chunk_retry_preserves_received_chunk(self, backup_service):
|
||||
"""同索引重传失败不得破坏已收到的分片"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
good = b"x" * 100
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(good),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
# 同索引用超大分片重试：必须报错，且已收到的分片原样保留
|
||||
with pytest.raises(BackupServiceError, match="size limit"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"y" * (CHUNK_SIZE + 1)),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
result = await backup_service.upload_complete(
|
||||
{"upload_id": upload_id}, owner="alice"
|
||||
)
|
||||
assert result["size"] == 100
|
||||
merged = Path(backup_service.backup_dir) / result["filename"]
|
||||
assert merged.read_bytes() == good
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_publish_failure_cleans_temp(self, backup_service, monkeypatch):
|
||||
"""原子改名失败时临时文件必须清理，已收分片不受影响"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": CHUNK_SIZE + 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * CHUNK_SIZE),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
def _boom(*args, **kwargs):
|
||||
raise OSError("simulated rename failure")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"astrbot.dashboard.services.chunked_upload_service.os.replace", _boom
|
||||
)
|
||||
with pytest.raises(OSError, match="simulated rename failure"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="1",
|
||||
chunk_file=_StubUploadFile(b"y" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
chunk_dir = backup_service.chunked_uploads.chunks_root / upload_id
|
||||
assert not list(chunk_dir.glob("*.tmp"))
|
||||
# 改名失败的分片未登记，已收到的第 0 片完好
|
||||
status = backup_service.upload_status({"upload_id": upload_id}, owner="alice")
|
||||
assert status["received_chunks"] == [0]
|
||||
assert (chunk_dir / "0.part").read_bytes() == b"x" * CHUNK_SIZE
|
||||
|
||||
await backup_service.cleanup_upload_session(upload_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_session_is_rejected_and_abortable(self, backup_service):
|
||||
"""过期会话不可用（get_session 强制过期），但 abort 清理仍有效"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
inner = backup_service.chunked_uploads.get_session(upload_id, owner="alice")
|
||||
inner.last_activity -= 7200 # 闲置两小时，已过 1 小时过期线
|
||||
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
backup_service.upload_status({"upload_id": upload_id}, owner="alice")
|
||||
with pytest.raises(BackupServiceError, match="not found or expired"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
# abort 是清理路径，对过期会话仍然有效
|
||||
await backup_service.upload_abort({"upload_id": upload_id}, owner="alice")
|
||||
assert upload_id not in backup_service.chunked_uploads.sessions
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_init_starts_cleanup_task(self, backup_service):
|
||||
"""备份上传初始化必须启动过期清理任务"""
|
||||
backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
task = backup_service.chunked_uploads._cleanup_task
|
||||
assert task is not None and not task.done()
|
||||
task.cancel()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_failure_keeps_session_for_retry(
|
||||
self, backup_service, monkeypatch
|
||||
):
|
||||
"""目录删除失败时会话保持注册，看门狗可重试，成功后正常摘除"""
|
||||
import shutil
|
||||
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
|
||||
calls = {"n": 0}
|
||||
real_rmtree = shutil.rmtree
|
||||
|
||||
def _flaky(path, *args, **kwargs):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
raise OSError("transient fs error")
|
||||
return real_rmtree(path, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"astrbot.dashboard.services.chunked_upload_service.shutil.rmtree", _flaky
|
||||
)
|
||||
await backup_service.chunked_uploads.cleanup_session(upload_id)
|
||||
assert upload_id in backup_service.chunked_uploads.sessions
|
||||
|
||||
await backup_service.chunked_uploads.cleanup_session(upload_id)
|
||||
assert upload_id not in backup_service.chunked_uploads.sessions
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_write_chunk_is_rejected(self, backup_service):
|
||||
"""save() 落盘字节数不足时不得发布分片，且不留临时文件"""
|
||||
session = backup_service.upload_init(
|
||||
{"filename": "b.zip", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
|
||||
with pytest.raises(BackupServiceError, match="size mismatch"):
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"x" * 50),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
# 分片未发布、无临时文件残留；补传完整数据后正常完成
|
||||
chunk_dir = backup_service.chunked_uploads.chunks_root / upload_id
|
||||
assert not list(chunk_dir.glob("*.tmp"))
|
||||
status = backup_service.upload_status({"upload_id": upload_id}, owner="alice")
|
||||
assert status["received_chunks"] == []
|
||||
|
||||
await backup_service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubUploadFile(b"y" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
result = await backup_service.upload_complete(
|
||||
{"upload_id": upload_id}, owner="alice"
|
||||
)
|
||||
assert result["size"] == 100
|
||||
|
||||
@@ -534,11 +534,16 @@ async def test_save_uploaded_file_rejects_oversized_saved_file(
|
||||
content_type = "application/octet-stream"
|
||||
content_length = None
|
||||
|
||||
async def save(self, path):
|
||||
async def save(self, path, *, max_bytes=None):
|
||||
saved = Path(path)
|
||||
saved.write_bytes(b"x")
|
||||
with saved.open("rb+") as f:
|
||||
f.truncate(chat_service.MAX_UPLOAD_FILE_SIZE_BYTES + 1)
|
||||
# Mirror save_upload_stream: enforce the cap mid-write and
|
||||
# remove the partial file on overflow.
|
||||
if max_bytes is not None and saved.stat().st_size > max_bytes:
|
||||
saved.unlink()
|
||||
raise chat_service.UploadTooLargeError(max_bytes)
|
||||
|
||||
with pytest.raises(ChatServiceError, match="File too large"):
|
||||
await chat_service_instance.save_uploaded_file(FakeUpload())
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Tests for the shared streaming upload save helper."""
|
||||
|
||||
import pytest
|
||||
|
||||
from astrbot.core.utils.upload import (
|
||||
COPY_BLOCK_SIZE,
|
||||
UploadTooLargeError,
|
||||
save_upload_stream,
|
||||
)
|
||||
|
||||
|
||||
class FakeUpload:
|
||||
"""Minimal async upload object: seek + bounded read."""
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self._data = data
|
||||
self._pos = 0
|
||||
|
||||
async def seek(self, pos: int) -> int:
|
||||
self._pos = pos
|
||||
return self._pos
|
||||
|
||||
async def read(self, size: int = -1) -> bytes:
|
||||
if size < 0:
|
||||
size = len(self._data) - self._pos
|
||||
block = self._data[self._pos : self._pos + size]
|
||||
self._pos += len(block)
|
||||
return block
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_upload_stream_writes_all_bytes(tmp_path):
|
||||
data = b"x" * (COPY_BLOCK_SIZE * 2 + 123)
|
||||
dest = tmp_path / "out.bin"
|
||||
|
||||
written = await save_upload_stream(FakeUpload(data), dest)
|
||||
|
||||
assert written == len(data)
|
||||
assert dest.read_bytes() == data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_upload_stream_respects_max_bytes(tmp_path):
|
||||
data = b"x" * (COPY_BLOCK_SIZE + 1)
|
||||
dest = tmp_path / "out.bin"
|
||||
|
||||
with pytest.raises(UploadTooLargeError):
|
||||
await save_upload_stream(FakeUpload(data), dest, max_bytes=COPY_BLOCK_SIZE)
|
||||
|
||||
# The partial file must not survive a rejected upload.
|
||||
assert not dest.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_upload_stream_preserves_existing_dest_on_failure(tmp_path):
|
||||
"""A failed overwrite must not destroy a pre-existing destination."""
|
||||
dest = tmp_path / "out.bin"
|
||||
dest.write_bytes(b"original")
|
||||
|
||||
with pytest.raises(UploadTooLargeError):
|
||||
await save_upload_stream(
|
||||
FakeUpload(b"x" * (COPY_BLOCK_SIZE + 1)),
|
||||
dest,
|
||||
max_bytes=COPY_BLOCK_SIZE,
|
||||
)
|
||||
|
||||
assert dest.read_bytes() == b"original"
|
||||
# Only the uniquely named temp file is removed; nothing is left behind.
|
||||
assert not list(tmp_path.glob("*.tmp"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_upload_stream_allows_exact_limit(tmp_path):
|
||||
data = b"x" * COPY_BLOCK_SIZE
|
||||
dest = tmp_path / "out.bin"
|
||||
|
||||
written = await save_upload_stream(
|
||||
FakeUpload(data), dest, max_bytes=COPY_BLOCK_SIZE
|
||||
)
|
||||
|
||||
assert written == COPY_BLOCK_SIZE
|
||||
assert dest.read_bytes() == data
|
||||
@@ -0,0 +1,111 @@
|
||||
from astrbot.dashboard.server import _check_body_limit
|
||||
from astrbot.dashboard.services.backup_service import CHUNK_SIZE
|
||||
from astrbot.dashboard.services.chat_service import MAX_UPLOAD_FILE_SIZE_BYTES
|
||||
|
||||
DEFAULT_LIMIT = 128 * 1024 * 1024
|
||||
MULTIPART = "multipart/form-data; boundary=----test"
|
||||
|
||||
|
||||
class TestCheckBodyLimit:
|
||||
def test_non_api_paths_pass(self):
|
||||
assert (
|
||||
_check_body_limit("/assets/index.js", None, "", default_limit=DEFAULT_LIMIT)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_multipart_without_content_length_is_rejected(self):
|
||||
# Multipart form parsing spools large bodies to disk before any
|
||||
# per-file size check can run, so length must be declared up front.
|
||||
result = _check_body_limit(
|
||||
"/api/v1/files", None, MULTIPART, default_limit=DEFAULT_LIMIT
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0] == 411
|
||||
|
||||
def test_leading_space_content_type_still_rejected(self):
|
||||
# The form parser strips surrounding whitespace, so a padded
|
||||
# Content-Type is still multipart; detection must match the parser.
|
||||
result = _check_body_limit(
|
||||
"/api/v1/files",
|
||||
None,
|
||||
" multipart/form-data; boundary=----test",
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0] == 411
|
||||
|
||||
def test_get_with_junk_multipart_content_type_passes(self):
|
||||
# Body-less methods never trigger form parsing, so a bogus
|
||||
# multipart Content-Type on a GET (seen in the wild from the SSE
|
||||
# log client) must not be rejected with 411.
|
||||
assert (
|
||||
_check_body_limit(
|
||||
"/api/v1/logs/live",
|
||||
None,
|
||||
MULTIPART,
|
||||
method="GET",
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_json_without_content_length_passes(self):
|
||||
# Stopgap scope: only multipart is bounded pre-parse (disk spool);
|
||||
# other lengthless bodies pass and are bounded at save time.
|
||||
assert (
|
||||
_check_body_limit(
|
||||
"/api/v1/chat", None, "application/json", default_limit=DEFAULT_LIMIT
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_over_default_limit_rejected(self):
|
||||
result = _check_body_limit(
|
||||
"/api/v1/config",
|
||||
DEFAULT_LIMIT + 1,
|
||||
"application/json",
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0] == 413
|
||||
|
||||
def test_exact_limit_file_passes_with_multipart_slack(self):
|
||||
# The /files override includes framing overhead, so a file exactly
|
||||
# at the 512 MiB cap is not rejected, while anything beyond the
|
||||
# slack still is.
|
||||
assert (
|
||||
_check_body_limit(
|
||||
"/api/v1/files",
|
||||
MAX_UPLOAD_FILE_SIZE_BYTES,
|
||||
MULTIPART,
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
is None
|
||||
)
|
||||
result = _check_body_limit(
|
||||
"/api/v1/files",
|
||||
MAX_UPLOAD_FILE_SIZE_BYTES + 2 * 1024 * 1024,
|
||||
MULTIPART,
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0] == 413
|
||||
|
||||
def test_chunk_route_uses_chunk_limit(self):
|
||||
assert (
|
||||
_check_body_limit(
|
||||
"/api/v1/files/upload/chunk",
|
||||
CHUNK_SIZE * 2,
|
||||
MULTIPART,
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
is None
|
||||
)
|
||||
result = _check_body_limit(
|
||||
"/api/v1/files/upload/chunk",
|
||||
CHUNK_SIZE * 2 + 1,
|
||||
MULTIPART,
|
||||
default_limit=DEFAULT_LIMIT,
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0] == 413
|
||||
@@ -0,0 +1,207 @@
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from astrbot.core.utils.upload import UploadTooLargeError
|
||||
from astrbot.dashboard.services.chat_service import (
|
||||
MAX_UPLOAD_FILE_SIZE_BYTES,
|
||||
ChatService,
|
||||
ChatServiceError,
|
||||
)
|
||||
|
||||
CHUNK_SIZE = 1024 * 1024 # ChunkedUploadService default chunk size
|
||||
|
||||
|
||||
class _StubChunk:
|
||||
"""Minimal upload adapter matching the save() contract of UploadFileAdapter."""
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self._data = data
|
||||
|
||||
async def save(self, destination, *, max_bytes=None):
|
||||
if max_bytes is not None and len(self._data) > max_bytes:
|
||||
raise UploadTooLargeError(max_bytes)
|
||||
Path(destination).write_bytes(self._data)
|
||||
return len(self._data)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def service(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(
|
||||
"astrbot.dashboard.services.chat_service.get_astrbot_data_path",
|
||||
lambda: str(tmp_path),
|
||||
)
|
||||
db = SimpleNamespace(inserted=None)
|
||||
|
||||
async def insert_attachment(path, type, mime_type):
|
||||
db.inserted = {"path": path, "type": type, "mime_type": mime_type}
|
||||
return SimpleNamespace(attachment_id="att-1", path=path)
|
||||
|
||||
db.insert_attachment = insert_attachment
|
||||
svc = ChatService(
|
||||
db,
|
||||
SimpleNamespace(
|
||||
conversation_manager=None,
|
||||
platform_message_history_manager=None,
|
||||
umop_config_router=None,
|
||||
),
|
||||
)
|
||||
yield svc
|
||||
# upload_init() lazily starts the expiry janitor; do not let it leak
|
||||
# across tests.
|
||||
task = svc.chunked_uploads._cleanup_task
|
||||
if task is not None and not task.done():
|
||||
task.cancel()
|
||||
|
||||
|
||||
class TestChatChunkedUpload:
|
||||
def test_upload_init_validates_input(self, service):
|
||||
session = service.upload_init(
|
||||
{"filename": "report.pdf", "total_size": CHUNK_SIZE * 2 + 100},
|
||||
owner="alice",
|
||||
)
|
||||
assert session["chunk_size"] == CHUNK_SIZE
|
||||
assert session["total_chunks"] == 3
|
||||
assert session["upload_id"]
|
||||
|
||||
with pytest.raises(ChatServiceError, match="Missing key: filename"):
|
||||
service.upload_init({"total_size": 100}, owner="alice")
|
||||
with pytest.raises(ChatServiceError, match="Invalid file size"):
|
||||
service.upload_init({"filename": "a.txt", "total_size": 0}, owner="alice")
|
||||
with pytest.raises(ChatServiceError, match="File too large"):
|
||||
service.upload_init(
|
||||
{"filename": "a.bin", "total_size": MAX_UPLOAD_FILE_SIZE_BYTES + 1},
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_flow_merges_and_creates_attachment(self, service):
|
||||
part0 = b"x" * CHUNK_SIZE
|
||||
part1 = b"y" * 100
|
||||
session = service.upload_init(
|
||||
{
|
||||
"filename": "notes.txt",
|
||||
"total_size": len(part0) + len(part1),
|
||||
"content_type": "text/plain",
|
||||
},
|
||||
owner="alice",
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
|
||||
# Out-of-order arrival must still merge in index order.
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="1",
|
||||
chunk_file=_StubChunk(part1),
|
||||
owner="alice",
|
||||
)
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubChunk(part0),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
result = await service.upload_complete({"upload_id": upload_id}, owner="alice")
|
||||
|
||||
assert result["attachment_id"] == "att-1"
|
||||
assert result["type"] == "file"
|
||||
assert service.db.inserted["mime_type"] == "text/plain"
|
||||
merged = Path(service.db.inserted["path"])
|
||||
assert merged.read_bytes() == part0 + part1
|
||||
# Session, chunk dir and the temp merge file are all gone.
|
||||
assert not service.chunked_uploads.sessions
|
||||
assert not list(service.chunked_uploads.chunks_root.iterdir())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sessions_bound_to_owner(self, service):
|
||||
session = service.upload_init(
|
||||
{"filename": "a.txt", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
|
||||
with pytest.raises(ChatServiceError, match="not found or expired"):
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubChunk(b"x" * 100),
|
||||
owner="mallory",
|
||||
)
|
||||
with pytest.raises(ChatServiceError, match="not found or expired"):
|
||||
service.upload_status({"upload_id": upload_id}, owner="mallory")
|
||||
with pytest.raises(ChatServiceError, match="not found or expired"):
|
||||
await service.upload_complete({"upload_id": upload_id}, owner="mallory")
|
||||
with pytest.raises(ChatServiceError, match="not found or expired"):
|
||||
await service.upload_abort({"upload_id": upload_id}, owner="mallory")
|
||||
|
||||
# The real owner is unaffected by the probing attempts.
|
||||
result = await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubChunk(b"x" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
assert result["received"] == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_status_reports_progress_without_extending_lifetime(self, service):
|
||||
session = service.upload_init(
|
||||
{"filename": "a.txt", "total_size": CHUNK_SIZE * 2}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="1",
|
||||
chunk_file=_StubChunk(b"x" * CHUNK_SIZE),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
inner = service.chunked_uploads.get_session(upload_id, owner="alice")
|
||||
inner.last_activity -= 100 # pretend the session has been idle
|
||||
status = service.upload_status({"upload_id": upload_id}, owner="alice")
|
||||
|
||||
assert status["received_chunks"] == [1]
|
||||
assert status["total_chunks"] == 2
|
||||
assert status["expires_in"] <= 3600 - 100 + 1
|
||||
assert (
|
||||
service.chunked_uploads.get_session(upload_id, owner="alice").last_activity
|
||||
== inner.last_activity
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_rejects_missing_chunks(self, service):
|
||||
session = service.upload_init(
|
||||
{"filename": "a.txt", "total_size": CHUNK_SIZE + 1}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubChunk(b"x" * CHUNK_SIZE),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
with pytest.raises(ChatServiceError, match="Chunks incomplete"):
|
||||
await service.upload_complete({"upload_id": upload_id}, owner="alice")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abort_cleans_up_session(self, service):
|
||||
session = service.upload_init(
|
||||
{"filename": "a.txt", "total_size": 100}, owner="alice"
|
||||
)
|
||||
upload_id = session["upload_id"]
|
||||
await service.upload_chunk(
|
||||
upload_id=upload_id,
|
||||
chunk_index_str="0",
|
||||
chunk_file=_StubChunk(b"x" * 100),
|
||||
owner="alice",
|
||||
)
|
||||
|
||||
await service.upload_abort({"upload_id": upload_id}, owner="alice")
|
||||
|
||||
with pytest.raises(ChatServiceError, match="not found or expired"):
|
||||
service.upload_status({"upload_id": upload_id}, owner="alice")
|
||||
assert not list(service.chunked_uploads.chunks_root.iterdir())
|
||||
# Aborting an already-gone session stays silent.
|
||||
await service.upload_abort({"upload_id": upload_id}, owner="alice")
|
||||
@@ -20,7 +20,7 @@ async def test_webchat_upload_uses_detected_image_type(tmp_path):
|
||||
content_type = "image/png"
|
||||
content_length = None
|
||||
|
||||
async def save(self, destination):
|
||||
async def save(self, destination, *, max_bytes=None):
|
||||
with open(destination, "wb") as output:
|
||||
output.write(image_buffer.getvalue())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user