diff --git a/linktools-ai/docs/capability-persistence.md b/linktools-ai/docs/capability-persistence.md index 82b0d9ff3b..ae0819efb4 100644 --- a/linktools-ai/docs/capability-persistence.md +++ b/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. diff --git a/linktools-ai/docs/domain-composition.md b/linktools-ai/docs/domain-composition.md index 8d4ee81ba7..a863b00278 100644 --- a/linktools-ai/docs/domain-composition.md +++ b/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. diff --git a/linktools-ai/src/linktools/ai/runtime/_feature_policy.py b/linktools-ai/src/linktools/ai/runtime/_feature_policy.py index 84625ebc1e..9d38058059 100644 --- a/linktools-ai/src/linktools/ai/runtime/_feature_policy.py +++ b/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,6 +130,8 @@ 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 @@ -137,9 +139,9 @@ def capture_persistence(capability: AbstractCapability[AgentContext[object]]) -> 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), ) diff --git a/tests/ai/test_custom_persistence_container.py b/tests/ai/test_custom_persistence_container.py new file mode 100644 index 0000000000..6ba46b1116 --- /dev/null +++ b/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