fix(ai): discover persistence from native capability captures

linktools-ai/docs/capability-persistence.md

@@ -163,3 +163,29 @@ values; do not rely on an embedded object ID to retain an external asset.
 Only supported transaction-store changes receive that adapter's publication and
 cache behavior. This protocol does not publish arbitrary application hot state.
 One transaction may still execute several SQL statements and network round trips.
+
+### Native containers and callback owners
+
+A custom native container must expose its children through the SDK `apply` and
+`visit_and_replace` contracts and preserve replacements when forwarding hooks.
+Persistence discovery uses the leaves visited during declaration capture; a
+custom container does not need a separate partition registration or a built-in
+container type. Each nested persistent leaf needs its own explicit native `id`.
+The outer registration supplies the revision, not a shared ID for every child.
+A directly registered persistent leaf without an ID may use its registration ID.
+Duplicate participant IDs and duplicate built-in persistence domains remain
+conflicts, including inside custom containers and wrappers.
+
+| Stage | Callback owner | Retained boundary |
+| --- | --- | --- |
+| Declaration capture | Reusable capability declaration | Discovers participant IDs/revisions and partitions; does not attach a run writer to the original declaration |
+| `bind_persistence(writer)` | Independent run-local capability | The writer belongs to this run and partition; binding preserves the commit policy and borrowed buffer |
+| Native `for_agent` / `for_run` replacement | Effective native capability instance | Future submissions use the replacement's callback; a write accepted during `for_run` before it returns retains the earlier owner |
+| Deferred flush | Callback captured by each accepted submission | Later native replacements do not retarget pending writes to a different instance |
+| Execution sealing / Runtime close | Existing execution and writer lifecycle | Sealing rejects later writes/bindings; close drains or reports unresolved work without taking ownership of borrowed resources |
+
+The container paths are exercised by `test_custom_persistence_container.py`.
+`test_capability_persistence_lifecycle.py` separately covers native replacements
+and two concurrent runs: their writers remain independent, and early/late writes
+reach the callbacks that accepted them. These are owner-local lifecycle checks,
+not evidence of cross-process callback replay.

linktools-ai/docs/domain-composition.md

@@ -89,6 +89,48 @@ forwarding facade just to reduce the number of `service` references.
   `connection=...` must instead be closed by that caller after Runtime closes.
   Native snapshots are not copied into a second Runtime checkpoint archive.
 
+## Resource ownership is selected per resource
+
+`Runtime.attach` borrows Persistence, not every object in `features`. Resource
+declarations still determine who initializes and closes each dependency:
+
+| Resource supplied to Runtime | Initialization and close owner | Existing public call shape |
+| --- | --- | --- |
+| Persistence opened by `Runtime.open` | Runtime | One Runtime context |
+| Persistence supplied to `Runtime.attach` | Host; it must already be open and writable | Outer Persistence context, inner Runtime context |
+| `AssetCapability(store)` in `RuntimeFeatures` | Runtime, even with `attach` or an already initialized store | Pass the declaration directly for ownership transfer |
+| Asset source in `features.borrow()` | Host; initialize before borrowers and close after all borrowers finish | Pass the live borrowed source, not an inspection resolution |
+| Default `StepPersistenceCapability()` adapter | Runtime creates/closes the adapter over Execution storage; its external SQL engine remains host-owned | No extra store context; requires an Execution SQL route |
+| Explicit StepStore in `StepPersistenceCapability(store=...)` or native `StepPersistence(store=...)` | Caller; Runtime does not close it | Keep the store usable through producer cleanup |
+
+Runtime releases owned dependencies only after producer cleanup succeeds. A
+producer close failure leaves those dependencies available for retrying
+`runtime.close()`; closing borrowed Persistence first would remove that recovery
+path. Partial initialization also uses the same cleanup ownership rather than
+transferring it to the failing capability. The existing
+`test_runtime_resource_ownership.py` checks these combinations, explicit stores
+inside wrappers, partial initialization and close retry; this table does not add
+a new lifecycle or automatic recovery mechanism.
+
+## Recovery uses the existing owner
+
+These cases have different identities and must not be combined into a generic
+"restore Memory" path:
+
+| Operation | State and resource owner | Existing contract coverage |
+| --- | --- | --- |
+| Retry an ordinary execution after Runtime reopens | Execution retains its admitted Memory scope; current capability declarations determine available tools. A removed Memory declaration stays absent, and a same-named business tool does not become a Memory tool. | `test_memory_capability_scope.py::test_memory_tools_retry_resolves_current_tools_but_retains_admitted_scope` |
+| Retry an imported execution | The new execution gets an isolated scope seeded with the captured content and versions, not later edits in the original execution's scope. | `test_imported_context_retry.py::test_retry_imported_memory_uses_fresh_isolated_scope` |
+| Continue a deferred imported execution | The same execution keeps its scope and intervening mutations; continuation does not reapply the seed over those mutations. | `test_memory_capture_replay.py::test_runtime_imported_memory_resumes_deferred_without_rollback` |
+| Reopen and retry a Memory seed | `PersistentMemoryStore` uses the original seed receipts; Persistence owns their atomic commit with the seeded paths. | `test_memory_capture_replay.py::test_memory_capture_receipt_replays_after_sqlite_reopen` and `test_memory_capture_retry_resolves_committed_partial_seed` |
+| Replace a custom native persistent capability during a run | The SDK selects the effective run instance. New submissions use that instance; already accepted writes retain their original callback owner until commit. | `test_capability_persistence_lifecycle.py::test_native_persistence_replacements_keep_current_and_accepted_write_owners` |
+
+The deferred-continuation case covers one Runtime lifetime. The seed-receipt
+case separately covers SQLite reopen; together they do not prove every
+cold-process continuation scenario. Accepted owner-local Python callbacks are
+not a cross-worker durable callback queue. Extending recovery requires a real
+caller contract, not another registry or a second copy of Memory state.
+
 ## Failure, cancellation, and recovery boundaries
 
 A wait timeout or a cancelled observer is not proof that the producer stopped.

