diffusers-workflow 0.4.0__py3-none-any.whl

dw_mcp/__init__.py

@@ -0,0 +1,6 @@
+"""MCP server for diffusers-workflow.
+
+A stdio MCP server that is an HTTP client of a running `dw.serve`. It owns
+no job state and no GPU worker - every tool is a call against the REST API
+that the web UI already uses.
+"""

dw_mcp/__main__.py

@@ -0,0 +1,133 @@
+"""`python -m dw_mcp` / `dw-mcp`: serve the tool surface over stdio."""
+
+import argparse
+import sys
+
+import httpx
+
+from dw_mcp.client import (
+    DwApiError,
+    DwClient,
+    is_loopback_url,
+    resolve_base_url,
+    resolve_token,
+)
+from dw_mcp.server import build_server
+
+
+def _refuse(message):
+    print(f"dw-mcp: {message}", file=sys.stderr)
+    return 2
+
+
+def main(argv=None, transport=None):
+    parser = argparse.ArgumentParser(
+        prog="dw-mcp",
+        description="MCP server for diffusers-workflow. Requires a running "
+        "dw.serve - start one with `dw-serve` first.",
+    )
+    parser.add_argument(
+        "--url",
+        default=None,
+        help="Base URL of the running dw.serve "
+        "(default: $DW_MCP_URL, else http://127.0.0.1:8765)",
+    )
+    parser.add_argument(
+        "--token",
+        default=None,
+        help="Bearer token the dw.serve was started with, if any "
+        "(default: $DW_API_TOKEN). Required when --url is not loopback.",
+    )
+    parser.add_argument(
+        "--workspace",
+        default=None,
+        help="Which of the server's workspaces to work in (default: "
+        "$DW_MCP_WORKSPACE, else the server's default). A name on the "
+        "server, not a directory here - list_workspaces shows them",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=30.0,
+        help="Seconds to wait on any one API request (default: 30)",
+    )
+    parser.add_argument(
+        "--no-probe",
+        action="store_true",
+        default=False,
+        help="Skip the startup GET /api/health that confirms the server is "
+        "reachable and the token is accepted",
+    )
+    args = parser.parse_args(argv)
+
+    base_url = resolve_base_url(args.url)
+    token = resolve_token(args.token)
+
+    # A remote dw.serve with no token would let anyone on that network run
+    # workflows as this user; refusing here is the client-side half of the
+    # server's own non-loopback-without-token warning.
+    if not is_loopback_url(base_url) and not token:
+        return _refuse(
+            f"{base_url} is not a loopback address and no token is set. "
+            "A dw.serve reachable from other machines must be started with "
+            "--token / DW_API_TOKEN, and the same token passed here with "
+            "--token or DW_API_TOKEN."
+        )
+
+    client = DwClient(
+        base_url=base_url,
+        timeout=args.timeout,
+        token=token,
+        transport=transport,
+        workspace=args.workspace,
+    )
+    try:
+        if not args.no_probe:
+            # A remote URL or token that is wrong is a misconfiguration the
+            # user must fix before anything works, so it is fatal. On
+            # loopback it usually means "dw.serve is not up yet", which a
+            # tool call reports for itself with the same message - warn and
+            # serve anyway rather than making the agent restart us.
+            code = _probe(client, base_url, required=not is_loopback_url(base_url))
+            if code:
+                return code
+        build_server(client).run(transport="stdio")
+    finally:
+        client.close()
+    return 0
+
+
+def _probe(client, base_url, required=True):
+    """One GET /api/health so a wrong URL or token fails here, once, with a
+    message - not on every tool call as an unexplained 401.
+
+    Returns 2 when the probe failed and `required`; 0 otherwise. A failed
+    probe always prints, so a loopback server that is merely not up yet
+    still says so.
+    """
+    try:
+        health = client.get_json("/api/health")
+    except (DwApiError, httpx.HTTPError) as e:
+        text = str(e)
+        if "401" in text or "token" in text.lower():
+            code = _refuse(
+                f"dw.serve at {base_url} requires a bearer token and rejected "
+                f"the one given (or none was given): {text}"
+            )
+        else:
+            code = _refuse(f"could not reach dw.serve at {base_url}: {text}")
+        return code if required else 0
+    print(
+        "dw-mcp: connected to {host} (dw {version}, {device}) at {url}".format(
+            host=health.get("hostname", "?"),
+            version=health.get("version", "?"),
+            device=health.get("device", "?"),
+            url=base_url,
+        ),
+        file=sys.stderr,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

dw_mcp/assets.py

@@ -0,0 +1,336 @@
+"""The asset library: the input media a workflow reaches by an 'asset:'
+reference, and the one way to get a local file into it.
+
+An agent authoring remotely can name assets that already exist on the box,
+but until it can put bytes there it can only write workflows it has no way
+to supply inputs for. `upload_asset` is that path: the file's bytes go up as
+the request body - the same route the browser's file picker uses - and what
+comes back is the reference, not a path, because the reference is what a
+workflow carries and the path means nothing on the machine the agent is on.
+"""
+
+import base64
+import os
+
+from dw_mcp.client import DwApiError, api_path
+
+# Twin of the server's own limit (dw/server/app.py). Checked here as well so
+# a 200MB file fails before it is read and pushed, not after
+MAX_UPLOAD_BYTES = 200 * 1024 * 1024
+
+# `content` bytes arrive base64-encoded in the tool call itself rather than
+# read off a disk, so they compete with everything else in the call for the
+# caller's context budget - a few MB is enough for a voice sample or a small
+# image, and nowhere near enough to use this as a bulk-transfer path (#203)
+MAX_INLINE_UPLOAD_BYTES = 4 * 1024 * 1024
+
+# What the library holds, and what the upload route accepts. Duplicated from
+# dw/security.py rather than imported: importing anything under dw/ pulls in
+# torch, which this pure HTTP client must not do
+ALLOWED_UPLOAD_EXTENSIONS = frozenset(
+    {
+        ".jpg",
+        ".jpeg",
+        ".png",
+        ".gif",
+        ".bmp",
+        ".webp",
+        ".mp4",
+        ".avi",
+        ".mkv",
+        ".mov",
+        ".webm",
+        ".wav",
+        ".mp3",
+        ".flac",
+        ".ogg",
+    }
+)
+
+
+def _remote_roots(client):
+    """The directories a remote read is confined to, or None when local.
+
+    The mirror of media.py's `_remote_root` (#113), for the other direction.
+    Only the mounted MCP surface is remote: there `upload_asset` runs inside
+    dw.serve, so `file_path` names a file on the operator's box rather than
+    on the calling agent's machine, and an unconfined read is an arbitrary
+    file read plus a path-existence oracle (#138). A stdio `dw-mcp` returns
+    None and keeps reading whatever the user can, because there "local file"
+    is genuinely their own.
+    """
+    if not getattr(client, "mounted", False):
+        return None
+
+    directories = (client.get_json("/api/server").get("directories")) or {}
+    roots = []
+    for key in ("workspace", "workflows", "assets", "outputs", "prompts"):
+        value = directories.get(key)
+        if not value:
+            continue
+        resolved = os.path.normpath(
+            os.path.realpath(os.path.abspath(os.path.expanduser(str(value))))
+        )
+        if resolved not in roots:
+            roots.append(resolved)
+    if not roots:
+        raise DwApiError(
+            "This server cannot say which directories it works in, so it "
+            "will not read a file off its own disk for you. Upload the "
+            "bytes through the web UI's file picker, pass them inline to "
+            "upload_asset as content=, or keep a generated file with "
+            "keep_output."
+        )
+    return roots
+
+
+def _confine_source(path, roots, named):
+    """Refuse a source outside `roots`, before anything looks at the file.
+
+    Ordered ahead of the existence and extension checks on purpose: a
+    refusal that depends on whether the file is there turns the tool into a
+    path-existence oracle for the whole box, which is the condition this
+    closes as much as the read itself (#138). Containment is on the resolved
+    real path, so a symlink cannot carry the read out.
+    """
+    probe = path
+    while not os.path.exists(probe) and os.path.dirname(probe) != probe:
+        probe = os.path.dirname(probe)
+    resolved = os.path.normpath(
+        os.path.join(os.path.realpath(probe), os.path.relpath(path, probe))
+    )
+    if any(resolved == root or resolved.startswith(root + os.sep) for root in roots):
+        return
+    raise DwApiError(
+        f"Refusing to read {named} - this MCP endpoint is served by "
+        f"dw.serve, so the file would be read off the server, where a "
+        f"source is confined to the directories it works in "
+        f"({', '.join(roots)}). A file that is already there is reachable "
+        f"as an 'asset:' reference; to put a new one there when it exists "
+        f"only on your own machine, call upload_asset with content= (its "
+        f"bytes, base64-encoded) instead of file_path, upload it through "
+        f"the web UI's file picker, or promote a generated file with "
+        f"keep_output."
+    )
+
+
+# What a compact asset entry keeps: enough to name and choose one, nothing
+# that only opening or previewing it needs
+ASSET_SUMMARY_FIELDS = ("name", "reference", "kind", "size", "origin")
+
+
+def _summarised_asset_entries(entries):
+    return [
+        {key: entry.get(key) for key in ASSET_SUMMARY_FIELDS if key in entry}
+        for entry in entries
+        if isinstance(entry, dict)
+    ]
+
+
+def list_assets(client, detail=False):
+    """The input media on the server, each with the 'asset:' reference a
+    workflow argument carries.
+
+    Spans the whole search path: each entry's 'origin' says whether it is
+    this workspace's own ('workspace'), the library every workspace shares
+    ('common'), or one a read-only examples tree brought with it.
+    'libraries' names the roots searched, in order, with 'writable' saying
+    which one an upload or delete can reach. 'shadowed' lists names a
+    nearer library hides - present in a farther root but not resolvable by
+    'asset:', since the nearer file wins; a shadowed entry carries no 'url'
+    because that URL would serve the file that shadows it.
+
+    Compact by default - each `assets`/`shadowed` entry cut to
+    `name`/`reference`/`kind`/`size`/`origin` - because a library that has
+    accumulated months of episode and regression assets answers in the tens
+    of KB otherwise for a question that is usually "what's available and
+    under what name" (#249). Pass `detail=True` for each entry's `folder`,
+    `mtime` and `url` too.
+    """
+    result = client.get_json("/api/assets")
+    if detail:
+        return result
+    assets = result.get("assets")
+    shadowed = result.get("shadowed")
+    if not isinstance(assets, list):
+        return result
+    return {
+        **result,
+        "assets": _summarised_asset_entries(assets),
+        "shadowed": _summarised_asset_entries(shadowed)
+        if isinstance(shadowed, list)
+        else shadowed,
+        "note": (
+            "Compact: each entry is name/reference/kind/size/origin only. "
+            "Call list_assets(detail=True) for folder, mtime and url too."
+        ),
+    }
+
+
+def delete_asset(client, name):
+    """Remove one file from the asset library.
+
+    Deletes from whichever library holds it - the workspace's own before
+    the shared one, the order 'asset:' resolves in. An asset from a
+    read-only examples library answers 403.
+    """
+    return client.delete_json(api_path("api", "assets", name))
+
+
+def keep_output(
+    client, name, asset_name=None, overwrite=False, shared=False, workspace=None
+):
+    """Keep a generated file as an input asset, under a stable name.
+
+    A run's files are named by the run that made them, which is the wrong
+    thing to build on: 'latest' moves and a pinned run id breaks when
+    outputs are pruned. Keeping one gives it an 'asset:' name that stays
+    put, so a later workflow can rely on it.
+
+    The copy happens on the server, inside the workspace - downloading a
+    render here only to upload it back would move the bytes twice for
+    nothing.
+
+    `shared` keeps it in the library every workspace under the server's
+    root shares instead, which is where something a later episode in its
+    own workspace has to reach belongs.
+    """
+    return client.post_json(
+        "/api/assets/keep",
+        {
+            "name": name,
+            "asset_name": asset_name,
+            "overwrite": overwrite,
+            "shared": shared,
+        },
+        workspace=workspace,
+    )
+
+
+def upload_asset(client, file_path=None, content=None, asset_name=None, shared=False):
+    """Put an image, video or audio file into the server's asset library and
+    get back the reference a workflow can use.
+
+    Takes exactly one of two sources:
+
+    `file_path` is read from the machine this MCP server runs on, which is
+    not necessarily the machine dw.serve runs on - that is the point of the
+    tool. Over a `dw.serve --mcp` endpoint that machine *is* the server, so
+    there `file_path` is confined to the directories the server works in,
+    and the refusal comes before the file is looked for so it cannot be used
+    to probe which paths exist (#138). A stdio `dw-mcp` is unconfined,
+    because there the file really is the caller's own.
+
+    `content` is the file's bytes, base64-encoded, for the case `file_path`
+    cannot reach: a headless agent driving a mounted `dw.serve --mcp`
+    endpoint has no filesystem in common with that server, so a voice
+    sample or small image that exists only on the agent's own machine has
+    no path either side can name. No path is confined or read for this
+    path - the bytes just go straight to the upload route - so it does not
+    reopen #138. It is capped well under `file_path`'s limit
+    (MAX_INLINE_UPLOAD_BYTES, a few MB) because these bytes ride in the tool
+    call itself rather than being streamed off disk, and `asset_name` is
+    required so the upload has a name and extension to be validated against.
+
+    `asset_name` is the name it is stored under - 'cast/priya-voice.wav'
+    rather than the random one an upload gets by default. A recurring cast
+    referenced as 'asset:uploads/084eaecc....wav' in every workflow cannot
+    be told apart without opening each file, which is the whole reason to
+    name one (2026-09-11). With `file_path`, the extension comes from the
+    uploaded file when the name has none; with `content`, `asset_name` must
+    carry its own extension since there is no file to take one from.
+
+    `shared` puts it in the library every workspace shares rather than in
+    the session's own, which is what a recurring cast needs: assets are
+    per workspace, so a cast uploaded while making episode one was
+    invisible from the workspace episode four was made in.
+    """
+    if (file_path is None) == (content is None):
+        raise DwApiError("Pass exactly one of file_path or content.")
+
+    if content is not None:
+        return _upload_inline(client, content, asset_name=asset_name, shared=shared)
+
+    path = os.path.abspath(os.path.expanduser(str(file_path)))
+    roots = _remote_roots(client)
+    if roots is not None:
+        _confine_source(path, roots, file_path)
+    if not os.path.isfile(path):
+        raise DwApiError(f"No such file: {file_path}")
+
+    extension = os.path.splitext(path)[1].lower()
+    if extension not in ALLOWED_UPLOAD_EXTENSIONS:
+        raise DwApiError(
+            f"{os.path.basename(path)} is not a kind the asset library takes "
+            f"({', '.join(sorted(ALLOWED_UPLOAD_EXTENSIONS))})."
+        )
+
+    size = os.path.getsize(path)
+    if size > MAX_UPLOAD_BYTES:
+        raise DwApiError(
+            f"{os.path.basename(path)} is {size} bytes, over the "
+            f"{MAX_UPLOAD_BYTES} byte upload limit."
+        )
+
+    try:
+        with open(path, "rb") as handle:
+            body = handle.read()
+    except OSError as e:
+        raise DwApiError(f"Could not read {file_path}: {e}")
+
+    params = {"filename": os.path.basename(path)}
+    if asset_name:
+        params["asset_name"] = asset_name
+    if shared:
+        params["shared"] = "true"
+    result = client.post_bytes("/api/uploads", body, params=params)
+    # 'path' from a server with no asset library is an absolute path on that
+    # machine; from one with a library it is already the reference. Report
+    # whichever it gave, named for what it is
+    return {
+        "reference": result.get("path"),
+        "url": result.get("url"),
+        "uploaded": os.path.basename(path),
+        "size": size,
+    }
+
+
+def _upload_inline(client, content, asset_name=None, shared=False):
+    """The `content=` path of upload_asset - bytes with no path behind
+    them, so nothing here is confined or read off any disk (#203)."""
+    if not asset_name:
+        raise DwApiError(
+            "content requires asset_name, so the upload has a name and an "
+            "extension to validate - there is no file to take either from."
+        )
+
+    extension = os.path.splitext(asset_name)[1].lower()
+    if extension not in ALLOWED_UPLOAD_EXTENSIONS:
+        raise DwApiError(
+            f"{asset_name} is not a kind the asset library takes "
+            f"({', '.join(sorted(ALLOWED_UPLOAD_EXTENSIONS))})."
+        )
+
+    try:
+        body = base64.b64decode(content, validate=True)
+    except Exception as e:
+        raise DwApiError(f"content could not be decoded as base64: {e}")
+
+    if len(body) > MAX_INLINE_UPLOAD_BYTES:
+        raise DwApiError(
+            f"content is {len(body)} bytes, over the "
+            f"{MAX_INLINE_UPLOAD_BYTES} byte limit for an inline upload. A "
+            f"file this large should be reached by file_path instead, from "
+            f"a machine that has it on disk."
+        )
+
+    params = {"filename": asset_name, "asset_name": asset_name}
+    if shared:
+        params["shared"] = "true"
+    result = client.post_bytes("/api/uploads", body, params=params)
+    return {
+        "reference": result.get("path"),
+        "url": result.get("url"),
+        "uploaded": asset_name,
+        "size": len(body),
+    }

dw_mcp/authoring.py

@@ -0,0 +1,114 @@
+"""Writing workflows. Validation is free and comes first; saving overwrites,
+so it is annotated destructive at the tool layer.
+
+Path confinement is the server's job (dw/security.py already refuses
+traversal and anything outside the workflow directory). Nothing here
+re-implements it - a second, subtly different check is how the two drift.
+"""
+
+from dw_mcp.client import DwApiError, api_path, coerce_json_object
+
+
+def validate_workflow(
+    client,
+    workflow=None,
+    name=None,
+    inline_workflow=None,
+    workflow_path=None,
+    workspace=None,
+    arguments=None,
+):
+    """Schema- and signature-check a workflow without queuing anything. This
+    is free (no GPU work) and should be called before any run or save. Give
+    either an inline definition (`workflow`, or `inline_workflow` - the same
+    thing `run_workflow` calls it) or the name of a stored one (`name`, or
+    `workflow_path` - the same thing `run_workflow` calls it), as
+    `list_workflows` reports it. A validated document can be handed straight
+    to `run_workflow` under either spelling.
+
+    `arguments` is the same dict `run_workflow` takes, checked against the
+    variables the workflow declares and against this workspace's libraries.
+    Without it the answer is about the stored definition and its stock
+    defaults - which is everything except the part the caller wrote
+    (2026-09-11).
+
+    `plan.cached_steps` is evaluated against the pinned workspace's own
+    `outputs/` - a run sitting in a different workspace, however identical
+    its arguments and seed, does not count as a hit. Pin `workspace` to the
+    one an earlier run actually used if you want to see it credited."""
+    workflow = coerce_json_object(workflow, "workflow")
+    inline_workflow = coerce_json_object(inline_workflow, "inline_workflow")
+    if workflow is not None and inline_workflow is not None:
+        raise DwApiError(
+            "`workflow` and `inline_workflow` are the same thing - provide only one."
+        )
+    if name is not None and workflow_path is not None:
+        raise DwApiError(
+            "`name` and `workflow_path` are the same thing - provide only one."
+        )
+    inline = workflow if workflow is not None else inline_workflow
+    stored = name if name is not None else workflow_path
+    if (inline is None) == (stored is None):
+        raise DwApiError(
+            "Provide exactly one of `workflow`/`inline_workflow` (an inline "
+            "definition) or `name`/`workflow_path` (a stored workflow)."
+        )
+    params = {"workspace": workspace} if workspace else None
+    payload = {"workflow_path": stored} if inline is None else {"workflow": inline}
+    if arguments is not None:
+        # An explicit {} still means "check a run with no values supplied" -
+        # distinct from omitting arguments entirely, which means "check the
+        # document" (#364). A falsy-but-not-None check here would drop that
+        # distinction before it ever reaches the server.
+        payload["arguments"] = arguments
+    # The server resolves a name against its own workflow directory, so
+    # validation sees the same base directory a run would
+    return client.post_json("/api/validate", payload, params=params)
+
+
+def save_workflow(client, name, workflow=None, patch=None):
+    """Write a workflow into the server's writable workflow directory,
+    overwriting any file already under that name there. A name that resolves
+    to one of the server's read-only sources (an examples directory) is not
+    overwritten: the copy lands in the writable directory and shadows it from
+    then on. The server validates before writing.
+
+    Give exactly one of `workflow` (the full document to write) or `patch`
+    (a JSON Merge Patch, RFC 7396, applied to the currently stored
+    definition): a dict whose keys overwrite the stored ones, recursively
+    for nested dicts, so bumping one argument means sending just that
+    argument rather than the whole document. A key set to `null` deletes it.
+    A list replaces the stored list whole - a merge patch has no notion of
+    list position, so changing one entry of a `for_each` list still means
+    sending that whole list."""
+    if (workflow is None) == (patch is None):
+        raise DwApiError(
+            "Provide exactly one of `workflow` (a full replacement) or "
+            "`patch` (a JSON merge patch onto the stored version)."
+        )
+    workflow = coerce_json_object(workflow, "workflow")
+    patch = coerce_json_object(patch, "patch")
+    if patch is not None:
+        current = client.get_json(api_path("api", "workflows", name))
+        workflow = _merge_patch(current, patch)
+    return client.put_json(api_path("api", "workflows", name), {"workflow": workflow})
+
+
+def _merge_patch(target, patch):
+    """RFC 7396 JSON Merge Patch: each dict key in `patch` merges
+    recursively into `target`; any other value replaces `target` outright;
+    `None` deletes the key from the result."""
+    if not isinstance(patch, dict):
+        return patch
+    result = dict(target) if isinstance(target, dict) else {}
+    for key, value in patch.items():
+        if value is None:
+            result.pop(key, None)
+        else:
+            result[key] = _merge_patch(result.get(key), value)
+    return result
+
+
+def delete_workflow(client, name):
+    """Remove a workflow from the server's workflow directory."""
+    return client.delete_json(api_path("api", "workflows", name))