diff --git a/src/xagent/web/api/websocket.py b/src/xagent/web/api/websocket.py index 727fcb35d3..c6e0c13816 100644 --- a/src/xagent/web/api/websocket.py +++ b/src/xagent/web/api/websocket.py @@ -3166,10 +3166,22 @@ def _finalize_resumed_task( result: Dict[str, Any], task_lease: TaskLease, prepared_outputs: _PreparedTaskFileOutputs, + turn_id: str | None = None, ) -> dict[str, Any]: - """Persist one fenced resumed result in a single worker transaction.""" + """Persist one fenced resumed result in a single worker transaction. + + ``turn_id``, when given, pops that turn's ephemeral connector secrets + (see ``task_orchestrator.finish_turn``'s matching handling) once this + exact call has committed a genuine COMPLETED/FAILED transition - not + when the lease release below is fenced out and the whole transaction + rolls back. A committed WAITING_FOR_USER/PAUSED transition instead + renews those secrets' TTL: it is the same turn resuming again later + under this same turn_id, now carrying a fresh interaction lifetime of + its own. + """ from ..models.agent import Agent from ..services.chat_history_service import persist_assistant_message_no_commit + from ..services.task_orchestrator import TERMINAL_TASK_STATUSES finalized: dict[str, Any] = { "task_title": None, @@ -3306,6 +3318,40 @@ def _finalize_resumed_task( finalized["lease_released"] = True finalized["final_status"] = final_task_status.value finalized["control_event_state"] = control_snapshot.as_dict() + if turn_id is not None: + if final_task_status in TERMINAL_TASK_STATUSES: + try: + from ..services.connector_runtime import ( + pop_ephemeral_runtime_values, + ) + + pop_ephemeral_runtime_values(turn_id) + except Exception: + logger.warning( + "connector runtime cleanup failed for task %s turn %s", + task_id, + turn_id, + exc_info=True, + ) + else: + # WAITING_FOR_USER/PAUSED: the same turn resuming again later + # under this same turn_id, with a fresh interaction lifetime - + # keep whatever ephemeral secrets it may still need from + # expiring on the original pause's clock (see + # connector_runtime.renew_ephemeral_runtime_values). + try: + from ..services.connector_runtime import ( + renew_ephemeral_runtime_values, + ) + + renew_ephemeral_runtime_values(turn_id) + except Exception: + logger.warning( + "connector runtime secret renewal failed for task %s turn %s", + task_id, + turn_id, + exc_info=True, + ) return finalized finally: try: @@ -3322,11 +3368,14 @@ def _settle_resumed_task_lease( lease: TaskLease, *, error_message: str | None, + turn_id: str | None = None, ) -> bool: """Delegate resume cleanup to the shared run/runner-fenced lifecycle.""" from ..services.task_orchestrator import settle_task_lease_isolated - return settle_task_lease_isolated(lease, error_message=error_message) + return settle_task_lease_isolated( + lease, error_message=error_message, turn_id=turn_id + ) async def execute_resume_background( @@ -3362,6 +3411,22 @@ async def execute_resume_background( if resume_owner_task is None: raise RuntimeError(f"Task {task_id} resume has no asyncio task") + # The turn id this task's ephemeral connector secrets (if any were ever + # stored) live under. A resume deliberately never rebinds this on the + # cached agent's tool_config (see WebToolConfig.invalidate_connector_ + # runtime_cache's docstring), so it is still the original pausing + # turn's id here - exactly what a terminal settlement below must pop + # under to actually free those secrets. + connector_runtime_turn_id: str | None = None + _resume_tool_config = getattr(agent_service, "tool_config", None) + _turn_id_getter = getattr( + _resume_tool_config, "get_connector_runtime_turn_id", None + ) + if callable(_turn_id_getter): + _candidate_turn_id = _turn_id_getter() + if isinstance(_candidate_turn_id, str): + connector_runtime_turn_id = _candidate_turn_id + lease_stop_event = preacquired_heartbeat_stop lease_heartbeat_task = preacquired_heartbeat_task lease: TaskLease | None = preacquired_lease @@ -3515,6 +3580,7 @@ async def mark_deferred_delivery_failed() -> bool: lambda acquired: _settle_resumed_task_lease( acquired, error_message="resume cancelled during lease acquisition", + turn_id=connector_runtime_turn_id, ), ) if prior_status_box: @@ -3815,6 +3881,7 @@ async def mark_deferred_delivery_failed() -> bool: result=result, task_lease=lease, prepared_outputs=outputs_for_finalizer, + turn_id=connector_runtime_turn_id, ) ) finally: @@ -4188,11 +4255,20 @@ async def finalize_resume_resources() -> None: and not lease_released and not defer_db_cleanup_to_ttl_recovery ): + # When this IS deferred, _finalize_resumed_task's + # turn_id-scoped pop never runs for this turn - and + # cannot safely run here either, for the same reason as + # task_orchestrator.py's matching branch: this coroutine + # does not know whether the task will land on a terminal + # status or resume again under this same turn_id. + # connector_runtime.py's _EPHEMERAL_RUNTIME_TTL_SECONDS + # bounds that leak instead. try: settled = await run_db_io_cancellation_safe( lambda: _settle_resumed_task_lease( lease, error_message=settlement_error, + turn_id=connector_runtime_turn_id, ) ) if settled: diff --git a/src/xagent/web/services/connector_runtime.py b/src/xagent/web/services/connector_runtime.py index 1165ea3565..f9d59fe290 100644 --- a/src/xagent/web/services/connector_runtime.py +++ b/src/xagent/web/services/connector_runtime.py @@ -9,6 +9,7 @@ import json import logging +import time from collections.abc import Callable, Collection from dataclasses import dataclass from threading import RLock @@ -45,6 +46,7 @@ from ..models.custom_api import CustomApi, UserCustomApi from ..models.mcp import MCPServer, UserMCPServer from ..models.task import Task, TaskConnectorRuntimeContext +from .task_interaction_service import _MAX_INTERACTION_TTL_SECONDS logger = logging.getLogger(__name__) @@ -112,9 +114,57 @@ class _ConnectorRuntimeResolverRegistration: # through the resolver hook or a deployment-owned distributed secret store. _EPHEMERAL_RUNTIME_VALUES: dict[str, dict[str, Any]] = {} _EPHEMERAL_RUNTIME_MANIFESTS: dict[str, dict[str, dict[str, set[str]]]] = {} +_EPHEMERAL_RUNTIME_STORED_AT: dict[str, float] = {} _EPHEMERAL_RUNTIME_VALUES_LOCK = RLock() _RUNTIME_RESOLVER_REGISTRATION: _ConnectorRuntimeResolverRegistration | None = None +# Not a real interaction deadline - matches task_interaction_service.py's +# _MAX_INTERACTION_TTL_SECONDS by reference (not by copying the value) so +# this can never silently drift shorter than the longest pause production +# already allows a native waiting-for-user interaction to stay open for. +# Its only job is to turn what would otherwise be a permanent leak into a +# bounded one: task_orchestrator.py's and websocket.py's deferred-to-TTL- +# recovery branches (lease lost, DB pool exhaustion, unhealthy heartbeat at +# shutdown) cannot safely pop a turn's secrets themselves, because at the +# point they bail out they genuinely do not know whether the task will land +# on a terminal status or resume again under the same turn_id - and +# task_lease_recovery.py's batch sweep, which does later decide that, +# has no way to look up which turn_id belongs to which recovered task_id. +_EPHEMERAL_RUNTIME_TTL_SECONDS = _MAX_INTERACTION_TTL_SECONDS + + +def _prune_expired_ephemeral_runtime_values_locked() -> None: + """Evict every stale entry. Caller must hold _EPHEMERAL_RUNTIME_VALUES_LOCK.""" + + now = time.monotonic() + stale_turn_ids = [ + turn_id + for turn_id, stored_at in _EPHEMERAL_RUNTIME_STORED_AT.items() + if now - stored_at > _EPHEMERAL_RUNTIME_TTL_SECONDS + ] + for turn_id in stale_turn_ids: + _evict_if_expired_locked(turn_id) + + +def _evict_if_expired_locked(turn_id: str) -> None: + """Evict one entry if it has aged past the TTL. Caller must hold the lock. + + The opportunistic reaper in ``store_ephemeral_runtime_values`` only runs + when *some* turn writes - a quiet process (no new turns starting) never + calls it, so an entry whose owning turn stalled out would otherwise stay + readable forever past its advertised TTL. Every read/pop path below + checks age itself first, so "expired" is enforced as an observable + property of this store, not merely a future cleanup side effect. + """ + + stored_at = _EPHEMERAL_RUNTIME_STORED_AT.get(turn_id) + if stored_at is not None and time.monotonic() - stored_at > ( + _EPHEMERAL_RUNTIME_TTL_SECONDS + ): + _EPHEMERAL_RUNTIME_VALUES.pop(turn_id, None) + _EPHEMERAL_RUNTIME_MANIFESTS.pop(turn_id, None) + _EPHEMERAL_RUNTIME_STORED_AT.pop(turn_id, None) + def set_connector_runtime_resolver( resolver: ConnectorRuntimeResolver | None, @@ -200,16 +250,42 @@ def store_ephemeral_runtime_values( for ref, sections in values_by_ref.items() } with _EPHEMERAL_RUNTIME_VALUES_LOCK: + # Opportunistic reaper: there is no dedicated background sweep for + # this process-local store, so each new store is what eventually + # reclaims an entry a settlement path could never safely pop (see + # _EPHEMERAL_RUNTIME_TTL_SECONDS above). + _prune_expired_ephemeral_runtime_values_locked() _EPHEMERAL_RUNTIME_VALUES[turn_id] = encoded _EPHEMERAL_RUNTIME_MANIFESTS[turn_id] = manifest + _EPHEMERAL_RUNTIME_STORED_AT[turn_id] = time.monotonic() def pop_ephemeral_runtime_values(turn_id: str) -> dict[str, Any] | None: with _EPHEMERAL_RUNTIME_VALUES_LOCK: + _evict_if_expired_locked(turn_id) _EPHEMERAL_RUNTIME_MANIFESTS.pop(turn_id, None) + _EPHEMERAL_RUNTIME_STORED_AT.pop(turn_id, None) return _EPHEMERAL_RUNTIME_VALUES.pop(turn_id, None) +def renew_ephemeral_runtime_values(turn_id: str) -> None: + """Refresh a still-live turn's stored-at timestamp, extending its TTL. + + Called when a turn is confirmed to be re-pausing under the same + ``turn_id`` (see ``task_orchestrator.finish_turn`` and + ``websocket._finalize_resumed_task``'s matching non-terminal branches): + a fresh pause carries its own new interaction lifetime, so the secrets + it may still need must not expire on the *original* pause's clock. A + no-op when the turn has nothing stored (never used ephemeral secrets, or + already reaped) - there is nothing to keep alive either way. + """ + + with _EPHEMERAL_RUNTIME_VALUES_LOCK: + _evict_if_expired_locked(turn_id) + if turn_id in _EPHEMERAL_RUNTIME_VALUES: + _EPHEMERAL_RUNTIME_STORED_AT[turn_id] = time.monotonic() + + def drop_ephemeral_runtime_values_for_testing(turn_id: str) -> None: """Simulate losing the secret values while keeping safe provenance.""" @@ -219,6 +295,7 @@ def drop_ephemeral_runtime_values_for_testing(turn_id: str) -> None: def get_ephemeral_runtime_values(turn_id: str) -> dict[str, Any] | None: with _EPHEMERAL_RUNTIME_VALUES_LOCK: + _evict_if_expired_locked(turn_id) values = _EPHEMERAL_RUNTIME_VALUES.get(turn_id) return dict(values) if isinstance(values, dict) else None @@ -227,6 +304,7 @@ def get_ephemeral_runtime_manifest( turn_id: str, ) -> dict[str, dict[str, set[str]]] | None: with _EPHEMERAL_RUNTIME_VALUES_LOCK: + _evict_if_expired_locked(turn_id) manifest = _EPHEMERAL_RUNTIME_MANIFESTS.get(turn_id) if not isinstance(manifest, dict): return None diff --git a/src/xagent/web/services/task_orchestrator.py b/src/xagent/web/services/task_orchestrator.py index e133ead477..1e42eec1f2 100644 --- a/src/xagent/web/services/task_orchestrator.py +++ b/src/xagent/web/services/task_orchestrator.py @@ -132,6 +132,17 @@ TaskStatus.PAUSED, ) +# A turn's outcome that ends here for good, as opposed to WAITING_FOR_USER/ +# PAUSED (the same turn resuming again later under the same turn_id) - the +# single source of truth for "is this outcome terminal" so a future status +# added to the finished/not-finished distinction only needs one edit. Every +# ``finish_turn`` branch that pops a turn's ephemeral connector secrets +# (COMPLETED, FAILED, RUNNING-fallback) commits into this set by +# construction; ``_finalize_resumed_task`` (websocket.py), the separate +# finalizer for the resume path, checks its own computed outcome against it +# directly since resume has no matching branch structure to fall out of. +TERMINAL_TASK_STATUSES = frozenset({TaskStatus.COMPLETED, TaskStatus.FAILED}) + def timezone_schedule_context(timezone: str | None) -> dict[str, Any] | None: """Build the schedule ``context`` carrying the caller's clock timezone. @@ -755,6 +766,7 @@ async def _schedule_committed_turn( lambda: settle_task_lease_isolated( claimed.task_lease, error_message="turn scheduling failed after claim commit", + turn_id=payload.turn_id, ) ) except Exception as terminal_error: @@ -1429,12 +1441,50 @@ def _get_agent_manager() -> Any: return get_agent_manager() +def _pop_ephemeral_runtime_values_best_effort(turn_id: str) -> None: + """Pop one turn's ephemeral connector secrets; never raise into a settler. + + Called only from a branch that has just committed (or reconciled) a + genuinely terminal outcome for this exact turn_id - see the callers in + ``finish_turn`` and ``settle_task_lease_isolated`` for why each call site + is safe to pop from. + """ + from .connector_runtime import pop_ephemeral_runtime_values + + try: + pop_ephemeral_runtime_values(turn_id) + except Exception: + logger.warning( + "connector runtime cleanup failed for turn %s", turn_id, exc_info=True + ) + + +def _renew_ephemeral_runtime_values_best_effort(turn_id: str) -> None: + """Renew one turn's ephemeral connector secrets; never raise into a settler. + + Called only from a branch that has just committed a genuine + PAUSED/WAITING_FOR_USER outcome for this exact turn_id - the same turn + resuming again later, not a finished one. + """ + from .connector_runtime import renew_ephemeral_runtime_values + + try: + renew_ephemeral_runtime_values(turn_id) + except Exception: + logger.warning( + "connector runtime secret renewal failed for turn %s", + turn_id, + exc_info=True, + ) + + def settle_task_lease_isolated( lease: TaskLease, *, error_message: str | None = None, client_error_message: str = CLIENT_SAFE_TASK_FAILURE, client_message_type: str = TASK_FAILURE_MESSAGE_TYPE, + turn_id: str | None = None, ) -> bool: """Settle exactly one run/runner lease in one worker-owned Session. @@ -1453,6 +1503,13 @@ def settle_task_lease_isolated( On checkout or commit failure the transaction is rolled back and the lease is intentionally retained for TTL recovery; this function never creates an ownerless RUNNING task. + + ``turn_id``, when given, pops that turn's ephemeral connector secrets the + moment this call itself commits a genuine FAILED transition - using the + same row this decision is made from, not a separate later read (see + ``finish_turn``'s matching ``turn_id`` handling for the reconciliation + path, which is where every other outcome, including this one when + ``error_message`` is ``None``, gets the same treatment). """ from ..models.database import get_session_local from .chat_history_service import persist_assistant_message_no_commit @@ -1481,6 +1538,8 @@ def settle_task_lease_isolated( ) settle_db.commit() invalidate_task_cache_best_effort(lease.task_id) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) return True # The task may already have committed a terminal/control state. @@ -1492,6 +1551,7 @@ def settle_task_lease_isolated( settle_db, lease.task_id, task_lease=lease, + turn_id=turn_id, ) if error_message is not None: return False @@ -1509,6 +1569,7 @@ def finish_turn( task_id: int, *, task_lease: TaskLease | None = None, + turn_id: str | None = None, ) -> bool: """Reconcile terminal fields and, when supplied, release one exact lease. @@ -1541,6 +1602,19 @@ def finish_turn( - other statuses (PAUSED / WAITING_FOR_USER): control status and ``output`` are preserved; the lease release clears any stale ``error_message`` left by an earlier failed attempt + + ``turn_id``, when given, pops that turn's ephemeral connector secrets + (see ``connector_runtime.store_ephemeral_runtime_values``) from exactly + the COMPLETED, FAILED, and RUNNING-fallback branches above - the ones + that read this same already-fenced ``fresh.status`` as genuinely + terminal. This is the sole place that decides both things, so there is + no separate later read of the row to race against a fast concurrent + resume: the live-other-owner skip (this coroutine is not the one + settling the turn) never pops or renews. The PAUSED / WAITING_FOR_USER + branch instead renews the same secrets' TTL - it is the same turn + resuming later under this same turn_id, now carrying a fresh interaction + lifetime of its own, so its secrets must not expire on the original + pause's clock. """ from ..models.chat_message import TaskChatMessage from .workforce_runtime import sync_workforce_run_status @@ -1609,6 +1683,8 @@ def commit_terminal(status: TaskStatus, *, changed: bool = True) -> bool: task_id, len(latest_assistant.content), ) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) return committed else: logger.warning( @@ -1619,10 +1695,13 @@ def commit_terminal(status: TaskStatus, *, changed: bool = True) -> bool: trigger_run_changed = sync_trigger_run_status( bg_db, fresh, TaskStatus.COMPLETED ) - return commit_terminal( + committed = commit_terminal( TaskStatus.COMPLETED, changed=run_changed or trigger_run_changed, ) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) + return committed if status == TaskStatus.FAILED: changed = False @@ -1644,8 +1723,13 @@ def commit_terminal(status: TaskStatus, *, changed: bool = True) -> bool: "finish_turn: task %s marked failed (cleared stale output)", task_id, ) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) return committed - return commit_terminal(TaskStatus.FAILED, changed=False) + committed = commit_terminal(TaskStatus.FAILED, changed=False) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) + return committed if status == TaskStatus.RUNNING: # Lease ownership guard: a live lease held by another worker @@ -1689,12 +1773,22 @@ def commit_terminal(status: TaskStatus, *, changed: bool = True) -> bool: "flipping to FAILED", task_id, ) + if turn_id is not None: + _pop_ephemeral_runtime_values_best_effort(turn_id) return committed # PAUSED / WAITING_FOR_USER / other: preserve the control status while # releasing this exact run's lease. Legacy callers still leave it alone. if task_lease is not None: - return commit_terminal(status) + committed = commit_terminal(status) + # A fresh pause carries its own new interaction lifetime; the + # ephemeral secrets it may still need must not expire on the + # ORIGINAL pause's clock (see connector_runtime. + # renew_ephemeral_runtime_values's docstring for why this can't + # just be left to the opportunistic reaper). + if committed and turn_id is not None: + _renew_ephemeral_runtime_values_best_effort(turn_id) + return committed return False @@ -1822,6 +1916,7 @@ async def _runner() -> None: error_message=( "task execution cancelled during lease acquisition" ), + turn_id=payload.turn_id, ), ) if lease is None: @@ -2041,6 +2136,14 @@ async def execute_owned_run() -> None: ) if not defer_settlement_to_ttl_recovery: + # When this IS deferred (lease lost, DB pool exhaustion, + # unhealthy heartbeat), finish_turn's turn_id-scoped pop + # below never runs for this turn - and cannot safely run + # here either: this coroutine deliberately does not know + # (and must not guess by querying) whether the task will + # land on a terminal status or resume again under this + # same turn_id. connector_runtime.py's + # _EPHEMERAL_RUNTIME_TTL_SECONDS bounds that leak instead. lease_settled = False try: settled = await run_db_io_cancellation_safe( @@ -2053,6 +2156,7 @@ async def execute_owned_run() -> None: or CLIENT_SAFE_TASK_FAILURE ), client_message_type=client_history_message_type, + turn_id=turn_id, ) ) # Gate on the returned value, not on "didn't raise": @@ -2132,18 +2236,17 @@ async def execute_owned_run() -> None: task_id, exc_info=True, ) - try: - from .connector_runtime import pop_ephemeral_runtime_values - - if turn_id is not None: - pop_ephemeral_runtime_values(turn_id) - except Exception: - logger.warning( - "connector runtime cleanup failed for task %s turn %s", - task_id, - turn_id, - exc_info=True, - ) + # Ephemeral per-turn connector secrets are popped from inside + # settle_task_lease_isolated/finish_turn above (via the turn_id + # passed into each settle call), the moment - and using the same + # already-fenced row read - that one of them decides this turn + # reached a genuinely terminal outcome. That keeps a paused turn + # resuming under this same turn_id (WAITING_FOR_USER / PAUSED) + # from losing values it still needs, without a second, separate + # status read here racing a fast concurrent resume. A bystander + # coroutine that never held ``lease`` (skipped the block above + # entirely) correctly never pops either: it was never + # authoritative for this turn's outcome. if mcp_runtime_authorization_policy is not None: try: await run_db_io_cancellation_safe( diff --git a/src/xagent/web/tools/config.py b/src/xagent/web/tools/config.py index 1080375c51..5c0cd2521c 100644 --- a/src/xagent/web/tools/config.py +++ b/src/xagent/web/tools/config.py @@ -2096,6 +2096,19 @@ def _load_connector_runtime_view(self) -> Dict[str, Any]: ) from exc return self._connector_runtime_view + def get_connector_runtime_turn_id(self) -> Optional[str]: + """The turn id this config's connector runtime/secrets are bound to. + + A resume must look up ephemeral per-turn connector secrets under the + SAME turn id the original pausing turn stored them under - which is + exactly ``_connector_runtime_turn_id`` here, precisely because a + resume deliberately never rebinds it: websocket.py's + ``execute_resume_background`` reads this to know which turn's secrets + a terminal settlement must pop. + """ + + return self._connector_runtime_turn_id + def set_connector_runtime_turn_id(self, turn_id: Optional[str]) -> bool: """Switch the per-turn connector runtime source for reused agents. diff --git a/tests/web/api/test_execution_scope_turn_wiring.py b/tests/web/api/test_execution_scope_turn_wiring.py index 03595bcf48..27a8632228 100644 --- a/tests/web/api/test_execution_scope_turn_wiring.py +++ b/tests/web/api/test_execution_scope_turn_wiring.py @@ -873,7 +873,10 @@ def finalize_resume(*args: Any, **kwargs: Any) -> dict[str, Any]: } def release_resume_lease( - acquired_lease: object, *, error_message: str | None + acquired_lease: object, + *, + error_message: str | None, + turn_id: str | None = None, ) -> None: assert acquired_lease is lease assert error_message is None diff --git a/tests/web/api/test_websocket_owner_actor.py b/tests/web/api/test_websocket_owner_actor.py index 847941bf5f..71413440a7 100644 --- a/tests/web/api/test_websocket_owner_actor.py +++ b/tests/web/api/test_websocket_owner_actor.py @@ -34,6 +34,7 @@ from xagent.core.execution_scope import ( ExecutionScope, ) +from xagent.core.tools.adapters.vibe.connector_runtime import ConnectorRef from xagent.web.api import websocket as websocket_api from xagent.web.api.websocket import ( ResumeReservationOutcome, @@ -58,8 +59,14 @@ from xagent.web.models.task_interaction import TaskInteractionRequest from xagent.web.models.uploaded_file import UploadedFile from xagent.web.models.user import User +from xagent.web.services import connector_runtime as connector_runtime_module from xagent.web.services import task_orchestrator from xagent.web.services.chat_history_service import DELIVERY_FAILED, DELIVERY_PENDING +from xagent.web.services.connector_runtime import ( + get_ephemeral_runtime_values, + pop_ephemeral_runtime_values, + store_ephemeral_runtime_values, +) from xagent.web.services.managed_file_ref import ( DurableObjectIntegrityError, DurableStorageOperationError, @@ -3745,6 +3752,147 @@ async def test_execute_resume_background_persists_assistant_for_live_turn( assert task.output == "Guidance applied" +@pytest.mark.asyncio +async def test_execute_resume_background_pops_ephemeral_secrets_after_completion( + db_session, +) -> None: + """execute_resume_background settles through its own _finalize_resumed_task, + a completely separate finalizer from task_orchestrator.finish_turn - so a + turn that paused on waiting_for_user and then genuinely completes on + resume must pop its ephemeral secrets from THIS finalizer too, under the + original pausing turn's id; otherwise they'd sit until the TTL reaper + eventually reclaims them instead of being freed at the terminal outcome.""" + owner = _user(db_session, "owner") + task = _task(db_session, owner.id, status=TaskStatus.WAITING_FOR_USER) + turn_id = "resume-secrets-turn-completes" + store_ephemeral_runtime_values( + turn_id, + {ConnectorRef("mcp", 1): {"secrets": {"authorization": "Bearer resume-token"}}}, + ) + assert get_ephemeral_runtime_values(turn_id) is not None + + tool_config = MagicMock() + tool_config.get_connector_runtime_turn_id.return_value = turn_id + agent = MagicMock(tool_config=tool_config) + agent.resume_execution_by_id = AsyncMock( + return_value={ + "status": "completed", + "success": True, + "output": "done", + "agent_result": {}, + } + ) + ws_manager = MagicMock(broadcast_to_task=AsyncMock()) + + with patch("xagent.web.api.websocket.manager", ws_manager): + _register_current_resume(int(task.id)) + await execute_resume_background( + task_id=int(task.id), + agent_service=agent, + task_owner_user_id=int(owner.id), + ) + + db_session.refresh(task) + assert task.status == TaskStatus.COMPLETED + assert get_ephemeral_runtime_values(turn_id) is None + assert pop_ephemeral_runtime_values(turn_id) is None + + +@pytest.mark.asyncio +async def test_execute_resume_background_keeps_ephemeral_secrets_when_resume_pauses_again( + db_session, +) -> None: + """A resume that itself pauses again on waiting_for_user is the same turn + continuing under the same turn_id, not a finished one - popping here + would strand that next resume with nothing to look up its own secrets + under.""" + owner = _user(db_session, "owner") + task = _task(db_session, owner.id, status=TaskStatus.WAITING_FOR_USER) + turn_id = "resume-secrets-turn-repauses" + store_ephemeral_runtime_values( + turn_id, + {ConnectorRef("mcp", 1): {"secrets": {"authorization": "Bearer resume-token"}}}, + ) + assert get_ephemeral_runtime_values(turn_id) is not None + + tool_config = MagicMock() + tool_config.get_connector_runtime_turn_id.return_value = turn_id + agent = MagicMock(tool_config=tool_config) + agent.resume_execution_by_id = AsyncMock( + return_value={ + "status": "waiting_for_user", + "success": False, + "output": "Please connect another app.", + "agent_result": {}, + } + ) + ws_manager = MagicMock(broadcast_to_task=AsyncMock()) + + with patch("xagent.web.api.websocket.manager", ws_manager): + _register_current_resume(int(task.id)) + await execute_resume_background( + task_id=int(task.id), + agent_service=agent, + task_owner_user_id=int(owner.id), + ) + + db_session.refresh(task) + assert task.status == TaskStatus.WAITING_FOR_USER + assert get_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values(turn_id) is not None + + +@pytest.mark.asyncio +async def test_execute_resume_background_renews_ephemeral_secrets_when_resume_pauses_again( + db_session, monkeypatch +) -> None: + """A resume that re-pauses carries a fresh interaction lifetime of its + own - _finalize_resumed_task must actually renew the secrets' TTL, not + merely leave them alone, or they'd still expire on the ORIGINAL pause's + clock even though this one is still active.""" + owner = _user(db_session, "owner") + task = _task(db_session, owner.id, status=TaskStatus.WAITING_FOR_USER) + turn_id = "resume-secrets-turn-renews" + store_ephemeral_runtime_values( + turn_id, + {ConnectorRef("mcp", 1): {"secrets": {"authorization": "Bearer resume-token"}}}, + ) + + real_monotonic = connector_runtime_module.time.monotonic + offset = {"value": connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1} + monkeypatch.setattr( + connector_runtime_module.time, + "monotonic", + lambda: real_monotonic() + offset["value"], + ) + + tool_config = MagicMock() + tool_config.get_connector_runtime_turn_id.return_value = turn_id + agent = MagicMock(tool_config=tool_config) + agent.resume_execution_by_id = AsyncMock( + return_value={ + "status": "waiting_for_user", + "success": False, + "output": "Please connect another app.", + "agent_result": {}, + } + ) + ws_manager = MagicMock(broadcast_to_task=AsyncMock()) + + with patch("xagent.web.api.websocket.manager", ws_manager): + _register_current_resume(int(task.id)) + await execute_resume_background( + task_id=int(task.id), + agent_service=agent, + task_owner_user_id=int(owner.id), + ) + + # Past the ORIGINAL store's TTL window, but well within the renewed one. + offset["value"] += connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1 + assert get_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values(turn_id) is not None + + @pytest.mark.asyncio async def test_execute_resume_background_persists_missing_checkpoint_failure( db_session, diff --git a/tests/web/services/test_connector_runtime_ephemeral.py b/tests/web/services/test_connector_runtime_ephemeral.py new file mode 100644 index 0000000000..bae7cfc331 --- /dev/null +++ b/tests/web/services/test_connector_runtime_ephemeral.py @@ -0,0 +1,161 @@ +"""The ephemeral per-turn connector secrets store's TTL-based lifetime. + +Some settlement paths (task_orchestrator.py's and websocket.py's deferred- +to-TTL-recovery branches: lease lost, DB pool exhaustion, unhealthy +heartbeat at shutdown) can never safely pop a turn's ephemeral secrets +themselves - at the point they bail out they genuinely do not know whether +the task will land on a terminal status or resume again under the same +turn_id, and task_lease_recovery.py's later batch sweep has no way to map a +recovered task_id back to the turn_id its secrets were stored under. Without +a bound, a turn that goes through one of those paths leaks its secrets for +the life of the process (see connector_runtime.py's `_EPHEMERAL_RUNTIME_ +VALUES`, an unbounded module-global dict). Expiry is enforced on every +read/pop, not just by the opportunistic reaper another turn's store call may +trigger - a stale entry that nothing ever looks up again still becomes +unreadable the moment it ages past the TTL, not merely "eventually reclaimed +whenever something else happens to store." A turn that instead resumes +again under its same turn_id (WAITING_FOR_USER/PAUSED) renews the TTL, so a +still-active pause's secrets don't expire on the clock of the FIRST pause +that stored them. +""" + +from __future__ import annotations + +from xagent.core.tools.adapters.vibe.connector_runtime import ConnectorRef +from xagent.web.services import connector_runtime as connector_runtime_module +from xagent.web.services.connector_runtime import ( + get_ephemeral_runtime_manifest, + get_ephemeral_runtime_values, + pop_ephemeral_runtime_values, + renew_ephemeral_runtime_values, + store_ephemeral_runtime_values, +) + +_VALUES = {ConnectorRef("mcp", 1): {"secrets": {"authorization": "Bearer t"}}} + + +def _advance_clock(monkeypatch, seconds: float) -> None: + """Move every future ``time.monotonic()`` read forward by ``seconds``.""" + + real_monotonic = connector_runtime_module.time.monotonic + offset = {"value": seconds} + monkeypatch.setattr( + connector_runtime_module.time, + "monotonic", + lambda: real_monotonic() + offset["value"], + ) + + +def test_ephemeral_values_survive_well_within_the_ttl(monkeypatch) -> None: + turn_id = "ephemeral-ttl-fresh" + store_ephemeral_runtime_values(turn_id, _VALUES) + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS / 2 + ) + + store_ephemeral_runtime_values("ephemeral-ttl-fresh-trigger", _VALUES) + + assert get_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values("ephemeral-ttl-fresh-trigger") is not None + + +def test_a_stale_entry_is_reclaimed_by_the_next_store_call(monkeypatch) -> None: + """A stale entry that nothing ever reads again gets swept off the module + dicts as soon as an unrelated turn's store runs its opportunistic prune - + the entry does not merely become unreadable, it is actually removed.""" + turn_id = "ephemeral-ttl-stale" + store_ephemeral_runtime_values(turn_id, _VALUES) + assert get_ephemeral_runtime_values(turn_id) is not None + + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS + 1 + ) + + other_turn_id = "ephemeral-ttl-unrelated-store" + store_ephemeral_runtime_values(other_turn_id, _VALUES) + + assert turn_id not in connector_runtime_module._EPHEMERAL_RUNTIME_VALUES + assert pop_ephemeral_runtime_values(other_turn_id) is not None + + +def test_expiry_also_drops_the_manifest(monkeypatch) -> None: + turn_id = "ephemeral-ttl-manifest" + store_ephemeral_runtime_values(turn_id, _VALUES) + assert get_ephemeral_runtime_manifest(turn_id) is not None + + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS + 1 + ) + store_ephemeral_runtime_values("ephemeral-ttl-manifest-trigger", _VALUES) + + assert get_ephemeral_runtime_manifest(turn_id) is None + pop_ephemeral_runtime_values("ephemeral-ttl-manifest-trigger") + + +def test_reads_treat_an_expired_entry_as_absent_with_no_intervening_store( + monkeypatch, +) -> None: + """A quiet process - no other turn ever stores anything after this one + expires - must still stop returning it. Expiry has to be an observable + property of get/pop themselves, not only a side effect the NEXT store + call happens to trigger; otherwise a deployment with infrequent traffic + could read a secret well past its advertised TTL.""" + turn_id = "ephemeral-ttl-quiet-process" + store_ephemeral_runtime_values(turn_id, _VALUES) + assert get_ephemeral_runtime_values(turn_id) is not None + + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS + 1 + ) + + assert get_ephemeral_runtime_values(turn_id) is None + assert get_ephemeral_runtime_manifest(turn_id) is None + assert pop_ephemeral_runtime_values(turn_id) is None + + +def test_renew_extends_the_ttl_for_a_still_live_entry(monkeypatch) -> None: + """A second pause under the same turn_id carries its own fresh + interaction lifetime - its secrets must not expire on the ORIGINAL + pause's clock.""" + turn_id = "ephemeral-ttl-renewed" + store_ephemeral_runtime_values(turn_id, _VALUES) + + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1 + ) + renew_ephemeral_runtime_values(turn_id) + + # Past the ORIGINAL store's TTL window, but well within the renewed one. + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1 + ) + assert get_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values(turn_id) is not None + + +def test_renew_does_not_resurrect_an_already_expired_entry(monkeypatch) -> None: + """Renewal is not a way to un-expire something that already aged out - + once gone, a late renewal call must stay a no-op, not silently bring the + secrets back for a turn nothing else remembers is still active.""" + turn_id = "ephemeral-ttl-renew-too-late" + store_ephemeral_runtime_values(turn_id, _VALUES) + + _advance_clock( + monkeypatch, connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS + 1 + ) + renew_ephemeral_runtime_values(turn_id) + + assert get_ephemeral_runtime_values(turn_id) is None + assert turn_id not in connector_runtime_module._EPHEMERAL_RUNTIME_VALUES + + +def test_renew_is_a_noop_for_a_turn_with_nothing_stored() -> None: + """A turn that never stored ephemeral secrets (or was already popped) has + nothing to keep alive - renewing it must not fabricate an entry.""" + turn_id = "ephemeral-ttl-renew-unknown" + + renew_ephemeral_runtime_values(turn_id) + + assert turn_id not in connector_runtime_module._EPHEMERAL_RUNTIME_STORED_AT + assert get_ephemeral_runtime_values(turn_id) is None diff --git a/tests/web/services/test_task_orchestrator.py b/tests/web/services/test_task_orchestrator.py index cd62784ae0..8693b90c1e 100644 --- a/tests/web/services/test_task_orchestrator.py +++ b/tests/web/services/test_task_orchestrator.py @@ -57,6 +57,7 @@ ) from xagent.web.models.user import User from xagent.web.models.workforce import Workforce, WorkforceRun +from xagent.web.services import connector_runtime as connector_runtime_module from xagent.web.services import task_orchestrator as task_orchestrator_module from xagent.web.services.assistant_history_safety import ( CLIENT_SAFE_FAILURE_MESSAGE_TYPE, @@ -1698,6 +1699,42 @@ def test_finish_turn_does_not_touch_a_new_run_owned_by_same_process( assert persisted.error_message is None +def test_finish_turn_waiting_for_user_renews_ephemeral_secrets( + db_session, monkeypatch +) -> None: + """A turn that pauses on waiting_for_user is the same turn resuming later + under its same turn_id, now carrying a fresh interaction lifetime of its + own - its ephemeral secrets must not expire on the ORIGINAL pause's + clock (see connector_runtime.renew_ephemeral_runtime_values).""" + user = _create_user(db_session) + task = _create_task(db_session, user.id, status=TaskStatus.WAITING_FOR_USER) + task.runner_id = get_runner_id() + task.run_id = "waiting-run" + task.lease_expires_at = datetime.now(timezone.utc) + timedelta(minutes=5) + db_session.commit() + lease = TaskLease( + task_id=int(task.id), runner_id=get_runner_id(), run_id="waiting-run" + ) + + turn_id = "finish-turn-renews-waiting" + _store_runtime_secret_for_turn(turn_id) + + real_monotonic = connector_runtime_module.time.monotonic + offset = {"value": connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1} + monkeypatch.setattr( + connector_runtime_module.time, + "monotonic", + lambda: real_monotonic() + offset["value"], + ) + + finish_turn(db_session, int(task.id), task_lease=lease, turn_id=turn_id) + + # Past the ORIGINAL store's TTL window, but well within the renewed one. + offset["value"] += connector_runtime_module._EPHEMERAL_RUNTIME_TTL_SECONDS - 1 + assert get_ephemeral_runtime_values(turn_id) is not None + assert pop_ephemeral_runtime_values(turn_id) is not None + + def test_finish_turn_cache_invalidation_failure_is_non_fatal_after_release( db_session, ) -> None: @@ -1949,8 +1986,14 @@ async def test_schedule_bg_skips_finish_turn_when_lease_acquire_fails( mock_exec.assert_not_awaited() mock_finish.assert_not_called() - assert get_ephemeral_runtime_values(payload.turn_id) is None - assert pop_ephemeral_runtime_values(payload.turn_id) is None + # This coroutine never acquired the lease, so it is a bystander that + # never becomes authoritative for the turn - the other worker holding + # the lease owns this turn_id's eventual settlement (and pop) too. A + # bystander popping here would be the exact bug this whole mechanism + # exists to avoid: it could evict values a turn still actively running + # elsewhere depends on. + assert get_ephemeral_runtime_values(payload.turn_id) is not None + assert pop_ephemeral_runtime_values(payload.turn_id) is not None @pytest.mark.asyncio @@ -2519,7 +2562,12 @@ async def test_schedule_bg_releases_lease_on_execute_task_background_exception( user = _create_user(db_session) task = _create_task(db_session, user.id, status=TaskStatus.RUNNING) - fake_lease = TaskLease(task_id=int(task.id), runner_id="test-runner") + task.runner_id = "test-runner" + task.run_id = "run-a" + db_session.commit() + fake_lease = TaskLease( + task_id=int(task.id), runner_id="test-runner", run_id="run-a" + ) payload = TaskTurnPayload("x") _store_runtime_secret_for_turn(payload.turn_id) assert get_ephemeral_runtime_values(payload.turn_id) is not None @@ -2539,6 +2587,10 @@ async def test_schedule_bg_releases_lease_on_execute_task_background_exception( ), patch( "xagent.web.services.task_orchestrator.settle_task_lease_isolated", + # wraps, not a bare Mock: the ephemeral-values assertion below + # needs the real settlement (and finish_turn's turn_id pop + # inside it) to actually run, not just be recorded as called. + wraps=settle_task_lease_isolated, ) as mock_settle, patch.object(background_task_manager, "register_task"), patch( @@ -2763,7 +2815,12 @@ async def test_schedule_bg_forwards_execution_message_to_execute_task_background user = _create_user(db_session) task = _create_task(db_session, user.id, status=TaskStatus.RUNNING) - fake_lease = TaskLease(task_id=int(task.id), runner_id="test-runner") + task.runner_id = "test-runner" + task.run_id = "run-a" + db_session.commit() + fake_lease = TaskLease( + task_id=int(task.id), runner_id="test-runner", run_id="run-a" + ) payload = TaskTurnPayload( transcript_message="summarize this", execution_message="summarize this\n\n[uploaded file: secret.txt]", @@ -2788,6 +2845,10 @@ async def test_schedule_bg_forwards_execution_message_to_execute_task_background ) as mock_exec, patch( "xagent.web.services.task_orchestrator.settle_task_lease_isolated", + # wraps, not a bare Mock: the ephemeral-values assertion below + # needs the real settlement (and finish_turn's turn_id pop + # inside it) to actually run, not just be recorded as called. + wraps=settle_task_lease_isolated, ), patch.object(background_task_manager, "register_task"), patch( @@ -2799,6 +2860,7 @@ async def test_schedule_bg_forwards_execution_message_to_execute_task_background task_id=int(task.id), task_owner_user_id=int(user.id), task_source=task.source, + run_id="run-a", payload=payload, force_fresh=False, context={"turn_id": "caller-turn", "existing": "value"}, @@ -2829,6 +2891,70 @@ async def test_schedule_bg_forwards_execution_message_to_execute_task_background assert pop_ephemeral_runtime_values(payload.turn_id) is None +@pytest.mark.asyncio +async def test_schedule_bg_keeps_ephemeral_values_when_turn_pauses_waiting_for_user( + db_session, +) -> None: + """A turn that pauses on waiting_for_user is the SAME turn resuming later + under its SAME turn_id, not a finished one - popping its ephemeral values + here would leave that resume with nothing to look up under the one + turn_id a resume deliberately never rebinds (see + WebToolConfig.get_connector_runtime_turn_id).""" + from xagent.web.api.websocket import background_task_manager + from xagent.web.services.task_lease_service import TaskLease + + user = _create_user(db_session) + task = _create_task(db_session, user.id, status=TaskStatus.RUNNING) + task_id = int(task.id) + fake_lease = TaskLease(task_id=task_id, runner_id="test-runner") + payload = TaskTurnPayload("x") + _store_runtime_secret_for_turn(payload.turn_id) + assert get_ephemeral_runtime_values(payload.turn_id) is not None + + async def pause_execution(**_kwargs) -> None: + def commit_waiting_status() -> None: + SessionLocal = database_module.get_session_local() + with SessionLocal() as waiting_db: + waiting_task = waiting_db.query(Task).filter(Task.id == task_id).one() + waiting_task.status = TaskStatus.WAITING_FOR_USER + waiting_task.control_state = "waiting_for_user" + waiting_db.commit() + + await asyncio.to_thread(commit_waiting_status) + + with ( + patch( + "xagent.web.services.task_orchestrator.acquire_task_lease_isolated", + return_value=fake_lease, + ), + patch( + "xagent.web.services.task_orchestrator.run_task_lease_heartbeat", + new=AsyncMock(), + ), + patch( + "xagent.web.api.websocket.execute_task_background", + new=AsyncMock(side_effect=pause_execution), + ), + patch( + "xagent.web.services.task_orchestrator.settle_task_lease_isolated", + ), + patch.object(background_task_manager, "register_task"), + ): + bg_task = _schedule_bg( + task_id=task_id, + task_owner_user_id=int(user.id), + task_source=task.source, + payload=payload, + force_fresh=False, + context=None, + ) + await bg_task + + assert get_ephemeral_runtime_values(payload.turn_id) is not None + # Cleanup: don't leak this test's entry into module-global state. + assert pop_ephemeral_runtime_values(payload.turn_id) is not None + + @pytest.mark.asyncio async def test_schedule_bg_acquires_expired_lease_on_first_try(db_session) -> None: """Expired lease columns are granted by acquire_task_lease's atomic WHERE."""