feat(interface): expose workflow metadata models

[andreatnvidia]

📋 Summary

Adds a public Pydantic contract for workflow-metadata.json so users and integrations can validate and serialize workflow state while retaining compatibility with existing metadata files.

🔗 Related Issue

Closes #900

🔄 Changes

• Export WorkflowMetadata and status-specific workflow stage metadata models from data_designer.interface.
• Validate metadata at workflow read and write boundaries while preserving legacy fields and compatible extensions during resume.
• Validate stage metadata inputs before filesystem changes and cover invalid resume behavior.
• Report distinct type and range errors for invalid stage num_records values.
• Add focused model, serialization, legacy compatibility, and workflow integration tests.

🧪 Testing

• make test passes (not run; the change is scoped to the interface package)
• .venv/bin/ruff check --fix .
• .venv/bin/ruff format .
• .venv/bin/pytest packages/data-designer/tests -q - 1141 passed, 1 skipped
• .venv/bin/pytest packages/data-designer/tests/interface/test_composite_workflow.py::test_composite_workflow_runs_linear_stages_with_disk_handoff -q - 1 passed
• .venv/bin/pytest packages/data-designer/tests/interface/test_composite_workflow.py::test_composite_workflow_resume_if_possible_delegates_matching_resumable_stage -q - 2 passed
• Unit tests added/updated
• E2E tests added/updated (N/A - no E2E behavior changed)

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (N/A - no architecture change)

---

Description updated with AI

[github-actions]

Fern preview: https://nvidia-preview-pr-902.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[greptile-apps]

Greptile Summary

Adds public Pydantic models for workflow metadata and validates metadata at workflow persistence boundaries while preserving compatible extension fields.

• Exports workflow and status-specific stage metadata models through the public interface.
• Validates workflow metadata during reads and atomic writes.
• Adds stage input validation and workflow metadata compatibility tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
packages/data-designer/src/data_designer/interface/workflow_metadata.py	Defines the public, status-discriminated Pydantic contract for workflow and stage metadata.
packages/data-designer/src/data_designer/interface/composite_workflow.py	Integrates metadata validation and extension preservation into workflow read, write, and resume paths.
packages/data-designer/src/data_designer/interface/init.py	Exposes the new metadata models through the package’s lazy public imports.
packages/data-designer/tests/interface/test_workflow_metadata.py	Tests public exports, status variants, serialization, legacy compatibility, extensions, and invalid metadata.
packages/data-designer/tests/interface/test_composite_workflow.py	Adds workflow integration coverage for metadata validation, legacy resume behavior, and extension preservation.

_{Reviews (3): Last reviewed commit: "fix: tighten workflow metadata contracts" | Re-trigger Greptile}

[github-actions]

Thanks for putting this together, @andreatnvidia — the backward-compatibility work here is careful, and it shows.

Summary

This adds a public Pydantic contract (WorkflowMetadata + four status-discriminated stage variants) for workflow-metadata.json, and wires validation into both the read (resume) and write boundaries of CompositeWorkflow. The implementation matches the stated intent: legacy files still load, extension fields round-trip, and invalid stage metadata degrades to "start fresh" under IF_POSSIBLE while hard-failing under ALWAYS.

One thing worth highlighting up front: I traced every required field on the new models back through #636 (workflow chaining) and #747 (stage resume). Every field that was introduced later — only stage_output_override_path — correctly got a default, and every field marked required has been written since the feature landed. That's the part of this PR most likely to break users silently, and it looks right.

Findings

Warnings — Worth addressing

packages/data-designer/src/data_designer/interface/composite_workflow.py:490-493 — Metadata validation on the failure path can shadow the real stage error

• What: The except Exception: handler now calls _write_workflow_metadata, which raises DataDesignerWorkflowError("Workflow metadata has invalid shape.") if validation fails. Because that happens before the bare raise, a validation failure would replace the original stage exception (a generation error, an OOM, a callback bug) with a message about metadata shape.
• Why: I couldn't find a reachable path today — the status: "running" update at line 406 populates every field FailedWorkflowStageMetadata requires, so the failed-stage dict is always valid. But that's an invariant held only by the ordering of two update() calls ~80 lines apart. The next person who adds a field to _base_stage_metadata or the completed-stage update without touching workflow_metadata.py converts every stage failure in that path into a misleading error, and the original traceback is gone. Debugging that is genuinely painful.
• Suggestion: Make the failure-path write non-fatal so the real exception always wins:

  except Exception:
      stage_metadata.update({"status": "failed", "duration_sec": time.monotonic() - start_time})
      try:
          _write_workflow_metadata(workflow_path, metadata)
      except DataDesignerWorkflowError:
          logger.warning("Could not persist failed-stage metadata for stage %r.", stage.name, exc_info=True)
      raise

  Alternatively, if the write must stay strict, chaining it (raise ... from exc) at least preserves the original. What do you think?

