provekit 0.1.0__py3-none-any.whl

provekit/services/providers/mcp_client.py

@@ -0,0 +1,289 @@
+"""Generic MCP client — the best test client for Model Context Protocol servers.
+
+Transports:      Streamable HTTP  ·  stdio (spawn a local server process)
+Spec generations: stateful 2025-11-25 (initialize handshake + Mcp-Session-Id)
+                  stateless 2026-07-28 (no session, server/discover) · auto-detect
+Auth:            static bearer (via headers) · OAuth 2.1 client-credentials grant;
+                 a 401 surfaces the resource-metadata hint so the user knows what's needed.
+Discovery:       tools · resources · prompts, all cursor-paginated.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import subprocess
+import threading
+
+import httpx
+
+from ..netguard import guard_stdio, guard_url
+
+_ACCEPT = "application/json, text/event-stream"
+_STATEFUL = "2025-11-25"
+_STATELESS = "2026-07-28"
+_lock = threading.Lock()
+_id = 0
+
+
+def _next_id() -> int:
+    global _id
+    with _lock:
+        _id += 1
+        return _id
+
+
+class MCPError(RuntimeError):
+    pass
+
+
+def _fetch_oauth_token(oauth: dict) -> str:
+    """Exchange client credentials for a bearer token (OAuth 2.1 client_credentials).
+    Interactive auth-code/PKCE flows need a browser and are out of scope here."""
+    token_url = oauth.get("token_url")
+    if not token_url:
+        raise MCPError("oauth config needs a token_url")
+    guard_url(token_url)
+    data = {"grant_type": "client_credentials",
+            "client_id": oauth.get("client_id", ""),
+            "client_secret": oauth.get("client_secret", "")}
+    if oauth.get("scope"):
+        data["scope"] = oauth["scope"]
+    if oauth.get("resource"):  # RFC 8707 resource indicator
+        data["resource"] = oauth["resource"]
+    r = httpx.post(token_url, data=data, timeout=30, follow_redirects=False)
+    r.raise_for_status()
+    tok = r.json().get("access_token")
+    if not tok:
+        raise MCPError("no access_token in the OAuth token response")
+    return tok
+
+
+class _HTTPTransport:
+    def __init__(self, url, headers, timeout, stateful: bool):
+        guard_url(url)
+        self.url = url.rstrip("/")
+        self.headers = {"Content-Type": "application/json", "Accept": _ACCEPT, **(headers or {})}
+        self.stateful = stateful
+        self.session_id: str | None = None
+        self._client = httpx.Client(timeout=timeout, follow_redirects=False)
+
+    def _hdrs(self):
+        h = dict(self.headers)
+        if self.session_id:
+            h["Mcp-Session-Id"] = self.session_id
+        return h
+
+    def request(self, payload) -> dict:
+        resp = self._client.post(self.url, json=payload, headers=self._hdrs())
+        sid = resp.headers.get("mcp-session-id")
+        if sid:
+            self.session_id = sid
+        if resp.status_code == 401:
+            meta = resp.headers.get("www-authenticate", "")
+            raise MCPError(f"401 Unauthorized — server requires OAuth. {meta}".strip())
+        resp.raise_for_status()
+        return _parse_http(resp)
+
+    def notify(self, payload) -> None:
+        self._client.post(self.url, json=payload, headers=self._hdrs())
+
+    def close(self):
+        self._client.close()
+
+
+class _StdioTransport:
+    """Newline-delimited JSON-RPC over a spawned process's stdin/stdout."""
+    def __init__(self, command, args, env, timeout):
+        guard_stdio()  # RCE gate: never spawn a local process in hosted mode
+        self.proc = subprocess.Popen(
+            [command, *(args or [])], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL, text=True, bufsize=1, env=env)
+
+    def request(self, payload) -> dict:
+        assert self.proc.stdin and self.proc.stdout
+        self.proc.stdin.write(json.dumps(payload) + "\n")
+        self.proc.stdin.flush()
+        want = payload.get("id")
+        for line in self.proc.stdout:
+            line = line.strip()
+            if not line:
+                continue
+            msg = json.loads(line)
+            if msg.get("id") == want:
+                return msg
+        raise MCPError("stdio server closed before replying")
+
+    def notify(self, payload) -> None:
+        assert self.proc.stdin
+        self.proc.stdin.write(json.dumps(payload) + "\n")
+        self.proc.stdin.flush()
+
+    def close(self):
+        try:
+            self.proc.terminate()
+            self.proc.wait(timeout=5)
+        except Exception:
+            self.proc.kill()
+
+
+def _parse_http(resp: httpx.Response) -> dict:
+    ctype = resp.headers.get("content-type", "")
+    if ctype.startswith("text/event-stream"):
+        for line in resp.text.splitlines():
+            if line.startswith("data:"):
+                chunk = json.loads(line[5:].strip())
+                if isinstance(chunk, dict) and ("result" in chunk or "error" in chunk):
+                    return chunk
+        raise MCPError("no JSON-RPC message in SSE stream")
+    return resp.json()
+
+
+class MCPSession:
+    """A short-lived MCP session over one transport.
+
+    HTTP:  MCPSession(url, headers=..., spec="auto"|"2025-11-25"|"2026-07-28", oauth=...)
+    stdio: MCPSession(command="python", args=["server.py"], env=...)
+    """
+
+    # Class-level defaults so a session built via __new__ (tests inject a fake transport)
+    # still has them.
+    _inited = False
+    _keep_open = False
+
+    def __init__(self, url: str | None = None, headers: dict | None = None, timeout: float = 30, *,
+                 command: str | None = None, args: list | None = None, env: dict | None = None,
+                 spec: str = "auto", oauth: dict | None = None):
+        headers = dict(headers or {})
+        if oauth:
+            headers.setdefault("Authorization", f"Bearer {_fetch_oauth_token(oauth)}")
+        self.spec = spec
+        self._inited = False
+        self._keep_open = False
+        self._stdio = command is not None
+        if self._stdio:
+            self._t = _StdioTransport(command, args, env, timeout)
+            self._stateful = True  # stdio is always a stateful session
+        else:
+            if not url:
+                raise MCPError("MCP session needs a url (HTTP) or command (stdio)")
+            self._stateful = spec != _STATELESS
+            self._t = _HTTPTransport(url, headers, timeout, self._stateful)
+
+    def _rpc(self, method: str, params: dict | None = None) -> dict:
+        payload = {"jsonrpc": "2.0", "id": _next_id(), "method": method, "params": params or {}}
+        data = self._t.request(payload)
+        if "error" in data:
+            raise MCPError(str(data["error"].get("message", "MCP error")))
+        return data.get("result", {})
+
+    def open(self) -> "MCPSession":
+        """Keep the transport open across calls until close(). list_tools()/call_tool()
+        otherwise close after each call, so a tool-calling loop would re-handshake every
+        time — and for stdio, spawn and kill a subprocess every time."""
+        self._keep_open = True
+        return self
+
+    def close(self) -> None:
+        self._keep_open = False
+        self._close()
+
+    @contextlib.contextmanager
+    def session(self):
+        """open() for the duration of a block. The session is spent afterwards, exactly as
+        it is after a bare call today."""
+        try:
+            yield self.open()
+        finally:
+            self.close()
+
+    def _close(self) -> None:
+        self._inited = False
+        self._t.close()
+
+    def _close_unless_kept(self) -> None:
+        if not self._keep_open:
+            self._close()
+
+    def _init(self) -> None:
+        if self._inited or not self._stateful:
+            return  # stateless generation: no handshake, no session id
+        try:
+            self._rpc("initialize", {
+                "protocolVersion": _STATEFUL if self.spec == "auto" else self.spec,
+                "capabilities": {},
+                "clientInfo": {"name": "provekit", "version": "0.1.0"},
+            })
+            self._t.notify({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}})
+            self._inited = True
+        except Exception:
+            # auto: a server that rejects initialize is treated as stateless.
+            if self.spec == "auto":
+                self._stateful = False
+
+    def _paginate(self, method: str, key: str) -> list[dict]:
+        items, cursor = [], None
+        while True:
+            res = self._rpc(method, {"cursor": cursor} if cursor else {})
+            items += res.get(key, [])
+            cursor = res.get("nextCursor")
+            if not cursor:
+                return items
+
+    def list_tools(self) -> list[dict]:
+        try:
+            self._init()
+            tools = self._paginate("tools/list", "tools")
+            return [{"name": t["name"], "description": t.get("description", ""),
+                     "input_schema": t.get("inputSchema") or {}} for t in tools]
+        finally:
+            self._close_unless_kept()
+
+    def list_resources(self) -> list[dict]:
+        try:
+            self._init()
+            res = self._paginate("resources/list", "resources")
+            return [{"uri": r.get("uri"), "name": r.get("name", ""),
+                     "mime_type": r.get("mimeType", "")} for r in res]
+        finally:
+            self._close_unless_kept()
+
+    def list_prompts(self) -> list[dict]:
+        try:
+            self._init()
+            pr = self._paginate("prompts/list", "prompts")
+            return [{"name": p["name"], "description": p.get("description", ""),
+                     "arguments": p.get("arguments", [])} for p in pr]
+        finally:
+            self._close_unless_kept()
+
+    def call_tool(self, name: str, args: dict) -> dict:
+        try:
+            self._init()
+            res = self._rpc("tools/call", {"name": name, "arguments": args or {}})
+        finally:
+            self._close_unless_kept()
+        if res.get("isError"):
+            raise MCPError(_text(res) or f"tool '{name}' failed")
+        return _unwrap(res)
+
+
+def _text(res: dict) -> str | None:
+    for block in res.get("content", []):
+        if block.get("type") == "text":
+            return block.get("text")
+    return None
+
+
+def _unwrap(res: dict):
+    sc = res.get("structuredContent")
+    if sc is not None:
+        if isinstance(sc, dict) and set(sc.keys()) == {"result"}:
+            return sc["result"]
+        return sc
+    txt = _text(res)
+    if txt is None:
+        return res.get("content", res)
+    try:
+        return json.loads(txt)
+    except json.JSONDecodeError:
+        return {"text": txt}

