"""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() class ChildAndSiblingModel(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": "child"}, {"name": "sibling", "args": {}, "id": "slow"}]) else: message = AIMessage(content="Work completed") return ChatResult(generations=[ChatGeneration(message=message)]) @pytest.mark.asyncio async def test_steer_withdrawn_before_delivery_releases_the_children_it_paused(): 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, withdrawn, sibling_release = asyncio.Event(), asyncio.Event(), asyncio.Event() @tool async def task() -> str: """Controlled child; a real one parks at this boundary before its next model call.""" entered.set() await withdrawn.wait() await control.child("child").boundary() return "Child result" @tool async def sibling() -> str: """Keeps the tool step open, so the steer is withdrawn before delivery.""" await sibling_release.wait() return "Sibling result" model = ChildAndSiblingModel(observed=[]) graph = create_agent(model, tools=[task, sibling], 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(entered.wait(), 5) await inputs.submit("update", "Handle this first", "steer") gate = control.child("child") async with asyncio.timeout(5): while gate.resume.is_set(): await asyncio.sleep(0.01) await inputs.cancel("update") withdrawn.set() sibling_release.set() result = await asyncio.wait_for(run, 5) assert any(message.name == "subagent-result" and "Child result" in message.content for message in result["messages"]) assert not any(message.additional_kwargs.get("run_input_id") for message in result["messages"]) assert (await inputs.read())["items"][0]["status"] == "cancelled" finally: withdrawn.set() sibling_release.set() if not run.done(): run.cancel() await asyncio.gather(run, return_exceptions=True) await control.close() @pytest.mark.asyncio async def test_children_stay_paused_while_a_delivered_steer_is_being_handled(): 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) gate = control.child("child") gate.resume.clear() assert (await control.deliver({}))["applied_run_inputs"]["ids"] == ["update"] # Later model turns of the same steer have nothing to deliver, yet the # user's decision on the paused children is still outstanding. assert await control.deliver({}) is None assert not gate.resume.is_set() await control.close() @pytest.mark.asyncio async def test_steer_withdrawn_during_a_steering_turn_releases_the_child_it_paused(): store = MemoryRunStore() await store.put("run", thread_id="thread", user_id="owner", status="running", owner_worker_id="worker") inputs = RunInputs(store, "run", "owner", "worker") # Delivered at the first model boundary, before any child exists, so the # child paused below is outside this steer's interrupted snapshot. await inputs.submit("first", "Delegate the new work", "steer") saver = InMemorySaver() control = RunInputControl(inputs, agent_runs=SimpleNamespace(), thread_id="thread", checkpointer=saver) entered, withdrawn, sibling_release = asyncio.Event(), asyncio.Event(), asyncio.Event() @tool async def task() -> str: """Controlled child; a real one parks at this boundary before its next model call.""" entered.set() await withdrawn.wait() await control.child("child").boundary() return "Child result" @tool async def sibling() -> str: """Keeps the tool step open, so the second steer is withdrawn before delivery.""" await sibling_release.wait() return "Sibling result" model = ChildAndSiblingModel(observed=[]) graph = create_agent(model, tools=[task, sibling], 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(entered.wait(), 5) assert control.steering_active and control.interrupted_children == set() await inputs.submit("second", "Handle this first", "steer") gate = control.child("child") async with asyncio.timeout(5): while gate.resume.is_set(): await asyncio.sleep(0.01) await inputs.cancel("second") withdrawn.set() sibling_release.set() result = await asyncio.wait_for(run, 5) assert any(message.name == "subagent-result" and "Child result" in message.content for message in result["messages"]) assert {item["id"]: item["status"] for item in (await inputs.read())["items"]}["second"] == "cancelled" finally: withdrawn.set() sibling_release.set() if not run.done(): run.cancel() await asyncio.gather(run, return_exceptions=True) await control.close() class ReplacementReusingCallIdModel(InputModel): def _generate(self, messages, stop=None, run_manager=None, **kwargs): self.observed.append(messages) calls = { 1: [{"name": "task", "args": {}, "id": "call_0"}], 2: [{"name": "manage_current_work", "args": {"action": "replace"}, "id": "replace"}], 3: [{"name": "task", "args": {}, "id": "call_0"}, {"name": "sibling", "args": {}, "id": "slow"}], }.get(len(self.observed), []) return ChatResult(generations=[ChatGeneration(message=AIMessage(content="" if calls else "Work completed", tool_calls=calls))]) @pytest.mark.asyncio async def test_steer_withdrawn_from_a_replacement_reusing_a_call_id_releases_it(): 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() old_finished, entered, replacement_entered, withdrawn, sibling_release = (asyncio.Event() for _ in range(5)) async def authorize_control(**kwargs): old_finished.set() async with asyncio.timeout(5): while not old_gate.closed: await asyncio.sleep(0.01) control = RunInputControl(inputs, agent_runs=SimpleNamespace(authorize_control=authorize_control), thread_id="thread", checkpointer=saver) started = [] @tool async def task() -> str: """The first call is replaced; the second reuses its provider call ID.""" started.append(True) if len(started) == 1: entered.set() await old_finished.wait() return "Old result" replacement_entered.set() await withdrawn.wait() await control.child("call_0").boundary() return "Replacement result" @tool async def sibling() -> str: """Keeps the tool step open, so the second steer is withdrawn before delivery.""" await sibling_release.wait() return "Sibling result" model = ReplacementReusingCallIdModel(observed=[]) middleware = [RunInputAdmissionMiddleware(), DanglingToolCallMiddleware(), RunInputMiddleware()] graph = create_agent(model, tools=[task, sibling], middleware=middleware, 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(entered.wait(), 5) old_gate = control.child("call_0") await inputs.submit("first", "Do this instead", "steer") await asyncio.wait_for(replacement_entered.wait(), 5) gate = control.child("call_0") assert gate is not old_gate and control.abandoned await inputs.submit("second", "Handle this first", "steer") async with asyncio.timeout(5): while gate.resume.is_set(): await asyncio.sleep(0.01) await inputs.cancel("second") withdrawn.set() sibling_release.set() result = await asyncio.wait_for(run, 5) assert any(message.name == "subagent-result" and "Replacement result" in message.content for message in result["messages"]) assert {item["id"]: item["status"] for item in (await inputs.read())["items"]}["second"] == "cancelled" finally: old_finished.set() withdrawn.set() sibling_release.set() if not run.done(): run.cancel() await asyncio.gather(run, return_exceptions=True) await control.close() @pytest.mark.asyncio async def test_withdrawn_steer_keeps_children_held_by_the_delivered_steer_paused(): 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("first", "Handle this first", "steer") control = RunInputControl(inputs, agent_runs=None, thread_id="thread", checkpointer=None) held = control.child("held") held.resume.clear() await control.deliver({}) released = control.child("released") released.resume.clear() assert await control.deliver({}) is None assert released.resume.is_set() assert not held.resume.is_set() 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)