packages/data-designer/src/data_designer/interface/composite_workflow.py:223-226 — New num_records type guard is untested, and its error message is now misleading

• What: The guard grew from num_records < 1 to also reject non-int and bool. The new test_composite_workflow_rejects_invalid_stage_metadata_before_artifacts parametrize covers on_success_version, allow_empty, sampling_strategy, and selection_strategy — but not num_records. I grepped the test suite and found no coverage for the num_records guard at all (grep -rn "must be at least 1" packages/data-designer/tests/ only hits max_attempts in test_data_designer.py). Separately, num_records="3" or num_records=True now raises "Stage num_records must be at least 1.", which describes a range problem rather than the type problem that actually occurred.
• Why: AGENTS.md lists "No untested code paths" as a structural invariant, and the parametrize is right next door — this is a two-line addition. The message matters because num_records=True passing a range check is exactly the confusing case the bool exclusion was added to catch; telling the user "must be at least 1" when they passed True sends them looking in the wrong place.
• Suggestion: Split the checks so each message names its own failure, and fold the cases into the existing parametrize:

  if num_records is not None:
      if not isinstance(num_records, int) or isinstance(num_records, bool):
          raise DataDesignerWorkflowError("Stage num_records must be an integer.")
      if num_records < 1:
          raise DataDesignerWorkflowError("Stage num_records must be at least 1.")

  ({"num_records": "3"}, "num_records must be an integer"),
  ({"num_records": True}, "num_records must be an integer"),
  ({"num_records": 0}, "num_records must be at least 1"),

packages/data-designer/src/data_designer/interface/composite_workflow.py:227-234 — The isinstance guard set is inconsistent with what actually gets serialized

• What: add_stage now type-checks on_success_version, allow_empty, sampling_strategy, and selection_strategy, but not name, output, on_success, depends_on, or output_processors — even though output_processors is the one that fails most opaquely (processor.model_dump → AttributeError in _base_stage_metadata, deep inside run()).
• Why: I think the guards are well-motivated, and I want to be fair about that: I checked each one, and passing a plain "ordered" string or a dict selection strategy was already broken before this PR (.value / .model_dump() would AttributeError), while on_success_version=1 would have newly failed pydantic's lax str validation with a confusing error. So these turn latent breakage into clear messages at the right boundary. The concern is that a partial guard set reads as a complete one — a future reader will reasonably assume add_stage validates its inputs, and it mostly doesn't.
• Suggestion: Either close the gap for the fields that feed serialization (output_processors especially), or add a brief comment stating the guards exist specifically to keep _base_stage_metadata output valid against WorkflowStageMetadata, so the scope is intentional rather than accidental. A comment would be my preference — it's cheaper and documents the real invariant. Also worth noting: the PR description says these validate "before filesystem changes," but add_stage never touches the filesystem (the first mkdir is in run()), so that rationale may be worth restating.

packages/data-designer/tests/interface/test_composite_workflow.py:599-631 — The rename folded away the clean resume-skip baseline

• What: test_composite_workflow_resume_if_possible_skips_completed_stages became ..._skips_legacy_completed_stages and now mutates the metadata file before resuming (adds workflow_extension, adds per-stage stage_extension, pops stage_output_override_path). It asserts three independent things: that resume skips completed stages, that unknown extras survive a round-trip, and that a missing post-#747 field is tolerated.
• Why: The plain "two completed stages resume without re-running" path no longer has a dedicated test. Other tests exercise resume-skip incidentally (..._skips_stage_with_output_processors, ..._uses_relative_metadata_paths_after_move, ..._preserves_completed_empty_skip), so it isn't uncovered — but if the legacy-compat assertions start failing, it's now ambiguous whether the regression is in resume or in extras handling, and the simplest regression case is the one you'd most want isolated.
• Suggestion: Keep the original test unmodified as the baseline, and add the mutation as a separate ..._skips_legacy_completed_stages. The duplication is small and the two failure modes stay distinguishable.