provekit/services/runstore.py

@@ -0,0 +1,93 @@
+"""Storage for paused flow-run contexts (breakpoint/step debugging).
+
+In-memory by default (single-process local use). When REDIS_URL is set, contexts live in
+Redis so step-debugging survives across uvicorn workers and restarts — pop uses GETDEL so
+two concurrent /continue calls can't resume the same run twice.
+"""
+from __future__ import annotations
+
+import json
+import time
+from functools import lru_cache
+
+from ..config import get_settings
+
+_TTL = 1800          # evict paused runs abandoned for 30 min
+_MAX = 200           # in-memory hard cap
+_MAX_BYTES = 1_000_000  # reject a ctx that serializes larger than ~1MB
+
+
+class _MemoryStore:
+    def __init__(self):
+        self._d: dict[str, dict] = {}
+        self._ts: dict[str, float] = {}
+
+    def store(self, rid: str, ctx: dict) -> None:
+        now = time.monotonic()
+        self._d[rid] = ctx
+        self._ts[rid] = now
+        for k in [k for k, t in list(self._ts.items()) if now - t > _TTL]:
+            self._d.pop(k, None); self._ts.pop(k, None)
+        if len(self._d) > _MAX:
+            for k in sorted(self._ts, key=self._ts.get)[: len(self._d) - _MAX]:
+                self._d.pop(k, None); self._ts.pop(k, None)
+
+    def get(self, rid: str) -> dict | None:
+        return self._d.get(rid)
+
+    def pop(self, rid: str) -> dict | None:
+        self._ts.pop(rid, None)
+        return self._d.pop(rid, None)
+
+    def drop(self, rid: str) -> None:
+        self._d.pop(rid, None); self._ts.pop(rid, None)
+
+
+class _RedisStore:
+    _PREFIX = "agm:run:"
+
+    def __init__(self, url: str):
+        import redis  # imported only when REDIS_URL is set
+        self._r = redis.Redis.from_url(url, decode_responses=True)
+
+    def _k(self, rid: str) -> str:
+        return self._PREFIX + rid
+
+    def store(self, rid: str, ctx: dict) -> None:
+        blob = json.dumps(ctx)
+        if len(blob) > _MAX_BYTES:
+            raise ValueError("run context too large to persist")
+        self._r.setex(self._k(rid), _TTL, blob)
+
+    def get(self, rid: str) -> dict | None:
+        blob = self._r.get(self._k(rid))
+        return json.loads(blob) if blob else None
+
+    def pop(self, rid: str) -> dict | None:
+        blob = self._r.getdel(self._k(rid))  # atomic — prevents double-resume
+        return json.loads(blob) if blob else None
+
+    def drop(self, rid: str) -> None:
+        self._r.delete(self._k(rid))
+
+
+@lru_cache
+def _store():
+    url = get_settings().redis_url
+    return _RedisStore(url) if url else _MemoryStore()
+
+
+def store_ctx(rid: str, ctx: dict) -> None:
+    _store().store(rid, ctx)
+
+
+def get_ctx(rid: str) -> dict | None:
+    return _store().get(rid)
+
+
+def pop_ctx(rid: str) -> dict | None:
+    return _store().pop(rid)
+
+
+def drop_ctx(rid: str) -> None:
+    _store().drop(rid)

