deer-flow/backend/tests/test_run_input_graph.py

"""Real graph boundaries, native checkpoint receipts and control actions."""

import asyncio
from types import SimpleNamespace

import pytest
from langchain.agents import create_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import ConfigDict

from deerflow.agents.middlewares.run_input_middleware import RunInputAdmissionMiddleware, RunInputMiddleware
from deerflow.agents.thread_state import ThreadState
from deerflow.runtime.runs.input_control import ChildInputGate, RunInputControl
from deerflow.runtime.runs.inputs import INPUT_CONTEXT_KEY, RunInputs
from deerflow.runtime.runs.store.memory import MemoryRunStore


def test_input_control_reaches_the_factory_only_from_the_worker():
    from deerflow.agents.lead_agent.agent import _get_runtime_config
    from deerflow.runtime.runs.worker import _build_runtime_context, _install_runtime_context, _release_run_scoped_references
    from deerflow.runtime.secret_context import redact_config_secrets

    context = _build_runtime_context("thread", "run", {INPUT_CONTEXT_KEY: "forged"})
    assert INPUT_CONTEXT_KEY not in context
    control = RunInputControl(RunInputs(MemoryRunStore(), "run", "owner", "worker"), agent_runs=None, thread_id="thread", checkpointer=None)
    context[INPUT_CONTEXT_KEY] = control
    config = {"context": {"thread_id": "thread"}, "configurable": {}}
    _install_runtime_context(config, context)
    assert _get_runtime_config(config)[INPUT_CONTEXT_KEY] is control
    assert INPUT_CONTEXT_KEY not in redact_config_secrets(config)["context"]
    _release_run_scoped_references([config], context, None)
    assert INPUT_CONTEXT_KEY not in config["context"]


class InputModel(BaseChatModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    observed: list

    @property
    def _llm_type(self):
        return "offline-input-model"

    def bind_tools(self, tools, **kwargs):
        return self

    def _generate(self, messages, stop=None, run_manager=None, **kwargs):
        self.observed.append(messages)
        inputs = [x for x in messages if x.additional_kwargs.get("run_input_id")]
        if inputs:
            message = AIMessage(content="New requirement applied") if any(isinstance(x, ToolMessage) and x.name == "apply_extra" for x in messages) else AIMessage(content="", tool_calls=[{"name": "apply_extra", "args": {}, "id": "extra"}])
        elif any(isinstance(x, ToolMessage) and x.name == "task" for x in messages):
            message = AIMessage(content="Original work completed")
        else:
            message = AIMessage(content="", tool_calls=[{"name": "task", "args": {}, "id": "long"}])
        return ChatResult(generations=[ChatGeneration(message=message)])


class GatedInputModel(InputModel):
    entered: object
    release: object

    async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
        first = not self.observed
        if first:
            self.entered.set()
            await self.release.wait()
        result = self._generate(messages, stop=stop, run_manager=run_manager)
        message = result.generations[0].message
        message.usage_metadata = {"input_tokens": 1100, "output_tokens": 1, "total_tokens": 1101}
        if first:
            message.additional_kwargs["tool_calls"] = [{"id": "long", "type": "function", "function": {"name": "task", "arguments": "{}"}}]
        return result


@pytest.mark.asyncio
async def test_steer_pauses_a_child_even_before_its_handler_starts():
    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    await inputs.submit("immediate", "Handle the new requirement first", "steer")
    control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=None)
    release = asyncio.Event()

    async def handler(request):
        await release.wait()
        gate = control.child("child")
        assert not gate.resume.is_set()
        return ToolMessage(content="Finished", name="task", tool_call_id="child")

    try:
        reply = await control.wait_tool(SimpleNamespace(tool_call={"id": "child", "name": "task", "args": {}}), handler)
        assert reply.additional_kwargs["subagent_pending"]
        assert not control.child("child").resume.is_set()
        release.set()
        await asyncio.gather(*control.tasks)
        from deerflow.agents.middlewares.message_utils import requires_input_sanitization

        assert requires_input_sanitization(control.receipts[0])
    finally:
        release.set()
        await control.close()