Suggestions — Take it or leave it

packages/data-designer/src/data_designer/interface/workflow_metadata.py:11-26 — WorkflowStageMetadata is exported but can't appear in stages

• What: WorkflowStageMetadata is public and its status field lists all five values, but WorkflowMetadata.stages is typed as the discriminated union of the four concrete subclasses. A user who constructs a WorkflowStageMetadata directly gets something that will never validate as a stages entry, and the five-value Literal on the base is dead — every concrete subclass narrows it.
• Why: It's useful as an isinstance target and a shared base, so exporting it is reasonable; the ambiguity is just about what it's for.
• Suggestion: A one-line docstring addition would settle it — something like "Common fields shared by all stage variants; use the status-specific subclasses to construct or match a stage." Optionally, exporting the union alias publicly (e.g. WorkflowStageMetadataVariant) would give integrations a name to annotate against, instead of the currently-private _WorkflowStageMetadataVariant that types a public field.

packages/data-designer/src/data_designer/interface/composite_workflow.py:898-905 — Full-metadata validation on every write is quadratic in stage count

• What: _write_workflow_metadata runs model_validate then model_dump over the entire metadata payload — including every prior stage's full config dump — and it's called up to three times per stage iteration. For an N-stage workflow that's O(N²) traversal over a growing payload, where it used to be a single json.dump.
• Why: Honestly, this is probably fine: the os.fsync in the same function almost certainly dominates, and workflows are typically well under ten stages. Flagging it only because it's new cost introduced by this change and config payloads can get large.
• Suggestion: No action needed unless you see it in practice. If it ever matters, validating only the stage dict that changed (rather than the whole document) would keep it linear.

packages/data-designer/src/data_designer/interface/workflow_metadata.py:11,79 — Deviates from the project's ConfigBase convention

• What: Both models subclass raw BaseModel with ConfigDict(extra="allow"), rather than the project's ConfigBase (packages/data-designer-config/src/data_designer/config/base.py), which sets protected_namespaces=(), use_enum_values=True, and extra="forbid".
• Why: extra="allow" is clearly deliberate and correct here — forward-compatibility is the whole point, and extra="forbid" would defeat it. But this is the first Pydantic model in the interface package, so it sets a precedent, and interface → config is a legal import direction if you wanted to reuse the base.
• Suggestion: Either is defensible. If you keep raw BaseModel, a short comment noting that extra="allow" is required for forward compatibility with future metadata fields would stop someone from "fixing" it to extra="forbid" later.

packages/data-designer/tests/interface/test_composite_workflow.py:372-393 — Two small polish items in the new parametrized test

• What: assert not stub_artifact_path.exists() passes trivially — add_stage never touches the filesystem, and the fixture is tmp_path / "artifacts" which nothing creates. Also compose_workflow(name="invalid-version") looks copy-pasted from a version-specific case; the test now covers four unrelated kwargs.
• Why: A vacuous assertion in a test named ..._before_artifacts implies coverage that isn't there — if add_stage ever started creating directories, this would catch it only by accident.
• Suggestion: Rename to something like invalid-stage-kwargs, and either drop the exists() assertion or move the guard to run() where artifact creation actually happens.

packages/data-designer/tests/interface/test_workflow_metadata.py:539-560 — pytest.raises(ValidationError) without match lets cases pass for the wrong reason

• What: The five parametrized rejection cases fail for three different reasons (unknown discriminator tag, missing required started-stage fields, wrong duration_sec type) but all assert only that some ValidationError was raised. {"status": "running"} in particular is rejected because base_stage_metadata lacks fingerprint/config/etc., which is easy to misread as "running is invalid."
• Suggestion: Adding a match per case (or an expected-loc assertion) would pin down what's actually being verified and make the intent readable at a glance.

Docs — new public API isn't documented anywhere

• What: Six new names are exported from data_designer.interface, but I found no mention of them in fern/ or architecture/. The closest page, fern/versions/latest/pages/concepts/workflow-chaining.mdx, references the metadata schema only to say it's experimental and subject to change.
• Why: Nothing in the docs is made false by this PR, so it isn't a correctness problem. But the PR's stated goal is "so users and integrations can validate and serialize workflow state" — and users can't do that if they don't know the models exist.
• Suggestion: A short snippet on the workflow-chaining page showing WorkflowMetadata.model_validate_json(path.read_text()) would make the feature discoverable. A follow-up issue is fine if you'd rather keep this PR tight.