provekit/services/sealing.py

@@ -0,0 +1,106 @@
+"""Encrypt connection secrets at rest (Fernet). Sealed values carry an "enc:" prefix;
+plaintext values from older databases are accepted on read and upgraded by the
+one-shot reseal pass at startup (and on any subsequent write).
+
+Key source: SECRET_KEY env var (any string — derived to a Fernet key), or, for local
+SQLite use, an auto-generated key file stored next to the database (.provekit.key).
+Rotation: set the new SECRET_KEY, then re-enter credentials (values sealed with the
+old key decrypt to "" with a logged warning rather than crashing).
+"""
+from __future__ import annotations
+
+import base64
+import hashlib
+import logging
+import os
+from functools import lru_cache
+from pathlib import Path
+
+from cryptography.fernet import Fernet, InvalidToken
+
+from ..config import get_settings
+from .masking import is_secret_header
+
+_PREFIX = "enc:"
+log = logging.getLogger("provekit.sealing")
+
+
+def _key_file() -> Path | None:
+    url = get_settings().database_url
+    if url.startswith("sqlite:///"):
+        db_path = Path(url.removeprefix("sqlite:///"))
+        parent = db_path.parent if str(db_path.parent) else Path(".")
+        return parent / ".provekit.key"
+    return None
+
+
+@lru_cache
+def _fernet() -> Fernet:
+    s = get_settings()
+    if s.secret_key:
+        return Fernet(base64.urlsafe_b64encode(hashlib.sha256(s.secret_key.encode()).digest()))
+    kf = _key_file()
+    if kf is None:
+        raise RuntimeError("SECRET_KEY must be set when not using a local SQLite database")
+    if kf.exists():
+        return Fernet(kf.read_bytes().strip())
+    key = Fernet.generate_key()
+    kf.parent.mkdir(parents=True, exist_ok=True)
+    kf.write_bytes(key)
+    try:
+        os.chmod(kf, 0o600)
+    except OSError:
+        pass
+    log.info("generated local encryption key at %s", kf)
+    return Fernet(key)
+
+
+def seal(value: str) -> str:
+    if not value or not isinstance(value, str) or value.startswith(_PREFIX):
+        return value
+    return _PREFIX + _fernet().encrypt(value.encode()).decode()
+
+
+def unseal(value):
+    if not isinstance(value, str) or not value.startswith(_PREFIX):
+        return value
+    try:
+        return _fernet().decrypt(value[len(_PREFIX):].encode()).decode()
+    except (InvalidToken, ValueError):
+        log.warning("could not decrypt a stored secret (SECRET_KEY changed?) — treating as unset")
+        return ""
+
+
+def seal_config(cfg: dict) -> dict:
+    """Encrypt the secret fields of a connection config: api_key, secret headers,
+    the MCP OAuth client_secret, and stdio env values (which routinely carry tokens)."""
+    out = dict(cfg)
+    if out.get("api_key"):
+        out["api_key"] = seal(out["api_key"])
+    hdrs = out.get("headers")
+    if isinstance(hdrs, dict):
+        out["headers"] = {k: (seal(v) if is_secret_header(k) and isinstance(v, str) and v else v)
+                          for k, v in hdrs.items()}
+    oauth = out.get("oauth")
+    if isinstance(oauth, dict) and oauth.get("client_secret"):
+        out["oauth"] = {**oauth, "client_secret": seal(oauth["client_secret"])}
+    env = out.get("env")
+    if isinstance(env, dict):
+        out["env"] = {k: (seal(v) if isinstance(v, str) and v else v) for k, v in env.items()}
+    return out
+
+
+def unseal_config(cfg: dict) -> dict:
+    out = dict(cfg)
+    if out.get("api_key"):
+        out["api_key"] = unseal(out["api_key"])
+    hdrs = out.get("headers")
+    if isinstance(hdrs, dict):
+        out["headers"] = {k: unseal(v) for k, v in hdrs.items()}
+    oauth = out.get("oauth")
+    if isinstance(oauth, dict) and "client_secret" in oauth:
+        out["oauth"] = {**oauth, "client_secret": unseal(oauth.get("client_secret"))}
+    env = out.get("env")
+    if isinstance(env, dict):
+        out["env"] = {k: unseal(v) for k, v in env.items()}
+    return out

