"""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() async def _start_parallel_tasks(inputs, control, handlers): """Steer only once every native handler runs, so each waiter sees all gates.""" started = asyncio.Event() running = set() def tracked(call_id): async def handler(request): running.add(call_id) if len(running) == len(handlers): started.set() return await handlers[call_id](request) return handler waiters = [asyncio.create_task(control.wait_tool(SimpleNamespace(tool_call={"id": call_id, "name": "task", "args": {}}), tracked(call_id))) for call_id in handlers] await asyncio.wait_for(started.wait(), 2) await inputs.submit("update", "Handle this first", "steer") return await asyncio.wait_for(asyncio.gather(*waiters), 5) @pytest.mark.asyncio async def test_parallel_pending_results_keep_their_own_ids(): 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=None) release = asyncio.Event() async def handler(request): await release.wait() return ToolMessage(content="Finished", name="task", tool_call_id=request.tool_call["id"]) try: replies = await _start_parallel_tasks(inputs, control, {"first": handler, "second": handler}) assert [reply.tool_call_id for reply in replies] == ["first", "second"] assert [reply.id for reply in replies] == [control.children[reply.tool_call_id][0].result_id for reply in replies] assert replies[0].id != replies[1].id finally: release.set() await control.close() @pytest.mark.asyncio async def test_finished_task_closes_only_its_own_gate(): 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=None) finished, release = asyncio.Event(), asyncio.Event() async def short_task(request): await finished.wait() return ToolMessage(content="Finished", name="task", tool_call_id="short") async def long_task(request): await release.wait() return ToolMessage(content="Finished", name="task", tool_call_id="long") try: await _start_parallel_tasks(inputs, control, {"short": short_task, "long": long_task}) finished.set() short_gate, long_gate = control.children["short"][0], control.children["long"][0] async with asyncio.timeout(2): while not short_gate.closed: await asyncio.sleep(0) assert not long_gate.closed and not long_gate.resume.is_set() control.interrupted_children = {"long"} assert control.resume_children("Include 2025") == ["long"] finally: finished.set() release.set() await control.close() class ParallelTaskModel(InputModel): def _generate(self, messages, stop=None, run_manager=None, **kwargs): self.observed.append(messages) if len(self.observed) == 1: message = AIMessage(content="", tool_calls=[{"name": "task", "args": {}, "id": "first"}, {"name": "task", "args": {}, "id": "second"}]) else: message = AIMessage(content="New requirement applied") return ChatResult(generations=[ChatGeneration(message=message)]) @pytest.mark.asyncio async def test_graph_keeps_a_pending_result_for_every_parallel_task(): from deerflow.agents.middlewares.dangling_tool_call_middleware import DanglingToolCallMiddleware 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 = asyncio.Event(), asyncio.Event() started = [] @tool async def task() -> str: """Controlled long operation; both parallel calls stay running.""" started.append(True) if len(started) == 2: entered.set() await release.wait() return "Long operation completed" model = ParallelTaskModel(observed=[]) # The real repair layer would otherwise hide a lost result behind a placeholder. graph = create_agent(model, tools=[task], middleware=[RunInputAdmissionMiddleware(), DanglingToolCallMiddleware(), 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: await asyncio.wait_for(entered.wait(), 5) await inputs.submit("update", "Handle this first", "steer") async with asyncio.timeout(5): while len(model.observed) < 2: await asyncio.sleep(0.01) pending = [message for message in model.observed[1] if isinstance(message, ToolMessage)] assert [message.tool_call_id for message in pending] == ["first", "second"] assert all(message.additional_kwargs.get("subagent_pending") for message in pending) checkpoint = await graph.aget_state(config) assert [message.tool_call_id for message in checkpoint.values["messages"] if isinstance(message, ToolMessage)] == ["first", "second"] finally: release.set() await asyncio.wait_for(asyncio.gather(run, return_exceptions=True), 5) 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)