lda-wf/src/wf_transport_rpc_http/methods/drafts.py

"""Draft JSON-RPC method registration.

Return annotations stay eagerly evaluated because fastapi-jsonrpc captures them
while registering nested handlers for response validation and OpenRPC output.
"""

import fastapi_jsonrpc as jsonrpc

from wf_api.models import (
    AuthoringContractInventoryPayload,
    CompileDraftWorkspaceResult,
    CreateArtifactFromWorkspaceResult,
    CreateDraftWorkspaceFromCapabilityResult,
    DeleteDraftWorkspaceResult,
    DraftWorkspaceResult,
    ListDraftWorkspacesResult,
    PatchDraftResult,
    ValidateDraftResult,
)
from wf_api.surface import RouteSource
from wf_server import WorkflowServer

from ..errors import WorkflowRpcError, raise_workflow_rpc_error
from ..models import (
    AddDraftStepParams,
    AddStepFromCapabilityParams,
    BindDraftParams,
    BranchDraftParams,
    CompileDraftWorkspaceParams,
    CreateArtifactFromWorkspaceParams,
    CreateDraftFromCapabilityParams,
    CreateEmptyDraftWorkspaceParams,
    CreateWrapperFromWorkspaceParams,
    DeleteDraftWorkspaceParams,
    GetDraftWorkspaceParams,
    HandleDraftParams,
    InspectDraftAuthoringContractParams,
    ListDraftWorkspacesParams,
    PatchDraftParams,
    PatchDraftWorkspaceParams,
    RemoveDraftBindingParams,
    RemoveDraftRouteParams,
    RemoveDraftStepParams,
    ReplaceDraftWorkspaceDocumentParams,
    SetDraftContractParams,
    SetDraftNameParams,
    SetDraftRouteParams,
    SetDraftStartParams,
    SetStepInputBindingsParams,
    SetStepInputMapParams,
    SetStepOutputBindingsParams,
    SetStepOutputMapParams,
    SetWorkflowOutputBindingsParams,
    SetWorkflowOutputMapParams,
    UpdateCapabilityStepParams,
    ValidateDraftParams,
    ValidateDraftWorkspaceParams,
)
from ..params import RpcParams