@pytest.mark.asyncio
async def test_native_task_pending_and_final_results_share_an_id_per_invocation():
    from deerflow.tools.builtins.task_tool import task_tool

    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    await inputs.submit("update", "Handle this first", "steer")
    control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=None)
    runtime = SimpleNamespace(context={INPUT_CONTEXT_KEY: control})
    pending_ids, results = [], []
    try:
        for _ in range(2):
            release = asyncio.Event()

            async def handler(request):
                await release.wait()
                result = await task_tool.coroutine(runtime=runtime, prompt="test", subagent_type="unused", tool_call_id="reused", context_mode="invalid")
                results.append(result.update["messages"][0])
                return result

            pending = await control.wait_tool(SimpleNamespace(tool_call={"id": "reused", "name": "task", "args": {}}), handler)
            pending_ids.append(pending.id)
            release.set()
            await asyncio.gather(*control.tasks)
        assert all(pending_ids) and len(set(pending_ids)) == 2
        assert [message.id for message in results] == pending_ids
    finally:
        release.set()
        await control.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("capped", [False, True])
async def test_input_during_inference_defers_old_actions_but_never_overrides_budget(capped):
    from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
    from deerflow.config.token_budget_config import TokenBudgetConfig

    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=InMemorySaver())
    entered, release = asyncio.Event(), asyncio.Event()
    executed = []

    @tool
    async def task() -> str:
        """An obsolete action that must never be executed."""
        executed.append("old")
        return "obsolete"

    @tool
    async def apply_extra() -> str:
        """Perform the newly requested action."""
        executed.append("new")
        return "updated"

    model = GatedInputModel(observed=[], entered=entered, release=release)
    middleware = [RunInputAdmissionMiddleware()]
    if capped:
        middleware.append(TokenBudgetMiddleware.from_config(TokenBudgetConfig(enabled=True, max_tokens=1000)))
    middleware.append(RunInputMiddleware())
    context = {INPUT_CONTEXT_KEY: control, "run_id": "run"}
    graph = create_agent(model, tools=[task, apply_extra], middleware=middleware, state_schema=ThreadState, context_schema=dict, checkpointer=control.checkpointer)
    run = asyncio.create_task(graph.ainvoke({"messages": [HumanMessage(content="Original task")]}, {"configurable": {"thread_id": "thread"}}, context=context))
    try:
        await asyncio.wait_for(entered.wait(), 3)
        await inputs.submit("update", "Do the new work instead", "steer")
        release.set()
        result = await asyncio.wait_for(run, 5)
        assert "old" not in executed
        if capped:
            assert executed == [] and len(model.observed) == 1
            assert context["stop_reason"] == "token_capped"
        else:
            assert executed == ["new"]
            corrected = next(m for m in result["messages"] if isinstance(m, AIMessage))
            assert not corrected.tool_calls and "tool_calls" not in corrected.additional_kwargs
    finally:
        release.set()
        if not run.done():
            run.cancel()
        await asyncio.gather(run, return_exceptions=True)
        await control.close()