packages/data-designer/src/data_designer/interface/composite_workflow.py:290 — Top-level extras survive resume but are dropped by a fresh run

• What: **(prior_metadata or {}) preserves unknown top-level keys across a resume. On resume=NEVER, prior_metadata is None, so those same extras are silently discarded.
• Why: Arguably correct — a fresh run is a new run — but the asymmetry isn't obvious from the call site, and an integration storing annotations there would lose them without warning.
• Suggestion: A brief comment on the spread explaining that extras are intentionally carried across resumes only would make the intent explicit.

Structural Impact

(graphify, 2.4s)

Risk: HIGH (103 import direction violation(s))

• 5 Python files, 89 AST entities, 5/77 clusters

Import Direction Violations (103)

Legal direction: interface -> engine -> config

• .create_report_section() (config) --calls--> .keys() (interface)
• inject_sampler_type_into_params() (config) --calls--> .items() (interface)
• _resolve_sampler_kwargs() (config) --calls--> .items() (interface)
• allowed_references() (config) --calls--> .keys() (interface)
• ._resolve_drop_column_names() (config) --calls--> .keys() (interface)
• +98 more

High-Connectivity Changes

• .items() (87 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• CompositeWorkflow (44 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• .keys() (39 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• composite_workflow.py (38 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• CompositeWorkflowResults (38 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• .run() (35 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• WorkflowMetadata (25 deps) in packages/data-designer/src/data_designer/interface/workflow_metadata.py
• SkippedStageResult (22 deps) in packages/data-designer/src/data_designer/interface/composite_workflow.py
• +47 more

Cross-Package Dependencies

• Lazily import interface exports when accessed. (interface) --rationale_for--> __getattr__() (config)
• Lazily import interface exports when accessed. (interface) --uses--> ResumeMode (config)
• Return available exports for tab-completion. (interface) --rationale_for--> __dir__() (config)
• Return available exports for tab-completion. (interface) --uses--> ResumeMode (config)
• _WorkflowStage (interface) --uses--> DatasetProfilerResults (config)
• _WorkflowStage (interface) --uses--> ProcessorConfig (config)
• +413 more

Reviewer note on the HIGH rating — I believe it is a false positive, and did not weight it in the verdict.

I checked the reported violations against the actual imports rather than taking the rating at face value:

• The 103 "import direction violations" are method-name collisions, not import edges. Every sampled violation resolves to a call to .keys() or .items() somewhere in the config package being attributed to CompositeWorkflowResults.keys() / .items() (composite_workflow.py:120-124). Those are ordinary dict-protocol methods; the extractor is matching on bare method name across package boundaries. The same artifact inflates the "god node" list — .items() at 87 deps and .keys() at 39 deps are name collisions, not this PR's blast radius.
• The actual new imports are all intra-package and legal. workflow_metadata.py imports only stdlib typing and pydantic. composite_workflow.py adds one interface → interface import. __init__.py adds interface → interface lazy entries. There is no new cross-package edge in any direction, let alone a reversed one.
• Import-time cost is unchanged. Worth checking given AGENTS.md's fast-imports invariant: workflow_metadata.py pulls in nothing beyond pydantic, which composite_workflow.py already imported (from pydantic import ValidationError). The interface/__init__.py additions go through _LAZY_IMPORTS, so importing data_designer.interface costs the same as before.

Where the report is useful: CompositeWorkflow (44 deps) and CompositeWorkflow.run() (35 deps) are genuinely central, and run() is where this PR concentrates its behavior change. That's what drove the extra scrutiny on the write-path and failure-path findings above, and the backward-compatibility trace through #636/#747.

Linting

I was not able to run ruff — this runner has no .venv and neither ruff nor pydantic is importable from the system Python, so .venv/bin/ruff check and --format --check both failed on a missing binary. The changes look consistent with the configured rule set (W, F, I, ICN, PIE, TID, UP006, UP007, UP045 at target-version = "py310"): both new files carry SPDX headers and from __future__ import annotations, imports are absolute and sorted, annotations use modern syntax, and every typing import in workflow_metadata.py is used. Worth confirming with make check-all in CI, since I could not verify formatting locally.

What Looks Good

• The backward-compatibility analysis is the strongest part of this PR. Getting stage_output_override_path a default while leaving genuinely-always-present fields required is exactly the right call, and it's the kind of thing that's easy to get wrong in the direction that silently breaks resume for existing users. I traced each required field back to #636/#747 and the split holds up.
• exclude_unset=True on both dump sites is a subtle, correct choice. It's what makes unknown extension fields and absent-legacy-fields round-trip instead of getting normalized away, and the tests pin both halves of that behavior (test_workflow_metadata_preserves_extra_fields, test_workflow_metadata_supports_legacy_completed_stage).
• The read-path error handling respects the existing IF_POSSIBLE / ALWAYS contract. Slotting ValidationError into the same warn-and-start-fresh vs. hard-fail structure as the existing JSON-decode and shape branches means resume semantics stay predictable, and surfacing the failing field path in the warning is a genuinely helpful touch for debugging.
• The discriminated union is the right modeling choice. Keying on status gives precise per-state field requirements instead of a single model where everything is optional, and the round-trip assertion in test_workflow_metadata_supports_stage_statuses is a nice way to cover serialization for all five statuses at once.

Verdict

Needs changes — nothing here is a correctness bug in the current code, and the backward-compatibility work is solid. The items I'd like to see addressed before merge:

• Make the failure-path _write_workflow_metadata call non-fatal (or chain it) so a metadata-shape problem can't shadow the real stage exception (composite_workflow.py:490-493).
• Add num_records cases to the new parametrize and split the type-vs-range error messages (composite_workflow.py:223-226).
• Either close the output_processors gap in the add_stage guards or comment the intended scope (composite_workflow.py:227-234).
• Restore a standalone clean resume-skip test alongside the legacy variant (test_composite_workflow.py:599).

Everything under Suggestions is genuinely optional. Since I couldn't run ruff locally, please confirm CI is green on make check-all.

---

This review was generated by an AI assistant.

[nabinchha]

Thanks for putting this together, @andreatnvidia!

Summary

This PR adds public, status-specific Pydantic models for workflow-metadata.json, integrates them at workflow read/write boundaries, and retains legacy completed-stage compatibility. The implementation largely matches the stated intent, but two paths still weaken the promised stage-level validation and extension preservation.

Findings

Warnings — Worth addressing

packages/data-designer/src/data_designer/interface/workflow_metadata.py:11 — The public stage model accepts incomplete status payloads

• What: WorkflowStageMetadata is publicly exported as the stage-level model, but it only validates common fields while allowing every status literal. For example, WorkflowStageMetadata.model_validate(...) accepts a status="completed" payload with none of fingerprint, num_records_actual, output_seed_path, or duration_sec; only the private _WorkflowStageMetadataVariant performs status-specific dispatch when a stage is nested inside WorkflowMetadata.
• Why: An integration validating a standalone metadata["stages"][i] through the obvious public API gets a false success for metadata that the enclosing workflow model correctly rejects. That leaves the new public stage-level contract inconsistent depending on whether callers validate the full file or one stage.
• Suggestion: Could we make the common base private and expose WorkflowStageMetadata as the discriminated status union (or an equivalent RootModel/public adapter), then add a direct test showing that a standalone completed stage without completion fields fails validation?

packages/data-designer/src/data_designer/interface/composite_workflow.py:407 — Partial-stage resume drops extension fields

• What: Top-level extras and extras on reused completed stages are preserved, but a matching running or failed stage takes the stage_resume = ResumeMode.ALWAYS path and rebuilds stage_metadata from _base_stage_metadata() plus the known running fields. Any model_extra fields from prior_stage_metadata disappear on the first metadata write.
• Why: Integrations that attach compatible stage-level observability/provenance fields lose them during the normal partial-resume path, despite extra="allow" and the PR's compatibility goal. The added integration test covers only a completed stage that is skipped, so this data-loss path is currently untested.
• Suggestion: Preserve the prior stage's extra fields when stage_resume == ResumeMode.ALWAYS (while deliberately excluding stale schema-owned status fields), and cover both running and failed resume cases with an extension field assertion.

Suggestions — Take it or leave it

packages/data-designer/src/data_designer/interface/workflow_metadata.py:17 — Add domain constraints to the validation contract

• What: The models currently accept impossible values such as negative stage indices, requested/actual/output record counts, and durations; Pydantic coercion also turns values such as allow_empty="yes" into booleans.
• Why: The persisted JSON is now presented as a public validation contract for observability and integrations, so accepting semantically impossible state makes a successful validation less meaningful.
• Suggestion: Consider strict model validation plus Field(ge=0)/Field(gt=0) constraints for indices, counts, and durations, with focused rejection tests. The persisted files already emit the corresponding concrete JSON types, so this should not affect valid legacy files.

fern/versions/latest/pages/concepts/workflow-chaining.mdx:6 — Document the newly public metadata surface

• What: The workflow guide still only says the metadata schema is experimental and does not mention WorkflowMetadata or the status-specific models.
• Why: The new API is otherwise difficult for users and integration authors to discover, and the page does not distinguish a typed contract for the current schema from a promise of long-term schema stability.
• Suggestion: Add a short metadata-inspection example using WorkflowMetadata.model_validate_json(...) and clarify that the model represents the current experimental schema.

What Looks Good

• The discriminated union inside WorkflowMetadata cleanly maps all five persisted statuses to precise runtime types.
• exclude_unset=True is a thoughtful compatibility choice for the newly optional legacy field, and extra="allow" preserves compatible top-level and reused-stage extensions.
• The read/write boundary normalization and focused regression coverage are solid; all 91 directly relevant tests pass, as do lint and format checks.

Verdict

Needs changes — Please make standalone stage validation status-aware and preserve extension fields when resuming running/failed stages. The numeric constraints and Fern documentation would also strengthen the public contract but are non-blocking suggestions.

---

This review was generated by an AI assistant.

[andreatnvidia]

Thanks, @nabinchha. I addressed both warnings in d000729a: standalone WorkflowStageMetadata validation is now status-aware, and partial running/failed resumes preserve extension fields without retaining stale schema-owned fields. I added regression coverage for both paths. Could you take another look when you have a chance?

[nabinchha]

Thanks for the quick follow-up, @andreatnvidia!

Summary

I re-reviewed the PR through d000729a, focusing first on the two warnings from my previous review. Those behaviors are fixed and well covered, but the fix now exposes two public representations of the same stage concept; simplifying that contract before release would make the API more consistent and avoid a later compatibility problem.

Findings

Warnings — Worth addressing

packages/data-designer/src/data_designer/interface/workflow_metadata.py:70 — One stage concept has two public type representations

• What: The follow-up exports both WorkflowStageMetadata, a RootModel wrapper, and WorkflowStageMetadataVariant, the discriminated union containing the actual stage models. Standalone validation returns the wrapper and requires .root, while WorkflowMetadata.stages contains the concrete variants directly. Callers therefore access the same logical value differently: WorkflowStageMetadata.model_validate(payload).root.status versus WorkflowMetadata.model_validate(payload).stages[0].status.
• Why: This is a new public contract, so both names and the wrapper-versus-variant distinction become compatibility commitments once released. The split makes annotations and runtime usage harder to understand, and removing it later would require changing public symbols or return shapes.
• Suggestion: Could we make WorkflowStageMetadata the public discriminated union, keep a private TypeAdapter, and expose a clearly named validate_workflow_stage_metadata(value) function for standalone parsing? Then standalone and nested parsing both return the same concrete status model, there is no .root, and WorkflowStageMetadataVariant is unnecessary.

Suggestions — Take it or leave it

packages/data-designer/src/data_designer/interface/composite_workflow.py:405 — Keep prior metadata typed through the resume path

• What: _read_prior_workflow_metadata() validates the file as WorkflowMetadata and immediately dumps it back to a dictionary. The partial-resume path then validates the selected stage again through WorkflowStageMetadata solely to recover its concrete model and model_extra.
• Why: The validate → dump → revalidate cycle adds representation changes and repeated parsing introduced by this PR, and it is what makes the resume code depend on the standalone wrapper.
• Suggestion: Return WorkflowMetadata | None from _read_prior_workflow_metadata() and return the already-validated concrete stage from _get_prior_stage_metadata(). The resume path can then read prior_stage.model_extra directly, converting to JSON only at the persistence boundary.

What Looks Good

• The previous standalone-validation and partial-resume extension bugs are both fixed, and the new tests exercise the behavior before the resumed create call as well as in the final persisted metadata.
• Using the active concrete model's model_extra is the right mechanism for distinguishing compatible extensions from fields owned by that status schema.
• Focused tests, the broader interface suite, lint/format, and all current PR checks are green.

Verdict

Needs changes — The behavioral fixes are solid, but I recommend collapsing WorkflowStageMetadata and WorkflowStageMetadataVariant into one public stage type before this new API becomes a compatibility commitment. Keeping validated models through resume is a related non-blocking simplification.

---

This review was generated by an AI assistant.