def register_methods(
    entrypoint: jsonrpc.Entrypoint,
    server: WorkflowServer,
) -> None:
    """Register draft and draft-workspace JSON-RPC methods."""

    @entrypoint.method(
        name="workflow.drafts.create_from_capability", errors=[WorkflowRpcError]
    )
    async def workflow_drafts_create_from_capability(
        params: CreateDraftFromCapabilityParams = RpcParams(),
    ) -> CreateDraftWorkspaceFromCapabilityResult:
        try:
            return await _create_from_capability(server, params)
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(name="workflow.drafts.patch", errors=[WorkflowRpcError])
    async def workflow_drafts_patch(
        params: PatchDraftParams = RpcParams(),
    ) -> PatchDraftResult:
        try:
            return await server.api.patch_draft(draft=params.draft, patch=params.patch)
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(name="workflow.drafts.validate", errors=[WorkflowRpcError])
    async def workflow_drafts_validate(
        params: ValidateDraftParams = RpcParams(),
    ) -> ValidateDraftResult:
        try:
            return await server.api.validate_draft(draft=params.draft)
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(name="workflow.draft_workspaces.list", errors=[WorkflowRpcError])
    async def workflow_draft_workspaces_list(
        params: ListDraftWorkspacesParams = RpcParams(),
    ) -> ListDraftWorkspacesResult:
        try:
            return await server.api.list_draft_workspaces()
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(name="workflow.draft_workspaces.get", errors=[WorkflowRpcError])
    async def workflow_draft_workspaces_get(
        params: GetDraftWorkspaceParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return 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)

    @entrypoint.method(
        name="workflow.draft_workspaces.inspect_authoring_contract",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_inspect_authoring_contract(
        params: InspectDraftAuthoringContractParams = RpcParams(),
    ) -> AuthoringContractInventoryPayload | DraftWorkspaceResult:
        try:
            return await server.api.inspect_draft_authoring_contract(
                workspace_id=params.workspace_id,
                revision=params.revision,
                selected_step_id=params.selected_step_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.create_from_capability",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_create_from_capability(
        params: CreateDraftFromCapabilityParams = RpcParams(),
    ) -> CreateDraftWorkspaceFromCapabilityResult:
        try:
            return await _create_from_capability(server, params)
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.create_empty",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_create_empty(
        params: CreateEmptyDraftWorkspaceParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.create_empty_draft_workspace(
                workspace_id=params.workspace_id,
                name=params.name,
                title=params.title,
                input_schema=params.input_schema,
                state_schema=params.state_schema,
                output_schema=params.output_schema,
                outcomes=params.outcomes,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.patch", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_patch(
        params: PatchDraftWorkspaceParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.patch_draft_workspace(
                workspace_id=params.workspace_id,
                revision=params.revision,
                patch=params.patch,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.replace_document",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_replace_document(
        params: ReplaceDraftWorkspaceDocumentParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.replace_draft_workspace_document(
                workspace_id=params.workspace_id,
                revision=params.revision,
                draft=params.draft,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_name", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_set_name(
        params: SetDraftNameParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_draft_name(
                workspace_id=params.workspace_id,
                revision=params.revision,
                name=params.name,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_start", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_set_start(
        params: SetDraftStartParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_draft_start(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_contract", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_set_contract(
        params: SetDraftContractParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_draft_contract(
                workspace_id=params.workspace_id,
                revision=params.revision,
                input_schema=params.input_schema,
                state_schema=params.state_schema,
                output_schema=params.output_schema,
                outcomes=params.outcomes,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_route", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_set_route(
        params: SetDraftRouteParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_draft_route(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                outcome=params.outcome,
                target=params.target,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_step_input_map",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_step_input_map(
        params: SetStepInputMapParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_step_input_map(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                input_map=params.input_map,
                merge=params.merge,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_step_input_bindings",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_step_input_bindings(
        params: SetStepInputBindingsParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_step_input_bindings(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                bindings=params.bindings,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.update_capability_step",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_update_capability_step(
        params: UpdateCapabilityStepParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.update_capability_step(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                update=params.update,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_step_output_bindings",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_step_output_bindings(
        params: SetStepOutputBindingsParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_step_output_bindings(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                bindings=params.bindings,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_step_output_map",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_step_output_map(
        params: SetStepOutputMapParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_step_output_map(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                output_map=params.output_map,
                merge=params.merge,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_workflow_output_map",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_workflow_output_map(
        params: SetWorkflowOutputMapParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_workflow_output_map(
                workspace_id=params.workspace_id,
                revision=params.revision,
                output_map=params.output_map,
                merge=params.merge,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.set_workflow_output_bindings",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_set_workflow_output_bindings(
        params: SetWorkflowOutputBindingsParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.set_workflow_output_bindings(
                workspace_id=params.workspace_id,
                revision=params.revision,
                bindings=params.bindings,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.bind",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_bind(
        params: BindDraftParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.bind_draft(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                source_path=params.source_path,
                target_path=params.target_path,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.add_step_from_capability",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_add_step_from_capability(
        params: AddStepFromCapabilityParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.add_step_from_capability(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                capability_name=params.capability_name,
                route_from_step=params.route_from_step,
                route_from_outcome=params.route_from_outcome,
                routes=params.routes,
                input_map=params.input_map,
                input_bindings=params.input_bindings,
                bind_outputs=params.bind_outputs,
                desc=params.desc,
                retry=params.retry,
                timeout_seconds=params.timeout_seconds,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.add_step",
        errors=[WorkflowRpcError],
    )
    async def workflow_draft_workspaces_add_step(
        params: AddDraftStepParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            incoming = (
                None
                if params.incoming is None
                else RouteSource(
                    step_id=params.incoming.step_id,
                    outcome=params.incoming.outcome,
                )
            )
            return await server.api.add_step(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                step=params.step,
                incoming=incoming,
                routes=params.routes,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.branch", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_branch(
        params: BranchDraftParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.branch_draft(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                routes=params.routes,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.handle", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_handle(
        params: HandleDraftParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.handle_draft(
                workspace_id=params.workspace_id,
                revision=params.revision,
                branches=[
                    {"step_id": b.step_id, "outcome": b.outcome}
                    for b in params.branches
                ],
                target=params.target,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.validate", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_validate(
        params: ValidateDraftWorkspaceParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.validate_draft_workspace(
                workspace_id=params.workspace_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.compile", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_compile(
        params: CompileDraftWorkspaceParams = RpcParams(),
    ) -> CompileDraftWorkspaceResult:
        try:
            return await server.api.compile_draft_workspace(
                workspace_id=params.workspace_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.delete", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_delete(
        params: DeleteDraftWorkspaceParams = RpcParams(),
    ) -> DeleteDraftWorkspaceResult:
        try:
            return await server.api.delete_draft_workspace(
                workspace_id=params.workspace_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.remove_route", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_remove_route(
        params: RemoveDraftRouteParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.remove_draft_route(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                outcome=params.outcome,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.remove_step", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_remove_step(
        params: RemoveDraftStepParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.remove_draft_step(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.remove_binding", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_remove_binding(
        params: RemoveDraftBindingParams = RpcParams(),
    ) -> DraftWorkspaceResult:
        try:
            return await server.api.remove_draft_binding(
                workspace_id=params.workspace_id,
                revision=params.revision,
                step_id=params.step_id,
                inputs=params.inputs,
                outputs=params.outputs,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.create_artifact", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_create_artifact(
        params: CreateArtifactFromWorkspaceParams = RpcParams(),
    ) -> CreateArtifactFromWorkspaceResult:
        try:
            return await server.api.create_artifact_from_workspace(
                workspace_id=params.workspace_id,
                artifact_id=params.artifact_id,
                version=params.version,
                title=params.title,
                outcomes=tuple(params.outcomes),
                kind=params.kind,
                description=params.description,
                required_capabilities=params.required_capabilities,
                source_bindings=params.source_bindings,
                created_from_catalog_version=params.created_from_catalog_version,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)

    @entrypoint.method(
        name="workflow.draft_workspaces.create_wrapper", errors=[WorkflowRpcError]
    )
    async def workflow_draft_workspaces_create_wrapper(
        params: CreateWrapperFromWorkspaceParams = RpcParams(),
    ) -> CreateArtifactFromWorkspaceResult:
        try:
            return await server.api.create_wrapper_from_workspace(
                workspace_id=params.workspace_id,
                artifact_id=params.artifact_id,
                version=params.version,
                title=params.title,
                outcomes=tuple(params.outcomes),
                description=params.description,
                required_capabilities=params.required_capabilities,
                source_bindings=params.source_bindings,
                created_from_catalog_version=params.created_from_catalog_version,
            )
        except (ValueError, KeyError, LookupError, FileNotFoundError) as exc:
            raise_workflow_rpc_error(exc)


async def _create_from_capability(
    server: WorkflowServer,
    params: CreateDraftFromCapabilityParams,
) -> CreateDraftWorkspaceFromCapabilityResult:
    return await server.api.create_draft_workspace_from_capability(
        workspace_id=params.workspace_id,
        capability_name=params.capability_name,
        name=params.name,
        title=params.title,
        input_schema=params.input_schema,
        state_schema=params.state_schema,
        output_schema=params.output_schema,
        input=params.input,
        output=params.output,
        input_map=params.input_map,
        output_map=params.output_map,
        error_message_source=params.error_message_source,
    )