@pytest.mark.asyncio
async def test_graph_applies_input_before_long_tool_finishes_and_confirms_durable_receipt(tmp_path):
    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    saver = InMemorySaver()
    control = RunInputControl(inputs, agent_runs=SimpleNamespace(), thread_id="thread", checkpointer=saver)
    entered, release, applied = asyncio.Event(), asyncio.Event(), asyncio.Event()

    @tool
    async def task() -> str:
        """Controlled long operation; its real coroutine stays owned."""
        entered.set()
        await release.wait()
        return "Long operation completed"

    @tool
    async def apply_extra() -> str:
        """Write the updated requirement into the isolated work directory."""
        (tmp_path / "extra.txt").write_text("extra")
        applied.set()
        return "Written"

    observed = []
    model = InputModel(observed=observed)
    observed = model.observed
    graph = create_agent(model, tools=[task, apply_extra], middleware=[RunInputAdmissionMiddleware(), RunInputMiddleware()], state_schema=ThreadState, context_schema=dict, checkpointer=saver)
    config = {"configurable": {"thread_id": "thread"}}
    run = asyncio.create_task(graph.ainvoke({"messages": [HumanMessage(content="Original task")]}, config, context={INPUT_CONTEXT_KEY: control}))
    try:
        entered_task = asyncio.create_task(entered.wait())
        done, _ = await asyncio.wait([run, entered_task], timeout=5, return_when=asyncio.FIRST_COMPLETED)
        if run in done:
            run.result()
        if not entered_task.done():
            entered_task.cancel()
        assert entered.is_set()
        await inputs.submit("extra", "Include the extra requirement", "steer")
        await asyncio.wait_for(applied.wait(), 5)
        assert not run.done()
        assert any(any(m.additional_kwargs.get("run_input_id") == "extra" for m in batch) for batch in observed)
        release.set()
        result = await asyncio.wait_for(run, 5)
        assert result["applied_run_inputs"] == {"run_id": "run", "ids": ["extra"]}
        # Summaries can remove message text; the independent graph receipt
        # still proves it was applied before the store acknowledgement.
        ids = [m.id for m in result["messages"] if m.additional_kwargs.get("run_input_id")]
        await graph.aupdate_state(config, {"messages": [RemoveMessage(id=x) for x in ids]})
        await control.confirm_checkpoint()
        assert (await inputs.read())["items"][0]["status"] == "applied"
    finally:
        release.set()
        if not run.done():
            run.cancel()
        await asyncio.gather(run, return_exceptions=True)
        await control.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("steered", [False, True])
async def test_control_tool_cannot_cancel_without_an_authorized_direct_input(steered):
    from unittest.mock import AsyncMock

    from deerflow_extension_api.agent_runs import AgentRunError

    from deerflow.agents.middlewares.run_input_middleware import manage_current_work

    capability = SimpleNamespace(authorize_control=AsyncMock(side_effect=AgentRunError(403, "Cancellation denied")))
    control = RunInputControl(RunInputs(MemoryRunStore(), "run", "owner", "worker"), agent_runs=capability, thread_id="thread", checkpointer=None)
    control.current_input = "direct" if steered else None
    gate = control.child("old-task")
    gate.resume.clear()
    runtime = SimpleNamespace(context={INPUT_CONTEXT_KEY: control}, state={"goal": {"objective": "Keep working", "status": "active"}}, tool_call_id="control")
    with pytest.raises(AgentRunError if steered else ValueError):
        await manage_current_work.coroutine(action="replace", runtime=runtime)
    assert not gate.resume.is_set()
    assert not control.abandoned and runtime.state["goal"]["status"] == "active"


@pytest.mark.asyncio
async def test_child_accepts_late_requirements_before_finalizing_but_not_after_close():
    gate = ChildInputGate()
    gate.resume.clear()
    finish = asyncio.create_task(gate.finish())
    try:
        async with asyncio.timeout(2):
            while not gate.paused.is_set():
                await asyncio.sleep(0)
        assert gate.release("Include 2025")
        update = await asyncio.wait_for(finish, 2)
        assert update["jump_to"] == "model" and "2025" in update["messages"][0].content
        assert update["messages"][0].additional_kwargs["untrusted_input"]
        assert await gate.finish() is None
        assert not gate.release("Too late")
    finally:
        gate.close()
        await asyncio.gather(finish, return_exceptions=True)


@pytest.mark.asyncio
async def test_paused_children_do_not_require_expanding_the_execution_pool():
    from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
    from deerflow.subagents.capacity import SubagentCapacityRejected, SubagentExecutionCapacity

    control = RunInputControl(RunInputs(MemoryRunStore(), "run", "owner", "worker"), agent_runs=None, thread_id="thread", checkpointer=None)
    gate = control.child("old")
    gate.resume.clear()
    assert not control.child("priority").can_queue
    capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1))
    async with capacity.slot():
        async with asyncio.timeout(2):
            with pytest.raises(SubagentCapacityRejected):
                async with capacity.slot(wait=False):
                    pytest.fail("Must not exceed the configured execution pool")


