lda-wf/tests/authoring/test_demo_workflow.py

from __future__ import annotations

import pytest
from pydantic import BaseModel

from examples.demo_workflow import build_demo_registry, build_demo_workflow
from wf_authoring import (
    NodeReturn,
    WorkflowBuilder,
    build_registry,
    context_path,
    input_from,
    input_path,
    node,
    output_to,
    state,
    state_path,
)
from wf_core import (
    END,
    FrameStatus,
    RunStatus,
    RuntimeContext,
    WorkflowExecutionError,
    execute_workflow,
    resume_workflow,
    step_workflow,
)
from wf_core.runtime.ops.runs import create_run_state


class DriveListFilesInput(BaseModel):
    folder_id: str


class DriveListFilesOutput(BaseModel):
    documents: list[str]


class SummarizeDocumentInput(BaseModel):
    document: str


class SummarizeDocumentOutput(BaseModel):
    item_summary: str


class CombineSummariesInput(BaseModel):
    item_summaries: list[str]


class CombineSummariesOutput(BaseModel):
    summary: str


class SendEmailInput(BaseModel):
    summary: str


class SendEmailOutput(BaseModel):
    email_status: str


class MarkEmailSkippedInput(BaseModel):
    pass


class MarkEmailSkippedOutput(BaseModel):
    email_status: str


@node(
    name="drive_list_files",
    input_model=DriveListFilesInput,
    output_model=DriveListFilesOutput,
)
def drive_list_files_spec(
    payload: DriveListFilesInput,
    ctx: RuntimeContext,
) -> DriveListFilesOutput:
    return DriveListFilesOutput(
        documents=[
            f"{payload.folder_id}/meeting-notes.md",
            f"{payload.folder_id}/weekly-report.md",
        ]
    )


@node(
    name="summarize_document",
    input_model=SummarizeDocumentInput,
    output_model=SummarizeDocumentOutput,
)
def summarize_document_spec(
    payload: SummarizeDocumentInput,
    ctx: RuntimeContext,
) -> SummarizeDocumentOutput:
    return SummarizeDocumentOutput(item_summary=f"Summary of {payload.document}")


@node(
    name="combine_summaries",
    input_model=CombineSummariesInput,
    output_model=CombineSummariesOutput,
)
def combine_summaries_spec(
    payload: CombineSummariesInput,
    ctx: RuntimeContext,
) -> CombineSummariesOutput:
    return CombineSummariesOutput(summary=" | ".join(payload.item_summaries))


@node(
    name="send_email",
    input_model=SendEmailInput,
    output_model=SendEmailOutput,
    outcomes=("sent",),
)
def send_email_spec(
    payload: SendEmailInput,
    ctx: RuntimeContext,
) -> NodeReturn[SendEmailOutput]:
    return NodeReturn(
        outcome="sent",
        output=SendEmailOutput(email_status=f"sent: {payload.summary}"),
    )


@node(
    name="mark_email_skipped",
    input_model=MarkEmailSkippedInput,
    output_model=MarkEmailSkippedOutput,
)
def mark_email_skipped_spec(
    payload: MarkEmailSkippedInput,
    ctx: RuntimeContext,
) -> MarkEmailSkippedOutput:
    return MarkEmailSkippedOutput(email_status="skipped")


