diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index bbdff52..0907902 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -220,6 +220,62 @@ jobs: packages: write steps: *build-image + release-notes: + name: Update GitHub release notes + needs: docker + if: github.server_url == 'https://github.com' && vars.REGISTRY == '' && startsWith(github.ref, 'refs/tags/v') + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + contents: write + steps: + - name: Append release commits + uses: actions/github-script@v9 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const tag = context.ref.replace('refs/tags/', ''); + const repo = context.repo; + const { data: release } = await github.rest.repos.getReleaseByTag({ ...repo, tag }); + const releases = await github.paginate(github.rest.repos.listReleases, { + ...repo, + per_page: 100, + }); + const published = releases.filter(r => !r.draft); + const baseTag = published[published.findIndex(r => r.id === release.id) + 1]?.tag_name; + const commits = baseTag + ? await github.paginate(github.rest.repos.compareCommitsWithBasehead, { + ...repo, + basehead: `${baseTag}...${tag}`, + per_page: 100, + }) + : await github.paginate(github.rest.repos.listCommits, { + ...repo, + sha: tag, + per_page: 100, + }); + const changelog = commits + .filter(c => c.parents.length <= 1) + .map(c => `- ${c.sha.substring(0, 7)} ${c.commit.message.split('\n')[0]}`) + .join('\n'); + + if (!changelog) { + core.setFailed('No commits found for the release notes.'); + return; + } + + const heading = baseTag ? `## Commits since ${baseTag}` : '## Commits'; + const existingBody = (release.body || '').trim(); + const body = existingBody + ? `${existingBody}\n\n---\n\n${heading}\n\n${changelog}` + : `${heading}\n\n${changelog}`; + + await github.rest.repos.updateRelease({ + ...repo, + release_id: release.id, + body, + }); + dockerhub-sync-readme: name: Sync Docker Hub README needs: test diff --git a/README.md b/README.md index 27b136c..3a45aa9 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,11 @@ media processing tasks, such as changing video codecs, containers, audio tracks, individual files or let Millwright handle new files automatically. > [!WARNING] -> Millwright is still in alpha phase and may have bugs or incomplete features. +> Millwright is in beta testing phase and may have bugs or incomplete features. +> +> I personally tested the app on my own library and it worked as expected at least for nvidia, qsv and software encoders. +> Unfortunately, I cannot guarantee that it will work perfectly for all setups or encoders without actual testing on +> those setups. **AMD** encoders may also work, but I haven't tested them as i dont have the necessary hardware. ### If you like my work, you might also like my other projects: @@ -62,7 +66,7 @@ bash <(curl -fsSL https://raw.githubusercontent.com/arabcoders/millwright/dev/in See the [installation guide](docs/installation.md) for requirements, storage, GPU access, and remote nodes. -After starting Millwright, open `http://:8000`, create your account, and [add your first library](docs/usage.md#add-a-library). +After starting Millwright, open `http://:8000`, create your account, claim the pending node on **Nodes**, and [add your first library](docs/usage.md#add-a-library). ## Security diff --git a/common/src/common/actions.py b/common/src/common/actions.py index 17fb9d7..5ce9eab 100644 --- a/common/src/common/actions.py +++ b/common/src/common/actions.py @@ -1,8 +1,9 @@ from __future__ import annotations +from pathlib import PurePath, PurePosixPath, PureWindowsPath from typing import TYPE_CHECKING, Annotated, Literal -from pydantic import BaseModel, Field, HttpUrl, field_validator +from pydantic import BaseModel, Field, HttpUrl, field_validator, model_validator from common.media.language import normalize_language from common.metadata import MetadataValues @@ -55,16 +56,40 @@ class Notify(PluginAction): on: list[Literal["job.succeeded", "job.failed", "job.cancelled"]] = Field(min_length=1) +class PathMapping(BaseModel): + source: str = Field(min_length=1, description="Path used by Millwright.") + destination: str = Field(min_length=1, description="Path used by the application.") + + @model_validator(mode="after") + def check_paths(self) -> PathMapping: + source = PurePath(self.source) + destination = ( + PureWindowsPath(self.destination) + if PureWindowsPath(self.destination).is_absolute() + else PurePosixPath(self.destination) + ) + if not source.is_absolute() or not destination.is_absolute(): + raise ValueError("Use absolute paths in mappings.") + if ".." in source.parts or ".." in destination.parts: + raise ValueError("Mapping paths can't contain '..'.") + return self + + class ExternalNotificationConfig(BaseModel): - application: Literal["radarr", "sonarr"] + application: Literal["radarr", "sonarr", "plex", "jellyfin", "emby"] url: HttpUrl = Field(description="Application URL, including any URL base.") api_key: str = Field( - min_length=1, repr=False, description="Application API key.", json_schema_extra={"ui": "password"} + min_length=1, repr=False, description="Application API key or token.", json_schema_extra={"ui": "password"} ) unmonitor_after_refresh: bool = Field( default=False, description="Disable future automatic upgrades for the movie or associated episodes after refresh succeeds.", ) + path_mappings: list[PathMapping] = Field( + default=[], + description="Add mappings if an application uses different paths for your files.", + json_schema_extra={"ui": "path_mappings"}, + ) @field_validator("url") @classmethod @@ -73,6 +98,15 @@ def check_url(cls, value: HttpUrl) -> HttpUrl: raise ValueError("Use an application URL without credentials, a query, or a fragment.") return value + @model_validator(mode="after") + def check_options(self) -> ExternalNotificationConfig: + if self.application not in {"radarr", "sonarr"} and self.unmonitor_after_refresh: + raise ValueError("Unmonitor after refresh is only supported for Radarr and Sonarr.") + sources = [PurePath(mapping.source) for mapping in self.path_mappings] + if len(sources) != len(set(sources)): + raise ValueError("Each Millwright path can appear only once.") + return self + class NotifyExternal(ExternalNotificationConfig, PluginAction): kind: Literal["notify_external"] = "notify_external" diff --git a/common/src/common/plan/models.py b/common/src/common/plan/models.py index b84bb17..4e1a7bd 100644 --- a/common/src/common/plan/models.py +++ b/common/src/common/plan/models.py @@ -57,24 +57,28 @@ class StreamPolicy(BaseModel): class StreamPolicyInputs(BaseModel): audio_keep_languages: list[str] = Field( - default=[], description="Comma-separated ISO-639-2 audio languages (eng, jpn). Empty keeps all" + default=[], description="Keep audio in these languages, such as eng or jpn. Leave empty to keep all languages." + ) + audio_drop_descriptive: bool = Field(default=False, description="Remove commentary and descriptive audio tracks.") + audio_keep_forced: bool = Field( + default=True, description="Keep forced audio tracks even if their language isn't selected." ) - audio_drop_descriptive: bool = Field(default=False, description="Drop commentary/description audio tracks") - audio_keep_forced: bool = Field(default=True, description="Keep forced audio tracks regardless of language") audio_keep_all_if_empty: bool = Field( - default=False, description="Keep all audio if the language policy would remove every track" + default=False, description="Keep all audio tracks if the other audio filters would remove them all." ) audio_max_per_language: int | None = Field( - default=None, ge=1, description="Keep at most this many audio tracks per language. Blank keeps all" + default=None, ge=1, description="Limit audio tracks per language. Leave empty for no limit." ) subtitle_keep_languages: list[str] = Field( default=[], - description="Comma-separated ISO-639-2 subtitle languages. Empty leaves subtitles untouched", + description="Keep subtitles in these languages, such as eng or jpn. Leave empty to keep all languages.", + ) + subtitle_strip_all: bool = Field(default=False, description="Remove subtitles, keeping forced tracks if enabled.") + subtitle_keep_forced: bool = Field( + default=True, description="Keep forced subtitles even if their language isn't selected." ) - subtitle_strip_all: bool = Field(default=False, description="Drop all subtitles except retained forced tracks") - subtitle_keep_forced: bool = Field(default=True, description="Keep forced subtitles regardless of language") - drop_data: bool = Field(default=False, description="Drop all data streams") - drop_cover_art: bool = Field(default=False, description="Drop embedded picture video streams") + drop_data: bool = Field(default=False, description="Remove data streams.") + drop_cover_art: bool = Field(default=False, description="Remove embedded cover art.") def stream_policy(self) -> StreamPolicy | None: audio = None @@ -101,10 +105,11 @@ def stream_policy(self) -> StreamPolicy | None: class AudioPlanInputs(StreamPolicyInputs): - audio_codec: str | None = Field(default=None, description="Codec for re-encoding matching audio tracks") - audio_bitrate: str | None = Field(default=None, description="Target bitrate for re-encoded audio tracks") + audio_codec: str | None = Field(default=None, description="Choose the output audio codec.") + audio_bitrate: str | None = Field(default=None, description="Set the output audio bitrate, such as 192k.") audio_when_codecs: list[str] = Field( - default=[], description="Re-encode audio with these source codecs. Empty matches all" + default=[], + description="Convert audio tracks with these source codecs. Leave empty to convert all selected tracks.", ) @field_validator("audio_codec", "audio_bitrate", mode="before") diff --git a/common/src/common/plugin.py b/common/src/common/plugin.py index 8609a59..068c0b7 100644 --- a/common/src/common/plugin.py +++ b/common/src/common/plugin.py @@ -12,7 +12,7 @@ TextField = Annotated[str, Field(json_schema_extra={"ui": "text"})] -FieldType = Literal["int", "float", "bool", "str", "text", "password", "enum", "strings"] +FieldType = Literal["int", "float", "bool", "str", "text", "password", "enum", "strings", "path_mappings"] def input_fields(model: type[BaseModel]) -> list[dict[str, Any]]: @@ -56,6 +56,9 @@ def _field_type(info: FieldInfo) -> FieldType: if get_origin(annotation) is Literal: return "enum" if get_origin(annotation) is list: + extra = info.json_schema_extra + if isinstance(extra, dict) and extra.get("ui") == "path_mappings": + return "path_mappings" return "strings" if annotation is bool: return "bool" diff --git a/common/src/common/process.py b/common/src/common/process.py index 5b188ae..22ef70b 100644 --- a/common/src/common/process.py +++ b/common/src/common/process.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Iterator + from collections.abc import Awaitable, Callable, Generator from pathlib import Path _TERM_GRACE_S = 10.0 @@ -24,7 +24,7 @@ @contextmanager def process_policy( *, create: Callable[..., Awaitable[asyncio.subprocess.Process]], stop: Callable[..., Awaitable[int]] -) -> Iterator[None]: +) -> Generator[None]: token = _POLICY.set((create, stop)) try: yield diff --git a/common/src/common/protocol/__init__.py b/common/src/common/protocol/__init__.py index c36c4a7..364bcc9 100644 --- a/common/src/common/protocol/__init__.py +++ b/common/src/common/protocol/__init__.py @@ -7,8 +7,7 @@ ActiveAttempt, AttemptMissing, Cancel, - Enrolled, - EnrollmentSaved, + CredentialsSaved, Error, Finished, HealthResult, @@ -21,7 +20,11 @@ NodeLogs, OfferedJob, Progress, + Registered, + RegistrationPending, SetLimits, + Unregister, + Unregistered, Welcome, WorkspaceChallenge, WorkspaceResponse, @@ -35,8 +38,7 @@ "ActiveAttempt", "AttemptMissing", "Cancel", - "Enrolled", - "EnrollmentSaved", + "CredentialsSaved", "Error", "Finished", "HealthResult", @@ -49,7 +51,11 @@ "NodeLogs", "OfferedJob", "Progress", + "Registered", + "RegistrationPending", "SetLimits", + "Unregister", + "Unregistered", "Welcome", "WorkspaceChallenge", "WorkspaceResponse", diff --git a/common/src/common/protocol/messages.py b/common/src/common/protocol/messages.py index d8b788c..f185e71 100644 --- a/common/src/common/protocol/messages.py +++ b/common/src/common/protocol/messages.py @@ -82,7 +82,7 @@ class Hello(_MessageBase): type: Literal["hello"] = "hello" node_id: str | None = Field(default=None, min_length=1) - token: str = Field(min_length=1) + token: str = "" capabilities: NodeCapabilities = Field(default_factory=NodeCapabilities) active_attempts: list[ActiveAttempt] @@ -109,19 +109,33 @@ class WorkspaceResponse(_MessageBase): _validate_path = field_validator("response_path")(_workspace_path) -class Enrolled(_MessageBase): - type: Literal["enrolled"] = "enrolled" +class RegistrationPending(_MessageBase): + type: Literal["registration_pending"] = "registration_pending" + + request_id: str = Field(min_length=1) + + +class Registered(_MessageBase): + type: Literal["registered"] = "registered" node_id: str = Field(min_length=1) token: str = Field(min_length=1) -class EnrollmentSaved(_MessageBase): - type: Literal["enrollment_saved"] = "enrollment_saved" +class CredentialsSaved(_MessageBase): + type: Literal["credentials_saved"] = "credentials_saved" node_id: str = Field(min_length=1) +class Unregister(_MessageBase): + type: Literal["unregister"] = "unregister" + + +class Unregistered(_MessageBase): + type: Literal["unregistered"] = "unregistered" + + class GPUUsage(_StrictModel): name: str = Field(min_length=1) percent: float | None = Field(default=None, ge=0, le=100) @@ -340,8 +354,11 @@ class Error(_MessageBase): Hello | WorkspaceChallenge | WorkspaceResponse - | Enrolled - | EnrollmentSaved + | RegistrationPending + | Registered + | CredentialsSaved + | Unregister + | Unregistered | Heartbeat | Progress | Finished @@ -361,8 +378,11 @@ class Error(_MessageBase): Hello, WorkspaceChallenge, WorkspaceResponse, - Enrolled, - EnrollmentSaved, + RegistrationPending, + Registered, + CredentialsSaved, + Unregister, + Unregistered, Heartbeat, Progress, Finished, diff --git a/common/src/common/workspace.py b/common/src/common/workspace.py index 9744fdf..91809f2 100644 --- a/common/src/common/workspace.py +++ b/common/src/common/workspace.py @@ -11,7 +11,7 @@ from common.fs.acquire import directory_fd if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Generator DEFAULT_MIN_FREE_GB = 10.0 DEFAULT_HARD_SIZE_CEILING = 1.25 @@ -41,7 +41,7 @@ def open_workspace_lock(path: Path, *, exclusive: bool = False) -> int: @contextmanager -def workspace_lock(path: Path, *, exclusive: bool = False) -> Iterator[int]: +def workspace_lock(path: Path, *, exclusive: bool = False) -> Generator[int]: fd = open_workspace_lock(path, exclusive=exclusive) try: # Cleanup can take the exclusive lock only after all writers release theirs. diff --git a/compose.yml b/compose.yml index a7be609..fdf4625 100644 --- a/compose.yml +++ b/compose.yml @@ -2,6 +2,7 @@ services: postgres: image: postgres:18 + container_name: millwright-postgres environment: POSTGRES_USER: millwright POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD in .env} @@ -16,6 +17,7 @@ services: restart: unless-stopped server: + container_name: millwright-server image: ghcr.io/arabcoders/millwright:latest user: "${UID:-1000}:${GID:-1000}" depends_on: @@ -37,6 +39,7 @@ services: restart: unless-stopped node: + container_name: millwright-node image: ghcr.io/arabcoders/millwright:latest user: "${UID:-1000}:${GID:-1000}" depends_on: @@ -45,7 +48,6 @@ services: environment: MW_ROLE: node MW_NODE_SERVER_URL: http://server:8000 - MW_NODE_ENROLLMENT_TOKEN: file:/run/node_enrollment_token MW_NODE_WORKSPACE_DIR: /workspace UMASK: ${UMASK:-0002} TZ: ${TZ:-UTC} @@ -53,7 +55,6 @@ services: - ./node-config:/config - ./node-data:/data - ./workspace:/workspace - - ./config/millwright/secrets/node_enrollment_token:/run/node_enrollment_token:ro restart: unless-stopped volumes: diff --git a/docs/development.md b/docs/development.md index d876c13..6370e01 100644 --- a/docs/development.md +++ b/docs/development.md @@ -32,10 +32,11 @@ Then start the server and a local node from the repository root: uv run python run.py --role all ``` -Open `http://localhost:8000`. Use `--role server` if you don't need a local node. +Open `http://localhost:8000` and claim the pending node on **Nodes**. Use `--role server` if you don't +need a local node. -For a separate source node, set `MW_NODE_SERVER_URL`, `MW_NODE_ENROLLMENT_TOKEN`, and -`MW_NODE_WORKSPACE_DIR`, then run `uv run python -m node` from `node/`. See +For a separate source node, set `MW_NODE_SERVER_URL` and `MW_NODE_WORKSPACE_DIR`, then run +`uv run python -m node` from `node/`. Claim it on **Nodes**. See [Add a node](installation.md#add-a-node) for connection and shared-storage setup. ## Use the CLI diff --git a/docs/events.md b/docs/events.md index b9bce10..3308d68 100644 --- a/docs/events.md +++ b/docs/events.md @@ -57,6 +57,7 @@ Some messages contain a batch. Handle each entry in `payload.events` as an ordin | `job.cancel_requested` | Show cancellation in progress | | `file.status` | Update the file list or show scan totals | | `node.status` | Update a node's connection state and active job counts | +| `node.registration` | Update registration availability and pending nodes | | `library.scan`, `library.deletion` | Show scan and deletion results | | `library.progress` | Update scanning and metadata-reading activity | | `scanner.state`, `probing.state` | Show scanning or metadata-reading status | @@ -67,6 +68,12 @@ Individual file updates contain `media_file_id`, `library_id`, and `file`. On `a the file from your client's list. Scan totals instead use `action: "scanned"` with `added`, `updated`, and `pruned` counts. +Node registration updates contain `enabled` and the complete `pending` list. Each pending request +contains `request_id` and `capabilities`. Load the initial state from `GET /api/nodes/registration`. +Claim a connected request with `POST /api/nodes/pending/{request_id}/claim`. A `202` response starts +registration. The node becomes available only after saving its credentials and verifying the workspace. +Use `PATCH /api/settings` with `node_registration_enabled` to allow or stop new requests. + For initial library activity, request `GET /api/libraries/progress`. Active metadata reads include `media_file_id`, `path`, `tool`, `started_at`, and nullable `percent`. These activity counters reset when the server restarts. diff --git a/docs/installation.md b/docs/installation.md index 411b7c1..449c081 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -11,14 +11,14 @@ Run this in the directory where you want to keep Millwright: bash <(curl -fsSL https://raw.githubusercontent.com/arabcoders/millwright/HEAD/install.sh) ``` -Enter your media folder when prompted. The installer starts PostgreSQL, the server, and a node. +Enter your media folder when prompted. The installer asks before starting compose file. Press Enter to skip startup. Rerunning it keeps your existing configuration and credentials. Open `http://YOUR_SERVER_IP:8000` and create your account. Use your server's IP address, or `http://localhost:8000` if your browser is on the same machine. -Then [add a library](usage.md#add-a-library), using `/media` for the folder you selected during -installation. Check **Nodes** to confirm the node is connected before starting a conversion. +Open **Nodes** and click **Claim** on the pending node. Then [add a library](usage.md#add-a-library), +using `/media` for the folder you selected during installation. ## Update Millwright @@ -87,12 +87,8 @@ workspace, stop the server and nodes and copy the contents of `workspace/` to th ### Connect the new node -Create a directory for the new node. Copy `config/millwright/secrets/node_enrollment_token` from -the server's installation directory into it, keeping the filename `node_enrollment_token`. -This file lets the new node connect. Keep it private. - -Save the following as `compose.yml` in the new node's directory. Replace `YOUR_SERVER_IP` with your -server's address and adjust the shared folder's host path if needed: +Create a directory for the new node and save the following as `compose.yml`. Replace +`YOUR_SERVER_IP` with your server's address and adjust the shared folder's host path if needed: ```yaml services: @@ -103,14 +99,12 @@ services: MW_ROLE: node UMASK: ${UMASK:-0002} MW_NODE_SERVER_URL: http://YOUR_SERVER_IP:8000 - MW_NODE_ENROLLMENT_TOKEN: file:/run/node_enrollment_token MW_NODE_WORKSPACE_DIR: /workspace TZ: ${TZ:-UTC} volumes: - ./node-config:/config - ./node-data:/data - /mnt/millwright-workspace:/workspace - - ./node_enrollment_token:/run/node_enrollment_token:ro restart: unless-stopped ``` @@ -127,8 +121,13 @@ Start the node: docker compose up -d --pull always ``` -Check **Nodes** in the web interface, then use **Nodes -> Edit** to set its processing limits. -Keep each node's `node-config` and `node-data` directories separate. Only the workspace is shared. +Open **Nodes** in the web interface and click **Claim**. Use **Edit** to set the node's processing limits. + +New node registration is enabled by default. The **Allow new nodes** switch on **Nodes** stops new +registration. + +Keep each node's `node-config` and `node-data` directories separate. The node saves its credentials +there and reconnects automatically after a restart. Only the workspace is shared. You can keep or remove the node on the server machine. For connections across an untrusted network, use [HTTPS](#reverse-proxies) for the server address. diff --git a/docs/plugins.md b/docs/plugins.md index 385e702..2a06fdb 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -274,13 +274,13 @@ return { } ``` -| Action | Result | -| --------------------- | ---------------------------------------------------------------------------------------------------------------------------- | -| `extract_subtitles` | Extract selected subtitles and publish sidecars beside the final media | -| `notify` | Send selected job outcomes to named notification targets | -| `notify_external` | Notify a configured external application after replacement. Supports Radarr and Sonarr refreshes, with optional unmonitoring | -| `preserve_timestamps` | Preserve the file times captured for processing | -| `embed_metadata` | Apply sidecar metadata during planning | +| Action | Result | +| --------------------- | ---------------------------------------------------------------------- | +| `extract_subtitles` | Extract selected subtitles and publish sidecars beside the final media | +| `notify` | Send selected job outcomes to named notification targets | +| `notify_external` | Notify a configured external application after replacement | +| `preserve_timestamps` | Preserve the file times captured for processing | +| `embed_metadata` | Apply sidecar metadata during planning | Required subtitle extraction happens before encoding. If removed tracks would otherwise be lost when discarding the original, their sidecars must be saved first. A collision leaves the original protected. diff --git a/docs/recipes.md b/docs/recipes.md index 5f639dd..e8c223c 100644 --- a/docs/recipes.md +++ b/docs/recipes.md @@ -15,9 +15,8 @@ To convert video to HEVC in an MKV container: 3. Connect its `plan` output to a **Done** result and `skip` to **Skip**. 4. Save the flow and preview it against a few files. -With no encoder selected, HEVC uses `libx265`. AVC and AV1 use `libx264` and `libsvtav1`. Quality mode -lets the encoder choose the bitrate. Try a sample to judge both quality and file size rather than -expecting a fixed percentage saving. +Leave the encoder unset to use `libx265` for HEVC, `libx264` for AVC, or `libsvtav1` for AV1. Use quality +mode to choose quality without a fixed bitrate. Try a sample to judge the quality and file size.
Import this example as JSON @@ -48,9 +47,8 @@ Save this as a JSON file and choose **Flows -> Import flow**: ### Target a source bitrate ratio -If you prefer a bitrate target, select bitrate mode and set a ratio. With the video bitrate source, -`0.5` targets half the reported video bitrate. If that bitrate is unavailable, it uses the whole-file -average, including audio and container data. The ratio doesn't halve the output file size. +Select bitrate mode and set a ratio. A ratio of `0.5` targets half the source bitrate, not half the file +size. The step inputs are: @@ -58,8 +56,7 @@ The step inputs are: {"target": "hevc", "rate_control": "bitrate", "ratio": 0.5, "bitrate_source": "video"} ``` -Use quality mode when the required source bitrate is unavailable. Quality and bitrate modes are -alternative ways of choosing the encoding target. +Use quality mode if the source bitrate is unavailable. ### Prefer a hardware encoder @@ -70,9 +67,9 @@ use this priority list: {"target": "hevc", "encoder_priority": ["hevc_qsv", "libx265"]} ``` -The choice is made before processing begins. It doesn't switch encoders halfway through a failed job. -A fixed encoder setting overrides the priority list. See [GPU access](installation.md#hardware-acceleration) -if your preferred encoder isn't available. +The first available encoder is used. A fixed encoder overrides this list. Failed encodes aren't retried +with another encoder. See [GPU access](installation.md#hardware-acceleration) if your preferred encoder +isn't available. ## Route by size, codec, or language @@ -106,8 +103,7 @@ Connect `plan` to **Done** and `skip` to **Skip**, as in the video example. Leav empty to convert all selected audio. Use the language and commentary options if you also want to narrow the track selection. -Audio re-encoding counts as transcoding for the library's original-retention setting, even though the -video is copied. +The [`transcode` trash mode](recovery.md#retain-originals) keeps the original when audio is re-encoded. ## Reject an output that shrinks too far @@ -123,25 +119,22 @@ This checks size, not visual quality. Rejected output doesn't replace your origi ## Embed sidecar metadata -Use **Embed sidecar metadata** to copy titles, descriptions, dates, and genres from Kodi NFO or yt-dlp -JSON files into MKV or MP4 media. Episode NFOs can also provide the show name, season, and episode number. -yt-dlp JSON can provide the artist from its channel or uploader. +Use **Embed sidecar metadata** to copy metadata from Kodi NFO or yt-dlp JSON files into MKV or MP4 files. For `Example.mkv`, put `Example.nfo` or `Example.info.json` beside it. Add **Embed sidecar metadata** -before the remux or video step, and connect its `next` output to that step. Millwright reads the -sidecars without changing them. +before the remux or video step, and connect its `next` output to that step. The step offers these settings: | Setting | Default | Choice | | ----------------------- | ---------------------- | --------------------------------------------------------------------- | -| `existing_tags` | `fill_missing` | Fill absent fields, or use `overwrite` to replace supplied fields | -| `source_priority` | `["nfo", "info_json"]` | The first source with a value wins for each field | -| `when_no_media_changes` | `skip` | Use `remux` to allow tag changes without an audio or video conversion | +| `existing_tags` | `fill_missing` | Fill missing tags, or use `overwrite` to replace supplied fields | +| `source_priority` | `["nfo", "info_json"]` | Prefer NFO or JSON when both supply a tag | +| `when_no_media_changes` | `skip` | Use `remux` to update tags without re-encoding | NFO files must describe a Kodi movie or episode. JSON must describe one yt-dlp video, not a playlist or -channel. Missing sidecars are allowed. If tags aren't added, check the job log for sidecar errors and -confirm that the step runs before conversion. Matching tags don't trigger another conversion. +channel. Missing sidecars are allowed. If tags aren't added, check the job log and confirm that the step +runs before conversion. MP4 tag support varies by player. Try a file in your player before applying this to a library. @@ -157,29 +150,34 @@ Choose the target and the outcomes you want to receive. For success, failure, an {"targets": ["ops"], "on": ["job.succeeded", "job.failed", "job.cancelled"]} ``` -Putting Notify before the conversion means it's already selected if that conversion fails. Disable the -step to stop notifications without deleting it. A notification can occasionally be delivered more than once. +Turn off **Enabled** to stop notifications without removing the step. Notifications can occasionally +arrive more than once. ## Notify external applications Use **Notify external application** to tell another application when Millwright replaces a file. -The supported applications are Radarr and Sonarr, which rescan the published media. This is useful when -a conversion changes the extension or media information the application last saw. +Supported applications are Radarr, Sonarr, Plex, Jellyfin, and Emby. -Add the step to the flow's successful path. Select the application, enter its base URL and API key, and -save the flow. Include any URL base, such as `https://example.org/sonarr`, but don't append `/api/v3`. +1. Add the step to the flow's successful path. +2. Select the application and enter its URL and API key or token. +3. Add **Path mappings** if the application uses different paths for your files. +4. Save the flow. -Millwright sends the refresh only after the converted file has replaced the original. It doesn't send -one for dry runs, skips, or failed conversions. Delivery status and errors appear under the job's -**Notification deliveries**. A notification failure doesn't mark the conversion itself as failed. +Use the application's base URL, such as `https://example.org/sonarr`, without an API endpoint. Add a +separate step for each application. -Enable **Unmonitor after refresh** only if you want to stop future automatic upgrades. Radarr unmonitors -the movie. Sonarr unmonitors the episodes in that file, not the whole season or series. +Each mapping pairs a **Millwright path** with an **Application path**. For example: -If lookup fails, check whether the media manager can recognize the published path and filename. Radarr -can use `{imdb-...}` or `{tmdb-...}` identifiers in the path or filename before trying its filename parser. +| Millwright path | Application path | +| ------------------- | ---------------- | +| `/mnt/movies` | `/media/Movies` | +| `/mnt/series` | `/media/TV` | -Flow and job exports can contain this API key. Keep them private. +For Radarr and Sonarr, enable **Unmonitor after refresh** to stop automatic upgrades for the converted +movie or episodes. -For the other direction, where a media manager tells Millwright about new imports, use -[Incoming webhooks](webhooks.md). +Check the job's **Notification deliveries** for status and errors. + +Flow and job exports can contain the API key or token. Keep them private. + +To receive imports from applications, use [Incoming webhooks](webhooks.md). diff --git a/docs/usage.md b/docs/usage.md index 696333d..a058665 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -190,8 +190,7 @@ describe files still in the library, and deleting history can change earlier siz - Create an Apprise target under **Notifications** and add a [notification step](recipes.md#job-notifications) to the flows you want to hear about. - Use [incoming webhooks](webhooks.md) to receive imports and renames from other applications. -- Use [Notify external applications](recipes.md#notify-external-applications) to tell Radarr or Sonarr - when a converted file is ready. +- [Notify applications after conversion](recipes.md#notify-external-applications). ## Command-line and API access diff --git a/install.sh b/install.sh index 97c8aa3..e5b1774 100755 --- a/install.sh +++ b/install.sh @@ -38,16 +38,19 @@ EOF fi mkdir -p config/millwright/secrets server-data node-config node-data workspace -for name in api_key node_enrollment_token; do - if [[ ! -f "config/millwright/secrets/$name" ]]; then - secret=$(openssl rand -hex 32) - printf '%s\n' "$secret" > "config/millwright/secrets/$name" - fi -done - -docker compose -f compose.yml up -d --pull always --wait +if [[ ! -f config/millwright/secrets/api_key ]]; then + secret=$(openssl rand -hex 32) + printf '%s\n' "$secret" > config/millwright/secrets/api_key +fi -printf '\nMillwright is ready on port 8000.\n' -printf 'Open http://:8000\n' -printf 'Installation directory: %s\n' "$PWD" -printf 'Create your account, then add /media as a library.\n' +printf '\nInstallation directory: %s\n' "$PWD" +if IFS= read -r -p 'Start Millwright containers now? [y/N] ' answer && [[ "$answer" =~ ^[Yy]([Ee][Ss])?$ ]]; then + docker compose -f compose.yml up -d --pull always --wait + printf '\nMillwright is ready on port 8000.\n' + printf 'Open http://:8000\n' + printf 'Create your account, claim the pending node on Nodes, then add /media as a library.\n' +else + printf '\nContainers were not started.\n' + printf 'To start them, run this from the installation directory:\n' + printf 'docker compose -f compose.yml up -d --pull always --wait\n' +fi diff --git a/node/config.example.toml b/node/config.example.toml index cdccfcb..d211ac6 100644 --- a/node/config.example.toml +++ b/node/config.example.toml @@ -2,16 +2,12 @@ # Config precedence: --config, $MW_NODE_CONFIG, $MW_CONFIG_PATH/node.toml, # then $XDG_CONFIG_HOME/millwright/node.toml (default: ~/.config/millwright/node.toml). # MW_NODE_ overrides each field. Lists use JSON arrays. -# The server assigns an ID and credential. State directory precedence: +# Claim the pending node on the server's Nodes page to assign an ID and credential. +# State directory precedence: # $MW_NODE_STATE_PATH, $MW_STATE_PATH, then $XDG_STATE_HOME/millwright # (default: ~/.local/state/millwright). -# enrollment_token accepts file: in config and environment values. # Unknown keys are ignored with a warning. -# Required for enrollment and re-enrollment. -enrollment_token = "change-me" -# enrollment_token = "file:/run/secrets/node_token" - # Required. Server URL. HTTP becomes WS, HTTPS becomes WSS, and /api/hub is added. # A full ws:// or wss:// endpoint ending in /api/hub also works. server_url = "http://localhost:8000" diff --git a/node/src/node/agent.py b/node/src/node/agent.py index b694af8..94a6f17 100644 --- a/node/src/node/agent.py +++ b/node/src/node/agent.py @@ -35,8 +35,7 @@ ActiveAttempt, AttemptMissing, Cancel, - Enrolled, - EnrollmentSaved, + CredentialsSaved, Error, Finished, Heartbeat, @@ -45,7 +44,11 @@ NodeCapabilities, OfferedJob, Progress, + Registered, + RegistrationPending, SetLimits, + Unregister, + Unregistered, Welcome, WorkspaceChallenge, WorkspaceResponse, @@ -108,8 +111,7 @@ async def close(self) -> None: ... @dataclass(kw_only=True) class AgentConfig: node_id: str | None = None - token: str - enrollment_token: str | None = None + token: str = "" credentials_path: Path | None = None server_url: str workspace_dir: str @@ -129,7 +131,7 @@ class _WebSocketConnection: def __init__(self, ws: ClientConnection, config: AgentConfig) -> None: self._ws = ws self._config = config - self._secrets = (config.token, config.enrollment_token or "") + self._secrets = (config.token,) async def send_json(self, data: dict) -> None: await self._ws.send(json.dumps(data)) @@ -151,7 +153,7 @@ async def recv_json(self) -> dict | None: frame.code if frame is not None else 1006, redact( reason, - secrets=(*self._secrets, self._config.token, self._config.enrollment_token or ""), + secrets=(*self._secrets, self._config.token), ), extra={"operation": "node.connection.closed"}, ) @@ -203,7 +205,6 @@ def __init__( self._transcode_limit = 1 self._lane_limits = {"healthcheck": 1, "mkv_stats": 1} self._active_lanes = {"healthcheck": 0, "mkv_stats": 0} - self._enroll = config.node_id is None self._active_transcodes = 0 self._cpu_sample: tuple[int, int] | None = None @@ -225,15 +226,12 @@ async def run(self) -> None: context = "The hub is temporarily unavailable. " elif isinstance(exc, websockets.exceptions.InvalidStatus): cause = f"HTTP {exc.response.status_code} rejected the WebSocket upgrade." - context = ( - "No enrollment token has been sent. " - "Check /api/hub routing and WebSocket upgrades on the server and reverse proxy. " - ) + context = "Check /api/hub routing and WebSocket upgrades on the server and reverse proxy. " else: userinfo = urlsplit(self._config.server_url).netloc.rpartition("@")[0] cause = error_message( exc, - secrets=(self._config.token, self._config.enrollment_token or "", userinfo), + secrets=(self._config.token, userinfo), ) if isinstance(exc, TimeoutError): context = "Connection timed out before the WebSocket upgrade completed. " @@ -265,7 +263,7 @@ async def run(self) -> None: reason = ( "Connection dropped without a WebSocket close handshake" if exc.rcvd is None and exc.sent is None - else redact(str(exc), secrets=(self._config.token, self._config.enrollment_token or "")) + else redact(str(exc), secrets=(self._config.token,)) ) if frame is not None and frame.code == 1012: LOG.info("Node hub is shutting down or restarting. Reconnecting in %.1f seconds.", backoff) @@ -275,12 +273,8 @@ async def run(self) -> None: LOG.error( "Node hub session ended unexpectedly.", extra={ - "error": error_message( - exc, secrets=(self._config.token, self._config.enrollment_token or "") - ), - "traceback": error_traceback( - exc, secrets=(self._config.token, self._config.enrollment_token or "") - ), + "error": error_message(exc, secrets=(self._config.token,)), + "traceback": error_traceback(exc, secrets=(self._config.token,)), }, ) finally: @@ -383,7 +377,7 @@ def _expire(self) -> None: async def _session(self, conn: HubConnection) -> bool: self._conn = conn - credentials = (self._config.token, self._config.enrollment_token or "") + credentials = (self._config.token,) hello = Hello( active_attempts=[ ActiveAttempt(job_id=job_id, attempt_id=attempt_id) @@ -394,10 +388,8 @@ async def _session(self, conn: HubConnection) -> bool: if not task.done() ) ], - node_id=None if self._enroll else self._config.node_id, - token=self._config.enrollment_token - if self._enroll and self._config.enrollment_token - else self._config.token, + node_id=self._config.node_id, + token=self._config.token, capabilities=self._config.capabilities.model_copy( update={ "max_transcodes": 1, @@ -430,9 +422,17 @@ async def _session(self, conn: HubConnection) -> bool: }, ) return False + if isinstance(first, RegistrationPending): + LOG.info("Node registration request %s is waiting to be claimed on the Nodes page.", first.request_id) + data = await conn.recv_json() + if data is None: + return False + first = parse_message(data) + if isinstance(first, Unregister): + await self._unregister() + await conn.send_json(Unregistered().model_dump(mode="json")) + return False if isinstance(first, Error): - if first.code == "auth" and self._config.node_id is not None and self._config.enrollment_token: - self._enroll = True LOG.warning( "Node hub handshake rejected at welcome: %s (%s)", redact(first.message, secrets=credentials), @@ -468,9 +468,11 @@ async def _session(self, conn: HubConnection) -> bool: extra={"error": error_message(exc, secrets=credentials)}, ) return False + if isinstance(first, Unregister): + await self._unregister() + await conn.send_json(Unregistered().model_dump(mode="json")) + return False if isinstance(first, Error): - if first.code == "auth" and self._config.node_id is not None and self._config.enrollment_token: - self._enroll = True LOG.warning( "Node hub handshake rejected after workspace: %s (%s)", redact(first.message, secrets=credentials), @@ -482,46 +484,42 @@ async def _session(self, conn: HubConnection) -> bool: }, ) return False - if isinstance(first, Enrolled): - if not self._enroll: - LOG.warning("Node hub handshake received unexpected Enrolled while not enrolling.") - return False + if isinstance(first, Registered): credentials = (*credentials, first.token) if self._config.credentials_path is not None: try: self._save_credentials(first.node_id, first.token) except OSError as exc: LOG.warning( - "Node hub enrollment failed while saving credentials.", + "Node registration failed while saving credentials.", extra={"error": error_message(exc, secrets=credentials)}, ) return False self._config.node_id = first.node_id self._config.token = first.token - self._enroll = False - await conn.send_json(EnrollmentSaved(node_id=first.node_id).model_dump(mode="json")) + await conn.send_json(CredentialsSaved(node_id=first.node_id).model_dump(mode="json")) try: data = await asyncio.wait_for(conn.recv_json(), timeout=_HANDSHAKE_TIMEOUT_S) except (TimeoutError, ValueError) as exc: LOG.warning( - "Node hub enrollment handshake failed.", + "Node registration handshake failed.", extra={"error": error_message(exc, secrets=credentials)}, ) return False if data is None: - LOG.warning("Node hub handshake ended without a welcome reply after enrollment.") + LOG.warning("Node hub handshake ended without a welcome reply after registration.") return False try: first = parse_message(data) except ValueError as exc: LOG.warning( - "Node hub enrollment handshake reply was invalid.", + "Node registration handshake reply was invalid.", extra={"error": error_message(exc, secrets=credentials)}, ) return False if isinstance(first, Error): LOG.warning( - "Node hub handshake rejected after enrollment: %s (%s)", + "Node hub handshake rejected after registration: %s (%s)", redact(first.message, secrets=credentials), redact(first.code, secrets=credentials), extra={ @@ -544,19 +542,7 @@ async def _session(self, conn: HubConnection) -> bool: return False authorized = {(attempt.job_id, attempt.attempt_id) for attempt in first.active_attempts} self._disconnecting = True - revoked = [] - for owners in (self._jobs, self._action_tasks): - for key, task in owners.items(): - if not task.done() and key not in authorized: - self._cancel_reasons[key] = _REVOKED_ERROR - if not task.cancelling(): - task.cancel() - revoked.append(task) - if revoked: - outcomes = await drain(asyncio.gather(*revoked, return_exceptions=True)) - for outcome in outcomes: - if isinstance(outcome, ChildNotReapedError): - raise outcome + await self._revoke_attempts(authorized) self._welcomed = True self._disconnecting = False if self._reconnect_timer is not None: @@ -582,6 +568,11 @@ async def _session(self, conn: HubConnection) -> bool: unreaped = False try: done, _ = await asyncio.wait({receiver, writer, heartbeat}, return_when=asyncio.FIRST_COMPLETED) + if receiver in done and await receiver: + writer.cancel() + heartbeat.cancel() + await asyncio.gather(writer, heartbeat) + await asyncio.wait_for(conn.send_json(Unregistered().model_dump(mode="json")), _HANDSHAKE_TIMEOUT_S) for task in (heartbeat, writer, receiver): if task in done: await task @@ -590,7 +581,7 @@ async def _session(self, conn: HubConnection) -> bool: raise finally: self._disconnecting = True - if not self._closed and self._reconnect_timer is None: + if not self._closed and self._config.node_id is not None and self._reconnect_timer is None: self._reconnect_timer = asyncio.get_running_loop().call_later(max(0.0, first.lease_ttl_s), self._expire) for task in (writer, heartbeat, receiver): task.cancel() @@ -608,6 +599,40 @@ async def _session(self, conn: HubConnection) -> bool: raise asyncio.CancelledError return True + async def _revoke_attempts(self, authorized: set[tuple[int, int]]) -> None: + revoked = [] + for owners in (self._jobs, self._action_tasks): + for key, task in owners.items(): + if not task.done() and key not in authorized: + self._cancel_reasons[key] = _REVOKED_ERROR + if not task.cancelling(): + task.cancel() + revoked.append(task) + if revoked: + outcomes = await drain(asyncio.gather(*revoked, return_exceptions=True)) + for outcome in outcomes: + if isinstance(outcome, ChildNotReapedError): + raise outcome + + async def _unregister(self) -> None: + self._disconnecting = True + await self._revoke_attempts(set()) + path = self._config.credentials_path + if path is not None: + path.unlink(missing_ok=True) + if path.parent.exists(): + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + self._config.node_id = None + self._config.token = "" + if self._reconnect_timer is not None: + self._reconnect_timer.cancel() + self._reconnect_timer = None + LOG.info("Node registration removed. Reconnecting as an unclaimed node.") + def _save_credentials(self, node_id: str, token: str) -> None: path = self._config.credentials_path if path is None: @@ -634,21 +659,22 @@ def _save_credentials(self, node_id: str, token: str) -> None: temp.unlink() raise - async def _receive_loop(self, conn: HubConnection) -> None: + async def _receive_loop(self, conn: HubConnection) -> bool: while True: data = await conn.recv_json() if data is None or self._closed or self._disconnecting: - return + return False try: msg = parse_message(data) except Exception as exc: LOG.warning( "Ignored invalid hub message.", - extra={ - "error": error_message(exc, secrets=(self._config.token, self._config.enrollment_token or "")) - }, + extra={"error": error_message(exc, secrets=(self._config.token,))}, ) continue + if isinstance(msg, Unregister): + await self._unregister() + return True if isinstance(msg, Error): reason = f"hub rejected message ({msg.code}): {msg.message}" raise RuntimeError(reason) @@ -1292,10 +1318,8 @@ async def _heartbeat_loop(self, interval_s: float) -> None: extra={ "operation": "node.heartbeat.failed", "node_id": self._config.node_id, - "error": error_message(exc, secrets=(self._config.token, self._config.enrollment_token or "")), - "traceback": error_traceback( - exc, secrets=(self._config.token, self._config.enrollment_token or "") - ), + "error": error_message(exc, secrets=(self._config.token,)), + "traceback": error_traceback(exc, secrets=(self._config.token,)), }, ) raise @@ -1350,7 +1374,7 @@ def _workspace_response(self, challenge: WorkspaceChallenge) -> dict: workspace_free_bytes=shutil.disk_usage(root).free, ).model_dump(mode="json") except (OSError, ValueError) as exc: - secrets = (self._config.token, self._config.enrollment_token or "") + secrets = (self._config.token,) cause = error_message(exc, secrets=secrets) LOG.warning( "Node workspace verification failed. Resolved workspace: %s. Challenge: %s. " diff --git a/node/src/node/cli.py b/node/src/node/cli.py index eead0b1..ebc16e7 100644 --- a/node/src/node/cli.py +++ b/node/src/node/cli.py @@ -26,7 +26,6 @@ from common.diagnostics import error_message from common.logging import setup_logging from common.protocol import NodeCapabilities -from common.secrets import resolve_secret from pydantic import TypeAdapter, ValidationError from node.agent import Agent, AgentConfig, safe_endpoint @@ -44,7 +43,6 @@ ENV_STATE = "MW_NODE_STATE_PATH" ENV_OVERRIDES = { - "enrollment_token": "MW_NODE_ENROLLMENT_TOKEN", "server_url": "MW_NODE_SERVER_URL", "workspace_dir": "MW_NODE_WORKSPACE_DIR", "ffmpeg_bin": "MW_NODE_FFMPEG_BIN", @@ -168,17 +166,13 @@ def _build_settings(data: dict[str, Any], path: Path) -> AgentConfig: data.pop("max_transcodes", None) data.pop("credentials_path", None) data.pop("token", None) - if isinstance(data.get("enrollment_token"), str): - data["enrollment_token"] = resolve_secret(data["enrollment_token"], base_dir=path.parent) credentials = _load_credentials() if credentials is not None: data["node_id"], data["token"] = credentials - else: - data["token"] = data.get("enrollment_token") missing = ", ".join( f"{key} (config key {key!r} or env {ENV_OVERRIDES[key]})" - for key in ("enrollment_token", "server_url", "workspace_dir") - if not data.get(key) and (key != "enrollment_token" or credentials is None) + for key in ("server_url", "workspace_dir") + if not data.get(key) ) if missing: msg = f"Node settings incomplete ({path}). Missing: {missing}" @@ -263,7 +257,7 @@ async def _amain(args: argparse.Namespace, dates: Dates) -> int: handlers[sig] = previous LOG.info( "Starting node %s. Hub endpoint: %s. Worker tags: %s", - config.node_id or "awaiting enrollment", + config.node_id or "awaiting claim", safe_endpoint(config.server_url), ", ".join(config.capabilities.worker_tags) or "none", ) diff --git a/node/src/node/files.py b/node/src/node/files.py index 959bee1..20500bd 100644 --- a/node/src/node/files.py +++ b/node/src/node/files.py @@ -8,7 +8,7 @@ from common.fs.effects import publish if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Generator from pathlib import Path from typing import BinaryIO @@ -47,7 +47,7 @@ def check_free_space(path: Path, min_free_gb: float) -> None: @contextmanager -def partial_output(path: Path) -> Iterator[_Output]: +def partial_output(path: Path) -> Generator[_Output]: if path.exists(): raise ValueError("subtitle artifact already exists") output = create_output(path.with_name(path.name + ".partial")) diff --git a/node/src/node/repair.py b/node/src/node/repair.py index 0484af3..10d8f14 100644 --- a/node/src/node/repair.py +++ b/node/src/node/repair.py @@ -19,7 +19,7 @@ from node.process import spawn_tool, stop_tool if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Collection + from collections.abc import AsyncGenerator, Callable, Collection from common.capabilities import MkvTools @@ -48,7 +48,7 @@ async def prepare_input( log: Callable[[str], None], min_free_gb: float = 0, timeout_s: float = 600, -) -> AsyncIterator[PreparedInput]: +) -> AsyncGenerator[PreparedInput]: ratio = mkv_repair_ratio(probe) if not ratio: yield PreparedInput(path=source, probe=probe) diff --git a/plugins/action.embed_metadata/plugin.py b/plugins/action.embed_metadata/plugin.py index 9d30bbf..6e56c0b 100644 --- a/plugins/action.embed_metadata/plugin.py +++ b/plugins/action.embed_metadata/plugin.py @@ -12,20 +12,20 @@ class Inputs(BaseModel): existing_tags: Literal["fill_missing", "overwrite"] = Field( - default="fill_missing", description="Fill absent tags or replace only supplied fields" + default="fill_missing", description="Fill missing tags or overwrite tags supplied by the sidecar." ) source_priority: list[MetadataSource] = Field( - default=["nfo", "info_json"], description="Source order. The first nonempty value wins for each field" + default=["nfo", "info_json"], description="Choose which sidecars to prefer when they provide the same tag." ) when_no_media_changes: Literal["skip", "remux"] = Field( - default="skip", description="Skip metadata-only work or permit changed tags alone to trigger processing" + default="skip", description="Skip files needing only tag changes, or remux them to update the tags." ) @field_validator("source_priority") @classmethod def check_priority(cls, value: list[MetadataSource]) -> list[MetadataSource]: if not value or len(value) != len(set(value)): - raise ValueError("source_priority must be nonempty and unique") + raise ValueError("Choose at least one metadata source without duplicates.") return value diff --git a/plugins/action.extract_subs/plugin.py b/plugins/action.extract_subs/plugin.py index 5376a07..b3d1ee9 100644 --- a/plugins/action.extract_subs/plugin.py +++ b/plugins/action.extract_subs/plugin.py @@ -10,8 +10,8 @@ class Inputs(BaseModel): - formats: list[str] = Field(default=["srt"], description="Subtitle output formats to extract: srt or vtt.") - languages: list[str] = Field(default=["eng"], description="ISO 639-2 language tags to extract.") + formats: list[str] = Field(default=["srt"], description="Extract subtitles as srt or vtt.") + languages: list[str] = Field(default=["eng"], description="Choose subtitle language codes, such as eng or jpn.") INPUTS = Inputs diff --git a/plugins/action.extract_subs/plugin.toml b/plugins/action.extract_subs/plugin.toml index db21b98..825f8f1 100644 --- a/plugins/action.extract_subs/plugin.toml +++ b/plugins/action.extract_subs/plugin.toml @@ -3,7 +3,6 @@ kind = "action" description = "Extract selected subtitle tracks beside converted media." help = """ Select formats and languages. The default is English SRT. -Extracted files are written beside the converted media. """ name = "Extract subtitles" author = "Millwright" diff --git a/plugins/action.notify/plugin.py b/plugins/action.notify/plugin.py index 42d26cb..5bd6127 100644 --- a/plugins/action.notify/plugin.py +++ b/plugins/action.notify/plugin.py @@ -12,11 +12,11 @@ class Inputs(BaseModel): enabled: bool = Field( default=True, - description="Enable notifications from this flow step. When off, this step adds no notification actions.", + description="Enable notifications from this step.", ) - targets: list[str] = Field(default=[], description="Notification target identifiers. Required when enabled.") + targets: list[str] = Field(default=[], description="Choose at least one notification target when enabled.") on: list[Literal["job.succeeded", "job.failed", "job.cancelled"]] = Field( - default=[], description="Job events that trigger notifications. Required when enabled." + default=[], description="Choose at least one job event to notify on when enabled." ) @model_validator(mode="after") diff --git a/plugins/action.notify_external/plugin.toml b/plugins/action.notify_external/plugin.toml index 2d0b8df..98624fa 100644 --- a/plugins/action.notify_external/plugin.toml +++ b/plugins/action.notify_external/plugin.toml @@ -1,6 +1,6 @@ id = "action.notify_external" kind = "action" -description = "Refresh Radarr or Sonarr after verified publication. Optionally unmonitor the movie or associated episodes after refresh succeeds." +description = "Notify applications after conversion." name = "Notify external application" author = "Millwright" url = "https://github.com/arabcoders/millwright" diff --git a/plugins/inspector.audio_languages/plugin.py b/plugins/inspector.audio_languages/plugin.py index 0dbdb7e..355a980 100644 --- a/plugins/inspector.audio_languages/plugin.py +++ b/plugins/inspector.audio_languages/plugin.py @@ -10,7 +10,9 @@ class Inputs(BaseModel): - expected: list[str] = Field(default=["eng", "und"], description="ISO 639-2 language tags allowed on audio streams.") + expected: list[str] = Field( + default=["eng", "und"], description="List allowed audio language codes, such as eng or und." + ) INPUTS = Inputs diff --git a/plugins/inspector.bitrate_estimate/plugin.py b/plugins/inspector.bitrate_estimate/plugin.py index 98f8b76..90e8545 100644 --- a/plugins/inspector.bitrate_estimate/plugin.py +++ b/plugins/inspector.bitrate_estimate/plugin.py @@ -10,16 +10,16 @@ class Inputs(BaseModel): source: Literal["video", "container", "average"] = Field( - default="video", description="Bitrate source: video, container, or their average when both are available" + default="video", description="Choose video bitrate, whole-file bitrate, or their average." ) avc_ratio: float = Field( - default=0.75, ge=0.05, le=1.0, description="AVC bitrate estimate as a source bitrate ratio." + default=0.75, ge=0.05, le=1.0, description="Estimate AVC bitrate as this fraction of the source bitrate." ) hevc_ratio: float = Field( - default=0.5, ge=0.05, le=1.0, description="HEVC bitrate estimate as a source bitrate ratio." + default=0.5, ge=0.05, le=1.0, description="Estimate HEVC bitrate as this fraction of the source bitrate." ) av1_ratio: float = Field( - default=0.45, ge=0.05, le=1.0, description="AV1 bitrate estimate as a source bitrate ratio." + default=0.45, ge=0.05, le=1.0, description="Estimate AV1 bitrate as this fraction of the source bitrate." ) @@ -29,9 +29,9 @@ class Inputs(BaseModel): def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: video = next(iter(ctx.probe.real_video), None) if video is None: - return {"type": "skip", "reason": "no video stream"} + return {"type": "skip", "reason": "No video stream found."} if ctx.probe.target_bitrate_kbps(inputs.source, 1) is None: - return {"type": "skip", "reason": "no derivable video bitrate"} + return {"type": "skip", "reason": "Cannot determine the source bitrate."} ratios = {"avc": inputs.avc_ratio, "hevc": inputs.hevc_ratio, "av1": inputs.av1_ratio} targets = ["avc", "hevc", "av1"] facets: dict[str, int] = {} diff --git a/plugins/inspector.bitrate_estimate/plugin.toml b/plugins/inspector.bitrate_estimate/plugin.toml index db5dfaf..cce695a 100644 --- a/plugins/inspector.bitrate_estimate/plugin.toml +++ b/plugins/inspector.bitrate_estimate/plugin.toml @@ -2,9 +2,7 @@ id = "inspector.bitrate_estimate" kind = "inspector" description = "Estimate AVC, HEVC, and AV1 bitrate and size." help = """ -Apply codec ratios to the selected bitrate source to estimate bitrate and size. -Video bitrate falls back to file size and duration. Skip files without video or -a derivable bitrate. +Set a bitrate ratio for each codec. Files without video or a usable bitrate are skipped. """ name = "Bitrate estimate" author = "Millwright" diff --git a/plugins/inspector.output_ratio/plugin.py b/plugins/inspector.output_ratio/plugin.py index 64f6812..1af2373 100644 --- a/plugins/inspector.output_ratio/plugin.py +++ b/plugins/inspector.output_ratio/plugin.py @@ -12,16 +12,13 @@ class Inputs(BaseModel): metric: Literal["size", "bitrate"] = Field( default="size", - description=( - "Compare output and source by file size or bitrate. Bitrate uses the first video stream's reported value, " - "or whole-file average bitrate when unavailable." - ), + description="Compare the output and input by file size or bitrate.", ) min_ratio: float = Field( default=0.5, ge=0.01, le=100.0, - description="Minimum output-to-source ratio. 0.5 permits output values down to half the source.", + description="Require output to be at least this fraction of the input. For example, 0.5 means half.", ) @@ -41,15 +38,15 @@ def metric_value(metric: str, probe: ProbeData) -> float | None: def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: if ctx.source_probe is None: - return {"type": "skip", "reason": "no source probe"} + return {"type": "skip", "reason": "Input file information is unavailable for this check."} output = metric_value(inputs.metric, ctx.probe) source = metric_value(inputs.metric, ctx.source_probe) if output is None or source is None or source == 0: - return {"type": "veto", "reason": f"cannot determine {inputs.metric} for guard comparison"} + return {"type": "veto", "reason": f"Cannot compare input and output {inputs.metric}."} ratio = output / source if ratio < inputs.min_ratio: return { "type": "veto", - "reason": f"output {inputs.metric} is {ratio:.2f}x of source, minimum is {inputs.min_ratio:.2f}x", + "reason": f"Output {inputs.metric} is {ratio:.2f} times the input, below the minimum of {inputs.min_ratio:.2f}.", } return {"type": "match"} diff --git a/plugins/inspector.output_ratio/plugin.toml b/plugins/inspector.output_ratio/plugin.toml index 341fdc7..05e7b03 100644 --- a/plugins/inspector.output_ratio/plugin.toml +++ b/plugins/inspector.output_ratio/plugin.toml @@ -2,9 +2,7 @@ id = "inspector.output_ratio" kind = "inspector" description = "Reject output below a minimum source size or bitrate ratio." help = """ -Reject output below the minimum output-to-source ratio (default: 0.5). -Compare file size or video bitrate. Bitrate falls back to size and duration. -Skip without a source probe. Reject if either metric is unavailable. +Place after a conversion step. Choose size or bitrate and set the minimum ratio. """ name = "Minimum output ratio" author = "Millwright" diff --git a/plugins/inspector.stream_guard/plugin.py b/plugins/inspector.stream_guard/plugin.py index e08d67a..b549372 100644 --- a/plugins/inspector.stream_guard/plugin.py +++ b/plugins/inspector.stream_guard/plugin.py @@ -11,17 +11,19 @@ class Inputs(BaseModel): - min_audio: int = Field(default=1, ge=0, description="Minimum number of audio streams. Zero disables the check.") + min_audio: int = Field( + default=1, ge=0, description="Require at least this many audio streams. Set to zero to disable." + ) min_subtitles: int = Field( - default=0, ge=0, description="Minimum number of subtitle streams. Zero disables the check." + default=0, ge=0, description="Require at least this many subtitle streams. Set to zero to disable." ) audio_languages: list[str] = Field( default=[], - description="Comma-separated ISO-639-2 audio languages. Blank counts all, and und includes missing tags.", + description=("Count these audio language codes. Leave empty to count all. Use und for untagged tracks."), ) subtitle_languages: list[str] = Field( default=[], - description="Comma-separated ISO-639-2 subtitle languages. Blank counts all, and und includes missing tags.", + description=("Count these subtitle language codes. Leave empty to count all. Use und for untagged tracks."), ) @@ -35,18 +37,18 @@ def _count(streams: list[ProbeStream], languages: list[str]) -> int: def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: if ctx.source_probe is None: - return {"type": "skip", "reason": "no source probe"} + return {"type": "skip", "reason": "Input file information is unavailable for this check."} audio = _count(ctx.probe.audio, inputs.audio_languages) subtitles = _count(ctx.probe.subtitle, inputs.subtitle_languages) if audio < inputs.min_audio: - reason = f"audio streams {audio} below minimum {inputs.min_audio}" + reason = f"Audio stream count is {audio}, below the minimum of {inputs.min_audio}" if inputs.audio_languages: reason += f" (languages: {', '.join(inputs.audio_languages)})" - return {"type": "veto", "reason": reason} + return {"type": "veto", "reason": reason + "."} if inputs.min_subtitles and subtitles < inputs.min_subtitles: - reason = f"subtitle streams {subtitles} below minimum {inputs.min_subtitles}" + reason = f"Subtitle stream count is {subtitles}, below the minimum of {inputs.min_subtitles}" if inputs.subtitle_languages: reason += f" (languages: {', '.join(inputs.subtitle_languages)})" - return {"type": "veto", "reason": reason} + return {"type": "veto", "reason": reason + "."} return {"type": "match"} diff --git a/plugins/inspector.stream_guard/plugin.toml b/plugins/inspector.stream_guard/plugin.toml index ce37796..90cec05 100644 --- a/plugins/inspector.stream_guard/plugin.toml +++ b/plugins/inspector.stream_guard/plugin.toml @@ -2,9 +2,8 @@ id = "inspector.stream_guard" kind = "inspector" description = "Reject output with too few audio or subtitle streams." help = """ -Run after transcoding. Check audio and subtitle counts separately. -Set a minimum to zero to disable its check. Optional comma-separated ISO-639-2 -filters count any listed language. Missing tags count as und. Blank filters count all streams. +Place after a conversion step. Set minimum audio and subtitle counts. +Use language filters to count only selected languages. """ name = "Stream guard" author = "Millwright" diff --git a/plugins/transcode.burn_subtitles/plugin.py b/plugins/transcode.burn_subtitles/plugin.py index 9101bb7..846efcb 100644 --- a/plugins/transcode.burn_subtitles/plugin.py +++ b/plugins/transcode.burn_subtitles/plugin.py @@ -12,32 +12,32 @@ class Inputs(BaseModel): - language: str = Field(default="eng", description="Lowercase ISO-639-2 language code, such as eng") + language: str = Field(default="eng", description="Use a lowercase ISO 639-2 language code, such as eng.") name_priority: list[str] = Field( default=["video_language", "video", "language"], - description="Name patterns in priority order: video_language=movie.eng.srt, video=movie.srt, language=eng.srt", + description="Rank filename patterns: video_language=movie.eng.srt, video=movie.srt, language=eng.srt.", ) format_priority: list[str] = Field( - default=["srt", "ass", "ssa", "vtt"], description="Subtitle extensions in priority order" + default=["srt", "ass", "ssa", "vtt"], description="Rank subtitle formats from srt, ass, ssa, and vtt." ) priority_order: Literal["name_first", "format_first"] = Field( - default="name_first", description="Rank matches by name first or format first" + default="name_first", description="Rank matches by name first or format first." ) on_multiple: Literal["first_match", "error"] = Field( - default="first_match", description="For multiple matches, choose by priority or return an error" + default="first_match", description="For multiple matches, choose by priority or report an error." ) on_missing: Literal["skip", "error"] = Field( - default="skip", description="When no sidecar matches, skip or return an error" + default="skip", description="When no sidecar matches, skip or report an error." ) - encoder: str = Field(default="libx264", description="FFmpeg video encoder used for burning subtitles") - container: Literal["mkv", "mp4"] = Field(default="mkv", description="Output container format") - quality: int = Field(default=23, ge=0, le=63, description="Video quality value passed to the encoder") + encoder: str = Field(default="libx264", description="Choose the FFmpeg video encoder for burning subtitles.") + container: Literal["mkv", "mp4"] = Field(default="mkv", description="Choose the output container.") + quality: int = Field(default=23, ge=0, le=63, description="Set the encoder's video quality value.") @field_validator("language") @classmethod def check_language(cls, value: str) -> str: if re.fullmatch(r"[a-z]{3}", value) is None: - raise ValueError("language must be lowercase ISO-639-2") + raise ValueError("Use a three-letter lowercase language code, such as eng.") return value @field_validator("name_priority") @@ -45,7 +45,9 @@ def check_language(cls, value: str) -> str: def check_names(cls, value: list[str]) -> list[str]: allowed = {"video_language", "video", "language"} if not value or len(value) != len(set(value)) or any(item not in allowed for item in value): - raise ValueError("name_priority must be nonempty, unique, and use allowed names") + raise ValueError( + "Choose at least one filename pattern without duplicates: video_language, video, or language." + ) return value @field_validator("format_priority") @@ -53,7 +55,7 @@ def check_names(cls, value: list[str]) -> list[str]: def check_formats(cls, value: list[str]) -> list[str]: allowed = {"srt", "ass", "ssa", "vtt"} if not value or len(value) != len(set(value)) or any(item not in allowed for item in value): - raise ValueError("format_priority must be nonempty, unique, and use allowed formats") + raise ValueError("Choose at least one subtitle format without duplicates: srt, ass, ssa, or vtt.") return value @@ -62,7 +64,7 @@ def check_formats(cls, value: list[str]) -> list[str]: def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: if not ctx.probe.real_video: - return {"type": "skip", "reason": "no video stream"} + return {"type": "skip", "reason": "No video stream found."} stem = PurePath(ctx.original_filename).stem candidates: list[tuple[int, int, str]] = [] seen: set[str] = set() @@ -85,11 +87,11 @@ def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: ) if len(candidates) > 1 and inputs.on_multiple == "error": matches = ", ".join(sorted(candidate[2] for candidate in candidates)) - return {"type": "error", "error": f"multiple matching subtitle sidecars: {matches}"} + return {"type": "error", "error": f"Multiple matching subtitle files: {matches}."} if not candidates: if inputs.on_missing == "error": - return {"type": "error", "error": "no matching subtitle sidecar"} - return {"type": "skip", "reason": "no matching subtitle sidecar"} + return {"type": "error", "error": "No matching subtitle file found."} + return {"type": "skip", "reason": "No matching subtitle file found."} selected = candidates[0][2] plan = FFmpegPlan( container=inputs.container, @@ -117,4 +119,4 @@ def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: ) name = inputs.name_priority[candidates[0][0]] extension = inputs.format_priority[candidates[0][1]] - return {"type": "plan", "plan": plan, "reason": f"selected {selected} ({name}, {extension} priority)"} + return {"type": "plan", "plan": plan, "reason": f"Selected {selected} ({name}, {extension} priority)."} diff --git a/plugins/transcode.burn_subtitles/plugin.toml b/plugins/transcode.burn_subtitles/plugin.toml index f06b72d..c13ac6e 100644 --- a/plugins/transcode.burn_subtitles/plugin.toml +++ b/plugins/transcode.burn_subtitles/plugin.toml @@ -2,11 +2,8 @@ id = "transcode.burn_subtitles" kind = "transcode" description = "Burn a matching sidecar subtitle into the video." help = """ -Place the sidecar beside the input and set its lowercase ISO-639-2 language code. -Rank matching names and formats by the configured priorities (default formats: SRT, ASS, SSA, VTT). -Skip files without video or with only cover art. Skip missing sidecars unless on_missing is error. -Re-encode video with the selected encoder and quality (defaults: libx264 and 23). -MP4 output drops embedded subtitle streams. +Place subtitle files beside the video. Choose a language and set filename and format priorities. +Burning subtitles re-encodes the video. MP4 output removes embedded subtitles. """ name = "Burn subtitles" author = "Millwright" diff --git a/plugins/transcode.remux/plugin.py b/plugins/transcode.remux/plugin.py index e88237f..0d58f33 100644 --- a/plugins/transcode.remux/plugin.py +++ b/plugins/transcode.remux/plugin.py @@ -12,8 +12,10 @@ class Inputs(AudioPlanInputs): - container: Literal["mkv", "mp4", "original"] = Field(default="mkv", description="Output container format") - size_ceiling_ratio: float = Field(default=1.05, gt=0, description="Maximum output-to-source file size ratio") + container: Literal["mkv", "mp4", "original"] = Field(default="mkv", description="Choose the output container.") + size_ceiling_ratio: float = Field( + default=1.05, gt=0, description="Reject output larger than this multiple of the input file size." + ) INPUTS = Inputs @@ -30,7 +32,7 @@ def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: media_changes = changes_media(plan, ctx.probe) plan = apply_metadata(plan, ctx) if not changes_media(plan, ctx.probe): - return {"type": "skip", "reason": "container and streams already satisfy the policy"} + return {"type": "skip", "reason": "The file already matches the requested settings."} reasons = change_reasons(plan, ctx.probe) # Generate missing packet timestamps during remuxing, not as a reason to remux again. plan.generate_timestamps = media_changes diff --git a/plugins/transcode.remux/plugin.toml b/plugins/transcode.remux/plugin.toml index 7d12832..272fc8e 100644 --- a/plugins/transcode.remux/plugin.toml +++ b/plugins/transcode.remux/plugin.toml @@ -2,8 +2,8 @@ id = "transcode.remux" kind = "transcode" description = "Change container or stream settings without re-encoding video." help = """ -Choose MKV (default), MP4, or the original container. Apply stream selection and optional -audio encoding. Skip files already matching the container, stream, and audio settings. +Choose MKV (default), MP4, or the original container. Select streams to keep and optionally +re-encode audio. Files already matching the requested settings are skipped. """ name = "Remux" author = "Millwright" diff --git a/plugins/transcode.video/plugin.py b/plugins/transcode.video/plugin.py index a1027a5..2aa2c39 100644 --- a/plugins/transcode.video/plugin.py +++ b/plugins/transcode.video/plugin.py @@ -13,37 +13,42 @@ class Inputs(AudioPlanInputs): - target: Literal["avc", "hevc", "av1"] = Field(default="hevc", description="Target video codec family") + target: Literal["avc", "hevc", "av1"] = Field(default="hevc", description="Choose the output video codec family.") rate_control: RateControlMode = Field( default=RateControlMode.QUALITY, - description="Quality sets encoder quality with variable size. Bitrate targets a source bitrate ratio", + description="Use encoder quality with variable output size, or target a source bitrate ratio.", ) quality: int | None = Field( - default=None, ge=0, le=63, description="Codec-specific video quality value. Used only in quality mode" + default=None, ge=0, le=63, description="Set the codec-specific quality value. Applies only in quality mode." ) ratio: float | None = Field( default=None, gt=0, le=1, - description="Target video bitrate as a fraction of the source. Used only in bitrate mode", + description="Set video bitrate as a fraction of the source. Applies only in bitrate mode.", ) - encoder: str | None = Field(default=None, description="FFmpeg video encoder override") + encoder: str | None = Field(default=None, description="Override the FFmpeg video encoder.") encoder_priority: list[str] = Field( default=[], - description="Encoder priority order. Override wins. Empty uses the software default", + description=( + "List encoders in priority order. The encoder override wins. Leave blank for the software default." + ), ) - pixel_format: str | None = Field(default=None, description="Pixel format for the encoded video") - ten_bit: bool = Field(default=False, description="Encode video with 10-bit precision") - preset: str | None = Field(default=None, description="Encoder speed and compression preset") - max_width: int | None = Field(default=None, ge=1, description="Maximum output video width in pixels") - max_height: int | None = Field(default=None, ge=1, description="Maximum output video height in pixels") - container: Literal["mkv", "mp4"] = Field(default="mkv", description="Output container format") + pixel_format: str | None = Field(default=None, description="Set the encoded video's pixel format.") + ten_bit: bool = Field(default=False, description="Encode video with 10-bit precision.") + preset: str | None = Field(default=None, description="Choose the encoder's speed and compression preset.") + max_width: int | None = Field(default=None, ge=1, description="Limit output video width in pixels.") + max_height: int | None = Field(default=None, ge=1, description="Limit output video height in pixels.") + container: Literal["mkv", "mp4"] = Field(default="mkv", description="Choose the output container.") bitrate_source: Literal["video", "container", "average"] = Field( - default="video", description="Bitrate mode source: video, container, or their average when both are available" + default="video", + description="In bitrate mode, choose video bitrate, whole-file bitrate, or their average.", + ) + size_ceiling_ratio: float = Field( + default=1.05, gt=0, description="Reject output larger than this multiple of the input file size." ) - size_ceiling_ratio: float = Field(default=1.05, gt=0, description="Maximum output-to-source file size ratio") - fix_timestamps: bool = Field(default=False, description="Generate missing timestamps for AVI and TS inputs") - extras: list[str] = Field(default=[], description="Additional FFmpeg arguments, one argument per entry") + fix_timestamps: bool = Field(default=False, description="Generate missing timestamps for AVI and TS inputs.") + extras: list[str] = Field(default=[], description="Add FFmpeg arguments, one argument per entry.") INPUTS = Inputs @@ -52,7 +57,7 @@ class Inputs(AudioPlanInputs): def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: video = ctx.probe.real_video if not video: - return {"type": "skip", "reason": "no video stream"} + return {"type": "skip", "reason": "No video stream found."} families = {"avc": {"avc", "h264", "avc1"}, "hevc": {"hevc", "h265"}, "av1": {"av1", "av01"}} defaults = {"avc": ("libx264", 0.75, 23), "hevc": ("libx265", 0.5, 21), "av1": ("libsvtav1", 0.45, 30)} encoder, ratio, quality = defaults[inputs.target] @@ -116,7 +121,7 @@ def decide(ctx: PluginContext, inputs: Inputs) -> dict[str, Any]: plan.video = None plan = apply_metadata(plan, ctx) if not changes_media(plan, ctx.probe): - return {"type": "skip", "reason": f"already {inputs.target}"} + return {"type": "skip", "reason": f"The file already matches the {inputs.target.upper()} settings."} reasons = change_reasons(plan, ctx.probe) return { "type": "plan", diff --git a/plugins/transcode.video/plugin.toml b/plugins/transcode.video/plugin.toml index 87e243d..7db8cba 100644 --- a/plugins/transcode.video/plugin.toml +++ b/plugins/transcode.video/plugin.toml @@ -1,13 +1,9 @@ id = "transcode.video" kind = "transcode" -description = "Encode AVC, HEVC, or AV1 with quality or bitrate rate control." +description = "Convert video to AVC, HEVC, or AV1." help = """ -Quality mode (default) uses encoder quality without a bitrate target. Bitrate mode uses -a source bitrate ratio without quality flags. The output size limit applies in both modes. -Optionally set encoder, pixel format, bit depth, preset, or maximum dimensions. -Skip matching files unless another video, audio, stream, or container change is requested. -Skip files without video or with only cover art. HEVC defaults to libx265, quality 21, -and bitrate ratio 0.5. AVC uses libx264. AV1 uses libsvtav1. Both have codec-specific defaults. +Choose quality or bitrate mode. Set an encoder or an encoder priority list if needed. +Files already matching your settings are skipped. """ name = "Video transcode" author = "Millwright" diff --git a/run.py b/run.py index bd8d4be..3f0774c 100644 --- a/run.py +++ b/run.py @@ -254,18 +254,12 @@ def node_environment(config_path: Path, *, port: int) -> dict[str, str]: var = f"MW_NODE_{key.upper()}" if var not in env and key not in settings: env[var] = value - if config.node_enrollment_token is not None and ( - "MW_NODE_ENROLLMENT_TOKEN" in env or "enrollment_token" not in settings - ): - # Shared file references resolve beside the server's selected config, - # which can differ from the node's config directory in all-in-one mode. - env["MW_NODE_ENROLLMENT_TOKEN"] = config.node_enrollment_token return env def healthcheck(command: Sequence[str], *, urlopen: Urlopen = urllib.request.urlopen) -> int: # A standalone node has no HTTP endpoint. Its PID 1 is the node process, - # so reading that command confirms it is running without testing enrollment. + # so reading that command confirms it is running without requiring registration. if any(Path(part).name == NODE_BIN for part in command[:2]): return 0 parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) diff --git a/server/pyproject.toml b/server/pyproject.toml index 9ea6fba..73cd419 100644 --- a/server/pyproject.toml +++ b/server/pyproject.toml @@ -6,7 +6,7 @@ requires-python = ">=3.14" dependencies = [ "millwright-common", "pydantic>=2.12", - "sqlalchemy>=2.0.45", + "sqlalchemy[asyncio]>=2.1", "psycopg[binary]>=3.3", "watchfiles>=1.2.0", "fastapi>=0.115", diff --git a/server/src/server/api/app.py b/server/src/server/api/app.py index ff55a73..1c951a4 100644 --- a/server/src/server/api/app.py +++ b/server/src/server/api/app.py @@ -304,8 +304,6 @@ def create_app( root.mkdir(parents=True, exist_ok=True) hub_kwargs.setdefault("timezone", ZoneInfo(config.tz)) hub_kwargs.setdefault("min_free_gb", operational.min_free_gb) - if config.node_enrollment_token is not None: - hub_kwargs.setdefault("enrollment_token", config.node_enrollment_token) hub = hub if hub is not None else NodeHub(sessionmaker, events=events, **hub_kwargs) if hub.events is None: hub.events = events @@ -453,6 +451,7 @@ async def prune_history() -> None: scheduler.add(partial(cleanup_trash, sessionmaker, root), every=3600, name="trash") def apply_settings(value: OperationalSettings) -> None: + hub.set_registration_enabled(value.node_registration_enabled) healthchecks.configure( max_concurrency=value.max_healthchecks, hwaccel=value.healthcheck_hwaccel, @@ -697,7 +696,6 @@ async def lifespan(_app: FastAPI): value for value in ( config.api_key, - config.node_enrollment_token, config.auth.oidc.secret if config.auth.oidc is not None else None, ) if value diff --git a/server/src/server/api/backups.py b/server/src/server/api/backups.py index 2b40073..306fe63 100644 --- a/server/src/server/api/backups.py +++ b/server/src/server/api/backups.py @@ -45,7 +45,7 @@ from server.workspace import recovery_jobs, wait_workspace_io if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncGenerator from sqlalchemy.ext.asyncio import async_sessionmaker @@ -64,7 +64,7 @@ class _RestoreOutcome: @asynccontextmanager -async def _pause_dry_runs(request: Request) -> AsyncIterator[None]: +async def _pause_dry_runs(request: Request) -> AsyncGenerator[None]: try: async with request.app.state.dry_runs.pause(): yield diff --git a/server/src/server/api/files.py b/server/src/server/api/files.py index d21abe2..f07a69c 100644 --- a/server/src/server/api/files.py +++ b/server/src/server/api/files.py @@ -8,7 +8,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response, status from pydantic import BaseModel, ConfigDict, Field, PositiveInt, field_validator from sqlalchemy import case, delete, or_, select -from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import JSONB, distinct_on from sqlalchemy.orm import load_only, with_expression from server.api.auth import PrincipalDep, current_principal @@ -220,7 +220,7 @@ async def get_files( Job.kind == "transcode", JobAttempt.result == "succeeded", ) - .distinct(Job.media_file_id) + .ext(distinct_on(Job.media_file_id)) .order_by(Job.media_file_id, JobAttempt.id.desc()) .options( load_only( diff --git a/server/src/server/api/nodes.py b/server/src/server/api/nodes.py index 02fb59c..f0daee3 100644 --- a/server/src/server/api/nodes.py +++ b/server/src/server/api/nodes.py @@ -9,7 +9,7 @@ from server.api.auth import current_principal from server.models import NodeRow from server.nodehub.store import list_nodes -from server.queue.store import NodeRemovalError, remove_node +from server.queue.store import NodeRemovalError from server.serialize import iso_utc if TYPE_CHECKING: @@ -67,6 +67,19 @@ async def get_nodes(request: Request) -> list[dict]: ] +@router.get("/registration", summary="Read live node registration requests") +async def get_registration(request: Request) -> dict: + return request.app.state.hub.registration_status() + + +@router.post("/pending/{request_id}/claim", status_code=status.HTTP_202_ACCEPTED, summary="Claim a pending node") +async def claim_node(request_id: str, request: Request) -> None: + async with request.app.state.reset_lock: + if not request.app.state.hub.claim_node(request_id): + raise HTTPException(status_code=404, detail="Pending node is no longer connected.") + LOG.info("Node registration request %s claimed.", request_id) + + @router.patch("/{node_id}", summary="Update a node's workload limits, name, or paused state") async def patch_node(node_id: int, body: NodePatchBody, request: Request) -> dict: db: async_sessionmaker = request.app.state.db @@ -101,14 +114,12 @@ async def patch_node(node_id: int, body: NodePatchBody, request: Request) -> dic return _node_dict(node, entry.in_flight if entry is not None else {}) -@router.delete("/{node_id}", status_code=status.HTTP_204_NO_CONTENT, summary="Remove an offline node") +@router.delete("/{node_id}", status_code=status.HTTP_204_NO_CONTENT, summary="Unregister a node without active work") async def delete_node(node_id: int, request: Request) -> None: - db: async_sessionmaker = request.app.state.db - async with db() as session: - try: - removed = await remove_node(session, node_id) - if not removed: - raise HTTPException(status_code=404, detail="node not found") - await session.commit() - except NodeRemovalError as exc: - raise HTTPException(status_code=409, detail=str(exc)) from exc + try: + async with request.app.state.reset_lock: + removed = await request.app.state.hub.remove_node(node_id) + if not removed: + raise HTTPException(status_code=404, detail="node not found") + except NodeRemovalError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc diff --git a/server/src/server/backup.py b/server/src/server/backup.py index 26a54d1..fe6e28f 100644 --- a/server/src/server/backup.py +++ b/server/src/server/backup.py @@ -434,7 +434,8 @@ async def restore_data(session: AsyncSession, data: dict, selection: dict) -> No raw = next(item for item in _items(data, "libraries") if item.get("name") == name) flow_name = raw.get("flow_name") flow_row = flow_rows[flow_name] if flow_name else None - fields = body.model_dump(exclude={"name"}) | {"flow_id": flow_row.id if flow_row else None} + fields = body.model_dump(exclude={"name"}) + fields["flow_id"] = flow_row.id if flow_row else None if name in existing_libraries: await update_library(session, existing_libraries[name].id, fields) else: @@ -534,7 +535,7 @@ def export_config(path: Path) -> str | None: value = oidc.get(key) if isinstance(value, str) and value.startswith("file:"): oidc[key] = resolve_secret(value, base_dir=path.parent) - for key in ("api_key", "node_enrollment_token", "database_url"): + for key in ("api_key", "database_url"): value = doc.get(key) if isinstance(value, str) and value.startswith("file:"): doc[key] = resolve_secret(value, base_dir=path.parent) @@ -548,10 +549,7 @@ def validate_config(content: str) -> None: isinstance(oidc.get(key), str) and oidc[key].startswith("file:") for key in ("client_id", "secret") ): raise ValueError("backup config contains unresolved secret references") - if any( - isinstance(data.get(key), str) and data[key].startswith("file:") - for key in ("node_enrollment_token", "api_key", "database_url") - ): + if any(isinstance(data.get(key), str) and data[key].startswith("file:") for key in ("api_key", "database_url")): raise ValueError("backup config contains unresolved secret references") ServerConfig.model_validate(data) @@ -561,7 +559,7 @@ def prepared_config(path: Path, content: str) -> tuple[str, dict[str, str]]: doc = tomlkit.parse(content) values = {} generation = secrets.token_hex(16) - for key in ("api_key", "node_enrollment_token", "database_url"): + for key in ("api_key", "database_url"): value = doc.get(key) if isinstance(value, str): name = f"{generation}-{key}" diff --git a/server/src/server/config.example.toml b/server/src/server/config.example.toml index fda8e47..3dbec59 100644 --- a/server/src/server/config.example.toml +++ b/server/src/server/config.example.toml @@ -17,9 +17,6 @@ # Instance API key for CLI commands and password resets. Keep it private. api_key = "file:secrets/api_key" -# Node enrollment secret. It does not grant API access. Keep it separate from api_key. -node_enrollment_token = "file:secrets/node_enrollment_token" - # IANA timezone for schedules, calendar statistics, and displayed dates. # Defaults to $TZ, then UTC. # tz = "Asia/Kuwait" diff --git a/server/src/server/core/config.py b/server/src/server/core/config.py index 429ed48..f4c1aff 100644 --- a/server/src/server/core/config.py +++ b/server/src/server/core/config.py @@ -61,8 +61,6 @@ class ServerConfig(BaseModel): enable_docs: bool = False workspace_dir: str | None = None trash_dir: str | None = None - # Shared enrollment token for /api/hub. Absent keeps auto-enrollment off. - node_enrollment_token: str | None = None api_key: str | None = None # None uses the state directory's uploaded_plugins subdirectory. uploaded_plugins_dir: str | None = None @@ -145,8 +143,6 @@ def load_config(path: Path) -> ServerConfig: for key in ("client_id", "secret"): if isinstance(oidc.get(key), str): oidc[key] = resolve_secret(oidc[key], base_dir=path.parent) - if isinstance(data.get("node_enrollment_token"), str): - data["node_enrollment_token"] = resolve_secret(data["node_enrollment_token"], base_dir=path.parent) if isinstance(data.get("api_key"), str): data["api_key"] = resolve_secret(data["api_key"], base_dir=path.parent) if isinstance(data.get("database_url"), str): @@ -176,16 +172,15 @@ def ensure_config(path: Path) -> None: if secret_dir.is_symlink() or stat.S_IMODE(secret_dir.stat().st_mode) != 0o700: msg = f"secrets directory must be private: {secret_dir}" raise ValueError(msg) - for name in ("api_key",) if existing else ("api_key", "node_enrollment_token"): - secret_path = secret_dir / name - if secret_path.is_symlink(): - msg = f"secret file must not be a symlink: {secret_path}" + secret_path = secret_dir / "api_key" + if secret_path.is_symlink(): + msg = f"secret file must not be a symlink: {secret_path}" + raise ValueError(msg) + if secret_path.exists(): + if not secret_path.is_file() or stat.S_IMODE(secret_path.stat().st_mode) != 0o600: + msg = f"secret file must be private: {secret_path}" raise ValueError(msg) - if secret_path.exists(): - if not secret_path.is_file() or stat.S_IMODE(secret_path.stat().st_mode) != 0o600: - msg = f"secret file must be private: {secret_path}" - raise ValueError(msg) - continue + else: fd = os.open(secret_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600) with os.fdopen(fd, "w", encoding="utf-8") as output: output.write(secrets.token_urlsafe(32) + "\n") diff --git a/server/src/server/core/logging.py b/server/src/server/core/logging.py index 2c63af8..4ece5ed 100644 --- a/server/src/server/core/logging.py +++ b/server/src/server/core/logging.py @@ -21,7 +21,7 @@ def filter(self, record: logging.LogRecord) -> bool: class _PlaybackStreamFilter(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: - if record.exc_info is not None and isinstance(record.exc_info[1], PlaybackStreamClosed): + if record.exc_info and isinstance(record.exc_info[1], PlaybackStreamClosed): # Uvicorn must still catch this exception and close the truncated transport. record.msg = "Playback response aborted: %s" record.args = (str(record.exc_info[1]),) diff --git a/server/src/server/external_notifications/__init__.py b/server/src/server/external_notifications/__init__.py index 90103f4..13ce734 100644 --- a/server/src/server/external_notifications/__init__.py +++ b/server/src/server/external_notifications/__init__.py @@ -1,10 +1,13 @@ from __future__ import annotations +from pathlib import PurePath, PurePosixPath, PureWindowsPath from typing import TYPE_CHECKING from common.actions import NotifyExternal from server.external_notifications.arr import notify_arr +from server.external_notifications.media_server import notify_media_server +from server.external_notifications.plex import notify_plex if TYPE_CHECKING: from collections.abc import Awaitable, Callable @@ -12,6 +15,24 @@ async def notify_external(payload: dict, save_progress: Callable[[dict], Awaitable[None]] | None = None) -> bool: action = NotifyExternal.model_validate(payload["action"]) + path = PurePath(payload["path"]) + if action.path_mappings: + matches = [mapping for mapping in action.path_mappings if path.parent.is_relative_to(mapping.source)] + if not matches: + message = f"No path mapping matches {str(path.parent)!r}." + raise ValueError(message) + mapping = max(matches, key=lambda item: len(PurePath(item.source).parts)) + relative = path.relative_to(mapping.source) + destination = ( + PureWindowsPath(mapping.destination) + if PureWindowsPath(mapping.destination).is_absolute() + else PurePosixPath(mapping.destination) + ) + path = destination.joinpath(*relative.parts) match action.application: case "radarr" | "sonarr": - return await notify_arr(action, payload, save_progress) + return await notify_arr(action, payload, str(path), save_progress) + case "plex" | "jellyfin" | "emby": + if action.application == "plex": + return await notify_plex(action, str(path.parent)) + return await notify_media_server(action, str(path.parent)) diff --git a/server/src/server/external_notifications/arr.py b/server/src/server/external_notifications/arr.py index c62be39..80aff65 100644 --- a/server/src/server/external_notifications/arr.py +++ b/server/src/server/external_notifications/arr.py @@ -49,9 +49,9 @@ class _Progress(BaseModel): async def notify_arr( - action: NotifyExternal, payload: dict, save_progress: Callable[[dict], Awaitable[None]] | None = None + action: NotifyExternal, payload: dict, path: str, save_progress: Callable[[dict], Awaitable[None]] | None = None ) -> bool: - path, filename = payload["path"], payload["original_filename"] + filename = payload["original_filename"] async with httpx.AsyncClient( base_url=str(action.url).rstrip("/") + "/api/v3/", headers={"X-Api-Key": action.api_key}, @@ -64,7 +64,7 @@ async def notify_arr( response.raise_for_status() progress = _Command.model_validate(response.json()) elif action.application == "radarr": - identifiers = {kind.lower(): value for kind, value in _IDENTIFIER.findall(path + "/" + filename)} + identifiers = {kind.lower(): value for kind, value in _IDENTIFIER.findall(payload["path"] + "/" + filename)} if identifiers: params = {"tmdbId": identifiers["tmdb"]} if "tmdb" in identifiers else {} response = await client.get("movie", params=params) diff --git a/server/src/server/external_notifications/media_server.py b/server/src/server/external_notifications/media_server.py new file mode 100644 index 0000000..63eb416 --- /dev/null +++ b/server/src/server/external_notifications/media_server.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import httpx + +if TYPE_CHECKING: + from common.actions import NotifyExternal + + +async def notify_media_server(action: NotifyExternal, folder: str) -> bool: + headers = ( + {"Authorization": f'MediaBrowser Token="{action.api_key}"'} + if action.application == "jellyfin" + else {"X-Emby-Token": action.api_key} + ) + async with httpx.AsyncClient(base_url=str(action.url).rstrip("/") + "/", headers=headers, timeout=15) as client: + response = await client.post( + "Library/Media/Updated", json={"Updates": [{"Path": folder, "UpdateType": "Modified"}]} + ) + response.raise_for_status() + return True diff --git a/server/src/server/external_notifications/plex.py b/server/src/server/external_notifications/plex.py new file mode 100644 index 0000000..5082762 --- /dev/null +++ b/server/src/server/external_notifications/plex.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from pathlib import PurePosixPath, PureWindowsPath +from typing import TYPE_CHECKING + +import httpx +from pydantic import BaseModel, Field + +if TYPE_CHECKING: + from common.actions import NotifyExternal + + +class _Location(BaseModel): + path: str + + +class _Section(BaseModel): + key: int = Field(gt=0) + locations: list[_Location] = Field(default=[], alias="Location") + + +class _Sections(BaseModel): + sections: list[_Section] = Field(default=[], alias="Directory") + + +class _Library(BaseModel): + container: _Sections = Field(alias="MediaContainer") + + +async def notify_plex(action: NotifyExternal, folder: str) -> bool: + async with httpx.AsyncClient( + base_url=str(action.url).rstrip("/") + "/", + headers={ + "X-Plex-Token": action.api_key, + "X-Plex-Client-Identifier": "millwright", + "X-Plex-Product": "Millwright", + "Accept": "application/json", + }, + timeout=15, + ) as client: + response = await client.get("library/sections") + response.raise_for_status() + library = _Library.model_validate(response.json()) + path_type = PureWindowsPath if PureWindowsPath(folder).is_absolute() else PurePosixPath + path = path_type(folder) + sections = [ + section + for section in library.container.sections + if any(path.is_relative_to(path_type(location.path)) for location in section.locations) + ] + if not sections: + message = f"No Plex library contains {folder!r}. Check the path mappings and Plex library folders." + raise ValueError(message) + for section in sections: + response = await client.get(f"library/sections/{section.key}/refresh", params={"path": folder}) + response.raise_for_status() + return True diff --git a/server/src/server/flows/batch.py b/server/src/server/flows/batch.py index 2df4bce..4123c81 100644 --- a/server/src/server/flows/batch.py +++ b/server/src/server/flows/batch.py @@ -30,7 +30,7 @@ from server.workspace import wait_workspace_io if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncGenerator, Callable from common.media.probe import ProbeData from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -557,7 +557,7 @@ async def cancel(self, report: BatchReport) -> None: await _drain(report.task) @asynccontextmanager - async def pause(self) -> AsyncIterator[None]: + async def pause(self) -> AsyncGenerator[None]: self._paused += 1 try: if self.active is not None: diff --git a/server/src/server/nodehub/hub.py b/server/src/server/nodehub/hub.py index 10accc5..557d79b 100644 --- a/server/src/server/nodehub/hub.py +++ b/server/src/server/nodehub/hub.py @@ -31,8 +31,7 @@ ActionTask, AttemptMissing, Cancel, - Enrolled, - EnrollmentSaved, + CredentialsSaved, Error, Finished, Heartbeat, @@ -42,6 +41,10 @@ NodeLogs, OfferedJob, Progress, + Registered, + RegistrationPending, + Unregister, + Unregistered, Welcome, WorkspaceChallenge, WorkspaceResponse, @@ -72,12 +75,14 @@ from server.nodehub.store import get_node, hash_token, register_node, set_node_state, touch_node from server.queue.store import ( TERMINAL_STATUSES, + NodeRemovalError, claim, current_attempt, finish, finish_artifact_check, get_job, release_transcode_decision, + remove_node, requeue, transition, ) @@ -221,6 +226,14 @@ class NodeEntry: offered_at: dict[int, datetime] = field(default_factory=dict) last_progress_at: dict[int, datetime] = field(default_factory=dict) disconnected: asyncio.Event = field(default_factory=asyncio.Event) + unregistered: asyncio.Event = field(default_factory=asyncio.Event) + + +@dataclass(kw_only=True, slots=True) +class PendingNode: + request_id: str + capabilities: NodeCapabilities + approved: asyncio.Future[bool] async def probe_for_file(session: AsyncSession, media_file: MediaFile) -> ProbeData | None: @@ -255,7 +268,6 @@ def __init__( dispatch_interval_s: float = 1.0, sweep_interval_s: float = 1.0, hello_timeout_s: float = 10.0, - enrollment_token: str | None = None, replace_original: bool = True, action_task_timeout_s: float = 30.0, ffprobe_bin: str = "ffprobe", @@ -289,7 +301,8 @@ def __init__( self.dispatch_interval_s = dispatch_interval_s self.sweep_interval_s = sweep_interval_s self.hello_timeout_s = hello_timeout_s - self.enrollment_token = enrollment_token + self.registration_enabled = True + self.pending_nodes: dict[str, PendingNode] = {} self.replace_original = replace_original self.action_task_timeout_s = action_task_timeout_s self.ffprobe_bin = ffprobe_bin @@ -335,13 +348,21 @@ async def serve_connection(self, conn: NodeConnection) -> None: hello = await self._read_hello(conn) if hello is None: return - node_name = hello.node_id or "enrolling node" + node_name = hello.node_id or "unregistered node" + claimed = False + if await self._authenticate_hello(hello) is None: + if hello.node_id is not None or hello.token: + await self._unregister_connection(conn, node_name) + return + claimed = await self._wait_for_claim(conn, hello) + if not claimed: + return while True: await self._reset_done.wait() async with self.dispatch_lock: if self._reset_active: continue - handshake = await self._handshake_locked(conn, hello) + handshake = await self._handshake_locked(conn, hello, claimed=claimed) if handshake is None: return hello, node_db_id, limits, paused = handshake @@ -895,7 +916,7 @@ async def ready() -> tuple[NodeEntry | None] | None: self.staging.busy_jobs.discard(job_id) @asynccontextmanager - async def _dispatch_candidates(self) -> AsyncIterator[AsyncIterator[int]]: + async def _dispatch_candidates(self) -> AsyncGenerator[AsyncIterator[int]]: async with self.sessionmaker() as session: upper = await session.scalar(select(func.max(Job.id))) or 0 eligible = ( @@ -1649,6 +1670,7 @@ def start(self) -> None: def begin_shutdown(self) -> None: if self._close_task is None: + self.set_registration_enabled(False) tasks = [*self._tasks, *self._action_waits] for task in tasks: task.cancel() @@ -1836,41 +1858,125 @@ async def _read_hello(self, conn: NodeConnection) -> Hello | None: LOG.warning("Expected node hello in workspace %s, received %s", self.workspace_dir, type(msg).__name__) await self._reject(conn, "handshake", "expected a hello message") return None - if await self._authenticate_hello(conn, msg) is None: - return None return msg - async def _authenticate_hello(self, conn: NodeConnection, msg: Hello) -> tuple[int | None, bool] | None: - node_db_id: int | None = None + async def _authenticate_hello(self, msg: Hello) -> int | None: async with self.sessionmaker() as session: node = await get_node(session, msg.node_id) if msg.node_id is not None else None - known = node is not None and hmac.compare_digest(node.token_hash, hash_token(msg.token)) - enrolled = self.enrollment_token is not None and hmac.compare_digest( - hash_token(msg.token), hash_token(self.enrollment_token) - ) - if not (known or (enrolled and msg.node_id is None)): - if node is None: - LOG.error("unknown node %s rejected", msg.node_id) - await self._reject(conn, "auth", "unknown node", 4401) + if node is not None and hmac.compare_digest(node.token_hash, hash_token(msg.token)): + return node.id + return None + + async def remove_node(self, node_id: int) -> bool: + async with self.dispatch_lock: + entry = next((item for item in self.entries.values() if item.node_db_id == node_id), None) + if entry is not None and entry.in_flight: + raise NodeRemovalError("Node has active work. Wait for it to finish before removing the node.") + async with self.sessionmaker() as session: + if not await remove_node(session, node_id): + return False + await session.commit() + if entry is not None: + entry.workspace_ready = False + entry.paused = True + if entry is not None and not entry.disconnected.is_set(): + await self._unregister_connection(entry.conn, entry.name, entry.unregistered) + return True + + async def _unregister_connection( + self, conn: NodeConnection, name: str, acknowledged: asyncio.Event | None = None + ) -> None: + try: + async with asyncio.timeout(self.hello_timeout_s): + await conn.send_json(Unregister().model_dump(mode="json")) + if acknowledged is not None: + await acknowledged.wait() else: - LOG.error("authentication failed for node %s", msg.node_id) - await self._reject(conn, "auth", "authentication failed", 4401) - return None - if node is not None: - node_db_id = node.id - return node_db_id, enrolled + data = await conn.recv_json() + if data is None or not isinstance(parse_message(data), Unregistered): + LOG.warning("Node %s did not acknowledge credential removal.", name) + except Exception as exc: + LOG.warning( + "Could not confirm credential removal for node %s: %s", name, error_message(exc) or type(exc).__name__ + ) + finally: + await asyncio.wait_for(conn.close(), self.hello_timeout_s) + + def set_registration_enabled(self, enabled: bool) -> None: + if self.registration_enabled == enabled: + return + self.registration_enabled = enabled + if not enabled: + for pending in self.pending_nodes.values(): + if not pending.approved.done(): + pending.approved.set_result(False) + self.pending_nodes.clear() + self._publish_registration() + + def _publish_registration(self) -> None: + if self.events is not None: + self.events.publish("node.registration", self.registration_status()) + + def registration_status(self) -> dict: + return { + "enabled": self.registration_enabled, + "pending": [ + {"request_id": item.request_id, "capabilities": item.capabilities.model_dump(mode="json")} + for item in self.pending_nodes.values() + ], + } + + def claim_node(self, request_id: str) -> bool: + pending = self.pending_nodes.pop(request_id, None) + if pending is None or pending.approved.done(): + return False + pending.approved.set_result(True) + self._publish_registration() + return True + + async def _wait_for_claim(self, conn: NodeConnection, msg: Hello) -> bool: + if not self.registration_enabled: + await self._reject(conn, "registration", "Node registration is disabled.") + return False + request_id = secrets.token_hex(8) + pending = PendingNode( + request_id=request_id, + capabilities=msg.capabilities, + approved=asyncio.get_running_loop().create_future(), + ) + self.pending_nodes[request_id] = pending + disconnected: asyncio.Task[dict | None] | None = None + try: + await conn.send_json(RegistrationPending(request_id=request_id).model_dump(mode="json")) + self._publish_registration() + disconnected = asyncio.create_task(conn.recv_json()) + await asyncio.wait((pending.approved, disconnected), return_when=asyncio.FIRST_COMPLETED) + if disconnected.done(): + disconnected.result() + return False + if not pending.approved.result(): + await self._reject(conn, "registration", "Node registration is disabled.") + return False + return True + finally: + self.pending_nodes.pop(request_id, None) + self._publish_registration() + if disconnected is not None: + disconnected.cancel() + with contextlib.suppress(asyncio.CancelledError): + await disconnected async def _handshake_locked( - self, conn: NodeConnection, msg: Hello + self, conn: NodeConnection, msg: Hello, *, claimed: bool = False ) -> tuple[Hello, int, dict[str, int], bool] | None: # A reset may have changed node credentials while this connection waited for its dispatch lock. - authentication = await self._authenticate_hello(conn, msg) - if authentication is None: + node_db_id = None if claimed else await self._authenticate_hello(msg) + if node_db_id is None and not claimed: + await self._unregister_connection(conn, msg.node_id or "unregistered node") return None - node_db_id, enrolled = authentication - issued_token = secrets.token_urlsafe(32) if enrolled else msg.token + issued_token = secrets.token_urlsafe(32) if claimed else msg.token challenge_id = secrets.token_hex(16) - directory = f"nodes/{node_db_id if node_db_id is not None else 'enrollment'}/{challenge_id}" + directory = f"nodes/{node_db_id if node_db_id is not None else 'registration'}/{challenge_id}" relative_response = f"{directory}/response" challenge = secrets.token_bytes(32) async with storage_scope(): @@ -1918,7 +2024,7 @@ async def _handshake_locked( finally: await storage_io(_cleanup_challenge, self.workspace_dir, directory) async with self.sessionmaker() as session: - if enrolled: + if claimed: node_name = f"node-{secrets.token_hex(8)}" node = await register_node( session, @@ -1928,7 +2034,7 @@ async def _handshake_locked( display_name=_hardware_name(msg.capabilities), ) node_db_id = node.id - await conn.send_json(Enrolled(node_id=node_name, token=issued_token).model_dump(mode="json")) + await conn.send_json(Registered(node_id=node_name, token=issued_token).model_dump(mode="json")) cause = None try: data = await asyncio.wait_for(conn.recv_json(), timeout=self.hello_timeout_s) @@ -1936,20 +2042,20 @@ async def _handshake_locked( except (TimeoutError, ValueError) as exc: cause = error_message(exc) saved = None - if not isinstance(saved, EnrollmentSaved) or saved.node_id != node_name: + if not isinstance(saved, CredentialsSaved) or saved.node_id != node_name: if cause is None: cause = ( f"Acknowledgement named node {saved.node_id} instead of {node_name}." - if isinstance(saved, EnrollmentSaved) - else f"Expected EnrollmentSaved, got {type(saved).__name__ if saved else 'no acknowledgement'}." + if isinstance(saved, CredentialsSaved) + else f"Expected CredentialsSaved, got {type(saved).__name__ if saved else 'no acknowledgement'}." ) LOG.warning( - "Could not confirm credential storage for enrolling node %s in workspace %s", + "Could not confirm credential storage for node %s in workspace %s", node_name, self.workspace_dir, extra={"error": cause}, ) - await self._reject(conn, "enrollment", "node did not save credentials") + await self._reject(conn, "registration", "Node did not save credentials.") return None msg = msg.model_copy(update={"node_id": node_name, "token": issued_token}) assert msg.node_id is not None @@ -1960,7 +2066,7 @@ async def _handshake_locked( .returning(NodeRow.id) ) if current is None: - await self._reject(conn, "auth", "authentication revoked during workspace verification", 4401) + await self._unregister_connection(conn, msg.node_id) return None node = await get_node(session, msg.node_id) assert node is not None @@ -2022,12 +2128,12 @@ async def _handshake_locked( return msg, node_db_id, limits, paused def _log_workspace_failure(self, msg: Hello, cause: str) -> None: - credentials = (msg.token, self.enrollment_token or "") + credentials = (msg.token,) LOG.warning( "Workspace verification failed for node %s. Resolved server workspace: %s. " "Check server workspace_dir/MW_WORKSPACE_DIR and node workspace_dir/MW_NODE_WORKSPACE_DIR. " "Both must access the same shared files. Paths can differ between hosts.", - redact(msg.node_id or "enrollment", secrets=credentials), + redact(msg.node_id or "registration", secrets=credentials), redact(str(self.workspace_dir), secrets=credentials), extra={"error": redact(cause, secrets=credentials)}, ) @@ -2046,6 +2152,9 @@ async def _receive_loop(self, entry: NodeEntry) -> None: Error(code="protocol", message="unparseable message").model_dump(mode="json") ) continue + if isinstance(msg, Unregistered): + entry.unregistered.set() + return if isinstance(msg, Heartbeat): await self._handle_heartbeat(entry, msg) elif isinstance(msg, AttemptMissing): @@ -3372,7 +3481,7 @@ def wake(self) -> None: self._wake.set() @asynccontextmanager - async def pause_for_reset(self) -> AsyncIterator[None]: + async def pause_for_reset(self) -> AsyncGenerator[None]: async with self.dispatch_lock: if self._close_task is not None: raise ValueError("server is shutting down") diff --git a/server/src/server/notifications.py b/server/src/server/notifications.py index a4dfb37..fa9688a 100644 --- a/server/src/server/notifications.py +++ b/server/src/server/notifications.py @@ -30,7 +30,7 @@ from server.workspace import wait_workspace_io if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Generator from apprise import Apprise from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -143,7 +143,7 @@ def _url_secrets(url: str) -> tuple[str, ...]: @contextmanager -def _provider_call(url: str) -> Iterator[None]: +def _provider_call(url: str) -> Generator[None]: tokens = ( _PROVIDER_CALL.set(True), _PROVIDER_SECRETS.set(_url_secrets(url)), diff --git a/server/src/server/probing.py b/server/src/server/probing.py index 0c77eb2..49f406d 100644 --- a/server/src/server/probing.py +++ b/server/src/server/probing.py @@ -400,7 +400,7 @@ async def _claim( Job.media_file_id.not_in(busy_files), eligible, ) - deadline = MediaFile.first_seen + literal(timedelta(seconds=1)) * Library.cooldown_s + deadline = MediaFile.first_seen + Library.cooldown_s * literal(timedelta(seconds=1)) cooling = and_( MediaFile.first_seen.is_not(None), Library.mkv_stats.is_(True), @@ -878,7 +878,7 @@ async def _remaining(self) -> dict[int, int]: .group_by(Library.id) .order_by(Library.id) ) - return dict(rows.tuples().all()) + return dict(rows.all()) def _pending_status(self, library_id: int) -> str: if library_id in self._task_libraries.values(): diff --git a/server/src/server/queue/store.py b/server/src/server/queue/store.py index 76c9e20..11cd8c6 100644 --- a/server/src/server/queue/store.py +++ b/server/src/server/queue/store.py @@ -218,12 +218,13 @@ async def remove_node(session: AsyncSession, node_id: int) -> bool: node = await session.get(NodeRow, node_id, with_for_update=True) if node is None: return False - if node.state != "offline": - msg = "only offline nodes may be removed" - raise NodeRemovalError(msg) - active = await session.scalar(select(Job.id).where(Job.node_id == node_id, Job.status.in_(_ACTIVE_VALUES)).limit(1)) + active_jobs = select(Job.id).where( + Job.node_id == node_id, + or_(Job.status.in_(_ACTIVE_VALUES), Job.id.in_(post_action_jobs())), + ) + active = await session.scalar(active_jobs.limit(1)) if active is not None: - msg = "node owns an active job" + msg = "Node has active work. Wait for it to finish before removing the node." raise NodeRemovalError(msg) removed = await session.scalar( update(NodeRow) @@ -231,14 +232,13 @@ async def remove_node(session: AsyncSession, node_id: int) -> bool: NodeRow.id == node_id, NodeRow.name == node.name, NodeRow.token_hash == node.token_hash, - NodeRow.state == "offline", - ~select(Job.id).where(Job.node_id == node_id, Job.status.in_(_ACTIVE_VALUES)).exists(), + ~active_jobs.exists(), ) .values(state=NodeRow.state, updated_at=NodeRow.updated_at) .returning(NodeRow.id) ) if removed is None: - raise NodeRemovalError("node changed, connected, or acquired work during removal") + raise NodeRemovalError("Node changed or acquired work during removal. Reload and try again.") await session.execute(update(Job).where(Job.node_id == node_id).values(node_id=None)) await session.execute(update(JobAttempt).where(JobAttempt.node_id == node_id).values(node_id=None)) await session.execute(delete(NodeRow).where(NodeRow.id == node_id)) diff --git a/server/src/server/retention.py b/server/src/server/retention.py index 1c42835..aaa5bc2 100644 --- a/server/src/server/retention.py +++ b/server/src/server/retention.py @@ -21,7 +21,7 @@ QUERY_BATCH = 500 -def protected_history_jobs() -> Select[tuple[int]]: +def protected_history_jobs() -> Select[int]: return select(Job.id).where( or_( Job.id.in_(recovery_jobs()), diff --git a/server/src/server/scanner/service.py b/server/src/server/scanner/service.py index cff9859..ab7c39e 100644 --- a/server/src/server/scanner/service.py +++ b/server/src/server/scanner/service.py @@ -184,6 +184,7 @@ async def check_intervals(self) -> None: intervals = { library.id: library.scan_interval_s for library in libraries if library.scan_interval_s is not None } + scheduled: dict[int, asyncio.Task[ScanReport | None]] = {} for identifier in self._due.keys() - intervals.keys(): self._due.pop(identifier, None) if self.events is not None: @@ -191,14 +192,37 @@ async def check_intervals(self) -> None: for identifier, interval in intervals.items(): if self._intervals.get(identifier) != interval or identifier not in self._due: self._due[identifier] = self._full_completed.get(identifier, now) + timedelta(seconds=interval) - if self._due[identifier] <= now and not self.paused: - self.queue_scan(identifier, "interval") + if self._due[identifier] <= now and not self.paused and self.queue_scan(identifier, "interval") == "queued": + scheduled[identifier] = self._scan_tasks[identifier] if self.events is not None: self.events.update_library_progress( identifier, scanner={"next_run_at": None if identifier in self._scan_tasks else iso_utc(self._due[identifier])}, ) self._intervals = intervals + if scheduled: + asyncio.gather(*scheduled.values(), return_exceptions=True).add_done_callback( + lambda completed: self._log_scan_batch(list(scheduled), completed.result()) + ) + + @staticmethod + def _log_scan_batch(library_ids: list[int], results: list[ScanReport | BaseException | None]) -> None: + reports = [result for result in results if isinstance(result, ScanReport)] + cancelled = sum(isinstance(result, asyncio.CancelledError) for result in results) + failed = len(results) - len(reports) - cancelled + cancelled += sum(report.cancelled for report in reports) + LOG.info( + "Scheduled scans finished for libraries %s: %s files found, %s changed, %s removed, " + "%s errors, %s failed, %s cancelled.", + library_ids, + sum(report.found for report in reports), + sum(report.added + report.updated for report in reports), + sum(report.pruned for report in reports), + sum(len(report.errors) for report in reports), + failed, + cancelled, + extra={"operation": "library_scan_batch", "library_ids": library_ids}, + ) async def update_watchers(self) -> None: async with self._watch_lock: @@ -323,7 +347,7 @@ async def _execute_scan(self, library: Library, started: float) -> ScanReport | self.events.update_library_progress(library.id, discovery={"reason": reason}) scope = "full" if target_paths is None else "scoped" paths = library.roots if target_paths is None else sorted(str(path) for path in target_paths) - LOG.info( + LOG.debug( "Scanning library %s (%s, paths %s), reason %s", library.id, scope, @@ -355,15 +379,34 @@ async def _execute_scan(self, library: Library, started: float) -> ScanReport | self._scan_progress( library.id, "cancelled" if report.cancelled else "failed" if report.errors else "completed", report ) - LOG.info( - "Scanned library %s (%s, paths %s): found %s files, %s changed, %s errors in %.2fs", + level = logging.DEBUG + if report.errors: + level = logging.WARNING + elif ( + reason not in {"interval", "watch"} + or report.added + or report.updated + or report.pruned + or report.cancelled + ): + level = logging.INFO + details = "" + if report.errors: + details = f" Errors: {report.errors}" + elif report.cancelled: + details = " Scan cancelled." + LOG.log( + level, + "Scanned library %s (%s, paths %s): found %s files, %s changed, %s removed, %s errors in %.2fs.%s", library.id, scope, paths, report.found, report.added + report.updated, + report.pruned, len(report.errors), time.perf_counter() - started, + details, extra={ "operation": "library_scan", "reason": reason, @@ -551,7 +594,7 @@ async def _requested_scan(self, library_id: int, reason: str) -> ScanReport | No raise except Exception as exc: reason = self._requests[library_id][0] - LOG.exception("scan failed for library %s", library_id) + LOG.exception("Scan failed for library %s: %s", library_id, exc) if library is None: self._scan_progress(library_id, "failed", error=f"{type(exc).__name__}: {exc}") if self.events is not None: diff --git a/server/src/server/settings.py b/server/src/server/settings.py index 50202d0..55aa142 100644 --- a/server/src/server/settings.py +++ b/server/src/server/settings.py @@ -42,6 +42,7 @@ class OperationalSettings(BaseModel): model_config = ConfigDict(extra="forbid") max_healthchecks: int = Field(default=1, ge=1) + node_registration_enabled: bool = True max_probes: int = Field(default=4, ge=1) max_mkv_stats: int = Field(default=1, ge=1) healthcheck_hwaccel: bool = True diff --git a/server/src/server/storage.py b/server/src/server/storage.py index aa0b7f5..8e21654 100644 --- a/server/src/server/storage.py +++ b/server/src/server/storage.py @@ -14,7 +14,7 @@ from common.process import drain if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncGenerator, Callable _ACTIVE_POOL: ContextVar[ProcessPoolExecutor | None] = ContextVar("storage_pool", default=None) _OPERATION = Lock() @@ -41,7 +41,7 @@ def _operate[T, **P](operation: Callable[P, T], /, *args: P.args, **kwargs: P.kw @asynccontextmanager -async def storage_scope() -> AsyncIterator[ProcessPoolExecutor]: +async def storage_scope() -> AsyncGenerator[ProcessPoolExecutor]: if (pool := _ACTIVE_POOL.get()) is not None: yield pool return diff --git a/server/src/server/trash.py b/server/src/server/trash.py index 27cf3e3..6970d80 100644 --- a/server/src/server/trash.py +++ b/server/src/server/trash.py @@ -28,7 +28,7 @@ from server.workspace import FINISHED_FINALIZATION_STATES, recovery_jobs, wait_workspace_io if TYPE_CHECKING: - from collections.abc import Callable, Generator + from collections.abc import Callable, Generator, Sequence from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.sql import Select @@ -358,7 +358,7 @@ def _manifest(meta: Path, root: Path, *, allow_missing_media: bool = False) -> t return data, media -def _check_trash_records(root: Path, rows: list[tuple[str, int, str, str]], active: set[str]) -> list: +def _check_trash_records(root: Path, rows: Sequence[tuple[str, int, str, str]], active: set[str]) -> list: row_results = [] for row_id, library_id, suffix, library_root in rows: leaf = trash_path(root, library_id, row_id, suffix).parent @@ -494,14 +494,12 @@ async def _refresh_trash_records(session: AsyncSession, root: Path) -> int: after = "" while True: rows = list( - ( - await session.execute( - select(TrashEntry.id, TrashEntry.library_id, TrashEntry.suffix, TrashEntry.library_root) - .where(TrashEntry.id > after) - .order_by(TrashEntry.id) - .limit(_TRASH_BATCH_SIZE) - ) - ).tuples() + await session.execute( + select(TrashEntry.id, TrashEntry.library_id, TrashEntry.suffix, TrashEntry.library_root) + .where(TrashEntry.id > after) + .order_by(TrashEntry.id) + .limit(_TRASH_BATCH_SIZE) + ) ) if not rows: break @@ -668,7 +666,7 @@ async def purge_trash( return count -def _finalization_trash() -> Select[tuple[str]]: +def _finalization_trash() -> Select[str]: return select(FinalizationJournal.trash_path).where( or_( FinalizationJournal.state.not_in(FINISHED_FINALIZATION_STATES), diff --git a/server/src/server/workspace.py b/server/src/server/workspace.py index 89f9942..6ca4af8 100644 --- a/server/src/server/workspace.py +++ b/server/src/server/workspace.py @@ -24,7 +24,7 @@ if TYPE_CHECKING: import threading - from collections.abc import Callable, Iterator + from collections.abc import Callable, Generator from typing import BinaryIO from sqlalchemy.ext.asyncio import AsyncSession @@ -106,7 +106,7 @@ def inspect_directory(jobs_fd: int, job_id: int, stopped: threading.Event) -> di } -def post_action_jobs(*, scoped: bool = False) -> Select[tuple[int]]: +def post_action_jobs(*, scoped: bool = False) -> Select[int]: from sqlalchemy import select from server.models import Job @@ -184,7 +184,7 @@ def _post_action_pending(*, scoped: bool) -> ColumnElement[bool]: def recovery_jobs( *, include_actions: bool = True, include_artifact_checks: bool = True, scoped: bool = False -) -> Select[tuple[int]]: +) -> Select[int]: from common.enums import JobStatus from sqlalchemy import and_, exists, or_, select from sqlalchemy.orm import aliased @@ -269,7 +269,7 @@ def recovery_jobs( ) -def unread_log_jobs() -> Select[tuple[int]]: +def unread_log_jobs() -> Select[int]: from sqlalchemy import and_, exists, select from sqlalchemy.orm import aliased @@ -385,7 +385,7 @@ def _check_directory(path: Path, error: str) -> None: @contextmanager -def _file(directory: int, name: str) -> Iterator[BinaryIO]: +def _file(directory: int, name: str) -> Generator[BinaryIO]: with os.fdopen(os.open(name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory), "rb") as handle: info = os.fstat(handle.fileno()) if not stat.S_ISREG(info.st_mode): @@ -622,7 +622,7 @@ def create_attempt_layout(root: Path, job_id: int, attempt_number: int, destinat @contextmanager -def lock_job(root: Path, job_id: int) -> Iterator[int]: +def lock_job(root: Path, job_id: int) -> Generator[int]: if type(job_id) is not int or job_id < 0: raise ValueError("invalid job workspace") resolve_workspace_path(root, f"jobs/{job_id}") @@ -803,7 +803,7 @@ def read_attempt_log(root: Path, relative_path: str) -> bytes: @contextmanager -def _detached_directory(name: str, directory: int) -> Iterator[int]: +def _detached_directory(name: str, directory: int) -> Generator[int]: parent_fd = os.dup(directory) private = f".{secrets.token_hex(8)}.cleanup" try: diff --git a/tests/conftest.py b/tests/conftest.py index 20edc11..73b7c24 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -147,7 +147,7 @@ def pytest_configure(config) -> None: if getattr(config.option, "basetemp", None) is None or hasattr(config, "workerinput"): config.option.basetemp = str(run_root / "pytest") - # Keep eager plugin-dir mkdirs under the per-run root, not the real XDG state home. + os.environ["MW_CONFIG_PATH"] = str(run_root / "config") os.environ["MW_STATE_PATH"] = str(run_root / "state") os.environ["MW_CACHE_PATH"] = str(run_root / "cache") diff --git a/tests/helpers.py b/tests/helpers.py index 70a055a..57a0803 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -22,7 +22,7 @@ from sqlalchemy.schema import CreateSchema, DropSchema if TYPE_CHECKING: - from collections.abc import AsyncIterator, Generator + from collections.abc import AsyncGenerator, Generator from server.fs.library import LibraryFS, ProbeResult from sqlalchemy.engine import Engine @@ -193,7 +193,7 @@ def make_test_disk_path(*parts: str) -> Path: @contextlib.asynccontextmanager async def database_context( name: str, *, tables: tuple[str, ...] | None = None, fresh: bool = False -) -> AsyncIterator[async_sessionmaker[AsyncSession]]: +) -> AsyncGenerator[async_sessionmaker[AsyncSession]]: with _database_factory(name, tables=tables, fresh=fresh) as (engine, factory): try: yield factory @@ -360,7 +360,7 @@ def cleanup_test_schemas() -> None: @contextlib.asynccontextmanager -async def plugin_host(roots: list[Path]) -> AsyncIterator[PluginHost]: +async def plugin_host(roots: list[Path]) -> AsyncGenerator[PluginHost]: host = PluginHost(roots) try: yield host diff --git a/tests/test_agent.py b/tests/test_agent.py index eee7385..b5cf8ac 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -5,7 +5,7 @@ import json import logging import os -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncGenerator, Callable from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta from dataclasses import dataclass @@ -28,13 +28,15 @@ ActiveAttempt, AttemptMissing, Cancel, - Enrolled, + Registered, + RegistrationPending, Error, Finished, JobOffer, NodeCapabilities, OfferedJob, SetLimits, + Unregister, Welcome, WorkspaceChallenge, parse_message, @@ -113,8 +115,7 @@ async def send_json(self, data: dict) -> None: def make_config(env: AgentEnv, **overrides) -> AgentConfig: fields: dict = { "node_id": None, - "token": "enroll-secret", - "enrollment_token": "enroll-secret", + "token": "", "credentials_path": env.workspace / "node-credentials.json", "server_url": "ws://hub/api/hub", "workspace_dir": str(env.workspace), @@ -244,7 +245,7 @@ async def actions(*_args): monkeypatch.setattr(agent_module, "execute", execute) monkeypatch.setattr(agent_module, "run_actions", actions) - incoming = [] + incoming: list[dict | None] = [] if phase != "hello": incoming.append(Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json")) if phase == "job": @@ -964,7 +965,7 @@ async def wait_until(predicate: Callable[[], object], wake: asyncio.Event, timeo @asynccontextmanager -async def driven_session(agent: Agent, conn: FakeHubConn) -> AsyncIterator[asyncio.Task]: +async def driven_session(agent: Agent, conn: FakeHubConn) -> AsyncGenerator[asyncio.Task]: task = asyncio.create_task(agent._session(conn)) try: yield task @@ -1031,10 +1032,11 @@ async def sleep(interval: float) -> None: @pytest.mark.asyncio -async def test_enrollment_persists_pair(env: AgentEnv) -> None: +async def test_registration_persists_pair(env: AgentEnv) -> None: conn = FakeHubConn( [ - Enrolled(node_id="node-a", token="node-token").model_dump(mode="json"), + RegistrationPending(request_id="request-a").model_dump(mode="json"), + Registered(node_id="node-a", token="node-token").model_dump(mode="json"), Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json"), None, ] @@ -1043,7 +1045,8 @@ async def test_enrollment_persists_pair(env: AgentEnv) -> None: agent = Agent(config) assert await asyncio.wait_for(agent._session(conn), timeout=5) assert conn.sent[0]["node_id"] is None - assert conn.sent[0]["token"] == "enroll-secret" + assert conn.sent[0]["token"] == "" + assert conn.sent[1] == {"v": 1, "type": "credentials_saved", "node_id": "node-a"} assert json.loads((env.workspace / "node-credentials.json").read_text()) == { "node_id": "node-a", "token": "node-token", @@ -1051,62 +1054,67 @@ async def test_enrollment_persists_pair(env: AgentEnv) -> None: @pytest.mark.asyncio -async def test_auth_rejection_enrolls(env: AgentEnv) -> None: +async def test_rejected_credentials_stay(env: AgentEnv) -> None: conn = FakeHubConn([Error(code="auth", message="rejected").model_dump(mode="json")]) config = make_config(env, node_id="node-a", token="old-token") agent = Agent(config) assert await agent._session(conn) is False assert config.node_id == "node-a" and config.token == "old-token" - assert agent._enroll is True - - next_conn = FakeHubConn( - [ - Enrolled(node_id="node-new", token="new-token").model_dump(mode="json"), - Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json"), - None, - ] - ) - assert await asyncio.wait_for(agent._session(next_conn), 5) - assert next_conn.sent[0]["node_id"] is None - assert next_conn.sent[0]["token"] == "enroll-secret" - assert next_conn.sent[1] == {"v": 1, "type": "enrollment_saved", "node_id": "node-new"} - assert json.loads((env.workspace / "node-credentials.json").read_text()) == { - "node_id": "node-new", - "token": "new-token", - } - - reconnect = FakeHubConn([Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json"), None]) - assert await asyncio.wait_for(agent._session(reconnect), 5) - assert reconnect.sent[0]["node_id"] == "node-new" - assert reconnect.sent[0]["token"] == "new-token" @pytest.mark.asyncio -async def test_uncommitted_credentials_reenroll(env: AgentEnv) -> None: +async def test_uncommitted_credentials_claim(env: AgentEnv) -> None: config = make_config(env) agent = Agent(config) - issued = FakeHubConn([Enrolled(node_id="node-issued", token="issued-token").model_dump(mode="json"), None]) + issued = FakeHubConn([Registered(node_id="node-issued", token="issued-token").model_dump(mode="json"), None]) assert await agent._session(issued) is False assert (config.node_id, config.token) == ("node-issued", "issued-token") - rejected = FakeHubConn([Error(code="auth", message="not committed").model_dump(mode="json")]) - assert await agent._session(rejected) is False - assert agent._enroll is True + revoked = FakeHubConn([Unregister().model_dump(mode="json")]) + assert await agent._session(revoked) is False + assert revoked.sent_of("unregistered") + assert config.node_id is None and config.token == "" + assert not (env.workspace / "node-credentials.json").exists() replacement = FakeHubConn( [ - Enrolled(node_id="node-replaced", token="replacement-token").model_dump(mode="json"), + RegistrationPending(request_id="request-b").model_dump(mode="json"), + Registered(node_id="node-replaced", token="replacement-token").model_dump(mode="json"), Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json"), None, ] ) assert await asyncio.wait_for(agent._session(replacement), timeout=5) + assert replacement.sent[0]["node_id"] is None + assert replacement.sent[0]["token"] == "" assert json.loads((env.workspace / "node-credentials.json").read_text()) == { "node_id": "node-replaced", "token": "replacement-token", } +@pytest.mark.asyncio +async def test_unregister_delete_failure(env: AgentEnv, monkeypatch: pytest.MonkeyPatch) -> None: + config = make_config(env, node_id="node-a", token="node-token") + agent = Agent(config) + agent._save_credentials("node-a", config.token) + path = env.workspace / "node-credentials.json" + unlink = Path.unlink + + def denied(current, *args, **kwargs): + if current == path: + raise PermissionError("Credential file is not writable.") + return unlink(current, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", denied) + conn = FakeHubConn([Unregister().model_dump(mode="json")]) + with pytest.raises(PermissionError): + await agent._session(conn) + assert not conn.sent_of("unregistered") + assert config.node_id == "node-a" and config.token == "node-token" + assert path.is_file() + + @pytest.mark.asyncio async def test_network_failure_keeps_identity(env: AgentEnv) -> None: conn = FakeHubConn([None]) @@ -1139,7 +1147,7 @@ def test_workspace_failure_diagnostics(env, monkeypatch, caplog, failure): if failure == "permission": def denied(_path): - raise PermissionError("Access denied for node-secret and enroll-secret") + raise PermissionError("Access denied for node-secret") monkeypatch.setattr(Path, "read_bytes", denied) monkeypatch.chdir(env.workspace.parent) @@ -1153,13 +1161,13 @@ def denied(_path): failed = next(record for record in caplog.records if getattr(record, "operation", None) == "node.workspace.failed") assert getattr(failed, "error") == response["error"] assert failed.levelno == logging.WARNING - for secret in (config.token, config.enrollment_token): + for secret in (config.token,): assert secret not in caplog.text + getattr(failed, "error") assert secret not in response["error"] @pytest.mark.asyncio -@pytest.mark.parametrize("phase", ["initial", "workspace", "enrollment"]) +@pytest.mark.parametrize("phase", ["initial", "workspace", "registration"]) @pytest.mark.parametrize("failure", ["missing", "unexpected", "invalid", "timeout", "auth"]) async def test_handshake_failure_reason(env, caplog, phase, failure): incoming = [] @@ -1170,17 +1178,16 @@ async def test_handshake_failure_reason(env, caplog, phase, failure): mode="json" ) ) - elif phase == "enrollment": - incoming.append(Enrolled(node_id="node-issued", token="issued-secret").model_dump(mode="json")) + elif phase == "registration": + incoming.append(Registered(node_id="node-issued", token="issued-secret").model_dump(mode="json")) reply = { "missing": None, "unexpected": SetLimits(max_transcodes=1, max_healthchecks=1, max_mkv_stats=1).model_dump(mode="json"), - "invalid": {"v": 1, "type": "welcome", "heartbeat_s": "node-secret enroll-secret issued-secret"}, + "invalid": {"v": 1, "type": "welcome", "heartbeat_s": "node-secret issued-secret"}, "timeout": None, "auth": Error( code="auth", - message="Authentication revoked. node-secret enroll-secret" - + (" issued-secret" if phase == "enrollment" else ""), + message="Authentication revoked. node-secret" + (" issued-secret" if phase == "registration" else ""), ).model_dump(mode="json"), }[failure] incoming.append(reply) @@ -1193,7 +1200,7 @@ async def recv_json(self): return data config = make_config(env, token="node-secret") - if phase != "enrollment": + if phase != "registration": config.node_id = "node-a" agent = Agent(config) assert await asyncio.wait_for(agent._session(DiagnosticConn(incoming)), 2) is False @@ -1208,22 +1215,12 @@ async def recv_json(self): if failure == "auth": rejected = next(record for record in records if getattr(record, "operation", None) == "node.handshake.rejected") assert getattr(rejected, "outcome") == "auth" - for secret in ("node-secret", "enroll-secret", "issued-secret"): + for secret in ("node-secret", "issued-secret"): assert secret not in caplog.text + "".join(getattr(record, "error", "") for record in records) @pytest.mark.asyncio -async def test_enrolled_reply_unexpected(env, caplog): - conn = FakeHubConn([Enrolled(node_id="node-a", token="issued-secret").model_dump(mode="json")]) - agent = Agent(make_config(env, node_id="node-a")) - assert await agent._session(conn) is False - records = [record for record in caplog.records if record.name == "node.agent"] - assert records and all(record.levelno == logging.WARNING for record in records) - assert "issued-secret" not in caplog.text - - -@pytest.mark.asyncio -@pytest.mark.parametrize("reason", ["Normal shutdown.", "Rejected old-secret node-secret enroll-secret"]) +@pytest.mark.parametrize("reason", ["Normal shutdown.", "Rejected old-secret node-secret"]) async def test_websocket_close_diagnostics(env, monkeypatch, caplog, reason): frame = Close(1000, reason) @@ -1252,24 +1249,24 @@ async def closed(): assert forwarded[0]["entries"][0]["message"] == closed_record.getMessage() if reason == "Normal shutdown.": assert reason in closed_record.getMessage() - for secret in ("old-secret", "node-secret", "enroll-secret"): + for secret in ("old-secret", "node-secret"): assert secret not in caplog.text @pytest.mark.asyncio -async def test_enrollment_save_diagnostics(env, monkeypatch, caplog): +async def test_registration_save_diagnostics(env, monkeypatch, caplog): agent = Agent(make_config(env, token="node-secret")) def denied(*_args): - raise PermissionError("Credentials denied for node-secret enroll-secret issued-secret") + raise PermissionError("Credentials denied for node-secret issued-secret") monkeypatch.setattr(agent, "_save_credentials", denied) - conn = FakeHubConn([Enrolled(node_id="node-a", token="issued-secret").model_dump(mode="json")]) + conn = FakeHubConn([Registered(node_id="node-a", token="issued-secret").model_dump(mode="json")]) assert await agent._session(conn) is False records = [record for record in caplog.records if record.name == "node.agent"] assert any(getattr(record, "error", None) for record in records) assert records and all(record.levelno == logging.WARNING for record in records) - for secret in ("node-secret", "enroll-secret", "issued-secret"): + for secret in ("node-secret", "issued-secret"): assert secret not in caplog.text + "".join(getattr(record, "error", "") for record in records) @@ -1621,7 +1618,7 @@ async def test_heartbeat_failure_disconnects( env: AgentEnv, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: sampled = asyncio.Event() - error = OSError("resource sample failed enroll-secret") + error = OSError("resource sample failed node-secret") async def sample(_previous): sampled.set() @@ -1633,7 +1630,7 @@ async def sleep(_interval): monkeypatch.setattr(agent_module, "resource_sample", sample) monkeypatch.setattr(agent_module.asyncio, "sleep", sleep) caplog.set_level(logging.ERROR, logger="node.agent") - agent = Agent(make_config(env)) + agent = Agent(make_config(env, token="node-secret")) conn = FakeHubConn([Welcome(timezone="UTC", active_attempts=[]).model_dump(mode="json")]) try: with pytest.raises(OSError) as raised: @@ -1647,8 +1644,8 @@ async def sleep(_interval): assert "resource sample failed" in getattr(failed, "error") traceback = getattr(failed, "traceback") assert traceback and traceback not in failed.getMessage() - assert "enroll-secret" not in traceback - assert "enroll-secret" not in caplog.text + getattr(failed, "error") + assert "node-secret" not in traceback + assert "node-secret" not in caplog.text + getattr(failed, "error") finally: await asyncio.wait_for(agent.aclose(), 2) diff --git a/tests/test_api.py b/tests/test_api.py index e3b5a6b..fd0dc63 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -8,11 +8,12 @@ import re import uuid import zipfile -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import UTC, datetime, timedelta from pathlib import Path +from types import SimpleNamespace from typing import TYPE_CHECKING, Literal from unittest.mock import AsyncMock @@ -24,7 +25,7 @@ from common.media.probe import probe_from_raw from common.plan.models import FFmpegPlan, VideoTarget -from common.protocol import NodeCapabilities +from common.protocol import Hello, NodeCapabilities from common.enums import JobStatus from tests.helpers import api_client from tests.helpers import create_test_app as create_app @@ -56,6 +57,7 @@ from server.staging import StagingService from common.fs.identity import source_identity from tests.helpers import database_context, graph_flow, make_test_temp_dir +from tests.hub_helpers import FakeConn if TYPE_CHECKING: from fastapi import FastAPI @@ -101,7 +103,8 @@ class ApiEnv: @asynccontextmanager -async def api_context(*, background: bool = False) -> AsyncIterator[ApiEnv]: +async def api_context(monkeypatch: pytest.MonkeyPatch, *, background: bool = False) -> AsyncGenerator[ApiEnv]: + monkeypatch.setattr("server.staging.shutil.disk_usage", lambda _: SimpleNamespace(free=20_000_000_000)) lib_dir = make_test_temp_dir("api-lib") async with database_context("api") as factory: app = create_app( @@ -174,9 +177,9 @@ async def drain_services() -> None: @pytest_asyncio.fixture -async def api_env(request: pytest.FixtureRequest) -> AsyncIterator[ApiEnv]: +async def api_env(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[ApiEnv]: # Only background integration tests start service loops. - async with api_context(background=getattr(request, "param", False)) as env: + async with api_context(monkeypatch, background=getattr(request, "param", False)) as env: yield env @@ -234,7 +237,7 @@ async def requeue_outcome(env: ApiEnv, job_id: int, *, event_type: str = "requeu @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["complete", "failure", "cancel"]) async def test_cleanup_joins_scans(monkeypatch: pytest.MonkeyPatch, outcome: str) -> None: - context = api_context() + context = api_context(monkeypatch) env = await context.__aenter__() started, closing_started, release = asyncio.Event(), asyncio.Event(), asyncio.Event() order: list[str] = [] @@ -442,7 +445,7 @@ async def test_queue_rotation(api_env: ApiEnv, monkeypatch) -> None: ) h.entries[entry.name] = entry async with api_env.factory() as session: - before_jobs = [dict(row) for row in (await session.execute(select(Job.__table__))).mappings()] + before_jobs = [dict(row) for row in (await session.execute(select(Job.__table__).order_by(Job.id))).mappings()] before_node = dict((await session.execute(select(NodeRow.__table__))).mappings().one()) before_events = await session.scalar(select(func.count(JobEvent.id))) h._wake.clear() @@ -471,7 +474,9 @@ def track(_conn, _cursor, statement, _parameters, _context, _many): assert "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY" in statements assert not any(statement.startswith(("INSERT", "UPDATE", "DELETE")) for statement in statements) async with api_env.factory() as session: - assert [dict(row) for row in (await session.execute(select(Job.__table__))).mappings()] == before_jobs + assert [ + dict(row) for row in (await session.execute(select(Job.__table__).order_by(Job.id))).mappings() + ] == before_jobs assert dict((await session.execute(select(NodeRow.__table__))).mappings().one()) == before_node assert await session.scalar(select(func.count(JobEvent.id))) == before_events @@ -3589,7 +3594,13 @@ async def test_attempt_history_and(api_env: ApiEnv) -> None: ] ) await session.flush() - attempts = list((await session.execute(select(JobAttempt).where(JobAttempt.job_id == job.id))).scalars()) + attempts = list( + ( + await session.execute( + select(JobAttempt).where(JobAttempt.job_id == job.id).order_by(JobAttempt.attempt_number) + ) + ).scalars() + ) session.add(JobEvent(job_id=job.id, attempt_id=attempts[1].id, type="finished", payload={})) await session.commit() job_id, first_id, second_id, third_id = job.id, attempts[0].id, attempts[1].id, attempts[2].id @@ -3664,6 +3675,49 @@ async def test_attempt_history_and(api_env: ApiEnv) -> None: assert (await api_env.client.get(f"/api/jobs/{job_id}/attempts/{second_id}/log")).status_code == 404 +@pytest.mark.asyncio +async def test_registration_claim(api_env: ApiEnv) -> None: + hub = api_env.app.state.hub + # Immediate delivery makes subscription coalescing independent of runner speed. + api_env.app.state.events.min_interval_s = 0 + conn = FakeConn( + [Hello(capabilities=NodeCapabilities(hostname="encoder"), active_attempts=[]).model_dump(mode="json")] + ) + task = asyncio.create_task(hub.serve_connection(conn)) + sub = api_env.app.state.events.subscribe() + try: + async with asyncio.timeout(3): + while not conn.messages("registration_pending"): + conn.changed.clear() + await conn.changed.wait() + request_id = conn.messages("registration_pending")[0]["request_id"] + response = await api_env.client.get("/api/nodes/registration") + assert response.status_code == 200 + pending = response.json()["pending"] + assert [(item["request_id"], item["capabilities"]["hostname"]) for item in pending] == [(request_id, "encoder")] + assert (await api_env.client.get("/api/nodes")).json() == [] + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=api_env.app), base_url=UI_BASE) as anonymous: + assert (await anonymous.get("/api/nodes/registration")).status_code == 401 + assert (await anonymous.post(f"/api/nodes/pending/{request_id}/claim")).status_code == 401 + async with asyncio.timeout(3): + event = await sub.next() + while event is not None and event.type != "node.registration": + event = await sub.next() + assert event is not None and event.payload["pending"][0]["request_id"] == request_id + response = await api_env.client.post(f"/api/nodes/pending/{request_id}/claim") + assert response.status_code == 202 + assert (await api_env.client.post(f"/api/nodes/pending/{request_id}/claim")).status_code == 404 + async with asyncio.timeout(3): + while not conn.messages("workspace_challenge"): + conn.changed.clear() + await conn.changed.wait() + assert not conn.messages("registered") + finally: + conn.disconnect() + await asyncio.wait_for(task, 3) + api_env.app.state.events.unsubscribe(sub) + + @pytest.mark.asyncio async def test_nodes_list(api_env: ApiEnv) -> None: await seed_node(api_env) @@ -3768,7 +3822,7 @@ async def test_remove_offline_node(api_env: ApiEnv) -> None: @pytest.mark.asyncio -async def test_reject_online_node(api_env: ApiEnv) -> None: +async def test_remove_online_node(api_env: ApiEnv) -> None: await seed_node(api_env) async with api_env.factory() as session: node = await session.scalar(select(NodeRow).where(NodeRow.name == "node-a")) @@ -3777,15 +3831,18 @@ async def test_reject_online_node(api_env: ApiEnv) -> None: await session.commit() node_id = node.id response = await api_env.client.delete(f"/api/nodes/{node_id}") - assert response.status_code == 409 + assert response.status_code == 204 + assert (await api_env.client.get("/api/nodes")).json() == [] @pytest.mark.asyncio -async def test_reject_active_node(api_env: ApiEnv) -> None: +@pytest.mark.parametrize("state", ["online", "offline"]) +async def test_reject_active_node(api_env: ApiEnv, state: str) -> None: await seed_node(api_env) async with api_env.factory() as session: node = await session.scalar(select(NodeRow).where(NodeRow.name == "node-a")) assert node is not None + node.state = state media_id, library_id = await seed_media(api_env) job = Job(media_file_id=media_id, library_id=library_id, status=JobStatus.QUEUED.value, node_id=node.id) session.add(job) @@ -3844,7 +3901,9 @@ async def test_system_timezone(api_env: ApiEnv, tmp_path: Path) -> None: @pytest.mark.asyncio async def test_stats_overview_scope(api_env: ApiEnv) -> None: - now = datetime.now(UTC) + now = T0 + api_env.app.state.stats.clock = lambda: now + api_env.app.state.cache.clock = lambda: now async with api_env.factory() as session: libraries = [Library(name=f"overview-{uuid.uuid4().hex}", roots=[str(api_env.lib_dir)]) for _ in range(3)] session.add_all(libraries) @@ -4806,7 +4865,7 @@ async def test_manual_flow_snapshot(api_env: ApiEnv, skipped: bool, binding: str response = await api_env.client.post( f"/api/files/{media_id}/transcode", - **({"json": {"flow_id": flow["id"], "expected_version": flow["version"]}} if binding != "default" else {}), + json={"flow_id": flow["id"], "expected_version": flow["version"]} if binding != "default" else None, ) assert response.status_code == 202 outcome = await requeue_outcome(api_env, response.json()["job_id"], event_type="request_result") @@ -4939,7 +4998,7 @@ async def test_manual_flow_retries( (api_env.lib_dir / "movie.avi").write_bytes(b"changed before preparation") response = await api_env.client.post( f"/api/files/{media_id}/transcode", - **({"json": {"flow_id": flow["id"], "expected_version": flow["version"]}} if selected_flow else {}), + json={"flow_id": flow["id"], "expected_version": flow["version"]} if selected_flow else None, params={"refresh": scenario == "unprepared"}, ) assert response.status_code == 202, response.text diff --git a/tests/test_auth.py b/tests/test_auth.py index e4d4e2d..6b46e8c 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -603,11 +603,10 @@ async def test_authentication_required(factory) -> None: @pytest.mark.asyncio async def test_api_key_with_oidc(factory) -> None: - app = create_app(factory, config=make_config(api_key="server-cli-key", node_enrollment_token="node-key")) + app = create_app(factory, config=make_config(api_key="server-cli-key")) async with make_ui(app) as ui: assert (await ui.get("/api/system")).status_code == 401 assert (await ui.get("/api/system", headers={"Authorization": "Bearer wrong"})).status_code == 401 - assert (await ui.get("/api/system", headers={"Authorization": "Bearer node-key"})).status_code == 401 headers = {"Authorization": "Bearer server-cli-key"} assert (await ui.get("/api/system", headers=headers)).status_code == 200 assert (await ui.post("/api/auth/logout", headers=headers)).status_code == 400 diff --git a/tests/test_backup_safety.py b/tests/test_backup_safety.py index b2ccde7..0e116bb 100644 --- a/tests/test_backup_safety.py +++ b/tests/test_backup_safety.py @@ -323,11 +323,14 @@ def swap_secrets(fd: int) -> None: @pytest.mark.parametrize("change", ["replacement", "hardlink"]) def test_published_secret_generation(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str) -> None: path = tmp_path / "config.toml" - restore_config(path, 'access_log = false\napi_key = "old-key"\nnode_enrollment_token = "old-token"\n') + monkeypatch.delenv("MW_DATABASE_URL", raising=False) + restore_config( + path, 'access_log = false\napi_key = "old-key"\ndatabase_url = "postgresql://user:old-password@db/millwright"\n' + ) before = path.read_bytes() directory = tmp_path / "secrets" original_files = {file.name: file.read_bytes() for file in directory.iterdir()} - foreign = tmp_path / "retained-token" + foreign = tmp_path / "retained-database-url" changed: Path | None = None fsync = os.fsync @@ -337,8 +340,8 @@ def change_published(fd: int) -> None: if partial.name == ".config.toml.partial": assert changed is None staged = tomllib.loads(partial.read_text()) - changed = tmp_path / staged["node_enrollment_token"].removeprefix("file:") - assert changed.read_bytes() == b"new-token\n" + changed = tmp_path / staged["database_url"].removeprefix("file:") + assert changed.read_bytes() == b"postgresql://user:new-password@db/millwright\n" if change == "replacement": changed.rename(foreign) changed.write_bytes(b"another writer's data") @@ -348,12 +351,15 @@ def change_published(fd: int) -> None: monkeypatch.setattr("server.backup.os.fsync", change_published) with pytest.raises(ValueError): - restore_config(path, 'access_log = false\napi_key = "new-key"\nnode_enrollment_token = "new-token"\n') + restore_config( + path, + 'access_log = false\napi_key = "new-key"\ndatabase_url = "postgresql://user:new-password@db/millwright"\n', + ) assert changed is not None assert path.read_bytes() == before config = load_config(path) - assert config.api_key == "old-key" and config.node_enrollment_token == "old-token" - assert foreign.read_bytes() == b"new-token\n" + assert config.api_key == "old-key" and config.database_url == "postgresql://user:old-password@db/millwright" + assert foreign.read_bytes() == b"postgresql://user:new-password@db/millwright\n" retained = original_files.copy() if change == "replacement": assert changed.read_bytes() == b"another writer's data" diff --git a/tests/test_backups.py b/tests/test_backups.py index fe82369..9d7974e 100644 --- a/tests/test_backups.py +++ b/tests/test_backups.py @@ -97,21 +97,18 @@ def test_restores_config_secrets(tmp_path: Path) -> None: exported = export_config(path) assert exported is not None assert tomllib.loads(exported)["api_key"] == original.api_key - assert tomllib.loads(exported)["node_enrollment_token"] == original.node_enrollment_token restore_config( path, - 'api_key = "other-key"\nnode_enrollment_token = "other-token"\n' + 'api_key = "other-key"\n' '[auth.oidc]\nissuer = "https://sso.example.test"\nclient_id = "client"\n' 'secret = "oidc-secret"\nredirect_uri = "https://example.test/cb"\n', ) restored = load_config(path) assert restored.api_key == "other-key" - assert restored.node_enrollment_token == "other-token" assert restored.auth.oidc is not None and restored.auth.oidc.secret == "oidc-secret" config = tomllib.loads(path.read_text(encoding="utf-8")) references = [ config["api_key"], - config["node_enrollment_token"], config["auth"]["oidc"]["secret"], config["auth"]["oidc"]["client_id"], ] @@ -132,7 +129,8 @@ def test_rejects_secret_symlink(tmp_path: Path) -> None: @pytest.mark.parametrize("stage", ["secret", "config"]) def test_config_publish_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, stage: str) -> None: path = tmp_path / "config.toml" - restore_config(path, 'api_key = "old-key"\nnode_enrollment_token = "old-token"\n') + monkeypatch.delenv("MW_DATABASE_URL", raising=False) + restore_config(path, 'api_key = "old-key"\ndatabase_url = "postgresql://user:old-password@db/millwright"\n') before = path.read_bytes() files = {file.name: file.read_bytes() for file in (tmp_path / "secrets").iterdir()} @@ -157,9 +155,10 @@ def fail_second(source: str | Path, destination: str | Path, **kwargs) -> None: monkeypatch.setattr("server.backup.os.link", fail_second) with pytest.raises(OSError): - restore_config(path, 'api_key = "new-key"\nnode_enrollment_token = "new-token"\n') + restore_config(path, 'api_key = "new-key"\ndatabase_url = "postgresql://user:new-password@db/millwright"\n') assert path.read_bytes() == before - assert load_config(path).api_key == "old-key" and load_config(path).node_enrollment_token == "old-token" + config = load_config(path) + assert config.api_key == "old-key" and config.database_url == "postgresql://user:old-password@db/millwright" assert {file.name: file.read_bytes() for file in (tmp_path / "secrets").iterdir()} == files @@ -487,10 +486,8 @@ async def backup_env( (tmp_path / "oidc.txt").write_text("secret-from-file\n", encoding="utf-8") (tmp_path / "client-id.txt").write_text("millwright\n", encoding="utf-8") (tmp_path / "api-key.txt").write_text("api-secret-from-file\n", encoding="utf-8") - (tmp_path / "enrollment.txt").write_text("node-private-token\n", encoding="utf-8") config_path.write_text( '# My server settings\napi_key = "file:api-key.txt"\n' - 'node_enrollment_token = "file:enrollment.txt"\n' '[auth.oidc]\nissuer = "https://sso.example.test"\n' 'client_id = "file:client-id.txt"\nsecret = "file:oidc.txt"\nredirect_uri = "https://example.test/cb"\n', encoding="utf-8", @@ -1078,7 +1075,6 @@ async def test_backup_selective_restore( assert 'client_id = "millwright"' in payload["config"] assert "# My server settings" in payload["config"] assert 'api_key = "api-secret-from-file"' in payload["config"] - assert 'node_enrollment_token = "node-private-token"' in payload["config"] preview = await client.post( "/api/backups/preview", files={"file": ("backup.mwbackup", exported.content)}, @@ -1139,7 +1135,6 @@ async def test_backup_selective_restore( ) restored_config = load_config(config_path) assert restored_config.api_key == "api-secret-from-file" - assert restored_config.node_enrollment_token == "node-private-token" assert restored_config.auth.oidc is not None and restored_config.auth.oidc.secret == "secret-from-file" assert config_path.stat().st_mode & 0o077 == 0 restored = await client.post( diff --git a/tests/test_changelog.py b/tests/test_changelog.py index 3907aa3..5629c7a 100644 --- a/tests/test_changelog.py +++ b/tests/test_changelog.py @@ -16,6 +16,7 @@ def git(repo: Path, *args: str, date: str = "2026-01-01T12:00:00+00:00") -> str: return subprocess.check_output( ["git", "-C", str(repo), *args], text=True, + timeout=5, env={ **os.environ, "GIT_AUTHOR_NAME": "History Author", @@ -29,7 +30,9 @@ def git(repo: Path, *args: str, date: str = "2026-01-01T12:00:00+00:00") -> str: @pytest.fixture -def repo(tmp_path: Path) -> Path: +def repo(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + monkeypatch.setenv("GIT_CONFIG_NOSYSTEM", "1") + monkeypatch.setenv("GIT_CONFIG_GLOBAL", os.devnull) git(tmp_path, "init", "--initial-branch=main") return tmp_path @@ -41,7 +44,9 @@ def commit(repo: Path, message: str, date: str = "2026-01-01T12:00:00+00:00") -> def generate(repo: Path, *args: str) -> list[dict]: output = repo / "exported/CHANGELOG.json" - subprocess.run([sys.executable, str(SCRIPT), "--repo", str(repo), "--output", str(output), *args], check=True) + subprocess.run( + [sys.executable, str(SCRIPT), "--repo", str(repo), "--output", str(output), *args], check=True, timeout=10 + ) return json.loads(output.read_text()) diff --git a/tests/test_config.py b/tests/test_config.py index af8be6c..f54432f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,12 +8,20 @@ from common.config import apply_env_overrides from common.enums import JobStatus from server.api.app import create_app -from server.core.config import ServerConfig, ensure_config, is_oidc_ready, load_config +from server.core.config import ENV_OVERRIDES, ServerConfig, ensure_config, is_oidc_ready, load_config from server.settings import OperationalSettings from tests.helpers import database_context, make_test_temp_dir from tests.test_staging import add_job, get_job +@pytest.fixture(autouse=True) +def clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for variable in ENV_OVERRIDES.values(): + if variable != "MW_DATABASE_URL": + monkeypatch.delenv(variable, raising=False) + monkeypatch.setenv("TZ", "UTC") + + @pytest.mark.parametrize("field", ["max_probes", "max_mkv_stats"]) def test_metadata_limits(field: str) -> None: with pytest.raises(ValidationError) as raised: @@ -31,18 +39,15 @@ def test_first_start_config(tmp_path: Path) -> None: path = tmp_path / "millwright" / "config.toml" ensure_config(path) config = load_config(path) - assert config.api_key and config.node_enrollment_token - assert config.api_key != config.node_enrollment_token + assert config.api_key content = path.read_text(encoding="utf-8") assert 'api_key = "file:secrets/api_key"' in content assert path.parent.stat().st_mode & 0o777 == 0o700 assert (path.parent / "secrets").stat().st_mode & 0o777 == 0o700 assert path.stat().st_mode & 0o777 == 0o600 - for name in ("api_key", "node_enrollment_token"): - assert (path.parent / "secrets" / name).stat().st_mode & 0o777 == 0o600 + assert (path.parent / "secrets" / "api_key").stat().st_mode & 0o777 == 0o600 ensure_config(path) assert load_config(path).api_key == config.api_key - assert load_config(path).node_enrollment_token == config.node_enrollment_token def test_existing_config_preserved(tmp_path: Path) -> None: @@ -176,15 +181,6 @@ def test_unknown_keys_warn(caplog: pytest.LogCaptureFixture) -> None: assert len(records) == 3 and all(record.levelno == logging.WARNING for record in records) -def test_node_enrollment_token() -> None: - default = load_config(write_config("")) - assert default.node_enrollment_token is None - secret = make_test_temp_dir("config") / "enrollment" - secret.write_text(" fleet-secret \n", encoding="utf-8") - config = load_config(write_config(f'node_enrollment_token = "file:{secret}"\n')) - assert config.node_enrollment_token == "fleet-secret" - - def test_missing_secret_file() -> None: missing = make_test_temp_dir("config") / "nope" with pytest.raises(ValueError): @@ -233,7 +229,6 @@ def test_env_overrides_toml(monkeypatch: pytest.MonkeyPatch) -> None: "enable_docs": "true", "workspace_dir": "/env/workspace", "trash_dir": "/env/trash", - "node_enrollment_token": "env-enrollment", "api_key": "env-api-key", "uploaded_plugins_dir": "/env/plugins", "tz": "Europe/Berlin", @@ -259,7 +254,6 @@ def test_env_overrides_toml(monkeypatch: pytest.MonkeyPatch) -> None: "enable_docs": True, "workspace_dir": "/env/workspace", "trash_dir": "/env/trash", - "node_enrollment_token": "env-enrollment", "api_key": "env-api-key", "uploaded_plugins_dir": "/env/plugins", "tz": "Europe/Berlin", @@ -311,18 +305,17 @@ def test_env_secret_overrides(monkeypatch: pytest.MonkeyPatch) -> None: def test_env_secret_resolution(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: path = tmp_path / "config.toml" path.write_text( - 'api_key = "file:missing"\nnode_enrollment_token = "file:missing"\n' + 'api_key = "file:missing"\n' + FULL_OIDC.replace("literal-secret", "file:missing").replace( 'client_id = "transcode"', 'client_id = "file:missing"' ), encoding="utf-8", ) - for field in ("API_KEY", "NODE_ENROLLMENT_TOKEN", "AUTH_OIDC_SECRET", "AUTH_OIDC_CLIENT_ID"): + for field in ("API_KEY", "AUTH_OIDC_SECRET", "AUTH_OIDC_CLIENT_ID"): (tmp_path / field).write_text(f" {field.lower()} \n", encoding="utf-8") monkeypatch.setenv(f"MW_{field}", f"file:{field}") config = load_config(path) assert config.api_key == "api_key" - assert config.node_enrollment_token == "node_enrollment_token" assert config.auth.oidc is not None assert config.auth.oidc.secret == "auth_oidc_secret" assert config.auth.oidc.client_id == "auth_oidc_client_id" diff --git a/tests/test_config_safety.py b/tests/test_config_safety.py index 4eff029..32b13c2 100644 --- a/tests/test_config_safety.py +++ b/tests/test_config_safety.py @@ -44,13 +44,10 @@ def test_env_validation_privacy( ) -> None: monkeypatch.setenv(variable, value) monkeypatch.setenv("MW_API_KEY", "audit-secret-marker") - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "audit-secret-marker") monkeypatch.setenv("MW_NODE_STATE_PATH", str(tmp_path / "state")) path = tmp_path / "config.toml" path.write_text( - f'enrollment_token = "enroll-secret"\nserver_url = "ws://hub"\nworkspace_dir = "{tmp_path / "workspace"}"\n' - if node - else "", + f'server_url = "ws://hub"\nworkspace_dir = "{tmp_path / "workspace"}"\n' if node else "", encoding="utf-8", ) with pytest.raises(ValueError) as error: @@ -114,24 +111,6 @@ async def test_trash_contains_workspace(tmp_path: Path) -> None: create_app(db, config=config) -def test_node_validation_privacy(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - marker = "audit-secret-marker" - for key in ("MW_NODE_ENROLLMENT_TOKEN", "MW_NODE_SERVER_URL", "MW_NODE_WORKSPACE_DIR"): - monkeypatch.delenv(key, raising=False) - monkeypatch.setenv("MW_NODE_STATE_PATH", str(tmp_path / "state")) - path = tmp_path / "node.toml" - path.write_text( - f'enrollment_token = ["{marker}"]\nserver_url = "ws://hub"\nworkspace_dir = "{tmp_path / "workspace"}"\n', - encoding="utf-8", - ) - with pytest.raises(ValueError) as error: - load_settings(path, must_exist=True) - leaked = marker in str(error.value) - assert leaked is False - traceback_leaked = marker in "".join(traceback.format_exception(error.value)) - assert traceback_leaked is False - - @pytest.mark.asyncio async def test_upload_contains_config(tmp_path: Path) -> None: root = tmp_path / "uploads" diff --git a/tests/test_executor_safety.py b/tests/test_executor_safety.py index b12e478..dd83e2f 100644 --- a/tests/test_executor_safety.py +++ b/tests/test_executor_safety.py @@ -699,13 +699,16 @@ async def test_log_fsync_preserves(tmp_path: Path, monkeypatch: pytest.MonkeyPat job = spec(tmp_path, min_free_gb=0) logs: list[_Log] = [] opening, syncing = _Log.__init__, os.fsync + fsync_failed = False def open_log(log, path): opening(log, path) logs.append(log) def failed_fsync(fd): + nonlocal fsync_failed if logs and not logs[0].file.closed and fd == logs[0].file.fileno(): + fsync_failed = True raise OSError("NFS log fsync rejected") return syncing(fd) @@ -737,5 +740,6 @@ async def probe(*_args): assert result.output_path == job.dst and result.output_probe is not None elif outcome == "aborted": assert result.error and "src.mkv" in result.error + assert fsync_failed assert logs and logs[0].file.closed assert Path(job.dst).exists() == (outcome in {"succeeded", "unreaped"}) diff --git a/tests/test_external_notifications.py b/tests/test_external_notifications.py index d0cc644..b66e30e 100644 --- a/tests/test_external_notifications.py +++ b/tests/test_external_notifications.py @@ -2,9 +2,12 @@ import asyncio import json +from pathlib import PurePath import httpx import pytest +from common.actions import ExternalNotificationConfig +from common.plugin import build_inputs from sqlalchemy import select from server.external_notifications import notify_external @@ -28,12 +31,253 @@ def notification(application: str = "radarr", *, unmonitor: bool = True) -> dict } +@pytest.mark.parametrize( + "options, error", + [ + ({"application": "plex", "unmonitor_after_refresh": True}, "Unmonitor after refresh"), + ({"application": "plex", "path_mappings": [{"destination": "/media"}]}, "source"), + ({"application": "plex", "path_mappings": [{"source": "/library"}]}, "destination"), + ( + {"application": "plex", "path_mappings": [{"source": "library", "destination": "/media"}]}, + "Use absolute paths", + ), + ( + {"application": "plex", "path_mappings": [{"source": "/library", "destination": "media"}]}, + "Use absolute paths", + ), + ( + {"application": "plex", "path_mappings": [{"source": "/library", "destination": "/media/../other"}]}, + "can't contain", + ), + ( + { + "application": "plex", + "path_mappings": [ + {"source": "/library", "destination": "/media"}, + {"source": "/library/", "destination": "/other"}, + ], + }, + "appear only once", + ), + ], +) +def test_external_options_rejected(options, error) -> None: + with pytest.raises(ValueError, match=error): + build_inputs(ExternalNotificationConfig, {"url": "https://server.test", "api_key": "key", **options}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mappings, root, folder, prefix", + [ + ([], "/library", "/library/Series", "/library/Ser"), + ( + [ + {"source": "/library/Movies", "destination": "/movies"}, + {"source": "/library/", "destination": "/media/"}, + ], + "/media", + "/media/Series", + "/media/Ser", + ), + ( + [ + {"source": "/library", "destination": "/wrong"}, + {"source": "/library/Series", "destination": "/media/Series"}, + ], + "/media", + "/media/Series", + "/media/Ser", + ), + ( + [ + {"source": "/library/Series", "destination": "/media/Series"}, + {"source": "/library", "destination": "/wrong"}, + ], + "/media", + "/media/Series", + "/media/Ser", + ), + ( + [{"source": "/library", "destination": "D:\\Media"}], + "d:\\media", + "D:\\Media\\Series", + "D:\\Media\\Ser", + ), + ], +) +async def test_plex_folder_refresh(monkeypatch, mappings, root, folder, prefix) -> None: + payload = notification("plex", unmonitor=False) + payload["action"]["path_mappings"] = mappings + requests = [] + + def respond(request): + requests.append(request) + assert request.headers["X-Plex-Token"] == "private-api-key" + assert request.headers["Accept"] == "application/json" + if request.url.path == "/base/library/sections": + return httpx.Response( + 200, + json={ + "MediaContainer": { + "Directory": [ + {"key": "7", "Location": [{"path": root}, {"path": folder}]}, + {"key": "8", "Location": [{"path": root}]}, + {"key": "9", "Location": [{"path": prefix}]}, + ] + } + }, + ) + assert request.method == "GET" + assert dict(request.url.params) == {"path": folder} + return httpx.Response(200) + + client = httpx.AsyncClient + monkeypatch.setattr( + "server.external_notifications.plex.httpx.AsyncClient", + lambda **kwargs: client(**kwargs, transport=httpx.MockTransport(respond)), + ) + assert await notify_external(payload) + assert [(request.method, request.url.path) for request in requests] == [ + ("GET", "/base/library/sections"), + ("GET", "/base/library/sections/7/refresh"), + ("GET", "/base/library/sections/8/refresh"), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mismatch", ["mapping", "library"]) +async def test_plex_path_mismatch(monkeypatch, mismatch) -> None: + payload = notification("plex", unmonitor=False) + payload["path"] = "/library/SeriesOther/Episode.mp4" + if mismatch == "mapping": + payload["action"]["path_mappings"] = [ + {"source": "/library/Series", "destination": "/media/Series"}, + {"source": "/library/Movies", "destination": "/media/Movies"}, + ] + requests = [] + + def respond(request): + requests.append(request) + assert request.url.path == "/base/library/sections" + return httpx.Response( + 200, + json={"MediaContainer": {"Directory": [{"key": "7", "Location": [{"path": "/library/Series"}]}]}}, + ) + + client = httpx.AsyncClient + monkeypatch.setattr( + "server.external_notifications.plex.httpx.AsyncClient", + lambda **kwargs: client(**kwargs, transport=httpx.MockTransport(respond)), + ) + with pytest.raises(ValueError, match="No path mapping matches" if mismatch == "mapping" else "No Plex library"): + await notify_external(payload) + assert len(requests) == (0 if mismatch == "mapping" else 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "application, path, folder", + [ + ("jellyfin", "/library/Series & Friends/Season 1/Episode.mp4", "/media/Series & Friends/Season 1"), + ("emby", "/library/Movies/A Movie/Movie.mp4", "/movies/A Movie"), + ], +) +async def test_media_folder_update(monkeypatch, application, path, folder) -> None: + payload = notification(application, unmonitor=False) + payload["path"] = path + payload["action"]["path_mappings"] = [ + {"source": "/library/Movies", "destination": "/movies"}, + {"source": "/library", "destination": "/media"}, + ] + requests = [] + + def respond(request): + requests.append(request) + assert request.method == "POST" + assert request.url.path == "/base/Library/Media/Updated" + assert not request.url.params + if application == "jellyfin": + assert request.headers["Authorization"] == 'MediaBrowser Token="private-api-key"' + else: + assert request.headers["X-Emby-Token"] == "private-api-key" + assert json.loads(request.content) == {"Updates": [{"Path": folder, "UpdateType": "Modified"}]} + return httpx.Response(204 if application == "jellyfin" else 200) + + client = httpx.AsyncClient + monkeypatch.setattr( + "server.external_notifications.media_server.httpx.AsyncClient", + lambda **kwargs: client(**kwargs, transport=httpx.MockTransport(respond)), + ) + assert await notify_external(payload) + assert len(requests) == 1 + + +@pytest.mark.asyncio +async def test_plex_retry_persisted(monkeypatch, caplog) -> None: + refreshes = [] + + def respond(request): + if request.url.path == "/base/library/sections": + return httpx.Response( + 200, + json={ + "MediaContainer": { + "Directory": [{"key": key, "Location": [{"path": "/library"}]} for key in ["7", "8"]] + } + }, + ) + refreshes.append(request.url.path) + if len(refreshes) == 2: + raise httpx.ConnectError("connection refused private-api-key", request=request) + return httpx.Response(200) + + client = httpx.AsyncClient + monkeypatch.setattr( + "server.external_notifications.plex.httpx.AsyncClient", + lambda **kwargs: client(**kwargs, transport=httpx.MockTransport(respond)), + ) + async with database_context("plex-retry") as factory: + async with factory() as session: + job_id = await make_job(session, notify=False) + job = await session.get(Job, job_id) + assert job is not None + job.status = "succeeded" + event = JobEvent(job_id=job_id, type="external_notification", payload=notification("plex", unmonitor=False)) + session.add(event) + await session.flush() + session.add(NotificationDelivery(job_id=job_id, job_event_id=event.id, target_name="external:plex")) + await session.commit() + await NotificationService(factory).deliver_pending() + async with factory() as session: + delivery = (await session.scalars(select(NotificationDelivery))).one() + assert (delivery.status, delivery.attempts) == ("pending", 1) + assert delivery.error and "connection refused" in delivery.error + assert "private-api-key" not in delivery.error + await NotificationService(factory).deliver_pending() + async with factory() as session: + delivery = (await session.scalars(select(NotificationDelivery))).one() + assert (delivery.status, delivery.attempts, delivery.error) == ("sent", 2, None) + job = await session.get(Job, job_id) + assert job is not None and job.status == "succeeded" + assert refreshes == [ + "/base/library/sections/7/refresh", + "/base/library/sections/8/refresh", + "/base/library/sections/7/refresh", + "/base/library/sections/8/refresh", + ] + assert all("private-api-key" not in str(record.__dict__) for record in caplog.records) + + @pytest.mark.asyncio @pytest.mark.parametrize("identifier", [None, "tmdb-42", "imdb-tt123"]) async def test_radarr_refresh_order(monkeypatch, identifier) -> None: payload = notification() if identifier: payload["path"] = f"/library/Movie {{{identifier}}}/Movie.mp4" + payload["action"]["path_mappings"] = [ + {"source": str(PurePath(payload["path"]).parent), "destination": "/radarr/Movie"} + ] requests = [] movie = {"id": 7, "tmdbId": 42, "imdbId": "tt123"} @@ -79,12 +323,19 @@ async def tick(_seconds): @pytest.mark.parametrize("failure", [None, "lookup", "refresh", "unmonitor"]) async def test_sonarr_delivery_persisted(monkeypatch, caplog, failure) -> None: payload = notification("sonarr") + payload["action"]["path_mappings"] = [ + {"source": "/other", "destination": "/unused"}, + {"source": "/library", "destination": "/sonarr"}, + ] requests = [] def respond(request): requests.append(request) if request.url.path.endswith("/parse"): - assert dict(request.url.params) == {"title": payload["original_filename"], "path": payload["path"]} + assert dict(request.url.params) == { + "title": payload["original_filename"], + "path": "/sonarr/Series/Series.S01E01E02.mp4", + } return httpx.Response( 200, json={ @@ -128,6 +379,8 @@ def respond(request): assert delivery.attempts == (3 if failure else 1) assert (delivery.error is not None) is bool(failure) assert "private-api-key" not in (delivery.error or "") + event = await session.get(JobEvent, delivery.job_event_id) + assert event is not None and event.payload["path"] == "/library/Series/Series.S01E01E02.mp4" job = await session.get(Job, job_id) assert job is not None and job.status == "succeeded" if failure in {"lookup", "refresh"}: diff --git a/tests/test_finalization.py b/tests/test_finalization.py index 2cefab6..aedd08c 100644 --- a/tests/test_finalization.py +++ b/tests/test_finalization.py @@ -1542,7 +1542,11 @@ async def test_recovery_batches(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, assert [job.status for job in jobs] == ["finalizing", *(["failed"] * 4)] journals = (await session.scalars(select(FinalizationJournal).order_by(FinalizationJournal.id))).all() assert [row.state for row in journals] == ["discard_prepared", *(["discard_rolled_back"] * 4)] - events = (await session.scalars(select(JobEvent).where(JobEvent.type == "finalization_rolled_back"))).all() + events = ( + await session.scalars( + select(JobEvent).where(JobEvent.type == "finalization_rolled_back").order_by(JobEvent.job_id) + ) + ).all() assert [event.job_id for event in events] == job_ids[1:] assert all((tmp_path / f"source-{index}.mkv").read_bytes() == b"original" for index in range(5)) diff --git a/tests/test_hub.py b/tests/test_hub.py index e4008b3..fa38431 100644 --- a/tests/test_hub.py +++ b/tests/test_hub.py @@ -36,6 +36,7 @@ NodeLog, NodeLogs, Progress, + Unregistered, WorkspaceResponse, parse_message, ) @@ -267,7 +268,8 @@ async def test_foreign_action_report(env): @pytest_asyncio.fixture -async def env(): +async def env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("server.staging.shutil.disk_usage", lambda _: SimpleNamespace(free=20_000_000_000)) async with database_context("hub") as factory: async with factory() as session: await register_node(session, "node-a", "tok") @@ -282,7 +284,8 @@ async def env(): @pytest_asyncio.fixture -async def writer_env(tmp_path: Path): +async def writer_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr("server.staging.shutil.disk_usage", lambda _: SimpleNamespace(free=20_000_000_000)) library = tmp_path / "library" library.mkdir() async with database_context("hub-writer") as factory: @@ -699,15 +702,23 @@ async def wait_for(pred, event, timeout=3): await event.wait() -async def serve(h, conn, *, enroll=False): +async def serve(h, conn, *, claim=False): task = asyncio.create_task(h.serve_connection(conn)) - await wait_for(lambda: conn.messages("workspace_challenge"), conn.changed) - answer_challenge(h, conn) - if enroll: - await wait_for(lambda: conn.messages("enrolled"), conn.changed) - conn.push({"v": 1, "type": "enrollment_saved", "node_id": conn.messages("enrolled")[0]["node_id"]}) - await wait_for(lambda: conn.messages("welcome"), conn.changed) - await asyncio.wait_for(conn.ready.wait(), 3) + try: + if claim: + await wait_for(lambda: conn.messages("registration_pending"), conn.changed) + assert h.claim_node(conn.messages("registration_pending")[0]["request_id"]) + await wait_for(lambda: conn.messages("workspace_challenge"), conn.changed) + answer_challenge(h, conn) + if claim: + await wait_for(lambda: conn.messages("registered"), conn.changed) + conn.push({"v": 1, "type": "credentials_saved", "node_id": conn.messages("registered")[0]["node_id"]}) + await wait_for(lambda: conn.messages("welcome"), conn.changed) + await asyncio.wait_for(conn.ready.wait(), 3) + except BaseException: + task.cancel() + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), 3) + raise return task @@ -1015,7 +1026,9 @@ async def test_challenge_revocation_checked(env): ) await session.commit() answer_challenge(h, conn) - await wait_for(lambda: conn.messages("welcome") or conn.messages("error") or connection.done(), conn.changed) + await wait_for(lambda: conn.messages("unregister"), conn.changed) + conn.push(Unregistered().model_dump(mode="json")) + await asyncio.wait_for(connection, 3) assert not conn.messages("welcome"), "revoked credentials queued a connection after workspace verification" assert "node-a" not in h.entries finally: @@ -1041,9 +1054,11 @@ async def test_removed_challenge_refused(env, replace): await register_node(session, "node-a", "replacement-token") await session.commit() answer_challenge(h, conn) + await wait_for(lambda: conn.messages("unregister"), conn.changed) + conn.push(Unregistered().model_dump(mode="json")) await asyncio.wait_for(connection, 3) assert not conn.messages("welcome") and not h.entries - assert conn.messages("error")[0]["code"] == "auth" and conn.closed == 4401 + assert conn.closed == 1000 async with env.factory() as session: stored = await get_node(session, "node-a") if replace: @@ -1526,8 +1541,10 @@ async def remove(): assert (env.workspace / offer["output_path"]).read_bytes() == b"output" else: assert removed and replacement_id != node_id + await wait_for(lambda: conn.messages("unregister"), conn.changed) + conn.push(Unregistered().model_dump(mode="json")) await asyncio.wait_for(connection, 3) - assert conn.messages("error")[-1]["code"] == "auth" + assert not conn.messages("welcome") replacement_conn = FakeConn([hello("node-b", "replacement-token")]) replacement_connection = await serve(h, replacement_conn) assert await asyncio.wait_for(h.dispatch_once(), 2) == 1 @@ -1706,7 +1723,7 @@ async def actions(*_args): @pytest.mark.asyncio async def test_post_cleanup_retry(env, monkeypatch) -> None: - from server.queue.store import ActiveJobError + from server.queue.store import ActiveJobError, NodeRemovalError, remove_node cleanup_started, release_cleanup = asyncio.Event(), asyncio.Event() async with post_action_case(env) as case: @@ -1735,6 +1752,8 @@ async def delayed_cleanup(*args): async with env.factory() as session: from server.queue.store import restore_retry_result + with pytest.raises(NodeRemovalError): + await remove_node(session, case.hub.entries["node-a"].node_db_id) request_id = await request_retry(session, case.jid) with pytest.raises(ActiveJobError): await retry_job(session, case.jid, request_id=request_id) @@ -4149,9 +4168,10 @@ async def test_idle_hello_availability(env): @pytest.mark.asyncio -async def test_hello_rejects_unlocked(env): +async def test_registration_rejects_unlocked(env): h = hub(env) - conn = FakeConn([hello(token="wrong")]) + h.set_registration_enabled(False) + conn = FakeConn([hello(name=None, token="")]) async with h.dispatch_lock: await asyncio.wait_for(h.serve_connection(conn), 3) assert conn.messages("error") and conn.closed is not None @@ -4214,8 +4234,10 @@ async def wait_reset(): current.token_hash = hash_token("replacement") await session.commit() if revoke: + await wait_for(lambda: conn.messages("unregister"), conn.changed) + conn.push(Unregistered().model_dump(mode="json")) await asyncio.wait_for(task, 3) - assert conn.messages("error")[0]["code"] == "auth" and conn.closed == 4401 + assert conn.closed == 1000 assert not conn.messages("workspace_challenge") and not h.entries else: await wait_for(lambda: conn.messages("workspace_challenge"), conn.changed) @@ -4234,10 +4256,12 @@ async def test_handshake_ready(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - n = await node(env) - assert n.state == "online" and n.workspace_ready and n.workspace_free_bytes == 123456 - c.disconnect() - await asyncio.wait_for(task, 2) + try: + n = await node(env) + assert n.state == "online" and n.workspace_ready and n.workspace_free_bytes == 123456 + finally: + c.disconnect() + await asyncio.wait_for(task, 2) n = await node(env) assert n.state == "offline" and not n.workspace_ready and n.workspace_free_bytes is None @@ -4307,19 +4331,176 @@ async def heartbeat_handler(entry, msg): @pytest.mark.asyncio -async def test_handshake_rejects(env): - async with env.factory() as s: - s.add(await register_node(s, "node-a", "right")) - await s.commit() - c = FakeConn([hello(token="wrong")]) - await hub(env).serve_connection(c) - assert c.messages("error")[0]["code"] == "auth" and c.closed == 4401 +@pytest.mark.parametrize("registration", ["invalid", "offline", "interrupted"]) +async def test_stale_credentials_clear(env, registration): + h = hub(env) + await file_job(env) + removed = registration != "invalid" + if removed: + + class RemovalConn(FakeConn): + async def send_json(self, data): + if data["type"] == "unregister": + raise OSError("Connection lost before unregister delivery.") + await super().send_json(data) + + original = RemovalConn([hello()]) + connection = await serve(h, original) if registration == "interrupted" else None + async with env.factory() as session: + saved = await get_node(session, "node-a") + assert saved is not None + node_id = saved.id + try: + assert await h.remove_node(node_id) + finally: + original.disconnect() + if connection is not None: + await asyncio.wait_for(connection, 2) + c = FakeConn([hello(token="tok" if removed else "wrong")]) + task = asyncio.create_task(h.serve_connection(c)) + retry = FakeConn([hello(name=None, token="")]) + retried = None + try: + await wait_for(lambda: c.messages("unregister"), c.changed) + assert not h.pending_nodes and not c.messages("workspace_challenge") + c.push(Unregistered().model_dump(mode="json")) + await asyncio.wait_for(task, 2) + retried = asyncio.create_task(h.serve_connection(retry)) + await wait_for(lambda: retry.messages("registration_pending"), retry.changed) + assert not retry.messages("registered") + await h.dispatch_once() + assert not retry.messages("job") + async with env.factory() as session: + saved = await get_node(session, "node-a") + if removed: + assert saved is None + else: + assert saved is not None and saved.token_hash == hash_token("tok") + assert len(list(await session.scalars(select(NodeRow)))) == (0 if removed else 1) + finally: + c.disconnect() + retry.disconnect() + await asyncio.wait_for(asyncio.gather(task, *([retried] if retried else [])), 2) + assert not h.pending_nodes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("registration", [True, False]) +async def test_online_unregister_reconnects(env, monkeypatch, registration): + from node.agent import Agent, AgentConfig + + h = hub(env) + h.set_registration_enabled(registration) + path = env.workspace.parent / "node-credentials.json" + config = AgentConfig( + node_id="node-a", + token="tok", + credentials_path=path, + server_url="ws://test", + workspace_dir=str(env.workspace), + min_free_gb=0, + ) + agent = Agent(config) + agent._save_credentials("node-a", config.token) + retained = env.workspace / "retained.log" + retained.write_bytes(b"retained log") + acknowledged, release = asyncio.Event(), asyncio.Event() + connections, tasks = [], [] + + def connect(): + server, client = FakeConn(), FakeConn() + connections.extend((server, client)) + + async def from_hub(data): + await FakeConn.send_json(server, data) + client.push(data) + + async def from_node(data): + await FakeConn.send_json(client, data) + if data["type"] == "unregistered": + acknowledged.set() + await asyncio.wait_for(release.wait(), 3) + server.push(data) + + async def close_hub(code=1000): + server.closed = code + client.disconnect() + + monkeypatch.setattr(server, "send_json", from_hub) + monkeypatch.setattr(client, "send_json", from_node) + monkeypatch.setattr(server, "close", close_hub) + tasks.extend((asyncio.create_task(h.serve_connection(server)), asyncio.create_task(agent._session(client)))) + return server, client + + server, client = connect() + removal = None + try: + await asyncio.wait_for(server.ready.wait(), 3) + node_id = h.entries["node-a"].node_db_id + jid = await file_job(env) + removal = asyncio.create_task(h.remove_node(node_id)) + await asyncio.wait_for(acknowledged.wait(), 3) + assert not removal.done() + assert await asyncio.wait_for(h.dispatch_once(), 2) == 0 + assert (await job(env, jid)).status == JobStatus.QUEUED + assert not path.exists() and config.node_id is None and config.token == "" + assert retained.read_bytes() == b"retained log" + async with env.factory() as session: + assert await get_node(session, "node-a") is None + release.set() + assert await asyncio.wait_for(removal, 3) + await asyncio.wait_for(asyncio.gather(*tasks), 3) + assert client.messages("unregistered") and not agent._closed + retry, retried_client = connect() + if registration: + await wait_for(lambda: retry.messages("registration_pending"), retry.changed) + assert len(h.pending_nodes) == 1 and not retry.messages("registered") + else: + await asyncio.wait_for(asyncio.gather(*tasks[-2:]), 3) + assert not h.pending_nodes and retry.messages("error")[0]["code"] == "registration" + assert retried_client.messages("hello")[0]["node_id"] is None + assert retried_client.messages("hello")[0]["token"] == "" + finally: + release.set() + for connection in connections: + connection.disconnect() + if removal is not None and not removal.done(): + removal.cancel() + await asyncio.wait_for(asyncio.gather(*tasks, *([removal] if removal else []), return_exceptions=True), 3) + await asyncio.wait_for(agent.aclose(), 3) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["cancel", "connection"]) +async def test_pending_cleanup(env, failure): + class BrokenConn(FakeConn): + async def recv_json(self): + data = await super().recv_json() + if data is None: + raise OSError("Connection lost.") + return data + + h = hub(env) + conn = BrokenConn([hello(name=None, token="")]) + task = asyncio.create_task(h.serve_connection(conn)) + try: + await wait_for(lambda: conn.messages("registration_pending"), conn.changed) + if failure == "cancel": + task.cancel() + else: + conn.disconnect() + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), 2) + assert not h.pending_nodes and not conn.messages("registered") + finally: + if not task.done(): + task.cancel() + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), 2) @pytest.mark.asyncio @pytest.mark.parametrize("failure", ["node", "missing", "digest", "permission", "verification", "closure"]) async def test_workspace_failure_diagnostics(env, monkeypatch, caplog, failure): - h = hub(env, enrollment_token="fleet-secret") + h = hub(env) conn = FakeConn([hello()]) task = asyncio.create_task(h.serve_connection(conn)) try: @@ -4333,13 +4514,13 @@ async def test_workspace_failure_diagnostics(env, monkeypatch, caplog, failure): elif failure == "permission": def denied(_path): - raise PermissionError("Response denied for tok fleet-secret") + raise PermissionError("Response denied for tok") monkeypatch.setattr(Path, "read_bytes", denied) else: def invalid(*_args): - raise ValueError("Unsafe shared mount for tok fleet-secret") + raise ValueError("Unsafe shared mount for tok") monkeypatch.setattr("server.nodehub.hub._check_challenge", invalid) elif failure == "closure": @@ -4351,7 +4532,7 @@ def invalid(*_args): response_path=challenge["response_path"], workspace_ready=failure != "node", workspace_free_bytes=0, - error="Shared mount missing for tok fleet-secret" if failure == "node" else None, + error="Shared mount missing for tok" if failure == "node" else None, ).model_dump(mode="json") ) await asyncio.wait_for(task, 2) @@ -4363,7 +4544,7 @@ def invalid(*_args): assert diagnostics if failure in {"missing", "digest"}: assert str(response_path) in diagnostics - assert "tok" not in caplog.text + diagnostics and "fleet-secret" not in caplog.text + diagnostics + assert "tok" not in caplog.text + diagnostics finally: conn.disconnect() if not task.done(): @@ -4373,10 +4554,10 @@ def invalid(*_args): @pytest.mark.asyncio async def test_workspace_creation_diagnostics(env, monkeypatch, caplog): - h = hub(env, enrollment_token="fleet-secret") + h = hub(env) def denied(*_args): - raise PermissionError("Challenge directory denied for tok fleet-secret") + raise PermissionError("Challenge directory denied for tok") monkeypatch.setattr("server.nodehub.hub._prepare_challenge", denied) conn = FakeConn([hello()]) @@ -4384,24 +4565,26 @@ def denied(*_args): assert conn.messages("error")[0]["code"] == "workspace" diagnostics = "\n".join(getattr(record, "error", "") for record in caplog.records) assert diagnostics - assert "tok" not in caplog.text + diagnostics and "fleet-secret" not in caplog.text + diagnostics + assert "tok" not in caplog.text + diagnostics @pytest.mark.asyncio -async def test_enrollment_issues_token(env): - h = hub(env, enrollment_token="fleet-secret") - greeting = hello(name=None, token="fleet-secret") +async def test_claim_issues_token(env): + h = hub(env) + greeting = hello(name=None, token="") greeting["capabilities"].update(cpu_model="Intel(R) Core(TM) i9-14900K", gpus=["Intel UHD", "NVIDIA RTX 4090"]) c = FakeConn([greeting]) - task = await serve(h, c, enroll=True) + task = await serve(h, c, claim=True) try: - enrolled = c.messages("enrolled")[0] - assert enrolled["node_id"].startswith("node-") - assert enrolled["token"] != "fleet-secret" + registered = c.messages("registered")[0] + assert registered["node_id"].startswith("node-") + assert registered["token"] async with env.factory() as session: - node = await get_node(session, enrolled["node_id"]) + node = await get_node(session, registered["node_id"]) assert node is not None - assert node.token_hash == hash_token(enrolled["token"]) + assert node.token_hash == hash_token(registered["token"]) + saved = await get_node(session, "node-a") + assert saved is not None and saved.token_hash == hash_token("tok") assert node.transcode_limit == 1 assert node.display_name.startswith("i9-igpu-nvidia-") assert len(node.display_name.rsplit("-", 1)[-1]) == 4 @@ -4411,47 +4594,81 @@ async def test_enrollment_issues_token(env): @pytest.mark.asyncio -async def test_enrollment_rejects_named(env): - h = hub(env, enrollment_token="fleet-secret") - c = FakeConn([hello(name="arbitrary", token="fleet-secret")]) - await h.serve_connection(c) - assert c.messages("error")[0]["code"] == "auth" - async with env.factory() as session: - assert await get_node(session, "arbitrary") is None +async def test_registration_toggle(env): + h = hub(env) + registered = FakeConn([hello()]) + active = await serve(h, registered) + pending = FakeConn([hello(name=None, token="")]) + waiting = asyncio.create_task(h.serve_connection(pending)) + retry = FakeConn([hello(name=None, token="")]) + retried = None + try: + await wait_for(lambda: pending.messages("registration_pending"), pending.changed) + request_id = pending.messages("registration_pending")[0]["request_id"] + h.set_registration_enabled(False) + await asyncio.wait_for(waiting, 2) + assert not h.pending_nodes and not h.claim_node(request_id) + assert pending.messages("error")[0]["code"] == "registration" + assert h.entries["node-a"].workspace_ready and not registered.messages("error") + await asyncio.wait_for(h.serve_connection(retry), 2) + assert not retry.messages("registration_pending") + h.set_registration_enabled(True) + retry = FakeConn([hello(name=None, token="")]) + retried = asyncio.create_task(h.serve_connection(retry)) + await wait_for(lambda: retry.messages("registration_pending"), retry.changed) + assert len(h.pending_nodes) == 1 + await asyncio.wait_for(h.aclose(), 2) + await asyncio.wait_for(retried, 2) + assert not h.pending_nodes and retry.messages("error")[0]["code"] == "registration" + finally: + registered.disconnect() + pending.disconnect() + retry.disconnect() + await asyncio.wait_for(asyncio.gather(active, waiting, *([retried] if retried else [])), 2) @pytest.mark.asyncio -async def test_enrollment_workspace_failure(env): - h = hub(env, enrollment_token="fleet-secret") - c = FakeConn([hello(name=None, token="fleet-secret")]) +async def test_registration_workspace_failure(env): + h = hub(env) + c = FakeConn([hello(name=None, token="")]) task = asyncio.create_task(h.serve_connection(c)) - await wait_for(lambda: c.messages("workspace_challenge"), c.changed) - challenge = c.messages("workspace_challenge")[0] - c.push( - WorkspaceResponse( - challenge_id=challenge["challenge_id"], - response_path=challenge["response_path"], - workspace_ready=False, - workspace_free_bytes=0, - ).model_dump(mode="json") - ) - await asyncio.wait_for(task, 2) - assert c.messages("error")[0]["code"] == "workspace" + try: + await wait_for(lambda: c.messages("registration_pending"), c.changed) + assert h.claim_node(c.messages("registration_pending")[0]["request_id"]) + await wait_for(lambda: c.messages("workspace_challenge"), c.changed) + challenge = c.messages("workspace_challenge")[0] + c.push( + WorkspaceResponse( + challenge_id=challenge["challenge_id"], + response_path=challenge["response_path"], + workspace_ready=False, + workspace_free_bytes=0, + ).model_dump(mode="json") + ) + await asyncio.wait_for(task, 2) + assert c.messages("error")[0]["code"] == "workspace" + finally: + c.disconnect() + if not task.done(): + task.cancel() + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), 2) @pytest.mark.asyncio @pytest.mark.parametrize("failure", ["closure", "parse", "timeout", "wrong_node"]) -async def test_enrollment_save_failure(env, caplog, failure): - class EnrollmentConn(FakeConn): +async def test_registration_save_failure(env, caplog, failure): + class RegistrationConn(FakeConn): async def recv_json(self): - if failure == "timeout" and self.messages("enrolled"): + if failure == "timeout" and self.messages("registered"): raise TimeoutError return await super().recv_json() - h = hub(env, enrollment_token="fleet-secret") - c = EnrollmentConn([hello(name=None, token="fleet-secret")]) + h = hub(env) + c = RegistrationConn([hello(name=None, token="")]) task = asyncio.create_task(h.serve_connection(c)) try: + await wait_for(lambda: c.messages("registration_pending"), c.changed) + assert h.claim_node(c.messages("registration_pending")[0]["request_id"]) await wait_for(lambda: c.messages("workspace_challenge"), c.changed) challenge = c.messages("workspace_challenge")[0] path = h.workspace_dir / challenge["response_path"] @@ -4465,12 +4682,14 @@ async def recv_json(self): workspace_free_bytes=10, ).model_dump(mode="json") ) - await wait_for(lambda: c.messages("enrolled"), c.changed) - issued_id = c.messages("enrolled")[0]["node_id"] + await wait_for(lambda: c.messages("registered"), c.changed) + issued_id = c.messages("registered")[0]["node_id"] + async with env.factory() as session: + assert await get_node(session, issued_id) is None if failure == "parse": - c.push({"v": 1, "type": "enrollment_saved"}) + c.push({"v": 1, "type": "credentials_saved"}) elif failure == "wrong_node": - c.push({"v": 1, "type": "enrollment_saved", "node_id": "other-node"}) + c.push({"v": 1, "type": "credentials_saved", "node_id": "other-node"}) elif failure == "closure": c.push(None) await asyncio.wait_for(task, 2) @@ -4555,17 +4774,19 @@ async def test_dispatch_paths(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - assert await h.dispatch_once() == 1 - offer = c.messages("job")[0]["job"] - assert offer["min_free_gb"] == h.staging.min_free_gb == 1 - assert offer["attempt_id"] > 0 and all( - not p.startswith("/") for p in (offer["source_path"], offer["output_path"], offer["log_path"]) - ) - assert str(env.library) not in json.dumps(offer) - a = (await attempt(env, jid))[0] - assert a.source_path == offer["source_path"] and a.output_path == offer["output_path"] - c.disconnect() - await asyncio.wait_for(task, 2) + try: + assert await h.dispatch_once() == 1 + offer = c.messages("job")[0]["job"] + assert offer["min_free_gb"] == h.staging.min_free_gb == 1 + assert offer["attempt_id"] > 0 and all( + not p.startswith("/") for p in (offer["source_path"], offer["output_path"], offer["log_path"]) + ) + assert str(env.library) not in json.dumps(offer) + a = (await attempt(env, jid))[0] + assert a.source_path == offer["source_path"] and a.output_path == offer["output_path"] + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -4751,16 +4972,18 @@ async def test_paused_skips_dispatch(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - h.set_paused(h.entries["node-a"].node_db_id, True) - assert await h.dispatch_once() == 0 - j = await job(env, jid) - assert j.status == JobStatus.QUEUED - assert j.warning is not None - h.set_paused(h.entries["node-a"].node_db_id, False) - assert await h.dispatch_once() == 1 - assert c.messages("job")[0]["job"]["job_id"] == jid - c.disconnect() - await asyncio.wait_for(task, 2) + try: + h.set_paused(h.entries["node-a"].node_db_id, True) + assert await h.dispatch_once() == 0 + j = await job(env, jid) + assert j.status == JobStatus.QUEUED + assert j.warning is not None + h.set_paused(h.entries["node-a"].node_db_id, False) + assert await h.dispatch_once() == 1 + assert c.messages("job")[0]["job"]["job_id"] == jid + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -4786,20 +5009,24 @@ async def test_waiting_warning_clears(env): h = hub(env) c = FakeConn([hello(encoders=("hevc_nvenc",))]) task = await serve(h, c) - assert await h.dispatch_once() == 0 - assert await h.dispatch_once() == 0 - waiting = await job(env, jid) - assert waiting.status == JobStatus.QUEUED and waiting.warning is not None - assert (await events(env, jid)).count("waiting") == 1 - c.disconnect() - await asyncio.wait_for(task, 2) + try: + assert await h.dispatch_once() == 0 + assert await h.dispatch_once() == 0 + waiting = await job(env, jid) + assert waiting.status == JobStatus.QUEUED and waiting.warning is not None + assert (await events(env, jid)).count("waiting") == 1 + finally: + c.disconnect() + await asyncio.wait_for(task, 2) c2 = FakeConn([hello()]) task2 = await serve(h, c2) - await h.dispatch_once() - j = await job(env, jid) - assert j.warning is None and j.status == JobStatus.CLAIMED - c2.disconnect() - await asyncio.wait_for(task2, 2) + try: + await h.dispatch_once() + j = await job(env, jid) + assert j.warning is None and j.status == JobStatus.CLAIMED + finally: + c2.disconnect() + await asyncio.wait_for(task2, 2) @pytest.mark.asyncio @@ -4849,12 +5076,14 @@ async def test_cancel_attempt(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - assert await h.cancel_job(jid) - msg = c.messages("cancel")[0] - assert msg["attempt_id"] == (await attempt(env, jid))[0].id - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + assert await h.cancel_job(jid) + msg = c.messages("cancel")[0] + assert msg["attempt_id"] == (await attempt(env, jid))[0].id + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5173,12 +5402,14 @@ async def test_stall_cancel(env): h = hub(env, stall_timeout_s=1) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - env.clock.advance(2) - await h.sweep_once() - assert c.messages("cancel")[0]["attempt_id"] > 0 - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + env.clock.advance(2) + await h.sweep_once() + assert c.messages("cancel")[0]["attempt_id"] > 0 + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5187,20 +5418,24 @@ async def test_disconnect_retains_attempt(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - first = (await attempt(env, jid))[0].id - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + first = (await attempt(env, jid))[0].id + finally: + c.disconnect() + await asyncio.wait_for(task, 2) assert (await job(env, jid)).status == JobStatus.CLAIMED assert jid in h.entries["node-a"].in_flight c2 = FakeConn([hello(active_attempts=[{"job_id": jid, "attempt_id": first}])]) task = await serve(h, c2) - assert await h.dispatch_once() == 0 - rows = await attempt(env, jid) - assert len(rows) == 1 and rows[0].id == first and rows[0].result is None - assert jid in h.entries["node-a"].in_flight and not c2.messages("job") - c2.disconnect() - await asyncio.wait_for(task, 2) + try: + assert await h.dispatch_once() == 0 + rows = await attempt(env, jid) + assert len(rows) == 1 and rows[0].id == first and rows[0].result is None + assert jid in h.entries["node-a"].in_flight and not c2.messages("job") + finally: + c2.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5210,15 +5445,17 @@ async def test_finished_failure(env, result): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - o = c.messages("job")[0]["job"] - await h._handle_finished(h.entries["node-a"], finished(jid, o["attempt_id"], o["log_path"], result)) - status = (await job(env, jid)).status - assert status == JobStatus.FAILED - assert not (h.workspace_dir / f"jobs/{jid}").exists() - assert (env.library / "movie.avi").read_bytes() == b"source" - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + o = c.messages("job")[0]["job"] + await h._handle_finished(h.entries["node-a"], finished(jid, o["attempt_id"], o["log_path"], result)) + status = (await job(env, jid)).status + assert status == JobStatus.FAILED + assert not (h.workspace_dir / f"jobs/{jid}").exists() + assert (env.library / "movie.avi").read_bytes() == b"source" + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5592,21 +5829,23 @@ async def test_finishes_cancelled(env): h = hub(env, max_attempts=3) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - offer = c.messages("job")[0]["job"] + try: + await h.dispatch_once() + offer = c.messages("job")[0]["job"] - assert await h.cancel_job(jid) - await h._handle_finished( - h.entries["node-a"], - finished(jid, offer["attempt_id"], offer["log_path"], "cancelled"), - ) + assert await h.cancel_job(jid) + await h._handle_finished( + h.entries["node-a"], + finished(jid, offer["attempt_id"], offer["log_path"], "cancelled"), + ) - assert (await job(env, jid)).status == JobStatus.CANCELLED - assert (await attempt(env, jid))[0].status == JobStatus.CANCELLED - assert await h.dispatch_once() == 0 - assert len(await attempt(env, jid)) == 1 - c.disconnect() - await asyncio.wait_for(task, 2) + assert (await job(env, jid)).status == JobStatus.CANCELLED + assert (await attempt(env, jid))[0].status == JobStatus.CANCELLED + assert await h.dispatch_once() == 0 + assert len(await attempt(env, jid)) == 1 + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5615,33 +5854,37 @@ async def test_duplicate_finished(env, monkeypatch): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - offer = c.messages("job")[0]["job"] - await h._handle_finished(h.entries["node-a"], finished(jid, offer["attempt_id"], offer["log_path"], "cancelled")) - calls = 0 + try: + await h.dispatch_once() + offer = c.messages("job")[0]["job"] + await h._handle_finished( + h.entries["node-a"], finished(jid, offer["attempt_id"], offer["log_path"], "cancelled") + ) + calls = 0 - async def probe(*args, **kwargs): - nonlocal calls - calls += 1 - raise AssertionError("duplicate was probed") + async def probe(*args, **kwargs): + nonlocal calls + calls += 1 + raise AssertionError("duplicate was probed") - monkeypatch.setattr("server.nodehub.hub.run_ffprobe", probe) - duplicate = finished( - jid, - offer["attempt_id"], - offer["log_path"], - "succeeded", - output_path=offer["output_path"], - output_probe=ProbeData.model_validate( - {"file_path": offer["output_path"], "duration_s": 1, "size_bytes": 1, "streams": []} - ), - ) - await h._handle_finished(h.entries["node-a"], duplicate) - assert calls == 0 and (await job(env, jid)).status == JobStatus.CANCELLED - assert "foreign_finished" in await events(env, jid) - assert (await events(env, jid)).count("cancelled") == 1 - c.disconnect() - await asyncio.wait_for(task, 2) + monkeypatch.setattr("server.nodehub.hub.run_ffprobe", probe) + duplicate = finished( + jid, + offer["attempt_id"], + offer["log_path"], + "succeeded", + output_path=offer["output_path"], + output_probe=ProbeData.model_validate( + {"file_path": offer["output_path"], "duration_s": 1, "size_bytes": 1, "streams": []} + ), + ) + await h._handle_finished(h.entries["node-a"], duplicate) + assert calls == 0 and (await job(env, jid)).status == JobStatus.CANCELLED + assert "foreign_finished" in await events(env, jid) + assert (await events(env, jid)).count("cancelled") == 1 + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -5650,17 +5893,23 @@ async def test_reconnect_keeps_locks(env): h = hub(env) old = FakeConn([hello()]) old_task = await serve(h, old) - await h.dispatch_once() - offer = old.messages("job")[0]["job"] - new = FakeConn([hello(active_attempts=[{"job_id": jid, "attempt_id": offer["attempt_id"]}])]) - new_task = await serve(h, new) - old.disconnect() - await asyncio.wait_for(old_task, 2) - assert h.entries["node-a"].conn is new and jid in h.entries["node-a"].in_flight - assert (await job(env, jid)).status == JobStatus.CLAIMED - assert "retried" not in await events(env, jid) - new.disconnect() - await asyncio.wait_for(new_task, 2) + try: + await h.dispatch_once() + offer = old.messages("job")[0]["job"] + new = FakeConn([hello(active_attempts=[{"job_id": jid, "attempt_id": offer["attempt_id"]}])]) + new_task = await serve(h, new) + try: + old.disconnect() + await asyncio.wait_for(old_task, 2) + assert h.entries["node-a"].conn is new and jid in h.entries["node-a"].in_flight + assert (await job(env, jid)).status == JobStatus.CLAIMED + assert "retried" not in await events(env, jid) + finally: + new.disconnect() + await asyncio.wait_for(new_task, 2) + finally: + old.disconnect() + await asyncio.wait_for(old_task, 2) @pytest.mark.asyncio @@ -5711,22 +5960,24 @@ async def test_heartbeat_filters_locks(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - entry = h.entries["node-a"] - before = await job(env, jid) - env.clock.advance(1) - await h._handle_heartbeat(entry, Heartbeat(cpu_percent=1, memory_percent=1, gpus=[])) - renewed = await job(env, jid) - assert renewed.lease_expires_at > before.lease_expires_at - async with env.factory() as session: - row = await session.get(Job, jid) - row.status = JobStatus.QUEUED.value - row.node_id = None - await session.commit() - await h._handle_heartbeat(entry, Heartbeat(cpu_percent=1, memory_percent=1, gpus=[])) - assert jid not in entry.in_flight and jid not in entry.offered_at and jid not in entry.last_progress_at - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + entry = h.entries["node-a"] + before = await job(env, jid) + env.clock.advance(1) + await h._handle_heartbeat(entry, Heartbeat(cpu_percent=1, memory_percent=1, gpus=[])) + renewed = await job(env, jid) + assert renewed.lease_expires_at > before.lease_expires_at + async with env.factory() as session: + row = await session.get(Job, jid) + row.status = JobStatus.QUEUED.value + row.node_id = None + await session.commit() + await h._handle_heartbeat(entry, Heartbeat(cpu_percent=1, memory_percent=1, gpus=[])) + assert jid not in entry.in_flight and jid not in entry.offered_at and jid not in entry.last_progress_at + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -6007,34 +6258,36 @@ async def test_probe_failure(env, monkeypatch): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - offer = c.messages("job")[0]["job"] - materialize_offer(h, offer) + try: + await h.dispatch_once() + offer = c.messages("job")[0]["job"] + materialize_offer(h, offer) - async def fail_probe(*args, **kwargs): - raise RuntimeError("bad artifact") + async def fail_probe(*args, **kwargs): + raise RuntimeError("bad artifact") - monkeypatch.setattr("server.nodehub.hub.run_ffprobe", fail_probe) - node_probe = ProbeData.model_validate( - {"file_path": offer["output_path"], "duration_s": 60, "size_bytes": 6, "streams": []} - ) - await h._handle_finished( - h.entries["node-a"], - finished( - jid, - offer["attempt_id"], - offer["log_path"], - "succeeded", - output_path=offer["output_path"], - output_probe=node_probe, - ), - ) - row = await job(env, jid) - assert row.status == JobStatus.FAILED and row.error is not None - assert not (h.workspace_dir / offer["output_path"]).exists() - assert (env.library / "movie.avi").exists() - c.disconnect() - await asyncio.wait_for(task, 2) + monkeypatch.setattr("server.nodehub.hub.run_ffprobe", fail_probe) + node_probe = ProbeData.model_validate( + {"file_path": offer["output_path"], "duration_s": 60, "size_bytes": 6, "streams": []} + ) + await h._handle_finished( + h.entries["node-a"], + finished( + jid, + offer["attempt_id"], + offer["log_path"], + "succeeded", + output_path=offer["output_path"], + output_probe=node_probe, + ), + ) + row = await job(env, jid) + assert row.status == JobStatus.FAILED and row.error is not None + assert not (h.workspace_dir / offer["output_path"]).exists() + assert (env.library / "movie.avi").exists() + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -6043,33 +6296,35 @@ async def test_post_failure(env, monkeypatch): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - offer = c.messages("job")[0]["job"] - output = materialize_offer(h, offer) + try: + await h.dispatch_once() + offer = c.messages("job")[0]["job"] + output = materialize_offer(h, offer) - async def fail_gate(*args, **kwargs): - raise RuntimeError("plugin exploded") + async def fail_gate(*args, **kwargs): + raise RuntimeError("plugin exploded") - monkeypatch.setattr(h, "_post_gate", fail_gate) - node_probe = ProbeData.model_validate( - {"file_path": offer["output_path"], "duration_s": 60, "size_bytes": 6, "streams": []} - ) - await h._handle_finished( - h.entries["node-a"], - finished( - jid, - offer["attempt_id"], - offer["log_path"], - "succeeded", - output_path=offer["output_path"], - output_probe=node_probe, - ), - ) - row = await job(env, jid) - assert row.status == JobStatus.FAILED and row.error is not None - assert not output.exists() - c.disconnect() - await asyncio.wait_for(task, 2) + monkeypatch.setattr(h, "_post_gate", fail_gate) + node_probe = ProbeData.model_validate( + {"file_path": offer["output_path"], "duration_s": 60, "size_bytes": 6, "streams": []} + ) + await h._handle_finished( + h.entries["node-a"], + finished( + jid, + offer["attempt_id"], + offer["log_path"], + "succeeded", + output_path=offer["output_path"], + output_probe=node_probe, + ), + ) + row = await job(env, jid) + assert row.status == JobStatus.FAILED and row.error is not None + assert not output.exists() + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -6437,33 +6692,37 @@ async def fail_rescan(*args, **kwargs): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - o = c.messages("job")[0]["job"] - output = materialize_offer(h, o) - probe = ProbeData.model_validate( - {"file_path": o["output_path"], "container": "mkv", "duration_s": 60, "size_bytes": 6, "streams": []} - ) - await h._handle_finished( - h.entries["node-a"], - finished(jid, o["attempt_id"], o["log_path"], "succeeded", output_path=o["output_path"], output_probe=probe), - ) - row = await job(env, jid) - assert row.status == JobStatus.SUCCEEDED - assert row.warning is not None - async with env.factory() as session: - assert ( - await session.scalar( - select(JobEvent.type).where( - JobEvent.attempt_id == o["attempt_id"], - JobEvent.type.in_(["artifact_check_pending", "artifact_check_finished"]), + try: + await h.dispatch_once() + o = c.messages("job")[0]["job"] + output = materialize_offer(h, o) + probe = ProbeData.model_validate( + {"file_path": o["output_path"], "container": "mkv", "duration_s": 60, "size_bytes": 6, "streams": []} + ) + await h._handle_finished( + h.entries["node-a"], + finished( + jid, o["attempt_id"], o["log_path"], "succeeded", output_path=o["output_path"], output_probe=probe + ), + ) + row = await job(env, jid) + assert row.status == JobStatus.SUCCEEDED + assert row.warning is not None + async with env.factory() as session: + assert ( + await session.scalar( + select(JobEvent.type).where( + JobEvent.attempt_id == o["attempt_id"], + JobEvent.type.in_(["artifact_check_pending", "artifact_check_finished"]), + ) ) + == "artifact_check_finished" ) - == "artifact_check_finished" - ) - assert (env.library / "movie.mkv").exists() - assert not (env.library / "movie.avi").exists() - c.disconnect() - await asyncio.wait_for(task, 2) + assert (env.library / "movie.mkv").exists() + assert not (env.library / "movie.avi").exists() + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.asyncio @@ -7807,8 +8066,11 @@ async def send_json(self, data): dispatch.cancel() else: release.set() - await asyncio.wait_for(asyncio.gather(dispatch, return_exceptions=True), 2) - assert dispatch.cancelled() is (failure == "cancel") + if failure == "cancel": + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(dispatch, 2) + else: + assert await asyncio.wait_for(dispatch, 2) is None assert not (env.workspace / f"jobs/{case.jid}").exists() assert not case.hub._action_reports assert (env.library / "movie.avi").read_bytes() == b"source" @@ -8057,7 +8319,11 @@ async def send_json(self, data): teardown_report() elif phase == "cancel": dispatch.cancel() - await asyncio.wait_for(asyncio.gather(dispatch, return_exceptions=True), 2) + if phase == "cancel" or failure == "cancelled": + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(dispatch, 2) + else: + assert await asyncio.wait_for(dispatch, 2) is None assert not output.exists() assert "actions_finished" in await events(env, case.jid) assert not case.hub._action_reports @@ -8368,16 +8634,18 @@ async def test_log_parent_symlink(env): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - o = c.messages("job")[0]["job"] - p = h.workspace_dir / o["log_path"] - external = make_test_temp_dir("outside") - p.parent.rename(external / "moved") - p.parent.symlink_to(external / "moved") - await h._handle_finished(h.entries["node-a"], finished(jid, o["attempt_id"], o["log_path"])) - assert (await job(env, jid)).status == JobStatus.FAILED - c.disconnect() - await asyncio.wait_for(task, 2) + try: + await h.dispatch_once() + o = c.messages("job")[0]["job"] + p = h.workspace_dir / o["log_path"] + external = make_test_temp_dir("outside") + p.parent.rename(external / "moved") + p.parent.symlink_to(external / "moved") + await h._handle_finished(h.entries["node-a"], finished(jid, o["attempt_id"], o["log_path"])) + assert (await job(env, jid)).status == JobStatus.FAILED + finally: + c.disconnect() + await asyncio.wait_for(task, 2) @pytest.mark.parametrize("preserve", [False, True]) @@ -8398,48 +8666,52 @@ async def test_finished_transcode(env, monkeypatch, preserve, missing_log): h = hub(env) c = FakeConn([hello()]) task = await serve(h, c) - await h.dispatch_once() - o = c.messages("job")[0]["job"] - output = materialize_offer(h, o) - output.write_bytes(b"remux") - if missing_log: - (h.workspace_dir / o["log_path"]).unlink() - from common.media.probe import ProbeData + try: + await h.dispatch_once() + o = c.messages("job")[0]["job"] + output = materialize_offer(h, o) + output.write_bytes(b"remux") + if missing_log: + (h.workspace_dir / o["log_path"]).unlink() + from common.media.probe import ProbeData - probe = ProbeData.model_validate( - {"file_path": o["output_path"], "container": "mkv", "duration_s": 60, "size_bytes": 5, "streams": []} - ) - await h._handle_finished( - h.entries["node-a"], - finished(jid, o["attempt_id"], o["log_path"], "succeeded", output_path=o["output_path"], output_probe=probe), - ) - assert (await job(env, jid)).status == JobStatus.SUCCEEDED - assert list(env.trash.rglob("*")) - async with env.factory() as s: - event = await s.scalar( - select(JobEvent).where(JobEvent.attempt_id == o["attempt_id"], JobEvent.type == "publication") + probe = ProbeData.model_validate( + {"file_path": o["output_path"], "container": "mkv", "duration_s": 60, "size_bytes": 5, "streams": []} ) - assert event.payload["preserve_timestamps"] is preserve - media = ( - await s.execute(select(MediaFile).where(MediaFile.path == str(env.library / "movie.mkv"))) - ).scalar_one() - assert media.size == 5 - assert media.container == "mkv" - assert media.video_codec == "h264" - assert media.health_status == HealthStatus.PENDING.value - assert media.transcode_status == "transcoded" - attempt = await s.get(JobAttempt, o["attempt_id"]) - assert attempt is not None and attempt.output_probe is not None - assert attempt.output_probe["file_path"] == str(env.library / "movie.mkv") - assert attempt.output_probe["container"] == "mkv" - assert attempt.output_probe["streams"][0]["codec"] == "h264" - probe = await s.get(ProbeRow, media.probe_id) - assert json.loads(probe.raw_json)["format"]["duration"] == "600.0" - healthcheck = ( - await s.execute(select(Job).where(Job.kind == "healthcheck", Job.media_file_id == media.id)) - ).scalar_one() - c.disconnect() - await asyncio.wait_for(task, 2) + await h._handle_finished( + h.entries["node-a"], + finished( + jid, o["attempt_id"], o["log_path"], "succeeded", output_path=o["output_path"], output_probe=probe + ), + ) + assert (await job(env, jid)).status == JobStatus.SUCCEEDED + assert list(env.trash.rglob("*")) + async with env.factory() as s: + event = await s.scalar( + select(JobEvent).where(JobEvent.attempt_id == o["attempt_id"], JobEvent.type == "publication") + ) + assert event.payload["preserve_timestamps"] is preserve + media = ( + await s.execute(select(MediaFile).where(MediaFile.path == str(env.library / "movie.mkv"))) + ).scalar_one() + assert media.size == 5 + assert media.container == "mkv" + assert media.video_codec == "h264" + assert media.health_status == HealthStatus.PENDING.value + assert media.transcode_status == "transcoded" + attempt = await s.get(JobAttempt, o["attempt_id"]) + assert attempt is not None and attempt.output_probe is not None + assert attempt.output_probe["file_path"] == str(env.library / "movie.mkv") + assert attempt.output_probe["container"] == "mkv" + assert attempt.output_probe["streams"][0]["codec"] == "h264" + probe = await s.get(ProbeRow, media.probe_id) + assert json.loads(probe.raw_json)["format"]["duration"] == "600.0" + healthcheck = ( + await s.execute(select(Job).where(Job.kind == "healthcheck", Job.media_file_id == media.id)) + ).scalar_one() + finally: + c.disconnect() + await asyncio.wait_for(task, 2) assert media.first_seen.replace(tzinfo=UTC) == first_seen assert healthcheck.status == JobStatus.QUEUED and healthcheck.hold_until is None diff --git a/tests/test_job_service.py b/tests/test_job_service.py index 777b48c..f7e2f4a 100644 --- a/tests/test_job_service.py +++ b/tests/test_job_service.py @@ -315,41 +315,42 @@ async def blocked(*args, **kwargs): raise monkeypatch.setattr(service.staging, "prune_job", blocked) - assert (await api_env.client.post(f"/api/jobs/{job_id}/retry")).status_code == 202 - await asyncio.wait_for(entered.wait(), 2) - if operation == "shutdown": - ending = asyncio.create_task(service.close()) - elif operation == "reset": - api_env.app.state.services_running = True - api_env.app.state.auth.config.api_key = "test-job-service" - api_env.client.headers["Authorization"] = "Bearer test-job-service" - ending = asyncio.create_task(api_env.client.post("/api/backups/reset", json={"keep": []})) - # Keep reset's service restarts deterministic and avoid starting unrelated loops. - for owner, name in ( - (api_env.app.state.scanner, "start"), - (api_env.app.state.healthchecks, "start"), - (api_env.app.state.staging, "start"), - (api_env.app.state.scheduler, "start"), - (api_env.app.state.notifications, "start"), - ): - monkeypatch.setattr(owner, name, lambda: None) - - async def start(): - pass - - monkeypatch.setattr(api_env.app.state.probes, "start", start) - monkeypatch.setattr(api_env.app.state.flow_enqueue, "start", start) - else: - api_env.app.state.services_running = True - async with api_env.factory() as session: - job = await session.get(Job, job_id) - assert job is not None - library_id = job.library_id - await api_env.app.state.library_deletion.start() - response = await api_env.client.delete(f"/api/libraries/{library_id}") - assert response.status_code == 202 - ending = api_env.app.state.library_deletion._tasks[library_id] + ending = None try: + assert (await api_env.client.post(f"/api/jobs/{job_id}/retry")).status_code == 202 + await asyncio.wait_for(entered.wait(), 2) + if operation == "shutdown": + ending = asyncio.create_task(service.close()) + elif operation == "reset": + api_env.app.state.services_running = True + api_env.app.state.auth.config.api_key = "test-job-service" + api_env.client.headers["Authorization"] = "Bearer test-job-service" + ending = asyncio.create_task(api_env.client.post("/api/backups/reset", json={"keep": []})) + # Keep reset's service restarts deterministic and avoid starting unrelated loops. + for owner, name in ( + (api_env.app.state.scanner, "start"), + (api_env.app.state.healthchecks, "start"), + (api_env.app.state.staging, "start"), + (api_env.app.state.scheduler, "start"), + (api_env.app.state.notifications, "start"), + ): + monkeypatch.setattr(owner, name, lambda: None) + + async def start(): + pass + + monkeypatch.setattr(api_env.app.state.probes, "start", start) + monkeypatch.setattr(api_env.app.state.flow_enqueue, "start", start) + else: + api_env.app.state.services_running = True + async with api_env.factory() as session: + job = await session.get(Job, job_id) + assert job is not None + library_id = job.library_id + await api_env.app.state.library_deletion.start() + response = await api_env.client.delete(f"/api/libraries/{library_id}") + assert response.status_code == 202 + ending = api_env.app.state.library_deletion._tasks[library_id] await asyncio.wait_for(draining.wait(), 2) assert not ending.done() release.set() @@ -366,8 +367,9 @@ async def start(): assert await session.get(Library, library_id) is None finally: release.set() - await asyncio.wait_for(asyncio.gather(ending, return_exceptions=True), 5) - await api_env.app.state.library_deletion.close() + if ending is not None: + await asyncio.wait_for(asyncio.gather(ending, return_exceptions=True), 5) + await asyncio.wait_for(api_env.app.state.library_deletion.close(), 5) @pytest.mark.asyncio diff --git a/tests/test_job_store.py b/tests/test_job_store.py index dacff73..9639e78 100644 --- a/tests/test_job_store.py +++ b/tests/test_job_store.py @@ -607,7 +607,7 @@ async def test_startup_requeues_batches(db: async_sessionmaker[AsyncSession], mo assert await requeue_startup(session, after(5)) == 5 await session.commit() async with db() as session: - statuses = dict((await session.execute(select(Job.id, Job.status))).tuples().all()) + statuses = dict((await session.execute(select(Job.id, Job.status))).all()) assert all(statuses[job_id] == "queued" for job_id in jobs) assert statuses[transcode.id] == "claimed" assert statuses[remote.id] == "claimed" @@ -665,7 +665,11 @@ async def test_claims_interrupted_healthcheck(db: async_sessionmaker[AsyncSessio assert claimed is not None and claimed.status == JobStatus.CLAIMED async with db() as session: - attempts = list((await session.scalars(select(JobAttempt).where(JobAttempt.job_id == job.id))).all()) + attempts = list( + await session.scalars( + select(JobAttempt).where(JobAttempt.job_id == job.id).order_by(JobAttempt.attempt_number) + ) + ) assert [(attempt.attempt_number, attempt.status, attempt.result) for attempt in attempts] == [ (1, JobStatus.FAILED.value, "interrupted"), (2, JobStatus.CLAIMED.value, None), diff --git a/tests/test_kv.py b/tests/test_kv.py index 10f9fb8..ed2361e 100644 --- a/tests/test_kv.py +++ b/tests/test_kv.py @@ -3,6 +3,7 @@ import asyncio import pytest +from pydantic import JsonValue from sqlalchemy import event, func, select from server.kv import KeyValueStore @@ -25,7 +26,7 @@ async def test_store_crud() -> None: assert not await store.delete("activity_paused") assert await store.get("activity_paused", default=True) is True - value = {"enabled": True, "counts": [1, 2], "label": "example"} + value: JsonValue = {"enabled": True, "counts": [1, 2], "label": "example"} await store.put("example", value) assert await store.get("example") == value await store.put("empty", None) diff --git a/tests/test_library_deletion.py b/tests/test_library_deletion.py index 6813918..f150184 100644 --- a/tests/test_library_deletion.py +++ b/tests/test_library_deletion.py @@ -341,10 +341,11 @@ async def closing(tasks): monkeypatch.setattr(runtime.service.scanner, "drain_library", held) monkeypatch.setattr(runtime.service, "_close", closing) - await runtime.service.request(library_id) - await asyncio.wait_for(started.wait(), 1) - stopping = asyncio.create_task(runtime.service.close()) + stopping = None try: + await runtime.service.request(library_id) + await asyncio.wait_for(started.wait(), 1) + stopping = asyncio.create_task(runtime.service.close()) await asyncio.wait_for(closing_started.wait(), 1) assert not stopping.done() release.set() @@ -359,7 +360,8 @@ async def closing(tasks): assert await session.get(Library, library_id) is None finally: release.set() - await asyncio.gather(stopping, return_exceptions=True) + if stopping is not None: + await asyncio.wait_for(asyncio.gather(stopping, return_exceptions=True), 5) @pytest.mark.asyncio diff --git a/tests/test_lifecycle_safety.py b/tests/test_lifecycle_safety.py index 5e30e37..830a550 100644 --- a/tests/test_lifecycle_safety.py +++ b/tests/test_lifecycle_safety.py @@ -188,11 +188,11 @@ async def blocked_pass() -> int: monkeypatch.setattr(service, "_pass", blocked_pass) monkeypatch.setattr(service, "start", type(service).start.__get__(service)) monkeypatch.setattr(service, "close", type(service).close.__get__(service)) - await service.start() - await asyncio.wait_for(started.wait(), 2) - first = asyncio.create_task(service.close()) - second = None + first = second = None try: + await service.start() + await asyncio.wait_for(started.wait(), 2) + first = asyncio.create_task(service.close()) await asyncio.wait_for(stopping.wait(), 2) second = asyncio.create_task(service.close()) checkpoint = asyncio.Event() @@ -205,6 +205,7 @@ async def blocked_pass() -> int: finally: release.set() await asyncio.wait_for(asyncio.gather(*(task for task in (first, second) if task is not None)), 2) + await asyncio.wait_for(service.close(), 2) @pytest.mark.asyncio @@ -471,7 +472,7 @@ def wait(self, timeout: float | None = None) -> int: server, node = ExitedServer("server"), Node("node") children = iter((server, node)) monkeypatch.setenv("MW_CONFIG_PATH", str(tmp_path)) - (tmp_path / "config.toml").write_text('node_enrollment_token = "test-token"\n') + (tmp_path / "config.toml").write_text('api_key = "test-key"\n') monkeypatch.setattr(run, "resolve_bin", lambda name: name) monkeypatch.setattr(run.subprocess, "Popen", lambda *_args, **_kwargs: next(children)) monkeypatch.setattr(run, "wait_healthy", lambda _port, *, stopped=None: True) diff --git a/tests/test_log_history.py b/tests/test_log_history.py index aaa24af..7c67b2b 100644 --- a/tests/test_log_history.py +++ b/tests/test_log_history.py @@ -12,12 +12,12 @@ from server.core.log_history import install_log_file, query_logs if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Generator from logging.handlers import RotatingFileHandler @contextmanager -def file_logs(path: Path, *, max_bytes: int = 10_000_000) -> Iterator[RotatingFileHandler]: +def file_logs(path: Path, *, max_bytes: int = 10_000_000) -> Generator[RotatingFileHandler]: handler = install_log_file(path, max_bytes=max_bytes) try: yield handler diff --git a/tests/test_node_cli.py b/tests/test_node_cli.py index a69023c..eee3ffa 100644 --- a/tests/test_node_cli.py +++ b/tests/test_node_cli.py @@ -10,7 +10,6 @@ from common.dates import Dates from node import cli from node.cli import ENV_OVERRIDES, _credentials_path, default_config_path, load_settings -from server.core.config import load_config from tests.helpers import make_test_temp_dir @@ -36,7 +35,7 @@ def config(body: str) -> Path: def body(root: Path) -> str: - return f'enrollment_token = "enroll-secret"\nserver_url = "http://hub"\nworkspace_dir = "{root}"\n' + return f'server_url = "http://hub"\nworkspace_dir = "{root}"\n' @pytest.mark.parametrize( @@ -79,26 +78,16 @@ def test_hub_url_errors(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, value: assert "private-value" not in str(error.value) -def test_credentials_load_on_restart(monkeypatch: pytest.MonkeyPatch) -> None: +def test_credentials_restart(monkeypatch: pytest.MonkeyPatch) -> None: root = make_test_temp_dir("workspace") state = make_test_temp_dir("state") monkeypatch.setenv("MW_NODE_STATE_PATH", str(state)) - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "enroll-secret") monkeypatch.setenv("MW_NODE_SERVER_URL", "ws://hub") monkeypatch.setenv("MW_NODE_WORKSPACE_DIR", str(root)) credentials = state / "node-credentials.json" credentials.write_text('{"node_id":"node-a1b2c3d4e5f60708","token":"node-token"}\n', encoding="utf-8") settings = load_settings(config("")) assert (settings.node_id, settings.token) == ("node-a1b2c3d4e5f60708", "node-token") - assert settings.enrollment_token == "enroll-secret" - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "rotated-secret") - settings = load_settings(config("")) - assert (settings.node_id, settings.token) == ("node-a1b2c3d4e5f60708", "node-token") - assert settings.enrollment_token == "rotated-secret" - monkeypatch.delenv("MW_NODE_ENROLLMENT_TOKEN") - settings = load_settings(config("")) - assert (settings.node_id, settings.token) == ("node-a1b2c3d4e5f60708", "node-token") - assert settings.enrollment_token is None @pytest.mark.parametrize("source", ["node", "shared", "xdg", "home"]) @@ -134,29 +123,6 @@ def test_missing_required() -> None: assert "MW_NODE_WORKSPACE_DIR" in str(exc.value) -def test_secret_file(monkeypatch: pytest.MonkeyPatch) -> None: - root = make_test_temp_dir("workspace") - secret = root / "token" - secret.write_text(" secret \n") - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", f"file:{secret}") - settings = load_settings(config(body(root).replace("enroll-secret", "file:missing"))) - assert settings.enrollment_token == "secret" - - -@pytest.mark.parametrize("source", ["literal", "file"]) -def test_shared_enrollment_override(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, source: str) -> None: - (tmp_path / "enrollment").write_text("shared-secret\n") - server_path = tmp_path / "config.toml" - server_path.write_text('node_enrollment_token = "file:missing-server-secret"\n') - node_path = tmp_path / "node.toml" - node_path.write_text(body(tmp_path / "workspace").replace("enroll-secret", "file:missing-node-secret")) - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "shared-secret" if source == "literal" else "file:enrollment") - server = load_config(server_path) - node = load_settings(node_path) - assert server.node_enrollment_token == node.enrollment_token == "shared-secret" - assert node.token == "shared-secret" and node.node_id is None - - def test_unknown_warns(caplog: pytest.LogCaptureFixture) -> None: root = make_test_temp_dir("workspace") with caplog.at_level(logging.WARNING, logger="node.cli"): @@ -173,7 +139,6 @@ def test_config_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: def test_env_optional_settings(monkeypatch: pytest.MonkeyPatch) -> None: root = make_test_temp_dir("workspace") values = { - "enrollment_token": "env-token", "server_url": "https://env-hub", "workspace_dir": str(root), "ffmpeg_bin": "/env/ffmpeg", @@ -189,8 +154,6 @@ def test_env_optional_settings(monkeypatch: pytest.MonkeyPatch) -> None: for field, value in values.items(): monkeypatch.setenv(ENV_OVERRIDES[field], value) settings = load_settings(config(body(root) + 'worker_tags = ["toml"]\nmin_free_gb = 20.0\n')) - assert settings.token == "env-token" - assert settings.enrollment_token == "env-token" assert settings.server_url == "wss://env-hub/api/hub" assert settings.workspace_dir == str(root) assert settings.ffmpeg_bin == "/env/ffmpeg" @@ -209,7 +172,6 @@ def test_optional_toml_preserved() -> None: settings = load_settings( config(body(root) + 'ffmpeg_bin = "custom-ffmpeg"\nworker_tags = ["gpu"]\nheartbeat_s = 8.0\n') ) - assert settings.node_id is None and settings.enrollment_token == "enroll-secret" assert settings.workspace_dir == str(root) assert settings.server_url == "ws://hub/api/hub" assert settings.ffmpeg_bin == "custom-ffmpeg" @@ -220,7 +182,7 @@ def test_optional_toml_preserved() -> None: def test_empty_env_overrides(monkeypatch: pytest.MonkeyPatch) -> None: root = make_test_temp_dir("workspace") - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "") + monkeypatch.setenv("MW_NODE_WORKSPACE_DIR", "") with pytest.raises(ValueError): load_settings(config(body(root))) diff --git a/tests/test_notification_delivery.py b/tests/test_notification_delivery.py index a34781d..fcd14c1 100644 --- a/tests/test_notification_delivery.py +++ b/tests/test_notification_delivery.py @@ -59,7 +59,9 @@ async def test_terminal_enqueue_is_transactional(db: async_sessionmaker[AsyncSes await session.flush() assert await finish(session, job_id, JobStatus.FAILED, now=NOW) deliveries = list((await session.execute(select(NotificationDelivery))).scalars()) - events = list((await session.execute(select(JobEvent).where(JobEvent.job_id == job_id))).scalars()) + events = list( + (await session.execute(select(JobEvent).where(JobEvent.job_id == job_id).order_by(JobEvent.id))).scalars() + ) assert [(item.target_name, item.job_event_id) for item in deliveries] == [("primary", events[-1].id)] await session.rollback() diff --git a/tests/test_playback.py b/tests/test_playback.py index e026db7..0a43d8d 100644 --- a/tests/test_playback.py +++ b/tests/test_playback.py @@ -62,7 +62,7 @@ class PlaybackEnv: @asynccontextmanager -async def playback_context() -> AsyncIterator[PlaybackEnv]: +async def playback_context() -> AsyncGenerator[PlaybackEnv]: lib_dir = make_test_temp_dir("playback-lib") files_dir = make_test_temp_dir("playback-files") workspace = make_test_temp_dir("playback-workspace") diff --git a/tests/test_plugin_host.py b/tests/test_plugin_host.py index 5941472..c532613 100644 --- a/tests/test_plugin_host.py +++ b/tests/test_plugin_host.py @@ -7,7 +7,7 @@ import os import shutil import sys -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from pathlib import Path from typing import Any @@ -188,7 +188,7 @@ async def roundtrip(proc: asyncio.subprocess.Process, request: dict[str, Any]) - @contextlib.asynccontextmanager -async def host_for(roots: list[Path], **kwargs: Any) -> AsyncIterator[PluginHost]: +async def host_for(roots: list[Path], **kwargs: Any) -> AsyncGenerator[PluginHost]: host = PluginHost(roots, **kwargs) try: yield host diff --git a/tests/test_plugin_migrations.py b/tests/test_plugin_migrations.py index e9b1628..6a99097 100644 --- a/tests/test_plugin_migrations.py +++ b/tests/test_plugin_migrations.py @@ -200,7 +200,13 @@ async def test_legacy_flow_http() -> None: assert canonical["nodes"][0]["config_version"] == 2 assert canonical["nodes"][0]["inputs"] == inputs assert canonical["actions"] == [ - {**action, "kind": "notify_external", "plugin_id": "action.notify_external", "config_version": 2} + { + **action, + "kind": "notify_external", + "plugin_id": "action.notify_external", + "config_version": 2, + "path_mappings": [], + } ] assert canonical["_ui"] == raw["_ui"] async with factory() as session: diff --git a/tests/test_plugin_safety.py b/tests/test_plugin_safety.py index 5778d8b..e46fd40 100644 --- a/tests/test_plugin_safety.py +++ b/tests/test_plugin_safety.py @@ -9,7 +9,7 @@ import pytest -from server.plugins.host import PluginHost +from server.plugins.host import EvaluatedDecision, PluginHost from server.plugins import host as plugin_host_module from server.plugins import upload as plugin_upload from server.plugins.upload import UploadError, install_plugin @@ -97,6 +97,7 @@ async def mutate() -> None: if describe: assert result is None else: + assert isinstance(result, EvaluatedDecision) assert result.type == "error" assert result.error is not None if operation == "close": diff --git a/tests/test_queue_staging_limits.py b/tests/test_queue_staging_limits.py index 511a55e..ee6223c 100644 --- a/tests/test_queue_staging_limits.py +++ b/tests/test_queue_staging_limits.py @@ -283,15 +283,16 @@ async def test_bulk_metadata_order(api_env: ApiEnv, monkeypatch) -> None: for index in range(150) ], ) - identifiers = list( + identifiers = sorted( await session.scalars( insert(Job).returning(Job.id), [ {"library_id": library.id, "media_file_id": media_id, "kind": "healthcheck", "status": "failed"} for media_id in media_ids ], - ) - )[::-1] + ), + reverse=True, + ) for identifier in (identifiers[0], identifiers[-1]): job = await session.get(Job, identifier) assert job is not None diff --git a/tests/test_run.py b/tests/test_run.py index 4c54893..000a825 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -28,7 +28,6 @@ def local_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[Pa "MW_NODE_CONFIG", "MW_NODE_SERVER_URL", "MW_NODE_WORKSPACE_DIR", - "MW_NODE_ENROLLMENT_TOKEN", "MW_TZ", "TZ", "IN_CONTAINER", @@ -37,7 +36,7 @@ def local_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[Pa monkeypatch.delenv(key, raising=False) monkeypatch.setenv("MW_CONFIG_PATH", str(tmp_path)) path = tmp_path / "config.toml" - path.write_text('node_enrollment_token = "configured-token"\n') + path.write_text("") yield path @@ -286,7 +285,6 @@ def test_container_permission_failure( (tmp_path / "node.toml").write_text(f'workspace_dir = "{workspace}"\n') monkeypatch.setenv("IN_CONTAINER", "1") monkeypatch.setenv("MW_NODE_SERVER_URL", "http://server:8000") - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "test-enrollment") directory = local_config.parent if blocked == "config" else workspace access = os.access # CI can run as root, so simulate the OS denial without depending on chmod bypasses. @@ -306,49 +304,6 @@ def test_container_permission_failure( assert f'Compose user: "{owner.st_uid}:{owner.st_gid}"' in message -def test_token_override_wins(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("MW_NODE_ENROLLMENT_TOKEN", raising=False) - monkeypatch.delenv("MW_PORT", raising=False) - monkeypatch.setattr(run, "resolve_bin", lambda name: f"/bin/{name}") - envs: list[dict[str, str] | None] = [] - - def fake_popen(argv: list[str], env: dict[str, str] | None = None) -> Any: - envs.append(env) - return FakeProc(argv[0]) - - monkeypatch.setattr(run.subprocess, "Popen", fake_popen) - monkeypatch.setattr(run, "wait_healthy", lambda *a, stopped=None: True) - monkeypatch.setattr(run, "supervise", lambda *a, **k: 0) - assert run.run_all() == 0 - assert envs[0] is None - node_env = envs[1] - assert node_env is not None and node_env["MW_NODE_ENROLLMENT_TOKEN"] == "configured-token" - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "tok") - assert run.run_all() == 0 - node_env = envs[3] - assert node_env is not None and node_env["MW_NODE_ENROLLMENT_TOKEN"] == "tok" - - -def test_local_enrollment(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: - monkeypatch.setenv("MW_CONFIG_PATH", str(tmp_path)) - monkeypatch.delenv("MW_NODE_ENROLLMENT_TOKEN", raising=False) - (tmp_path / "secrets").mkdir() - (tmp_path / "secrets" / "node_enrollment_token").write_text("fleet-token\n", encoding="utf-8") - (tmp_path / "config.toml").write_text('node_enrollment_token = "file:secrets/node_enrollment_token"\n') - monkeypatch.setattr(run, "resolve_bin", lambda name: f"/bin/{name}") - envs = [] - - def fake_popen(argv, env=None): - envs.append(env) - return FakeProc(argv[0]) - - monkeypatch.setattr(run.subprocess, "Popen", fake_popen) - monkeypatch.setattr(run, "wait_healthy", lambda *a, stopped=None: True) - monkeypatch.setattr(run, "supervise", lambda *a, **k: 0) - assert run.run_all() == 0 - assert envs[1]["MW_NODE_ENROLLMENT_TOKEN"] == "fleet-token" - - def test_all_role_dev(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MW_PORT", raising=False) monkeypatch.setattr(run, "resolve_bin", lambda name: f"/bin/{name}") @@ -376,7 +331,6 @@ def fake_popen(argv: list[str], env: dict[str, str] | None = None) -> Any: def test_dev_drains_children(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: monkeypatch.setenv("MW_CONFIG_PATH", str(tmp_path)) - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "test-token") monkeypatch.setattr(run, "resolve_bin", lambda name: name) monkeypatch.setattr(run, "wait_healthy", lambda *args, **kwargs: True) children = [] @@ -420,16 +374,13 @@ def supervise(*args, **kwargs): def test_node_inherits_server(local_config: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: root = tmp_path / "shared" - local_config.write_text(f'workspace_dir = "{root}"\nnode_enrollment_token = "custom-token"\n') + local_config.write_text(f'workspace_dir = "{root}"\n') env = run.node_environment(local_config, port=9000) assert env["MW_NODE_WORKSPACE_DIR"] == str(root) - assert env["MW_NODE_ENROLLMENT_TOKEN"] == "custom-token" assert env["MW_NODE_SERVER_URL"] == "http://127.0.0.1:9000" monkeypatch.setenv("MW_WORKSPACE_DIR", str(tmp_path / "overridden")) - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "overridden-token") env = run.node_environment(local_config, port=9000) assert env["MW_NODE_WORKSPACE_DIR"] == str(tmp_path / "overridden") - assert env["MW_NODE_ENROLLMENT_TOKEN"] == "overridden-token" def test_node_inherits_timezone(local_config: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -443,53 +394,27 @@ def test_node_inherits_timezone(local_config: Path, monkeypatch: pytest.MonkeyPa def test_node_overrides_preserved(local_config: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - (tmp_path / "node.toml").write_text( - 'workspace_dir = "/node/workspace"\nenrollment_token = "node-token"\nserver_url = "https://custom"\n' - ) + (tmp_path / "node.toml").write_text('workspace_dir = "/node/workspace"\nserver_url = "https://custom"\n') env = run.node_environment(local_config, port=8000) - assert all(key not in env for key in ("MW_NODE_ENROLLMENT_TOKEN", "MW_NODE_WORKSPACE_DIR", "MW_NODE_SERVER_URL")) + assert all(key not in env for key in ("MW_NODE_WORKSPACE_DIR", "MW_NODE_SERVER_URL")) monkeypatch.setenv("MW_NODE_WORKSPACE_DIR", "/explicit/workspace") - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "explicit-token") monkeypatch.setenv("MW_NODE_SERVER_URL", "https://explicit") env = run.node_environment(local_config, port=8000) assert env["MW_NODE_WORKSPACE_DIR"] == "/explicit/workspace" - assert env["MW_NODE_ENROLLMENT_TOKEN"] == "explicit-token" assert env["MW_NODE_SERVER_URL"] == "https://explicit" def test_explicit_node_config(local_config: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: path = tmp_path / "custom-node.toml" - path.write_text('enrollment_token = "custom-enrollment"\n') + path.write_text('workspace_dir = "/custom/workspace"\nserver_url = "https://custom"\n') monkeypatch.setenv("MW_NODE_CONFIG", str(path)) env = run.node_environment(local_config, port=8000) - assert "MW_NODE_ENROLLMENT_TOKEN" not in env - assert env["MW_NODE_SERVER_URL"] == "http://127.0.0.1:8000" - - -def test_shared_secret_resolution(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - from node.cli import load_settings - - server_dir = tmp_path / "server" - server_dir.mkdir() - path = server_dir / "config.toml" - path.write_text("") - (server_dir / "enrollment").write_text("shared-secret\n") - node_path = tmp_path / "node.toml" - node_path.write_text('enrollment_token = "file:missing-secret"\n') - monkeypatch.setenv("MW_NODE_CONFIG", str(node_path)) - monkeypatch.setenv("MW_NODE_ENROLLMENT_TOKEN", "file:enrollment") - env = run.node_environment(path, port=8000) - assert env["MW_NODE_ENROLLMENT_TOKEN"] == "shared-secret" - for key in ("MW_NODE_ENROLLMENT_TOKEN", "MW_NODE_SERVER_URL", "MW_NODE_WORKSPACE_DIR"): - monkeypatch.setenv(key, env[key]) - settings = load_settings(node_path) - assert settings.enrollment_token == "shared-secret" - assert settings.server_url == "ws://127.0.0.1:8000/api/hub" + assert all(key not in env for key in ("MW_NODE_WORKSPACE_DIR", "MW_NODE_SERVER_URL")) def test_aio_explicit_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: config = tmp_path / "custom.toml" - config.write_text(f'workspace_dir = "{tmp_path / "custom-workspace"}"\nnode_enrollment_token = "custom-token"\n') + config.write_text(f'workspace_dir = "{tmp_path / "custom-workspace"}"\n') monkeypatch.setenv("MW_PORT", "9000") monkeypatch.setattr(run, "resolve_bin", lambda name: f"/bin/{name}") spawns = [] @@ -503,7 +428,6 @@ def test_aio_explicit_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> assert health_ports == [9100] assert spawns[1][1]["MW_NODE_SERVER_URL"] == "http://127.0.0.1:9100" assert spawns[1][1]["MW_NODE_WORKSPACE_DIR"] == str(tmp_path / "custom-workspace") - assert spawns[1][1]["MW_NODE_ENROLLMENT_TOKEN"] == "custom-token" @pytest.mark.parametrize( diff --git a/tests/test_scanner.py b/tests/test_scanner.py index 673b15b..dbe7061 100644 --- a/tests/test_scanner.py +++ b/tests/test_scanner.py @@ -1330,7 +1330,7 @@ def capture(type_, payload, key=None, *, bulk=False): else: payload = {"action": "pruned", "media_file_id": media.id, "library_id": scan_env.library_id} assert removals == [(payload, f"file.status:file:{media.id}")] - assert job_removals == job_ids + assert sorted(job_removals) == sorted(job_ids) events = [] while subscription.backlog: events.append(await asyncio.wait_for(subscription.next(), 1)) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index f8cb4b3..230ce9e 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -397,7 +397,7 @@ async def test_scan_scope_logging(factory, staging, tmp_path, scan_log, caplog, identifier = await add_library(factory, tmp_path, "logging") scheduler = ScannerService(factory, staging=staging) media = tmp_path / "movie.mkv" - with caplog.at_level(logging.INFO, logger=scheduler_service.__name__): + with caplog.at_level(logging.DEBUG, logger=scheduler_service.__name__): scheduler.queue_scan(identifier, "watch" if scoped else "manual", target_paths={media} if scoped else None) await asyncio.wait_for(scheduler._scan_tasks[identifier], 2) records = [record for record in caplog.records if record.name == scheduler_service.__name__] diff --git a/tests/test_settings.py b/tests/test_settings.py index d5bf053..7909d3a 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -20,11 +20,17 @@ async def test_settings_persist(database_factory: async_sessionmaker[AsyncSessio async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver") as client: assert (await client.get("/api/settings")).status_code == 401 assert (await client.patch("/api/settings", json={"max_probes": 2})).status_code == 401 + assert (await client.patch("/api/settings", json={"node_registration_enabled": False})).status_code == 401 client.headers["Authorization"] = "Bearer test-key" original = (await client.get("/api/settings")).json() response = await client.patch( "/api/settings", - json={"max_probes": 2, "healthcheck_location": "node", "scanner": {"ignore_dirs": ["Extras"]}}, + json={ + "max_probes": 2, + "healthcheck_location": "node", + "scanner": {"ignore_dirs": ["Extras"]}, + "node_registration_enabled": False, + }, ) assert response.status_code == 200, response.text updated = response.json() @@ -33,11 +39,13 @@ async def test_settings_persist(database_factory: async_sessionmaker[AsyncSessio assert app.state.scanner.scanner.ignore_dirs == ["Extras"] assert app.state.events.healthcheck_slots == 0 assert app.state.events.healthchecks_on_nodes + assert not app.state.hub.registration_enabled assert await app.state.store.get(SETTINGS_KEY) == updated restarted = create_app(database_factory, config=config, cache_root=tmp_path / "restarted") async with restarted.router.lifespan_context(restarted): assert restarted.state.settings.current.model_dump(mode="json") == updated assert restarted.state.scanner.scanner.ignore_dirs == ["Extras"] + assert not restarted.state.hub.registration_enabled @pytest.mark.asyncio diff --git a/tests/test_staging.py b/tests/test_staging.py index a49249c..6cbdb50 100644 --- a/tests/test_staging.py +++ b/tests/test_staging.py @@ -7,7 +7,7 @@ import threading import uuid from contextlib import asynccontextmanager -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from datetime import UTC, datetime, timedelta T0 = datetime(2026, 1, 1, 12, tzinfo=UTC) @@ -43,6 +43,11 @@ from tests.helpers import request_retry, database_context, make_test_disk_path, make_test_temp_dir +@pytest.fixture(autouse=True) +def workspace_capacity(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("server.staging.shutil.disk_usage", lambda _: SimpleNamespace(free=20_000_000_000)) + + @pytest_asyncio.fixture async def db() -> AsyncIterator[async_sessionmaker[AsyncSession]]: # Parallel staging jobs need independent database transactions. @@ -206,7 +211,7 @@ async def prune(): @asynccontextmanager -async def subscribed() -> AsyncIterator[tuple[EventBus, Subscription]]: +async def subscribed() -> AsyncGenerator[tuple[EventBus, Subscription]]: bus = EventBus() sub = bus.subscribe() try: diff --git a/tests/test_staging_safety.py b/tests/test_staging_safety.py index 63f6794..5e8c4a0 100644 --- a/tests/test_staging_safety.py +++ b/tests/test_staging_safety.py @@ -382,8 +382,8 @@ async def test_continuation_refuses_changes(staged_env, change: str): @pytest.mark.asyncio -@pytest.mark.parametrize("state", ["drained", "exhausted", "unconfirmed", "unread", "transcode", "unavailable"]) -async def test_prune_healthcheck_logs(staged_env, state: str): +@pytest.mark.parametrize("workload", ["healthcheck", "transcode"]) +async def test_prune_healthcheck_logs(staged_env, workload: str): factory, service, job_id, source = staged_env relative = workspace.create_attempt_layout(service.workspace_dir, job_id, 1, ".mkv")["log"] log = service.workspace_dir / relative @@ -396,18 +396,14 @@ async def test_prune_healthcheck_logs(staged_env, state: str): attempt_number=1, status=JobStatus.FAILED, result="interrupted", - workload="transcode" if state == "transcode" else "healthcheck", + workload=workload, log_path=relative, - log_data=None if state == "unread" else log.read_bytes(), + log_data=log.read_bytes(), finished_at=utcnow(), - error="max attempts exceeded" - if state == "exhausted" - else ("healthcheck process exit is unconfirmed" if state == "unconfirmed" else "requeued"), + error="requeued", ) session.add(attempt) await session.flush() - if state == "unavailable": - session.add(JobEvent(job_id=job_id, attempt_id=attempt.id, type="healthcheck_log_unavailable")) await session.commit() eligible = False assert await service.prune_attempt(attempt.id) is eligible @@ -415,7 +411,7 @@ async def test_prune_healthcheck_logs(staged_env, state: str): async with factory() as session: saved = await session.get(JobAttempt, attempt.id) assert saved is not None and saved.result == "interrupted" - assert saved.log_data == (None if state == "unread" else b"healthcheck log") + assert saved.log_data == b"healthcheck log" assert source.read_bytes() == b"source" diff --git a/tests/test_workspace_recheck.py b/tests/test_workspace_recheck.py index b850d60..6a25ee9 100644 --- a/tests/test_workspace_recheck.py +++ b/tests/test_workspace_recheck.py @@ -287,7 +287,11 @@ async def continue_subtitle(*args, **kwargs) -> str: job = await session.get(Job, job_id) assert job is not None and job.status == JobStatus.RUNNING stages = list( - (await session.scalars(select(JobEvent).where(JobEvent.job_id == job_id, JobEvent.type == "stage"))) + await session.scalars( + select(JobEvent) + .where(JobEvent.job_id == job_id, JobEvent.type == "stage") + .order_by(JobEvent.id) + ) ) assert len(stages) == 2 and stages[-1].payload["plan"] == plan.model_dump(mode="json") finally: diff --git a/ui/app/components/NodeUsageGauges.vue b/ui/app/components/NodeUsageGauges.vue index 2588750..ec4cc4b 100644 --- a/ui/app/components/NodeUsageGauges.vue +++ b/ui/app/components/NodeUsageGauges.vue @@ -2,7 +2,7 @@

Last known

- + +
+
+ + + + + + + +
+ Add mapping +