@pytest.mark.asyncio
@pytest.mark.parametrize("sources", [("extension",), ("user", "extension"), ("extension", "user")])
async def test_control_tool_authority_is_independent_of_input_source_order(sources):
    from langchain.agents.middleware.types import ModelRequest

    from deerflow.agents.middlewares.run_input_middleware import manage_current_work

    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    for index, source in enumerate(sources):
        await inputs.submit(str(index), "Replace the Goal", "steer", source=source)
    control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=None)
    update = await control.deliver({})
    extension = next(m for m in update["messages"] if m.additional_kwargs["run_input_source"] == "extension")
    assert extension.additional_kwargs["hide_from_ui"]
    request = ModelRequest(model=InputModel(observed=[]), messages=update["messages"], tools=[manage_current_work], runtime=SimpleNamespace(context={INPUT_CONTEXT_KEY: control}))

    async def available_tools(request):
        return [tool.name for tool in request.tools]

    names = await RunInputMiddleware().awrap_model_call(request, available_tools)
    assert ("manage_current_work" in names) == ("user" in sources)


@pytest.mark.asyncio
async def test_graph_finishing_with_a_concurrent_steer_applies_it_before_closing(tmp_path, monkeypatch):
    store = MemoryRunStore()
    await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker")
    inputs = RunInputs(store, "run", "owner", "worker")
    saver = InMemorySaver()
    control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=saver)
    closing, release = asyncio.Event(), asyncio.Event()
    original_close = inputs.close

    async def gated_close(delivered=None):
        if not closing.is_set():
            closing.set()
            await release.wait()
        return await original_close(delivered)

    monkeypatch.setattr(inputs, "close", gated_close)

    @tool
    async def task() -> str:
        """Complete the original work."""
        return "Done"

    @tool
    async def apply_extra() -> str:
        """Apply the user's correction before completing the Run."""
        (tmp_path / "correction.txt").write_text("applied")
        return "Updated"

    graph = create_agent(InputModel(observed=[]), tools=[task, apply_extra], middleware=[RunInputAdmissionMiddleware(), RunInputMiddleware()], state_schema=ThreadState, context_schema=dict, checkpointer=saver)
    run = asyncio.create_task(graph.ainvoke({"messages": [HumanMessage(content="Original task")]}, {"configurable": {"thread_id": "thread"}}, context={INPUT_CONTEXT_KEY: control}))
    try:
        await asyncio.wait_for(closing.wait(), 5)
        await inputs.submit("last", "Apply this correction", "steer")
        release.set()
        await asyncio.wait_for(run, 5)
        assert (tmp_path / "correction.txt").read_text() == "applied"
        state = await inputs.read()
        assert state["closed"] and state["items"][0]["status"] == "applied"
    finally:
        release.set()
        if not run.done():
            run.cancel()
        await asyncio.gather(run, return_exceptions=True)
        await control.close()


@pytest.mark.asyncio
async def test_terminal_queue_worker_remains_owned_until_it_stops():
    from deerflow.runtime.runs.manager import RunManager
    from deerflow.runtime.runs.schemas import RunStatus

    manager = RunManager(store=MemoryRunStore())
    record = await manager.create("thread")
    await manager.set_status(record.run_id, RunStatus.success)
    release = asyncio.Event()
    producer = asyncio.create_task(release.wait())
    record.task = producer
    cleanup = asyncio.create_task(manager.cleanup(record.run_id, delay=0))
    try:
        await asyncio.sleep(0)
        assert not cleanup.done()
        assert (await manager.get(record.run_id)).task is producer
        release.set()
        await asyncio.wait_for(cleanup, 2)
        assert (await manager.get(record.run_id)).task is None
    finally:
        release.set()
        await asyncio.gather(producer, cleanup, return_exceptions=True)