From 4a53a36e32a7245e4ac2d5e80115ce7e31f8bd6c Mon Sep 17 00:00:00 2001 From: Chau Vu <58875435+chauvuusvn@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:06:20 +0700 Subject: [PATCH] fix(extensions): roll back shared state mutations on budget check failure (#7902) * fix(extensions): roll back shared state mutations on budget check failure In deerflow_extension_agent_teams.shared_state, publish() and update() mutated team['public_state'] in-memory before invoking check_budget(). When check_budget() raised a budget-exceeded ValueError (e.g. 16,000 JSON bytes or 24,000 shared context bytes), the state mutation was not rolled back, causing: 1. Permanent state lockout on rejected publish (the oversized entry remained in state['entries']). 2. False optimistic concurrency conflicts on update retry (the entry revision was prematurely bumped and content altered). Wrap check_budget() in try-except in both functions: - On publish failure: pop the appended entry and decrement revision. - On update failure: restore original entry dict and decrement revision. - Added regression test covering both rollback scenarios. * test(extensions): request plugin fixture for isolated regression run * fix(extensions): add rollback identity guard and pin retry assertions --- backend/tests/test_agent_teams_plugin.py | 45 +++++++++++++++++++ .../shared_state.py | 17 ++++++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/backend/tests/test_agent_teams_plugin.py b/backend/tests/test_agent_teams_plugin.py index 96ca8b8c86..4870af6e00 100644 --- a/backend/tests/test_agent_teams_plugin.py +++ b/backend/tests/test_agent_teams_plugin.py @@ -2650,3 +2650,48 @@ async def test_history_cancelled_results_are_unavailable_before_and_after_admiss assert legacy["result"] == unavailable filtered = await service.search_results({"team_id": team["id"], "query": notice}, caller) assert filtered["items"] == [] + + +def test_shared_state_budget_rejection_rolls_back_in_memory_mutations(plugin): + from deerflow_extension_agent_teams import shared_state + + team = { + "id": "team_test", + "goal": "Goal", + "members": [{"id": "m1", "name": "Agent 1", "agent": "a1", "thread_id": "t1"}], + "jobs": [{"id": "j1", "thread_id": "t1", "status": "running", "input": {}}], + } + caller = SimpleNamespace(thread_id="t1") + for i in range(12): + shared_state.publish(team, {"team_id": "team_test", "request_id": f"req_{i}", "kind": "objective", "title": f"T{i}", "content": "A" * 950}, caller) + assert len(team["public_state"]["entries"]) == 12 + assert team["public_state"]["revision"] == 12 + + with pytest.raises(ValueError, match="16,000 UTF-8 JSON bytes"): + shared_state.publish(team, {"team_id": "team_test", "request_id": "req_overflow", "kind": "objective", "title": "T_overflow", "content": "A" * 950}, caller) + assert len(team["public_state"]["entries"]) == 12 + assert team["public_state"]["revision"] == 12 + + # Subsequent publish on the same dict succeeds after the rejection + res_pub = shared_state.publish(team, {"team_id": "team_test", "request_id": "req_small", "kind": "objective", "title": "T_small", "content": "small"}, caller) + assert res_pub["revision"] == 13 + assert len(team["public_state"]["entries"]) == 13 + + entry_0 = team["public_state"]["entries"][0] + rev_before = entry_0["revision"] + title_before = entry_0["title"] + team["goal"] = "G" * 30000 + update_payload = {"team_id": "team_test", "entry_id": entry_0["id"], "expected_revision": rev_before, "title": "New", "content": "New", "status": "active"} + with pytest.raises(ValueError, match="shared context budget"): + shared_state.update(team, update_payload, caller) + assert entry_0["title"] == title_before + assert entry_0["revision"] == rev_before + assert team["public_state"]["revision"] == 13 + + # Retrying the exact same update succeeds after resetting team["goal"] + team["goal"] = "Goal" + res_update = shared_state.update(team, update_payload, caller) + assert res_update["updated"] is True + assert res_update["conflict"] is False + assert entry_0["title"] == "New" + assert entry_0["revision"] == rev_before + 1 diff --git a/examples/deerflow-extension-agent-teams/deerflow_extension_agent_teams/shared_state.py b/examples/deerflow-extension-agent-teams/deerflow_extension_agent_teams/shared_state.py index 634a2305e8..64dc8075d7 100644 --- a/examples/deerflow-extension-agent-teams/deerflow_extension_agent_teams/shared_state.py +++ b/examples/deerflow-extension-agent-teams/deerflow_extension_agent_teams/shared_state.py @@ -116,7 +116,13 @@ def publish(team, payload, context): } state["entries"].append(entry) state["revision"] += 1 - check_budget(team) + try: + check_budget(team) + except BaseException: + if state["entries"] and state["entries"][-1] is entry: + state["entries"].pop() + state["revision"] -= 1 + raise return {"team_id": team["id"], "revision": state["revision"], "entry": {key: entry[key] for key in PUBLIC_FIELDS}} @@ -136,7 +142,14 @@ def update(team, payload, context): if entry["_update"] != stamp: if entry["revision"] != expected: return {"team_id": team["id"], "revision": state["revision"], "updated": False, "conflict": True, "entry": {key: entry[key] for key in PUBLIC_FIELDS}} + original_entry = dict(entry) entry.update(title=title, content=content, status=status, revision=expected + 1, updated_by=member_id, updated_at=datetime.now(UTC).isoformat(), source=source, _update=stamp) state["revision"] += 1 - check_budget(team) + try: + check_budget(team) + except BaseException: + entry.clear() + entry.update(original_entry) + state["revision"] -= 1 + raise return {"team_id": team["id"], "revision": state["revision"], "updated": True, "conflict": False, "entry": {key: entry[key] for key in PUBLIC_FIELDS}}