provekit/services/templates.py

@@ -0,0 +1,69 @@
+"""Bundled flow templates — the searchable gallery served to the New-flow picker."""
+from __future__ import annotations
+
+import json
+import pathlib
+from functools import lru_cache
+
+from . import testfile
+
+_DIR = pathlib.Path(__file__).resolve().parent.parent / "templates" / "flows"
+
+# A small, hand-picked showcase surfaced at the top of the New-flow picker. The gallery is a
+# regular 24-domain × 18-pattern grid (432 entries), which reads as a wall of near-duplicates;
+# these are eight *distinct patterns* across eight *distinct domains* — each a different
+# capability (branching, self-critique, structured output, safety gate, numeric gate,
+# language routing) — so a newcomer sees a few great starting points instead of everything.
+_FEATURED = [
+    "customer-support-triage-route",
+    "sales-lead-draft-critique-revise",
+    "healthcare-intake-extract-to-json",
+    "content-moderation-moderation-gate",
+    "legal-intake-summarize-then-decide",
+    "e-commerce-sentiment-route",
+    "recruiting-score-gate",
+    "travel-language-route",
+]
+
+
+@lru_cache
+def _manifest() -> list[dict]:
+    f = _DIR / "_manifest.json"
+    return json.loads(f.read_text()) if f.exists() else []
+
+
+def featured() -> list[dict]:
+    """The curated showcase, in order. Skips any slug missing from the manifest so a
+    renamed template can't break the picker."""
+    by_slug = {m["slug"]: m for m in _manifest()}
+    return [by_slug[s] for s in _FEATURED if s in by_slug]
+
+
+def search(q: str = "", limit: int = 60) -> list[dict]:
+    """Filter templates by a case-insensitive substring over name/description/category."""
+    items = _manifest()
+    q = (q or "").strip().lower()
+    if q:
+        items = [m for m in items
+                 if q in m["name"].lower() or q in m["description"].lower() or q in m["category"].lower()]
+    return items[:limit]
+
+
+def categories() -> list[str]:
+    return sorted({m["category"] for m in _manifest()})
+
+
+def total() -> int:
+    return len(_manifest())
+
+
+def load(slug: str) -> dict | None:
+    """Return the parsed .provekit flow doc for a template slug (None if unknown)."""
+    # Reject separators/traversal (incl. Windows drive-relative "C:foo"), then confirm the
+    # resolved path stays inside the templates dir before reading.
+    if not slug or any(ch in slug for ch in ("/", "\\", ":")) or ".." in slug:
+        return None
+    f = (_DIR / f"{slug}.yaml").resolve()
+    if _DIR not in f.parents or not f.exists():
+        return None
+    return testfile.load(f.read_text())

