mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
fix(memory): cancel buffered extraction when agent is deleted or cleared (#5123)
* fix(memory): cancel buffered extraction when agent is deleted or cleared * fix(memory): cancel buffered work before agent delete Address review: cancel before/after delete to close the rmtree race, scope user_id=None cancels to the legacy root only, and import memory helpers at module scope. Signed-off-by: SPEC <zt1y17@soton.ac.uk> * fix(memory): close remaining cancel races from review Post-clear cancel, legacy-only all_agents scope, always cancel even when memory is disabled, and fold cancel+delete into one offloaded thread. Signed-off-by: SPEC <zt1y17@soton.ac.uk> * docs(memory): align cancel_by_agent None-scope with legacy root Document that user_id=None cancels only the legacy no-user bucket, matching clear/storage semantics, not the whole process-local queue. Signed-off-by: SPEC <zt1y17@soton.ac.uk> * test(memory): fix cancel_by_agent docstring regression assertion Signed-off-by: SPEC <zt1y17@soton.ac.uk> * fix(memory): address final cancel review nits Type the delete helper with AgentStore, replace docstring pinning with a kwargs mapping test, and document scoped cancel + residual window in AGENTS.md. Signed-off-by: SPEC <zt1y17@soton.ac.uk> * fix(memory): resolve agent store inside delete worker thread get_agent_store() does blocking config/FS work; keep it off the event loop so test_delete_agent_does_not_block_event_loop and backend-blocking-io CI pass. Signed-off-by: SPEC <zt1y17@soton.ac.uk> --------- Signed-off-by: SPEC <zt1y17@soton.ac.uk>
This commit is contained in:
SPEC
GitHub
parent
281f04b9eb
commit
822c7bca4b
@ -8,6 +8,7 @@ from typing import Literal
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from deerflow.agents.memory.manager import get_memory_manager
|
||||
from deerflow.config.agents_api_config import get_agents_api_config
|
||||
from deerflow.config.agents_config import (
|
||||
AgentConfig,
|
||||
@ -566,12 +567,11 @@ async def delete_agent(name: str) -> None:
|
||||
_validate_agent_name(name)
|
||||
name = _normalize_agent_name(name)
|
||||
user_id = get_effective_user_id()
|
||||
try:
|
||||
# Off the event loop: file rmtree or a DB delete plus memory cleanup.
|
||||
def _delete_agent() -> AgentDeleteOutcome:
|
||||
return get_agent_store().delete(name, user_id=user_id)
|
||||
|
||||
outcome = await asyncio.to_thread(_delete_agent)
|
||||
try:
|
||||
# Off the event loop: resolve store + cancel → delete → cancel-on-success
|
||||
# (get_agent_store / memory manager do blocking config and FS I/O).
|
||||
outcome = await asyncio.to_thread(_delete_agent_with_memory_cancel, name, user_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete agent '{name}': {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail=f"Failed to delete agent: {str(e)}")
|
||||
@ -590,3 +590,45 @@ async def delete_agent(name: str) -> None:
|
||||
)
|
||||
|
||||
logger.info(f"Deleted agent '{name}'")
|
||||
|
||||
|
||||
def _delete_agent_with_memory_cancel(name: str, user_id: str | None) -> AgentDeleteOutcome:
|
||||
"""Cancel buffered memory, delete the agent, then cancel again on success.
|
||||
|
||||
Runs entirely in a worker thread so blocking store/memory/config I/O stays
|
||||
off the event loop. Pre-delete cancel closes the debounce-timer race during
|
||||
rmtree; post-success cancel covers work enqueued mid-delete. A rejected
|
||||
delete may still drop buffered work; that update is re-fed on the next turn.
|
||||
"""
|
||||
store = get_agent_store()
|
||||
_cancel_pending_memory_for_agent(name, user_id)
|
||||
outcome = store.delete(name, user_id=user_id)
|
||||
if outcome == "deleted":
|
||||
_cancel_pending_memory_for_agent(name, user_id)
|
||||
return outcome
|
||||
|
||||
|
||||
def _cancel_pending_memory_for_agent(name: str, user_id: str | None) -> None:
|
||||
"""Best-effort cancel of buffered memory extraction for one agent scope.
|
||||
|
||||
Always attempts cancellation even if memory is currently disabled: settings
|
||||
are hot-reloadable and disabling does not destroy an already-live queue.
|
||||
Failure must never fail agent deletion: a dropped update is re-fed on the
|
||||
next conversation turn per the queue contract.
|
||||
|
||||
Callers must run this off the event loop (blocking config/manager I/O).
|
||||
"""
|
||||
try:
|
||||
cancelled = get_memory_manager().cancel_by_agent(name, user_id=user_id)
|
||||
if cancelled:
|
||||
logger.info(
|
||||
"Cancelled %d pending memory update(s) for agent '%s'",
|
||||
cancelled,
|
||||
name,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to cancel pending memory updates for agent '%s' (non-fatal)",
|
||||
name,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
@ -3,6 +3,7 @@
|
||||
**Components**:
|
||||
- `updater.py` - LLM-based memory updates with fact extraction, whitespace-normalized fact deduplication, optimistic revision checks, and repository change sets
|
||||
- `queue.py` - Debounced update queue (per-thread deduplication, configurable wait time); captures `user_id` at enqueue time so it survives the `threading.Timer` boundary
|
||||
- `manager.py` / DeerMem `cancel_by_agent` - Scoped cancellation of **pending** debounce contexts (used by agent delete and `clear_memory`). `user_id=None` means the **legacy no-user root only**, never every user in the process; `agent_name=None` cancels every agent bucket inside that user scope. Contexts already pulled out of `_items` by an in-flight `_process_queue` worker are deliberately left alone — do not "fix" that residual by interrupting mid-LLM extraction; a durable outbox would be required for that. There is no whole-queue cancel form; broader sweeps must iterate known user scopes.
|
||||
- `prompt.py` - Prompt templates for memory updates
|
||||
- `storage.py` - File repository with one user-global summary JSON, agent-owned single-fact Markdown, target-only journaled changes, strict fact validation, shared-user plus per-fact optimistic revisions, lock-protected migration, deep-copy caching, and a RetrievalPort adapter boundary
|
||||
- `retrieval.py` - Built-in scope-aware SQLite FTS5/BM25 adapter; it stores only rebuildable derived data and can be disabled with an empty `retrieval_adapter`. Chinese jieba tokenization is optional via the backend `memory-zh` extra; without it the adapter uses SQLite unicode tokenization and the substring fallback. A corrupt persistent derived database is deleted and recreated once before falling back to substring retrieval. The Gateway closes the derived SQLite connection after its shutdown flush; reads and writes remain serialized by the adapter lock, with connection pooling deferred as a performance follow-up.
|
||||
|
||||
@ -454,16 +454,39 @@ class DeerMem(MemoryManager):
|
||||
# NotImplementedError) -- they are dead contract (zero callers; /memory/export
|
||||
# routes via get_memory), so DeerMem no longer repeats the raise.
|
||||
|
||||
def cancel_by_agent(
|
||||
self,
|
||||
agent_name: str | None = None,
|
||||
*,
|
||||
user_id: str | None = None,
|
||||
) -> int:
|
||||
"""Drop pending debounce-queue contexts for a deleted or cleared scope.
|
||||
|
||||
``user_id=None`` matches the legacy no-user storage root only (same as
|
||||
``clear_memory`` / ``clear_*_memory_data``), not every queued user.
|
||||
"""
|
||||
if agent_name is None:
|
||||
return self._queue.cancel_by_agent(user_id=user_id, all_agents=True)
|
||||
return self._queue.cancel_by_agent(
|
||||
_resolve_agent_name(agent_name),
|
||||
user_id=user_id,
|
||||
all_agents=False,
|
||||
)
|
||||
|
||||
def clear_memory(
|
||||
self,
|
||||
*,
|
||||
user_id: str | None = None,
|
||||
agent_name: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
# Cancel same-scope pending extraction before and after clearing so a
|
||||
# stale debounce timer cannot rewrite facts during/after the clear.
|
||||
self.cancel_by_agent(agent_name, user_id=user_id)
|
||||
if agent_name is None:
|
||||
memory_data = _call_backend(lambda: self._updater.clear_all_memory_data(user_id=user_id))
|
||||
else:
|
||||
memory_data = _call_backend(lambda: self._updater.clear_memory_data(agent_name=_resolve_agent_name(agent_name), user_id=user_id))
|
||||
self.cancel_by_agent(agent_name, user_id=user_id)
|
||||
return _compat_document(memory_data)
|
||||
|
||||
def import_memory(
|
||||
|
||||
@ -400,6 +400,51 @@ class MemoryUpdateQueue:
|
||||
# before _process_queue completes. Acceptable for best-effort memory updates.
|
||||
self._schedule_timer(0)
|
||||
|
||||
def cancel_by_agent(
|
||||
self,
|
||||
agent_name: str | None = None,
|
||||
*,
|
||||
user_id: str | None = None,
|
||||
all_agents: bool = False,
|
||||
) -> int:
|
||||
"""Drop pending contexts for a scope without processing them.
|
||||
|
||||
Matches ``(agent_name, user_id)`` against items still sitting in
|
||||
``_items``. Contexts already pulled out by an in-flight
|
||||
:meth:`_process_queue` worker are deliberately left alone -- interrupting
|
||||
mid-LLM-call belongs to a durable outbox, not this in-memory debounce
|
||||
queue.
|
||||
|
||||
Args:
|
||||
agent_name: Canonical agent bucket to cancel. Ignored when
|
||||
``all_agents`` is True.
|
||||
user_id: When set, only that user's pending contexts are eligible.
|
||||
When omitted, only the legacy ``user_id is None`` root is
|
||||
cancelled (mirrors storage), whether ``all_agents`` is set or not.
|
||||
all_agents: When True, ignore ``agent_name`` and cancel every agent
|
||||
bucket in the matched user scope (used by ``clear_memory`` with
|
||||
``agent_name=None``).
|
||||
|
||||
Returns:
|
||||
Number of pending contexts removed.
|
||||
"""
|
||||
with self._lock:
|
||||
before = len(self._items)
|
||||
|
||||
def _keep(context: ConversationContext) -> bool:
|
||||
# Scope matches storage: user_id=None is the legacy no-user root
|
||||
# only (None == None), never "every user".
|
||||
if not all_agents and context.agent_name != agent_name:
|
||||
return True
|
||||
return context.user_id != user_id
|
||||
|
||||
self._items = [context for context in self._items if _keep(context)]
|
||||
removed = before - len(self._items)
|
||||
if removed and not self._items and self._timer is not None:
|
||||
self._timer.cancel()
|
||||
self._timer = None
|
||||
return removed
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear the queue without processing.
|
||||
|
||||
|
||||
@ -314,6 +314,36 @@ class MemoryManager(BaseModel):
|
||||
"""
|
||||
raise NotImplementedError(f"clear_memory not supported by {type(self).__name__}")
|
||||
|
||||
def cancel_by_agent(
|
||||
self,
|
||||
agent_name: str | None = None,
|
||||
*,
|
||||
user_id: str | None = None,
|
||||
) -> int:
|
||||
"""Cancel buffered memory-extraction work for a scope.
|
||||
|
||||
Backends without a debounce queue have nothing to cancel and inherit
|
||||
the default ``0``. DeerMem drops matching pending contexts so a deleted
|
||||
or cleared agent cannot be resurrected by a late timer fire.
|
||||
|
||||
Scope (must stay symmetric with ``clear_memory`` / storage buckets):
|
||||
|
||||
- ``user_id`` selects the user bucket. ``user_id=None`` means the
|
||||
**legacy no-user root only**, never "every user in the process".
|
||||
- ``agent_name=None`` cancels every agent bucket inside that user
|
||||
scope (including the default/global bucket).
|
||||
- An explicit ``agent_name`` cancels only that agent's pending
|
||||
contexts inside the same user scope.
|
||||
|
||||
There is no "cancel the whole process-local queue" form of this
|
||||
method; callers that need a broader sweep must iterate known user
|
||||
scopes explicitly.
|
||||
|
||||
Returns:
|
||||
Number of pending contexts cancelled. Default: ``0``.
|
||||
"""
|
||||
return 0
|
||||
|
||||
def import_memory(
|
||||
self,
|
||||
memory_data: dict[str, Any],
|
||||
|
||||
209
backend/tests/test_memory_cancel_by_agent.py
Normal file
209
backend/tests/test_memory_cancel_by_agent.py
Normal file
@ -0,0 +1,209 @@
|
||||
"""Regression for #5037: scoped cancellation of buffered memory extraction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.queue import ConversationContext
|
||||
from deerflow.agents.memory.manager import MemoryManager, get_memory_manager, reset_memory_manager
|
||||
from deerflow.config.memory_config import MemoryConfig, get_memory_config, set_memory_config
|
||||
|
||||
|
||||
def test_deermem_cancel_by_agent_uses_canonical_bucket(tmp_path) -> None:
|
||||
mem = DeerMem(backend_config={"storage_path": str(tmp_path)})
|
||||
with patch.object(mem._queue, "_schedule_timer"):
|
||||
mem._queue.add(thread_id="t1", messages=["m"], agent_name="research-agent", user_id="u1")
|
||||
mem._queue.add(thread_id="t2", messages=["m"], agent_name="other", user_id="u1")
|
||||
|
||||
removed = mem.cancel_by_agent("Research-Agent", user_id="u1")
|
||||
|
||||
assert removed == 1
|
||||
assert mem._queue.pending_count == 1
|
||||
assert mem._queue._items[0].agent_name == "other"
|
||||
|
||||
|
||||
def test_deermem_clear_memory_cancels_before_and_after_clear(tmp_path) -> None:
|
||||
mem = DeerMem(backend_config={"storage_path": str(tmp_path)})
|
||||
mem._queue._items = [
|
||||
ConversationContext(thread_id="t1", messages=["m"], agent_name="research-agent", user_id="u1"),
|
||||
ConversationContext(thread_id="t2", messages=["m"], agent_name="other", user_id="u1"),
|
||||
]
|
||||
mem._updater = MagicMock()
|
||||
|
||||
def _clear(**kwargs):
|
||||
mem._queue._items.append(ConversationContext(thread_id="t-mid", messages=["m"], agent_name="research-agent", user_id="u1"))
|
||||
return {"facts": []}
|
||||
|
||||
mem._updater.clear_memory_data.side_effect = _clear
|
||||
|
||||
mem.clear_memory(agent_name="research-agent", user_id="u1")
|
||||
|
||||
assert [c.agent_name for c in mem._queue._items] == ["other"]
|
||||
mem._updater.clear_memory_data.assert_called_once_with(agent_name="research-agent", user_id="u1")
|
||||
|
||||
|
||||
def test_deermem_clear_all_cancels_all_pending_for_user(tmp_path) -> None:
|
||||
mem = DeerMem(backend_config={"storage_path": str(tmp_path)})
|
||||
mem._queue._items = [
|
||||
ConversationContext(thread_id="t1", messages=["m"], agent_name="a", user_id="u1"),
|
||||
ConversationContext(thread_id="t2", messages=["m"], agent_name="b", user_id="u1"),
|
||||
ConversationContext(thread_id="t3", messages=["m"], agent_name="a", user_id="u2"),
|
||||
]
|
||||
mem._updater = MagicMock()
|
||||
mem._updater.clear_all_memory_data.return_value = {"facts": []}
|
||||
|
||||
mem.clear_memory(user_id="u1")
|
||||
|
||||
assert mem._queue.pending_count == 1
|
||||
assert mem._queue._items[0].user_id == "u2"
|
||||
mem._updater.clear_all_memory_data.assert_called_once_with(user_id="u1")
|
||||
|
||||
|
||||
def test_base_memory_manager_cancel_by_agent_defaults_to_zero() -> None:
|
||||
class _Bare(MemoryManager):
|
||||
def add(self, thread_id, messages, *, agent_name=None, user_id=None, trace_id=None) -> None:
|
||||
return None
|
||||
|
||||
def get_context(self, user_id, *, agent_name=None, thread_id=None) -> str:
|
||||
return ""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, backend_config, *, mode="middleware", **host_hooks):
|
||||
return cls(backend_config=backend_config or {}, mode=mode)
|
||||
|
||||
assert _Bare().cancel_by_agent("x", user_id="u") == 0
|
||||
|
||||
|
||||
def test_deermem_cancel_by_agent_forwards_scoped_queue_kwargs(tmp_path) -> None:
|
||||
"""Manager mapping: None agent → all_agents; named agent → canonical bucket."""
|
||||
mem = DeerMem(backend_config={"storage_path": str(tmp_path)})
|
||||
queue = MagicMock()
|
||||
queue.cancel_by_agent.return_value = 0
|
||||
mem._queue = queue
|
||||
|
||||
mem.cancel_by_agent(None, user_id=None)
|
||||
queue.cancel_by_agent.assert_called_with(user_id=None, all_agents=True)
|
||||
|
||||
mem.cancel_by_agent("Research-Agent", user_id="u1")
|
||||
queue.cancel_by_agent.assert_called_with("research-agent", user_id="u1", all_agents=False)
|
||||
|
||||
|
||||
def test_delete_agent_cancels_before_and_after_successful_delete(tmp_path) -> None:
|
||||
"""Cancel must run before store.delete so a timer cannot resurrect mid-rmtree."""
|
||||
from app.gateway.routers import agents as agents_router
|
||||
|
||||
orig = get_memory_config()
|
||||
reset_memory_manager()
|
||||
set_memory_config(
|
||||
MemoryConfig(
|
||||
enabled=True,
|
||||
manager_class="deermem",
|
||||
backend_config={"storage_path": str(tmp_path / "memory")},
|
||||
)
|
||||
)
|
||||
try:
|
||||
manager = get_memory_manager()
|
||||
with patch.object(manager._queue, "_schedule_timer"):
|
||||
manager._queue.add(thread_id="t1", messages=["m"], agent_name="gone", user_id="user-1")
|
||||
manager._queue.add(thread_id="t2", messages=["m"], agent_name="keep", user_id="user-1")
|
||||
|
||||
store = MagicMock()
|
||||
order: list[str] = []
|
||||
|
||||
def _delete(name, *, user_id=None):
|
||||
order.append("delete")
|
||||
manager._queue._items.append(ConversationContext(thread_id="t-mid", messages=["m"], agent_name="gone", user_id="user-1"))
|
||||
return "deleted"
|
||||
|
||||
store.delete.side_effect = _delete
|
||||
real_cancel = agents_router._cancel_pending_memory_for_agent
|
||||
|
||||
def tracked_cancel(name, user_id):
|
||||
order.append("cancel")
|
||||
return real_cancel(name, user_id)
|
||||
|
||||
with (
|
||||
patch.object(agents_router, "_require_agents_api_enabled"),
|
||||
patch.object(agents_router, "_validate_agent_name"),
|
||||
patch.object(agents_router, "_normalize_agent_name", side_effect=lambda n: n.lower()),
|
||||
patch.object(agents_router, "get_effective_user_id", return_value="user-1"),
|
||||
patch.object(agents_router, "get_agent_store", return_value=store),
|
||||
patch.object(agents_router, "_cancel_pending_memory_for_agent", side_effect=tracked_cancel),
|
||||
):
|
||||
asyncio.run(agents_router.delete_agent("Gone"))
|
||||
|
||||
assert order == ["cancel", "delete", "cancel"]
|
||||
assert manager._queue.pending_count == 1
|
||||
assert manager._queue._items[0].agent_name == "keep"
|
||||
store.delete.assert_called_once_with("gone", user_id="user-1")
|
||||
finally:
|
||||
set_memory_config(orig)
|
||||
reset_memory_manager()
|
||||
|
||||
|
||||
def test_delete_agent_still_cancels_when_memory_disabled(tmp_path) -> None:
|
||||
"""Disabling memory must not skip cancel of an already-live queue."""
|
||||
from app.gateway.routers import agents as agents_router
|
||||
|
||||
orig = get_memory_config()
|
||||
reset_memory_manager()
|
||||
set_memory_config(
|
||||
MemoryConfig(
|
||||
enabled=True,
|
||||
manager_class="deermem",
|
||||
backend_config={"storage_path": str(tmp_path / "memory")},
|
||||
)
|
||||
)
|
||||
try:
|
||||
manager = get_memory_manager()
|
||||
with patch.object(manager._queue, "_schedule_timer"):
|
||||
manager._queue.add(thread_id="t1", messages=["m"], agent_name="gone", user_id="user-1")
|
||||
|
||||
# Hot-disable after work was queued; delete must still cancel.
|
||||
set_memory_config(MemoryConfig(enabled=False, manager_class="deermem"))
|
||||
|
||||
store = MagicMock()
|
||||
store.delete.return_value = "deleted"
|
||||
|
||||
with (
|
||||
patch.object(agents_router, "_require_agents_api_enabled"),
|
||||
patch.object(agents_router, "_validate_agent_name"),
|
||||
patch.object(agents_router, "_normalize_agent_name", side_effect=lambda n: n.lower()),
|
||||
patch.object(agents_router, "get_effective_user_id", return_value="user-1"),
|
||||
patch.object(agents_router, "get_agent_store", return_value=store),
|
||||
):
|
||||
asyncio.run(agents_router.delete_agent("Gone"))
|
||||
|
||||
assert manager._queue.pending_count == 0
|
||||
finally:
|
||||
set_memory_config(orig)
|
||||
reset_memory_manager()
|
||||
|
||||
|
||||
def test_delete_agent_still_cancels_before_rejected_delete() -> None:
|
||||
"""Pre-delete cancel is intentional even when delete later 404s."""
|
||||
from app.gateway.routers import agents as agents_router
|
||||
|
||||
store = MagicMock()
|
||||
store.delete.return_value = "missing"
|
||||
manager = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(agents_router, "_require_agents_api_enabled"),
|
||||
patch.object(agents_router, "_validate_agent_name"),
|
||||
patch.object(agents_router, "_normalize_agent_name", side_effect=lambda n: n.lower()),
|
||||
patch.object(agents_router, "get_effective_user_id", return_value="user-1"),
|
||||
patch.object(agents_router, "get_agent_store", return_value=store),
|
||||
patch.object(agents_router, "get_memory_manager", return_value=manager),
|
||||
):
|
||||
try:
|
||||
asyncio.run(agents_router.delete_agent("ghost"))
|
||||
raise AssertionError("expected 404")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 404
|
||||
|
||||
manager.cancel_by_agent.assert_called_once_with("ghost", user_id="user-1")
|
||||
@ -79,6 +79,7 @@ def test_minimal_backend_onboards_via_factory_with_only_add_get_context():
|
||||
manager.get_memory(user_id="u")
|
||||
with pytest.raises(NotImplementedError):
|
||||
manager.clear_memory(user_id="u")
|
||||
assert manager.cancel_by_agent("x", user_id="u") == 0
|
||||
with pytest.raises(NotImplementedError):
|
||||
manager.import_memory({}, user_id="u")
|
||||
with pytest.raises(NotImplementedError):
|
||||
|
||||
@ -403,3 +403,104 @@ def test_flush_sync_skips_inter_item_delay_on_drain_path() -> None:
|
||||
# No inter-item rate-limit sleep on the drain path.
|
||||
mock_sleep.assert_not_called()
|
||||
assert mock_updater.update_memory.call_count == 3
|
||||
|
||||
|
||||
def test_cancel_by_agent_drops_matching_pending_and_preserves_others() -> None:
|
||||
"""#5037: deleting/clearing an agent must drop its debounce buffer only."""
|
||||
queue = _queue()
|
||||
with patch.object(queue, "_schedule_timer"):
|
||||
queue.add(thread_id="t1", messages=["keep"], agent_name="alice", user_id="u1")
|
||||
queue.add(thread_id="t2", messages=["drop"], agent_name="bob", user_id="u1")
|
||||
queue.add(thread_id="t3", messages=["other-user"], agent_name="bob", user_id="u2")
|
||||
|
||||
existing_timer = MagicMock()
|
||||
queue._timer = existing_timer
|
||||
|
||||
removed = queue.cancel_by_agent("bob", user_id="u1")
|
||||
|
||||
assert removed == 1
|
||||
assert queue.pending_count == 2
|
||||
assert {(c.agent_name, c.user_id) for c in queue._items} == {("alice", "u1"), ("bob", "u2")}
|
||||
existing_timer.cancel.assert_not_called()
|
||||
|
||||
|
||||
def test_cancel_by_agent_all_agents_for_user_cancels_timer_when_empty() -> None:
|
||||
queue = _queue()
|
||||
with patch.object(queue, "_schedule_timer"):
|
||||
queue.add(thread_id="t1", messages=["a"], agent_name="alice", user_id="u1")
|
||||
queue.add(thread_id="t2", messages=["b"], agent_name="bob", user_id="u1")
|
||||
queue.add(thread_id="t3", messages=["c"], agent_name="alice", user_id="u2")
|
||||
|
||||
existing_timer = MagicMock()
|
||||
queue._timer = existing_timer
|
||||
|
||||
removed = queue.cancel_by_agent(user_id="u1", all_agents=True)
|
||||
|
||||
assert removed == 2
|
||||
assert queue.pending_count == 1
|
||||
assert queue._items[0].user_id == "u2"
|
||||
existing_timer.cancel.assert_not_called()
|
||||
|
||||
removed_rest = queue.cancel_by_agent(user_id="u2", all_agents=True)
|
||||
assert removed_rest == 1
|
||||
assert queue.pending_count == 0
|
||||
existing_timer.cancel.assert_called_once_with()
|
||||
assert queue._timer is None
|
||||
|
||||
|
||||
def test_cancel_by_agent_without_user_id_only_drops_legacy_scope() -> None:
|
||||
"""user_id=None mirrors storage: legacy root only, not every user."""
|
||||
queue = _queue()
|
||||
with patch.object(queue, "_schedule_timer"):
|
||||
queue.add(thread_id="t1", messages=["legacy"], agent_name="bob", user_id=None)
|
||||
queue.add(thread_id="t2", messages=["alice"], agent_name="bob", user_id="u1")
|
||||
queue.add(thread_id="t3", messages=["other"], agent_name="alice", user_id=None)
|
||||
|
||||
removed = queue.cancel_by_agent("bob", user_id=None)
|
||||
|
||||
assert removed == 1
|
||||
assert {(c.agent_name, c.user_id) for c in queue._items} == {("bob", "u1"), ("alice", None)}
|
||||
|
||||
|
||||
def test_cancel_all_agents_without_user_id_only_drops_legacy_scope() -> None:
|
||||
queue = _queue()
|
||||
with patch.object(queue, "_schedule_timer"):
|
||||
queue.add(thread_id="t1", messages=["legacy-a"], agent_name="a", user_id=None)
|
||||
queue.add(thread_id="t2", messages=["legacy-b"], agent_name="b", user_id=None)
|
||||
queue.add(thread_id="t3", messages=["named"], agent_name="a", user_id="u1")
|
||||
|
||||
removed = queue.cancel_by_agent(user_id=None, all_agents=True)
|
||||
|
||||
assert removed == 2
|
||||
assert queue.pending_count == 1
|
||||
assert queue._items[0].user_id == "u1"
|
||||
|
||||
|
||||
def test_cancel_by_agent_does_not_touch_in_flight_batch() -> None:
|
||||
"""Contexts already pulled out of `_items` keep running (#5037 residual)."""
|
||||
mock_updater = MagicMock()
|
||||
mock_updater.update_memory.return_value = True
|
||||
queue = _queue(mock_updater)
|
||||
with patch.object(queue, "_schedule_timer"):
|
||||
queue.add(thread_id="t1", messages=["in-flight"], agent_name="bob", user_id="u1")
|
||||
queue.add(thread_id="t2", messages=["still-pending"], agent_name="bob", user_id="u1")
|
||||
|
||||
with queue._lock:
|
||||
in_flight = queue._items[:1]
|
||||
queue._items = queue._items[1:]
|
||||
queue._processing = True
|
||||
|
||||
removed = queue.cancel_by_agent("bob", user_id="u1")
|
||||
assert removed == 1
|
||||
assert queue.pending_count == 0
|
||||
# Simulate the worker finishing the already-pulled context.
|
||||
mock_updater.update_memory(
|
||||
messages=in_flight[0].messages,
|
||||
thread_id=in_flight[0].thread_id,
|
||||
agent_name=in_flight[0].agent_name,
|
||||
signals=frozenset(),
|
||||
user_id=in_flight[0].user_id,
|
||||
trace_id=None,
|
||||
bypass_watermark=False,
|
||||
)
|
||||
mock_updater.update_memory.assert_called_once()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user