def build_authoring_demo_workflow():
    declared = build_demo_workflow()
    builder = WorkflowBuilder(
        name=declared.name,
        input_schema=declared.input_schema,
        state_schema=declared.state_schema,
        output_schema=declared.output_schema,
        start="list_files",
    )

    list_files = builder.use(
        drive_list_files_spec,
        id="list_files",
        input=[input_from(input_path("folder_id"), "folder_id")],
        output=[output_to("documents", state_path("documents"))],
        desc="List files from a Google Drive folder",
    )
    summarize_each = builder.foreach(
        id="summarize_each",
        over=state_path("documents"),
        as_="document",
        mode="serial",
        on_item_error="fail",
    )
    summarize_one = builder.use(
        summarize_document_spec,
        id="summarize_one",
        input=[input_from(context_path("document"), "document")],
        output=[output_to("item_summary", state_path("item_summaries"))],
        desc="Summarize one document",
    )
    combine_summaries = builder.use(
        combine_summaries_spec,
        id="combine_summaries",
        input=[input_from(state_path("item_summaries"), "item_summaries")],
        output=[output_to("summary", state_path("summary"))],
        desc="Combine item summaries into one final summary",
    )
    should_email = builder.condition(
        id="should_email",
        check=state("should_email").eq(True),
    )
    send_email = builder.use(
        send_email_spec,
        id="send_email",
        input=[input_from(state_path("summary"), "summary")],
        output=[output_to("email_status", state_path("email_status"))],
        desc="Send the summary by email",
    )
    approve_email = builder.interrupt(
        id="approve_email",
        kind="approval",
        request=[
            input_from(state_path("summary"), "summary"),
            input_from(input_path("folder_id"), "folder_id"),
        ],
        resume=[
            output_to("approved", state_path("approved")),
            output_to("comment", state_path("approval_comment")),
        ],
        outcomes=["submitted", "cancelled"],
    )
    skip_email = builder.use(
        mark_email_skipped_spec,
        id="skip_email",
        output=[output_to("email_status", state_path("email_status"))],
        desc="Record that email delivery was skipped",
    )

    builder.connect(list_files, "ok", summarize_each)
    builder.connect(summarize_each, "loop", summarize_one)
    builder.connect(summarize_each, "done", combine_summaries)
    builder.connect(summarize_one, "ok", END)
    builder.connect(combine_summaries, "ok", should_email)
    builder.connect(should_email, "true", approve_email)
    builder.connect(should_email, "false", skip_email)
    builder.connect(approve_email, "submitted", send_email)
    builder.connect(approve_email, "cancelled", skip_email)
    builder.connect(send_email, "sent", END)
    builder.connect(skip_email, "ok", END)

    registry = build_registry(
        drive_list_files_spec,
        summarize_document_spec,
        combine_summaries_spec,
        send_email_spec,
        mark_email_skipped_spec,
    )
    return builder.compile(), registry


def _strip_schema_titles(value: object) -> object:
    if isinstance(value, dict):
        normalized = {
            key: _strip_schema_titles(inner)
            for key, inner in value.items()
            if key not in {"title", "items"}
        }
        return normalized
    if isinstance(value, list):
        return [_strip_schema_titles(item) for item in value]
    return value


def test_interrupt_then_resume_to_send_email() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()

    interrupted_run = execute_workflow(
        workflow,
        {"folder_id": "demo-folder", "should_email": True},
        registry,
    )

    assert interrupted_run.status == RunStatus.INTERRUPTED
    assert interrupted_run.current_node_id == "approve_email"
    assert interrupted_run.interrupt is not None
    assert interrupted_run.interrupt.kind == "approval"
    assert interrupted_run.state["summary"].startswith("Summary of demo-folder/")

    resumed_run = resume_workflow(
        workflow,
        interrupted_run,
        registry,
        resume_payload={"approved": True, "comment": "Looks good to send."},
        resume_outcome="submitted",
    )

    assert resumed_run.status == RunStatus.COMPLETED
    assert resumed_run.current_node_id == END
    assert resumed_run.output["email_status"].startswith("sent:")
    assert resumed_run.state["approved"] is True
    assert resumed_run.state["approval_comment"] == "Looks good to send."


def test_non_interrupt_path_skips_email() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()

    run = execute_workflow(
        workflow,
        {"folder_id": "demo-folder", "should_email": False},
        registry,
    )

    assert run.status == RunStatus.COMPLETED
    assert run.output["email_status"] == "skipped"
    assert run.interrupt is None


def test_workflow_input_schema_rejects_wrong_type() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()

    with pytest.raises(
        WorkflowExecutionError,
        match=r"workflow input\['should_email'\].*not of type 'boolean'",
    ):
        execute_workflow(
            workflow,
            {"folder_id": "demo-folder", "should_email": "false"},
            registry,
        )