provekit/services/testfile.py

@@ -0,0 +1,122 @@
+"""The .provekit file format — plain-text, git-diffable tests and flows.
+
+A *test* file bundles one request (prompt | tool | agent), its assertions, and an
+optional dataset of input rows — enough to re-run it anywhere, including CI.
+A *flow* file carries a visual flow's graph. Connections are referenced BY NAME:
+credentials never enter the file; they resolve against the local workspace at
+import/run time.
+
+Canonical test shape (YAML, stable field order):
+
+    version: 1
+    kind: test
+    name: Support bot regression
+    connection: OpenAI (prod)
+    request:
+      type: prompt
+      model: gpt-4o-mini
+      system: You are a support assistant.
+      user: "{{message}}"
+    assertions:
+      - type: contains
+        value: refund
+    dataset:
+      - name: angry customer
+        variables: {message: "URGENT: refund now"}
+"""
+from __future__ import annotations
+
+import yaml
+
+VERSION = 1
+KINDS = ("test", "flow")
+
+# Stable, human-friendly field order per request type.
+_REQ_ORDER = {
+    "prompt": ["type", "model", "system", "user", "messages", "temperature", "max_tokens"],
+    "tool": ["type", "tool", "args"],
+    "agent": ["type", "method", "path", "headers", "body", "stream"],
+}
+_ASSERT_ORDER = ["type", "name", "path", "value", "schema", "criteria", "model"]
+
+# Never serialized: secrets and workspace-local ids.
+_STRIP = {"api_key", "connection_id", "assertions", "_k", "provider", "base_url", "url"}
+
+
+def _ordered(d: dict, order: list[str]) -> dict:
+    out = {k: d[k] for k in order if k in d and d[k] not in (None, "", {}, [])}
+    out.update({k: v for k, v in d.items() if k not in out and k not in order})
+    return out
+
+
+def _clean_request(request: dict) -> dict:
+    req = {k: v for k, v in request.items() if k not in _STRIP}
+    return _ordered(req, _REQ_ORDER.get(request.get("type", ""), ["type"]))
+
+
+def dump_test(name: str, request: dict, connection_name: str | None = None,
+              dataset: list[dict] | None = None) -> str:
+    doc: dict = {"version": VERSION, "kind": "test", "name": name or "test"}
+    if connection_name:
+        doc["connection"] = connection_name
+    doc["request"] = _clean_request(request)
+    asserts = request.get("assertions") or []
+    if asserts:
+        doc["assertions"] = [_ordered(dict(a), _ASSERT_ORDER) for a in asserts]
+    if dataset:
+        doc["dataset"] = [{"name": r.get("name") or f"row {i + 1}", "variables": r.get("variables") or {}}
+                          for i, r in enumerate(dataset)]
+    return yaml.safe_dump(doc, sort_keys=False, allow_unicode=True, width=100)
+
+
+def dump_flow(name: str, description: str, nodes: list, edges: list,
+              connection_names: dict[int, str] | None = None) -> str:
+    """connection_names maps connection_id -> name so flow files stay portable."""
+    names = connection_names or {}
+    out_nodes = []
+    for n in nodes or []:
+        node = dict(n)
+        cfg = dict(node.get("config") or {})
+        cid = cfg.pop("connection_id", None)
+        if cid is not None and cid in names:
+            cfg["connection"] = names[cid]
+        node["config"] = cfg
+        out_nodes.append(node)
+    doc: dict = {"version": VERSION, "kind": "flow", "name": name or "flow"}
+    if description:
+        doc["description"] = description
+    doc["nodes"] = out_nodes
+    doc["edges"] = edges or []
+    return yaml.safe_dump(doc, sort_keys=False, allow_unicode=True, width=100)
+
+
+def load(text: str) -> dict:
+    """Parse + validate an .provekit document. Returns the normalized dict; the caller
+    resolves `connection` names to ids against its own workspace."""
+    try:
+        doc = yaml.safe_load(text)
+    except yaml.YAMLError as exc:
+        raise ValueError(f"not valid YAML: {exc}") from exc
+    if not isinstance(doc, dict):
+        raise ValueError("expected a YAML mapping at the top level")
+    if doc.get("version") != VERSION:
+        raise ValueError(f"unsupported version {doc.get('version')!r} (expected {VERSION})")
+    kind = doc.get("kind")
+    if kind not in KINDS:
+        raise ValueError(f"unsupported kind {kind!r} (expected one of {KINDS})")
+    if kind == "test":
+        req = doc.get("request")
+        if not isinstance(req, dict) or req.get("type") not in ("prompt", "tool", "agent"):
+            raise ValueError("test needs a request with type prompt | tool | agent")
+        if "api_key" in req:
+            raise ValueError("credentials do not belong in .provekit files — use a connection name")
+        for a in doc.get("assertions") or []:
+            if not isinstance(a, dict) or not a.get("type"):
+                raise ValueError("each assertion needs a type")
+        for r in doc.get("dataset") or []:
+            if not isinstance(r, dict):
+                raise ValueError("each dataset row must be a mapping")
+    else:  # flow
+        if not isinstance(doc.get("nodes"), list) or not isinstance(doc.get("edges"), list):
+            raise ValueError("flow needs nodes and edges lists")
+    return doc