diff --git a/docs/system-specs/modules/subagent.md b/docs/system-specs/modules/subagent.md index 61b625bce17..fb06c7d44a9 100644 --- a/docs/system-specs/modules/subagent.md +++ b/docs/system-specs/modules/subagent.md @@ -546,10 +546,133 @@ On startup, `SubagentManager` scans `~/.kiro/crew/subagents/` and reconciles: - Created on: process death without result, delivery failure, timeout (`cause` = `error` / `timeout` / `cancelled` / `reaped` / `gateway_restart`), **and on successful delivery** (`cause="delivered"`, via `mark_delivered`) so `result.txt` - is retained for the grace window instead of deleted immediately. + is retained for the grace window instead of deleted immediately. The generic + writer snapshots any non-empty session ID, provider, and CWD from readable + state; live abnormal-exit values captured immediately after session acquisition + (before resume validation and context construction) override that fallback. + Cancel recovery can acquire multiple sessions under one run ID, so persistence + atomically records complete per-session cleanup generations in an owner-only + record below the file-gated `trust/` root, outside the agent-writable run folder. + Every read and write first applies the repository's fail-loud owner-only directory + restriction, including inheritable Windows DACLs, and re-locks an existing record + because tightening its parent does not retrofit an older file ACL. Only a missing + record reads as empty; I/O, parse, or schema failures propagate, so an append can + never rewrite unreadable history as a fresh one-generation record. Prune catches + those failures per tombstone and continues later entries without altering the + corrupt record; shared-session setup logs them as best-effort and keeps the live + handle instead of falling into dedicated fallback. + Each generation also records the run's retention intent and continuation owner + key before the later best-effort combined state update. Protected generation + ownership is authoritative even when readable agent-writable state supplies an + empty or conflicting key; only the referenced owner's current readable state + decides retention. The generation `keep` value is a fallback when local state + lacks that field. Session acquisition publishes the new generation + synchronously in memory before submitting durable work, so executor + saturation or cancellation cannot prevent a terminal tombstone from seeing it. + The in-memory fallback is append-only; the worker deduplicates only while + serializing the protected record off-loop and never replaces the live list, so + an older writer cannot discard a concurrent recovery SID. On the dedicated arm, + durable generation persistence follows the cancellation-drained provenance + write; both dedicated and shared identity workers are shielded and fully drained + before cancellation is re-raised, so restart cannot precede protected authority. + Cancellation cannot skip required model fields, and the already-published + memory record still feeds the terminal tombstone. Slow storage therefore cannot + stall chat/heartbeat or hide a just-acquired session from a contending tombstone, + and a transient SID1 generation-write failure cannot be lost when SID2 later + succeeds. Shared-handle ownership and provider references are attached + immediately after handle creation, before any cancellable persistence + await, so force-reap always destroys the shared handle rather than resetting a + nonexistent dedicated session. Identity persistence errors are logged without + triggering dedicated fallback or abandoning the live handle. Event-loop tombstone snapshots are memory-only: they acquire the identity + lock non-blocking and use already-published in-memory generations, never reading + the protected durable record. Executor-owned prune independently merges that + record after restart before cleanup. The protected record and in-memory fallback + are evicted when prune or explicit folder deletion succeeds. Tombstones expose + the latest identity in compatibility fields and snapshot the full list for + diagnostics/restart hints, but those agent-folder fields cannot authorize + provider deletion. Prune's deletion set comes only from the protected record + and synchronous live gateway publication, and reclaims every trusted generation. + Retention fallback selects a protected generation matching the current readable-state + SID, or the latest protected generation when the SID is absent or mismatched; + agent-writable state/tombstone SID and owner fields never select a victim identity + or suppress trusted owner/`keep` metadata. +- For readable state, only a literal current `keep is True` is retained; strings + such as `"false"` are non-retention rather than truthy policy. `true` preserves + the identity folder for restart registry rebuild; release writes `false`, + allowing prune to retry provider cleanup. + Every completed plain run records `false`; a readable legacy/failed-write + record with no key is treated as non-retained and prunes at the normal cutoff. + Readable `true` always defers disk prune; release or the conversation TTL writes + `false` and owns deletion. This arbitration is part of the cleanup fix rather + than a separate retention feature: once durable identity makes provider files + reachable, a stale `keep=False` prune racing a continuation promotion could + destroy the newly reachable resume material. Within the single gateway process, + promotion and prune arbitrate under per-agent short-held state transactions. + When a continuation tombstone points at an original owner, prune pre-resolves + that owner and acquires both per-agent locks in stable order, then re-reads under + lock; unrelated agents never contend. Promotion writes `true` before that locked + read, or prune keeps arbitration through provider cleanup and folder removal so a + later promotion returns retryable instead of racing deletion. On the event loop, + promotion probes arbitration and the per-agent off-loop-writer lock non-blocking. + Contention returns retryable `conversation_busy`, so a later retry writes + `keep=True` only after every older writer completes. + Off-loop promotion lets `update_state` acquire that non-reentrant writer lock + normally, avoiding self-deadlock while preserving serialization. + Transient persistence errors likewise return retryable without dispatch. Retry + restores the exact pre-attempt SessionManager and TTL-registry ownership; an + already-retained conversation is never unmarked. The facade returns the result + directly, so concurrent callers carry independent outcomes without a hidden + clear/call/read side channel. A process crash leaves no half-committed claim + format: the next prune re-reads the current owner state under arbitration. + If state and a rewritten tombstone both lack a top-level SID, prune derives + retention and owner from the latest valid durable cleanup generation instead of + treating the record as non-retained. Provider cleanup runs before lock release; + folder and protected-record removal follow only when every trusted generation + reports success. Unsupported providers and transient deletion failures preserve + both retry surfaces for later sweeps, capped at 90 days so a permanently missing + cleanup route cannot accumulate private run folders forever. A legacy SID present + only in agent-folder state/tombstone likewise preserves the folder inside that + window: it cannot authorize deletion, but the lookup gives a later trusted + migration time to reclaim the transcript. At the hard ceiling, only the run + folder and protected metadata are reaped; untrusted identity is never used for + provider-file deletion. Restart registry rebuild likewise accepts only literal + `keep is True`, requires the state SID and conversation owner to match a + protected/live generation, and sources provider/CWD from that trusted record; + agent-folder state cannot seed a victim SID into the later TTL release path. An + explicitly injected noncanonical state reader is an application-owned trusted + seam; the canonical disk reader never takes that compatibility fallback. + A continuation follows its original conversation's + readable `keep` value directly: `false` or a missing key is non-retention, while + unreadable owner state receives the bounded grace below instead of inheriting + the continuation's stale local `true`. Registry rebuild, prune, and TTL sweep + share `subagent_id_from_conversation_key`; malformed keys are dropped per entry + so one corrupt record cannot abort later cleanup. - Pruned by reaper: `delivered` tombstones after `agent.subagent_result_ttl_secs` (default 1h); all other tombstones after 7 days. `prune_stale_tombstones` takes - a per-cause cutoff for this. + a per-cause cutoff for this and treats timestamps exactly at the cutoff as + eligible, avoiding platform clock-resolution gaps. Tombstone `died` must be numeric, positive, and + non-future; string, NaN, infinity, future, and oversized values fall back to the + validated tombstone-file mtime, or the current sweep time when no valid bounded + time exists. That final fallback preserves unknown-retention grace across wall-clock + rollback instead of making the record immediately eligible. + Missing, malformed, deeply nested, or non-object tombstones are skipped for that + entry without aborting later entries in the sweep. + Missing, malformed, deeply nested, Unicode-invalid, or non-object `state.json` + is unreadable. An acquisition-time `keep=false` generation remains unknown in + this branch because a later promotion may have landed only in the now-unreadable + state; only `keep=true` may collapse uncertainty, since it can only preserve data. + A tombstone with a SID publishes only a SID-less in-memory retention hint, so the + SessionManager file-deletion exemption performs no tombstone read on the gateway + event loop without laundering agent-folder identity into provider-deletion + authority; the executor-owned restart scan + rehydrates that hint from durable tombstones before registry rebuild completes. + A readable legacy continuation whose owner + state is unreadable uses its resolved local-state SID for the same bounded + protection, even when its pre-upgrade tombstone has no SID. A malformed + continuation owner ID is likewise bounded as unknown retention for that entry; + it cannot abort processing of later tombstones. At the cutoff, unknown intent receives one extra 24-hour grace anchored on tombstone death time; + after that bounded window, trusted cleanup-generation metadata drives best-effort + provider cleanup while tombstone metadata drives folder removal eligibility. - `spawn_status` falls back to persistence layer for completed/tombstoned agents, reading the retained `result.txt` (and honoring offset/limit/grep). diff --git a/src/kiro_crew/subagent_manager/continuation.py b/src/kiro_crew/subagent_manager/continuation.py index fa7ccd86191..4ace8524c1b 100644 --- a/src/kiro_crew/subagent_manager/continuation.py +++ b/src/kiro_crew/subagent_manager/continuation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING +from .. import subagent_persistence as persistence from ._component import ManagerComponent if TYPE_CHECKING: @@ -34,6 +35,7 @@ class ContinuationCoordinator(ManagerComponent): """Own continuation transitions while state remains facade-owned.""" + _persistence = persistence __slots__ = () def _conversation_busy_impl(self, conv_key: str) -> SubagentInfo | None: @@ -79,37 +81,47 @@ def _conversation_busy_impl(self, conv_key: str) -> SubagentInfo | None: return None def _keep_recorded_on_disk_impl(self, key: str) -> bool: - """Disk-truth continuable check for SessionManager's cache (#1115). + """Retention guard for subagent conversations without loop-side disk probes. - True iff *key* is a subagent conversation whose run's ``state.json`` - records ``keep`` — the single persisted source of retention intent. - Non-subagent keys short-circuit without touching disk. + Readable ``state.json`` remains the persisted retention authority. If + state is unreadable, the session-acquisition/tombstone path's synchronous + in-memory identity publication fails safe without reading ``tombstone.json`` + on the gateway event loop. """ - if not key.startswith("subagent:"): + conv_id = self._persistence.subagent_id_from_conversation_key(key) + if conv_id is None: return False - conv_id = key[len("subagent:") :] try: - state = read_state(conv_id) or {} - except Exception: - return False - return bool(state.get("keep")) + state = read_state(conv_id) + except OSError: + state = None + if isinstance(state, dict): + return state.get("keep") is True + return self._persistence.has_live_cleanup_identity(conv_id) def _promote_conversation_impl( - self, conv_id: str, conv_key: str, last_used: float | None = None - ) -> None: - """Single choke point for promoting a conversation's retention (#1115). - - Writes all three retention surfaces together so they cannot drift: - ``keep=True`` in state.json (the persisted source of truth), the - SessionManager continuable cache (file-deletion exemption), and the - TTL registry entry (sweep ownership). - """ + self, + conv_id: str, + conv_key: str, + last_used: float | None = None, + ) -> Any: + """Atomically promote all retention surfaces for a conversation.""" try: - update_state(conv_id, keep=True) - except Exception: + result = self._persistence.promote_retention(conv_id, state_writer=update_state) + except OSError: logger.debug("promote: failed to persist keep for %s", conv_id, exc_info=True) + # Preserve the established fail-safe marker for direct callers, but + # report retryable failure so continue_conversation rolls it back. + self._manager._sessions.mark_continuable(conv_key) + self._manager._conversations[conv_key] = ( + last_used if last_used is not None else time.time() + ) + return self._persistence.RetentionPromotionResult.RETRYABLE + if result is not self._persistence.RetentionPromotionResult.PROMOTED: + return result self._manager._sessions.mark_continuable(conv_key) self._manager._conversations[conv_key] = last_used if last_used is not None else time.time() + return result def _scan_keep_states_impl(self) -> list[tuple[str, str, str, str, str, float]]: """Blocking scan for keep runs (#1114): read every ``state.json`` @@ -128,20 +140,46 @@ def _scan_keep_states_impl(self) -> list[tuple[str, str, str, str, str, float]]: try: if not d.is_dir(): continue - state = read_state(d.name) or {} - if not state.get("keep"): + state = read_state(d.name) + if not isinstance(state, dict): + tombstone = self._persistence.read_tombstone(d.name) or {} + sid = str(tombstone.get("session_id") or "") + if sid: + # Agent-folder tombstones can preserve retention/exemption + # hints, never provider-deletion authority. + self._persistence.publish_live_cleanup_hint(d.name) + continue + if state.get("keep") is not True: continue conv_key = str(state.get("conversation_key") or "") or f"subagent:{d.name}" - conv_id = conv_key[len("subagent:") :] + conv_id = self._persistence.subagent_id_from_conversation_key(conv_key) + if conv_id is None: + continue sid = str(state.get("session_id") or "") + trusted_identity = self._persistence.trusted_cleanup_identity_record( + d.name, + sid, + conv_key, + ) + if trusted_identity is None: + # The canonical disk reader consumes agent-writable state and + # must match protected authority. Tests/embedders may inject a + # different reader as their own trusted authority; retaining + # that established seam does not make on-disk state trusted. + if read_state is self._persistence.read_state: + continue + trusted_identity = { + "provider": state.get("provider"), + "cwd": state.get("cwd"), + } last_used = float(state.get("updated_at") or state.get("started") or 0.0) out.append( ( conv_id, conv_key, sid, - str(state.get("provider") or PROVIDER_LABEL_DEFAULT), - str(state.get("cwd") or ""), + str(trusted_identity.get("provider") or PROVIDER_LABEL_DEFAULT), + str(trusted_identity.get("cwd") or ""), last_used, ) ) @@ -301,7 +339,34 @@ def continue_conversation_impl( # (#1115): state.json keep=True (tombstone pruner skips deletion), # the SessionManager continuable cache, and the TTL registry entry. # The conversation TTL sweep / spawn_release owns deletion from here. - self._manager._promote_conversation(conv_id, conv_key) + # Snapshot existing ownership: a retryable promotion attempt must undo + # only state it introduced, never erase an earlier keep/continuation. + was_continuable = self._manager._sessions.is_continuable(conv_key) + had_previous_last_used = conv_key in self._manager._conversations + previous_last_used = self._manager._conversations.get(conv_key, 0.0) + # The facade forwards the enum result directly. Legacy tests and external + # monkeypatches that return None/non-enum retain the historical promoted + # behavior; only an explicit RETRYABLE outcome alters dispatch. + promotion = self._manager._promote_conversation( # type: ignore[func-returns-value] + conv_id, conv_key + ) + if promotion is self._persistence.RetentionPromotionResult.RETRYABLE: + if not was_continuable: + self._manager._sessions.unmark_continuable(conv_key) + if not had_previous_last_used: + self._manager._conversations.pop(conv_key, None) + else: + self._manager._conversations[conv_key] = previous_last_used + return SubagentInfo( + id=uuid.uuid4().hex[:8], + task=_redact(task), + done=True, + parent_session_key=parent_session_key, + error=( + "conversation_busy: retention promotion is temporarily " + f"unavailable for {conv_id}; retry the continuation" + ), + ) inc_memory, inc_lessons, inc_project = self._manager._inherited_context_groups(conv_id) # A continuation has to run WHERE THE RUN RAN. `spawn` resolves an empty # cwd to the pool project before it validates the agent name, so a run @@ -734,7 +799,11 @@ def _sweep_conversations_impl(self, now: float) -> None: if self._manager._conversation_busy(conv_key) is not None: self._manager._conversations[conv_key] = now # active — refresh continue - conv_id = conv_key[len("subagent:") :] + conv_id = self._persistence.subagent_id_from_conversation_key(conv_key) + if conv_id is None: + logger.warning("Dropping malformed conversation registry key %r", conv_key) + self._manager._conversations.pop(conv_key, None) + continue ok, detail = self._manager.release_conversation(conv_id) logger.info( "Conversation %s expired after %ds idle: %s", diff --git a/src/kiro_crew/subagent_manager/run.py b/src/kiro_crew/subagent_manager/run.py index 4360a7f901a..ccfcd8ebdde 100644 --- a/src/kiro_crew/subagent_manager/run.py +++ b/src/kiro_crew/subagent_manager/run.py @@ -2,8 +2,13 @@ from __future__ import annotations +import asyncio from typing import TYPE_CHECKING +from ..subagent_persistence import ( + publish_live_cleanup_identity, + remember_live_cleanup_identity, +) from ._component import ManagerComponent if TYPE_CHECKING: @@ -25,6 +30,7 @@ FALLBACK_CANDIDATE_ATTEMPTS, FALLBACK_STORY_ATTR, HOOK_EVENT_POST_TOOL_USE, + PROVIDER_LABEL_DEFAULT, TOOL_AUTO_APPROVE, TOOL_DENY, TRANSIENT_RETRIES, @@ -51,7 +57,6 @@ annotate_model_fallback, append_fallback_story, apply_completion_keep, - asyncio, cap_result_file, configured_fallback_chain, evict_completed_agents, @@ -73,6 +78,8 @@ class RunEventCoordinator(ManagerComponent): """Own run transitions while state remains facade-owned.""" + _publish_identity = staticmethod(publish_live_cleanup_identity) + _remember_identity = staticmethod(remember_live_cleanup_identity) __slots__ = () def _effective_turn_limit_impl(self, info: SubagentInfo) -> int: @@ -210,6 +217,54 @@ def _settled( info._state_drain_active = False raise + async def _remember_identity_off_loop( + self, + info: SubagentInfo, + *, + session_id: str, + provider: str, + cwd: str = "", + keep: bool | None = None, + conversation_key: str = "", + ) -> None: + """Persist protected cleanup identity off-loop and drain cancellation. + + The live generation is already published synchronously, but restart + durability depends on this protected write. ``to_thread`` workers survive + cancellation, so shield the worker and delay re-raising until it settles; + otherwise terminal teardown can finish and restart before the authority + record exists. + """ + writer = asyncio.ensure_future( + asyncio.to_thread( + self._remember_identity, + info.id, + session_id=session_id, + provider=provider, + cwd=cwd, + keep=keep, + conversation_key=conversation_key, + ) + ) + try: + await asyncio.shield(writer) + except asyncio.CancelledError: + info._state_drain_active = True + try: + while not writer.done(): + try: + await asyncio.wait({writer}) + except asyncio.CancelledError: + pass + if not writer.cancelled(): + try: + writer.result() + except Exception: + pass + finally: + info._state_drain_active = False + raise + def update_completion_keep_impl(self, mode: str, max_chars: int) -> None: """Update the live completion-keep mode and char budget. @@ -561,8 +616,15 @@ def _emit_queue_depth_impl(self, parent_session_key: str, batch_id: str = "") -> ) loop.create_task(self._manager._fire_event("subagent_queued", info, {"queued": depth})) - async def _run_inner_impl(self, info: SubagentInfo, session_key: str) -> None: + async def _run_inner_impl( + self, + info: SubagentInfo, + session_key: str, + ) -> None: """Inner execution — called within timeout wrapper.""" + setattr(info, "_session_id", "") + setattr(info, "_session_provider", "") + setattr(info, "_session_cwd", "") # Mark the real start of execution BEFORE any await so the startup # watchdog measures from here, not from registration (which may include # an arbitrary spawn-approval wait). Must be the first statement. @@ -730,23 +792,48 @@ async def _run_inner_impl(self, info: SubagentInfo, session_key: str) -> None: approval_policy=parent_policy, **extra_kwargs, ) - # Fail CLOSED on a continuation that did not actually resume: - # get_or_create silently falls back to a FRESH session when - # session/load fails (lock held, corrupt files, backend refusal). - # Executing the follow-up on that fresh session would silently run - # it context-free — worse than an honest error the parent can react - # to (re-spawn with a summary). conversation_key is only set by - # continue_conversation, so first spawns are unaffected. - if info.conversation_key and not _resumed: - raise RuntimeError( - "resume_failed: session/load did not restore conversation " - f"{info.conversation_key} — refusing to execute the " - "follow-up without its prior context. The conversation " - "may be locked by a live process or its files corrupt; " - "re-spawn with a fresh task carrying a summary." - ) - # Detect CC provider to skip permission event loop is_cc = self._manager._is_cc_provider(client) + + # Capture cleanup identity immediately after successful session + # acquisition. Every later step can fail and tombstone the run, so + # delaying this until the state write leaves provider files unidentified. + try: + cleanup_session_id = ( + str(client.session_id or "") if hasattr(client, "session_id") else "" + ) + cleanup_provider = self._manager._provider_label_of(client) + cleanup_cwd = "" + if is_cc: + cleanup_cwd = info.cwd + if not cleanup_cwd: + inner = getattr(client, "client", None) + work_dir = getattr(inner, "_work_dir", None) + if work_dir: + cleanup_cwd = str(work_dir) + setattr(info, "_session_id", cleanup_session_id) + setattr(info, "_session_provider", cleanup_provider) + setattr(info, "_session_cwd", cleanup_cwd) + self._publish_identity( + info.id, + session_id=cleanup_session_id, + provider=cleanup_provider, + cwd=cleanup_cwd, + keep=info.keep, + conversation_key=session_key if info.keep else "", + ) + except Exception: + logger.debug("Failed to capture live cleanup identity for %s", info.id, exc_info=True) + + # Fail CLOSED on a continuation that did not actually resume. Identity is + # already captured so the abnormal tombstone can reclaim the fresh session. + if info.conversation_key and not _resumed: + raise RuntimeError( + "resume_failed: session/load did not restore conversation " + f"{info.conversation_key} — refusing to execute the " + "follow-up without its prior context. The conversation " + "may be locked by a live process or its files corrupt; " + "re-spawn with a fresh task carrying a summary." + ) # Intentionally check info.agent (not resolved `agent`) so only # explicitly requested agents skip _SYSTEM_PREFIX (defense-in-depth). named_agent = bool(info.agent and _AGENT_NAME_RE.fullmatch(info.agent)) @@ -847,6 +934,22 @@ async def _run_inner_impl(self, info: SubagentInfo, session_key: str) -> None: logger.debug("Provenance write skipped (unreadable state) for %s", info.id) except Exception: logger.debug("Failed to persist model provenance for %s", info.id, exc_info=True) + # The live generation was published synchronously at acquisition so + # terminal tombstones are cancellation-safe. Persist the sidecar only + # after the drained provenance write: cancellation during provenance must + # not prevent its required model fields from landing, and cancellation + # here still leaves the live tombstone snapshot complete. + try: + await self._remember_identity_off_loop( + info, + session_id=str(getattr(info, "_session_id", "")), + provider=str(getattr(info, "_session_provider", "")), + cwd=str(getattr(info, "_session_cwd", "")), + keep=info.keep, + conversation_key=session_key if info.keep else "", + ) + except Exception: + logger.debug("Failed to persist cleanup identity for %s", info.id, exc_info=True) await self._manager._fire_event( "subagent_spawn", info, @@ -882,44 +985,23 @@ async def _run_inner_impl(self, info: SubagentInfo, session_key: str) -> None: except Exception: logger.debug("Failed to record PID for %s", info.id, exc_info=True) - # Record session_id and provider type for session file cleanup + # Persist the cleanup identity captured immediately after session + # acquisition, together with mutable retention intent. try: - session_id = client.session_id if hasattr(client, "session_id") else "" - provider_type = self._manager._provider_label_of(client) state_update: dict[str, object] = { - "session_id": session_id, - "provider": provider_type, + "session_id": str(getattr(info, "_session_id", "")), + "provider": str(getattr(info, "_session_provider", "")), # Model provenance (requested_model/resolved_model) is NOT # re-written here: the crash-safe write BEFORE the # subagent_spawn event above is the single owner of those two # fields on the spawn path, and a transient failure there is - # handled by that write's own bounded retry (#5394). This write - # still performs the same read-merge-rewrite either way, so the - # point is one authoritative writer, not saved I/O. The CC-path - # refinement below still updates resolved_model when it first - # becomes known. - # keep marks this run's session files as resume material: the - # orphan reconciler and tombstone pruner skip file deletion - # for keep runs (restart-safe — read from disk, not memory). + # handled by that write's own bounded retry (#5394). "keep": info.keep, "conversation_key": session_key if info.keep else "", } - # Store CWD for CC cleanup (needed to derive project-key path). - # info.cwd is only set when a caller passes an explicit cwd - # override (disabled by default), so for the common case derive - # the project dir from the provider's own work dir — that is the - # same path sent as ACP `cwd`, hence the encoded project key under - # ~/.claude/projects. Without this, CC cleanup is skipped (no cwd) - # and the transcript leaks. - if is_cc: - cc_cwd = info.cwd - if not cc_cwd: - inner = getattr(client, "client", None) - work_dir = getattr(inner, "_work_dir", None) - if work_dir: - cc_cwd = str(work_dir) - if cc_cwd: - state_update["cwd"] = cc_cwd + cleanup_cwd = str(getattr(info, "_session_cwd", "")) + if cleanup_cwd: + state_update["cwd"] = cleanup_cwd # Same off-loop, drained write as the PID record above (#6288, # #7302). This one also carries `keep`, the field the two remaining # on-loop writers (promote / release) contend for -- taking the @@ -1656,7 +1738,10 @@ def _should_use_session_sharing_impl(self, info: SubagentInfo) -> bool: return self._manager._sessions.is_session_sharing_eligible(info.parent_session_key) async def _create_shared_session_impl( - self, info: SubagentInfo, session_key: str, agent: str + self, + info: SubagentInfo, + session_key: str, + agent: str, ) -> "LLMProvider": """Create a subagent session on the parent's AcpRuntime. @@ -1667,6 +1752,7 @@ async def _create_shared_session_impl( AcpSessionProvider. Marks info._session_sharing=True so cleanup calls provider.shutdown() instead of SessionManager.release/reset. """ + runtime = self._manager._get_parent_runtime(info.parent_session_key) if runtime is None: runtime = await self._manager._sessions.get_subagent_runtime(info.parent_session_key) @@ -1677,21 +1763,54 @@ async def _create_shared_session_impl( agent=agent or None, ) provider = AcpSessionProvider(handle, runtime) - # This consumer implements the low-fidelity child downgrade (interactive - # approver when configured, reject when headless) — opt in so the - # handle-level fail-close gate yields those events instead of rejecting. + # The handle exists now. Publish ownership before any cancellable await so + # force-reap always takes the shared-session branch and destroys this handle + # instead of resetting a nonexistent dedicated session. provider.child_fidelity_aware = True info._session_sharing = True info._shared_provider = provider + # Capture cleanup identity before persistence or later setup can fail, + # otherwise the live handle becomes an untracked ghost. + cleanup_session_id = str(handle.session_id or "") + cleanup_provider = PROVIDER_LABEL_DEFAULT + setattr(info, "_session_id", cleanup_session_id) + setattr(info, "_session_provider", cleanup_provider) + self._publish_identity( + info.id, + session_id=cleanup_session_id, + provider=cleanup_provider, + keep=info.keep, + conversation_key=session_key if info.keep else "", + ) + try: + await self._remember_identity_off_loop( + info, + session_id=cleanup_session_id, + provider=cleanup_provider, + keep=info.keep, + conversation_key=session_key if info.keep else "", + ) + except (OSError, ValueError, RecursionError): + logger.debug( + "Shared-session identity persistence failed for %s", + info.id, + exc_info=True, + ) if runtime.pid: info._pid = runtime.pid - # Same off-loop, drained write as the non-shared spawn path's PID - # record (#6288, #7302). Unguarded here as before: this method has no - # best-effort contract, so an _atomic_write failure still propagates - # to the caller rather than yielding a session with no recorded pid. - await self._manager._write_state_off_loop( - info, "PID record", pid=runtime.pid, pid_recorded_at=time.time() - ) + try: + # Keep the shared handle alive on a storage error, but route the + # write through the run-owned off-loop drain so cancellation + # cannot detach a stale whole-file writer (#6288, #7302). + await self._manager._write_state_off_loop( + info, "PID record", pid=runtime.pid, pid_recorded_at=time.time() + ) + except Exception: + logger.debug( + "Shared-session PID persistence failed for %s", + info.id, + exc_info=True, + ) logger.info( "Subagent %s using session sharing on runtime PID %s (session %s, key %s)", info.id, diff --git a/src/kiro_crew/subagent_persistence.py b/src/kiro_crew/subagent_persistence.py index f24de9ec216..26084bce9f3 100644 --- a/src/kiro_crew/subagent_persistence.py +++ b/src/kiro_crew/subagent_persistence.py @@ -17,8 +17,11 @@ import threading import time import weakref +from enum import Enum from pathlib import Path +from typing import Callable +from kiro_crew import platform_compat from kiro_crew.acp.types import PROVIDER_LABEL_DEFAULT from kiro_crew.config.paths import data_home, kiro_sessions_dir from kiro_crew.jsonl_util import rotate_jsonl_at @@ -33,6 +36,265 @@ # dashboard/handlers/usage.py is the reference implementation. _SUBAGENTS_DIR: Path | None = None +# Promotion and prune arbitrate per agent; the weak lock registry is defined +# beside the state-writer registry below so unrelated conversations never couple. +SUBAGENT_CONVERSATION_PREFIX = "subagent:" +_CLEANUP_IDENTITIES_FILE = "cleanup-identities.json" +_CLEANUP_IDENTITIES_TRUST_DIR = "subagent-cleanup-identities" +_CLEANUP_IDENTITY_LOCK = threading.Lock() +_LIVE_CLEANUP_IDENTITIES: dict[str, list[dict[str, object]]] = {} +_LIVE_CLEANUP_HINTS: set[str] = set() + + +class RetentionPromotionResult(Enum): + """Outcome of the non-blocking promotion transaction.""" + + PROMOTED = "promoted" + RETRYABLE = "retryable" + + +def _cleanup_identities_path(agent_id: str) -> Path: + # Validate with the canonical agent-directory guard, but keep this durable + # cleanup authority OUTSIDE the agent-writable run folder. ``trust`` is on + # the shared file gate's read+write sensitive floor, so a subagent cannot + # replace another session ID and trick prune into deleting its transcript. + _agent_dir(agent_id) + return ( + _subagents_dir().parent + / "trust" + / _CLEANUP_IDENTITIES_TRUST_DIR + / agent_id + / _CLEANUP_IDENTITIES_FILE + ) + + +def _protect_cleanup_identities_path(agent_id: str) -> Path: + """Return the cleanup record only after fail-loud owner-only lockdown.""" + path = _cleanup_identities_path(agent_id) + for protected_dir in (path.parents[2], path.parents[1], path.parent): + platform_compat.make_owner_only_dir(protected_dir) + # ``make_owner_only_dir`` is best-effort by contract. Cleanup identity + # authorizes transcript deletion, so failure here must abort the read or + # write on Windows as well as POSIX instead of trusting a permissive ACL. + platform_compat.restrict_dir_to_owner(protected_dir) + if path.exists(): + # Tightening a parent does not retrofit an existing Windows file DACL. + platform_compat.restrict_to_owner(path) + return path + + +def _delete_cleanup_identities_file(agent_id: str) -> None: + """Remove the protected generation record after its run folder is gone.""" + shutil.rmtree(_cleanup_identities_path(agent_id).parent, ignore_errors=True) + + +def _read_cleanup_identities_file(agent_id: str) -> list[dict[str, object]]: + path = _protect_cleanup_identities_path(agent_id) + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return [] + if not isinstance(payload, dict) or not isinstance(payload.get("identities"), list): + raise ValueError("invalid protected cleanup identity payload") + raw = payload["identities"] + records: list[dict[str, object]] = [] + for item in raw: + if not isinstance(item, dict): + raise ValueError("invalid protected cleanup identity record") + sid = item.get("session_id") + if not isinstance(sid, str) or not sid: + raise ValueError("invalid protected cleanup identity SID") + record: dict[str, object] = {"session_id": sid} + provider = item.get("provider") + cwd = item.get("cwd") + keep = item.get("keep") + conversation_key = item.get("conversation_key") + if "provider" in item and (not isinstance(provider, str) or not provider): + raise ValueError("invalid protected cleanup identity provider") + if "cwd" in item and (not isinstance(cwd, str) or not cwd): + raise ValueError("invalid protected cleanup identity CWD") + if "keep" in item and not isinstance(keep, bool): + raise ValueError("invalid protected cleanup identity retention") + if "conversation_key" in item and ( + not isinstance(conversation_key, str) or not conversation_key + ): + raise ValueError("invalid protected cleanup identity owner") + if isinstance(provider, str): + record["provider"] = provider + if isinstance(cwd, str): + record["cwd"] = cwd + if isinstance(keep, bool): + record["keep"] = keep + if isinstance(conversation_key, str): + record["conversation_key"] = conversation_key + records.append(record) + return records + + +def _merge_cleanup_identity_records( + *groups: object, +) -> list[dict[str, object]]: + """Merge identity groups left-to-right, keyed by SID with richer fields kept.""" + by_sid: dict[str, dict[str, object]] = {} + for group in groups: + if not isinstance(group, (list, tuple)): + continue + for item in group: + if not isinstance(item, dict): + continue + sid = item.get("session_id") + if not isinstance(sid, str) or not sid: + continue + record = by_sid.setdefault(sid, {"session_id": sid}) + provider = item.get("provider") + cwd = item.get("cwd") + keep = item.get("keep") + conversation_key = item.get("conversation_key") + if isinstance(provider, str) and provider: + record["provider"] = provider + if isinstance(cwd, str) and cwd: + record["cwd"] = cwd + if isinstance(keep, bool): + record["keep"] = keep + if isinstance(conversation_key, str) and conversation_key: + record["conversation_key"] = conversation_key + return list(by_sid.values()) + + +def _cleanup_identity_record( + *, + session_id: str, + provider: str, + cwd: str, + keep: bool | None, + conversation_key: str, +) -> dict[str, object]: + record: dict[str, object] = {"session_id": session_id} + if provider: + record["provider"] = provider + if cwd: + record["cwd"] = cwd + if isinstance(keep, bool): + record["keep"] = keep + if conversation_key: + record["conversation_key"] = conversation_key + return record + + +def publish_live_cleanup_identity( + agent_id: str, + *, + session_id: str = "", + provider: str = "", + cwd: str = "", + keep: bool | None = None, + conversation_key: str = "", +) -> None: + """Publish one generation in memory without waiting or filesystem I/O.""" + if not session_id: + return + record = _cleanup_identity_record( + session_id=session_id, + provider=provider, + cwd=cwd, + keep=keep, + conversation_key=conversation_key, + ) + # list.append is atomic under the interpreter lock. Deliberately do not + # acquire _CLEANUP_IDENTITY_LOCK: its owner may be blocked in fsync, while + # this path runs on the gateway event loop and must publish before a queued + # to_thread call can be cancelled. Deduplication belongs to snapshots and + # sidecar serialization; replacing the list here can lose a concurrent SID. + records = _LIVE_CLEANUP_IDENTITIES.setdefault(agent_id, []) + records.append(record) + + +def remember_live_cleanup_identity( + agent_id: str, + *, + session_id: str = "", + provider: str = "", + cwd: str = "", + keep: bool | None = None, + conversation_key: str = "", +) -> None: + """Durably append one complete live session generation for later cleanup.""" + if not session_id: + return + publish_live_cleanup_identity( + agent_id, + session_id=session_id, + provider=provider, + cwd=cwd, + keep=keep, + conversation_key=conversation_key, + ) + with _CLEANUP_IDENTITY_LOCK: + durable = _read_cleanup_identities_file(agent_id) + records = _merge_cleanup_identity_records( + durable, + _LIVE_CLEANUP_IDENTITIES.get(agent_id, []), + ) + path = _protect_cleanup_identities_path(agent_id) + _atomic_write(path, {"identities": records}) + platform_compat.restrict_to_owner(path) + + +def _live_cleanup_identities(agent_id: str) -> list[dict[str, object]]: + """Snapshot already-published identities without performing filesystem I/O.""" + if not _CLEANUP_IDENTITY_LOCK.acquire(blocking=False): + # The writer publishes its merged in-memory fallback before fsync, so a + # contending event-loop tombstone can snapshot it without waiting. + fallback = _LIVE_CLEANUP_IDENTITIES.get(agent_id, []) + return _merge_cleanup_identity_records(fallback) + try: + fallback = _LIVE_CLEANUP_IDENTITIES.get(agent_id, []) + return _merge_cleanup_identity_records(fallback) + finally: + _CLEANUP_IDENTITY_LOCK.release() + + +def publish_live_cleanup_hint(agent_id: str) -> None: + """Mark run identity as cleanup-owned without trusting agent SID fields.""" + _LIVE_CLEANUP_HINTS.add(agent_id) + + +def trusted_cleanup_identity_record( + agent_id: str, + session_id: str, + conversation_key: str, +) -> dict[str, object] | None: + """Return the trusted generation matching restart state, or ``None``. + + Agent-folder state may request registry rebuild, but it cannot choose the SID, + owner, provider, or CWD that the TTL release path will later clean. Those + fields must match/source from gateway-published live or protected authority. + """ + default_key = f"{SUBAGENT_CONVERSATION_PREFIX}{agent_id}" + records = _merge_cleanup_identity_records( + _read_cleanup_identities_file(agent_id), + _live_cleanup_identities(agent_id), + ) + for record in records: + record_sid = record.get("session_id") + record_key = record.get("conversation_key") or default_key + if record_sid == session_id and record_key == conversation_key: + return record + return None + + +def has_live_cleanup_identity(agent_id: str) -> bool: + """Return whether cleanup retention is hinted, without disk I/O.""" + return agent_id in _LIVE_CLEANUP_HINTS or bool(_live_cleanup_identities(agent_id)) + + +def subagent_id_from_conversation_key(key: str) -> str | None: + """Return the subagent owner ID encoded by *key*, or ``None``.""" + if not key.startswith(SUBAGENT_CONVERSATION_PREFIX): + return None + owner_id = key[len(SUBAGENT_CONVERSATION_PREFIX) :] + return owner_id or None + def _subagents_dir() -> Path: """Subagents registry directory, resolved against the live data home.""" @@ -133,15 +395,23 @@ def create_agent_folder( def read_state(agent_id: str) -> dict | None: - """Read state.json. Returns None on missing/corrupt.""" + """Read state.json. Returns None on missing, corrupt, or non-object data.""" try: p = _agent_dir(agent_id) / "state.json" - except ValueError: + state = json.loads(p.read_text(encoding="utf-8")) + except (OSError, ValueError, RecursionError): return None + return state if isinstance(state, dict) else None + + +def read_tombstone(agent_id: str) -> dict | None: + """Read tombstone.json as an object. Return None on missing/invalid data.""" try: - return json.loads(p.read_text(encoding="utf-8")) - except (FileNotFoundError, json.JSONDecodeError, OSError): + p = _agent_dir(agent_id) / "tombstone.json" + tombstone = json.loads(p.read_text(encoding="utf-8")) + except (OSError, ValueError, RecursionError): return None + return tombstone if isinstance(tombstone, dict) else None # ── per-agent write serialization ──────────────────────────────────── @@ -163,21 +433,19 @@ def read_state(agent_id: str) -> dict | None: #: await DETACHES the worker rather than stopping it, so it finishes carrying a #: read that is already stale (#6308). #: -#: SCOPE -- OFF-LOOP WRITERS ONLY. Serializing every writer would mean a -#: loop-side caller waiting on a pool thread's fsync, i.e. a new blocking call -#: on the event loop, which this repo's ``no-blocking-call-on-event-loop`` anchor -#: forbids and which a bounded wait only shrinks rather than removes. So the lock -#: is taken only by callers that are NOT on the loop, where blocking is legal; -#: on-loop callers keep exactly their pre-existing behaviour (an unserialized -#: rewrite). That closes pool-vs-pool interleaving completely and leaves the -#: loop-side half untouched -- see :func:`update_state` for what remains open and -#: the issue tracking the loop-side offload (#6288's class). +#: SCOPE -- ordinary ``update_state`` callers take the lock OFF-LOOP only. +#: Serializing every loop-side write by waiting would block the event loop behind +#: a pool thread's fsync, which the no-blocking anchor forbids. Retention promotion +#: instead probes the same per-agent lock non-blocking and returns RETRYABLE when +#: busy; once acquired, its existing on-loop keep write cannot be overwritten by +#: an older pool writer. Other on-loop callers keep their pre-existing unlocked +#: behavior -- see :func:`update_state` for the remaining #6308 limitation. #: -#: The acquire is UNBOUNDED, and can be, precisely because no on-loop caller ever -#: waits on it: the only threads that block here are pool workers whose own write -#: (read + fsync + rename) already exposes them to a wedged FS, so waiting on the -#: lock adds no parking risk the write itself did not already carry. No holder is -#: ever an on-loop caller. +#: The ordinary acquire is UNBOUNDED, and can be, because no on-loop caller reaches +#: it: only pool workers block there, and their own read + fsync + rename already +#: exposes them to a wedged filesystem. Promotion may hold the same lock around its +#: existing loop-side write, but its acquire is always non-blocking and never parks +#: the event loop. #: #: In-process only, mirroring the per-key ``threading.Lock`` registry this repo #: already uses to serialize file read-modify-write (``learn._lock_for``, @@ -201,6 +469,8 @@ def read_state(agent_id: str) -> dict | None: #: it -- and agent ids are per-run uuids, so nothing accumulates either. _STATE_LOCKS: "weakref.WeakValueDictionary[str, _AgentLock]" = weakref.WeakValueDictionary() _STATE_LOCKS_GUARD = threading.Lock() +_RETENTION_LOCKS: "weakref.WeakValueDictionary[str, _AgentLock]" = weakref.WeakValueDictionary() +_RETENTION_LOCKS_GUARD = threading.Lock() class _AgentLock: @@ -236,20 +506,65 @@ def _on_event_loop() -> bool: return True -def _lock_for_agent(agent_id: str) -> "_AgentLock": - """Return the process-wide ``state.json`` lock holder for *agent_id*. - - The caller MUST keep the returned object alive for its whole critical - section -- see :data:`_STATE_LOCKS`. - """ - with _STATE_LOCKS_GUARD: - holder = _STATE_LOCKS.get(agent_id) +def _lock_for_registry(agent_id, registry, guard): # type: ignore[no-untyped-def] + with guard: + holder = registry.get(agent_id) if holder is None: holder = _AgentLock() - _STATE_LOCKS[agent_id] = holder + registry[agent_id] = holder return holder +def _try_acquire_registry_lock(agent_id, registry, guard): # type: ignore[no-untyped-def] + if not guard.acquire(blocking=False): + return None + try: + holder = registry.get(agent_id) + if holder is None: + holder = _AgentLock() + registry[agent_id] = holder + finally: + guard.release() + if not holder.lock.acquire(blocking=False): + return None + return holder + + +def _lock_for_agent(agent_id: str) -> "_AgentLock": + """Return the process-wide ``state.json`` lock holder for *agent_id*.""" + return _lock_for_registry(agent_id, _STATE_LOCKS, _STATE_LOCKS_GUARD) + + +def _try_acquire_state_lock(agent_id: str) -> "_AgentLock | None": + """Return the held per-agent writer lock, or None without blocking.""" + return _try_acquire_registry_lock(agent_id, _STATE_LOCKS, _STATE_LOCKS_GUARD) + + +def _retention_lock_for_agent(agent_id: str) -> "_AgentLock": + """Return the process-wide retention arbitration holder for *agent_id*.""" + return _lock_for_registry(agent_id, _RETENTION_LOCKS, _RETENTION_LOCKS_GUARD) + + +def _try_acquire_retention_lock(agent_id: str) -> "_AgentLock | None": + """Return held per-agent retention arbitration, or None without blocking.""" + return _try_acquire_registry_lock( + agent_id, _RETENTION_LOCKS, _RETENTION_LOCKS_GUARD + ) + + +def _acquire_retention_locks(*agent_ids: str) -> list["_AgentLock"]: + """Acquire retention locks in stable order for off-loop prune work.""" + holders = [_retention_lock_for_agent(agent_id) for agent_id in sorted(set(agent_ids))] + for holder in holders: + holder.lock.acquire() + return holders + + +def _release_retention_locks(holders: list["_AgentLock"]) -> None: + for holder in reversed(holders): + holder.lock.release() + + def update_state(agent_id: str, **fields: object) -> bool: """Merge *fields* into state.json (atomic rewrite). @@ -264,22 +579,19 @@ def update_state(agent_id: str, **fields: object) -> bool: :data:`_STATE_LOCKS`), so two pool writers can no longer rewrite a snapshot that predates the other's write. - KNOWN LIMITATION: an ON-LOOP caller does not take the lock, because waiting + KNOWN LIMITATION: ordinary ON-LOOP callers do not take the lock, because waiting on a pool thread's fsync from the event loop is exactly the blocking call the - repo's anchor forbids. What closes the interleave for those callers instead - is that every off-loop writer inside a RUN is drained on cancellation - (``_write_state_off_loop``, #6298 / #6308): a pool writer cannot outlive the - run it belongs to, and the on-loop retention (``keep``) writes are reached - only through a ``_conversation_busy`` gate that refuses while a run is in - flight. The drain is bounded, so on a wedged FS a worker IS abandoned -- but - that same gate then holds the conversation until the abandoned worker - settles, so the two on-loop ``keep`` writes are deferred past it rather than - silently undone. Separately, an on-loop caller still pays this function's - fsync ON the loop. Every writer inside a run now goes off-loop through the - helper (#7302); the two that remain on the loop are the retention writers - ``_promote_conversation`` / ``release_conversation``, which are synchronous - functions whose other work (the ``SessionMap`` mutation) is required to stay - on the loop -- moving those is the rest of #7302. + repo's anchor forbids. Every writer inside a run now goes off-loop through + ``_write_state_off_loop`` and is drained on cancellation (#6298 / #6308 / + #7302); an abandoned writer holds the conversation until it settles, so the + on-loop retention writes are deferred past it. Retention promotion adds a + second defense: on the event loop it probes the same per-agent lock + non-blocking and returns RETRYABLE on contention, while off-loop promotion + lets ``update_state`` acquire the lock normally. Thus no stale writer can + roll back ``keep=True`` and no loop-side caller waits for a pool writer's + fsync. The remaining on-loop callers are the synchronous retention writers; + they still pay their own fsync on the loop, and moving that I/O while keeping + their ``SessionMap`` mutation on-loop is the rest of #7302. """ p = _agent_dir(agent_id) / "state.json" # Off-loop callers serialize; on-loop callers keep pre-existing behaviour. @@ -292,9 +604,12 @@ def update_state(agent_id: str, **fields: object) -> bool: try: try: state = json.loads(p.read_text(encoding="utf-8")) - except (FileNotFoundError, json.JSONDecodeError, OSError): + except (OSError, ValueError, RecursionError): logger.debug("update_state: cannot read state for %s, skipping", agent_id) return False + if not isinstance(state, dict): + logger.debug("update_state: non-object state for %s, skipping", agent_id) + return False state.update(fields) state["updated_at"] = time.time() _atomic_write(p, state) @@ -304,6 +619,39 @@ def update_state(agent_id: str, **fields: object) -> bool: return True +def promote_retention( + agent_id: str, + *, + state_writer: Callable[..., bool] = update_state, +) -> RetentionPromotionResult: + """Atomically promote *agent_id* against concurrent in-process prune. + + Promotion acquires per-agent retention arbitration and, on the event loop, + probes the state-writer lock non-blocking. Contention returns ``RETRYABLE`` + without parking the loop. Off-loop callers let ``update_state`` acquire its + normal non-reentrant state lock. Thus promotion either writes ``keep=True`` + before prune's locked re-read or waits for a later retry; a process crash + leaves no half-committed claim format, and the next prune re-evaluates state. + """ + retention_holder = _try_acquire_retention_lock(agent_id) + if retention_holder is None: + return RetentionPromotionResult.RETRYABLE + state_holder = None + if _on_event_loop(): + state_holder = _try_acquire_state_lock(agent_id) + if state_holder is None: + retention_holder.lock.release() + return RetentionPromotionResult.RETRYABLE + try: + if not state_writer(agent_id, keep=True): + return RetentionPromotionResult.RETRYABLE + return RetentionPromotionResult.PROMOTED + finally: + if state_holder is not None: + state_holder.lock.release() + retention_holder.lock.release() + + # ── result streaming ───────────────────────────────────────────────── @@ -338,6 +686,16 @@ def write_tombstone( """Write ``tombstone.json`` for an abnormally exited agent.""" d = _agent_dir(agent_id) state = read_state(agent_id) or {} + cleanup_identity = { + key: state[key] + for key in ("session_id", "provider", "cwd") + if state.get(key) + } + live_cleanup_identities = _live_cleanup_identities(agent_id) + latest_live_identity = live_cleanup_identities[-1] if live_cleanup_identities else {} + generation_metadata: dict[str, object] = {} + if live_cleanup_identities: + generation_metadata["cleanup_identities"] = live_cleanup_identities tombstone = { "id": agent_id, "task": state.get("task", ""), @@ -349,12 +707,21 @@ def write_tombstone( "recovery_action": recovery_action, "result_available": _check_result_available(d / "result.txt"), "result_path": str(d / "result.txt"), + **cleanup_identity, + **latest_live_identity, + **generation_metadata, **extra, } try: _atomic_write(d / "tombstone.json", tombstone) except OSError: logger.warning("write_tombstone failed for %s", agent_id, exc_info=True) + state_sid = cleanup_identity.get("session_id") + if not live_cleanup_identities and isinstance(state_sid, str) and state_sid: + # The run folder is agent-writable. Preserve only the retention/exemption + # signal here; provider-deletion authority requires gateway-published live + # identity or the protected durable record. + publish_live_cleanup_hint(agent_id) def mark_delivered(agent_id: str) -> None: @@ -445,9 +812,14 @@ def record_slow_command(agent_id: str, **fields: object) -> None: def delete_agent_folder(agent_id: str) -> None: - """Remove the entire agent directory.""" + """Remove the entire agent directory and its in-process identity fallback.""" d = _agent_dir(agent_id) - shutil.rmtree(d, ignore_errors=True) + with _CLEANUP_IDENTITY_LOCK: + shutil.rmtree(d, ignore_errors=True) + if not d.exists(): + _delete_cleanup_identities_file(agent_id) + _LIVE_CLEANUP_IDENTITIES.pop(agent_id, None) + _LIVE_CLEANUP_HINTS.discard(agent_id) # ── list orphans ───────────────────────────────────────────────────── @@ -475,6 +847,118 @@ def list_orphans() -> list[dict]: # ── prune ──────────────────────────────────────────────────────────── +# Preserve unknown retention state long enough to outlive the six-hour +# continuable-conversation window, then reclaim it using tombstone metadata. +_UNREADABLE_STATE_GRACE_SECS = 24 * 3600 +_UNRECLAIMABLE_LOOKUP_MAX_AGE_SECS = 90 * 86400 + + +def _tombstone_died(ts: dict[str, object], path: Path, now: float) -> int | float: + """Return a finite, positive, non-future death time with bounded fallback.""" + died = ts.get("died") + if ( + isinstance(died, (int, float)) + and not isinstance(died, bool) + and 0 < died <= now + ): + return died + try: + fallback = path.stat().st_mtime + except OSError: + return now + if 0 < fallback <= now: + return fallback + return now + + +def _should_defer_tombstone_cleanup( + *, + retention_state: dict[str, object], + retention_unknown: bool, + cleanup_session_id: object, + died: object, + cutoff: float, + now: float, +) -> bool: + """Return whether prune must preserve provider files and identity folder.""" + if retention_unknown: + if not cleanup_session_id or not ( + isinstance(died, (int, float)) + and not isinstance(died, bool) + and 0 < died <= now + ): + return False + return died >= cutoff - _UNREADABLE_STATE_GRACE_SECS + return retention_state.get("keep") is True + + +def _cleanup_identity_fallback_record( + agent_id: str, session_id: object +) -> dict[str, object] | None: + """Return the matching generation, or latest when top-level SID is absent.""" + records = _merge_cleanup_identity_records( + _read_cleanup_identities_file(agent_id), + _live_cleanup_identities(agent_id), + ) + if not records: + return None + if isinstance(session_id, str) and session_id: + for record in reversed(records): + if record.get("session_id") == session_id: + return record + # Agent-writable state cannot suppress protected retention authority by + # naming an unrelated SID. Fall back to the latest trusted generation; + # this record can preserve ownership but never expands deletion authority. + return records[-1] + + +def _cleanup_retention_fallback( + agent_id: str, session_id: object +) -> tuple[bool | None, str, str]: + """Return trusted fallback retention, owner, and SID.""" + record = _cleanup_identity_fallback_record(agent_id, session_id) + if record is None: + return None, "", "" + keep = record.get("keep") + conversation_key = record.get("conversation_key") + record_sid = record.get("session_id") + return ( + keep if isinstance(keep, bool) else None, + conversation_key if isinstance(conversation_key, str) else "", + record_sid if isinstance(record_sid, str) else "", + ) + + +def _tombstone_cleanup_identities(agent_id: str) -> list[tuple[str, str, str]]: + """Return gateway-authorized cleanup generations only. + + Tombstone and state identity fields live in the agent-writable run folder. + They remain compatibility/display hints, but cannot expand the set of provider + transcripts prune may delete. Durable protected records and synchronous live + gateway publication are the only deletion authorities. + """ + records = _merge_cleanup_identity_records( + _read_cleanup_identities_file(agent_id), + _live_cleanup_identities(agent_id), + ) + identities: list[tuple[str, str, str]] = [] + for identity_record in records: + sid = identity_record.get("session_id") + if not isinstance(sid, str) or not sid: + continue + record_provider = identity_record.get("provider") + record_cwd = identity_record.get("cwd") + identities.append( + ( + sid, + record_provider + if isinstance(record_provider, str) and record_provider + else PROVIDER_LABEL_DEFAULT, + record_cwd if isinstance(record_cwd, str) else "", + ) + ) + return identities + def prune_stale_tombstones(max_age_days: int = 7, delivered_ttl_secs: int = 3600) -> int: """Delete tombstoned folders past their retention window. Returns count pruned. @@ -499,27 +983,190 @@ def prune_stale_tombstones(max_age_days: int = 7, delivered_ttl_secs: int = 3600 if not ts_path.exists(): continue try: - ts = json.loads(ts_path.read_text(encoding="utf-8")) + ts = read_tombstone(d.name) + if ts is None: + logger.debug("prune: skipping corrupt tombstone in %s", d.name) + continue cutoff = delivered_cutoff if ts.get("cause") == "delivered" else default_cutoff - if ts.get("died", 0) < cutoff: - # Best-effort session cleanup — must not block folder removal + died = _tombstone_died(ts, ts_path, now) + if died <= cutoff: + tombstone_session_id = ts.get("session_id", "") + state_hint = read_state(d.name) + claim_hint = d.name + hinted_key = "" + if isinstance(state_hint, dict): + hinted_session_id = state_hint.get("session_id", "") + _, fallback_key, fallback_sid = _cleanup_retention_fallback( + d.name, + hinted_session_id, + ) + # Protected acquisition-time ownership outranks all + # agent-writable state, including an empty or conflicting + # conversation key. Legacy runs without a generation retain + # the compatibility hint from state. + hinted_key = ( + fallback_key + if fallback_sid + else str(state_hint.get("conversation_key") or "") + ) + else: + hinted_record = _cleanup_identity_fallback_record(d.name, "") + if hinted_record is not None: + hinted_key = str(hinted_record.get("conversation_key") or "") + hinted_owner = subagent_id_from_conversation_key(hinted_key) + if hinted_owner: + try: + _agent_dir(hinted_owner) + except ValueError: + pass + else: + claim_hint = hinted_owner + locked_agent_ids = {d.name, claim_hint} + retention_holders = _acquire_retention_locks(*locked_agent_ids) try: state = read_state(d.name) - session_id = ts.get("session_id") or (state.get("session_id", "") if state else "") - provider = ts.get("provider") or (state.get("provider", "acp") if state else "acp") - cwd = ts.get("cwd") or (state.get("cwd", "") if state else "") - # keep=True conversations retain their session files as - # resume material for spawn_continue; the conversation - # TTL sweep (SubagentManager reaper) owns their deletion. - _keep = bool(state.get("keep")) if state else False - if session_id and not _keep: - _cleanup_session_files_sync(session_id, provider, cwd=cwd) - except Exception: - logger.debug("prune: session cleanup failed for %s", d.name, exc_info=True) - shutil.rmtree(d, ignore_errors=True) - pruned += 1 - except (json.JSONDecodeError, OSError): - logger.debug("prune: skipping corrupt tombstone in %s", d.name) + claim_agent_id = d.name + lookup_session_id = tombstone_session_id + if isinstance(state, dict): + lookup_session_id = state.get("session_id") or tombstone_session_id + retention_unknown = False + retention_state: dict[str, object] = state + session_id = state.get("session_id", "") + # Every completed plain run records keep=False. A continuation + # follows readable original-owner state; unreadable owner + # intent receives bounded grace. + state_key = str(state.get("conversation_key") or "") + ( + fallback_keep, + fallback_key, + fallback_sid, + ) = _cleanup_retention_fallback( + d.name, + session_id, + ) + if fallback_sid: + # State is agent-writable. Once gateway-published + # identity exists, it cannot erase or redirect the + # continuation owner used for retention arbitration. + conversation_key = fallback_key + if not session_id: + session_id = fallback_sid + if "keep" not in state and fallback_keep is not None: + retention_state = dict(state) + retention_state["keep"] = fallback_keep + else: + conversation_key = state_key + owner_id = subagent_id_from_conversation_key(conversation_key) + if owner_id and owner_id != d.name: + try: + _agent_dir(owner_id) + except ValueError: + retention_unknown = True + else: + claim_agent_id = owner_id + owner_state = read_state(owner_id) + if isinstance(owner_state, dict): + retention_state = owner_state + else: + retention_unknown = True + else: + fallback_record = _cleanup_identity_fallback_record(d.name, "") + retention_state = {} + retention_unknown = True + session_id = "" + conversation_key = "" + if fallback_record is not None: + session_id = fallback_record.get("session_id", session_id) + sidecar_keep = fallback_record.get("keep") + if isinstance(sidecar_keep, bool): + retention_state["keep"] = sidecar_keep + # A durable false predates any later promotion + # that landed only in now-unreadable state. It + # cannot authorize immediate deletion; keep the + # bounded unknown-retention grace. True is safe + # to honor because it only preserves material. + if sidecar_keep: + retention_unknown = False + conversation_key = str( + fallback_record.get("conversation_key") or "" + ) + owner_id = subagent_id_from_conversation_key(conversation_key) + if owner_id and owner_id != d.name: + try: + _agent_dir(owner_id) + except ValueError: + retention_unknown = True + else: + claim_agent_id = owner_id + owner_state = read_state(owner_id) + if isinstance(owner_state, dict): + retention_state = owner_state + retention_unknown = False + else: + retention_unknown = True + if claim_agent_id not in locked_agent_ids: + continue + defer_cleanup = _should_defer_tombstone_cleanup( + retention_state=retention_state, + retention_unknown=retention_unknown, + cleanup_session_id=session_id, + died=died, + cutoff=cutoff, + now=now, + ) + if defer_cleanup: + continue + # Hold arbitration through cleanup and rmtree. A promotion + # arriving after the keep=False decision returns retryable + # instead of writing keep=True just before deletion. + cleanup_identities = _tombstone_cleanup_identities(d.name) + within_retry_window = ( + died >= now - _UNRECLAIMABLE_LOOKUP_MAX_AGE_SECS + ) + # Legacy/pre-upgrade runs can carry a SID only in the + # agent-writable state/tombstone. It is not safe deletion + # authority, but the folder is useful for a later trusted + # migration. Bound that lookup window so an unavailable + # migration cannot accumulate private run folders forever. + if ( + lookup_session_id + and not cleanup_identities + and within_retry_window + ): + continue + cleanup_succeeded = True + for cleanup_sid, cleanup_provider, cleanup_cwd in cleanup_identities: + try: + if ( + _cleanup_session_files_sync( + cleanup_sid, + cleanup_provider, + cwd=cleanup_cwd, + ) + is False + ): + cleanup_succeeded = False + except Exception: + cleanup_succeeded = False + logger.debug( + "prune: session cleanup failed for %s (%s)", + d.name, + cleanup_sid, + exc_info=True, + ) + if not cleanup_succeeded and within_retry_window: + continue + shutil.rmtree(d, ignore_errors=True) + if not d.exists(): + with _CLEANUP_IDENTITY_LOCK: + _delete_cleanup_identities_file(d.name) + _LIVE_CLEANUP_IDENTITIES.pop(d.name, None) + _LIVE_CLEANUP_HINTS.discard(d.name) + pruned += 1 + finally: + _release_retention_locks(retention_holders) + except (OSError, ValueError, RecursionError): + logger.debug("prune: tombstone processing failed for %s", d.name) return pruned @@ -528,21 +1175,26 @@ def prune_stale_tombstones(max_age_days: int = 7, delivered_ttl_secs: int = 3600 def _cleanup_session_files_sync( session_id: str, provider: str = PROVIDER_LABEL_DEFAULT, *, cwd: str = "" -) -> None: - """Delete LLM provider session files for a completed subagent. +) -> bool: + """Delete provider session files, returning whether cleanup completed. Synchronous — used during tombstone pruning (which runs in the reaper loop). - Best-effort: logs warnings on failure, never raises. + A false result preserves the run folder and protected identity record so a + future provider cleanup implementation or transient filesystem recovery can + retry instead of making the leaked transcript permanently unreachable. Only the kiro-cli backend stores transcripts where this function can reach them. Any other *provider* is logged and its files are left in place, since reporting success without deleting anything hides the leak. """ - if not session_id or session_id in (".", ".."): - return + if not session_id: + return True + if session_id in (".", ".."): + return False try: if provider == PROVIDER_LABEL_DEFAULT: sessions_dir = kiro_sessions_dir() + succeeded = True for suffix in (".json", ".jsonl"): target = sessions_dir / f"{session_id}{suffix}" if not _is_safe_path(target, sessions_dir): @@ -550,31 +1202,33 @@ def _cleanup_session_files_sync( "_cleanup_session_files_sync: path traversal blocked for %s", target, ) - return + return False try: target.unlink(missing_ok=True) except OSError: + succeeded = False logger.warning( "_cleanup_session_files_sync: failed to delete %s", target, exc_info=True, ) - else: - # Every other backend owns its own session storage, which this - # function has no route to. Say so rather than returning as if the - # files had been removed. - logger.debug( - "_cleanup_session_files_sync: no cleanup route for provider %s; " - "session %s files retained", - provider, - session_id, - ) + return succeeded + # Every other backend owns its own session storage, which this function + # has no route to. Report failure so prune retains retry metadata. + logger.debug( + "_cleanup_session_files_sync: no cleanup route for provider %s; " + "session %s files retained", + provider, + session_id, + ) + return False except Exception: logger.warning( "_cleanup_session_files_sync: unexpected error cleaning session %s", session_id, exc_info=True, ) + return False # ── helpers ────────────────────────────────────────────────────────── diff --git a/test/test_continuable_followups.py b/test/test_continuable_followups.py index b67de1d40e8..665ca0f17ab 100644 --- a/test/test_continuable_followups.py +++ b/test/test_continuable_followups.py @@ -22,7 +22,12 @@ import kiro_crew.subagent as subagent_mod from kiro_crew.subagent import SubagentInfo, SubagentManager -from kiro_crew.subagent_persistence import create_agent_folder, read_state, update_state +from kiro_crew.subagent_persistence import ( + create_agent_folder, + read_state, + remember_live_cleanup_identity, + update_state, +) # Subagent-registry isolation is provided globally by the autouse # ``_isolate_subagents_dir`` fixture in ``conftest.py``. @@ -154,6 +159,14 @@ def _keep_run(self, run_id: str, *, sid: str = "sid-x", conv_key: str = "") -> N cwd="", conversation_key=conv_key, ) + remember_live_cleanup_identity( + run_id, + session_id=sid, + provider="acp", + cwd="", + keep=True, + conversation_key=conv_key, + ) def test_scan_keep_states_finds_only_keep_runs(self) -> None: mgr = _manager() @@ -175,6 +188,56 @@ def test_scan_uses_recorded_conversation_key(self) -> None: found = mgr._scan_keep_states() assert found[0][1] == "subagent:origrun" + def test_agent_folder_identity_rehydrates_hint_without_deletion_authority( + self, + ) -> None: + import kiro_crew.subagent_persistence as sp + + mgr = _manager() + + state_run = "forged-state-hint" + create_agent_folder(state_run, task="t") + update_state(state_run, session_id="sid-victim-state", provider="acp") + sp.write_tombstone(state_run, cause="timeout", recovery_action="pending") + assert sp.has_live_cleanup_identity(state_run) is True + assert sp._tombstone_cleanup_identities(state_run) == [] + + tombstone_run = "forged-tombstone-hint" + create_agent_folder(tombstone_run, task="t") + sp.write_tombstone( + tombstone_run, + cause="timeout", + recovery_action="pending", + session_id="sid-victim-tombstone", + provider="acp", + ) + (sp._agent_dir(tombstone_run) / "state.json").write_text("{corrupt") + assert mgr._scan_keep_states() == [] + assert sp.has_live_cleanup_identity(tombstone_run) is True + assert sp._tombstone_cleanup_identities(tombstone_run) == [] + + for run_id, state_sid, state_key, state_keep in ( + ("forged-retained-sid", "sid-victim", "", True), + ("forged-retained-owner", "sid-owned", "subagent:victim-owner", True), + ("string-false-retention", "sid-owned", "", "false"), + ): + create_agent_folder(run_id, task="t") + remember_live_cleanup_identity( + run_id, + session_id="sid-owned", + provider="acp", + keep=True, + ) + update_state( + run_id, + session_id=state_sid, + provider="acp", + keep=state_keep, + conversation_key=state_key, + ) + + assert mgr._scan_keep_states() == [] + @pytest.mark.asyncio async def test_rebuild_seeds_map_registry_and_cache(self) -> None: sessions = _mock_sessions() @@ -224,6 +287,14 @@ def _write_keep_state(run_id: str, conv_key: str, ts: float) -> None: conversation_key=conv_key, updated_at=ts, ) p.write_text(_json.dumps(state), encoding="utf-8") + remember_live_cleanup_identity( + run_id, + session_id="sid-x", + provider="acp", + cwd="", + keep=True, + conversation_key=conv_key, + ) _write_keep_state("origrun", "", 1000.0) _write_keep_state("contrun", "subagent:origrun", 2000.0) @@ -267,11 +338,29 @@ def test_keep_recorded_on_disk(self) -> None: create_agent_folder("keeprun1", task="t") update_state("keeprun1", keep=True) create_agent_folder("plainrun", task="t") + create_agent_folder("stringfalse", task="t") + update_state("stringfalse", keep="false") assert mgr._keep_recorded_on_disk("subagent:keeprun1") is True assert mgr._keep_recorded_on_disk("subagent:plainrun") is False + assert mgr._keep_recorded_on_disk("subagent:stringfalse") is False assert mgr._keep_recorded_on_disk("subagent:missing") is False assert mgr._keep_recorded_on_disk("dashboard:tab1") is False + def test_unreadable_state_with_tombstone_identity_fails_safe(self) -> None: + from kiro_crew.subagent_persistence import _agent_dir, write_tombstone + + mgr = _manager() + create_agent_folder("corrupt-keep", task="t") + update_state("corrupt-keep", session_id="sid-safe", provider="acp", keep=True) + write_tombstone("corrupt-keep", cause="timeout", recovery_action="pending") + (_agent_dir("corrupt-keep") / "state.json").write_text("{corrupt") + + with patch( + "kiro_crew.subagent_persistence.read_tombstone", + side_effect=AssertionError("event-loop tombstone read"), + ): + assert mgr._keep_recorded_on_disk("subagent:corrupt-keep") is True + def test_manager_installs_fallback_on_sessions(self) -> None: sessions = _mock_sessions() mgr = _manager(sessions) diff --git a/test/test_session_cleanup.py b/test/test_session_cleanup.py index eeb7dc57b24..c248877f2d6 100644 --- a/test/test_session_cleanup.py +++ b/test/test_session_cleanup.py @@ -31,6 +31,7 @@ create_agent_folder, prune_stale_tombstones, read_state, + remember_live_cleanup_identity, update_state, write_tombstone, ) @@ -553,9 +554,15 @@ def test_pruning_deletes_session_files(self, tmp_path, agent_root, session_id): if d.exists(): shutil.rmtree(d) - # Create subagent folder with session_id in state + # Create a plain subagent folder with explicit non-retention and session identity. create_agent_folder(agent_id, task="old task") - update_state(agent_id, session_id=session_id, provider="acp") + update_state(agent_id, session_id=session_id, provider="acp", keep=False) + remember_live_cleanup_identity( + agent_id, + session_id=session_id, + provider="acp", + keep=False, + ) # Write an old tombstone (8 days ago) write_tombstone(agent_id, cause="timeout", recovery_action="delivered") diff --git a/test/test_session_sharing.py b/test/test_session_sharing.py index 3b2a6989004..01525284501 100644 --- a/test/test_session_sharing.py +++ b/test/test_session_sharing.py @@ -16,7 +16,7 @@ from kiro_crew.acp.session_provider import AcpSessionProvider from kiro_crew.acp.types import AcpEvent from kiro_crew.providers.base import EVENT_COMPLETE, EVENT_TEXT_CHUNK -from kiro_crew.subagent import SubagentManager +from kiro_crew.subagent import SubagentInfo, SubagentManager # ``SubagentManager.spawn`` refuses -- registering no task -- while the host # looks short of memory, which is the runner's state, not this test's input. @@ -281,6 +281,118 @@ async def test_shared_session_sets_flags(self): assert info._shared_provider is not None assert isinstance(info._shared_provider, AcpSessionProvider) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "identity_error", + [ + OSError("sidecar unavailable"), + ValueError("malformed protected record"), + RecursionError("nested protected record"), + ], + ) + async def test_shared_identity_persistence_failure_keeps_live_handle( + self, identity_error: Exception + ): + sessions = _mock_sessions(sharing_eligible=True) + manager = SubagentManager( + sessions=sessions, + ctx_builder=_mock_ctx_builder_auto(), + is_yolo=lambda: True, + ) + info = SubagentInfo( + id="shared-persist-fail", + task="t", + parent_session_key="dashboard:slot1", + ) + + with patch( + "kiro_crew.subagent.asyncio.to_thread", + AsyncMock(side_effect=identity_error), + ), patch("kiro_crew.subagent.update_state", side_effect=OSError("disk full")): + provider = await manager._create_shared_session( + info, + "subagent:shared-persist-fail", + "kirocrew", + ) + + assert info._session_sharing is True + assert info._shared_provider is provider + assert info._session_id == "shared-session-abc" + assert info._session_provider == "acp" + assert info._pid == 12345 + sessions.get_or_create.assert_not_awaited() + runtime = await sessions.get_subagent_runtime("dashboard:slot1") + runtime.create_session.assert_awaited_once() + await provider.shutdown() + + @pytest.mark.asyncio + async def test_cancelled_identity_writer_still_tombstones_live_session(self): + from kiro_crew.subagent_persistence import ( + _cleanup_identities_path, + create_agent_folder, + read_tombstone, + ) + + sessions = _mock_sessions(sharing_eligible=True) + manager = SubagentManager( + sessions=sessions, + ctx_builder=_mock_ctx_builder_auto(), + is_yolo=lambda: True, + ) + info = SubagentInfo( + id="shared-persist-cancel", + task="t", + parent_session_key="dashboard:slot1", + ) + create_agent_folder(info.id, task=info.task) + + # Model executor saturation: outer cancellation lands after the durable + # writer is submitted but before it starts. The awaiter must shield and + # drain that worker before propagating cancellation. + entered = asyncio.Event() + release = asyncio.Event() + + async def gated_to_thread(func, /, *args, **kwargs): # type: ignore[no-untyped-def] + entered.set() + await release.wait() + return func(*args, **kwargs) + + with patch( + "kiro_crew.subagent.asyncio.to_thread", + side_effect=gated_to_thread, + ): + task = asyncio.ensure_future( + manager._create_shared_session( + info, + "subagent:shared-persist-cancel", + "kirocrew", + ) + ) + await entered.wait() + task.cancel() + await asyncio.sleep(0) + assert not task.done(), "identity writer detached on cancellation" + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert _cleanup_identities_path(info.id).exists() + assert info._session_sharing is True + assert info._shared_provider is not None + manager._agents[info.id] = info + manager._tasks[info.id] = MagicMock(done=MagicMock(return_value=False)) + manager._tasks[info.id].cancel = MagicMock() + with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"): + await manager._force_reap(info.id, info, 1.0, reason="reaped") + + tombstone = read_tombstone(info.id) or {} + assert tombstone["session_id"] == "shared-session-abc" + assert tombstone["provider"] == "acp" + runtime = await sessions.get_subagent_runtime("dashboard:slot1") + handle = runtime.create_session.return_value + handle.destroy.assert_awaited_once() + sessions.reset.assert_not_awaited() + @pytest.mark.asyncio async def test_shared_session_cleanup_destroys_handle(self): """On completion, shared session calls provider.shutdown() not reset().""" diff --git a/test/test_subagent_continuable.py b/test/test_subagent_continuable.py index 5db67497fee..6703a780dc4 100644 --- a/test/test_subagent_continuable.py +++ b/test/test_subagent_continuable.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import threading import time from unittest.mock import AsyncMock, MagicMock, patch @@ -56,6 +57,7 @@ async def _empty_stream(*_args: object, **_kwargs: object): # type: ignore[no-u sessions.get_agent = MagicMock(return_value="") sessions.mark_continuable = MagicMock() sessions.unmark_continuable = MagicMock() + sessions.is_continuable = MagicMock(return_value=False) sessions.resumable_sid = MagicMock(return_value="sid-123") sessions.forget_conversation = MagicMock(return_value="sid-123") sessions.conversation_provider = MagicMock(return_value="acp") @@ -202,6 +204,87 @@ def test_gone_conversation_refused(self) -> None: assert info is not None and info.done assert info.error.startswith("conversation_gone") + def test_promotion_write_failure_is_retryable(self) -> None: + sessions = _mock_sessions() + manager = _manager(sessions) + with patch( + "kiro_crew.subagent_persistence.promote_retention", side_effect=OSError("disk busy") + ), patch.object(manager, "spawn") as spawn: + info = manager.continue_conversation("retryrun", "follow-up") + assert info.done + assert info.error.startswith("conversation_busy") + assert "subagent:retryrun" not in manager._conversations + sessions.unmark_continuable.assert_called_once_with("subagent:retryrun") + spawn.assert_not_called() + + def test_promotion_skipped_state_write_is_retryable(self) -> None: + sessions = _mock_sessions() + manager = _manager(sessions) + with patch("kiro_crew.subagent.update_state", return_value=False), patch.object( + manager, "spawn" + ) as spawn: + info = manager.continue_conversation("skiprun", "follow-up") + assert info.done + assert info.error.startswith("conversation_busy") + assert "subagent:skiprun" not in manager._conversations + spawn.assert_not_called() + + def test_retryable_promotion_preserves_existing_retention(self) -> None: + import kiro_crew.subagent_persistence as sp + + sessions = _mock_sessions() + sessions.is_continuable.return_value = True + manager = _manager(sessions) + manager._conversations["subagent:kept-run"] = 123.0 + with patch( + "kiro_crew.subagent_persistence.promote_retention", + return_value=sp.RetentionPromotionResult.RETRYABLE, + ), patch.object(manager, "spawn") as spawn: + info = manager.continue_conversation("kept-run", "follow-up") + assert info.done + assert info.error.startswith("conversation_busy") + assert manager._conversations["subagent:kept-run"] == 123.0 + sessions.unmark_continuable.assert_not_called() + spawn.assert_not_called() + + def test_concurrent_promotions_use_direct_return_values(self) -> None: + import kiro_crew.subagent_persistence as sp + + manager = _manager(_mock_sessions()) + for agent_id in ("retry-thread", "promoted-thread"): + sp.create_agent_folder(agent_id, task="t") + + results: dict[str, sp.RetentionPromotionResult] = {} + errors: list[BaseException] = [] + + def promote(agent_id: str) -> None: + try: + results[agent_id] = manager._promote_conversation( + agent_id, f"subagent:{agent_id}" + ) + except BaseException as exc: + errors.append(exc) + + def state_writer(agent_id: str, **_fields: object) -> bool: + return agent_id != "retry-thread" + + threads = [ + threading.Thread(target=promote, args=(agent_id,)) + for agent_id in ("retry-thread", "promoted-thread") + ] + with patch("kiro_crew.subagent.update_state", side_effect=state_writer): + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert all(not thread.is_alive() for thread in threads) + assert errors == [] + assert results == { + "retry-thread": sp.RetentionPromotionResult.RETRYABLE, + "promoted-thread": sp.RetentionPromotionResult.PROMOTED, + } + @pytest.mark.asyncio async def test_continue_seeds_from_state_json(self) -> None: """Retain-by-default: a run with no map entry seeds from state.json.""" @@ -212,15 +295,14 @@ async def test_continue_seeds_from_state_json(self) -> None: state = {"session_id": "sid-from-state", "provider": "acp", "cwd": "/tmp/x"} with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"), \ patch("kiro_crew.subagent.read_state", return_value=state), \ - patch("kiro_crew.subagent.update_state") as upd: + patch.object(manager, "_promote_conversation", return_value=object()) as promote: info = manager.continue_conversation("origrun2", "follow-up") assert info is not None and not info.error, info.error await manager._tasks[info.id] sessions.seed_conversation.assert_called_once_with( "subagent:origrun2", "sid-from-state", provider="acp", cwd="/tmp/x" ) - # Promotion: original run's state marked keep so the pruner retains it. - upd.assert_any_call("origrun2", keep=True) + promote.assert_called_once_with("origrun2", "subagent:origrun2") def test_continue_seed_with_missing_files_is_gone(self) -> None: """Seeded sid whose files are gone (map self-prunes) → conversation_gone.""" @@ -238,7 +320,9 @@ def test_continue_seed_with_missing_files_is_gone(self) -> None: async def test_continue_dispatches_new_run_on_same_key(self) -> None: sessions = _mock_sessions(resumed=True) manager = _manager(sessions) - with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"): + with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"), patch.object( + manager, "_promote_conversation", return_value=object() + ): info = manager.continue_conversation("origrun1", "follow-up work") assert info is not None and not info.error, info.error assert info.id != "origrun1" # new run id @@ -253,15 +337,27 @@ async def test_continuation_fails_closed_when_not_resumed(self) -> None: """session/load falling back to a fresh session must NOT execute the follow-up context-free — the run fails with a typed resume_failed.""" sessions = _mock_sessions(resumed=False) + provider = sessions.get_or_create.return_value[0] + provider.session_id = "sid-resume-fresh" manager = _manager(sessions) - with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"): + with patch("kiro_crew.subagent.Stats"), patch("kiro_crew.subagent.sel"), patch.object( + manager, "_promote_conversation", return_value=object() + ): info = manager.continue_conversation("origrun9", "follow-up work") assert info is not None and not info.error, info.error await manager._tasks[info.id] assert info.done assert "resume_failed" in info.error + # The fresh session must still be reclaimable even though execution + # fails before context construction or the state identity write. + import json + + import kiro_crew.subagent_persistence as sp + + ts = json.loads((sp._agent_dir(info.id) / "tombstone.json").read_text()) + assert ts["session_id"] == "sid-resume-fresh" + assert ts["provider"] == "acp" # The prompt must never have been sent on the fresh session. - provider = sessions.get_or_create.return_value[0] provider.stream.assert_not_called() @@ -401,6 +497,15 @@ def test_sweep_expires_only_idle_past_ttl(self) -> None: assert "subagent:old1" not in manager._conversations assert "subagent:new1" in manager._conversations + def test_sweep_drops_malformed_registry_key(self) -> None: + manager = _manager() + now = time.time() + manager._conversations["malformed"] = now - 7 * 3600 + with patch.object(manager, "release_conversation") as release: + manager._sweep_conversations(now) + assert "malformed" not in manager._conversations + release.assert_not_called() + def test_sweep_refreshes_busy_conversation(self) -> None: sessions = _mock_sessions() manager = _manager(sessions) @@ -547,6 +652,707 @@ async def test_cancelled_teardown_still_honours_keep_transcript( class TestPersistenceGuards: + def test_prune_lock_blocks_concurrent_false_to_true_promotion( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + owner_id = "promotion-race" + continuation_id = "promotion-race-child" + sp.create_agent_folder(owner_id, task="original") + sp.update_state(owner_id, session_id="sid-race", provider="acp", keep=False) + sp.create_agent_folder(continuation_id, task="continuation") + sp.update_state( + continuation_id, + session_id="sid-race", + provider="acp", + keep=True, + conversation_key=f"subagent:{owner_id}", + ) + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-race", + provider="acp", + keep=True, + conversation_key=f"subagent:{owner_id}", + ) + sp.write_tombstone( + continuation_id, cause="delivered", recovery_action="none" + ) + d = sp._agent_dir(continuation_id) + ts_path = d / "tombstone.json" + ts = json.loads(ts_path.read_text()) + ts["died"] = 1 + ts_path.write_text(json.dumps(ts)) + + observed_false = threading.Event() + allow_claim = threading.Event() + continuation_done = threading.Event() + original_decision = sp._should_defer_tombstone_cleanup + result: list[SubagentInfo] = [] + manager = _manager() + + def hold_after_false(**kwargs): # type: ignore[no-untyped-def] + observed_false.set() + assert allow_claim.wait(timeout=5) + return original_decision(**kwargs) + + def continue_run() -> None: + result.append(manager.continue_conversation(owner_id, "follow-up")) + continuation_done.set() + + with patch.object(sp, "_should_defer_tombstone_cleanup", hold_after_false), patch.object( + sp, "_cleanup_session_files_sync" + ): + prune_thread = threading.Thread( + target=sp.prune_stale_tombstones, + kwargs={"max_age_days": 0, "delivered_ttl_secs": 0}, + ) + continuation_thread: threading.Thread | None = None + try: + prune_thread.start() + assert observed_false.wait(timeout=5) + + continuation_thread = threading.Thread(target=continue_run) + continuation_thread.start() + assert continuation_done.wait(timeout=1) + assert result[0].error.startswith("conversation_busy") + finally: + allow_claim.set() + prune_thread.join(timeout=5) + if continuation_thread is not None: + continuation_thread.join(timeout=5) + + assert not prune_thread.is_alive() + assert continuation_thread is not None + assert not continuation_thread.is_alive() + assert continuation_done.is_set() + assert len(result) == 1 + assert result[0].done + manager._sessions.resumable_sid.return_value = None + retry = manager.continue_conversation(owner_id, "follow-up") + assert retry.done + assert retry.error.startswith("conversation_gone") + assert not d.exists() + assert not (sp.read_state(owner_id) or {}).get("keep") + + def test_unrelated_retention_lock_does_not_block_promotion(self, tmp_path) -> None: # type: ignore[no-untyped-def] + import kiro_crew.subagent_persistence as sp + + blocked_id = "retention-lock-blocked" + promoted_id = "retention-lock-independent" + for agent_id in (blocked_id, promoted_id): + sp.create_agent_folder(agent_id, task="original") + sp.update_state(agent_id, session_id=f"sid-{agent_id}", provider="acp", keep=False) + + holder = sp._retention_lock_for_agent(blocked_id) + holder.lock.acquire() + try: + result = asyncio.run(self._promote_on_loop(sp, promoted_id)) + finally: + holder.lock.release() + + assert result is sp.RetentionPromotionResult.PROMOTED + assert (sp.read_state(promoted_id) or {}).get("keep") is True + + def test_same_agent_retention_lock_is_retryable(self, tmp_path) -> None: # type: ignore[no-untyped-def] + import kiro_crew.subagent_persistence as sp + + agent_id = "promotion-retention-contention" + sp.create_agent_folder(agent_id, task="original") + sp.update_state(agent_id, session_id="sid-contention", provider="acp", keep=False) + + holder = sp._retention_lock_for_agent(agent_id) + holder.lock.acquire() + try: + result = asyncio.run(self._promote_on_loop(sp, agent_id)) + finally: + holder.lock.release() + + assert result is sp.RetentionPromotionResult.RETRYABLE + assert not (sp.read_state(agent_id) or {}).get("keep") + + def test_promotion_retries_while_off_loop_writer_lock_is_held( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import kiro_crew.subagent_persistence as sp + + agent_id = "promotion-writer-contention" + sp.create_agent_folder(agent_id, task="original") + sp.update_state(agent_id, session_id="sid-writer", provider="acp", keep=False) + holder = sp._lock_for_agent(agent_id) + holder.lock.acquire() + try: + started = time.monotonic() + result = asyncio.run(self._promote_on_loop(sp, agent_id)) + elapsed = time.monotonic() - started + finally: + holder.lock.release() + + assert result is sp.RetentionPromotionResult.RETRYABLE + assert elapsed < 0.5 + assert not (sp.read_state(agent_id) or {}).get("keep") + + result = asyncio.run(self._promote_on_loop(sp, agent_id)) + assert result is sp.RetentionPromotionResult.PROMOTED + assert (sp.read_state(agent_id) or {}).get("keep") is True + + def test_off_loop_promotion_uses_writer_lock_without_self_deadlock( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import kiro_crew.subagent_persistence as sp + + agent_id = "promotion-off-loop" + sp.create_agent_folder(agent_id, task="original") + sp.update_state(agent_id, session_id="sid-off-loop", provider="acp", keep=False) + results: list[sp.RetentionPromotionResult] = [] + + worker = threading.Thread( + target=lambda: results.append(sp.promote_retention(agent_id)) + ) + worker.start() + worker.join(timeout=2) + + assert not worker.is_alive(), "off-loop promotion self-deadlocked" + assert results == [sp.RetentionPromotionResult.PROMOTED] + assert (sp.read_state(agent_id) or {}).get("keep") is True + + @staticmethod + async def _promote_on_loop(sp, agent_id): # type: ignore[no-untyped-def] + return sp.promote_retention(agent_id) + + def test_invalid_owner_id_does_not_abort_prune_sweep( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + malformed_id = "a-malformed-owner" + valid_id = "z-valid-run" + sp.create_agent_folder(malformed_id, task="malformed") + sp.update_state( + malformed_id, + session_id="sid-malformed", + provider="acp", + keep=True, + conversation_key="subagent:../invalid", + ) + sp.create_agent_folder(valid_id, task="valid") + sp.update_state(valid_id, session_id="sid-valid", provider="acp", keep=False) + sp.remember_live_cleanup_identity( + malformed_id, + session_id="sid-malformed", + provider="acp", + keep=True, + conversation_key="subagent:../invalid", + ) + sp.remember_live_cleanup_identity( + valid_id, + session_id="sid-valid", + provider="acp", + keep=False, + ) + + for agent_id in (malformed_id, valid_id): + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + path = sp._agent_dir(agent_id) / "tombstone.json" + tombstone = json.loads(path.read_text()) + tombstone["died"] = 1 + path.write_text(json.dumps(tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 2 + + assert not sp._agent_dir(malformed_id).exists() + assert not sp._agent_dir(valid_id).exists() + assert cleanup.call_count == 2 + + def test_deep_tombstone_does_not_abort_prune_sweep( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + deep_id = "a-deep-tombstone" + valid_id = "z-valid-after-deep" + for agent_id in (deep_id, valid_id): + sp.create_agent_folder(agent_id, task=agent_id) + sp.update_state( + agent_id, + session_id=f"sid-{agent_id}", + provider="acp", + keep=False, + ) + sp.remember_live_cleanup_identity( + agent_id, + session_id=f"sid-{agent_id}", + provider="acp", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + + deep_path = sp._agent_dir(deep_id) / "tombstone.json" + deep_path.write_text("[" * 1100 + "0" + "]" * 1100, encoding="utf-8") + valid_path = sp._agent_dir(valid_id) / "tombstone.json" + valid_tombstone = json.loads(valid_path.read_text()) + valid_tombstone["died"] = 1 + valid_path.write_text(json.dumps(valid_tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + + assert sp._agent_dir(deep_id).exists() + assert not sp._agent_dir(valid_id).exists() + cleanup.assert_called_once_with(f"sid-{valid_id}", "acp", cwd="") + + def test_corrupt_protected_record_does_not_abort_prune_sweep( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + corrupt_id = "a-corrupt-protected" + valid_id = "z-valid-after-protected" + for agent_id in (corrupt_id, valid_id): + sp.create_agent_folder(agent_id, task=agent_id) + sp.update_state( + agent_id, + session_id=f"sid-{agent_id}", + provider="acp", + keep=False, + ) + sp.remember_live_cleanup_identity( + agent_id, + session_id=f"sid-{agent_id}", + provider="acp", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + tombstone_path = sp._agent_dir(agent_id) / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = 1 + tombstone_path.write_text(json.dumps(tombstone)) + + corrupt_record = sp._cleanup_identities_path(corrupt_id) + corrupt_record.write_text("{malformed") + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + + assert sp._agent_dir(corrupt_id).exists() + assert corrupt_record.read_text() == "{malformed" + assert not sp._agent_dir(valid_id).exists() + cleanup.assert_called_once_with(f"sid-{valid_id}", "acp", cwd="") + + def test_cancel_recovery_reclaims_every_session_generation( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + from unittest.mock import call + + import kiro_crew.subagent_persistence as sp + + agent_id = "cancel-recovery-generations" + sp.create_agent_folder(agent_id, task="t") + sp.update_state(agent_id, session_id="sid-state", provider="acp", keep=False) + real_atomic_write = sp._atomic_write + failed_once = False + + def flaky_sidecar_write(path, data): # type: ignore[no-untyped-def] + nonlocal failed_once + if path.name == sp._CLEANUP_IDENTITIES_FILE and not failed_once: + failed_once = True + raise OSError("transient sidecar failure") + return real_atomic_write(path, data) + + with patch.object(sp, "_atomic_write", side_effect=flaky_sidecar_write): + with pytest.raises(OSError, match="transient sidecar failure"): + sp.remember_live_cleanup_identity( + agent_id, session_id="sid-1", provider="acp", cwd="/first" + ) + sp.remember_live_cleanup_identity( + agent_id, session_id="sid-2", provider="acp", cwd="/second" + ) + sp.write_tombstone(agent_id, cause="cancelled", recovery_action="none") + + d = sp._agent_dir(agent_id) + ts_path = d / "tombstone.json" + tombstone = json.loads(ts_path.read_text()) + assert tombstone["session_id"] == "sid-2" + assert [item["session_id"] for item in tombstone["cleanup_identities"]] == [ + "sid-1", + "sid-2", + ] + + # Shutdown may clear an exclusion tombstone to re-admit orphan recovery. + # The protected generation record survives that and process-memory reset. + # Replacement tombstone creation stays memory-only; executor-owned prune + # merges protected generations without trusting the state-only SID. + sp.clear_tombstone(agent_id) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + assert not ts_path.exists() + assert sp._cleanup_identities_path(agent_id).exists() + sp.write_tombstone( + agent_id, cause="gateway_restart", recovery_action="notified" + ) + tombstone = json.loads(ts_path.read_text()) + assert "cleanup_identities" not in tombstone + assert tombstone["session_id"] == "sid-state" + tombstone["died"] = 1 + ts_path.write_text(json.dumps(tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + + assert cleanup.call_args_list == [ + call("sid-1", "acp", cwd="/first"), + call("sid-2", "acp", cwd="/second"), + ] + assert not d.exists() + assert not sp._cleanup_identities_path(agent_id).exists() + + def test_agent_writable_identity_files_cannot_delete_unrelated_session( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + agent_id = "untrusted-sidecar" + own_sid = "sid-owned-by-run" + victim_sid = "sid-owned-by-another-run" + sp.create_agent_folder(agent_id, task="t") + sp.update_state( + agent_id, + session_id=own_sid, + provider="acp", + keep=False, + ) + sp.remember_live_cleanup_identity( + agent_id, + session_id=own_sid, + provider="acp", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + agent_dir = sp._agent_dir(agent_id) + protected_path = sp._cleanup_identities_path(agent_id) + assert not protected_path.is_relative_to(agent_dir) + assert "trust" in protected_path.parts + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = 1 + tombstone["session_id"] = victim_sid + tombstone["cleanup_identities"] = [ + {"session_id": victim_sid, "provider": "acp"} + ] + tombstone_path.write_text(json.dumps(tombstone)) + + # These are the identity files a subagent can write. Durable cleanup + # authority lives under the protected trust root, so neither forged + # spelling may add the victim SID to the provider-deletion set. + (agent_dir / sp._CLEANUP_IDENTITIES_FILE).write_text( + json.dumps( + {"identities": [{"session_id": victim_sid, "provider": "acp"}]} + ) + ) + sessions_dir = tmp_path / "sessions" + sessions_dir.mkdir() + own_file = sessions_dir / f"{own_sid}.json" + victim_file = sessions_dir / f"{victim_sid}.json" + own_file.write_text("own") + victim_file.write_text("victim") + + with patch.object(sp, "kiro_sessions_dir", return_value=sessions_dir): + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + + assert not own_file.exists() + assert victim_file.read_text() == "victim" + assert not agent_dir.exists() + + def test_failed_provider_cleanup_preserves_folder_and_identity_for_retry( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + import time + + import kiro_crew.subagent_persistence as sp + + agent_id = "unsupported-cleanup-retry" + sp.create_agent_folder(agent_id, task="t") + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-claude-retry", + provider="claude_code", + cwd="/project", + keep=False, + ) + sp.update_state(agent_id, keep=False) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + agent_dir = sp._agent_dir(agent_id) + protected_path = sp._cleanup_identities_path(agent_id) + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = time.time() - 1 + tombstone_path.write_text(json.dumps(tombstone)) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + + # Claude Code has no cleanup route yet: keep both retry surfaces. + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + assert agent_dir.exists() + assert protected_path.exists() + + # Once a provider cleanup implementation succeeds, the same record is + # enough to complete prune and reap both surfaces. + with patch.object(sp, "_cleanup_session_files_sync", return_value=True): + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + assert not agent_dir.exists() + assert not protected_path.exists() + + def test_legacy_sid_without_trusted_generation_preserves_lookup_for_retry( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + import time + + import kiro_crew.subagent_persistence as sp + + agent_id = "legacy-untrusted-sid" + sp.create_agent_folder(agent_id, task="t") + sp.update_state( + agent_id, + session_id="sid-legacy", + provider="acp", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + agent_dir = sp._agent_dir(agent_id) + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = time.time() - 1 + tombstone_path.write_text(json.dumps(tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 0 + assert agent_dir.exists() + cleanup.assert_not_called() + + # A migration or later gateway-owned publication makes cleanup safe. + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-legacy", + provider="acp", + keep=False, + ) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + with patch.object(sp, "_cleanup_session_files_sync", return_value=True): + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + assert not agent_dir.exists() + + @pytest.mark.parametrize("trusted_generation", [False, True]) + def test_unreclaimable_lookup_has_ninety_day_hard_ceiling( + self, tmp_path, trusted_generation: bool + ) -> None: # type: ignore[no-untyped-def] + import json + import time + + import kiro_crew.subagent_persistence as sp + + agent_id = f"hard-ceiling-{trusted_generation}" + sp.create_agent_folder(agent_id, task="t") + sp.update_state( + agent_id, + session_id="sid-hard-ceiling", + provider="claude_code" if trusted_generation else "acp", + keep=False, + ) + if trusted_generation: + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-hard-ceiling", + provider="claude_code", + cwd="/project", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + agent_dir = sp._agent_dir(agent_id) + protected_path = sp._cleanup_identities_path(agent_id) + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = ( + time.time() - sp._UNRECLAIMABLE_LOOKUP_MAX_AGE_SECS - 1 + ) + tombstone_path.write_text(json.dumps(tombstone)) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + assert not agent_dir.exists() + assert not protected_path.exists() + + def test_unreadable_state_does_not_trust_stale_false_generation( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + import time + + import kiro_crew.subagent_persistence as sp + + agent_id = "stale-false-after-promotion" + sp.create_agent_folder(agent_id, task="t") + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-promoted", + provider="acp", + keep=False, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + # Promotion lands in state, but the acquisition-time generation remains + # false. Corrupt state must not turn that stale false into delete authority. + sp.update_state(agent_id, keep=True) + agent_dir = sp._agent_dir(agent_id) + (agent_dir / "state.json").write_text("{corrupt") + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = time.time() - 1 + tombstone.pop("cleanup_identities", None) + tombstone_path.write_text(json.dumps(tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 0 + assert agent_dir.exists() + cleanup.assert_not_called() + + # Unknown retention remains bounded rather than immortal. + tombstone["died"] = 1 + tombstone_path.write_text(json.dumps(tombstone)) + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + cleanup.assert_called_once_with("sid-promoted", "acp", cwd="") + + def test_cleanup_store_restriction_failure_aborts_before_access( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + agent_id = "cleanup-store-lockdown" + sp.create_agent_folder(agent_id, task="t") + protected_path = sp._cleanup_identities_path(agent_id) + protected_path.parent.mkdir(parents=True) + original = json.dumps({"identities": [{"session_id": "sid-original"}]}) + protected_path.write_text(original) + + with patch.object( + sp.platform_compat, "make_owner_only_dir" + ), patch.object( + sp.platform_compat, + "restrict_dir_to_owner", + side_effect=OSError("DACL refused"), + ): + with pytest.raises(OSError, match="DACL refused"): + sp._read_cleanup_identities_file(agent_id) + with pytest.raises(OSError, match="DACL refused"): + sp.remember_live_cleanup_identity( + agent_id, session_id="sid-forged", provider="acp" + ) + + assert protected_path.read_text() == original + + def test_cleanup_store_parse_failure_never_rewrites_history( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import kiro_crew.subagent_persistence as sp + + agent_id = "cleanup-store-corrupt" + sp.create_agent_folder(agent_id, task="t") + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-original", + provider="acp", + keep=False, + ) + protected_path = sp._cleanup_identities_path(agent_id) + malformed = "{malformed" + protected_path.write_text(malformed) + + with pytest.raises(ValueError): + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-new", + provider="acp", + keep=False, + ) + + assert protected_path.read_text() == malformed + + def test_tombstone_sid_cannot_override_latest_protected_retention( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + agent_id = "forged-tombstone-retention" + sp.create_agent_folder(agent_id, task="t") + sp.update_state(agent_id, provider="acp") + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-owned", + provider="acp", + keep=True, + ) + sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") + agent_dir = sp._agent_dir(agent_id) + tombstone_path = agent_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = 1 + tombstone["session_id"] = "sid-victim" + tombstone.pop("cleanup_identities", None) + tombstone_path.write_text(json.dumps(tombstone)) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 0 + assert agent_dir.exists() + cleanup.assert_not_called() + def test_tombstone_prune_keeps_files_for_keep_runs(self, tmp_path) -> None: # type: ignore[no-untyped-def] import json @@ -565,10 +1371,282 @@ def test_tombstone_prune_keeps_files_for_keep_runs(self, tmp_path) -> None: # t ts.pop("session_id", None) ts_path.write_text(json.dumps(ts)) with patch.object(sp, "_cleanup_session_files_sync") as cleanup: - pruned = sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) - assert pruned >= 1 + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + assert d.exists() cleanup.assert_not_called() + def test_continuation_prune_honors_original_release(self, tmp_path) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + original_id = "original-run" + continuation_id = "continuation-run" + sp.create_agent_folder(original_id, task="original") + sp.update_state(original_id, session_id="sid-c", provider="acp", keep=True) + sp.create_agent_folder(continuation_id, task="continuation") + sp.update_state( + continuation_id, + session_id="sid-c", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-c", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.update_state(original_id, keep=False) + sp.write_tombstone(continuation_id, cause="delivered", recovery_action="none") + d = sp._agent_dir(continuation_id) + ts_path = d / "tombstone.json" + ts = json.loads(ts_path.read_text()) + ts["died"] = 0 + ts_path.write_text(json.dumps(ts)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + pruned = sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) + assert pruned == 1 + assert not d.exists() + cleanup.assert_called_once_with("sid-c", "acp", cwd="") + + def test_continuation_partial_state_uses_sidecar_retention_fallback( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + original_id = "partial-state-owner" + continuation_id = "partial-state-continuation" + conversation_key = f"subagent:{original_id}" + sp.create_agent_folder(original_id, task="original") + sp.update_state(original_id, session_id="sid-partial", provider="acp", keep=True) + sp.create_agent_folder(continuation_id, task="continuation") + # Session acquisition publishes cleanup identity and retention together + # before the later best-effort combined state update. Simulate that update + # failing by leaving the otherwise readable state without either field. + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-partial", + provider="acp", + keep=True, + conversation_key=conversation_key, + ) + sp.update_state(continuation_id, provider="acp") + sp.write_tombstone(continuation_id, cause="delivered", recovery_action="none") + sp._LIVE_CLEANUP_IDENTITIES.clear() + d = sp._agent_dir(continuation_id) + ts_path = d / "tombstone.json" + ts = json.loads(ts_path.read_text()) + ts["died"] = 0 + # Restart can rewrite a cleared tombstone before state identity lands; + # durable generation metadata must still carry retention and owner. + ts.pop("session_id", None) + ts.pop("cleanup_identities", None) + ts_path.write_text(json.dumps(ts)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + assert d.exists() + cleanup.assert_not_called() + + # Current readable owner state remains authoritative over stale + # sidecar keep=True once release records keep=False. + sp.update_state(original_id, keep=False) + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + assert not d.exists() + cleanup.assert_called_once_with("sid-partial", "acp", cwd="") + + @pytest.mark.parametrize( + ("writable_key", "writable_sid"), + [ + ("", "sid-continuation"), + ("subagent:forged-retention-owner", "sid-continuation"), + ("", "sid-forged"), + ], + ) + def test_protected_continuation_owner_outranks_writable_state( + self, tmp_path, writable_key: str, writable_sid: str + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + owner_id = "protected-retention-owner" + continuation_id = f"protected-owner-{bool(writable_key)}" + conversation_key = f"subagent:{owner_id}" + sp.create_agent_folder(owner_id, task="original") + sp.update_state(owner_id, session_id="sid-owner", provider="acp", keep=True) + sp.create_agent_folder(continuation_id, task="continuation") + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-continuation", + provider="acp", + keep=True, + conversation_key=conversation_key, + ) + # Agent-writable state must not erase or redirect the protected owner. + sp.update_state( + continuation_id, + session_id=writable_sid, + provider="acp", + keep=False, + conversation_key=writable_key, + ) + sp.write_tombstone(continuation_id, cause="delivered", recovery_action="none") + continuation_dir = sp._agent_dir(continuation_id) + tombstone_path = continuation_dir / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = 0 + tombstone_path.write_text(json.dumps(tombstone)) + with sp._CLEANUP_IDENTITY_LOCK: + sp._LIVE_CLEANUP_IDENTITIES.clear() + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 0 + assert continuation_dir.exists() + cleanup.assert_not_called() + + sp.update_state(owner_id, keep=False) + assert sp.prune_stale_tombstones( + max_age_days=0, delivered_ttl_secs=0 + ) == 1 + + assert not continuation_dir.exists() + cleanup.assert_called_once_with("sid-continuation", "acp", cwd="") + + def test_unreadable_state_and_empty_tombstone_use_sidecar_retention( + self, tmp_path + ) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + owner_id = "unreadable-sidecar-owner" + child_id = "unreadable-sidecar-child" + conversation_key = f"subagent:{owner_id}" + sp.create_agent_folder(owner_id, task="original") + sp.update_state(owner_id, session_id="sid-owner", provider="acp", keep=True) + sp.create_agent_folder(child_id, task="continuation") + sp.remember_live_cleanup_identity( + child_id, + session_id="sid-sidecar-only", + provider="acp", + keep=True, + conversation_key=conversation_key, + ) + sp.write_tombstone(child_id, cause="delivered", recovery_action="none") + (sp._agent_dir(child_id) / "state.json").write_text("{corrupt") + sp._LIVE_CLEANUP_IDENTITIES.clear() + tombstone_path = sp._agent_dir(child_id) / "tombstone.json" + tombstone = json.loads(tombstone_path.read_text()) + tombstone["died"] = 0 + tombstone.pop("session_id", None) + tombstone.pop("cleanup_identities", None) + tombstone_path.write_text(json.dumps(tombstone)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + cleanup.assert_not_called() + + sp.update_state(owner_id, keep=False) + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + + cleanup.assert_called_once_with("sid-sidecar-only", "acp", cwd="") + assert not sp._agent_dir(child_id).exists() + + def test_continuation_prune_owner_missing_keep_is_nonretained(self, tmp_path) -> None: # type: ignore[no-untyped-def] + import json + + import kiro_crew.subagent_persistence as sp + + original_id = "owner-missing-keep" + continuation_id = "continuation-missing-keep" + sp.create_agent_folder(original_id, task="original") + sp.update_state(original_id, session_id="sid-m", provider="acp") + sp.create_agent_folder(continuation_id, task="continuation") + sp.update_state( + continuation_id, + session_id="sid-m", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-m", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.update_state(original_id, provider="acp") + sp.write_tombstone(continuation_id, cause="delivered", recovery_action="none") + d = sp._agent_dir(continuation_id) + ts_path = d / "tombstone.json" + ts = json.loads(ts_path.read_text()) + ts["died"] = 0 + ts_path.write_text(json.dumps(ts)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + pruned = sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) + assert pruned == 1 + assert not d.exists() + cleanup.assert_called_once_with("sid-m", "acp", cwd="") + + @pytest.mark.parametrize("tombstone_has_sid", [True, False]) + def test_continuation_prune_unreadable_owner_gets_bounded_grace( + self, tmp_path, tombstone_has_sid: bool + ) -> None: # type: ignore[no-untyped-def] + import json + import time + + import kiro_crew.subagent_persistence as sp + + original_id = "owner-corrupt" + continuation_id = "continuation-corrupt-owner" + sp.create_agent_folder(original_id, task="original") + sp.update_state(original_id, session_id="sid-u", provider="acp", keep=True) + (sp._agent_dir(original_id) / "state.json").write_text("{corrupt") + sp.create_agent_folder(continuation_id, task="continuation") + sp.update_state( + continuation_id, + session_id="sid-u", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.remember_live_cleanup_identity( + continuation_id, + session_id="sid-u", + provider="acp", + keep=True, + conversation_key=f"subagent:{original_id}", + ) + sp.write_tombstone(continuation_id, cause="delivered", recovery_action="none") + d = sp._agent_dir(continuation_id) + ts_path = d / "tombstone.json" + ts = json.loads(ts_path.read_text()) + ts["died"] = time.time() - (12 * 3600) + if not tombstone_has_sid: + ts.pop("session_id", None) + ts_path.write_text(json.dumps(ts)) + + with patch.object(sp, "_cleanup_session_files_sync") as cleanup: + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + assert d.exists() + cleanup.assert_not_called() + ts["died"] = time.time() - (2 * 86400) + ts_path.write_text(json.dumps(ts)) + assert sp.prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + assert not d.exists() + cleanup.assert_called_once_with("sid-u", "acp", cwd="") + def test_tombstone_prune_cleans_files_for_plain_runs(self, tmp_path) -> None: # type: ignore[no-untyped-def] import json @@ -576,7 +1654,13 @@ def test_tombstone_prune_cleans_files_for_plain_runs(self, tmp_path) -> None: # agent_id = "plainrun" sp.create_agent_folder(agent_id, task="t") - sp.update_state(agent_id, session_id="sid-p", provider="acp") + sp.update_state(agent_id, session_id="sid-p", provider="acp", keep=False) + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-p", + provider="acp", + keep=False, + ) sp.write_tombstone(agent_id, cause="delivered", recovery_action="none") d = sp._agent_dir(agent_id) ts_path = d / "tombstone.json" diff --git a/test/test_subagent_persistence.py b/test/test_subagent_persistence.py index 19b36fd2903..e33d9bfe51f 100644 --- a/test/test_subagent_persistence.py +++ b/test/test_subagent_persistence.py @@ -4,18 +4,23 @@ import json import os +import threading import time import pytest from kiro_crew.subagent_persistence import ( + _CLEANUP_IDENTITY_LOCK, + _LIVE_CLEANUP_IDENTITIES, create_agent_folder, delete_agent_folder, list_orphans, mark_delivered, prune_stale_tombstones, read_state, + read_tombstone, record_slow_command, + remember_live_cleanup_identity, update_state, write_result_chunk, write_tombstone, @@ -28,9 +33,11 @@ @pytest.fixture() def agent_root(tmp_path, monkeypatch): - """Point persistence at a temp directory.""" - monkeypatch.setattr("kiro_crew.subagent_persistence._SUBAGENTS_DIR", tmp_path) - return tmp_path + """Point persistence at a registry below this test's temp directory.""" + root = tmp_path / "subagents" + root.mkdir() + monkeypatch.setattr("kiro_crew.subagent_persistence._SUBAGENTS_DIR", root) + return root # ── create_agent_folder ────────────────────────────────────────────── @@ -73,6 +80,21 @@ def test_preserves_existing_fields(self, agent_root): assert state["task"] == "original" assert state["pid"] == 99 + def test_corrupt_or_nonobject_state_skips_update(self, agent_root, monkeypatch): + import kiro_crew.subagent_persistence as sp + + create_agent_folder("u-corrupt", task="t") + path = agent_root / "u-corrupt" / "state.json" + path.write_bytes(b"\xff") + assert update_state("u-corrupt", pid=1) is False + + path.write_text("[]", encoding="utf-8") + assert update_state("u-corrupt", pid=1) is False + + path.write_text("{}", encoding="utf-8") + monkeypatch.setattr(sp.json, "loads", lambda _text: (_ for _ in ()).throw(RecursionError())) + assert update_state("u-corrupt", pid=1) is False + def test_missing_folder_logs_no_crash(self, agent_root): # Should not raise update_state("nonexistent", pid=1) @@ -97,6 +119,50 @@ def test_corrupt_json_returns_none(self, agent_root): (folder / "state.json").write_text("{corrupt") assert read_state("bad") is None + def test_invalid_encoding_depth_and_nonobject_return_none( + self, agent_root, monkeypatch + ): + import kiro_crew.subagent_persistence as sp + + folder = agent_root / "bad-shapes" + folder.mkdir() + path = folder / "state.json" + path.write_bytes(b"\xff") + assert read_state("bad-shapes") is None + + path.write_text("[]", encoding="utf-8") + assert read_state("bad-shapes") is None + + path.write_text("{}", encoding="utf-8") + monkeypatch.setattr( + sp.json, + "loads", + lambda _text: (_ for _ in ()).throw(RecursionError()), + ) + assert read_state("bad-shapes") is None + + +class TestReadTombstone: + def test_reads_only_valid_tombstone_objects(self, agent_root): + create_agent_folder("rt1", task="t") + assert read_tombstone("rt1") is None + + write_tombstone( + "rt1", + cause="timeout", + recovery_action="notified", + session_id="sid-read", + ) + assert (read_tombstone("rt1") or {})["session_id"] == "sid-read" + + path = agent_root / "rt1" / "tombstone.json" + path.write_text("{corrupt", encoding="utf-8") + assert read_tombstone("rt1") is None + path.write_text("[]", encoding="utf-8") + assert read_tombstone("rt1") is None + path.write_text("[" * 1100 + "0" + "]" * 1100, encoding="utf-8") + assert read_tombstone("rt1") is None + # ── write_result_chunk ─────────────────────────────────────────────── @@ -129,6 +195,172 @@ def test_extra_fields_included(self, agent_root): assert ts["pid"] == 999 assert ts["turns"] == 12 + def test_live_identity_snapshot_never_reads_sidecar(self, agent_root): + from unittest.mock import patch + + import kiro_crew.subagent_persistence as sp + + agent_id = "t-memory-only-snapshot" + create_agent_folder(agent_id, task="t") + with _CLEANUP_IDENTITY_LOCK: + _LIVE_CLEANUP_IDENTITIES[agent_id] = [ + {"session_id": "sid-memory", "provider": "acp"} + ] + with patch.object( + sp, + "_read_cleanup_identities_file", + side_effect=AssertionError("event-loop snapshot read sidecar"), + ): + write_tombstone(agent_id, cause="error", recovery_action="none") + assert (read_tombstone(agent_id) or {})["session_id"] == "sid-memory" + + def test_live_identity_snapshot_does_not_wait_and_delete_clears_fallback( + self, agent_root + ): + agent_id = "t-identity-lock" + create_agent_folder(agent_id, task="t") + with _CLEANUP_IDENTITY_LOCK: + _LIVE_CLEANUP_IDENTITIES[agent_id] = [ + {"session_id": "sid-lock", "provider": "acp"} + ] + write_tombstone(agent_id, cause="error", recovery_action="none") + assert (read_tombstone(agent_id) or {})["session_id"] == "sid-lock" + + delete_id = "t-delete-fallback" + create_agent_folder(delete_id, task="t") + with _CLEANUP_IDENTITY_LOCK: + _LIVE_CLEANUP_IDENTITIES[delete_id] = [{"session_id": "sid-delete"}] + delete_agent_folder(delete_id) + assert delete_id not in _LIVE_CLEANUP_IDENTITIES + + def test_live_identity_published_before_sidecar_read(self, agent_root): + from unittest.mock import patch + + import kiro_crew.subagent_persistence as sp + + agent_id = "t-prepublish" + create_agent_folder(agent_id, task="t") + entered_read = threading.Event() + allow_read = threading.Event() + errors: list[BaseException] = [] + original_read = sp._read_cleanup_identities_file + + def blocked_read(run_id: str): + entered_read.set() + assert allow_read.wait(timeout=5) + return original_read(run_id) + + def remember() -> None: + try: + sp.remember_live_cleanup_identity( + agent_id, session_id="sid-prepublish", provider="acp" + ) + except BaseException as exc: # surfaced after unconditional join + errors.append(exc) + + worker = threading.Thread(target=remember) + with patch.object(sp, "_read_cleanup_identities_file", blocked_read): + worker.start() + assert entered_read.wait(timeout=5) + try: + write_tombstone(agent_id, cause="error", recovery_action="none") + assert (read_tombstone(agent_id) or {})["session_id"] == "sid-prepublish" + sp.publish_live_cleanup_identity( + agent_id, + session_id="sid-concurrent", + provider="acp", + ) + finally: + allow_read.set() + worker.join(timeout=5) + assert not worker.is_alive() + assert errors == [] + durable = json.loads(sp._cleanup_identities_path(agent_id).read_text()) + assert [item["session_id"] for item in durable["identities"]] == [ + "sid-prepublish", + "sid-concurrent", + ] + + def test_protected_record_stays_inside_test_temp_root(self, agent_root): + import kiro_crew.subagent_persistence as sp + + agent_id = "protected-test-root" + create_agent_folder(agent_id, task="t") + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-contained", + provider="acp", + keep=False, + ) + protected_path = sp._cleanup_identities_path(agent_id) + assert protected_path.is_relative_to(agent_root.parent) + assert protected_path.exists() + + def test_concurrent_publish_during_sidecar_write_is_not_overwritten( + self, agent_root + ): + from unittest.mock import patch + + import kiro_crew.subagent_persistence as sp + + agent_id = "t-append-only-publish" + create_agent_folder(agent_id, task="t") + entered_write = threading.Event() + allow_write = threading.Event() + errors: list[BaseException] = [] + original_write = sp._atomic_write + + def blocked_write(path, data): # type: ignore[no-untyped-def] + if path.name == sp._CLEANUP_IDENTITIES_FILE: + entered_write.set() + assert allow_write.wait(timeout=5) + return original_write(path, data) + + def remember_first() -> None: + try: + sp.remember_live_cleanup_identity( + agent_id, + session_id="sid-first", + provider="acp", + ) + except BaseException as exc: + errors.append(exc) + + worker = threading.Thread(target=remember_first) + with patch.object(sp, "_atomic_write", side_effect=blocked_write): + worker.start() + assert entered_write.wait(timeout=5) + try: + sp.publish_live_cleanup_identity( + agent_id, + session_id="sid-second", + provider="acp", + ) + write_tombstone(agent_id, cause="error", recovery_action="none") + finally: + allow_write.set() + worker.join(timeout=5) + + assert not worker.is_alive() + assert errors == [] + tombstone = read_tombstone(agent_id) or {} + assert [item["session_id"] for item in tombstone["cleanup_identities"]] == [ + "sid-first", + "sid-second", + ] + assert [ + item["session_id"] for item in sp._live_cleanup_identities(agent_id) + ] == ["sid-first", "sid-second"] + + def test_snapshots_cleanup_identity_from_readable_state(self, agent_root): + create_agent_folder("t3", task="t") + update_state("t3", session_id="sid-state", provider="claude_code", cwd="/project") + write_tombstone("t3", cause="gateway_restart", recovery_action="delivered") + ts = json.loads((agent_root / "t3" / "tombstone.json").read_text(encoding="utf-8")) + assert ts["session_id"] == "sid-state" + assert ts["provider"] == "claude_code" + assert ts["cwd"] == "/project" + # ── delete_agent_folder ────────────────────────────────────────────── @@ -189,6 +421,243 @@ def test_keeps_recent_tombstones(self, agent_root): prune_stale_tombstones(max_age_days=7) assert (agent_root / "new1").exists() + @pytest.mark.parametrize( + "died_case", + ["string", "nan", "infinity", "future", "oversized"], + ) + def test_invalid_tombstone_died_uses_mtime_fallback( + self, agent_root, monkeypatch, died_case + ): + agent_id = f"invalid-died-{died_case}" + create_agent_folder(agent_id, task="t") + update_state(agent_id, session_id="sid-died", provider="acp", keep=False) + remember_live_cleanup_identity( + agent_id, + session_id="sid-died", + provider="acp", + keep=False, + ) + write_tombstone( + agent_id, + cause="delivered", + recovery_action="notified", + session_id="sid-died", + ) + ts_path = agent_root / agent_id / "tombstone.json" + ts = json.loads(ts_path.read_text(encoding="utf-8")) + ts["died"] = { + "string": "invalid", + "nan": float("nan"), + "infinity": float("inf"), + "future": time.time() + 86400, + "oversized": 10**400, + }[died_case] + ts_path.write_text(json.dumps(ts), encoding="utf-8") + fallback = time.time() - (2 * 86400) + os.utime(ts_path, (fallback, fallback)) + + cleaned: list[str] = [] + monkeypatch.setattr( + "kiro_crew.subagent_persistence._cleanup_session_files_sync", + lambda sid, provider, *, cwd="": cleaned.append(sid), + ) + assert prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + assert cleaned == ["sid-died"] + assert not (agent_root / agent_id).exists() + + def test_mtime_fallback_at_cutoff_is_eligible(self, agent_root, monkeypatch): + import kiro_crew.subagent_persistence as sp + + agent_id = "mtime-at-cutoff" + create_agent_folder(agent_id, task="t") + update_state(agent_id, session_id="sid-cutoff", provider="acp", keep=False) + remember_live_cleanup_identity( + agent_id, + session_id="sid-cutoff", + provider="acp", + keep=False, + ) + write_tombstone( + agent_id, + cause="delivered", + recovery_action="notified", + session_id="sid-cutoff", + ) + ts_path = agent_root / agent_id / "tombstone.json" + ts = json.loads(ts_path.read_text(encoding="utf-8")) + ts["died"] = 0 + ts_path.write_text(json.dumps(ts), encoding="utf-8") + os.utime(ts_path, (100.0, 100.0)) + monkeypatch.setattr(sp.time, "time", lambda: 100.0) + + cleaned: list[str] = [] + monkeypatch.setattr( + sp, + "_cleanup_session_files_sync", + lambda sid, provider, *, cwd="": cleaned.append(sid), + ) + assert prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 1 + assert cleaned == ["sid-cutoff"] + assert not (agent_root / agent_id).exists() + + def test_future_died_and_mtime_preserve_unreadable_state_grace( + self, agent_root, monkeypatch + ): + """Clock rollback must not turn unknown retention into immediate cleanup.""" + import kiro_crew.subagent_persistence as sp + + agent_id = "future-died-and-mtime" + create_agent_folder(agent_id, task="t") + update_state(agent_id, session_id="sid-future", provider="acp", keep=False) + write_tombstone( + agent_id, + cause="delivered", + recovery_action="notified", + session_id="sid-future", + ) + state_path = agent_root / agent_id / "state.json" + state_path.write_text("{corrupt", encoding="utf-8") + ts_path = agent_root / agent_id / "tombstone.json" + ts = json.loads(ts_path.read_text(encoding="utf-8")) + ts["died"] = 200.0 + ts_path.write_text(json.dumps(ts), encoding="utf-8") + os.utime(ts_path, (200.0, 200.0)) + monkeypatch.setattr(sp.time, "time", lambda: 100.0) + + cleaned: list[str] = [] + monkeypatch.setattr( + sp, + "_cleanup_session_files_sync", + lambda sid, provider, *, cwd="": cleaned.append(sid), + ) + assert prune_stale_tombstones(max_age_days=0, delivered_ttl_secs=0) == 0 + assert cleaned == [] + assert (agent_root / agent_id).exists() + + @pytest.mark.parametrize( + ("state_case", "expect_cleanup"), + [ + pytest.param("partial-nonkeep", True, id="partial-state-cleans"), + pytest.param( + "initial-keep-write-missed", True, id="missing-keep-in-readable-state-cleans" + ), + pytest.param("corrupt", False, id="corrupt-state-skips-cleanup"), + pytest.param("non-object", False, id="non-object-state-skips-cleanup"), + pytest.param("promoted-keep", False, id="promoted-state-retains"), + pytest.param("string-false", True, id="string-false-cleans"), + ], + ) + def test_prune_uses_trusted_identity_with_safe_live_state( + self, agent_root, monkeypatch, state_case, expect_cleanup + ): + """Gateway publication supplies identity; live state owns retention intent.""" + agent_id = "retention-case" + session_id = "session-retention" + create_agent_folder(agent_id, task="t") + if state_case == "initial-keep-write-missed": + update_state(agent_id, provider="acp") + else: + update_state(agent_id, provider="acp", keep=False) + remember_live_cleanup_identity( + agent_id, + session_id=session_id, + provider="acp", + keep=False, + ) + write_tombstone( + agent_id, + cause="timeout", + recovery_action="notified", + session_id=session_id, + ) + + state_path = agent_root / agent_id / "state.json" + if state_case == "corrupt": + state_path.write_text("{corrupt", encoding="utf-8") + elif state_case == "non-object": + state_path.write_text("[]", encoding="utf-8") + elif state_case == "promoted-keep": + update_state(agent_id, keep=True) + elif state_case == "string-false": + update_state(agent_id, keep="false") + + ts_path = agent_root / agent_id / "tombstone.json" + ts = json.loads(ts_path.read_text(encoding="utf-8")) + ts["died"] = time.time() - ((7 * 86400) + (12 * 3600)) + ts_path.write_text(json.dumps(ts), encoding="utf-8") + cleaned: list[str] = [] + monkeypatch.setattr( + "kiro_crew.subagent_persistence._cleanup_session_files_sync", + lambda sid, provider, *, cwd="": cleaned.append(sid), + ) + + invalid_state = state_case in {"corrupt", "non-object"} + bounded_grace = invalid_state + if bounded_grace: + assert prune_stale_tombstones(max_age_days=7) == 0 + assert cleaned == [] + assert (agent_root / agent_id).exists() + quarantined = json.loads(ts_path.read_text(encoding="utf-8")) + assert "state_unreadable_at" not in quarantined + quarantined["died"] = time.time() - (9 * 86400) + ts_path.write_text(json.dumps(quarantined), encoding="utf-8") + + assert prune_stale_tombstones(max_age_days=7) == 1 + assert cleaned == [session_id] + assert not (agent_root / agent_id).exists() + elif state_case == "promoted-keep": + assert prune_stale_tombstones(max_age_days=7) == 0 + assert cleaned == [] + assert (agent_root / agent_id).exists() + update_state(agent_id, keep=False) + assert prune_stale_tombstones(max_age_days=7) == 1 + assert cleaned == [session_id] + assert not (agent_root / agent_id).exists() + else: + assert prune_stale_tombstones(max_age_days=7) == 1 + assert cleaned == ([session_id] if expect_cleanup else []) + assert not (agent_root / agent_id).exists() + + @pytest.mark.parametrize( + ("state_case", "expected_cleanup"), + [ + pytest.param("corrupt", [], id="unreadable-without-tombstone-id"), + pytest.param("readable", ["sid-state-only"], id="state-only-id-no-grace"), + ], + ) + def test_extra_grace_requires_trusted_cleanup_identity( + self, agent_root, monkeypatch, state_case, expected_cleanup + ): + agent_id = f"no-tombstone-id-{state_case}" + create_agent_folder(agent_id, task="t") + write_tombstone(agent_id, cause="timeout", recovery_action="notified") + update_state(agent_id, session_id="sid-state-only", provider="acp") + if state_case == "readable": + remember_live_cleanup_identity( + agent_id, + session_id="sid-state-only", + provider="acp", + keep=False, + ) + state_path = agent_root / agent_id / "state.json" + if state_case == "corrupt": + state_path.write_text("{corrupt", encoding="utf-8") + + ts_path = agent_root / agent_id / "tombstone.json" + ts = json.loads(ts_path.read_text(encoding="utf-8")) + assert "session_id" not in ts + ts["died"] = time.time() - (8 * 86400) + ts_path.write_text(json.dumps(ts), encoding="utf-8") + cleaned: list[str] = [] + monkeypatch.setattr( + "kiro_crew.subagent_persistence._cleanup_session_files_sync", + lambda sid, provider, *, cwd="": cleaned.append(sid), + ) + + assert prune_stale_tombstones(max_age_days=7) == 1 + assert cleaned == expected_cleanup + assert not (agent_root / agent_id).exists() + def test_keeps_non_tombstoned_folders(self, agent_root): create_agent_folder("running1", task="t") prune_stale_tombstones(max_age_days=7) @@ -441,6 +910,7 @@ async def test_pid_recorded_after_session_create(self, agent_root): provider.start = AsyncMock() provider.shutdown = AsyncMock() provider.context_usage_pct = lambda: 0.0 + provider.session_id = "session-live" async def _stream(*_a, **_kw): yield LLMEvent(kind=EVENT_COMPLETE) @@ -466,6 +936,9 @@ async def _stream(*_a, **_kw): state = json.loads((agent_root / info.id / "state.json").read_text(encoding="utf-8")) assert state["pid"] == 42 + assert state["session_id"] == "session-live" + assert info._session_id == "session-live" + assert info._session_provider == state["provider"] assert "pid_recorded_at" in state assert isinstance(state["pid_recorded_at"], float) @@ -539,7 +1012,19 @@ async def test_tombstone_on_timeout(self, agent_root): manager = SubagentManager(sessions=sessions, ctx_builder=MagicMock()) info = SubagentInfo(id="timeout1", task="t", parent_session_key="dashboard:default") + info._session_id = "session-live" + info._session_provider = "claude_code" + info._session_cwd = "/project" + info.keep = True create_agent_folder("timeout1", task="t") + from kiro_crew.subagent_persistence import remember_live_cleanup_identity + + remember_live_cleanup_identity( + "timeout1", + session_id="session-live", + provider="claude_code", + cwd="/project", + ) manager._agents["timeout1"] = info manager._running_count = 1 @@ -549,6 +1034,10 @@ async def test_tombstone_on_timeout(self, agent_root): ts = json.loads((agent_root / "timeout1" / "tombstone.json").read_text(encoding="utf-8")) assert ts["cause"] == "timeout" + assert ts["session_id"] == "session-live" + assert ts["provider"] == "claude_code" + assert ts["cwd"] == "/project" + assert "keep" not in ts @pytest.mark.asyncio async def test_timeout_skipped_when_already_reaped(self, agent_root):