mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(extensions): isolate team workers and rotate admissions (#6754)
* fix(extensions): isolate team workers and rotate admissions * fix(extensions): index and batch pending team snapshots * fix(agent-teams): retain unfinished owner rotation rounds Carry the remaining owner queue across planning passes, including exhausted owners and newcomers. Read legacy rows completely before updating their pending projection and distinguish dispatch reservations from sorted entries. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
d18d2937e7
commit
8cec469738
@ -319,15 +319,18 @@ async def test_active_budget_retains_real_and_ambiguous_admissions(plugin, monke
|
||||
raise TimeoutError("admitted but acknowledgement lost")
|
||||
|
||||
monkeypatch.setattr(runs, "start", uncertain)
|
||||
first_admissions = None
|
||||
for _ in range(2):
|
||||
await service.tick()
|
||||
assert len(runs.starts) == 8
|
||||
assert jobs[-1]["id"] not in runs.runs
|
||||
for job in jobs[:8]:
|
||||
if first_admissions is None:
|
||||
first_admissions = set(runs.runs)
|
||||
assert set(runs.runs) == first_admissions
|
||||
for job in jobs:
|
||||
stored = await service.db("get", job["owner"], job["team_id"])
|
||||
assert all(item["input"] is not None for item in stored["jobs"])
|
||||
healthy = await service.db("get", "bob", jobs[-1]["team_id"])
|
||||
assert healthy["jobs"][0]["status"] == "queued" and healthy["jobs"][0]["input"] is None
|
||||
saved = next(item for item in stored["jobs"] if item["id"] == job["id"])
|
||||
assert (saved["input"] is not None) == (job["id"] in first_admissions)
|
||||
assert jobs[-1]["id"] in first_admissions # Bob receives a fair admission.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -1048,3 +1051,619 @@ async def test_malformed_unicode_response_uses_the_friendly_validation_error(plu
|
||||
job = (await service.db("get", "alice", team["id"]))["jobs"][0]
|
||||
assert job["status"] == "waiting_input"
|
||||
assert job["resume"] is None
|
||||
|
||||
|
||||
async def scheduling_team(actions, runs, owner, key, members=2):
|
||||
return await actions["create"]({"request_id": key, "name": key, "goal": "Check scheduling", "members": [{"name": f"Member {i}", "agent": "researcher"} for i in range(members)]}, context(runs, owner))
|
||||
|
||||
|
||||
async def scheduling_send(actions, runs, owner, team, member, key):
|
||||
return await actions["send"]({"team_id": team["id"], "member_id": member["id"], "text": "Check", "request_id": key + member["id"]}, context(runs, owner))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["create", "send", "connect", "cancel", "resume", "delete"])
|
||||
async def test_slow_team_does_not_block_another_teams_control(plugin, monkeypatch, operation):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
slow = await scheduling_team(actions, runs, "alice", "Slow")
|
||||
healthy = await scheduling_team(actions, runs, "bob", "Healthy")
|
||||
slow_job = await scheduling_send(actions, runs, "alice", slow, slow["members"][0], "slow")
|
||||
healthy_job = None
|
||||
if operation in ("cancel", "resume"):
|
||||
healthy_job = await scheduling_send(actions, runs, "bob", healthy, healthy["members"][0], "healthy")
|
||||
await service.tick()
|
||||
if operation == "resume":
|
||||
runs.finish(healthy_job["id"], interrupted=True)
|
||||
await service.tick()
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
get = runs.get
|
||||
|
||||
async def stalled(**kwargs):
|
||||
if kwargs["thread_id"] == slow["members"][0]["thread_id"]:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return await get(**kwargs)
|
||||
|
||||
monkeypatch.setattr(runs, "get", stalled)
|
||||
tick = asyncio.create_task(service.tick())
|
||||
control = None
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 2)
|
||||
if operation == "create":
|
||||
control = asyncio.create_task(scheduling_team(actions, runs, "bob", "New"))
|
||||
elif operation == "send":
|
||||
control = asyncio.create_task(scheduling_send(actions, runs, "bob", healthy, healthy["members"][1], "second"))
|
||||
else:
|
||||
payload = {"team_id": healthy["id"]}
|
||||
if healthy_job:
|
||||
payload["job_id"] = healthy_job["id"]
|
||||
if operation == "resume":
|
||||
payload.update(response={"approved": True}, request_id="answer")
|
||||
control = asyncio.create_task(actions[operation](payload, context(runs, "bob")))
|
||||
await asyncio.wait_for(asyncio.shield(control), 2)
|
||||
assert not release.is_set()
|
||||
assert runs.runs[slow_job["id"]].status == "running"
|
||||
finally:
|
||||
release.set()
|
||||
await tick
|
||||
if control:
|
||||
await control
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_admissions_are_round_robin_by_owner_not_team_count(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
owners = {}
|
||||
for owner, count, members in (("alice", 4, 2), ("bob", 1, 8)):
|
||||
for index in range(count):
|
||||
team = await scheduling_team(actions, runs, owner, f"{owner}-{index}", members)
|
||||
for member in team["members"]:
|
||||
owners[member["thread_id"]] = owner
|
||||
await scheduling_send(actions, runs, owner, team, member, "first")
|
||||
await service.tick()
|
||||
started = [owners[thread] for thread, _ in runs.starts]
|
||||
assert len(started) == 8
|
||||
assert started.count("alice") == started.count("bob") == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_owner_rotation_survives_one_slot_at_a_time(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
owners = {}
|
||||
for owner in ("alice", "bob"):
|
||||
team = await scheduling_team(actions, runs, owner, owner, 8)
|
||||
for member in team["members"]:
|
||||
owners[member["thread_id"]] = owner
|
||||
await scheduling_send(actions, runs, owner, team, member, "first")
|
||||
await service.tick()
|
||||
admitted = []
|
||||
for _ in range(4):
|
||||
first = next(run for run in runs.runs.values() if run.status == "running")
|
||||
runs.finish(first.run_id)
|
||||
before = len(runs.starts)
|
||||
await service.tick()
|
||||
assert len(runs.starts) == before + 1
|
||||
admitted.append(owners[runs.starts[-1][0]])
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
assert admitted == ["alice", "bob", "alice", "bob"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_owner_rotation_finishes_the_round_after_single_job_owners_leave_pending(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
owners = {}
|
||||
for owner, members in [("alice", 8), *((f"owner-{i}", 2) for i in range(1, 17))]:
|
||||
team = await scheduling_team(actions, runs, owner, owner, members)
|
||||
for member in team["members"][: 8 if owner == "alice" else 1]:
|
||||
owners[member["thread_id"]] = owner
|
||||
await scheduling_send(actions, runs, owner, team, member, "first")
|
||||
await service.tick()
|
||||
assert {owners[thread] for thread, _ in runs.starts} == {"alice", *(f"owner-{i}" for i in range(1, 8))}
|
||||
admitted = []
|
||||
for _ in range(10):
|
||||
first = next(run for run in runs.runs.values() if run.status == "running")
|
||||
runs.finish(first.run_id)
|
||||
before = len(runs.starts)
|
||||
await service.tick()
|
||||
assert len(runs.starts) == before + 1
|
||||
admitted.append(owners[runs.starts[-1][0]])
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
assert admitted == [*(f"owner-{i}" for i in range(8, 17)), "alice"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("cancel_waiting_owner", [False, True])
|
||||
async def test_owner_rotation_appends_new_owners_and_drops_cancelled_queues(plugin, cancel_waiting_owner):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
owners, waiting = {}, None
|
||||
for owner, members in [("alice", 8), *((f"owner-{i}", 2) for i in range(1, 9))]:
|
||||
team = await scheduling_team(actions, runs, owner, owner, members)
|
||||
for member in team["members"][: 8 if owner == "alice" else 1]:
|
||||
owners[member["thread_id"]] = owner
|
||||
sent = await scheduling_send(actions, runs, owner, team, member, "first")
|
||||
if owner == "owner-8":
|
||||
waiting = (team, sent)
|
||||
await service.tick()
|
||||
late = await scheduling_team(actions, runs, "late", "late")
|
||||
owners[late["members"][0]["thread_id"]] = "late"
|
||||
await scheduling_send(actions, runs, "late", late, late["members"][0], "first")
|
||||
if cancel_waiting_owner:
|
||||
team, sent = waiting
|
||||
await actions["cancel"]({"team_id": team["id"], "job_id": sent["id"]}, context(runs, "owner-8"))
|
||||
await actions["delete"]({"team_id": team["id"]}, context(runs, "owner-8"))
|
||||
admitted = []
|
||||
for _ in range(4):
|
||||
first = next(run for run in runs.runs.values() if run.status == "running")
|
||||
runs.finish(first.run_id)
|
||||
before = len(runs.starts)
|
||||
await service.tick()
|
||||
assert len(runs.starts) == before + 1
|
||||
admitted.append(owners[runs.starts[-1][0]])
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
assert admitted == (["alice", "late", "alice", "alice"] if cancel_waiting_owner else ["owner-8", "alice", "late", "alice"])
|
||||
# With only Alice's backlog left, spare slots must remain usable.
|
||||
for run in list(runs.runs.values()):
|
||||
if run.status == "running":
|
||||
runs.finish(run.run_id)
|
||||
await service.tick()
|
||||
assert len(runs.starts) == (16 if cancel_waiting_owner else 17)
|
||||
for run in list(runs.runs.values()):
|
||||
if run.status == "running":
|
||||
runs.finish(run.run_id)
|
||||
await service.tick()
|
||||
assert not any(run.status == "running" for run in runs.runs.values())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_loop_keeps_cancelling_other_teams_while_one_call_stalls(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
slow = await scheduling_team(actions, runs, "alice", "Slow")
|
||||
healthy = await scheduling_team(actions, runs, "bob", "Healthy")
|
||||
await scheduling_send(actions, runs, "alice", slow, slow["members"][0], "slow")
|
||||
sent = await scheduling_send(actions, runs, "bob", healthy, healthy["members"][0], "healthy")
|
||||
entered, healthy_running, cancelled, release = (asyncio.Event() for _ in range(4))
|
||||
get, cancel = runs.get, runs.cancel
|
||||
|
||||
async def stalled(**kwargs):
|
||||
if kwargs["thread_id"] == slow["members"][0]["thread_id"]:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
else:
|
||||
healthy_running.set()
|
||||
return await get(**kwargs)
|
||||
|
||||
async def record_cancel(**kwargs):
|
||||
await cancel(**kwargs)
|
||||
cancelled.set()
|
||||
|
||||
monkeypatch.setattr(runs, "get", stalled)
|
||||
monkeypatch.setattr(runs, "cancel", record_cancel)
|
||||
await service.start(None)
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 3)
|
||||
await asyncio.wait_for(healthy_running.wait(), 3)
|
||||
await asyncio.wait_for(actions["cancel"]({"team_id": healthy["id"], "job_id": sent["id"]}, context(runs, "bob")), 2)
|
||||
await asyncio.wait_for(cancelled.wait(), 3)
|
||||
assert not release.is_set()
|
||||
# Shutdown cancels and drains gated workers; it never needs the host gate.
|
||||
await asyncio.wait_for(service.stop(), 3)
|
||||
assert not service.workers and not service.in_flight
|
||||
assert not service.available and service.process_lock is None
|
||||
finally:
|
||||
release.set()
|
||||
await service.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overlapping_ticks_keep_pre_freeze_reservations_within_eight(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
for index in range(9):
|
||||
owner = f"owner-{index}"
|
||||
team = await scheduling_team(actions, runs, owner, owner)
|
||||
await scheduling_send(actions, runs, owner, team, team["members"][0], "first")
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
get_state = runs.get_state
|
||||
|
||||
async def gated(**kwargs):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return await get_state(**kwargs)
|
||||
|
||||
monkeypatch.setattr(runs, "get_state", gated)
|
||||
first = asyncio.create_task(service.tick())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 3)
|
||||
await asyncio.wait_for(asyncio.gather(*(service.tick() for _ in range(3))), 2)
|
||||
assert not runs.starts # No input has frozen; reservations still own slots.
|
||||
finally:
|
||||
release.set()
|
||||
await first
|
||||
assert len(runs.starts) == 8
|
||||
assert not service.in_flight and not service.workers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_more_than_eight_owners_rotate_into_the_next_free_slot(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
jobs = []
|
||||
for index in range(9):
|
||||
owner = f"owner-{index}"
|
||||
team = await scheduling_team(actions, runs, owner, owner)
|
||||
jobs.append(await scheduling_send(actions, runs, owner, team, team["members"][0], "first"))
|
||||
await service.tick()
|
||||
assert set(runs.runs) == {job["id"] for job in jobs[:8]}
|
||||
runs.finish(jobs[0]["id"])
|
||||
await service.tick()
|
||||
assert jobs[-1]["id"] in runs.runs
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_resumes_wait_for_slots_and_cancel_without_admitting_an_answer(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
parked = await scheduling_team(actions, runs, "alice", "Parked")
|
||||
sent = await scheduling_send(actions, runs, "alice", parked, parked["members"][0], "park")
|
||||
await service.tick()
|
||||
runs.finish(sent["id"], interrupted=True)
|
||||
await service.tick()
|
||||
busy = await scheduling_team(actions, runs, "bob", "Busy", 8)
|
||||
for member in busy["members"]:
|
||||
await scheduling_send(actions, runs, "bob", busy, member, "busy")
|
||||
await service.tick()
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
payload = {"team_id": parked["id"], "job_id": sent["id"], "response": {"approved": True}, "request_id": "answer"}
|
||||
await actions["resume"](payload, context(runs, "alice"))
|
||||
await service.tick()
|
||||
assert len(runs.runs) == 9 # The finished interrupted run plus eight busy runs.
|
||||
assert (await service.db("get", "alice", parked["id"]))["jobs"][0]["status"] == "resuming"
|
||||
await actions["cancel"]({"team_id": parked["id"], "job_id": sent["id"]}, context(runs, "alice"))
|
||||
await service.tick()
|
||||
assert len(runs.runs) == 9
|
||||
assert (await service.db("get", "alice", parked["id"]))["jobs"][0]["status"] == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frozen_slots_survive_disconnection_and_service_restart(plugin, monkeypatch):
|
||||
from deerflow_extension_agent_teams.service import Teams
|
||||
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
team = await scheduling_team(actions, runs, "alice", "Busy", 8)
|
||||
for member in team["members"]:
|
||||
await scheduling_send(actions, runs, "alice", team, member, "busy")
|
||||
await service.tick()
|
||||
runs.denied = True
|
||||
await service.tick()
|
||||
assert not service.handles and len(runs.starts) == 8
|
||||
replacement = Teams(service.store.path)
|
||||
healthy_runs = Runs()
|
||||
# The real start loads the durable slot owners even before they reconnect.
|
||||
await replacement.start(None)
|
||||
try:
|
||||
healthy = await scheduling_team({"create": replacement.create}, healthy_runs, "bob", "Healthy")
|
||||
await scheduling_send({"send": replacement.send}, healthy_runs, "bob", healthy, healthy["members"][0], "new")
|
||||
await replacement.tick()
|
||||
assert not healthy_runs.starts
|
||||
runs.denied = False
|
||||
await replacement.connect({"team_id": team["id"]}, context(runs))
|
||||
for run in list(runs.runs.values()):
|
||||
runs.finish(run.run_id)
|
||||
await replacement.tick()
|
||||
await replacement.tick()
|
||||
assert len(healthy_runs.starts) == 1
|
||||
finally:
|
||||
await replacement.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_receipts_for_a_shared_source_are_serialized_across_teams(plugin):
|
||||
from deerflow_extension_agent_teams.service import identity
|
||||
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
await runs.create_thread(assistant_id="lead_agent", thread_id="shared", metadata={})
|
||||
requests = []
|
||||
for owner in ("alice", "bob"):
|
||||
team = await scheduling_team(actions, runs, owner, owner)
|
||||
sent = await service.db("change", owner, team["id"], lambda t: service.enqueue(t, member_id=t["members"][0]["id"], content="Check", request_id="shared-request", source="shared"))
|
||||
requests.append((owner, team, sent))
|
||||
await service.tick()
|
||||
for _, _, job in requests:
|
||||
runs.finish(job["id"])
|
||||
await service.tick()
|
||||
await service.tick()
|
||||
receipts = []
|
||||
for owner, team, request in requests:
|
||||
stored = await service.db("get", owner, team["id"])
|
||||
receipt = next(j for j in stored["jobs"] if j["id"] == identity(request["id"], "receipt"))
|
||||
receipts.append((receipt["created_at"], owner, team, receipt))
|
||||
receipts.sort(key=lambda entry: entry[0])
|
||||
first, later = receipts
|
||||
assert first[3]["id"] in runs.runs and later[3]["id"] not in runs.runs
|
||||
assert later[3]["input"] is None # No competing admission/checkpoint on the source.
|
||||
runs.finish(first[3]["id"])
|
||||
await service.tick()
|
||||
await service.tick()
|
||||
assert later[3]["id"] in runs.runs
|
||||
assert sum(thread == "shared" for thread, _ in runs.starts) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_snapshot_failure_releases_the_process_lock(plugin, monkeypatch):
|
||||
from deerflow_extension_agent_teams.service import Teams
|
||||
|
||||
_, _, service = plugin
|
||||
db = service.db
|
||||
|
||||
async def fail_snapshot(method, *args):
|
||||
if method == "scheduling_snapshots":
|
||||
raise OSError("snapshot failed")
|
||||
return await db(method, *args)
|
||||
|
||||
monkeypatch.setattr(service, "db", fail_snapshot)
|
||||
with pytest.raises(OSError, match="snapshot failed"):
|
||||
await service.start(None)
|
||||
assert service.process_lock is None and not service.task
|
||||
second = Teams(service.store.path)
|
||||
await second.start(None)
|
||||
await second.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_unfinished_thread_occupant_does_not_depend_on_timestamps(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
team = await create(actions, runs)
|
||||
member = team["members"][0]
|
||||
first = await scheduling_send(actions, runs, "alice", team, member, "first")
|
||||
later = await scheduling_send(actions, runs, "alice", team, member, "later")
|
||||
await service.db("change", "alice", team["id"], lambda t: t["jobs"][1].pop("created_at"))
|
||||
await service.tick()
|
||||
assert first["id"] in runs.runs and later["id"] not in runs.runs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_and_delete_before_worker_admission_release_only_that_reservation(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
jobs = []
|
||||
for index in range(9):
|
||||
owner = f"owner-{index}"
|
||||
team = await scheduling_team(actions, runs, owner, owner)
|
||||
sent = await scheduling_send(actions, runs, owner, team, team["members"][0], "first")
|
||||
jobs.append((owner, team, sent))
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
work_team = service.work_team
|
||||
ready = 0
|
||||
|
||||
async def gated(*args):
|
||||
nonlocal ready
|
||||
ready += 1
|
||||
if ready == 8:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
await work_team(*args)
|
||||
|
||||
monkeypatch.setattr(service, "work_team", gated)
|
||||
tick = asyncio.create_task(service.tick())
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), 3)
|
||||
owner, team, sent = jobs[0]
|
||||
await actions["cancel"]({"team_id": team["id"], "job_id": sent["id"]}, context(runs, owner))
|
||||
await actions["delete"]({"team_id": team["id"]}, context(runs, owner))
|
||||
assert not runs.starts
|
||||
finally:
|
||||
release.set()
|
||||
await tick
|
||||
assert sent["id"] not in runs.runs
|
||||
assert len(runs.starts) == 8 and jobs[-1][2]["id"] in runs.runs
|
||||
assert not service.workers and not service.in_flight
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_with_lost_ack_keeps_its_slot_before_any_new_owner_admission(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
parked = await scheduling_team(actions, runs, "alice", "Parked")
|
||||
sent = await scheduling_send(actions, runs, "alice", parked, parked["members"][0], "park")
|
||||
await service.tick()
|
||||
runs.finish(sent["id"], interrupted=True)
|
||||
await service.tick()
|
||||
busy = await scheduling_team(actions, runs, "bob", "Busy", 8)
|
||||
for member in busy["members"][:7]:
|
||||
await scheduling_send(actions, runs, "bob", busy, member, "busy")
|
||||
await service.tick()
|
||||
start = runs.resume
|
||||
|
||||
async def lost_ack(**kwargs):
|
||||
await start(**kwargs)
|
||||
raise TimeoutError("resume admitted; acknowledgement lost")
|
||||
|
||||
monkeypatch.setattr(runs, "resume", lost_ack)
|
||||
await actions["resume"]({"team_id": parked["id"], "job_id": sent["id"], "response": {"approved": True}, "request_id": "answer"}, context(runs, "alice"))
|
||||
await service.tick()
|
||||
stored = (await service.db("get", "alice", parked["id"]))["jobs"][0]
|
||||
assert stored["resume"]["started"]
|
||||
later = await scheduling_team(actions, runs, "carol", "Later")
|
||||
queued = await scheduling_send(actions, runs, "carol", later, later["members"][0], "new")
|
||||
await service.tick()
|
||||
assert queued["id"] not in runs.runs
|
||||
assert sum(run.status == "running" for run in runs.runs.values()) == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduler_does_not_read_or_decode_completed_history_per_team(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
for index in range(100):
|
||||
owner = f"archive-{index // 20}"
|
||||
team = await scheduling_team(actions, runs, owner, f"History {index}")
|
||||
await scheduling_send(actions, runs, owner, team, team["members"][0], "completed")
|
||||
await service.db("change", owner, team["id"], lambda t: t["jobs"][0].update(status="completed", input={"messages": []}))
|
||||
connections, decoded = [], []
|
||||
connect, decode = service.store.connection, json.loads
|
||||
|
||||
def record_connection():
|
||||
db = connect()
|
||||
connections.append(db)
|
||||
return db
|
||||
|
||||
def record_decode(*args, **kwargs):
|
||||
decoded.append(1)
|
||||
return decode(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(service.store, "connection", record_connection)
|
||||
monkeypatch.setattr(json, "loads", record_decode)
|
||||
await service.tick(wait=False)
|
||||
assert len(connections) == 1
|
||||
assert not decoded and not runs.starts
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleted_team_does_not_leave_a_repeated_missing_key_lookup(plugin, monkeypatch):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
team = await create(actions, runs)
|
||||
await service.db("delete", "alice", team["id"])
|
||||
methods = []
|
||||
db = service.db
|
||||
|
||||
async def record(method, *args):
|
||||
methods.append(method)
|
||||
return await db(method, *args)
|
||||
|
||||
monkeypatch.setattr(service, "db", record)
|
||||
await service.tick(wait=False)
|
||||
await service.tick(wait=False)
|
||||
assert "get" not in methods
|
||||
assert not service.workers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_projection_tracks_queue_wait_finish_and_requeue(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
team = await create(actions, runs)
|
||||
key = ("alice", team["id"])
|
||||
assert key not in await service.db("scheduling_snapshots")
|
||||
sent = await scheduling_send(actions, runs, "alice", team, team["members"][0], "first")
|
||||
assert key in await service.db("scheduling_snapshots")
|
||||
await service.tick()
|
||||
runs.finish(sent["id"], interrupted=True)
|
||||
await service.tick()
|
||||
waiting = (await service.db("scheduling_snapshots"))[key]["jobs"][0]
|
||||
assert waiting["status"] == "waiting_input" and waiting["input"] is not None
|
||||
assert not service.charged(waiting) # Keep the thread head, release the slot.
|
||||
await actions["cancel"]({"team_id": team["id"], "job_id": sent["id"]}, context(runs))
|
||||
await service.tick()
|
||||
assert key not in await service.db("scheduling_snapshots")
|
||||
assert (await actions["get"]({"team_id": team["id"]}, context(runs)))["connected"]
|
||||
runs.threads[team["members"][0]["thread_id"]].update(next=[], interrupts=[])
|
||||
await scheduling_send(actions, runs, "alice", team, team["members"][0], "second")
|
||||
assert key in await service.db("scheduling_snapshots")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_projection_keeps_a_receipt_created_by_terminal_result(plugin):
|
||||
_, actions, service = plugin
|
||||
runs = Runs()
|
||||
team = await create(actions, runs)
|
||||
sender, recipient = team["members"]
|
||||
sent = await actions["send"]({"team_id": team["id"], "member_id": recipient["id"], "text": "Check", "request_id": "peer"}, context(runs, thread=sender["thread_id"]))
|
||||
await service.tick()
|
||||
runs.finish(sent["id"])
|
||||
await service.tick()
|
||||
jobs = (await service.db("scheduling_snapshots"))[("alice", team["id"])]["jobs"]
|
||||
assert jobs[0]["status"] == "completed" and jobs[-1]["kind"] == "receipt"
|
||||
assert jobs[-1]["status"] == "queued"
|
||||
|
||||
|
||||
def test_pending_projection_backfills_legacy_rows_and_uses_an_index(plugin, tmp_path):
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
from deerflow_extension_agent_teams.store import Store
|
||||
|
||||
path = tmp_path / "legacy.sqlite"
|
||||
states = ["completed", "queued", "waiting_input", "resuming", "unknown-legacy-state"]
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("CREATE TABLE teams (id TEXT PRIMARY KEY, owner TEXT NOT NULL, data TEXT NOT NULL)")
|
||||
for index, status in enumerate(states):
|
||||
job = {"status": status, "input": {"messages": []}, "resume": {"started": False}}
|
||||
db.execute("INSERT INTO teams VALUES (?,?,?)", (str(index), f"owner-{index}", json.dumps({"jobs": [job]})))
|
||||
store = Store(path)
|
||||
snapshots = store.scheduling_snapshots()
|
||||
assert set(snapshots) == {(f"owner-{i}", str(i)) for i in range(1, len(states))}
|
||||
with closing(store.connection()) as db:
|
||||
plan = db.execute("EXPLAIN QUERY PLAN SELECT owner, id, data FROM teams WHERE has_pending_jobs=1 ORDER BY rowid").fetchall()
|
||||
assert any("teams_pending_jobs" in row[3] for row in plan)
|
||||
# Reopening must use the projection instead of re-decoding legacy history.
|
||||
again = Store(path)
|
||||
assert again.scheduling_snapshots() == snapshots
|
||||
|
||||
|
||||
def test_legacy_backfill_finishes_reading_before_updating_the_scanned_table(plugin, tmp_path, monkeypatch):
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
from deerflow_extension_agent_teams.store import Store
|
||||
|
||||
class BackfillCursor(sqlite3.Cursor):
|
||||
def fetchall(self):
|
||||
rows = super().fetchall()
|
||||
self.connection.scanning = False
|
||||
return rows
|
||||
|
||||
def __next__(self):
|
||||
try:
|
||||
return super().__next__()
|
||||
except StopIteration:
|
||||
self.connection.scanning = False
|
||||
raise
|
||||
|
||||
class BackfillConnection(sqlite3.Connection):
|
||||
scanning = False
|
||||
|
||||
def execute(self, sql, parameters=()):
|
||||
if sql.startswith("UPDATE teams"):
|
||||
assert not self.scanning, "Updating teams before its SELECT has finished"
|
||||
cursor = self.cursor(BackfillCursor).execute(sql, parameters)
|
||||
if sql == "SELECT id, data FROM teams":
|
||||
self.scanning = True
|
||||
return cursor
|
||||
|
||||
path = tmp_path / "legacy-scan.sqlite"
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("CREATE TABLE teams (id TEXT PRIMARY KEY, owner TEXT NOT NULL, data TEXT NOT NULL)")
|
||||
for index in range(100):
|
||||
db.execute("INSERT INTO teams VALUES (?,?,?)", (str(index), "alice", json.dumps({"jobs": [{"status": "completed" if index % 2 else "queued"}], "padding": "x" * 4000})))
|
||||
monkeypatch.setattr(Store, "connection", lambda store: sqlite3.connect(store.path, factory=BackfillConnection))
|
||||
assert set(Store(path).scheduling_snapshots()) == {("alice", str(i)) for i in range(0, 100, 2)}
|
||||
|
||||
|
||||
def test_failed_projection_backfill_rolls_back_schema_and_rows(plugin, tmp_path):
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
from deerflow_extension_agent_teams.store import Store
|
||||
|
||||
path = tmp_path / "invalid-legacy.sqlite"
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
db.execute("CREATE TABLE teams (id TEXT PRIMARY KEY, owner TEXT NOT NULL, data TEXT NOT NULL)")
|
||||
db.execute("INSERT INTO teams VALUES ('first','alice',?)", (json.dumps({"jobs": [{"status": "queued"}]}),))
|
||||
db.execute("INSERT INTO teams VALUES ('invalid','bob','invalid json')")
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
Store(path)
|
||||
with closing(sqlite3.connect(path)) as db, db:
|
||||
assert "has_pending_jobs" not in {row[1] for row in db.execute("PRAGMA table_info(teams)")}
|
||||
db.execute("UPDATE teams SET data=? WHERE id='invalid'", (json.dumps({"jobs": [{"status": "completed"}]}),))
|
||||
assert set(Store(path).scheduling_snapshots()) == {("alice", "first")}
|
||||
|
||||
@ -5,7 +5,9 @@ imports public `deerflow_extension_api` contracts and its declared LangChain
|
||||
dependency; do not import host `app.*` or `deerflow.*` internals. Browser code is a
|
||||
manifest-listed native module with no host React dependency.
|
||||
|
||||
Every storage operation is owner-scoped. Native mentions use Gateway's stamped
|
||||
Public reads and all team mutations are owner-scoped. The service-only snapshot
|
||||
spans owners for scheduling bookkeeping and is never exposed as a public read.
|
||||
Native mentions use Gateway's stamped
|
||||
runtime identity and delegated capability, never an owner from mention content.
|
||||
Thread sharing is not team sharing. Keep peer content on human-message channels.
|
||||
|
||||
@ -15,9 +17,25 @@ Unknown admission outcomes must retain ownership of the pending job. Never
|
||||
automatically approve interruptions, infer success from a terminal run alone, or
|
||||
serialize a host capability. Only an authenticated action or native mention
|
||||
rebinds after restart.
|
||||
Release a new job's temporary active-budget reservation only if its input remains
|
||||
unfrozen after advancement; conversation waits must not block unrelated teams.
|
||||
Keep frozen inputs charged even when admission acknowledgement is lost.
|
||||
Use team-scoped locks for every execution/control mutation, including native
|
||||
mentions; hold the scheduling lock only for snapshot/slot bookkeeping, never for
|
||||
host calls. Background ticks launch independent team workers and shutdown drains
|
||||
all workers. Preserve first-unfinished order within each thread, including source
|
||||
receipts shared across teams; legacy timestamps must not reorder a team's queue.
|
||||
Rotate new admission attempts by owner, carrying the cursor between passes.
|
||||
Reserve before dispatch and count in-flight reservations against eight slots.
|
||||
Terminal jobs and waiting_input jobs release their slots, even if the original
|
||||
input remains frozen; waiting_input still retains its thread's queue head.
|
||||
Unfrozen conversation waits also release their temporary reservation. Pending
|
||||
answers with resume.started=false wait for capacity. Other unfinished frozen
|
||||
jobs remain charged through disconnection/startup without a capability.
|
||||
Persist resume.started before attempting an answer,
|
||||
and keep uncertain admissions charged. Legacy resumes without the marker are
|
||||
conservatively charged. Cancelling an unstarted answer never admits it.
|
||||
Maintain the indexed has_pending_jobs projection in the same transaction as team
|
||||
data, including receipt creation. It tracks all unfinished jobs, not only charged
|
||||
jobs. Each planning pass uses one batch query; terminal history must not be
|
||||
parsed and missing keys must not remain in a separate scheduling registry.
|
||||
|
||||
Reconnect ensures every member thread before persisting readiness, serialized
|
||||
with deletion and task updates. Count only requests toward the 100-request cap;
|
||||
|
||||
@ -90,17 +90,34 @@ own teams. Sharing a host conversation does **not** share its team's private dat
|
||||
or raw model invocation. Capabilities are process-local, never stored in SQLite.
|
||||
Gateway rechecks the user, permission ceiling, current permissions and thread
|
||||
access on every call. PAT/internal callers without this capability cannot run teams.
|
||||
- Different member conversations can run concurrently; admissions are bounded to
|
||||
eight active plugin jobs shared by all owners in this process, without per-owner
|
||||
quotas or fair scheduling. A busy owner can keep other owners' jobs queued.
|
||||
The tick loop holds a shared lock while checking/admitting jobs, with a ten-second
|
||||
timeout per job advancement; slow host calls delay create, send, cancel, resume,
|
||||
delete and reconnect for all teams until the pass releases the lock. Per-owner
|
||||
fairness and finer-grained locking are outside this example's current scope.
|
||||
Each conversation is serialized, including result
|
||||
receipts. A busy host conversation stays queued without cancelling its run.
|
||||
Queued jobs blocked by an unanswered question or interrupt do not consume the
|
||||
active-job budget before admission; an unknown admission outcome still does.
|
||||
- Admissions are bounded to eight active plugin jobs across this database.
|
||||
New requests and pending answers rotate by owner, one admission attempt per
|
||||
owner per round. The unfinished round carries over between passes, including
|
||||
when a served owner's queue becomes empty. Unavailable queues drop out and new
|
||||
owners join the tail. Spare capacity can be used by any queued owner; existing
|
||||
runs are not preempted when another owner joins.
|
||||
- Each team has its own execution/control lock and background worker. Slow host
|
||||
calls in one team do not hold another team's create, send, cancel, resume,
|
||||
delete or reconnect operations, nor stop its background polling. Operations
|
||||
within the same team remain serialized, with a ten-second timeout per job
|
||||
advancement. Shutdown cancels and drains all workers before releasing the
|
||||
database service lock.
|
||||
- Each conversation is serialized, including result receipts targeting a source
|
||||
shared across teams. The first unfinished job retains the conversation until
|
||||
it finishes or is cancelled. A busy host conversation stays queued without
|
||||
cancelling its run. Unadmitted conversation waits release their temporary
|
||||
reservation. An admitted job in `waiting_input` also releases its slot even
|
||||
though its original input remains frozen; it still retains the conversation's
|
||||
queue head. Pending answers wait for capacity. Other unfinished frozen jobs,
|
||||
including unknown admission outcomes and resumes already attempted, keep their
|
||||
slots while disconnected and after service restart. Reconnect the owning team
|
||||
to reconcile them. Cancelling an unstarted answer does not admit it.
|
||||
- Each planning pass uses one indexed batch query for teams with unfinished jobs.
|
||||
Completed history is not decoded. The plugin database adds a pending-job flag
|
||||
and index, reads legacy rows completely before transactionally backfilling
|
||||
them at initialization, and updates the flag with each team mutation. The
|
||||
projection includes waiting-input jobs for conversation ordering even though
|
||||
those jobs do not occupy active slots.
|
||||
- Each owner can have 20 teams. A team accepts at most 100 requests, with separate
|
||||
space for up to 100 result receipts (200 total job entries); a handoff chain
|
||||
accepts at most 12 requests. Start a new user request for another chain, or a
|
||||
@ -152,7 +169,14 @@ real ToolNode dispatch, mention middleware in the Lead middleware pipeline,
|
||||
host-compatible mention labels at the UTF-16 boundary with stable routing IDs,
|
||||
serialization, interrupts, cancellation, restart/reconnect (including incomplete
|
||||
creation), request capacity independent of receipts, ambiguous admission, and
|
||||
lifecycle locking. Host-run transport is controlled for deterministic tests.
|
||||
lifecycle locking. Event-gated scheduling regressions also cover cross-team control
|
||||
isolation, continued background cancellation, owner rotation with single-slot
|
||||
release and mixed backlogs, exhausted/cancelled queues and new owners, overlapping
|
||||
ticks, pending answers, shared-source receipts, worker
|
||||
shutdown and durable slots across restart. Storage regressions cover indexed batch
|
||||
reads, terminal-history exclusion, complete legacy reads before backfill, rollback
|
||||
and projection updates when queuing, cancelling, finishing or creating receipts.
|
||||
Host-run transport is controlled for deterministic tests.
|
||||
Clarification tests execute the host's real tool and middleware in an Agent graph,
|
||||
including repeated questions, peer/mention/result-receipt paths, lost response
|
||||
acknowledgements and restart. Response regressions cover the 7,999/8,000/8,001-byte
|
||||
|
||||
@ -5,17 +5,18 @@ import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import uuid
|
||||
from collections import deque
|
||||
from collections.abc import Mapping
|
||||
from contextlib import suppress
|
||||
from datetime import UTC, datetime
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
from deerflow_extension_api import AgentRunError
|
||||
from deerflow_extension_api.agent_runs import AGENT_RUNS_CONTEXT_KEY
|
||||
|
||||
from .store import Store
|
||||
from .store import TERMINAL, Store
|
||||
|
||||
NAMESPACE = "community.agent-teams"
|
||||
TERMINAL = {"completed", "failed", "cancelled"}
|
||||
LIMIT = 100
|
||||
|
||||
|
||||
@ -76,7 +77,11 @@ class Teams:
|
||||
def __init__(self, path):
|
||||
self.store = Store(path)
|
||||
self.handles = {}
|
||||
self.lock = asyncio.Lock()
|
||||
self.locks = WeakValueDictionary()
|
||||
self.scheduling_lock = asyncio.Lock()
|
||||
self.workers = {}
|
||||
self.in_flight = set()
|
||||
self.owner_rotation = deque()
|
||||
self.task = None
|
||||
self.process_lock = None
|
||||
self.available = False
|
||||
@ -97,6 +102,12 @@ class Teams:
|
||||
return db
|
||||
|
||||
self.process_lock = await asyncio.to_thread(acquire)
|
||||
try:
|
||||
await self.db("scheduling_snapshots")
|
||||
except BaseException:
|
||||
await asyncio.to_thread(self.process_lock.close)
|
||||
self.process_lock = None
|
||||
raise
|
||||
self.available = True
|
||||
self.task = asyncio.create_task(self.loop())
|
||||
|
||||
@ -106,6 +117,7 @@ class Teams:
|
||||
self.task.cancel()
|
||||
with suppress(asyncio.CancelledError, Exception):
|
||||
await self.task
|
||||
await self.cancel_workers()
|
||||
self.handles.clear()
|
||||
if self.process_lock:
|
||||
await asyncio.to_thread(self.process_lock.close)
|
||||
@ -114,18 +126,32 @@ class Teams:
|
||||
async def loop(self):
|
||||
try:
|
||||
while True:
|
||||
await self.tick()
|
||||
await self.tick(wait=False)
|
||||
await asyncio.sleep(1)
|
||||
finally:
|
||||
self.available = False
|
||||
await self.cancel_workers()
|
||||
self.handles.clear()
|
||||
|
||||
def team_lock(self, owner, team_id):
|
||||
# Waiters/holders retain the lock; deleted teams do not leak lock entries.
|
||||
return self.locks.setdefault((owner, team_id), asyncio.Lock())
|
||||
|
||||
async def cancel_workers(self):
|
||||
workers = list(self.workers.values())
|
||||
for worker in workers:
|
||||
worker.cancel()
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
self.workers.clear()
|
||||
self.in_flight.clear()
|
||||
|
||||
def bind(self, context, team_id):
|
||||
if not self.available:
|
||||
raise ValueError("Team service unavailable; check Gateway startup diagnostics")
|
||||
if context.agent_runs is None:
|
||||
raise ValueError("Full Agent control requires an authenticated Gateway session")
|
||||
self.handles[(context.principal.user_id, team_id)] = context.agent_runs
|
||||
key = (context.principal.user_id, team_id)
|
||||
self.handles[key] = context.agent_runs
|
||||
|
||||
async def create(self, payload, context):
|
||||
fields(payload, "request_id name goal members")
|
||||
@ -141,7 +167,7 @@ class Teams:
|
||||
raise ValueError("Member names must be distinct")
|
||||
owner = context.principal.user_id
|
||||
team_id = identity(owner, text(payload, "request_id", 128))
|
||||
async with self.lock:
|
||||
async with self.team_lock(owner, team_id):
|
||||
self.bind(context, team_id)
|
||||
try:
|
||||
team = await self.db("create", owner, {"id": team_id, "name": name, "goal": goal, "members": roster, "messages": [], "jobs": [], "ready": False})
|
||||
@ -207,8 +233,9 @@ class Teams:
|
||||
|
||||
async def connect(self, payload, context):
|
||||
fields(payload, "team_id")
|
||||
async with self.lock:
|
||||
team = await self.db("get", context.principal.user_id, text(payload, "team_id", 64))
|
||||
owner, team_id = context.principal.user_id, text(payload, "team_id", 64)
|
||||
async with self.team_lock(owner, team_id):
|
||||
team = await self.db("get", owner, team_id)
|
||||
if context.agent_runs is None:
|
||||
raise ValueError("Agent control unavailable")
|
||||
for member in team["members"]:
|
||||
@ -255,7 +282,7 @@ class Teams:
|
||||
owner, team_id = context.principal.user_id, text(payload, "team_id", 64)
|
||||
member_id, content, key = text(payload, "member_id", 64), text(payload, "text"), text(payload, "request_id", 128)
|
||||
source = getattr(context, "thread_id", None)
|
||||
async with self.lock:
|
||||
async with self.team_lock(owner, team_id):
|
||||
team = await self.db("get", owner, team_id)
|
||||
if source and source not in {m["thread_id"] for m in team["members"]}:
|
||||
raise ValueError("Only team member conversations can send peer requests")
|
||||
@ -295,17 +322,17 @@ class Teams:
|
||||
# Gateway stamps user_id from its authenticated principal before
|
||||
# running middleware. Never derive team ownership from a mention or
|
||||
# from thread sharing: a shared conversation does not share a team.
|
||||
team = await self.db("get", owner, team_id)
|
||||
member = next((m for m in team["members"] if m["id"] == member_id), None)
|
||||
if member is None or member["thread_id"] == source:
|
||||
continue
|
||||
runs = handle.for_plugin(NAMESPACE)
|
||||
# Structured references and labels are untrusted. The host capability
|
||||
# must authorize BOTH conversations before the plugin uses the owner.
|
||||
await runs.get_state(thread_id=source)
|
||||
await runs.create_thread(assistant_id=member["agent"], thread_id=member["thread_id"], metadata={})
|
||||
content = text({"text": message.content}, "text")
|
||||
async with self.lock:
|
||||
async with self.team_lock(owner, team_id):
|
||||
team = await self.db("get", owner, team_id)
|
||||
member = next((m for m in team["members"] if m["id"] == member_id), None)
|
||||
if member is None or member["thread_id"] == source:
|
||||
continue
|
||||
runs = handle.for_plugin(NAMESPACE)
|
||||
# Structured references and labels are untrusted. The host capability
|
||||
# must authorize BOTH conversations before the plugin uses the owner.
|
||||
await runs.get_state(thread_id=source)
|
||||
await runs.create_thread(assistant_id=member["agent"], thread_id=member["thread_id"], metadata={})
|
||||
content = text({"text": message.content}, "text")
|
||||
self.handles[(owner, team_id)] = runs
|
||||
await self.db("change", owner, team_id, lambda t: self.enqueue(t, member_id=member_id, content=content, request_id=f"mention:{message.id}:{member_id}", source=source))
|
||||
queued += 1
|
||||
@ -331,7 +358,7 @@ class Teams:
|
||||
|
||||
async def control(self, payload, context, *, cancel):
|
||||
owner, team_id, job_id = context.principal.user_id, text(payload, "team_id", 64), text(payload, "job_id", 64)
|
||||
async with self.lock:
|
||||
async with self.team_lock(owner, team_id):
|
||||
team = await self.db("get", owner, team_id)
|
||||
self.bind(context, team_id)
|
||||
current = next((j for j in team["jobs"] if j["id"] == job_id), None)
|
||||
@ -362,9 +389,11 @@ class Teams:
|
||||
raise ValueError("Request unavailable")
|
||||
if cancel:
|
||||
if job["status"] not in TERMINAL:
|
||||
if job["resume"] and not job["resume"].get("started", True):
|
||||
job["resume"] = None
|
||||
job["status"] = "cancelling" if job["input"] else "cancelled"
|
||||
elif job["status"] == "waiting_input":
|
||||
job["resume"] = {"response": payload["response"], "key": identity(job_id, text(payload, "request_id", 128))}
|
||||
job["resume"] = {"response": payload["response"], "key": identity(job_id, text(payload, "request_id", 128)), "started": False}
|
||||
if response_input:
|
||||
job["resume"].update(input=response_input, key=identity(job_id, clarification["request_id"], payload["request_id"]))
|
||||
job["status"] = "resuming"
|
||||
@ -377,7 +406,7 @@ class Teams:
|
||||
async def delete(self, payload, context):
|
||||
fields(payload, "team_id")
|
||||
owner, team_id = context.principal.user_id, text(payload, "team_id", 64)
|
||||
async with self.lock:
|
||||
async with self.team_lock(owner, team_id):
|
||||
team = await self.db("get", owner, team_id)
|
||||
if any(j["status"] not in TERMINAL for j in team["jobs"]):
|
||||
raise ValueError("Cancel or complete pending requests first")
|
||||
@ -385,47 +414,115 @@ class Teams:
|
||||
self.handles.pop((owner, team_id), None)
|
||||
return {"deleted": True}
|
||||
|
||||
async def tick(self):
|
||||
async with self.lock:
|
||||
active = 0
|
||||
for owner, team_id in self.handles:
|
||||
current = await self.db("get", owner, team_id)
|
||||
active += sum(j["status"] in ("running", "resuming", "cancelling") or (j["status"] == "queued" and j["input"] is not None) for j in current["jobs"])
|
||||
for (owner, team_id), runs in list(self.handles.items()):
|
||||
team = await self.db("get", owner, team_id)
|
||||
@staticmethod
|
||||
def charged(job):
|
||||
if job["status"] in TERMINAL or job["status"] == "waiting_input":
|
||||
return False
|
||||
if job["status"] == "resuming" and job["resume"] and not job["resume"].get("started", True):
|
||||
return False
|
||||
return job["input"] is not None
|
||||
|
||||
async def plan(self, attempted, eligible):
|
||||
# Only SQLite snapshots and slot bookkeeping live under this lock.
|
||||
# Host calls run in independent workers, serialized by team.
|
||||
async with self.scheduling_lock:
|
||||
teams = await self.db("scheduling_snapshots")
|
||||
heads, entries = {}, []
|
||||
for key, team in teams.items():
|
||||
occupied = set()
|
||||
for job in team["jobs"]:
|
||||
if job["status"] in TERMINAL or job["thread_id"] in occupied:
|
||||
continue
|
||||
occupied.add(job["thread_id"])
|
||||
reserved = job["status"] == "queued" and job["input"] is None
|
||||
if reserved:
|
||||
if active >= 8:
|
||||
continue
|
||||
active += 1
|
||||
try:
|
||||
async with asyncio.timeout(10):
|
||||
await self.advance(owner, team, job, runs)
|
||||
except AgentRunError as exc:
|
||||
if exc.status_code in (403, 503):
|
||||
self.handles.pop((owner, team_id), None)
|
||||
entries.append((job.get("created_at") or "", key, job))
|
||||
entries.sort(key=lambda entry: entry[0])
|
||||
for _, key, job in entries:
|
||||
# Receipts can target a conversation shared across teams/owners.
|
||||
heads.setdefault(job["thread_id"], (key, job["id"]))
|
||||
if eligible is None:
|
||||
eligible = set(heads.values())
|
||||
charged = {(key, j["id"]) for key, t in teams.items() for j in t["jobs"] if self.charged(j)} | self.in_flight
|
||||
groups, pending = {}, {}
|
||||
for _, key, job in entries:
|
||||
entry = (key, job["id"])
|
||||
if entry not in eligible or entry in attempted or heads[job["thread_id"]] != entry or job["status"] == "waiting_input":
|
||||
continue
|
||||
if key not in self.handles or key in self.workers or self.team_lock(*key).locked():
|
||||
continue
|
||||
if entry in charged:
|
||||
groups.setdefault(key, []).append(job["id"])
|
||||
else:
|
||||
pending.setdefault(key[0], deque()).append((key, job["id"]))
|
||||
# Carry the unfinished round, even when the last admitted owner has
|
||||
# no pending job left. Drop unavailable queues and append newcomers.
|
||||
owners = deque(owner for owner in self.owner_rotation if owner in pending)
|
||||
retained = set(owners)
|
||||
owners.extend(owner for owner in pending if owner not in retained)
|
||||
while owners and len(charged) < 8:
|
||||
owner = owners.popleft()
|
||||
key, job_id = pending[owner].popleft()
|
||||
entry = (key, job_id)
|
||||
charged.add(entry)
|
||||
groups.setdefault(key, []).append(job_id)
|
||||
if pending[owner]:
|
||||
owners.append(owner)
|
||||
self.owner_rotation = owners
|
||||
tasks = []
|
||||
for key, job_ids in groups.items():
|
||||
reservations = {(key, job_id) for job_id in job_ids}
|
||||
attempted.update(reservations)
|
||||
self.in_flight.update(reservations)
|
||||
task = asyncio.create_task(self.work_team(key, job_ids, self.handles[key]))
|
||||
self.workers[key] = task
|
||||
tasks.append(task)
|
||||
return tasks, eligible
|
||||
|
||||
async def tick(self, *, wait=True):
|
||||
# Tests/manual ticks drain this pass. The service loop never waits for a
|
||||
# slow team's worker before it can poll another team again.
|
||||
attempted, eligible = set(), None
|
||||
while True:
|
||||
tasks, eligible = await self.plan(attempted, eligible)
|
||||
if not tasks or not wait:
|
||||
return
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def work_team(self, key, job_ids, runs):
|
||||
owner, team_id = key
|
||||
try:
|
||||
for job_id in job_ids:
|
||||
try:
|
||||
async with self.team_lock(*key):
|
||||
if not self.available or self.handles.get(key) is not runs:
|
||||
break
|
||||
if exc.status_code != 409:
|
||||
await self.finish(owner, team, job, "failed", "Agent unavailable or request rejected")
|
||||
except (TimeoutError, OSError):
|
||||
# Retain the frozen input and key: the host may have
|
||||
# accepted the request before the transport failed.
|
||||
continue
|
||||
except Exception:
|
||||
# An unknown host error may follow successful admission.
|
||||
# Disconnect rather than releasing the member's slot.
|
||||
self.handles.pop((owner, team_id), None)
|
||||
break
|
||||
finally:
|
||||
# A conversation wait did not attempt admission. Frozen
|
||||
# inputs retain the slot, including unknown outcomes.
|
||||
if reserved and job["input"] is None:
|
||||
active -= 1
|
||||
team = await self.db("get", *key)
|
||||
job = next(j for j in team["jobs"] if j["id"] == job_id)
|
||||
if job["status"] in TERMINAL:
|
||||
continue
|
||||
try:
|
||||
async with asyncio.timeout(10):
|
||||
await self.advance(owner, team, job, runs)
|
||||
except AgentRunError as exc:
|
||||
if exc.status_code in (403, 503):
|
||||
self.handles.pop(key, None)
|
||||
break
|
||||
if exc.status_code != 409:
|
||||
await self.finish(owner, team, job, "failed", "Agent unavailable or request rejected")
|
||||
except (TimeoutError, OSError):
|
||||
pass # Frozen input/key retain the slot and ownership.
|
||||
except Exception:
|
||||
self.handles.pop(key, None)
|
||||
break
|
||||
finally:
|
||||
async with self.scheduling_lock:
|
||||
self.in_flight.discard((key, job_id))
|
||||
except Exception:
|
||||
# Persistence/read failures disconnect only this worker's team.
|
||||
self.handles.pop(key, None)
|
||||
finally:
|
||||
async with self.scheduling_lock:
|
||||
self.in_flight.difference_update((key, job_id) for job_id in job_ids)
|
||||
self.workers.pop(key, None)
|
||||
|
||||
async def save_job(self, owner, team_id, job):
|
||||
def update(team):
|
||||
@ -466,6 +563,10 @@ class Teams:
|
||||
job["status"] = "running"
|
||||
await self.save_job(owner, team["id"], job)
|
||||
if job["status"] == "resuming" or (job["status"] == "cancelling" and job["resume"]):
|
||||
# A pending answer acquires a slot before attempting admission. An
|
||||
# unknown acknowledgement must remain charged across reconnects.
|
||||
job["resume"]["started"] = True
|
||||
await self.save_job(owner, team["id"], job)
|
||||
continuation = job["resume"].get("input")
|
||||
if continuation:
|
||||
state = await runs.get_state(thread_id=thread)
|
||||
|
||||
@ -5,6 +5,12 @@ import sqlite3
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
|
||||
TERMINAL = {"completed", "failed", "cancelled"}
|
||||
|
||||
|
||||
def has_pending_jobs(team):
|
||||
return int(any(job["status"] not in TERMINAL for job in team["jobs"]))
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self, path):
|
||||
@ -13,11 +19,25 @@ class Store:
|
||||
raise ValueError("storage_path must be an absolute deployment-owned path")
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with closing(self.connection()) as db, db:
|
||||
db.execute("CREATE TABLE IF NOT EXISTS teams (id TEXT PRIMARY KEY, owner TEXT NOT NULL, data TEXT NOT NULL)")
|
||||
# Serialize schema discovery and backfill with concurrent writers.
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
db.execute("CREATE TABLE IF NOT EXISTS teams (id TEXT PRIMARY KEY, owner TEXT NOT NULL, data TEXT NOT NULL, has_pending_jobs INTEGER NOT NULL DEFAULT 0)")
|
||||
if "has_pending_jobs" not in {row[1] for row in db.execute("PRAGMA table_info(teams)")}:
|
||||
db.execute("ALTER TABLE teams ADD COLUMN has_pending_jobs INTEGER NOT NULL DEFAULT 0")
|
||||
rows = db.execute("SELECT id, data FROM teams").fetchall()
|
||||
for team_id, data in rows:
|
||||
db.execute("UPDATE teams SET has_pending_jobs=? WHERE id=?", (has_pending_jobs(json.loads(data)), team_id))
|
||||
db.execute("CREATE INDEX IF NOT EXISTS teams_pending_jobs ON teams(has_pending_jobs) WHERE has_pending_jobs=1")
|
||||
|
||||
def connection(self):
|
||||
return sqlite3.connect(self.path, timeout=5)
|
||||
|
||||
def scheduling_snapshots(self):
|
||||
# Service-only, indexed projection. Completed history is never decoded;
|
||||
# all unfinished work is retained, regardless of capability availability.
|
||||
with closing(self.connection()) as db:
|
||||
return {(owner, team_id): json.loads(data) for owner, team_id, data in db.execute("SELECT owner, id, data FROM teams WHERE has_pending_jobs=1 ORDER BY rowid")}
|
||||
|
||||
def list(self, owner):
|
||||
with closing(self.connection()) as db:
|
||||
return [json.loads(row[0]) for row in db.execute("SELECT data FROM teams WHERE owner=? ORDER BY rowid DESC", (owner,))]
|
||||
@ -37,7 +57,7 @@ class Store:
|
||||
return json.loads(row[0])
|
||||
if db.execute("SELECT count(*) FROM teams WHERE owner=?", (owner,)).fetchone()[0] >= 20:
|
||||
raise ValueError("Team limit reached (20)")
|
||||
db.execute("INSERT INTO teams VALUES (?,?,?)", (team["id"], owner, json.dumps(team)))
|
||||
db.execute("INSERT INTO teams (id, owner, data, has_pending_jobs) VALUES (?,?,?,?)", (team["id"], owner, json.dumps(team), has_pending_jobs(team)))
|
||||
return team
|
||||
|
||||
def change(self, owner, team_id, update):
|
||||
@ -48,7 +68,7 @@ class Store:
|
||||
raise ValueError("Team unavailable")
|
||||
team = json.loads(row[0])
|
||||
result = update(team)
|
||||
db.execute("UPDATE teams SET data=? WHERE owner=? AND id=?", (json.dumps(team), owner, team_id))
|
||||
db.execute("UPDATE teams SET data=?, has_pending_jobs=? WHERE owner=? AND id=?", (json.dumps(team), has_pending_jobs(team), owner, team_id))
|
||||
return result
|
||||
|
||||
def delete(self, owner, team_id):
|
||||
|
||||
@ -36,7 +36,7 @@ upstream dependency chain resolves a patched version without it; regenerate
|
||||
| `pnpm typecheck` | TypeScript type check (`tsc --noEmit`) |
|
||||
| `pnpm start` | Start production server |
|
||||
|
||||
Unit tests live under `tests/unit/` and mirror the `src/` layout (e.g., `tests/unit/core/api/stream-mode.test.ts` tests `src/core/api/stream-mode.ts`). Powered by Rstest; import source modules via the `@/` path alias.
|
||||
Rstest unit tests under `tests/unit/` mirror `src/` (e.g., `tests/unit/core/api/stream-mode.test.ts` tests `src/core/api/stream-mode.ts`). Import source modules via `@/`.
|
||||
|
||||
Webpack is the default development bundler. Use `DEER_FLOW_DEV_BUNDLER=turbo` with `pnpm dev` to opt in to Turbopack when diagnosing a local Next.js bundler issue.
|
||||
|
||||
@ -70,7 +70,7 @@ Frontend (Next.js) ──▶ LangGraph SDK ──▶ LangGraph Backend (lead_age
|
||||
└── Tools & Skills
|
||||
```
|
||||
|
||||
The frontend is a stateful chat application. Users create **threads** (conversations), send messages, set thread-scoped `/goal` completion conditions, and receive streamed AI responses. The backend orchestrates agents that can produce **artifacts** (files/code), **todos**, and goal state updates.
|
||||
Users create **threads** (conversations), send messages, set thread-scoped `/goal` completion conditions, and receive streamed responses, agent **artifacts** (files/code), **todos** and goal state updates.
|
||||
|
||||
### Source Layout (`src/`)
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user