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,

src/wf_transport_rpc_http/client/drafts.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any, Literal, cast
+from typing import Any, Literal, cast, overload
 
 from wf_api import CapabilityStepUpdate
 from wf_api.models import (
@@ -11,6 +11,7 @@ from wf_api.models import (
     CreateDraftWorkspaceFromCapabilityResult,
     DeleteDraftWorkspaceResult,
     DraftWorkspaceResult,
+    DraftWorkspaceWithDocument,
     ListDraftWorkspacesResult,
     PatchDraftResult,
     ValidateDraftResult,
@@ -22,13 +23,38 @@ from wf_core.models.steps import InputBinding, OutputBinding, StepInputBinding
 from .base import RpcCaller
 
 
+@overload
 async def _call_draft_workspace(
     caller: RpcCaller,
     method: str,
     params: dict[str, Any],
-) -> DraftWorkspaceResult:
+    *,
+    include_draft: Literal[True],
+) -> DraftWorkspaceWithDocument: ...
+
+
+@overload
+async def _call_draft_workspace(
+    caller: RpcCaller,
+    method: str,
+    params: dict[str, Any],
+    *,
+    include_draft: Literal[False] = False,
+) -> DraftWorkspaceResult: ...
+
+
+async def _call_draft_workspace(
+    caller: RpcCaller,
+    method: str,
+    params: dict[str, Any],
+    *,
+    include_draft: bool = False,
+) -> DraftWorkspaceResult | DraftWorkspaceWithDocument:
     """Call one server-validated draft method with its canonical client type."""
-    return cast(DraftWorkspaceResult, await caller._call(method, params))
+    result = await caller._call(method, params)
+    if include_draft:
+        return cast(DraftWorkspaceWithDocument, result)
+    return cast(DraftWorkspaceResult, result)
 
 
 class RpcDraftClientMixin:
@@ -64,17 +90,34 @@ class RpcDraftClientMixin:
             await self._call("workflow.draft_workspaces.list", {}),
         )
 
+    @overload
+    async def get_draft_workspace(
+        self: RpcCaller,
+        *,
+        workspace_id: str,
+        include_draft: Literal[True],
+    ) -> DraftWorkspaceWithDocument: ...
+
+    @overload
+    async def get_draft_workspace(
+        self: RpcCaller,
+        *,
+        workspace_id: str,
+        include_draft: Literal[False] = False,
+    ) -> DraftWorkspaceResult: ...
+
     async def get_draft_workspace(
         self: RpcCaller,
         *,
         workspace_id: str,
         include_draft: bool = False,
-    ) -> DraftWorkspaceResult:
+    ) -> DraftWorkspaceResult | DraftWorkspaceWithDocument:
         """Return the remote workspace summary or revision-conflict payload."""
         return await _call_draft_workspace(
             self,
             "workflow.draft_workspaces.get",
             {"workspace_id": workspace_id, "include_draft": include_draft},
+            include_draft=include_draft,
         )
 
     async def inspect_draft_authoring_contract(

src/wf_transport_rpc_http/methods/drafts.py

@@ -4,6 +4,8 @@ Return annotations stay eagerly evaluated because fastapi-jsonrpc captures them
 while registering nested handlers for response validation and OpenRPC output.
 """
 
+from typing import cast
+
 import fastapi_jsonrpc as jsonrpc
 
 from wf_api.models import (
@@ -108,9 +110,12 @@ def register_methods(
         params: GetDraftWorkspaceParams = RpcParams(),
     ) -> DraftWorkspaceResult:
         try:
-            return await server.api.get_draft_workspace(
-                workspace_id=params.workspace_id,
-                include_draft=params.include_draft,
+            return cast(
+                DraftWorkspaceResult,
+                await server.api.get_draft_workspace(
+                    workspace_id=params.workspace_id,
+                    include_draft=params.include_draft,
+                ),
             )
         except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
             raise_workflow_rpc_error(exc)