def test_node_output_schema_rejects_wrong_type() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()

    def bad_list_files(
        payload: dict[str, object], ctx: RuntimeContext
    ) -> dict[str, object]:
        return {"outcome": "ok", "output": {"documents": "not-a-list"}}

    registry["drive_list_files"] = bad_list_files

    with pytest.raises(
        WorkflowExecutionError,
        match=r"node output for list_files\['documents'\].*not of type 'array'",
    ):
        execute_workflow(
            workflow,
            {"folder_id": "demo-folder", "should_email": False},
            registry,
        )


def test_stepwise_execution_reaches_interrupt() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()
    run = create_run_state(
        workflow,
        {"folder_id": "demo-folder", "should_email": True},
    )

    workflow.validate_structure().raise_for_errors()

    while run.status not in {RunStatus.INTERRUPTED, RunStatus.COMPLETED}:
        step_workflow(workflow, run, registry)

    assert run.status == RunStatus.INTERRUPTED
    assert run.current_node_id == "approve_email"
    assert any(entry.step_type == "foreach" for entry in run.trace)
    assert len(run.trace) == 9


def test_foreach_stress_with_many_documents() -> None:
    workflow = build_demo_workflow()
    registry = build_demo_registry()

    document_count = 25

    def many_files(payload: dict[str, object], ctx: object) -> dict[str, object]:
        folder_id = payload["folder_id"]
        return {
            "outcome": "ok",
            "output": {
                "documents": [
                    f"{folder_id}/doc-{index:02d}.md" for index in range(document_count)
                ]
            },
        }

    registry["drive_list_files"] = many_files

    run = execute_workflow(
        workflow,
        {"folder_id": "bulk-folder", "should_email": False},
        registry,
    )

    assert run.status == RunStatus.COMPLETED
    assert len(run.state["documents"]) == document_count
    assert len(run.state["item_summaries"]) == document_count
    assert (
        len(
            [
                frame
                for frame in run.frames.values()
                if frame.kind == "foreach_iteration"
                and frame.status == FrameStatus.COMPLETED
            ]
        )
        == document_count
    )
    assert len([entry for entry in run.trace if entry.step_type == "foreach"]) == (
        document_count + 1
    )


def test_builder_compiles_same_workflow_as_declared_demo() -> None:
    declared = build_demo_workflow()
    built, _registry = build_authoring_demo_workflow()

    assert _strip_schema_titles(
        built.model_dump(by_alias=True)
    ) == _strip_schema_titles(declared.model_dump(by_alias=True))


def test_builder_compiled_workflow_executes_like_declared_demo() -> None:
    declared = build_demo_workflow()
    declared_registry = build_demo_registry()
    built, built_registry = build_authoring_demo_workflow()

    declared_run = execute_workflow(
        declared,
        {"folder_id": "demo-folder", "should_email": False},
        declared_registry,
    )
    built_run = execute_workflow(
        built,
        {"folder_id": "demo-folder", "should_email": False},
        built_registry,
    )

    assert built_run.status == declared_run.status
    assert built_run.output == declared_run.output
    assert built_run.state == declared_run.state
    assert built_run.current_node_id == declared_run.current_node_id


def test_async_node_spec_cannot_export_sync_registry_handler() -> None:
    class AsyncInput(BaseModel):
        value: str

    class AsyncOutput(BaseModel):
        echoed: str

    @node(
        name="async_echo",
        input_model=AsyncInput,
        output_model=AsyncOutput,
        is_async=True,
    )
    async def async_echo(
        payload: AsyncInput,
        ctx: RuntimeContext,
    ) -> AsyncOutput:
        return AsyncOutput(echoed=payload.value)

    try:
        async_echo.to_registry_handler()
    except TypeError as exc:
        assert "async" in str(exc)
    else:
        raise AssertionError("expected async node export to fail for sync registry")