linktools-ai/src/linktools/ai/runtime/_feature_policy.py

@@ -114,8 +114,8 @@ def consume_runtime_declaration(capability: AbstractCapability[AgentContext[obje
             return rebound
         return capability
 
-    captured_persistence_ids: set[str] = set()
     capability_revisions: dict[str, int] = {}
+    captured_capabilities: list[AbstractCapability[AgentContext[object]]] = []
     runtime_capability_values = []
     for candidate in candidates:
         if candidate.scope != "runtime":
@@ -130,16 +130,18 @@ def capture_persistence(capability: AbstractCapability[AgentContext[object]]) ->
                 rebound.wrapped = wrapped
                 return rebound
             selected = consume_runtime_declaration(capability)
+            if selected is not None:
+                captured_capabilities.append(selected)
             if selected is None or not isinstance(selected, CapabilityPersistence):
                 return selected
             identity = selected.id
             if identity is None:
                 if capability is not candidate.value:
                     raise AIError(ErrorCode.CAPABILITY_RESOLUTION_INVALID, "Persistent capability leaves require a registered identity")
                 identity = candidate.id
-            if identity in captured_persistence_ids:
+            identity, = capability_persistence_ids(((identity, selected),))
+            if identity in capability_revisions:
                 raise AIError(ErrorCode.CAPABILITY_CONFLICT)
-            captured_persistence_ids.add(identity)
             capability_revisions[identity] = candidate.revision
             bound = copy(selected)
             bound.id = identity
@@ -150,6 +152,7 @@ def capture_persistence(capability: AbstractCapability[AgentContext[object]]) ->
             runtime_capability_values.append(selected)
     if not execution_declared:
         runtime_capability_values.append(execution_capability)
+        captured_capabilities.append(execution_capability)
     runtime_capabilities = tuple(runtime_capability_values)
     runtime_ids: set[str] = set()
 
@@ -161,11 +164,8 @@ def register_runtime_id(capability: AbstractCapability[AgentContext[object]]) ->
 
     for capability in runtime_capabilities:
         capability.apply(register_runtime_id)
-    validate_persistence_capabilities(runtime_capabilities)
-    capability_ids = capability_persistence_ids(
-        tuple((capability.id or "", capability) for capability in runtime_capabilities),
-    )
+    validate_persistence_capabilities(captured_capabilities)
     return _RuntimeFeaturePolicy(
         runtime_capabilities, execution_capability, session_capability,
-        memory_scope, tool_responses, capability_revisions, tuple(capability_ids),
+        memory_scope, tool_responses, capability_revisions, tuple(capability_revisions),
     )

tests/ai/test_custom_persistence_container.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Custom native capability trees retain their declared persistence participants."""
+
+from collections.abc import Callable, Mapping
+from copy import copy
+
+import pytest
+from pydantic_ai.capabilities import AbstractCapability, CombinedCapability, WrapperCapability
+from pydantic_ai.messages import ModelResponse, TextPart
+from pydantic_ai.tools import RunContext
+
+from linktools.ai.capability import AgentContext, TraceCapability
+from linktools.ai.core import CommitPolicy, ExecutionStatus, JsonValue
+from linktools.ai.errors import AIError, ErrorCode
+from linktools.ai.persistence import Persistence, PersistenceStorage
+from linktools.ai.persistence.capabilities import CapabilityWriter
+from linktools.ai.persistence.state import StateTransaction, StoredRecord
+from linktools.ai.runtime import Runtime, RuntimeFeatures
+
+from ._runtime_test_helpers import _Models, _UsageFunctionModel, model_features
+
+
+class _RunNote(AbstractCapability[AgentContext[None]]):
+    id = "custom.note"
+    commit = CommitPolicy.sync()
+    writer: CapabilityWriter | None = None
+
+    def bind_persistence(self, writer: CapabilityWriter) -> "_RunNote":
+        bound = copy(self)
+        bound.writer = writer
+        return bound
+
+    async def before_run(self, ctx: RunContext[AgentContext[None]]) -> None:
+        assert self.writer is not None
+        await self.writer.submit(key="started", payload={"execution_id": ctx.deps.execution_id})
+
+    async def persist_in_transaction(self, transaction: StateTransaction, payload: Mapping[str, JsonValue]) -> None:
+        await transaction.insert_record(StoredRecord(
+            key=b"started", scope_key=None, parent_key=None, kind="custom-note",
+            sort_key=b"started", state=None, storage_revision=0,
+            lease_owner=None, lease_fence=0, lease_expires_at=None, data=payload,
+        ))
+
+
+class _NoteGroup(AbstractCapability[AgentContext[None]]):
+    def __init__(self, note: AbstractCapability[AgentContext[None]]) -> None:
+        self.note = note
+
+    def apply(self, visitor: Callable[[AbstractCapability[AgentContext[None]]], None]) -> None:
+        self.note.apply(visitor)
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractCapability[AgentContext[None]]], AbstractCapability[AgentContext[None]] | None],
+    ) -> "_NoteGroup | None":
+        note = self.note.visit_and_replace(visitor)
+        if note is None:
+            return None
+        if note is self.note:
+            return self
+        bound = copy(self)
+        bound.note = note
+        return bound
+
+    async def before_run(self, ctx: RunContext[AgentContext[None]]) -> None:
+        await self.note.before_run(ctx)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("container", ("custom", "wrapped_custom", "wrapped_combined"))
+async def test_native_container_persists_its_registered_leaf(container: str) -> None:
+    async def model(messages, info):
+        return ModelResponse(parts=[TextPart("done")])
+
+    note = _RunNote()
+    if container == "wrapped_combined":
+        group = WrapperCapability(CombinedCapability([note]))
+    elif container == "wrapped_custom":
+        group = WrapperCapability(_NoteGroup(note))
+    else:
+        group = _NoteGroup(note)
+    features = RuntimeFeatures("custom-container")
+    features.agent("default", model="default", allow_tools=())
+    features.capability(group, id="custom.notes", revision=3, scope="runtime")
+    async with Persistence.open("custom-container", storage=PersistenceStorage.in_memory()) as persistence:
+        async with Runtime.attach(
+            persistence=persistence,
+            features=model_features(_Models(_UsageFunctionModel(model)), features),
+        ) as runtime:
+            execution = await runtime.agents.get().start("question")
+            result = (await execution.wait(timeout_seconds=5)).result
+            assert result.status is ExecutionStatus.SUCCEEDED
+            store = await persistence.capability_store(note.id)
+            rows = await store.read(lambda transaction: transaction.scan_records())
+            assert len(rows) == 1
+            assert rows[0].data == {"execution_id": execution.execution_id}
+            assert note.writer is None
+
+
+@pytest.mark.asyncio
+async def test_custom_native_container_rejects_duplicate_persistent_ids() -> None:
+    features = RuntimeFeatures("duplicate-container")
+    features.capability(
+        _NoteGroup(CombinedCapability([_RunNote(), _RunNote()])),
+        id="custom.notes", revision=3, scope="runtime",
+    )
+    async with Persistence.open("duplicate-container", storage=PersistenceStorage.in_memory()) as persistence:
+        with pytest.raises(AIError) as conflict:
+            async with Runtime.attach(persistence=persistence, features=features):
+                pytest.fail("duplicate persistent capability IDs were accepted")
+        assert conflict.value.code is ErrorCode.CAPABILITY_CONFLICT
+
+
+@pytest.mark.asyncio
+async def test_custom_native_container_rejects_duplicate_persistence_domains() -> None:
+    first, second = TraceCapability(), TraceCapability()
+    first.id, second.id = "trace.first", "trace.second"
+    features = RuntimeFeatures("duplicate-domains")
+    features.capability(_NoteGroup(CombinedCapability([first, second])), id="custom.traces", scope="runtime")
+    async with Persistence.open("duplicate-domains", storage=PersistenceStorage.in_memory()) as persistence:
+        with pytest.raises(AIError) as conflict:
+            async with Runtime.attach(persistence=persistence, features=features):
+                pytest.fail("duplicate persistent domain declarations were accepted")
+        assert conflict.value.code is ErrorCode.CAPABILITY_CONFLICT