type: narrow draft workspace result contracts

src/wf_api/drafts.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
-from typing import Any
+from typing import Any, Literal, cast, overload
 
 from jsonschema import Draft202012Validator, SchemaError
 
@@ -57,6 +57,7 @@ from .models import (
     CompileDraftWorkspaceSuccess,
     DeleteDraftWorkspaceResult,
     DraftWorkspaceResult,
+    DraftWorkspaceWithDocument,
     InvalidDraftResult,
     JsonProjector,
     ListDraftWorkspacesResult,
@@ -318,19 +319,38 @@ class WorkflowDraftApi:
             title=title,
         )
 
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[True],
+    ) -> DraftWorkspaceWithDocument: ...
+
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[False] = False,
+    ) -> DraftWorkspaceResult: ...
+
     async def get_draft_workspace(
         self,
         *,
         workspace_id: str,
         include_draft: bool = False,
-    ) -> DraftWorkspaceResult:
-        return _PROJECT_DRAFT_WORKSPACE(
+    ) -> DraftWorkspaceResult | DraftWorkspaceWithDocument:
+        projected = _PROJECT_DRAFT_WORKSPACE(
             get_draft_workspace_record(
                 self._draft_store(),
                 workspace_id=workspace_id,
                 include_draft=include_draft,
             )
         )
+        if include_draft:
+            return cast(DraftWorkspaceWithDocument, projected)
+        return projected
 
     async def delete_draft_workspace(
         self, *, workspace_id: str

src/wf_api/models/__init__.py

@@ -73,6 +73,7 @@ from .drafts import (
     DraftDiagnosticPayload,
     DraftWorkspaceResult,
     DraftWorkspaceSummary,
+    DraftWorkspaceWithDocument,
     InvalidDraftResult,
     ListDraftWorkspacesResult,
     PatchDraftResult,
@@ -159,6 +160,7 @@ __all__ = [
     "DraftDiagnosticPayload",
     "DraftWorkspaceResult",
     "DraftWorkspaceSummary",
+    "DraftWorkspaceWithDocument",
     "HealthResult",
     "GuidedResultPayload",
     "InterruptPayload",

src/wf_api/models/drafts.py

@@ -40,6 +40,21 @@ class DraftWorkspaceResult(TypedDict):
     draft: NotRequired[JsonObject]
 
 
+class DraftWorkspaceWithDocument(TypedDict):
+    """Persisted workspace envelope that includes the requested draft document."""
+
+    # Keep this as a sibling rather than inheriting DraftWorkspaceResult. The
+    # latter's optional key is not a valid base for a required key, and the
+    # sibling keeps the generated OpenRPC schema stable.
+    workspace_id: str
+    revision: int
+    title: str | None
+    status: Literal["valid", "invalid"]
+    diagnostics: list[DraftDiagnosticPayload]
+    summary: DraftWorkspaceSummary
+    draft: JsonObject
+
+
 class ListDraftWorkspacesResult(TypedDict):
     """All persisted draft-workspace summaries."""
 

src/wf_api/service.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
-from typing import Any
+from typing import Any, Literal, overload
 
 from wf_artifacts import ArtifactKind, compile_workflow_draft
 from wf_artifacts.drafts.models import DraftStep
@@ -32,6 +32,7 @@ from .models import (
     DeleteDeploymentResult,
     DeleteDraftWorkspaceResult,
     DraftWorkspaceResult,
+    DraftWorkspaceWithDocument,
     InspectCapabilityResult,
     ListArtifactsResult,
     ListCapabilitiesResult,
@@ -358,12 +359,28 @@ class WorkflowApi:
             outcomes=outcomes,
         )
 
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[True],
+    ) -> DraftWorkspaceWithDocument: ...
+
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[False] = False,
+    ) -> DraftWorkspaceResult: ...
+
     async def get_draft_workspace(
         self,
         *,
         workspace_id: str,
         include_draft: bool = False,
-    ) -> DraftWorkspaceResult:
+    ) -> DraftWorkspaceResult | DraftWorkspaceWithDocument:
         return await self.drafts.get_draft_workspace(
             workspace_id=workspace_id,
             include_draft=include_draft,

src/wf_api/surface.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Protocol
+from typing import Any, Literal, Protocol, overload
 
 from wf_artifacts import ArtifactKind
 from wf_artifacts.drafts.models import DraftStep
@@ -22,6 +22,7 @@ from .models import (
     DeleteDeploymentResult,
     DeleteDraftWorkspaceResult,
     DraftWorkspaceResult,
+    DraftWorkspaceWithDocument,
     InspectCapabilityResult,
     InspectRegistryEntryResult,
     InspectSourceResult,
@@ -101,12 +102,28 @@ class WorkflowDraftSurface(Protocol):
 
     async def list_draft_workspaces(self) -> ListDraftWorkspacesResult: ...
 
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[True],
+    ) -> DraftWorkspaceWithDocument: ...
+
+    @overload
+    async def get_draft_workspace(
+        self,
+        *,
+        workspace_id: str,
+        include_draft: Literal[False] = False,
+    ) -> DraftWorkspaceResult: ...
+
     async def get_draft_workspace(
         self,
         *,
         workspace_id: str,
         include_draft: bool = False,
-    ) -> DraftWorkspaceResult: ...
+    ) -> DraftWorkspaceResult | DraftWorkspaceWithDocument: ...
 
     async def inspect_draft_authoring_contract(
         self,