mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
perf(runtime): bound gateway memory after terminal runs (#5112)
* fix(runtime): clean up terminal run records * perf(sandbox): bound local path caches * perf(runtime): release terminal run cycles * fix(runtime): address terminal cleanup review * fix(runtime): clean up after end publish failure * fix(runtime): guard terminal cleanup from cancellation * fix(runtime): discard fenced journal buffers * fix(runtime): harden abort and teardown paths
This commit is contained in:
parent
a4f6665ef4
commit
45adb8fbb5
@ -330,6 +330,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
|
||||
> [!IMPORTANT]
|
||||
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a valid reconnect cursor has been trimmed, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination.
|
||||
>
|
||||
> After a run publishes its terminal stream marker, its process-local `RunRecord` remains available for the existing five-minute grace period before cleanup; durable run history remains available through `RunStore`, while the stream bridge retains its delivery tail on its separate cleanup schedule.
|
||||
>
|
||||
> Run cancellation may land on any Gateway worker. A non-owning worker now persists the interrupt or rollback request for the live owner, which observes it during lease renewal and performs the normal cancellation flow; load-balancer routing alone no longer produces a 409. The first accepted action wins even if a retry lands on the owner, and accepted cancellation competes atomically with owner completion. Dead owners still follow lease takeover and orphan recovery. Cancellation latency is therefore bounded by the lease heartbeat interval.
|
||||
>
|
||||
> With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation.
|
||||
|
||||
@ -57,6 +57,8 @@ The first `RunManager.list_by_thread()` hydration page uses a 100-row floor or
|
||||
the number of required IDs, whichever is larger; missing exact runs use targeted
|
||||
`get()` calls.
|
||||
|
||||
**Terminal run cleanup explicitly breaks graph-scoped references while preserving the existing `RunRecord` grace period.** Every `agent.astream()` iterator is closed in `_stream_once`, including abort/exception/early-break paths. A close failure after an abort is warning-only and cannot replace the user-requested `interrupted` outcome; normal-completion close failures still surface, and an in-flight stream exception remains authoritative over a secondary close failure. Journal construction and cancellable preflight work (including MCP task projection and the prior-finalization wait) live inside the worker's guarded body, so cancellation before agent startup still terminalizes the run and closes its stream. `run_agent()` wraps the complete terminal-finalization sequence in an outer teardown guard, so cancellation or failure from any terminal-stage await cannot skip `RunJournal.close()`, removal of the journal, `__pregel_runtime`, and internal runtime-context values from every runnable config, or release of local graph/payload references. That guard schedules bridge cleanup, run-record cleanup, and cyclic GC even when interruption happens before the terminal stream marker or terminal publication itself fails, so neither a cancelled observer nor a delivery-backend outage can strand process-local run state. `RunJournal.flush()` clears its `_pending_progress_task` after awaiting or cancelling it; ordinary `close()` detaches the event store/progress reporter and clears callback bookkeeping only after that flush succeeds, preserving the buffer for retry on a transient store failure. A fenced worker instead calls `close(flush=False)`, which cancels pending journal work and detaches without initiating another event-store write after lease ownership is lost; its final detach runs even if a second cancellation interrupts pending-task shutdown. `RunManager.cleanup(run_id)` retains the process-local `RunRecord`, completed task, and request payload for its default 300-second local join/status window before releasing them. Durable history remains in `RunStore`; `StreamBridge` data keeps its separate 60-second late-subscriber window, and both cleanup coroutines run in a fresh empty `contextvars.Context`. A contextless full cyclic-GC pass, coalesced to at most once every 10 seconds and dispatched through the default executor, bounds the lifetime of unreachable LangGraph callback/loop cycles without synchronously walking the heap in the event-loop timer; passes taking at least 100 ms are logged at INFO because CPython GC may still impose interpreter-level pauses.
|
||||
|
||||
**Where things live**:
|
||||
- `runtime/checkpoint_mode.py` — mode + snapshot-frequency freeze, marker injection, delta detection, compatibility gate, both error types
|
||||
- `runtime/checkpoint_state.py` — `CheckpointStateAccessor`, `build_state_mutation_graph`, `RollbackPoint`
|
||||
|
||||
@ -234,7 +234,8 @@ class RunJournal(BaseCallbackHandler):
|
||||
super().__init__()
|
||||
self.run_id = run_id
|
||||
self.thread_id = thread_id
|
||||
self._store = event_store
|
||||
self._store: RunEventStore | None = event_store
|
||||
self._closed = False
|
||||
self._track_tokens = track_token_usage
|
||||
self._flush_threshold = flush_threshold
|
||||
self._progress_reporter = progress_reporter
|
||||
@ -635,6 +636,8 @@ class RunJournal(BaseCallbackHandler):
|
||||
self._persist_tool_result_message(message)
|
||||
|
||||
def _put(self, *, event_type: str, category: str, content: str | dict = "", metadata: dict | None = None) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
self._buffer.append(
|
||||
{
|
||||
"thread_id": self.thread_id,
|
||||
@ -676,7 +679,10 @@ class RunJournal(BaseCallbackHandler):
|
||||
|
||||
async def _flush_async(self, batch: list[dict]) -> None:
|
||||
try:
|
||||
await self._store.put_batch(batch)
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
await store.put_batch(batch)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to flush %d events for run %s — returning to buffer",
|
||||
@ -886,26 +892,96 @@ class RunJournal(BaseCallbackHandler):
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Force flush remaining buffer. Called in worker's finally block."""
|
||||
if self._closed:
|
||||
return
|
||||
if self._pending_flush_tasks:
|
||||
await asyncio.gather(*tuple(self._pending_flush_tasks), return_exceptions=True)
|
||||
while self._pending_progress_task is not None and not self._pending_progress_task.done():
|
||||
while self._pending_progress_task is not None:
|
||||
pending_progress_task = self._pending_progress_task
|
||||
if pending_progress_task.done():
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
break
|
||||
if self._pending_progress_delayed:
|
||||
self._pending_progress_task.cancel()
|
||||
await asyncio.gather(self._pending_progress_task, return_exceptions=True)
|
||||
pending_progress_task.cancel()
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
self._progress_dirty = False
|
||||
self._pending_progress_delayed = False
|
||||
break
|
||||
await asyncio.gather(self._pending_progress_task, return_exceptions=True)
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
if self._pending_progress_task is pending_progress_task:
|
||||
self._pending_progress_task = None
|
||||
|
||||
while self._buffer:
|
||||
batch = self._buffer[: self._flush_threshold]
|
||||
del self._buffer[: self._flush_threshold]
|
||||
try:
|
||||
await self._store.put_batch(batch)
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
await store.put_batch(batch)
|
||||
except Exception:
|
||||
self._buffer = batch + self._buffer
|
||||
raise
|
||||
|
||||
def _detach_runtime_dependencies(self) -> None:
|
||||
"""Drop every external or potentially cyclic run-scoped reference."""
|
||||
self._closed = True
|
||||
self._store = None
|
||||
self._progress_reporter = None
|
||||
self._buffer.clear()
|
||||
self._pending_flush_tasks.clear()
|
||||
self._pending_progress_task = None
|
||||
self._pending_progress_delayed = False
|
||||
self._progress_dirty = False
|
||||
self._tokens_by_model.clear()
|
||||
self._counted_llm_run_ids.clear()
|
||||
self._counted_external_source_ids.clear()
|
||||
self._counted_message_llm_run_ids.clear()
|
||||
self._llm_start_times.clear()
|
||||
self._seen_llm_starts.clear()
|
||||
self._current_run_tool_call_names.clear()
|
||||
self._persisted_tool_message_identities.clear()
|
||||
self._produced_artifacts.clear()
|
||||
self._produced_artifact_keys.clear()
|
||||
self._last_ai_msg = None
|
||||
self._first_human_msg = None
|
||||
self._llm_error_fallback_message = None
|
||||
|
||||
async def close(self, *, flush: bool = True) -> None:
|
||||
"""Release run-scoped references, optionally flushing buffered events."""
|
||||
if self._closed:
|
||||
return
|
||||
if flush:
|
||||
# A failed terminal write returns its batch to ``_buffer``. Keep the
|
||||
# store and all buffered state attached so a later close/flush can retry
|
||||
# instead of silently discarding the tail of the run event stream.
|
||||
await self.flush()
|
||||
self._detach_runtime_dependencies()
|
||||
return
|
||||
|
||||
# A worker that lost its lease must detach without starting another
|
||||
# durable write. Drop dependencies before cancelling already-scheduled
|
||||
# work so tasks that have not begun observe the detached state. The
|
||||
# final detach must survive a second cancellation while those tasks stop.
|
||||
self._closed = True
|
||||
self._store = None
|
||||
self._progress_reporter = None
|
||||
try:
|
||||
pending_flush_tasks = tuple(self._pending_flush_tasks)
|
||||
for task in pending_flush_tasks:
|
||||
task.cancel()
|
||||
if pending_flush_tasks:
|
||||
await asyncio.gather(*pending_flush_tasks, return_exceptions=True)
|
||||
pending_progress_task = self._pending_progress_task
|
||||
if pending_progress_task is not None:
|
||||
pending_progress_task.cancel()
|
||||
await asyncio.gather(pending_progress_task, return_exceptions=True)
|
||||
finally:
|
||||
self._detach_runtime_dependencies()
|
||||
|
||||
def _schedule_progress_flush(self) -> None:
|
||||
"""Best-effort throttled progress snapshot for active run visibility."""
|
||||
if self._progress_reporter is None:
|
||||
|
||||
@ -17,14 +17,17 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import gc
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Coroutine
|
||||
from contextlib import asynccontextmanager
|
||||
from contextvars import Context
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
@ -92,6 +95,141 @@ logger = logging.getLogger(__name__)
|
||||
_checkpoint_locks_guard = threading.Lock()
|
||||
_checkpoint_locks_by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Lock]] = weakref.WeakKeyDictionary()
|
||||
|
||||
# Completed LangGraph runs can leave callback Contexts and AsyncPregelLoop
|
||||
# instances in unreachable reference cycles. They are collectable, but a busy
|
||||
# Gateway may promote those cycles into older GC generations faster than the
|
||||
# automatic collector revisits them, producing a rising post-GC heap floor.
|
||||
# Coalesce terminal full collections so the cycles have a bounded lifetime
|
||||
# without paying for a stop-the-world collection on every run.
|
||||
_TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS = 10.0
|
||||
_TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS = 0.1
|
||||
_terminal_cycle_collection_guard = threading.Lock()
|
||||
_terminal_cycle_collection_last_at = time.monotonic()
|
||||
_terminal_cycle_collection_scheduled_loops: weakref.WeakSet[asyncio.AbstractEventLoop] = weakref.WeakSet()
|
||||
|
||||
|
||||
def _create_contextless_task(coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
|
||||
"""Schedule terminal housekeeping without retaining the run's ContextVars."""
|
||||
return asyncio.create_task(coro, context=Context())
|
||||
|
||||
|
||||
def _schedule_terminal_cycle_collection() -> None:
|
||||
"""Coalesce full cyclic-GC passes after completed LangGraph runs."""
|
||||
loop = asyncio.get_running_loop()
|
||||
with _terminal_cycle_collection_guard:
|
||||
if loop in _terminal_cycle_collection_scheduled_loops:
|
||||
return
|
||||
elapsed = time.monotonic() - _terminal_cycle_collection_last_at
|
||||
delay = max(0.0, _TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS - elapsed)
|
||||
_terminal_cycle_collection_scheduled_loops.add(loop)
|
||||
|
||||
async def _collect() -> None:
|
||||
global _terminal_cycle_collection_last_at
|
||||
|
||||
try:
|
||||
with _terminal_cycle_collection_guard:
|
||||
now = time.monotonic()
|
||||
if now - _terminal_cycle_collection_last_at < _TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS:
|
||||
return
|
||||
_terminal_cycle_collection_last_at = now
|
||||
|
||||
started_at = time.perf_counter()
|
||||
# Do not run a heap walk synchronously from the event-loop timer.
|
||||
# CPython's collector can still contend for the GIL, so surface slow
|
||||
# passes below rather than claiming this removes every pause.
|
||||
collected = await loop.run_in_executor(None, gc.collect)
|
||||
duration = time.perf_counter() - started_at
|
||||
if duration >= _TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS:
|
||||
logger.info(
|
||||
"Terminal cyclic GC collected %d object(s) in %.3f seconds",
|
||||
collected,
|
||||
duration,
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"Terminal cyclic GC collected %d object(s) in %.3f seconds",
|
||||
collected,
|
||||
duration,
|
||||
)
|
||||
finally:
|
||||
with _terminal_cycle_collection_guard:
|
||||
_terminal_cycle_collection_scheduled_loops.discard(loop)
|
||||
|
||||
def _start_collection() -> None:
|
||||
_create_contextless_task(_collect())
|
||||
|
||||
# A blank Context prevents the timer itself from retaining the completed
|
||||
# run. The loop owns the TimerHandle; the WeakSet never keeps a test or
|
||||
# short-lived embedded-client event loop alive.
|
||||
loop.call_later(delay, _start_collection, context=Context())
|
||||
|
||||
|
||||
async def _close_agent_stream(stream: Any) -> None:
|
||||
"""Close a LangGraph stream deterministically after completion or early exit."""
|
||||
close = getattr(stream, "aclose", None)
|
||||
if close is None:
|
||||
return
|
||||
result = close()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
|
||||
def _remove_callback(config: dict[str, Any], handler: Any) -> None:
|
||||
callbacks = config.get("callbacks")
|
||||
if isinstance(callbacks, list):
|
||||
callbacks[:] = [callback for callback in callbacks if callback is not handler]
|
||||
return
|
||||
remove_handler = getattr(callbacks, "remove_handler", None)
|
||||
if callable(remove_handler):
|
||||
try:
|
||||
remove_handler(handler)
|
||||
except Exception:
|
||||
logger.debug("could not detach terminal callback", exc_info=True)
|
||||
|
||||
|
||||
def _release_run_scoped_references(
|
||||
configs: list[dict[str, Any]],
|
||||
runtime_context: dict[str, Any] | None,
|
||||
journal: Any | None,
|
||||
) -> None:
|
||||
"""Remove worker-owned graph references once durable finalization is done."""
|
||||
internal_context_keys = {
|
||||
"__run_journal",
|
||||
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
||||
}
|
||||
try:
|
||||
from deerflow.extensions import EXTENSION_SNAPSHOT_CONTEXT_KEY
|
||||
|
||||
internal_context_keys.add(EXTENSION_SNAPSHOT_CONTEXT_KEY)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from deerflow_extension_api import EXTENSION_TASK_STORE_KEY
|
||||
|
||||
internal_context_keys.add(EXTENSION_TASK_STORE_KEY)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
seen_configs: set[int] = set()
|
||||
handlers = [journal] if journal is not None else []
|
||||
for runnable_config in configs:
|
||||
if not isinstance(runnable_config, dict) or id(runnable_config) in seen_configs:
|
||||
continue
|
||||
seen_configs.add(id(runnable_config))
|
||||
configurable = runnable_config.get("configurable")
|
||||
if isinstance(configurable, dict):
|
||||
configurable.pop("__pregel_runtime", None)
|
||||
context = runnable_config.get("context")
|
||||
if isinstance(context, dict):
|
||||
for key in internal_context_keys:
|
||||
context.pop(key, None)
|
||||
for handler in handlers:
|
||||
_remove_callback(runnable_config, handler)
|
||||
|
||||
if isinstance(runtime_context, dict):
|
||||
for key in internal_context_keys:
|
||||
runtime_context.pop(key, None)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _checkpoint_thread_lock(thread_id: str) -> AsyncIterator[None]:
|
||||
@ -618,6 +756,11 @@ async def run_agent(
|
||||
accessor: CheckpointStateAccessor | None = None
|
||||
rollback_point: RollbackPoint | None = None
|
||||
journal = None
|
||||
runtime_ctx: dict[str, Any] | None = None
|
||||
runtime: Any | None = None
|
||||
agent: Any | None = None
|
||||
runnable_configs: list[dict[str, Any]] = [config]
|
||||
goal_evaluator_model: Any | None = None
|
||||
delivery_content: dict[str, Any] | None = None
|
||||
produced_output_paths: list[str] | None = None
|
||||
# Journal construction moved ahead of preflight so every terminal run can
|
||||
@ -633,20 +776,6 @@ async def run_agent(
|
||||
subagent_events: _SubagentEventBuffer | None = None
|
||||
started = False
|
||||
|
||||
if ctx.mcp_task_repo is not None and record.user_id is not None:
|
||||
try:
|
||||
task_rows = await ctx.mcp_task_repo.list_by_thread(
|
||||
thread_id,
|
||||
user_id=record.user_id,
|
||||
limit=20,
|
||||
)
|
||||
graph_input = {
|
||||
**graph_input,
|
||||
"background_tasks": _project_background_tasks(task_rows),
|
||||
}
|
||||
except Exception:
|
||||
logger.warning("Run %s: failed to project MCP task state", run_id, exc_info=True)
|
||||
|
||||
async def _finish_cancellation(
|
||||
action: str,
|
||||
*,
|
||||
@ -711,6 +840,22 @@ async def run_agent(
|
||||
progress_reporter=lambda snapshot: run_manager.update_run_progress(run_id, **snapshot),
|
||||
)
|
||||
|
||||
# Keep cancellable preflight work under the worker's terminal guard so
|
||||
# cancellation cannot strand a pending RunRecord or stream subscriber.
|
||||
if ctx.mcp_task_repo is not None and record.user_id is not None:
|
||||
try:
|
||||
task_rows = await ctx.mcp_task_repo.list_by_thread(
|
||||
thread_id,
|
||||
user_id=record.user_id,
|
||||
limit=20,
|
||||
)
|
||||
graph_input = {
|
||||
**graph_input,
|
||||
"background_tasks": _project_background_tasks(task_rows),
|
||||
}
|
||||
except Exception:
|
||||
logger.warning("Run %s: failed to project MCP task state", run_id, exc_info=True)
|
||||
|
||||
await run_manager.wait_for_prior_finalizing(
|
||||
thread_id,
|
||||
run_id,
|
||||
@ -868,6 +1013,7 @@ async def run_agent(
|
||||
# the agent name that this run will actually execute.
|
||||
config.setdefault("run_name", resolve_root_run_name(config, record.assistant_id))
|
||||
initial_runnable_config = RunnableConfig(**config)
|
||||
runnable_configs.append(initial_runnable_config)
|
||||
|
||||
def _continuation_runnable_config() -> RunnableConfig:
|
||||
continuation_config = dict(config)
|
||||
@ -876,7 +1022,9 @@ async def run_agent(
|
||||
configurable.pop("checkpoint_id", None)
|
||||
configurable.pop("checkpoint_map", None)
|
||||
continuation_config["configurable"] = configurable
|
||||
return RunnableConfig(**continuation_config)
|
||||
continuation = RunnableConfig(**continuation_config)
|
||||
runnable_configs.append(continuation)
|
||||
return continuation
|
||||
|
||||
agent_factory_kwargs: dict[str, Any] = {"config": initial_runnable_config}
|
||||
if ctx.app_config is not None and _agent_factory_supports_app_config(agent_factory):
|
||||
@ -933,6 +1081,7 @@ async def run_agent(
|
||||
# captured for rollback.
|
||||
pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(resumed_messages)})
|
||||
initial_runnable_config = RunnableConfig(**config)
|
||||
runnable_configs.append(initial_runnable_config)
|
||||
|
||||
runtime_ctx[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids)
|
||||
_install_runtime_context(config, runtime_ctx)
|
||||
@ -967,8 +1116,6 @@ async def run_agent(
|
||||
# the finally block so buffered steps survive abort/exception paths too.
|
||||
subagent_events = _SubagentEventBuffer(event_store, thread_id, run_id)
|
||||
|
||||
goal_evaluator_model: Any | None = None
|
||||
|
||||
def _get_goal_evaluator_model() -> Any:
|
||||
nonlocal goal_evaluator_model
|
||||
if goal_evaluator_model is None:
|
||||
@ -986,45 +1133,77 @@ async def run_agent(
|
||||
if len(lg_modes) == 1 and not stream_subgraphs:
|
||||
# Single mode, no subgraphs: astream yields raw chunks
|
||||
single_mode = lg_modes[0]
|
||||
async for chunk in agent.astream(input_payload, config=stream_config, stream_mode=single_mode):
|
||||
if record.abort_event.is_set():
|
||||
logger.info("Run %s abort requested — stopping", run_id)
|
||||
break
|
||||
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
||||
sse_event = _lg_mode_to_sse_event(single_mode)
|
||||
await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode))
|
||||
if single_mode == "custom":
|
||||
await subagent_events.add(chunk)
|
||||
stream = agent.astream(input_payload, config=stream_config, stream_mode=single_mode)
|
||||
broke_on_abort = False
|
||||
try:
|
||||
async for chunk in stream:
|
||||
if record.abort_event.is_set():
|
||||
broke_on_abort = True
|
||||
logger.info("Run %s abort requested — stopping", run_id)
|
||||
break
|
||||
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
||||
sse_event = _lg_mode_to_sse_event(single_mode)
|
||||
await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode))
|
||||
if single_mode == "custom":
|
||||
await subagent_events.add(chunk)
|
||||
finally:
|
||||
close_error = sys.exception()
|
||||
try:
|
||||
await _close_agent_stream(stream)
|
||||
except Exception:
|
||||
abort_requested = broke_on_abort or record.abort_event.is_set()
|
||||
if close_error is None and not abort_requested:
|
||||
raise
|
||||
if abort_requested:
|
||||
logger.warning("Could not close aborted agent stream for run %s", run_id, exc_info=True)
|
||||
else:
|
||||
logger.debug("Could not close agent stream for run %s", run_id, exc_info=True)
|
||||
return
|
||||
# Multiple modes or subgraphs: astream yields tuples
|
||||
async for item in agent.astream(
|
||||
stream = agent.astream(
|
||||
input_payload,
|
||||
config=stream_config,
|
||||
stream_mode=lg_modes,
|
||||
subgraphs=stream_subgraphs,
|
||||
):
|
||||
if record.abort_event.is_set():
|
||||
logger.info("Run %s abort requested — stopping", run_id)
|
||||
break
|
||||
)
|
||||
broke_on_abort = False
|
||||
try:
|
||||
async for item in stream:
|
||||
if record.abort_event.is_set():
|
||||
broke_on_abort = True
|
||||
logger.info("Run %s abort requested — stopping", run_id)
|
||||
break
|
||||
|
||||
mode, chunk, namespace = _unpack_stream_item(item, lg_modes, stream_subgraphs)
|
||||
if mode is None:
|
||||
continue
|
||||
mode, chunk, namespace = _unpack_stream_item(item, lg_modes, stream_subgraphs)
|
||||
if mode is None:
|
||||
continue
|
||||
|
||||
if not namespace:
|
||||
# Only root-graph frames may decide the parent run's error
|
||||
# fallback: a delegated subagent's marked fallback is the
|
||||
# executor's to map (task_failed), not this run's.
|
||||
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
||||
await _publish_stream_item(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
mode=mode,
|
||||
chunk=chunk,
|
||||
namespace=namespace,
|
||||
file_tool_chunk_batcher=file_tool_chunk_batcher,
|
||||
subagent_events=subagent_events,
|
||||
)
|
||||
if not namespace:
|
||||
# Only root-graph frames may decide the parent run's error
|
||||
# fallback: a delegated subagent's marked fallback is the
|
||||
# executor's to map (task_failed), not this run's.
|
||||
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
||||
await _publish_stream_item(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
mode=mode,
|
||||
chunk=chunk,
|
||||
namespace=namespace,
|
||||
file_tool_chunk_batcher=file_tool_chunk_batcher,
|
||||
subagent_events=subagent_events,
|
||||
)
|
||||
finally:
|
||||
close_error = sys.exception()
|
||||
try:
|
||||
await _close_agent_stream(stream)
|
||||
except Exception:
|
||||
abort_requested = broke_on_abort or record.abort_event.is_set()
|
||||
if close_error is None and not abort_requested:
|
||||
raise
|
||||
if abort_requested:
|
||||
logger.warning("Could not close aborted agent stream for run %s", run_id, exc_info=True)
|
||||
else:
|
||||
logger.debug("Could not close agent stream for run %s", run_id, exc_info=True)
|
||||
finally:
|
||||
stream_error = sys.exception()
|
||||
if file_tool_chunk_batcher is not None:
|
||||
@ -1145,219 +1324,257 @@ async def run_agent(
|
||||
)
|
||||
|
||||
finally:
|
||||
if record.ownership_lost:
|
||||
logger.warning(
|
||||
"Skipping durable finalization for run %s because this worker no longer owns its lease",
|
||||
run_id,
|
||||
)
|
||||
|
||||
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
|
||||
if not record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, True)
|
||||
try:
|
||||
if not checkpoint_rollback_completed:
|
||||
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
rollback_point=rollback_point,
|
||||
snapshot_capture_failed=snapshot_capture_failed,
|
||||
)
|
||||
if checkpoint_rollback_completed:
|
||||
await _publish_restored_checkpoint_values(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
accessor=accessor,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
||||
except Exception:
|
||||
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
|
||||
|
||||
# Persist any subagent step events still buffered (#3779) — including on
|
||||
# abort/exception paths, where the stream loop broke before its own flush.
|
||||
if not record.ownership_lost and subagent_events is not None:
|
||||
await subagent_events.flush()
|
||||
|
||||
if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None:
|
||||
try:
|
||||
await record_workspace_changes(
|
||||
event_store,
|
||||
thread_id,
|
||||
try:
|
||||
if record.ownership_lost:
|
||||
logger.warning(
|
||||
"Skipping durable finalization for run %s because this worker no longer owns its lease",
|
||||
run_id,
|
||||
pre_run_workspace_snapshot,
|
||||
user_id=workspace_changes_user_id,
|
||||
extra_excluded_dir_names=workspace_excluded_dir_names,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to record workspace changes for run %s", run_id, exc_info=True)
|
||||
|
||||
# Flush buffered journal events before the terminal receipt. The
|
||||
# receipt uses a run-scoped idempotent write shared with recovery, then
|
||||
# the staged terminal status is persisted. This ordering closes the
|
||||
# crash window where a terminal run could otherwise outlive its receipt.
|
||||
# A fenced worker leaves receipt recovery to the peer that claimed it.
|
||||
if not record.ownership_lost and journal is not None:
|
||||
try:
|
||||
await journal.flush()
|
||||
except Exception:
|
||||
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
|
||||
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
|
||||
if not record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, True)
|
||||
try:
|
||||
if not checkpoint_rollback_completed:
|
||||
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
||||
accessor=accessor,
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
rollback_point=rollback_point,
|
||||
snapshot_capture_failed=snapshot_capture_failed,
|
||||
)
|
||||
if checkpoint_rollback_completed:
|
||||
await _publish_restored_checkpoint_values(
|
||||
bridge=bridge,
|
||||
run_id=run_id,
|
||||
accessor=accessor,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
||||
except Exception:
|
||||
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
|
||||
|
||||
if delivery_content is None:
|
||||
if produced_output_paths is None:
|
||||
produced_output_paths = await _produced_output_paths(
|
||||
# Persist any subagent step events still buffered (#3779) — including on
|
||||
# abort/exception paths, where the stream loop broke before its own flush.
|
||||
if not record.ownership_lost and subagent_events is not None:
|
||||
await subagent_events.flush()
|
||||
|
||||
if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None:
|
||||
try:
|
||||
await record_workspace_changes(
|
||||
event_store,
|
||||
thread_id,
|
||||
run_id,
|
||||
pre_run_workspace_snapshot,
|
||||
thread_id=thread_id,
|
||||
user_id=workspace_changes_user_id,
|
||||
extra_excluded_dir_names=workspace_excluded_dir_names,
|
||||
)
|
||||
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
|
||||
receipt_persisted = await _persist_delivery_receipt(
|
||||
event_store,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
content=delivery_content,
|
||||
)
|
||||
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
|
||||
await run_manager.set_status(
|
||||
run_id,
|
||||
RunStatus.error,
|
||||
error=_DELIVERY_RECEIPT_FAILED_ERROR,
|
||||
persist=False,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to record workspace changes for run %s", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Advance the final completion fields and timestamp without
|
||||
# terminalizing the durable row. That active row continues to
|
||||
# fence peer checkpoint writers through the duration write.
|
||||
completion_data = journal.get_completion_data()
|
||||
await run_manager.update_finalizing_progress(run_id, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
|
||||
# Flush buffered journal events before the terminal receipt. The
|
||||
# receipt uses a run-scoped idempotent write shared with recovery, then
|
||||
# the staged terminal status is persisted. This ordering closes the
|
||||
# crash window where a terminal run could otherwise outlive its receipt.
|
||||
# A fenced worker leaves receipt recovery to the peer that claimed it.
|
||||
if not record.ownership_lost and journal is not None:
|
||||
try:
|
||||
await journal.flush()
|
||||
except Exception:
|
||||
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
|
||||
|
||||
# Keep the durable run row active through its final duration checkpoint
|
||||
# write. A peer Gateway admits history migration from the durable row,
|
||||
# not this worker's staged terminal status; terminalizing first would
|
||||
# let that migration read an unfinished lifetime and race this write.
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
|
||||
try:
|
||||
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
|
||||
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
|
||||
# Match legacy history semantics: turn_duration is the whole
|
||||
# RunRecord lifetime in integer seconds, including admission
|
||||
# delay. Persist zero for sub-second successful turns.
|
||||
duration = max(0, int((updated - created).total_seconds()))
|
||||
await _persist_run_duration(
|
||||
checkpointer=checkpointer,
|
||||
if delivery_content is None:
|
||||
if produced_output_paths is None:
|
||||
produced_output_paths = await _produced_output_paths(
|
||||
pre_run_workspace_snapshot,
|
||||
thread_id=thread_id,
|
||||
user_id=workspace_changes_user_id,
|
||||
extra_excluded_dir_names=workspace_excluded_dir_names,
|
||||
)
|
||||
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
|
||||
receipt_persisted = await _persist_delivery_receipt(
|
||||
event_store,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
duration_seconds=duration,
|
||||
content=delivery_content,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
|
||||
|
||||
if not record.ownership_lost and event_store is not None:
|
||||
try:
|
||||
# Even after bounded receipt retries are exhausted, persist the
|
||||
# real worker outcome. Leaving a successful row inflight would
|
||||
# let lease recovery rewrite it as an error with a synthetic
|
||||
# zero receipt.
|
||||
if record.abort_event.is_set():
|
||||
await run_manager.persist_current_status(run_id)
|
||||
else:
|
||||
cancel_action = await run_manager.set_status_if_not_cancelled(
|
||||
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
|
||||
await run_manager.set_status(
|
||||
run_id,
|
||||
record.status,
|
||||
error=record.error,
|
||||
stop_reason=record.stop_reason,
|
||||
RunStatus.error,
|
||||
error=_DELIVERY_RECEIPT_FAILED_ERROR,
|
||||
persist=False,
|
||||
)
|
||||
if cancel_action is not None:
|
||||
await _finish_cancellation(cancel_action)
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Advance the final completion fields and timestamp without
|
||||
# terminalizing the durable row. That active row continues to
|
||||
# fence peer checkpoint writers through the duration write.
|
||||
completion_data = journal.get_completion_data()
|
||||
await run_manager.update_finalizing_progress(run_id, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
# Keep the durable run row active through its final duration checkpoint
|
||||
# write. A peer Gateway admits history migration from the durable row,
|
||||
# not this worker's staged terminal status; terminalizing first would
|
||||
# let that migration read an unfinished lifetime and race this write.
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
|
||||
try:
|
||||
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
|
||||
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
|
||||
# Match legacy history semantics: turn_duration is the whole
|
||||
# RunRecord lifetime in integer seconds, including admission
|
||||
# delay. Persist zero for sub-second successful turns.
|
||||
duration = max(0, int((updated - created).total_seconds()))
|
||||
await _persist_run_duration(
|
||||
checkpointer=checkpointer,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
duration_seconds=duration,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
|
||||
|
||||
if not record.ownership_lost and event_store is not None:
|
||||
try:
|
||||
# Even after bounded receipt retries are exhausted, persist the
|
||||
# real worker outcome. Leaving a successful row inflight would
|
||||
# let lease recovery rewrite it as an error with a synthetic
|
||||
# zero receipt.
|
||||
if record.abort_event.is_set():
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
|
||||
else:
|
||||
cancel_action = await run_manager.set_status_if_not_cancelled(
|
||||
run_id,
|
||||
record.status,
|
||||
error=record.error,
|
||||
stop_reason=record.stop_reason,
|
||||
)
|
||||
if cancel_action is not None:
|
||||
await _finish_cancellation(cancel_action)
|
||||
await run_manager.persist_current_status(run_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
|
||||
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Persist token usage + convenience fields to RunStore
|
||||
completion_data = completion_data or journal.get_completion_data()
|
||||
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
|
||||
if not record.ownership_lost and journal is not None and persist_completion:
|
||||
try:
|
||||
# Persist token usage + convenience fields to RunStore
|
||||
completion_data = completion_data or journal.get_completion_data()
|
||||
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
|
||||
try:
|
||||
await run_manager.wait_for_prior_finalizing(thread_id, run_id)
|
||||
if not await run_manager.has_later_started_run(thread_id, run_id):
|
||||
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id=thread_id, app_config=ctx.app_config, graph_input=graph_input)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id)
|
||||
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
|
||||
try:
|
||||
await run_manager.wait_for_prior_finalizing(thread_id, run_id)
|
||||
if not await run_manager.has_later_started_run(thread_id, run_id):
|
||||
await _ensure_interrupted_title(checkpointer=checkpointer, thread_id=thread_id, app_config=ctx.app_config, graph_input=graph_input)
|
||||
except Exception:
|
||||
logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Sync title from checkpoint to threads_meta.display_name
|
||||
if started and not record.ownership_lost and checkpointer is not None and thread_store is not None:
|
||||
try:
|
||||
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
|
||||
if ckpt_tuple is not None:
|
||||
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
|
||||
title = ckpt.get("channel_values", {}).get("title")
|
||||
if title:
|
||||
await thread_store.update_display_name(thread_id, title)
|
||||
except Exception:
|
||||
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
|
||||
# Sync title from checkpoint to threads_meta.display_name
|
||||
if started and not record.ownership_lost and checkpointer is not None and thread_store is not None:
|
||||
try:
|
||||
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
|
||||
if ckpt_tuple is not None:
|
||||
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
|
||||
title = ckpt.get("channel_values", {}).get("title")
|
||||
if title:
|
||||
await thread_store.update_display_name(thread_id, title)
|
||||
except Exception:
|
||||
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
|
||||
|
||||
# Update threads_meta status based on run outcome
|
||||
if started and not record.ownership_lost and thread_store is not None:
|
||||
try:
|
||||
final_status = "idle" if record.status == RunStatus.success else record.status.value
|
||||
await thread_store.update_status(thread_id, final_status)
|
||||
except Exception:
|
||||
logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id)
|
||||
# Update threads_meta status based on run outcome
|
||||
if started and not record.ownership_lost and thread_store is not None:
|
||||
try:
|
||||
final_status = "idle" if record.status == RunStatus.success else record.status.value
|
||||
await thread_store.update_status(thread_id, final_status)
|
||||
except Exception:
|
||||
logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id)
|
||||
|
||||
if not record.ownership_lost and ctx.on_run_completed is not None:
|
||||
try:
|
||||
await ctx.on_run_completed(record)
|
||||
except Exception:
|
||||
logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True)
|
||||
if not record.ownership_lost and ctx.on_run_completed is not None:
|
||||
try:
|
||||
await ctx.on_run_completed(record)
|
||||
except Exception:
|
||||
logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True)
|
||||
|
||||
if task_info is not None and task_store is not None:
|
||||
# Keep the finalizing barrier held until stop observers finish, so
|
||||
# a same-thread replacement cannot overlap this task's lifecycle.
|
||||
if task_info is not None and task_store is not None:
|
||||
# Keep the finalizing barrier held until stop observers finish, so
|
||||
# a same-thread replacement cannot overlap this task's lifecycle.
|
||||
try:
|
||||
await notify_task_stop(
|
||||
extensions,
|
||||
task_store,
|
||||
task_info,
|
||||
lead_task_outcome(
|
||||
aborted=(record.abort_event.is_set() or record.status == RunStatus.interrupted),
|
||||
succeeded=record.status == RunStatus.success,
|
||||
),
|
||||
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Extension task-stop notification failed for run %s (non-fatal)",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
except BaseException as exc:
|
||||
# Cancellation here must not strand the finalizing barrier or
|
||||
# leave stream consumers waiting for the end frame.
|
||||
deferred_stop_interrupt = exc
|
||||
logger.warning(
|
||||
"Extension task-stop notification interrupted for run %s; completing cleanup first",
|
||||
run_id,
|
||||
)
|
||||
if record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, False)
|
||||
|
||||
await bridge.publish_end(run_id)
|
||||
|
||||
if deferred_stop_interrupt is not None:
|
||||
raise deferred_stop_interrupt
|
||||
finally:
|
||||
try:
|
||||
await notify_task_stop(
|
||||
extensions,
|
||||
task_store,
|
||||
task_info,
|
||||
lead_task_outcome(
|
||||
aborted=(record.abort_event.is_set() or record.status == RunStatus.interrupted),
|
||||
succeeded=record.status == RunStatus.success,
|
||||
),
|
||||
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
||||
if journal is not None:
|
||||
try:
|
||||
await journal.close(flush=not record.ownership_lost)
|
||||
except Exception:
|
||||
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
||||
finally:
|
||||
_release_run_scoped_references(
|
||||
runnable_configs,
|
||||
runtime_ctx,
|
||||
journal,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Extension task-stop notification failed for run %s (non-fatal)",
|
||||
run_id,
|
||||
exc_info=True,
|
||||
)
|
||||
except BaseException as exc:
|
||||
# Cancellation here must not strand the finalizing barrier or
|
||||
# leave stream consumers waiting for the end frame.
|
||||
deferred_stop_interrupt = exc
|
||||
logger.warning(
|
||||
"Extension task-stop notification interrupted for run %s; completing cleanup first",
|
||||
run_id,
|
||||
)
|
||||
if record.finalizing:
|
||||
await run_manager.set_finalizing(run_id, False)
|
||||
# Drop graph and per-run payload references before the terminal
|
||||
# worker task itself becomes collectable.
|
||||
agent = None
|
||||
accessor = None
|
||||
runtime = None
|
||||
runtime_ctx = None
|
||||
rollback_point = None
|
||||
subagent_events = None
|
||||
goal_evaluator_model = None
|
||||
task_store = None
|
||||
task_info = None
|
||||
pre_run_workspace_snapshot = None
|
||||
delivery_content = None
|
||||
produced_output_paths = None
|
||||
graph_input = {}
|
||||
|
||||
await bridge.publish_end(run_id)
|
||||
asyncio.create_task(bridge.cleanup(run_id, delay=60))
|
||||
|
||||
if deferred_stop_interrupt is not None:
|
||||
raise deferred_stop_interrupt
|
||||
# Durable finalization and terminal publication may depend on
|
||||
# external backends, but local housekeeping must always run.
|
||||
_create_contextless_task(bridge.cleanup(run_id, delay=60))
|
||||
# Preserve the existing five-minute grace period for local
|
||||
# join/status paths, then release the terminal record, completed
|
||||
# task, and request payload. Durable run history remains available
|
||||
# through RunStore.
|
||||
_create_contextless_task(run_manager.cleanup(run_id))
|
||||
_schedule_terminal_cycle_collection()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -6,7 +6,7 @@
|
||||
**Authorization gate** (`sandbox:execute`, RFC #4063 Phase 3): every sandbox-backed tool call passes through the gate in `deerflow/authz/sandbox_authz.py` - a binary `authorize(principal, "sandbox", "execute", target="*")` check before either reusing a persisted sandbox id or calling `provider.acquire`. Rechecking reuse is required because authorization config and user roles can change while the sandbox remains cached. Sync tool invocations call `authorize_sandbox_execution`; async tool invocations await `authorize_sandbox_execution_async` exactly once. A task-local `ContextVar` scopes that single decision across the complete composed tool invocation, including `ReadBeforeWriteMiddleware`'s pre-write inspection, tool body, and post-read mark; the value is copied into `asyncio.to_thread` workers. Authorization denial is converted to the normal error `ToolMessage` at the composed middleware boundary and is explicitly excluded from the gate's generic fail-open handlers. Async config loading and provider class discovery/import are offloaded before `aauthorize()` so reused sandbox calls do not hash config files or import custom modules on the event loop; provider construction remains on the running event loop because async providers may initialize loop-affine clients. The gate lives at the single tool initialization entry point (`ensure_sandbox_initialized` / `ensure_sandbox_initialized_async` in `tools.py`), while `SandboxMiddleware.before_agent` / `abefore_agent` apply the matching sync/async check to eager acquisition. Deny raises `SandboxAuthorizationError` (`sandbox/exceptions.py`), which propagates out of ordinary tool execution as a friendly error `ToolMessage` ("sandbox execution is not permitted for your role") - the eager path catches it and skips acquisition instead, deferring the deny to the first sandbox-touching tool call so both paths share the same semantics. Provider errors (authorization calls and provider resolution) follow `authorization.fail_closed` / `fail_open`; no readable `config.yaml` or `authorization.enabled: false` makes the gate a no-op (`safe_app_config` tolerates missing config). Gateway auxiliary sync paths (uploads/artifacts routers) call `try_acquire_sandbox_for_request` (`app/gateway/authz.py`), which gates via `authorize_sandbox_for_request` and skips the sync on deny - the upload/artifact edit itself still succeeds. Tests: `tests/test_sandbox_authorization.py` and `tests/blocking_io/test_sandbox_authorization.py`.
|
||||
**Environment policy** (`sandbox/env_policy.py`): `execute_command` no longer inherits the full `os.environ`. `build_sandbox_env()` scrubs secret-looking names (`*KEY*`/`*SECRET*`/`*TOKEN*`/`*PASS*`/`*CREDENTIAL*`) from the inherited environment before layering injected request secrets on top, so platform credentials (e.g. `OPENAI_API_KEY`) never leak into skill subprocesses. Benign vars (`PATH`, `HOME`, `LANG`, `VIRTUAL_ENV`, ...) are preserved.
|
||||
**Implementations**:
|
||||
- `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-user/thread `LocalSandbox` (id `local:{user_id}:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Shared runs use category mappings; a policy-scoped run replaces them with one `/mnt/skills` root mapping to the coherent thread view, so structured file tools resolve through one managed boundary. This is not a host filesystem security boundary: an enabled host `bash` subprocess can use canonical paths without `PathMapping`, so `supports_agent_skill_isolation` is dynamic and explicit Agent policies fail closed while host bash is enabled. On Windows, Git Bash/MSYS argument-conversion exclusions are limited to safe non-root virtual path prefixes; do not restore a blanket conversion disable, because host-native CLI launchers need normal MSYS path conversion for their own installation paths.
|
||||
- `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-user/thread `LocalSandbox` (id `local:{user_id}:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Shared runs use category mappings; a policy-scoped run replaces them with one `/mnt/skills` root mapping to the coherent thread view, so structured file tools resolve through one managed boundary. This is not a host filesystem security boundary: an enabled host `bash` subprocess can use canonical paths without `PathMapping`, so `supports_agent_skill_isolation` is dynamic and explicit Agent policies fail closed while host bash is enabled. Host-to-virtual output masking scans dynamic per-user/per-thread roots directly instead of compiling path-specific regexes, so evicted thread IDs do not remain in Python's global regex caches; a separate 256-entry root cache prevents repeated `realpath()` walks for every glob/grep match while bounding dynamic-path retention, and only the small process-stable skill/integration source set uses a bounded compiled cache. On Windows, Git Bash/MSYS argument-conversion exclusions are limited to safe non-root virtual path prefixes; do not restore a blanket conversion disable, because host-native CLI launchers need normal MSYS path conversion for their own installation paths.
|
||||
- `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker-based isolation. Active-cache and warm-pool entries are checked with the backend during acquire/reuse; definitively dead containers are dropped from all in-process maps so the thread can discover or create a fresh sandbox instead of reusing a stale client. Backend health-check failures are treated as unknown, not dead; local discovery likewise treats an unverifiable container as not adoptable and falls through to create rather than failing acquire. `get()` remains an in-memory lookup for event-loop-safe tool paths — it never touches the ownership store (that would be blocking IO on the event loop); ownership is published on acquire/reclaim and refreshed off the event loop by the dedicated renewal thread (`_renew_owned_leases`). `reset()` closes the per-instance acquire serializer so replacing the singleton cannot retain its executor workers; full remote sandbox teardown remains `shutdown()`. `uses_thread_data_mounts` defaults to backend detection (`LocalContainerBackend=True`, remote/provisioner backends=False), while the optional `sandbox.thread_data_mounts` boolean takes precedence for deployments that guarantee the Gateway and sandbox share the same thread user-data directories. Setting it `true` skips upload-time sandbox acquire/sync; a false positive leaves uploads unavailable to the sandbox. An explicit Agent policy uses four thread projection category mounts and a distinct deterministic sandbox identity, preventing reuse of an older container created with shared mounts. `skills.container_path` is a provider-startup snapshot shared by mount construction, sandbox identity, the remote Gateway request, and provisioner validation; custom roots are identity-scoped so a container or Pod created for one destination cannot be reused after the root changes. The Gateway and provisioner independently require one canonical absolute root that does not overlap reserved platform mounts, and both derive the four category allowlist entries from that root. The provisioner accepts all four category overrides; when all are present it suppresses the default hostPath or skills-PVC mount. With `USERDATA_PVC_NAME`, the thread projection categories use subpaths on that shared data PVC. Readiness probes and `agent_sandbox` clients classify loopback/private IPs, single-label cluster hosts, and Docker/Podman internal hostnames as direct control-plane destinations and set `trust_env=False`; external FQDNs and public IPs retain environment proxy support.
|
||||
- `E2BSandboxProvider` (`packages/harness/deerflow/community/e2b_sandbox/`) provides E2B remote isolation.
|
||||
New unrestricted sandboxes receive a one-shot upload from the enabled-only
|
||||
|
||||
@ -16,7 +16,7 @@ from typing import NamedTuple
|
||||
from deerflow.config.paths import VIRTUAL_PATH_PREFIX
|
||||
from deerflow.sandbox.env_policy import build_sandbox_env
|
||||
from deerflow.sandbox.local.list_dir import list_dir
|
||||
from deerflow.sandbox.path_patterns import build_output_mask_pattern
|
||||
from deerflow.sandbox.path_patterns import replace_output_path_matches
|
||||
from deerflow.sandbox.sandbox import Sandbox, _validate_extra_env
|
||||
from deerflow.sandbox.search import GrepMatch, find_glob_matches, find_grep_matches
|
||||
|
||||
@ -224,9 +224,9 @@ class LocalSandbox(Sandbox):
|
||||
self._agent_written_paths: set[str] = set()
|
||||
|
||||
# ``path_mappings`` is set once in ``__init__`` and never mutated, so the
|
||||
# sorted views and compiled path-rewrite patterns below are stable for the
|
||||
# sandbox's lifetime. Caching them avoids re-sorting and re-compiling these
|
||||
# regexes on every bash/read_file/write_file call (the agent's hot path).
|
||||
# sorted views and resolved roots below are stable for the sandbox's
|
||||
# lifetime. Caching them avoids repeated filesystem resolution and sorting
|
||||
# on every bash/read_file/write_file call (the agent's hot path).
|
||||
|
||||
@cached_property
|
||||
def _command_pattern(self) -> re.Pattern[str] | None:
|
||||
@ -248,31 +248,10 @@ class LocalSandbox(Sandbox):
|
||||
patterns = [re.escape(m.container_path) + r"(?=/|$|[^\w./-])(?:/[^\s\"';&|<>()]*)?" for m in mappings]
|
||||
return re.compile("|".join(f"({p})" for p in patterns))
|
||||
|
||||
@cached_property
|
||||
def _reverse_output_patterns(self) -> list[re.Pattern[str]]:
|
||||
"""Compiled matchers for local paths in command output (longest local path first)."""
|
||||
# The rule — segment boundary plus path tail — is owned by
|
||||
# ``deerflow.sandbox.path_patterns`` and shared with
|
||||
# ``sandbox.tools._compiled_mask_patterns``, the other site that rewrites host
|
||||
# paths back to virtual ones. Its rationale (why the boundary class is
|
||||
# text-oriented rather than shell-oriented like ``_command_pattern``, why ``$``
|
||||
# is load-bearing) lives with the owner rather than in a second copy here, which
|
||||
# is what let the two drift before (#4035 added the boundary here and missed
|
||||
# that site; #4053 added it there).
|
||||
#
|
||||
# What is specific to this site: without the boundary the regex yields the bare
|
||||
# root, which then *equals* the mount root and so satisfies
|
||||
# ``_reverse_resolve_path``'s own ``+ "/"`` guard — the sibling is rewritten to a
|
||||
# container path that forward resolution refuses to map back. And bases stay
|
||||
# separator-*sensitive*: they come from ``Path.resolve()`` and already carry the
|
||||
# platform's separator, so relaxing them would widen what this masks.
|
||||
return [build_output_mask_pattern(self._resolved_local_paths[m]) for m in self._mappings_by_local_specificity]
|
||||
|
||||
@cached_property
|
||||
def _resolved_local_paths(self) -> dict[PathMapping, str]:
|
||||
"""Filesystem-resolved local root per mapping. ``Path.resolve()`` hits the
|
||||
disk, and the mounted directories don't move, so resolve once and reuse."""
|
||||
return {m: str(Path(m.local_path).resolve()) for m in self.path_mappings}
|
||||
"""Filesystem-resolved local root per mapping, computed once."""
|
||||
return {m: os.path.realpath(m.local_path) for m in self.path_mappings}
|
||||
|
||||
@cached_property
|
||||
def _mappings_by_container_specificity(self) -> list[PathMapping]:
|
||||
@ -291,7 +270,7 @@ class LocalSandbox(Sandbox):
|
||||
mapping (i.e. the one whose local_path is the longest prefix of the
|
||||
resolved path), similar to how ``_resolve_path`` handles container paths.
|
||||
"""
|
||||
resolved = str(Path(resolved_path).resolve())
|
||||
resolved = os.path.realpath(resolved_path)
|
||||
|
||||
best_mapping: PathMapping | None = None
|
||||
best_prefix_len = -1
|
||||
@ -342,15 +321,16 @@ class LocalSandbox(Sandbox):
|
||||
return ResolvedPath(path_str, None)
|
||||
|
||||
mapping, relative = mapping_match
|
||||
local_root = Path(self._resolved_local_paths[mapping])
|
||||
resolved_path = (local_root / relative).resolve() if relative else local_root
|
||||
|
||||
local_root = self._resolved_local_paths[mapping]
|
||||
resolved_path = os.path.realpath(os.path.join(local_root, relative)) if relative else local_root
|
||||
try:
|
||||
resolved_path.relative_to(local_root)
|
||||
except ValueError as exc:
|
||||
raise PermissionError(errno.EACCES, "Access denied: path escapes mounted directory", path_str) from exc
|
||||
inside_root = os.path.normcase(os.path.commonpath([local_root, resolved_path])) == os.path.normcase(local_root)
|
||||
except ValueError:
|
||||
inside_root = False
|
||||
if not inside_root:
|
||||
raise PermissionError(errno.EACCES, "Access denied: path escapes mounted directory", path_str)
|
||||
|
||||
return ResolvedPath(str(resolved_path), mapping)
|
||||
return ResolvedPath(resolved_path, mapping)
|
||||
|
||||
def _resolve_path(self, path: str) -> str:
|
||||
return self._resolve_path_with_mapping(path).path
|
||||
@ -369,7 +349,7 @@ class LocalSandbox(Sandbox):
|
||||
Container path if mapping exists, otherwise original path
|
||||
"""
|
||||
normalized_path = path.replace("\\", "/")
|
||||
path_str = str(Path(normalized_path).resolve())
|
||||
path_str = os.path.realpath(normalized_path)
|
||||
|
||||
# Try each mapping (longest local path first for more specific matches)
|
||||
for mapping in self._mappings_by_local_specificity:
|
||||
@ -404,16 +384,16 @@ class LocalSandbox(Sandbox):
|
||||
Returns:
|
||||
Output with local paths resolved to container paths
|
||||
"""
|
||||
# Patterns are compiled once per sandbox (longest local path first for
|
||||
# correct prefix matching) and reused across calls.
|
||||
# Scan directly instead of compiling one regex per thread root. Python's
|
||||
# global regex caches outlive an evicted LocalSandbox and otherwise keep
|
||||
# high-cardinality thread paths resident.
|
||||
result = output
|
||||
for pattern in self._reverse_output_patterns:
|
||||
|
||||
def replace_match(match: re.Match) -> str:
|
||||
matched_path = match.group(0)
|
||||
return self._reverse_resolve_path(matched_path)
|
||||
|
||||
result = pattern.sub(replace_match, result)
|
||||
for mapping in self._mappings_by_local_specificity:
|
||||
result = replace_output_path_matches(
|
||||
result,
|
||||
self._resolved_local_paths[mapping],
|
||||
self._reverse_resolve_path,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
@ -786,7 +766,7 @@ class LocalSandbox(Sandbox):
|
||||
if "/" not in child_rel and mapping.container_path.rstrip("/") not in existing_dirs:
|
||||
# Verify the host path exists so we don't add phantom entries
|
||||
try:
|
||||
if Path(mapping.local_path).resolve().is_dir():
|
||||
if os.path.isdir(os.path.realpath(mapping.local_path)):
|
||||
result.append(f"{mapping.container_path}/")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@ -1,16 +1,15 @@
|
||||
"""Shared construction of the host→virtual output-masking regexes.
|
||||
"""Shared host→virtual output-path matching rules.
|
||||
|
||||
The boundary and tail are deliberately private: ``build_output_mask_pattern`` is
|
||||
the only supported way to spell this rule, so a third site cannot import the
|
||||
pieces and hand-roll a variant that drifts from the other two.
|
||||
The boundary and tail are deliberately private. Callers use either
|
||||
``build_output_mask_pattern`` for low-cardinality stable roots or
|
||||
``replace_output_path_matches`` for high-cardinality dynamic roots, so a third
|
||||
site cannot hand-roll a variant that drifts from the other two.
|
||||
|
||||
Two independent call sites rewrite host paths back to their virtual form in
|
||||
text that flows to the model: ``LocalSandbox._reverse_output_patterns`` (bash
|
||||
output) and ``sandbox.tools._compiled_mask_patterns`` (glob/grep/ls results).
|
||||
They must agree on where a host base is allowed to end, because both feed the
|
||||
same downstream contract — a match that stops short of a real segment boundary
|
||||
is rewritten to a container path that forward resolution then refuses to map
|
||||
back.
|
||||
text that flows to the model: ``LocalSandbox`` and ``sandbox.tools``. They must
|
||||
agree on where a host base is allowed to end, because both feed the same
|
||||
downstream contract — a match that stops short of a real segment boundary is
|
||||
rewritten to a container path that forward resolution then refuses to map back.
|
||||
|
||||
Keeping one copy of that rule per file is what let it drift: #4035 added the
|
||||
segment boundary to the reverse patterns and missed the masking patterns, and
|
||||
@ -24,6 +23,7 @@ The two sites are *not* identical, and the difference is deliberate — see
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
|
||||
# Only match where a host base ends at a real path-segment boundary, so a mount
|
||||
# root does not match inside a sibling that merely shares its prefix
|
||||
@ -43,6 +43,9 @@ _SEGMENT_BOUNDARY = r"(?=/|$|[^\w./-])"
|
||||
# path embedded in a larger line is not over-consumed.
|
||||
_PATH_TAIL = r"(?:[/\\][^\s\"';&|<>()]*)?"
|
||||
|
||||
_SEGMENT_BOUNDARY_CHAR = re.compile(r"[^\w./-]")
|
||||
_PATH_TAIL_TERMINATORS = frozenset("\"';&|<>()")
|
||||
|
||||
|
||||
def build_output_mask_pattern(base: str, *, separator_agnostic: bool = False) -> re.Pattern[str]:
|
||||
"""Compile the matcher for one host ``base`` in model-visible output.
|
||||
@ -54,9 +57,9 @@ def build_output_mask_pattern(base: str, *, separator_agnostic: bool = False) ->
|
||||
path with ``/``. ``sandbox.tools`` needs this because it derives its
|
||||
bases from ``_path_variants`` (which yields Windows-style spellings)
|
||||
and matches them against output whose separators it does not
|
||||
control. ``LocalSandbox`` does not: its bases come from
|
||||
``Path.resolve()``, so they already carry the running platform's
|
||||
separator, and relaxing them would widen what it masks.
|
||||
control. ``LocalSandbox`` does not: its bases come from filesystem
|
||||
resolution on the running platform, and relaxing them would widen
|
||||
what it masks.
|
||||
|
||||
Returns:
|
||||
A compiled pattern matching ``base`` at a segment boundary, plus an
|
||||
@ -66,3 +69,64 @@ def build_output_mask_pattern(base: str, *, separator_agnostic: bool = False) ->
|
||||
if separator_agnostic:
|
||||
escaped = escaped.replace(r"\\", r"[/\\]")
|
||||
return re.compile(escaped + _SEGMENT_BOUNDARY + _PATH_TAIL)
|
||||
|
||||
|
||||
def replace_output_path_matches(
|
||||
output: str,
|
||||
base: str,
|
||||
replacement: str | Callable[[str], str],
|
||||
*,
|
||||
separator_agnostic: bool = False,
|
||||
) -> str:
|
||||
"""Replace ``base`` path matches without compiling a path-specific regex.
|
||||
|
||||
Dynamic thread roots are high-cardinality. Compiling one regex per root
|
||||
leaves those roots in Python's global ``re`` caches after DeerFlow evicts
|
||||
the owning sandbox. This scanner preserves the same boundary and path-tail
|
||||
contract while keeping no process-level reference to ``base``.
|
||||
"""
|
||||
if not output or not base:
|
||||
return output
|
||||
|
||||
searchable_output = output.replace("\\", "/") if separator_agnostic and "\\" in output else output
|
||||
searchable_base = base.replace("\\", "/") if separator_agnostic and "\\" in base else base
|
||||
chunks: list[str] = []
|
||||
copied_until = 0
|
||||
search_from = 0
|
||||
|
||||
while True:
|
||||
match_start = searchable_output.find(searchable_base, search_from)
|
||||
if match_start < 0:
|
||||
break
|
||||
|
||||
base_end = match_start + len(searchable_base)
|
||||
match_end = base_end
|
||||
if base_end < len(searchable_output):
|
||||
next_char = searchable_output[base_end]
|
||||
if next_char in "/\\":
|
||||
match_end += 1
|
||||
while match_end < len(output):
|
||||
char = output[match_end]
|
||||
if char.isspace() or char in _PATH_TAIL_TERMINATORS:
|
||||
break
|
||||
match_end += 1
|
||||
elif _SEGMENT_BOUNDARY_CHAR.fullmatch(next_char) is None:
|
||||
search_from = match_start + 1
|
||||
continue
|
||||
|
||||
matched_path = output[match_start:match_end]
|
||||
if callable(replacement):
|
||||
replaced_path = replacement(matched_path)
|
||||
else:
|
||||
relative = matched_path[len(base) :].lstrip("/\\")
|
||||
replaced_path = f"{replacement}/{relative}" if relative else replacement
|
||||
|
||||
chunks.append(output[copied_until:match_start])
|
||||
chunks.append(replaced_path)
|
||||
copied_until = match_end
|
||||
search_from = match_end
|
||||
|
||||
if not chunks:
|
||||
return output
|
||||
chunks.append(output[copied_until:])
|
||||
return "".join(chunks)
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import ntpath
|
||||
import os
|
||||
import posixpath
|
||||
import re
|
||||
@ -32,7 +33,7 @@ from deerflow.sandbox.exceptions import (
|
||||
)
|
||||
from deerflow.sandbox.file_operation_lock import get_file_operation_lock
|
||||
from deerflow.sandbox.overwrite import unwrap_sandbox
|
||||
from deerflow.sandbox.path_patterns import build_output_mask_pattern
|
||||
from deerflow.sandbox.path_patterns import build_output_mask_pattern, replace_output_path_matches
|
||||
from deerflow.sandbox.sandbox import Sandbox
|
||||
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
||||
from deerflow.sandbox.search import GrepMatch
|
||||
@ -429,20 +430,17 @@ def _get_custom_mount_for_path(path: str):
|
||||
def _extract_thread_id_from_thread_data(thread_data: "ThreadDataState | None") -> str | None:
|
||||
"""Extract thread_id from thread_data by inspecting workspace_path.
|
||||
|
||||
The workspace_path has the form
|
||||
``{base_dir}/threads/{thread_id}/user-data/workspace``, so
|
||||
``Path(workspace_path).parent.parent.name`` yields the thread_id.
|
||||
The workspace path ends with
|
||||
``threads/{thread_id}/user-data/workspace``.
|
||||
"""
|
||||
if thread_data is None:
|
||||
return None
|
||||
workspace_path = thread_data.get("workspace_path")
|
||||
if not workspace_path:
|
||||
return None
|
||||
try:
|
||||
# {base_dir}/threads/{thread_id}/user-data/workspace → parent.parent = threads/{thread_id}
|
||||
return Path(workspace_path).parent.parent.name
|
||||
except Exception:
|
||||
return None
|
||||
normalized = workspace_path.replace("\\", "/").rstrip("/")
|
||||
parts = normalized.rsplit("/", 3)
|
||||
return parts[-3] if len(parts) == 4 and parts[-3] else None
|
||||
|
||||
|
||||
def _get_acp_workspace_host_path(thread_id: str | None = None) -> str | None:
|
||||
@ -640,6 +638,13 @@ def _join_path_preserving_style(base: str, relative: str) -> str:
|
||||
return f"{stripped_base}{separator}{normalized_relative}"
|
||||
|
||||
|
||||
def _path_parent(path: str) -> str:
|
||||
"""Return a lexical parent without interning high-cardinality path parts."""
|
||||
if "\\" in path and "/" not in path:
|
||||
return ntpath.dirname(path)
|
||||
return posixpath.dirname(path)
|
||||
|
||||
|
||||
def _sanitize_error(error: Exception, runtime: Runtime | None = None) -> str:
|
||||
"""Sanitize an error message to avoid leaking host filesystem paths.
|
||||
|
||||
@ -742,10 +747,10 @@ def _thread_virtual_to_actual_mappings(thread_data: ThreadDataState) -> dict[str
|
||||
mappings[f"{VIRTUAL_PATH_PREFIX}/outputs"] = outputs
|
||||
|
||||
# Also map the virtual root when all known dirs share the same parent.
|
||||
actual_dirs = [Path(p) for p in (workspace, uploads, outputs) if p]
|
||||
actual_dirs = [p for p in (workspace, uploads, outputs) if p]
|
||||
if actual_dirs:
|
||||
common_parent = str(Path(actual_dirs[0]).parent)
|
||||
if all(str(path.parent) == common_parent for path in actual_dirs):
|
||||
common_parent = _path_parent(actual_dirs[0])
|
||||
if all(_path_parent(path) == common_parent for path in actual_dirs):
|
||||
mappings[VIRTUAL_PATH_PREFIX] = common_parent
|
||||
|
||||
return mappings
|
||||
@ -756,37 +761,38 @@ def _thread_actual_to_virtual_mappings(thread_data: ThreadDataState) -> dict[str
|
||||
return {actual: virtual for virtual, actual in _thread_virtual_to_actual_mappings(thread_data).items()}
|
||||
|
||||
|
||||
@lru_cache(maxsize=512)
|
||||
def _compiled_mask_patterns(sources: tuple[tuple[str, str], ...]) -> tuple[tuple[re.Pattern[str], str, str], ...]:
|
||||
"""Compile the host→virtual masking patterns once per source set.
|
||||
@lru_cache(maxsize=256)
|
||||
def _mask_source_roots(host_base: str) -> tuple[str, ...]:
|
||||
"""Return lexical and filesystem-resolved spellings without ``pathlib``.
|
||||
|
||||
``sources`` is an ordered tuple of ``(host_base, virtual_base)`` pairs
|
||||
(skills, then ACP workspace, then per-thread user-data mappings sorted by
|
||||
host-path length, longest first). The patterns derive only from
|
||||
config-stable + per-thread inputs, so they're cached and reused instead of
|
||||
being rebuilt — ``re.escape`` + ``re.compile`` + ``Path.resolve`` (a
|
||||
syscall) — on every call. ``mask_local_paths_in_output`` runs once per
|
||||
glob/grep match, so without this the same patterns are recompiled per
|
||||
match.
|
||||
``PurePath`` interns every path component. That is useful for long-lived
|
||||
paths but wasteful for one-off thread IDs, because evicting DeerFlow's LRU
|
||||
does not shrink the interpreter's intern/allocator high-water mark. The
|
||||
bounded cache avoids repeating ``realpath`` walks for every glob/grep match
|
||||
without retaining an unbounded set of thread roots.
|
||||
"""
|
||||
# The segment boundary and path tail are shared with
|
||||
# ``LocalSandbox._reverse_output_patterns`` — see
|
||||
# ``deerflow.sandbox.path_patterns``, which owns that rule so the two copies
|
||||
# cannot drift again (#4035 fixed one and missed the other; #4053 fixed the
|
||||
# other).
|
||||
raw = os.path.normpath(host_base)
|
||||
resolved = os.path.realpath(host_base)
|
||||
return (raw,) if resolved == raw else (raw, resolved)
|
||||
|
||||
|
||||
@lru_cache(maxsize=16)
|
||||
def _compiled_mask_patterns(sources: tuple[tuple[str, str], ...]) -> tuple[tuple[re.Pattern[str], str, str], ...]:
|
||||
"""Compile patterns for the small, process-stable source set.
|
||||
|
||||
Per-user and per-thread sources must never be passed here; they are handled
|
||||
by the non-retaining scanner in ``replace_output_path_matches``.
|
||||
"""
|
||||
# The segment boundary and path tail are owned by
|
||||
# ``deerflow.sandbox.path_patterns`` so the static regex path and dynamic
|
||||
# scanner cannot drift.
|
||||
#
|
||||
# ``separator_agnostic=True`` is the one thing this site does differently:
|
||||
# its bases come from ``_path_variants``, which yields Windows-style
|
||||
# spellings, and they are matched against output whose separators this layer
|
||||
# does not control.
|
||||
# output separators are outside this layer's control.
|
||||
compiled: list[tuple[re.Pattern[str], str, str]] = []
|
||||
for host_base, virtual_base in sources:
|
||||
seen: set[str] = set()
|
||||
# Same base set as ``_path_variants(raw) | _path_variants(resolved)``;
|
||||
# ordered deterministically so the cached tuple is stable (variants of
|
||||
# one host map to the same virtual and don't overlap after substitution,
|
||||
# so order within a source is irrelevant to the result).
|
||||
for root in (str(Path(host_base)), str(Path(host_base).resolve())):
|
||||
for root in _mask_source_roots(host_base):
|
||||
for variant in sorted(_path_variants(root)):
|
||||
if variant in seen:
|
||||
continue
|
||||
@ -806,11 +812,12 @@ def mask_local_paths_in_output(output: str, thread_data: ThreadDataState | None)
|
||||
# custom/integration skills, then ACP workspace, then user-data mappings (longest
|
||||
# host path first). Custom mount host paths are masked by
|
||||
# LocalSandbox._reverse_resolve_paths_in_output().
|
||||
sources: list[tuple[str, str]] = []
|
||||
stable_sources: list[tuple[str, str]] = []
|
||||
dynamic_sources: list[tuple[str, str]] = []
|
||||
|
||||
skills_host = _get_skills_host_path()
|
||||
if skills_host:
|
||||
sources.append((skills_host, _get_skills_container_path()))
|
||||
stable_sources.append((skills_host, _get_skills_container_path()))
|
||||
|
||||
# Per-user custom skills: mask host paths under the user's custom
|
||||
# skills directory back to /mnt/skills/custom. The sandbox's
|
||||
@ -826,27 +833,28 @@ def mask_local_paths_in_output(output: str, thread_data: ThreadDataState | None)
|
||||
integrations_dir = get_paths().integration_skills_dir()
|
||||
if user_custom_dir.exists():
|
||||
skills_container = _get_skills_container_path()
|
||||
sources.append((str(user_custom_dir), f"{skills_container}/custom"))
|
||||
dynamic_sources.append((str(user_custom_dir), f"{skills_container}/custom"))
|
||||
if integrations_dir.exists():
|
||||
skills_container = _get_skills_container_path()
|
||||
sources.append((str(integrations_dir), f"{skills_container}/integrations"))
|
||||
stable_sources.append((str(integrations_dir), f"{skills_container}/integrations"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
acp_host = _get_acp_workspace_host_path(_extract_thread_id_from_thread_data(thread_data))
|
||||
if acp_host:
|
||||
sources.append((acp_host, _ACP_WORKSPACE_VIRTUAL_PATH))
|
||||
dynamic_sources.append((acp_host, _ACP_WORKSPACE_VIRTUAL_PATH))
|
||||
|
||||
if thread_data is not None:
|
||||
mappings = _thread_actual_to_virtual_mappings(thread_data)
|
||||
for actual_base, virtual_base in sorted(mappings.items(), key=lambda item: len(item[0]), reverse=True):
|
||||
sources.append((actual_base, virtual_base))
|
||||
dynamic_sources.append((actual_base, virtual_base))
|
||||
|
||||
if not sources:
|
||||
if not stable_sources and not dynamic_sources:
|
||||
return output
|
||||
|
||||
result = output
|
||||
for pattern, base, virtual in _compiled_mask_patterns(tuple(sources)):
|
||||
static_patterns = _compiled_mask_patterns(tuple(stable_sources)) if stable_sources else ()
|
||||
for pattern, base, virtual in static_patterns:
|
||||
|
||||
def replace_match(match: re.Match, _base: str = base, _virtual: str = virtual) -> str:
|
||||
matched_path = match.group(0)
|
||||
@ -857,6 +865,15 @@ def mask_local_paths_in_output(output: str, thread_data: ThreadDataState | None)
|
||||
|
||||
result = pattern.sub(replace_match, result)
|
||||
|
||||
for host_base, virtual_base in dynamic_sources:
|
||||
for root in _mask_source_roots(host_base):
|
||||
result = replace_output_path_matches(
|
||||
result,
|
||||
root,
|
||||
virtual_base,
|
||||
separator_agnostic=True,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
|
||||
@ -187,9 +187,11 @@ class _RunRecorder(_Recorder):
|
||||
class _OkAgent:
|
||||
def __init__(self) -> None:
|
||||
self.runtime_context = None
|
||||
self.runtime_task_store = None
|
||||
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
self.runtime_context = (config or {}).get("context")
|
||||
self.runtime_task_store = self.runtime_context.get(EXTENSION_TASK_STORE_KEY)
|
||||
yield {"messages": []}
|
||||
|
||||
|
||||
@ -231,7 +233,7 @@ async def test_run_agent_uses_the_run_bound_snapshot_for_lifecycle_and_task_stor
|
||||
("stop", record.run_id, "completed"),
|
||||
]
|
||||
assert recorder.start_infos[0].agent_name == "custom-agent"
|
||||
assert agent.runtime_context[EXTENSION_TASK_STORE_KEY] is recorder.start_stores[0]
|
||||
assert agent.runtime_task_store is recorder.start_stores[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@ -3089,6 +3089,7 @@ async def test_run_agent_full_mode_rejects_delta_before_graph_invocation():
|
||||
wait_for_prior_finalizing=AsyncMock(),
|
||||
set_status=set_status,
|
||||
set_status_if_not_cancelled=AsyncMock(side_effect=set_status_if_not_cancelled),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
record = RunRecord(
|
||||
run_id="run-checkpoint-mode",
|
||||
@ -3168,6 +3169,7 @@ async def test_run_agent_full_mode_checks_selected_checkpoint_before_graph():
|
||||
wait_for_prior_finalizing=AsyncMock(),
|
||||
set_status=set_status,
|
||||
set_status_if_not_cancelled=AsyncMock(side_effect=set_status_if_not_cancelled),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
record = RunRecord(
|
||||
run_id="run-selected-checkpoint-mode",
|
||||
|
||||
@ -637,6 +637,9 @@ async def test_run_agent_does_not_stream_continuation_after_abort(monkeypatch):
|
||||
async def set_finalizing(self, _run_id, finalizing):
|
||||
record.finalizing = finalizing
|
||||
|
||||
async def cleanup(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
class FakeBridge:
|
||||
async def publish(self, *_args, **_kwargs):
|
||||
return None
|
||||
@ -722,6 +725,9 @@ async def test_run_agent_reuses_goal_evaluator_model_for_goal_loop(monkeypatch):
|
||||
async def set_finalizing(self, _run_id, finalizing):
|
||||
record.finalizing = finalizing
|
||||
|
||||
async def cleanup(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
class FakeBridge:
|
||||
async def publish(self, *_args, **_kwargs):
|
||||
return None
|
||||
@ -906,6 +912,9 @@ async def test_run_agent_strips_branch_checkpoint_for_goal_continuation(monkeypa
|
||||
async def set_finalizing(self, _run_id, finalizing):
|
||||
record.finalizing = finalizing
|
||||
|
||||
async def cleanup(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
class FakeBridge:
|
||||
async def publish(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
@ -1,10 +1,10 @@
|
||||
"""Issue #3647 — LocalSandbox must compile its path-rewrite regexes once per
|
||||
sandbox (cached), not on every bash/read_file/write_file call, while keeping
|
||||
the exact same rewriting behavior.
|
||||
"""Issue #3647 — LocalSandbox caches stable forward-path data while dynamic
|
||||
host-path masking avoids process-global regex retention.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@ -32,16 +32,30 @@ def test_patterns_are_compiled_once_and_cached(tmp_path):
|
||||
# Each cached_property returns the identical object across accesses.
|
||||
assert sb._command_pattern is sb._command_pattern
|
||||
assert sb._content_pattern is sb._content_pattern
|
||||
assert sb._reverse_output_patterns is sb._reverse_output_patterns
|
||||
# Two mappings -> two reverse-output patterns.
|
||||
assert len(sb._reverse_output_patterns) == 2
|
||||
assert sb._resolved_local_paths is sb._resolved_local_paths
|
||||
|
||||
|
||||
def test_reverse_output_does_not_compile_regexes_for_unique_sandbox_paths(tmp_path):
|
||||
"""Evicted sandboxes must not leave their thread paths in ``re``'s caches."""
|
||||
resolved_tmp_path = str(tmp_path.resolve())
|
||||
for index in range(32):
|
||||
workspace = tmp_path / f"thread-{index}" / "workspace"
|
||||
workspace.mkdir(parents=True)
|
||||
sandbox = LocalSandbox(
|
||||
id=f"thread-{index}",
|
||||
path_mappings=[PathMapping(container_path="/mnt/user-data/workspace", local_path=str(workspace))],
|
||||
)
|
||||
|
||||
assert sandbox._reverse_resolve_paths_in_output(f"wrote {workspace}/result.txt") == "wrote /mnt/user-data/workspace/result.txt"
|
||||
|
||||
assert all(resolved_tmp_path not in str(cache_key) for cache_key in re._cache)
|
||||
|
||||
|
||||
def test_empty_mappings_yield_no_pattern(tmp_path):
|
||||
sb = LocalSandbox(id="empty", path_mappings=[])
|
||||
assert sb._command_pattern is None
|
||||
assert sb._content_pattern is None
|
||||
assert sb._reverse_output_patterns == []
|
||||
assert sb._resolved_local_paths == {}
|
||||
# No mappings -> command/content pass through unchanged.
|
||||
assert sb._resolve_paths_in_command("echo hello") == "echo hello"
|
||||
assert sb._resolve_paths_in_content("plain text") == "plain text"
|
||||
@ -147,15 +161,8 @@ def test_reverse_resolve_path_matches_windows_backslash_containment(monkeypatch)
|
||||
(real username, full directory tree) instead of the virtual
|
||||
``/mnt/user-data/...`` path.
|
||||
|
||||
CI runs only on ``ubuntu-latest`` (``os.sep == "/"``), where the pre-fix and
|
||||
post-fix code are observationally identical -- neither the hardcoded ``"/"``
|
||||
nor ``os.sep`` behave any differently there, so a test that just calls
|
||||
``_reverse_resolve_path`` on real POSIX paths cannot discriminate. To force
|
||||
the Windows code path independent of host OS, ``os.sep`` is monkeypatched to
|
||||
``"\\"`` and both the module's ``Path`` name and the sandbox's cached
|
||||
``_resolved_local_paths`` are stubbed to return backslash-joined strings --
|
||||
exactly what real ``WindowsPath.resolve()`` produces -- without touching the
|
||||
real filesystem or requiring an actual Windows host.
|
||||
CI runs only on POSIX, so the test stubs ``os.path.realpath`` and the cached
|
||||
roots to reproduce Windows' backslash-joined result.
|
||||
"""
|
||||
sb = LocalSandbox(
|
||||
id="windows-sep-test",
|
||||
@ -170,21 +177,7 @@ def test_reverse_resolve_path_matches_windows_backslash_containment(monkeypatch)
|
||||
# otherwise perform and pin it directly to the Windows-resolved root.
|
||||
sb._resolved_local_paths = {mapping: "C:\\Users\\test\\workspace"}
|
||||
|
||||
class _FakeWindowsPath:
|
||||
"""Stand-in for ``Path`` inside ``_reverse_resolve_path``. Mimics
|
||||
``WindowsPath.resolve()`` -- a backslash-joined ``str()`` -- without
|
||||
touching the real filesystem, so this runs identically on Linux CI."""
|
||||
|
||||
def __init__(self, raw: str) -> None:
|
||||
self._raw = raw
|
||||
|
||||
def resolve(self) -> _FakeWindowsPath:
|
||||
return _FakeWindowsPath(self._raw.replace("/", "\\"))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self._raw
|
||||
|
||||
monkeypatch.setattr(local_sandbox_module, "Path", _FakeWindowsPath)
|
||||
monkeypatch.setattr(local_sandbox_module.os.path, "realpath", lambda raw: raw.replace("/", "\\"))
|
||||
|
||||
result = sb._reverse_resolve_path("C:\\Users\\test\\workspace\\sub\\f.txt")
|
||||
|
||||
|
||||
@ -4,6 +4,7 @@ Uses MemoryRunEventStore as the backend for direct event inspection.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import weakref
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
@ -19,6 +20,124 @@ def test_run_journal_is_marked_as_loop_bound():
|
||||
assert RunJournal.deerflow_loop_bound is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_flushes_and_detaches_runtime_dependencies():
|
||||
class ProgressReporter:
|
||||
async def __call__(self, snapshot):
|
||||
del snapshot
|
||||
|
||||
store = MemoryRunEventStore()
|
||||
reporter = ProgressReporter()
|
||||
store_ref = weakref.ref(store)
|
||||
reporter_ref = weakref.ref(reporter)
|
||||
journal = RunJournal(
|
||||
"r-close",
|
||||
"t-close",
|
||||
store,
|
||||
progress_reporter=reporter,
|
||||
flush_threshold=100,
|
||||
)
|
||||
journal.record_middleware("test", name="test", hook="after", action="record", changes={})
|
||||
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._progress_reporter is None
|
||||
assert journal._buffer == []
|
||||
assert journal._pending_flush_tasks == set()
|
||||
del store, reporter
|
||||
await asyncio.sleep(0)
|
||||
assert store_ref() is None
|
||||
assert reporter_ref() is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_preserves_buffer_and_dependencies_when_flush_fails():
|
||||
class FailOnceRunEventStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.put_batch_calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
if self.put_batch_calls == 1:
|
||||
raise RuntimeError("transient store failure")
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = FailOnceRunEventStore()
|
||||
journal = RunJournal("r-close-retry", "t-close-retry", store, flush_threshold=100)
|
||||
journal.record_middleware("test", name="test", hook="after", action="record", changes={})
|
||||
|
||||
with pytest.raises(RuntimeError, match="transient store failure"):
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is False
|
||||
assert journal._store is store
|
||||
assert len(journal._buffer) == 1
|
||||
|
||||
await journal.close()
|
||||
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._buffer == []
|
||||
events = await store.list_events("t-close-retry", "r-close-retry")
|
||||
assert [event["event_type"] for event in events] == ["middleware:test"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_without_flush_discards_buffer_and_detaches_runtime_dependencies():
|
||||
class TrackingRunEventStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.put_batch_calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
return await super().put_batch(events)
|
||||
|
||||
store = TrackingRunEventStore()
|
||||
journal = RunJournal("r-close-discard", "t-close-discard", store, flush_threshold=100)
|
||||
journal.record_middleware("test", name="test", hook="after", action="record", changes={})
|
||||
|
||||
await journal.close(flush=False)
|
||||
|
||||
assert store.put_batch_calls == 0
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._buffer == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_close_without_flush_detaches_when_cancellation_interrupts_pending_task_cleanup():
|
||||
store = MemoryRunEventStore()
|
||||
journal = RunJournal("r-close-cancelled", "t-close-cancelled", store, flush_threshold=100)
|
||||
journal.record_middleware("test", name="test", hook="after", action="record", changes={})
|
||||
first_cancellation_seen = asyncio.Event()
|
||||
|
||||
async def stubborn_pending_flush() -> None:
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
first_cancellation_seen.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
pending_flush = asyncio.create_task(stubborn_pending_flush())
|
||||
journal._pending_flush_tasks.add(pending_flush)
|
||||
close_task = asyncio.create_task(journal.close(flush=False))
|
||||
await asyncio.wait_for(first_cancellation_seen.wait(), timeout=1)
|
||||
|
||||
close_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await close_task
|
||||
|
||||
assert pending_flush.done()
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._buffer == []
|
||||
assert journal._pending_flush_tasks == set()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def journal_setup():
|
||||
store = MemoryRunEventStore()
|
||||
@ -996,12 +1115,23 @@ class TestProgressSnapshots:
|
||||
parent_run_id=None,
|
||||
tags=["lead_agent"],
|
||||
)
|
||||
pending_task = j._pending_progress_task
|
||||
assert pending_task is not None
|
||||
pending_task_ref = weakref.ref(pending_task)
|
||||
|
||||
await asyncio.wait_for(j.flush(), timeout=0.2)
|
||||
|
||||
assert snapshots[-1]["total_tokens"] == 15
|
||||
assert snapshots[-1]["llm_call_count"] == 1
|
||||
assert snapshots[-1]["last_ai_message"] == "First"
|
||||
assert j._pending_progress_task is None
|
||||
|
||||
# The journal must not keep the cancelled task (and its traceback
|
||||
# frame) alive until cyclic GC. Dropping this last local reference
|
||||
# should release it immediately.
|
||||
del pending_task
|
||||
await asyncio.sleep(0)
|
||||
assert pending_task_ref() is None
|
||||
|
||||
|
||||
class TestChatModelStartHumanMessage:
|
||||
|
||||
@ -1,6 +1,10 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import logging
|
||||
import threading
|
||||
import weakref
|
||||
from contextlib import suppress
|
||||
from contextvars import ContextVar
|
||||
from types import SimpleNamespace
|
||||
from typing import Annotated, Any, NotRequired, TypedDict
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
@ -20,6 +24,7 @@ from deerflow.config.run_ownership_config import RunOwnershipConfig
|
||||
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
|
||||
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, RunManager
|
||||
from deerflow.runtime.runs.schemas import RunStatus
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
@ -48,6 +53,61 @@ class FakeCheckpointer:
|
||||
self.aput_writes = AsyncMock()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled():
|
||||
class CleanupTrackingRunManager(RunManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.cleanup_calls: list[tuple[str, float]] = []
|
||||
|
||||
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
|
||||
self.cleanup_calls.append((run_id, delay))
|
||||
|
||||
projection_started = asyncio.Event()
|
||||
|
||||
class BlockingTaskRepository:
|
||||
async def list_by_thread(self, thread_id, *, user_id, limit):
|
||||
del thread_id, user_id, limit
|
||||
projection_started.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
run_manager = CleanupTrackingRunManager()
|
||||
record = await run_manager.create("thread-mcp-projection-cancelled", user_id="alice")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
agent_factory = MagicMock(side_effect=AssertionError("cancelled preflight built the agent"))
|
||||
run_task = asyncio.create_task(
|
||||
run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(
|
||||
checkpointer=None,
|
||||
event_store=MemoryRunEventStore(),
|
||||
mcp_task_repo=BlockingTaskRepository(),
|
||||
),
|
||||
agent_factory=agent_factory,
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
)
|
||||
await asyncio.wait_for(projection_started.wait(), timeout=1)
|
||||
|
||||
run_task.cancel("MCP projection interrupted")
|
||||
await run_task
|
||||
await asyncio.sleep(0)
|
||||
|
||||
agent_factory.assert_not_called()
|
||||
assert record.status == RunStatus.interrupted
|
||||
assert record.finalizing is False
|
||||
bridge.publish_end.assert_awaited_once_with(record.run_id)
|
||||
bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
|
||||
assert run_manager.cleanup_calls == [(record.run_id, 300)]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pending_cancel_stops_waiting_for_prior_finalization():
|
||||
run_manager = RunManager()
|
||||
@ -678,6 +738,416 @@ async def test_run_agent_threads_explicit_app_config_into_config_only_factory():
|
||||
bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_schedules_terminal_run_record_cleanup():
|
||||
class CleanupTrackingRunManager(RunManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.cleanup_calls: list[tuple[str, float]] = []
|
||||
|
||||
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
|
||||
self.cleanup_calls.append((run_id, delay))
|
||||
|
||||
run_manager = CleanupTrackingRunManager()
|
||||
record = await run_manager.create("thread-terminal-cleanup")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
yield {"messages": []}
|
||||
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert run_manager.cleanup_calls == [(record.run_id, 300)]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_schedules_terminal_cleanup_when_publish_end_fails(monkeypatch):
|
||||
import deerflow.runtime.runs.worker as worker_module
|
||||
|
||||
class CleanupTrackingRunManager(RunManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.cleanup_calls: list[tuple[str, float]] = []
|
||||
|
||||
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
|
||||
self.cleanup_calls.append((run_id, delay))
|
||||
|
||||
run_manager = CleanupTrackingRunManager()
|
||||
record = await run_manager.create("thread-terminal-publish-failure")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(side_effect=RuntimeError("end publication unavailable")),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
schedule_collection = MagicMock()
|
||||
monkeypatch.setattr(worker_module, "_schedule_terminal_cycle_collection", schedule_collection)
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
yield {"messages": []}
|
||||
|
||||
with pytest.raises(RuntimeError, match="end publication unavailable"):
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
bridge.publish_end.assert_awaited_once_with(record.run_id)
|
||||
bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
|
||||
assert run_manager.cleanup_calls == [(record.run_id, 300)]
|
||||
schedule_collection.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_schedules_terminal_cleanup_when_completion_hook_is_cancelled(monkeypatch):
|
||||
import deerflow.runtime.runs.worker as worker_module
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
class CleanupTrackingRunManager(RunManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.cleanup_calls: list[tuple[str, float]] = []
|
||||
|
||||
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
|
||||
self.cleanup_calls.append((run_id, delay))
|
||||
|
||||
completion_hook_entered = asyncio.Event()
|
||||
|
||||
async def block_completion(_record) -> None:
|
||||
completion_hook_entered.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
run_manager = CleanupTrackingRunManager()
|
||||
record = await run_manager.create("thread-terminal-completion-cancelled")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
schedule_collection = MagicMock()
|
||||
monkeypatch.setattr(worker_module, "_schedule_terminal_cycle_collection", schedule_collection)
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, stream_mode, subgraphs
|
||||
callbacks = config.get("callbacks") or []
|
||||
captured["journal"] = next(callback for callback in callbacks if isinstance(callback, RunJournal))
|
||||
yield {"messages": []}
|
||||
|
||||
config: dict[str, Any] = {}
|
||||
run_task = asyncio.create_task(
|
||||
run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(
|
||||
checkpointer=None,
|
||||
event_store=MemoryRunEventStore(),
|
||||
on_run_completed=block_completion,
|
||||
),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config=config,
|
||||
)
|
||||
)
|
||||
await asyncio.wait_for(completion_hook_entered.wait(), timeout=1)
|
||||
run_task.cancel("completion hook interrupted")
|
||||
with pytest.raises(asyncio.CancelledError, match="completion hook interrupted"):
|
||||
await run_task
|
||||
await asyncio.sleep(0)
|
||||
|
||||
journal = captured["journal"]
|
||||
assert "__pregel_runtime" not in config["configurable"]
|
||||
assert journal not in config["callbacks"]
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60)
|
||||
assert run_manager.cleanup_calls == [(record.run_id, 300)]
|
||||
schedule_collection.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_closes_stream_when_abort_breaks_iteration():
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-stream-close")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
|
||||
class CloseTrackingStream:
|
||||
def __init__(self) -> None:
|
||||
self.yielded = False
|
||||
self.closed = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self.yielded:
|
||||
raise StopAsyncIteration
|
||||
self.yielded = True
|
||||
record.abort_event.set()
|
||||
return {"messages": []}
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
stream = CloseTrackingStream()
|
||||
|
||||
class DummyAgent:
|
||||
def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
return stream
|
||||
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
stream_modes=["values"],
|
||||
)
|
||||
|
||||
assert stream.closed is True
|
||||
assert record.status == RunStatus.interrupted
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]])
|
||||
@pytest.mark.parametrize("abort_before_break", [True, False], ids=["early-break", "exhaustion-race"])
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_ignores_stream_close_failure_after_abort(stream_modes, abort_before_break, caplog):
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-stream-close-failure")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
|
||||
class CloseFailingStream:
|
||||
def __init__(self) -> None:
|
||||
self.yielded = False
|
||||
self.close_attempted = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self.yielded:
|
||||
if not abort_before_break:
|
||||
record.abort_event.set()
|
||||
raise StopAsyncIteration
|
||||
self.yielded = True
|
||||
if abort_before_break:
|
||||
record.abort_event.set()
|
||||
chunk = {"messages": []}
|
||||
return chunk if len(stream_modes) == 1 else ("values", chunk)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.close_attempted = True
|
||||
raise RuntimeError("stream close failed")
|
||||
|
||||
stream = CloseFailingStream()
|
||||
|
||||
class DummyAgent:
|
||||
def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
return stream
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.runtime.runs.worker"):
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
stream_modes=stream_modes,
|
||||
)
|
||||
|
||||
assert stream.close_attempted is True
|
||||
assert record.status == RunStatus.interrupted
|
||||
assert record.error is None
|
||||
assert "Could not close aborted agent stream" in caplog.text
|
||||
bridge.publish_end.assert_awaited_once_with(record.run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_terminal_cleanup_tasks_do_not_inherit_run_context():
|
||||
marker: ContextVar[str | None] = ContextVar("run_cleanup_marker", default=None)
|
||||
seen: dict[str, str | None] = {}
|
||||
bridge_cleaned = asyncio.Event()
|
||||
manager_cleaned = asyncio.Event()
|
||||
|
||||
class CleanupTrackingRunManager(RunManager):
|
||||
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
|
||||
seen["manager"] = marker.get()
|
||||
await super().cleanup(run_id, delay=0)
|
||||
manager_cleaned.set()
|
||||
|
||||
class CleanupTrackingBridge:
|
||||
async def publish(self, *args, **kwargs) -> None:
|
||||
pass
|
||||
|
||||
async def publish_end(self, run_id: str) -> None:
|
||||
pass
|
||||
|
||||
async def cleanup(self, run_id: str, *, delay: float = 0) -> None:
|
||||
seen["bridge"] = marker.get()
|
||||
bridge_cleaned.set()
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, config, stream_mode, subgraphs
|
||||
yield {"messages": []}
|
||||
|
||||
run_manager = CleanupTrackingRunManager()
|
||||
record = await run_manager.create("thread-contextless-cleanup")
|
||||
token = marker.set("run-context")
|
||||
try:
|
||||
await run_agent(
|
||||
CleanupTrackingBridge(),
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(bridge_cleaned.wait(), manager_cleaned.wait()),
|
||||
timeout=1,
|
||||
)
|
||||
finally:
|
||||
marker.reset(token)
|
||||
|
||||
assert seen == {"bridge": None, "manager": None}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_terminal_cycle_collection_is_coalesced_contextless_and_off_loop(monkeypatch, caplog):
|
||||
import deerflow.runtime.runs.worker as worker_module
|
||||
|
||||
marker: ContextVar[str | None] = ContextVar("terminal_gc_marker", default=None)
|
||||
loop_thread_id = threading.get_ident()
|
||||
seen: list[tuple[str | None, int]] = []
|
||||
loop = asyncio.get_running_loop()
|
||||
collected = asyncio.Event()
|
||||
|
||||
def collect() -> int:
|
||||
seen.append((marker.get(), threading.get_ident()))
|
||||
loop.call_soon_threadsafe(collected.set)
|
||||
return 7
|
||||
|
||||
monkeypatch.setattr(worker_module, "_TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS", 0.0)
|
||||
monkeypatch.setattr(worker_module, "_TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS", 0.0)
|
||||
monkeypatch.setattr(worker_module, "_terminal_cycle_collection_last_at", 0.0)
|
||||
monkeypatch.setattr(worker_module.gc, "collect", collect)
|
||||
with worker_module._terminal_cycle_collection_guard:
|
||||
worker_module._terminal_cycle_collection_scheduled_loops.discard(loop)
|
||||
|
||||
caplog.set_level(logging.INFO, logger=worker_module.__name__)
|
||||
token = marker.set("run-context")
|
||||
try:
|
||||
worker_module._schedule_terminal_cycle_collection()
|
||||
worker_module._schedule_terminal_cycle_collection()
|
||||
await asyncio.wait_for(collected.wait(), timeout=1)
|
||||
|
||||
async def wait_until_finished() -> None:
|
||||
while True:
|
||||
with worker_module._terminal_cycle_collection_guard:
|
||||
if loop not in worker_module._terminal_cycle_collection_scheduled_loops:
|
||||
return
|
||||
await asyncio.sleep(0)
|
||||
|
||||
await asyncio.wait_for(wait_until_finished(), timeout=1)
|
||||
finally:
|
||||
marker.reset(token)
|
||||
|
||||
assert len(seen) == 1
|
||||
assert seen[0][0] is None
|
||||
assert seen[0][1] != loop_thread_id
|
||||
assert "Terminal cyclic GC collected 7 object(s)" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_releases_terminal_runtime_callbacks():
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-runtime-release")
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
del graph_input, stream_mode, subgraphs
|
||||
captured["config"] = config
|
||||
callbacks = config.get("callbacks") or []
|
||||
captured["journal"] = next(callback for callback in callbacks if isinstance(callback, RunJournal))
|
||||
yield {"messages": []}
|
||||
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(
|
||||
checkpointer=None,
|
||||
event_store=MemoryRunEventStore(),
|
||||
),
|
||||
agent_factory=lambda **_kwargs: DummyAgent(),
|
||||
graph_input={},
|
||||
config={},
|
||||
)
|
||||
|
||||
stream_config = captured["config"]
|
||||
journal = captured["journal"]
|
||||
assert "__pregel_runtime" not in stream_config["configurable"]
|
||||
assert "__run_journal" not in stream_config["context"]
|
||||
assert journal not in (stream_config.get("callbacks") or [])
|
||||
assert journal._closed is True
|
||||
assert journal._store is None
|
||||
assert journal._progress_reporter is None
|
||||
|
||||
journal_ref = weakref.ref(journal)
|
||||
captured.clear()
|
||||
del journal
|
||||
await asyncio.sleep(0)
|
||||
assert journal_ref() is None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context():
|
||||
run_manager = RunManager()
|
||||
@ -722,7 +1192,7 @@ async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context()
|
||||
)
|
||||
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
captured["context"] = config["context"]
|
||||
captured["pre_existing_message_ids"] = config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY]
|
||||
yield {"messages": []}
|
||||
|
||||
def factory(*, config):
|
||||
@ -738,8 +1208,7 @@ async def test_run_agent_threads_pre_existing_message_ids_into_runtime_context()
|
||||
config={},
|
||||
)
|
||||
|
||||
context = captured["context"]
|
||||
assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset({"h1", "a1"})
|
||||
assert captured["pre_existing_message_ids"] == frozenset({"h1", "a1"})
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@ -819,7 +1288,7 @@ async def test_run_agent_overrides_spoofed_pre_existing_message_ids_without_snap
|
||||
|
||||
class DummyAgent:
|
||||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||||
captured["context"] = config["context"]
|
||||
captured["pre_existing_message_ids"] = config["context"][CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY]
|
||||
yield {"messages": []}
|
||||
|
||||
def factory(*, config):
|
||||
@ -835,8 +1304,7 @@ async def test_run_agent_overrides_spoofed_pre_existing_message_ids_without_snap
|
||||
config={"context": {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}},
|
||||
)
|
||||
|
||||
context = captured["context"]
|
||||
assert context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] == frozenset()
|
||||
assert captured["pre_existing_message_ids"] == frozenset()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@ -3107,3 +3575,55 @@ async def test_worker_skips_execution_and_finalization_after_ownership_loss():
|
||||
thread_store.update_status.assert_not_awaited()
|
||||
on_run_completed.assert_not_awaited()
|
||||
bridge.publish_end.assert_awaited_once_with(record.run_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_worker_discards_buffered_journal_events_after_ownership_loss(monkeypatch):
|
||||
"""A fenced worker detaches its journal without appending buffered events."""
|
||||
|
||||
class TrackingRunEventStore(MemoryRunEventStore):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.put_batch_calls = 0
|
||||
|
||||
async def put_batch(self, events):
|
||||
self.put_batch_calls += 1
|
||||
return await super().put_batch(events)
|
||||
|
||||
journals: list[RunJournal] = []
|
||||
|
||||
class BufferedRunJournal(RunJournal):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.record_middleware("buffered", name="test", hook="after", action="record", changes={})
|
||||
journals.append(self)
|
||||
|
||||
monkeypatch.setattr("deerflow.runtime.journal.RunJournal", BufferedRunJournal)
|
||||
|
||||
event_store = TrackingRunEventStore()
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-lease-lost-buffered")
|
||||
record.ownership_lost = True
|
||||
record.abort_event.set()
|
||||
record.status = RunStatus.error
|
||||
bridge = SimpleNamespace(
|
||||
publish=AsyncMock(),
|
||||
publish_end=AsyncMock(),
|
||||
cleanup=AsyncMock(),
|
||||
)
|
||||
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
record,
|
||||
ctx=RunContext(checkpointer=None, event_store=event_store),
|
||||
agent_factory=MagicMock(side_effect=AssertionError("fenced worker started the agent")),
|
||||
graph_input={"messages": []},
|
||||
config={},
|
||||
)
|
||||
|
||||
assert event_store.put_batch_calls == 0
|
||||
assert len(journals) == 1
|
||||
assert journals[0]._closed is True
|
||||
assert journals[0]._store is None
|
||||
assert journals[0]._buffer == []
|
||||
|
||||
@ -20,6 +20,8 @@ from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.sandbox import path_patterns as path_patterns_module
|
||||
from deerflow.sandbox.local import local_sandbox as local_sandbox_module
|
||||
from deerflow.sandbox.local.local_sandbox import LocalSandbox, PathMapping
|
||||
from deerflow.sandbox.path_patterns import build_output_mask_pattern
|
||||
from deerflow.sandbox.tools import _compiled_mask_patterns
|
||||
@ -98,13 +100,42 @@ def test_boundary_still_rejects_prefix_siblings_and_accepts_real_segments() -> N
|
||||
assert pattern.search("/host/skills2/file.md") is None
|
||||
|
||||
|
||||
def test_local_sandbox_reverse_patterns_route_through_the_helper(tmp_path: Path) -> None:
|
||||
"""Call-site wiring: a re-inlined copy that *diverges* from the shared rule goes red.
|
||||
def test_direct_replacer_matches_the_shared_boundary_and_tail_contract() -> None:
|
||||
replacer = getattr(path_patterns_module, "replace_output_path_matches", None)
|
||||
assert replacer is not None
|
||||
|
||||
It does not (and cannot) catch a byte-identical re-inline — that is not yet a
|
||||
defect. What it catches is the shape of the actual regression: #4035 changed
|
||||
one copy of the rule and left the other behind.
|
||||
"""
|
||||
assert replacer("see /host/skills/a.md", "/host/skills", "/mnt/skills", separator_agnostic=True) == "see /mnt/skills/a.md"
|
||||
assert replacer("see \\host\\skills\\a.md", "/host/skills", "/mnt/skills", separator_agnostic=True) == "see /mnt/skills/a.md"
|
||||
assert replacer("see /host/skills-extra/a.md", "/host/skills", "/mnt/skills", separator_agnostic=True) == "see /host/skills-extra/a.md"
|
||||
assert replacer("root /host/skills, done", "/host/skills", "/mnt/skills", separator_agnostic=True) == "root /mnt/skills, done"
|
||||
|
||||
|
||||
def test_separator_agnostic_replacer_avoids_normalization_without_backslashes() -> None:
|
||||
class ReplaceTrackingString(str):
|
||||
def __init__(self, value: str) -> None:
|
||||
del value
|
||||
self.replace_calls = 0
|
||||
|
||||
def replace(self, old: str, new: str, count: int = -1) -> str:
|
||||
self.replace_calls += 1
|
||||
return super().replace(old, new, count)
|
||||
|
||||
output = ReplaceTrackingString("see /host/skills/a.md")
|
||||
base = ReplaceTrackingString("/host/skills")
|
||||
|
||||
result = path_patterns_module.replace_output_path_matches(
|
||||
output,
|
||||
base,
|
||||
"/mnt/skills",
|
||||
separator_agnostic=True,
|
||||
)
|
||||
|
||||
assert result == "see /mnt/skills/a.md"
|
||||
assert output.replace_calls == 0
|
||||
assert base.replace_calls == 0
|
||||
|
||||
|
||||
def test_local_sandbox_reverse_mask_routes_through_the_direct_helper(tmp_path: Path, monkeypatch) -> None:
|
||||
local = tmp_path / "skills"
|
||||
local.mkdir()
|
||||
sandbox = LocalSandbox(
|
||||
@ -113,7 +144,17 @@ def test_local_sandbox_reverse_patterns_route_through_the_helper(tmp_path: Path)
|
||||
)
|
||||
|
||||
resolved = str(Path(local).resolve())
|
||||
assert [p.pattern for p in sandbox._reverse_output_patterns] == [build_output_mask_pattern(resolved).pattern]
|
||||
calls: list[tuple[str, str]] = []
|
||||
original = path_patterns_module.replace_output_path_matches
|
||||
|
||||
def recording_replacer(output, base, replacement, **kwargs):
|
||||
calls.append((output, base))
|
||||
return original(output, base, replacement, **kwargs)
|
||||
|
||||
monkeypatch.setattr(local_sandbox_module, "replace_output_path_matches", recording_replacer)
|
||||
|
||||
assert sandbox._reverse_resolve_paths_in_output(f"read {resolved}/SKILL.md") == "read /mnt/skills/SKILL.md"
|
||||
assert calls == [(f"read {resolved}/SKILL.md", resolved)]
|
||||
|
||||
|
||||
def test_tools_mask_patterns_route_through_the_helper(tmp_path: Path) -> None:
|
||||
|
||||
@ -16,6 +16,7 @@ from deerflow.sandbox.tools import (
|
||||
_is_acp_workspace_path,
|
||||
_is_custom_mount_path,
|
||||
_is_skills_path,
|
||||
_mask_source_roots,
|
||||
_reject_path_traversal,
|
||||
_resolve_acp_workspace_path,
|
||||
_resolve_and_validate_user_data_path,
|
||||
@ -234,6 +235,55 @@ def test_mask_local_paths_compiled_patterns_are_cached() -> None:
|
||||
assert first is second # cache hit -> identical object, not rebuilt
|
||||
|
||||
|
||||
def test_mask_local_paths_cache_does_not_retain_per_thread_sources() -> None:
|
||||
"""Unique thread paths must not become process-lifetime regex-cache keys."""
|
||||
_compiled_mask_patterns.cache_clear()
|
||||
|
||||
with (
|
||||
patch("deerflow.sandbox.tools._get_skills_container_path", return_value="/mnt/skills"),
|
||||
patch("deerflow.sandbox.tools._get_skills_host_path", return_value="/srv/deer-flow/skills"),
|
||||
patch("deerflow.sandbox.tools._get_acp_workspace_host_path", return_value=None),
|
||||
):
|
||||
for index in range(32):
|
||||
root = f"/tmp/deer-flow/threads/thread-{index}/user-data"
|
||||
thread_data = {
|
||||
"workspace_path": f"{root}/workspace",
|
||||
"uploads_path": f"{root}/uploads",
|
||||
"outputs_path": f"{root}/outputs",
|
||||
}
|
||||
masked = mask_local_paths_in_output(f"created {root}/workspace/result.txt", thread_data)
|
||||
assert masked == "created /mnt/user-data/workspace/result.txt"
|
||||
|
||||
cache_info = _compiled_mask_patterns.cache_info()
|
||||
assert cache_info.currsize == 1
|
||||
assert cache_info.misses == 1
|
||||
|
||||
|
||||
def test_mask_local_paths_caches_dynamic_source_resolution_per_root() -> None:
|
||||
"""A batched glob/grep must not repeat ``realpath`` for every match."""
|
||||
_compiled_mask_patterns.cache_clear()
|
||||
_mask_source_roots.cache_clear()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch("deerflow.sandbox.tools._get_skills_container_path", return_value="/mnt/skills"),
|
||||
patch("deerflow.sandbox.tools._get_skills_host_path", return_value="/srv/deer-flow/skills"),
|
||||
patch("deerflow.sandbox.tools._get_acp_workspace_host_path", return_value=None),
|
||||
patch("deerflow.config.paths.get_paths", side_effect=RuntimeError("skip user paths")),
|
||||
patch("deerflow.sandbox.tools.os.path.realpath", side_effect=lambda path: path) as realpath,
|
||||
):
|
||||
for _ in range(200):
|
||||
masked = mask_local_paths_in_output("created /tmp/deer-flow/threads/t1/user-data/workspace/result.txt", _THREAD_DATA)
|
||||
assert masked == "created /mnt/user-data/workspace/result.txt"
|
||||
|
||||
# One stable skills root plus the workspace/uploads/outputs/thread roots.
|
||||
assert realpath.call_count == 5
|
||||
assert _mask_source_roots.cache_info().maxsize == 256
|
||||
finally:
|
||||
_compiled_mask_patterns.cache_clear()
|
||||
_mask_source_roots.cache_clear()
|
||||
|
||||
|
||||
def test_mask_local_paths_stable_across_repeated_and_batched_calls() -> None:
|
||||
"""Masking is identical whether applied once or repeatedly (per-match path)."""
|
||||
output = "a /tmp/deer-flow/threads/t1/user-data/workspace/x.txt and /tmp/deer-flow/threads/t1/user-data/outputs/y.log"
|
||||
|
||||
@ -67,6 +67,9 @@ class _FakeRunManager:
|
||||
async def update_run_completion(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
async def cleanup(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _FakeBridge:
|
||||
def __init__(self) -> None:
|
||||
|
||||
@ -342,6 +342,9 @@ class _IntegrationRunManager:
|
||||
async def set_finalizing(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
async def cleanup(self, *_args, **_kwargs):
|
||||
return None
|
||||
|
||||
|
||||
def _collect_ids(payload: object) -> set[str]:
|
||||
"""All string ``id`` values anywhere in a serialized stream payload."""
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user