diff --git a/linktools-ai/docs/measurements/history-window-read-io.md b/linktools-ai/docs/measurements/history-window-read-io.md new file mode 100644 index 0000000000..56fbb720bf --- /dev/null +++ b/linktools-ai/docs/measurements/history-window-read-io.md @@ -0,0 +1,90 @@ +# Exact History window reads + +## Scope + +Baseline: `52a9edb6ebee3d197edc10e7d00baf45cb19fa43`. +The isolated synthetic fixture performs two deterministic model requests and +twelve effect-free tool deliveries. It then requests the final tool's exact +History projection with `agent_run_seq=1`, `model_request_seq=1`, and `limit=2`. +Both content modes must equal the corresponding items in the complete History. + +**The measurements below are for one exact History query, not for executing the +whole fixture.** They are SQLite SQL attempts and transaction events, not disk +flushes, network bytes, or MySQL latency. Initialization and the full History +reference query are outside the measurement window. A task-local context marker +includes database work owned by the measured query and excludes earlier worker +cleanup. The optional listener stores counts only, without SQL parameters. + +## Minimal change + +`TranscriptRepository._read_range_start` previously read a transcript head and +its seek hint separately when starting a nonempty window after message zero. +It now reads those two records together in the same existing transaction. +The head remains authoritative; the seek record remains only a hint. Missing or +malformed hints retain their existing fallback to validated transcript facts. +No public interface, durable format, cache policy, lock, or transaction boundary +changes. Empty windows and already-captured heads keep their existing paths. + +| Read configuration | SELECT before → after | BEGIN before → after | Commit before → after | +| --- | --- | --- | --- | +| No ordinary cache, live warm | 12 → 11 | 7 → 7 | 7 → 7 | +| Ordinary cache, live warm | 10 → 9 | 5 → 5 | 5 → 5 | +| No ordinary cache, reopened cold | 14 → 13 | 9 → 9 | 8 → 8 | +| Ordinary cache, reopened cold | 14 → 13 | 9 → 9 | 9 → 9 | +| No ordinary cache, reopened warm | 12 → 11 | 7 → 7 | 7 → 7 | +| Ordinary cache, reopened warm | 10 → 9 | 5 → 5 | 5 → 5 | + +Both `include_content=False` and `True` give the same counts. The uncached cold +reader has one rolled-back content-discovery read before opening its archive; +that existing retry is unchanged. No measured read performs a write. + +A separate counterfactual grouped post-authorization execution root and child +metadata. It reduced one uncached read transaction but forced a source root +read on the ordinary-cache path, cancelling the SELECT saving there. That +counterfactual was removed rather than changing the existing cache policy. + +## Verification and reproduction + +The public regression covers exact item equality, order-preserving one-item +pagination, unknown tool selection, missing execution, tenant denial, both +content modes, and close/reopen with fresh caches. A positive-offset transcript +test covers present, missing, and malformed hints while requiring one bounded +head/seek metadata query. Existing range tests cover append cutoffs, sparse +windows, malformed content, paused object I/O, cancellation, and concurrent +collection. The sparse-window test no longer freezes which private helper is +called; its output, boundary, fallback and separate SQL skipped-page coverage +remain. + +The final focused set passed **155 tests** through `manage.py` (40.06 seconds), +including the seven new cases. The repository's configured Ruff selections and +`git diff --check` passed. Independent static review found no must-fix in this +slice. The same new positive-offset budget check fails on the unchanged baseline +with two metadata queries instead of one; its returned messages still match. +These are upstream synthetic checks, not application or real-MySQL acceptance. + +From the repository root, with the normal development environment installed: + +```sh +HISTORY_IO_OUTPUT=/tmp/history-read-counts.jsonl HISTORY_IO_LABEL=current python - <<'PY' +import manage +manage._run_pytest('linktools-ai', { + 'paths': ['tests/ai/test_history_exact_read_io.py'], +}, manage._check_environment(), 'all') +PY +``` + +For a matched baseline, copy only the same new test file into a clean worktree +at the baseline SHA and run its public-runtime test selector. The new +positive-offset I/O-budget test is expected to fail on the baseline because +that version issues separate head and hint queries. Use the same Python and +dependency environment for both worktrees. + +## Remaining boundary + +This is one fewer SELECT for each affected transcript window, not a reduction +of the fixed two-turn workload's 38 write transactions. It does not claim that +exact tool selection avoids all unrelated message materialization: +`include_content=False` still removes content after History projection. Broader +query selection, whole-workload transaction costs, long-session retention, +and designated application acceptance remain separate work. Do not extrapolate +these counts to a different pin, tool policy, graph, or History query shape. diff --git a/linktools-ai/docs/measurements/transaction-boundaries.md b/linktools-ai/docs/measurements/transaction-boundaries.md new file mode 100644 index 0000000000..ea1146e2cb --- /dev/null +++ b/linktools-ai/docs/measurements/transaction-boundaries.md @@ -0,0 +1,333 @@ +# Short transaction boundaries + +## Status and objective + +The write-transaction assessment is a measured baseline and correctness inventory; +its goal is to identify justified consolidation, not require a lower count. +The implemented narrow slice below only removes redundant reads when a local +terminal publication is already confirmed and retained. The source baseline is +`776caf18b1c11d1e2f2122e79748521863c0c217` (tree +`5ec1b6f1178b401edf7fbca15b05d3eeeeb33873`). + +The goal is to eliminate repeated short transactions for the same business +operation, not to reach an arbitrary transaction count. Each candidate must say +who needs the data, when it must be durable, which changes must be atomic, and +whether it moves reads, adds writes, or extends lock/transaction lifetime. + +## Reproduction and measured baseline + +Use the repository's installed check environment and `manage.py` test helper: + +```bash +TURN_IO_OUTPUT=/tmp/transaction-baseline.jsonl TURN_IO_LABEL=776caf18-baseline \ +PYTHONPATH="$PWD/linktools-ai/src:$PWD/linktools/src:$PWD" python - <<'PY' +import manage +manage._run_pytest('linktools-ai', + {'paths': ['tests/ai/test_upstream_runtime_measurement.py']}, + manage._check_environment(), 'all') +PY +``` + +The offline fixture runs two session turns, each with three tool calls and a final +text answer: eight model requests, six tool effects and sixteen history messages. +It uses Native SQL checkpoints/frontiers, synchronous Execution persistence, +retained conversation history, Trace off, and no turn coordination. Schema setup +and session creation are outside the measurement. The window includes worker +completion and retention release, not just the first publicly available result. +Replay must cause no additional model or tool effects; the second turn must see +the first turn's tool history. + +A fresh SQLite run on 2026-11-14 used Python 3.12.14, SQLite 3.53.1, +SQLAlchemy 2.0.51, aiosqlite 0.22.1, pydantic-ai-slim 2.68.0 and +pydantic-ai-harness 0.68.0: + +| Operation | Turn 1 | Turn 2 | Total | +| --- | ---: | ---: | ---: | +| SELECT | 14 | 17 | 31 | +| INSERT | 65 | 64 | 129 | +| UPDATE | 22 | 22 | 44 | +| Commit | 23 | 23 | 46 | +| Commit containing writes | 19 | 19 | 38 | +| Bound text/blob bytes | 98,276 | 101,907 | 200,183 | + +Bound bytes are SQL argument traffic, not unique stored bytes or wire protocol +size. This is a small upstream synthetic SQLite case, not a server latency +benchmark or an application acceptance run. Do not compare a future candidate +with a different workload, dependency environment, Trace policy or source pin. + +## Business ownership of the write boundaries + +The async-owner labels in the probe account for all 38 committed write transactions: + +| Business boundary | Two-turn count | Why the boundary currently exists | +| --- | ---: | --- | +| Execution admission | 2 | Durable idempotency and session admission before returning the handle | +| Attempt activation | 2 | Freeze accepted input context and claim producer authority before dispatch | +| Native model request/response | 16 | Recovery frontier before provider work and checkpoint before dependent work | +| Native run completion | 2 | SDK run outcome; not yet Runtime output/cleanup success | +| Tool admission/settlement | 12 | Before/after external effect; settlement already includes the Native event | +| Execution/session terminal publication | 2 | Result, execution status and conversation publication commit together | +| Retention closure | 2 | Only after owned work drains; current dependency holds still need checking | + +The eight other commits are read-only captures/queries: per turn, session-start +admission hints, committed conversation context, the first authorized wait read, +and terminal preparation each use one transaction. All 46 transactions have +async-owner labels; the JSON still reports unattributed counters for future paths. There were no open +transactions at either measured window end. Counts are observations, not long-lived +API assertions. The fixture now records each transaction and statement's async +owner path, SQL verb/table, bound bytes and commit/rollback outcome. It emits no +parameter values or private request data. SQLAlchemy commit callbacks describe +observed commit attempts; this successful fixture is not an unknown-commit proof. + +### Measured operation attribution + +| Boundary label | Commits | SELECT | INSERT | UPDATE | +| --- | ---: | ---: | ---: | ---: | +| `execution.read_session_start` | 2 | 2 | 0 | 0 | +| `execution.reserve_start_checkpoint` | 2 | 6 | 11 | 2 | +| `conversation.load_session_context` | 2 | 2 | 0 | 0 | +| `execution.load_result_authorized` | 2 | 2 | 0 | 0 | +| `execution.commit_prepared_agent_attempt_checkpoint` | 2 | 5 | 0 | 2 | +| `native.commit[run_started,model_request_started;snapshot=complete]` | 2 | 0 | 10 | 2 | +| `native.commit[model_request_completed;snapshot=interrupted]` | 6 | 0 | 24 | 6 | +| `execution.commit_tool_admission_boundary` | 6 | 0 | 18 | 6 | +| `execution.commit_tool_terminal_boundary` | 6 | 0 | 12 | 6 | +| `native.commit[model_request_started;snapshot=complete]` | 6 | 0 | 24 | 6 | +| `native.commit[model_request_completed;snapshot=complete]` | 2 | 0 | 8 | 2 | +| `native.commit[run_completed;snapshot=complete]` | 2 | 0 | 4 | 2 | +| `execution.commit_terminal_checkpoint` | 4 | 11 | 18 | 8 | +| `retention.release_execution_handoff` | 2 | 3 | 0 | 2 | + +The terminal label includes both its preparation read and its write transaction. +Native labels distinguish complete/interrupted snapshot state, not semantic +equivalence of their complete contents. The full JSON has per-statement table names. + +### Existing effect-free tool policy (not an optimization comparison) + +The default fixture remains `effect_policy="replay_safe"`. Its twelve tool +transactions are not a charge imposed on every read-only tool. A diagnostic run +of the same source with `TURN_IO_TOOL_POLICY=none` exercises the existing +explicit effect-free contract: 34 total commits / 26 write commits, SELECT 31, +INSERT 99, UPDATE 32, and 151,505 bound text/blob bytes. The eight-model/six-tool/ +sixteen-history and replay assertions still pass. This changes the policy, not +the implementation, and must not be reported as a 38-to-26 optimization. + +Built-in `list_files`, `read_file`, `search_files`, `read_memory` and +`search_memory` already declare `effect_policy="none"`. New effect-free calls +skip ToolOperation admission/settlement; an already recorded call still follows +its original replay decision. MCP read-only hints do not grant this contract. +Effectful settlement already atomically writes the business result and Native +tool event. A failed event publication can require a separate outcome-preserving +transaction after rollback; that recovery path must not be deleted as a duplicate. + +## Confirmed terminal wait read reuse + +A wait that starts after this process has confirmed a complete terminal publication +can use the same retained publication helper already used after wait/poll wakeups. +Previously the first read always reloaded the source result, even while that exact +publication was retained. The change reuses the existing helper at first entry; +it adds no cache, lifetime extension, write batching or new durable owner. + +The helper still applies the current authorization policy. The wait still waits +for an owned worker to settle, checks worker failure and uses the original result +consumption/handoff path. No receipt means the original source read, including +cold/remote and unconfirmed/unknown cases. External result bodies still use their +original readers. The consumer scope is entered before choosing a receipt, so +local lifetime cleanup cannot overtake consumption. + +A focused SQLite probe applies the same completed-worker, explicit dependency-hold +fixture to `fc92b8d6` and the candidate. It uses the built-in tenant policy and an +inline output. The hold keeps the real publication available; it is test setup, +not a new production hold. Counts cover `wait()`, not setup or hold acquisition: + +| Case | Before SELECT / read commits | After SELECT / read commits | Outcome | +| --- | ---: | ---: | --- | +| Confirmed retained publication, success | 3 / 2 | 0 / 0 | Same output | +| No retained publication, success | 3 / 2 | 3 / 2 | Original source fallback | +| Retained publication, authorization revoked | 1 / 1 | 0 / 0 | Denied | +| Retained publication, worker failure | 1 / 1 | 0 / 0 | Original failure | +| No publication, revoked authorization or worker failure | 1 / 1 | 1 / 1 | Original rejection | + +These are read-only transactions; no write commit is removed. The original fixed +8-model/6-tool/16-history probe remains **46 total / 38 write commits**, SELECT 31, +INSERT 129, UPDATE 44 and 200,183 bound bytes. This slice does not solve the write +transaction target. + +One availability change is intentional: an inline result with a confirmed retained +publication can now be consumed when the redundant source-result query is +unavailable. Authorization and worker-failure checks still apply. Without that +receipt the same source failure still propagates. This is not permission to +invent a result after an unknown commit or to suppress necessary body reads. + +Reproduce the retained-wait probe using the normal repository environment: + +```bash +RECEIPT_WAIT_IO_OUTPUT=/tmp/receipt-wait.jsonl \ +PYTHONPATH="$PWD/linktools-ai/src:$PWD/linktools/src:$PWD" python - <<'PYTEST' +import manage +manage._run_pytest('linktools-ai', {'paths': [ + 'tests/ai/test_confirmed_turn_receipts_independent.py' + '::test_wait_started_after_publication_preserves_read_guards', +]}, manage._check_environment(), 'all') +PYTEST +``` + +The opt-in JSON logs counters only, not parameters. Focused verification covers +receipt presence/absence, current authorization rejection, worker failure, source +unavailability, real dependency holds and the existing retention/context contracts. +The current slice passed 65 receipt/retention/exact-context tests and the fixed +probe; the final receipt test with opt-in measurement passed all 15 cases. + +## Candidate boundaries and proof obligations + +### Duplicate owner rows in terminal content reads + +A matched probe at `082b09527f417fcea3c7adc9f6684e768592f626` +identified one narrower duplication inside the existing terminal write query: +the same Execution record was selected as both a content owner and an additional +record. The query now omits the overlapping additional role only for the same +store instance, key and non-payload route. The fresh owner read, referenced +content, current-transaction comparison and physical CAS remain unchanged. +The original large-additional-batch fallback is retained before deduplication, +including its separate typed-payload routing. + +In the fixed two-turn fixture, returned/decoded Execution owner rows change +**4 → 2**. Their compact UTF-8 JSON representations change **13,766 → 6,883 +bytes**. These are logical row JSON bytes, not disk or network measurements. +SELECT 31, INSERT 129, UPDATE 44 and **46 total / 38 write commits** are unchanged. +No immutable body read or durability boundary was removed. + +The phase audit also rules out two tempting but incorrect shortcuts here: +conversation preparation already batches transcript metadata and its descriptor; +the subsequent write reads require their own current snapshot. Additional +`get_records` calls after terminal prefetch were transaction-cache hits, not +extra SQL. The locally held Execution projection had physical revision 2 while +the current source was revision 14, so it could not replace the fresh source. + +Four new regressions cover large owner rows, missing keys, different partitions +sharing a key, later owner batches, extras-only reads and the mixed typed-payload +fallback threshold. The unfixed query fails the two overlap cases. The candidate +passes 49 focused content-source, group/CAS, shared-terminal and measurement +cases, preserving the original model/tool/history/replay assertions. This is a +returned-row reduction only; the write-transaction objective remains open. + +### Admission and activation + +These are currently separated by history/metadata capture, preload/input-context +preparation and worker scheduling. There is no permission to await external input +or invoke user code inside a write transaction. A possible pure-input fast path +would have to prepare outside the transaction and atomically validate the captured +session cursor/revision with idempotency, admission, input context and authority. +It must preserve: + +- cancellation or forced session close after `start()` but before first dispatch; +- accepted metadata stability while later turns see updated metadata; +- a second Runtime claiming an admitted execution before a delayed worker, with + one provider invocation and one attempt sequence; +- same-version restore, CAS rejection and unknown-commit reconciliation. + +`test_initial_attempt_capture.py` covers these distinctions. No such fast path +has been implemented by this inventory. Claiming ACTIVE during `start()` closes +the existing delayed-worker takeover window. Conversely, keeping ADMITTED until +the first model checkpoint would move the authority claim past model materialization, +model context entry and capability preparation. The current contract does not +establish that all of these steps are pure or safe to run twice. Moving pure input +preparation alone does not eliminate either durable authority boundary. + +### Native run completion and Execution terminal publication + +`AtomicStepPersistence.after_run` has an equality/step check for snapshot reuse, +but this measured fixture still records `run_completed` with a complete snapshot +(two INSERTs and one UPDATE per turn). The response and terminal snapshot +difference was checked with a temporary diagnostic: messages and encoded bytes +are identical, while step moves from 3 to 4. Suppressing that descriptor in a +counterfactual reduces one INSERT and about 440 bound bytes per turn, but reduces +no transaction. Message bodies are already shared. Step participates in snapshot +integrity and public events, so this is not a proposed production change. Runtime subsequently reads the run +and checkpoint, checks unresolved tools, validates output and exits model and +workspace contexts. Those steps may fail or be cancelled. Native completion is +therefore not interchangeable with Execution success. + +A staged built-in-only terminal plan would need an explicit, exact publication +contract for both successful and failed Runtime outcomes. It cannot report a +staged record as durable, hide it from a supported observer, or guess that output +validation/cleanup succeeds. Unknown commit, deferred-tool and error paths must +retain their own correct outcomes. No staged plan is implemented here. The existing Native commit contract confirms +snapshot state and then emits `SnapshotSavedEvent`; replacing that commit with +staging would announce uncommitted data as durable. Shared lifecycle registration +only prepares event identity/clock, and its eligibility check does not exclude +`on_event` consumers. The generic capability terminal buffer verifies its own +operation receipt, not Native `StepWriteReceipt`, and supplies no pending Native +read view. A new staged view/confirmation/exit protocol is therefore not included +as a nominally small transaction reduction. + +### Terminal preparation read + +The successful Native-reference fixture spends one read-only transaction and one +SELECT per turn preparing conversation publication, then opens the terminal write +transaction. The observed head determines chunk positions (first turn `[0, 8]`, +second turn `[8, 8]`), target counts, quality and replay validation. Preparing those +values from assumed absence would be incorrect. Moving all preparation inside the +write transaction would move work and transaction lifetime rather than prove a +net improvement. Separating pure body preparation from final metadata resolution +would affect the transcript, checkpoint and terminal owners plus UNKNOWN handling; +no such refactor is included merely to remove these two read-only commits. + +### Adjacent boundaries: consolidation assessment + +The assessment also checks adjacent writes, rather than assuming separate owners +always need separate transactions: + +| Adjacent boundaries | Assessment under the current contract | +| --- | --- | +| Attempt activation → first model request | Activation claims producer authority before model/context/capability preparation. Delaying the claim permits duplicate preparation; moving the request earlier changes the accepted dispatch frontier. No transparent consolidation established. | +| Model response → first tool admission | `AtomicStepPersistence.after_node_run` commits the response snapshot and then emits `SnapshotSavedEvent` before later tool dispatch. Deferral would either announce an uncommitted checkpoint or change the observable callback/cancellation boundary. | +| Tool settlement → next model request | The known tool result must survive cancellation or producer loss even when no next request occurs. Settlement already includes its Native event. The next dispatch separately prepares its request and verifies producer authority; it cannot be predeclared solely because a tool completed. | +| Final model response → Native completion | The response checkpoint is confirmed before the later `after_run` terminal event/snapshot. An apparently identical message body does not establish identical snapshot step or outcome, and the confirmation callback is observable. | +| Native completion → Execution/session terminal | Runtime output validation and context cleanup still intervene and can fail; Native completion is not sufficient evidence of Runtime success. | +| Execution/session terminal → retention closure | Result visibility does not prove producer drain or absence of late dependent holds. Closing retention at publication would change the lifetime guarantee. | + +These are existing protocol obligations, not a claim that a deliberate future +protocol redesign is impossible. A redesign must specify its new observation, +cancellation and recovery semantics rather than present their removal as batching. + +### Current-master compatibility recheck + +On 2026-11-15 the fixed two-turn probe was repeated with Python 3.11.18 against a +local, uncommitted integration of PR head +`11e78c31be33124f5fc69cebf6b419d9b7c68a50` and master +`18165a4a64ff42942248990093c464baddcff415`. It passed: 31 SELECT, +129 INSERT, 44 UPDATE, 46 commits, including 38 write commits. All transaction +attribution counters were zero for unattributed/open transactions, and the +8 model requests, 6 tool effects, 16 history messages and no-duplicate-effects +replay assertions passed. This is a compatibility recheck of the stated fixture, +not an exact-head CI result or a database-server benchmark. + +### Necessary boundaries and non-goals + +- Do not span a model request, tool execution, user hook or external I/O with a + database transaction. +- Do not remove a model-response checkpoint before dependent tools or validation. +- Do not merge tool admission with settlement across the external effect. +- Do not remove the producer/CAS fence because a guard UPDATE leaves the logical + execution record unchanged. Native publication's guard already runs in the + same group transaction as its event/snapshot writes; it is not an extra + standalone transaction. Replacing it with a prior read loses serialization + against concurrent cancellation or producer replacement. +- Do not close retention before late dependency holds and producer drain settle. +- Keep `auto_recover`, split-store behavior, authorization, cancellation and + observation visibility unchanged. Useful small query projections remain + allowed; this work does not require a new model-fact storage owner. +- Keep broader IO and designated application validation separate from this + upstream probe. Do not claim they pass because this fixture passes. + +A candidate is deliverable only after matched measurement, focused correctness +checks, independent review and natural CI for its exact commit. A lower count +without those semantics is not an optimization. + +This inventory does not prove 38 is a universal lower bound. It finds no further +justified write-transaction consolidation within the examined unchanged contracts. +An unchanged count is a valid assessment result, not an unfinished numerical +target. Broader workload attribution and designated server/application validation +remain separate acceptance items; this conclusion neither completes nor removes +them. Any new workload or proposed protocol change requires its own evidence. diff --git a/linktools-ai/src/linktools/ai/execution/_service.py b/linktools-ai/src/linktools/ai/execution/_service.py index be899a7f6e..8166c3189d 100644 --- a/linktools-ai/src/linktools/ai/execution/_service.py +++ b/linktools-ai/src/linktools/ai/execution/_service.py @@ -260,7 +260,9 @@ async def release_dependency_hold( key = tenant_id, execution_id state, owner = await self._handoff.release_hold(key, hold_id) if owner and state is not None: - await self._run_handoff_cleanup(execution_id, tenant_id, state) + # The durable hold still prevents cleanup. Relinquish the local + # claim before I/O so failure cannot strand later consumers. + await self._handoff.finish_release(key, state, succeeded=False) current = await self._repositories.executions.release_dependency_hold( execution_id, tenant_id=tenant_id, @@ -902,6 +904,11 @@ async def resolve_task_effect( raise AIError(ErrorCode.STORAGE_CONFLICT) return project_execution_view(current) if current.status is ExecutionStatus.FAILED: + if resolution.kind == "failed": + if (current.error_code != resolution.error_code.value + or dict(current.safe_error_details) != dict(resolution.safe_details)): + raise AIError(ErrorCode.STORAGE_CONFLICT) + return project_execution_view(current) if ( resolution.kind != "applied" or applied_digest is None @@ -919,6 +926,17 @@ async def resolve_task_effect( if resolution.kind == "unknown": return project_execution_view(current) + if resolution.kind == "failed": + await self.fail_task( + execution_id, + principal=principal, + error=AIError(resolution.error_code, retryable=False, + safe_details=resolution.safe_details), + attempt=project_execution_view(current), + ) + with self._confirm_task_terminal(execution_id): + return await self.inspect(execution_id, principal=principal) + if resolution.kind == "not_applied": now = datetime.now(timezone.utc) if ( diff --git a/linktools-ai/src/linktools/ai/history/_projection.py b/linktools-ai/src/linktools/ai/history/_projection.py index e219474163..08af6445e9 100644 --- a/linktools-ai/src/linktools/ai/history/_projection.py +++ b/linktools-ai/src/linktools/ai/history/_projection.py @@ -1683,6 +1683,7 @@ async def candidates() -> AsyncGenerator[tuple[int, object, int, _ProjectedHisto responses, request_steps, replayed_parts = _request_associations(events) tool_metadata = self._tool_call_metadata(events) selected: list[_HistoryOccurrence] = [] + original_responses: dict[int, ModelResponse] = {} for message_index, message, projected_offset, value in window: replayed_indices = replayed_parts.get(message_index + 1, ()) if replayed_indices and ( @@ -1710,12 +1711,17 @@ async def candidates() -> AsyncGenerator[tuple[int, object, int, _ProjectedHisto if len(original_sequences) > 1: raise AIError(ErrorCode.STORAGE_INTEGRITY_ERROR) if original_sequences: - original_messages = [original async for original in self._store_for(source.record).iter_message_range( - agent_run_id=agent_run_id, start=original_sequences[0] - 1, end=original_sequences[0], - )] - if len(original_messages) != 1 or not isinstance(original_messages[0], ModelResponse): - raise AIError(ErrorCode.STORAGE_INTEGRITY_ERROR) - original_calls = [part for part in original_messages[0].parts + original_sequence = original_sequences[0] + original_response = original_responses.get(original_sequence) + if original_response is None: + original_messages = [original async for original in self._store_for(source.record).iter_message_range( + agent_run_id=agent_run_id, start=original_sequence - 1, end=original_sequence, + )] + if len(original_messages) != 1 or not isinstance(original_messages[0], ModelResponse): + raise AIError(ErrorCode.STORAGE_INTEGRITY_ERROR) + original_response = original_messages[0] + original_responses[original_sequence] = original_response + original_calls = [part for part in original_response.parts if isinstance(part, ToolCallPart) and part.tool_call_id == value.tool_call_id] if (len(original_calls) != 1 or original_calls[0].tool_name != value.tool_name or original_calls[0].args_as_dict() != value.content): diff --git a/linktools-ai/src/linktools/ai/persistence/_history.py b/linktools-ai/src/linktools/ai/persistence/_history.py index 5487f8e855..477f00885f 100644 --- a/linktools-ai/src/linktools/ai/persistence/_history.py +++ b/linktools-ai/src/linktools/ai/persistence/_history.py @@ -1980,12 +1980,26 @@ async def _read_range_start( observed_head: TranscriptHeadRecord | None = None, ) -> tuple[StoredFact, ...]: """Materialize only the first bounded fact page inside the metadata read.""" - head = await self._resolve_observed_head(transaction, owner_id, observed_head) + batch_head_and_seek = observed_head is None and 0 < start < end + if batch_head_and_seek: + head_key = self._head_key(owner_id) + block = (start // _TRANSCRIPT_SEEK_BLOCK) * _TRANSCRIPT_SEEK_BLOCK + seek_key = self._seek_key(owner_id, block) + records = await transaction.get_records((head_key, seek_key)) + head_record = records.get(head_key) + head = self._validate_observed_head( + owner_id, None if head_record is None else self._decode_head(head_record), + ) + sequence = self._seek_record_sequence(records.get(seek_key), message_index=start) + else: + head = await self._resolve_observed_head(transaction, owner_id, observed_head) + sequence = None if end > head.message_count: raise AIError(ErrorCode.STORAGE_INTEGRITY_ERROR) if start == end: return () - sequence = await self._seek_fact_sequence(transaction, owner_id, start) + if not batch_head_and_seek: + sequence = await self._seek_fact_sequence(transaction, owner_id, start) return tuple(await transaction.list_facts(FactQuery( self._transcript_stream(owner_id), after_sequence=sequence, @@ -2004,8 +2018,12 @@ async def _seek_fact_sequence( return None block = (message_index // _TRANSCRIPT_SEEK_BLOCK) * _TRANSCRIPT_SEEK_BLOCK record = await transaction.get_record(self._seek_key(owner_id, block)) + return self._seek_record_sequence(record, message_index=message_index) + + def _seek_record_sequence(self, record: StoredRecord | None, *, message_index: int) -> int | None: if record is None: return None + block = (message_index // _TRANSCRIPT_SEEK_BLOCK) * _TRANSCRIPT_SEEK_BLOCK try: seek = _decode_enveloped_domain(record.data, TranscriptSeekRecord) except (AIError, TypeError, ValueError): diff --git a/linktools-ai/src/linktools/ai/persistence/_sql.py b/linktools-ai/src/linktools/ai/persistence/_sql.py index e986d574e4..e3096d96ef 100644 --- a/linktools-ai/src/linktools/ai/persistence/_sql.py +++ b/linktools-ai/src/linktools/ai/persistence/_sql.py @@ -249,6 +249,13 @@ async def read_content_references( if sum(map(len, extras.values())) > 32: await self.get_records(extras) extras = {} + else: + # Owner rows already use the same transaction's canonical read. + # Preserve the large-batch and typed-payload fallback routes. + extras = {store: tuple(key for key in requested + if store is not owner_store or key not in keys or owner._is_payload_key(key)) + for store, requested in extras.items()} + extras = {store: requested for store, requested in extras.items() if requested} extra_stores = {self.transaction(store)._store_key: store for store in extras} found: dict[StateStore, dict[bytes, StoredRecord]] = {owner_store: {}, content_store: {}} found.update({store: {} for store in extras if store not in found}) diff --git a/linktools-ai/src/linktools/ai/runtime/_execution.py b/linktools-ai/src/linktools/ai/runtime/_execution.py index 215c4787b1..ad5fd26204 100644 --- a/linktools-ai/src/linktools/ai/runtime/_execution.py +++ b/linktools-ai/src/linktools/ai/runtime/_execution.py @@ -2440,11 +2440,7 @@ async def load_confirmed_terminal() -> ExecutionRecord | ExecutionResultRecord: return publication.execution async def wait_once() -> ExecutionResult: - execution = await self.manager.load_result_authorized( - execution_id, - principal, - AuthorizationAction.EXECUTION_READ, - ) + execution = await load_confirmed_terminal() self._live_broker.abandon_prepared_local_producer(execution_id) status = execution.status poll_backoff = 1.0 diff --git a/linktools-ai/src/linktools/ai/runtime/_local.py b/linktools-ai/src/linktools/ai/runtime/_local.py index f988d76855..d584770121 100644 --- a/linktools-ai/src/linktools/ai/runtime/_local.py +++ b/linktools-ai/src/linktools/ai/runtime/_local.py @@ -5028,6 +5028,10 @@ async def sink(emission: "LiveDelta | DurableBoundary") -> None: agent_run_id, ) except asyncio.CancelledError as error: + # Cancellation can arrive outside the backend's request path. Once + # settlement starts, a later control request must only drain it. + if self._tasks.get(execution_id) is asyncio.current_task(): + self._worker_cancel_requests.add(execution_id) cleanup_details: dict[str, JsonValue] = {} if ( isinstance(error.__cause__, AIError) diff --git a/linktools-ai/src/linktools/ai/task/_handler.py b/linktools-ai/src/linktools/ai/task/_handler.py index ec4358a608..5384df5843 100644 --- a/linktools-ai/src/linktools/ai/task/_handler.py +++ b/linktools-ai/src/linktools/ai/task/_handler.py @@ -15,6 +15,7 @@ CorrelationData, TaskStatus, canonical_sha256, + canonical_json_bytes, normalize_json_value, normalize_correlation, ) @@ -172,17 +173,35 @@ async def publish( @dataclass(frozen=True, slots=True) class TaskEffectResolution: - """A read-only reconciliation result for an external task effect.""" + """An observed external Task effect outcome. + + ``failed`` explicitly accepts a stopped Task's partial effects while + recording failure. Only the admitted principal through the authorized graph control + operation accepts it; automatic reconciliation leaves that decision recovery-required. + ``safe_details`` must contain only safe diagnostic JSON, at most 8 KiB. + """ kind: str value: JsonValue | None = None + error_code: ErrorCode | None = None + safe_details: Mapping[str, JsonValue] = field(default_factory=dict) def __post_init__(self) -> None: - if self.kind not in {"applied", "not_applied", "unknown"}: + if self.kind not in {"applied", "not_applied", "unknown", "failed"}: raise ValueError("task effect resolution kind is invalid") if self.kind != "applied" and self.value is not None: raise ValueError("only an applied effect can carry a value") object.__setattr__(self, "value", normalize_json_value(self.value)) + if self.kind == "failed": + if not isinstance(self.error_code, ErrorCode): + raise ValueError("a failed effect requires an error code") + elif self.error_code is not None or self.safe_details: + raise ValueError("only a failed effect can carry error details") + details = AIError(self.error_code or ErrorCode.TASK_NODE_FAILED, + safe_details=self.safe_details).safe_details + if len(canonical_json_bytes(details)) > 8192: + raise ValueError("task effect error details exceed 8192 bytes") + object.__setattr__(self, "safe_details", ImmutableJsonMapping(details)) @dataclass(frozen=True, slots=True) diff --git a/linktools-ai/src/linktools/ai/task/_local.py b/linktools-ai/src/linktools/ai/task/_local.py index f126877490..b748d79c99 100644 --- a/linktools-ai/src/linktools/ai/task/_local.py +++ b/linktools-ai/src/linktools/ai/task/_local.py @@ -556,7 +556,7 @@ async def resolve_effect( and state.status is TaskStatus.READY ) or ( - resolution.kind == "applied" + resolution.kind in {"applied", "failed"} and state.status in {TaskStatus.SUCCEEDED, TaskStatus.FAILED} ) ) @@ -570,6 +570,26 @@ async def resolve_effect( ): raise AIError(ErrorCode.TASK_FENCE_STALE) + if resolution.kind == "failed": + # A recovery snapshot can be visible while sibling cleanup is + # still draining. Do not publish failure before owned work stops. + async with self._lock: + active = self._graphs.get((tenant_id, graph_id)) + draining = None if active is None else active.task + if draining is not None: + await asyncio.shield(draining) + siblings = await self._repository.graph_state(graph_id, tenant_id=tenant_id) + if siblings is None: + raise AIError(ErrorCode.STORAGE_NOT_FOUND) + if node.failure_policy == "propagate" and any( + sibling.node_id != node_id and sibling.execution_id is not None + and sibling.status not in _TERMINAL | {TaskStatus.RECOVERY_REQUIRED} + for sibling in siblings.node_states + ): + # Cancellation is a separate owner-controlled operation. A + # graph snapshot cannot authorize cancelling a newer attempt. + raise AIError(ErrorCode.TASK_NOT_READY) + invocation = TaskNodeInvocation( node, graph_id, diff --git a/linktools-ai/src/linktools/ai/task/_service_impl.py b/linktools-ai/src/linktools/ai/task/_service_impl.py index dfc6ea56b2..516598417e 100644 --- a/linktools-ai/src/linktools/ai/task/_service_impl.py +++ b/linktools-ai/src/linktools/ai/task/_service_impl.py @@ -1375,6 +1375,12 @@ async def resolve_effect( tenant_id, graph_state, ) + if ( + request.resolution.kind == "failed" + and principal_identity_payload(request.principal) + != principal_identity_payload(admission.principal) + ): + raise AIError(ErrorCode.AUTHORIZATION_DENIED) if self._preflight is not None: self._preflight.validate_effect_resolution( node, @@ -1392,6 +1398,10 @@ async def resolve_effect( "resolution": { "kind": request.resolution.kind, "value": request.resolution.value, + **({ + "error_code": request.resolution.error_code, + "safe_details": dict(request.resolution.safe_details), + } if request.resolution.kind == "failed" else {}), }, } ) diff --git a/linktools-ai/src/linktools/ai/task/engine/_runner.py b/linktools-ai/src/linktools/ai/task/engine/_runner.py index e48b89840d..3460275f2a 100644 --- a/linktools-ai/src/linktools/ai/task/engine/_runner.py +++ b/linktools-ai/src/linktools/ai/task/engine/_runner.py @@ -2229,11 +2229,13 @@ async def _reconcile_custom_execution( "reconcile_result_type": type(resolution).__name__, }, ) - if resolution.kind == "unknown": + if resolution.kind in {"unknown", "failed"}: + # Accepting partial effects as a terminal failure requires the + # authorized graph control request, not automatic reconciliation. raise TaskNodeRunError( ErrorCode.TASK_EFFECT_UNKNOWN, execution_id, - safe_details={"reconcile_outcome": "unknown"}, + safe_details={"reconcile_outcome": resolution.kind}, ) if node.effect_policy == "non_replay_safe": diff --git a/tests/ai/persistence/test_execution_start_service.py b/tests/ai/persistence/test_execution_start_service.py index 615ae90eb9..7cd5bddf23 100644 --- a/tests/ai/persistence/test_execution_start_service.py +++ b/tests/ai/persistence/test_execution_start_service.py @@ -236,6 +236,10 @@ async def test_execution_start_claim_has_one_launcher_winner() -> None: @pytest.mark.asyncio async def test_task_start_holds_immediate_terminal_execution_before_return() -> None: + from linktools.ai.core import StopReason, UsageMetrics + from linktools.ai.persistence.contracts import ResultRecord + from linktools.ai.storage import StoredPayload + state = Persistence(PersistenceStorage.in_memory()) await state.initialize(namespace="service-start-handoff", tenant_id="tenant") release_started = asyncio.Event() @@ -250,13 +254,25 @@ async def launch( execution: ExecutionRecord, ) -> None: await super().launch(request, execution) + current = await self._repository.get(execution.execution_id, tenant_id=self._repository.tenant_id) + await self._repository.compare_and_swap( + current.execution_id, tenant_id=self._repository.tenant_id, + expected_revision=current.revision, + next_record=replace(current, status=ExecutionStatus.SUCCEEDED, + revision=current.revision + 1, completed_at=current.updated_at, + result=ResultRecord(StoredPayload.inline_json("answer"), StopReason.END_TURN, + UsageMetrics(), current.updated_at)), + ) await self.service.request_terminal_handoff( execution.execution_id, tenant_id=self._repository.tenant_id, ) - async def release_terminal(execution_id: str, *, tenant_id: str) -> bool: - del execution_id, tenant_id + async def release_terminal(execution_id: str, *, tenant_id: str, terminal_commit: bool = False) -> bool: + current = await state.execution.executions.get(execution_id, tenant_id=tenant_id) + assert current is not None and current.status is ExecutionStatus.SUCCEEDED + assert terminal_commit is current.conversation + assert not current.dependency_hold_ids release_started.set() await release_finished.wait() return True @@ -282,9 +298,11 @@ async def release_terminal(execution_id: str, *, tenant_id: str) -> bool: hold_id=hold_id, ) ) - await release_started.wait() - release_finished.set() - await cleanup + try: + await asyncio.wait_for(release_started.wait(), 5) + finally: + release_finished.set() + await asyncio.wait_for(cleanup, 5) finally: await state.close() diff --git a/tests/ai/test_canonical_history_read_boundaries.py b/tests/ai/test_canonical_history_read_boundaries.py index 72d58461c8..4605cb3fad 100644 --- a/tests/ai/test_canonical_history_read_boundaries.py +++ b/tests/ai/test_canonical_history_read_boundaries.py @@ -244,7 +244,10 @@ async def test_history_and_trace_pagination_preserve_cross_run_and_child_order(t @pytest.mark.asyncio -async def test_history_replay_projection_omits_only_duplicate_calls_at_original_coordinates(tmp_path: Path) -> None: +@pytest.mark.parametrize("changed_replay", (False, True)) +async def test_history_replay_projection_omits_only_duplicate_calls_at_original_coordinates( + tmp_path: Path, changed_replay: bool, +) -> None: from pydantic_ai.messages import ToolCallPart, ToolReturnPart storage = Persistence(PersistenceStorage.sqlite(tmp_path / "replayed-calls.db")) @@ -259,17 +262,21 @@ async def test_history_replay_projection_omits_only_duplicate_calls_at_original_ run_id, agent_conversation_id=conversation_id, metadata={"agent_run_seq": "1"}, ), execution_id="root") original = ToolCallPart("echo", {"value": "original"}, tool_call_id="replayed") + second = ToolCallPart("echo", {"value": "second"}, tool_call_id="replayed-second") repeated = ModelResponse(parts=[ - original, TextPart("new text"), ToolCallPart("echo", {"value": "new"}, tool_call_id="new"), + original, + replace(second, args={"value": "changed"}) if changed_replay else second, + TextPart("new text"), + ToolCallPart("echo", {"value": "new"}, tool_call_id="new"), ]) messages = ( - ModelRequest(parts=[UserPromptPart("question")]), ModelResponse(parts=[original]), + ModelRequest(parts=[UserPromptPart("question")]), ModelResponse(parts=[original, second]), ModelRequest(parts=[ToolReturnPart("echo", "original", tool_call_id="replayed")]), repeated, ) repository = archive.transcript_repository prepared = await repository.prepare_observation( run_id, messages[:3], first_message_index=0, pending=repeated, - pending_keys=("part:0", "part:1", "part:2"), + pending_keys=("part:0", "part:1", "part:2", "part:3"), ) await archive.state_store.mutate(lambda transaction: repository.commit_observation(transaction, prepared)) async def append_request(request: int, message: int, call_id: str) -> None: @@ -277,7 +284,7 @@ async def append_request(request: int, message: int, call_id: str) -> None: run_id, "MODEL_REQUEST_SUCCEEDED", request, agent_conversation_id=conversation_id, metadata={ "linktools.ai.model_request_seq": str(request), "linktools.ai.message_seq": str(message), - **({"linktools.ai.replayed_tool_call_indices": "[0]"} if request == 2 else {}), + **({"linktools.ai.replayed_tool_call_indices": "[0,1]"} if request == 2 else {}), }, ), execution_id="root") await archive.append_event(StepEvent( @@ -286,14 +293,26 @@ async def append_request(request: int, message: int, call_id: str) -> None: metadata={"linktools.ai.model_request_seq": str(request)}, ), execution_id="root") await append_request(1, 2, "replayed") + await archive.append_event(StepEvent( + run_id, "TOOL_CALL_STARTED", 1, tool_name="echo", tool_call_id="replayed-second", + agent_conversation_id=conversation_id, + metadata={"linktools.ai.model_request_seq": "1"}, + ), execution_id="root") reader = StepExecutionHistoryReader( namespace="history", executions=_Executions(_record("root")), store=archive, cursor_signer=HmacCursorSigner("history", b"history"), ) + if changed_replay: + from linktools.ai.errors import AIError, ErrorCode + + with pytest.raises(AIError) as invalid: + await reader.history("root", tenant_id="tenant", cursor=None, limit=100) + assert invalid.value.code is ErrorCode.STORAGE_INTEGRITY_ERROR + return pending = await reader.history("root", tenant_id="tenant", cursor=None, limit=100) frozen = await reader.history("root", tenant_id="tenant", cursor=None, limit=1) assert [(item.item_kind, item.part_index) for item in pending.items if item.message_seq == 4] == [ - ("assistant", 1), ("tool_call", 2), + ("assistant", 2), ("tool_call", 3), ] assert all(item.model_request_seq is None for item in pending.items if item.message_seq == 4) completed = await repository.prepare_observation( @@ -317,13 +336,13 @@ async def append_request(request: int, message: int, call_id: str) -> None: if cursor is None: break assert [(item.item_kind, item.message_seq, item.part_index, item.model_request_seq) for item in items] == [ - ("user", 1, 0, None), ("tool_call", 2, 0, 1), ("tool_result", 3, None, 1), - ("assistant", 4, 1, 2), ("tool_call", 4, 2, 2), + ("user", 1, 0, None), ("tool_call", 2, 0, 1), ("tool_call", 2, 1, 1), ("tool_result", 3, None, 1), + ("assistant", 4, 2, 2), ("tool_call", 4, 3, 2), ] assert (await reader.history("root", tenant_id="tenant", cursor=None, limit=10, message_seq=4, part_index=0)).items == () selected = await reader.history("root", tenant_id="tenant", cursor=None, limit=10, - message_seq=4, part_index=2) + message_seq=4, part_index=3) assert selected.items == (items[-1],) stored = [message async for message in archive.iter_message_range(agent_run_id=run_id, start=3, end=4)] assert stored == [repeated] diff --git a/tests/ai/test_confirmed_turn_receipts_independent.py b/tests/ai/test_confirmed_turn_receipts_independent.py index f5300b23d0..b3926187dc 100644 --- a/tests/ai/test_confirmed_turn_receipts_independent.py +++ b/tests/ai/test_confirmed_turn_receipts_independent.py @@ -3,10 +3,16 @@ """Confirmed publication receipts preserve live holds and public read boundaries.""" import asyncio +import json +import os +from collections import Counter +from pathlib import Path from dataclasses import replace import pytest from pydantic_ai.messages import ModelResponse, TextPart +from sqlalchemy import event +from sqlalchemy.engine import Engine from linktools.ai.core import ExecutionStatus, StopReason, UsageMetrics from linktools.ai.errors import AIError, ErrorCode @@ -187,3 +193,98 @@ def worker_failure(execution_id, *, tenant_id): if not task.done(): task.cancel() await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.parametrize("receipt_available", (True, False)) +@pytest.mark.parametrize("blocker", (None, "authorization", "worker_failure", "source_unavailable")) +async def test_wait_started_after_publication_preserves_read_guards( + tmp_path, monkeypatch, receipt_available, blocker, +): + entered, release = asyncio.Event(), asyncio.Event() + + async def respond(messages, info): + entered.set() + await release.wait() + return ModelResponse(parts=[TextPart("finished")]) + + async with Persistence.open(_NAMESPACE, storage=_storage("sqlite", tmp_path)) as state: + async with Runtime.attach(persistence=state, features=model_features( + _Models(_UsageFunctionModel(respond)), _features(), + )) as runtime: + session = await runtime.agents.get().create_session("late-receipt-wait") + execution = await session.start("question") + await asyncio.wait_for(entered.wait(), 5) + service = runtime._execution_service + backend = service._backend + await service.acquire_dependency_hold( + execution.execution_id, tenant_id=runtime.tenant_id, hold_id="consumer", + ) + worker = backend._tasks[execution.execution_id] + release.set() + await asyncio.wait_for(asyncio.shield(worker), 5) + publication = state.retention.terminal_publication( + execution.execution_id, tenant_id=runtime.tenant_id, + ) + assert publication is not None + if not receipt_available: + monkeypatch.setattr(service, "_retained_publication", None) + + # A retained result still requires the current authorization decision. + authorizations = [] + original_authorize = service.manager.authorize_loaded_execution + + async def authorize(value, *args, **kwargs): + authorizations.append(value) + if blocker == "authorization": + raise AIError(ErrorCode.AUTHORIZATION_DENIED) + await original_authorize(value, *args, **kwargs) + + monkeypatch.setattr(service.manager, "authorize_loaded_execution", authorize) + if blocker == "worker_failure": + failure = AIError(ErrorCode.STORAGE_INTEGRITY_ERROR) + monkeypatch.setattr(backend, "worker_failure", lambda *args, **kwargs: failure) + + if blocker == "source_unavailable": + async def unavailable(*args, **kwargs): + raise AIError(ErrorCode.STORAGE_UNAVAILABLE) + monkeypatch.setattr(state.execution.executions, "get_result_record", unavailable) + + async def wait_result(): + operation = service.wait( + execution.execution_id, principal=runtime.default_principal, timeout_seconds=5, + ) + output = os.environ.get("RECEIPT_WAIT_IO_OUTPUT") + if output is None: + return await operation + counts = Counter() + def statement(_connection, _cursor, sql, _parameters, _context, _many): + counts[sql.lstrip().split()[0].upper()] += 1 + def committed(_connection): + counts["COMMIT"] += 1 + event.listen(Engine, "before_cursor_execute", statement) + event.listen(Engine, "commit", committed) + try: + return await operation + finally: + event.remove(Engine, "before_cursor_execute", statement) + event.remove(Engine, "commit", committed) + with Path(output).open("a", encoding="utf-8") as stream: + stream.write(json.dumps({"receipt_available": receipt_available, + "blocker": blocker, "sql": dict(counts)}) + "\n") + + if blocker is None or blocker == "source_unavailable" and receipt_available: + result = await wait_result() + assert result.status is ExecutionStatus.SUCCEEDED + assert result.output == {"text": "finished"} + else: + with pytest.raises(AIError) as raised: + await wait_result() + assert raised.value.code is ( + ErrorCode.AUTHORIZATION_DENIED if blocker == "authorization" + else ErrorCode.STORAGE_UNAVAILABLE if blocker == "source_unavailable" + else ErrorCode.STORAGE_INTEGRITY_ERROR + ) + if blocker != "source_unavailable" or receipt_available: + assert authorizations and all(value.execution_id == execution.execution_id for value in authorizations) + retained = await state.execution.executions.get(execution.execution_id, tenant_id=runtime.tenant_id) + assert not retained.retention_closed and retained.dependency_hold_ids == ("consumer",) diff --git a/tests/ai/test_content_reference_owner_reads.py b/tests/ai/test_content_reference_owner_reads.py new file mode 100644 index 0000000000..811eb1faaf --- /dev/null +++ b/tests/ai/test_content_reference_owner_reads.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Overlapping content-read roles keep one physical owner and exact partitions.""" + +from dataclasses import replace +import hashlib +from pathlib import Path + +import pytest + +from linktools.ai.persistence import _sql +from linktools.ai.persistence._archive import ContentRecord, archive_scope_key +from linktools.ai.persistence.state import StoredRecord, record_key, active_state_scope + +from .test_state_group_record_batch import _database, _key, _record + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner_count", (0, 1, 40)) +async def test_content_read_overlap_preserves_current_owner_and_missing_partition( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, owner_count: int, +) -> None: + async with _database(tmp_path / "overlap.db") as (_, first, second): + owners = tuple(replace(_record(f"owner-{index}"), data={"body": "x" * 65536, "index": index}) + for index in range(owner_count)) + common = owners[-1].key if owners else _key("extra-only") + extra = _record("extra", key=common) + foreign = replace(extra, data={"partition": "second"}) + await first.mutate(lambda tx: tx.insert_records(owners or (extra,))) + await second.mutate(lambda tx: tx.insert_record(foreign)) + _, content = await first.storage_group.open_archive(create=True) + missing = _key("missing") + decoded = [] + original = _sql._record_from_row + + def capture(row): + value = original(row) + decoded.append((row["store_key"], value.key)) + return value + + monkeypatch.setattr(_sql, "_record_from_row", capture) + + async def read(group): + values = await active_state_scope().record_reader.read_content_references( + owner_store=first, owner_keys=(*(record.key for record in owners), missing) if owners else (), + content_store=content, content_scope_key=archive_scope_key("overlap", "tenant"), + content_paths=tuple(("unused", str(index)) for index in range(16)), + additional_records={first: (common, missing), second: (common, missing)}, + ) + first_tx, second_tx = group.transaction(first), group.transaction(second) + assert await first_tx.get_record(common) == (owners[-1] if owners else extra) + assert await second_tx.get_record(common) == foreign + assert first_tx.cached_records((missing,)) == {missing: None} + assert second_tx.cached_records((missing,)) == {missing: None} + if owners: + assert values[first] == {record.key: record for record in owners} + before = len(decoded) + await group.get_records({first: (common, missing), second: (common, missing)}) + assert len(decoded) == before + + await first.storage_group.read_many((first, second, content), read) + assert decoded.count((first.store_key, common)) == 1 + assert decoded.count((second.store_key, common)) == 1 + + +@pytest.mark.asyncio +async def test_large_additional_batch_keeps_typed_payload_route(tmp_path: Path) -> None: + async with _database(tmp_path / "payload.db") as (_, first, _second): + _, archive = await first.storage_group.open_archive(create=True) + body = b"canonical payload" + value = ContentRecord("execution-payload", hashlib.sha256(body).hexdigest(), len(body), "raw", body=body) + key = record_key("overlap", "tenant", "archive", "archive_content", (value.identity,)) + payload = StoredRecord(key, archive_scope_key("overlap", "tenant"), None, "archive_content", + value.identity.encode(), None, 0, None, 0, None, value.to_json(), payload_bytes=body) + owners = tuple(_record(f"owner-{index}") for index in range(32)) + await archive.mutate(lambda tx: tx.insert_records((*owners, payload))) + + async def read(group): + await active_state_scope().record_reader.read_content_references( + owner_store=archive, owner_keys=tuple(record.key for record in owners), content_store=archive, + content_scope_key=archive_scope_key("overlap", "tenant"), content_paths=(("unused",),), + additional_records={archive: (*(record.key for record in owners), key)}, + ) + tx = group.transaction(archive) + assert await tx.get_record(key) == payload + assert await tx.get_records(tuple(record.key for record in owners)) == { + record.key: record for record in owners} + + await first.storage_group.read_many((archive,), read) diff --git a/tests/ai/test_dependency_hold_handoff_reads.py b/tests/ai/test_dependency_hold_handoff_reads.py index 359ecb1ca2..b9434a080b 100644 --- a/tests/ai/test_dependency_hold_handoff_reads.py +++ b/tests/ai/test_dependency_hold_handoff_reads.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- """Hold release reuses committed terminal status, not retention authority.""" +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock @@ -43,3 +44,25 @@ async def test_hold_release_failure_does_not_request_cleanup() -> None: await service.release_dependency_hold("execution", tenant_id="tenant", hold_id="hold") service.load_execution.assert_not_awaited() service.request_terminal_handoff.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", (RuntimeError("conflict"), asyncio.CancelledError())) +async def test_failed_durable_release_does_not_strand_local_consumers(failure: BaseException) -> None: + service = object.__new__(TaskExecutionService) + service._repositories = SimpleNamespace(executions=SimpleNamespace( + release_dependency_hold=AsyncMock(side_effect=failure))) + service._handoff = HandoffGate() + service.request_terminal_handoff = AsyncMock() + key = ("tenant", "execution") + await service._handoff.acquire_hold(key, "hold") + await service._handoff.request_release(key) + with pytest.raises(type(failure)): + await service.release_dependency_hold("execution", tenant_id="tenant", hold_id="hold") + # A failed persistence await must not leave cleanup owning the local gate. + consumer = await asyncio.wait_for(service._handoff.enter(key), 1) + await service._handoff.acquire_hold(key, "other-hold") + assert not await service._handoff.leave(key, consumer) + _, owner = await service._handoff.release_hold(key, "other-hold") + assert owner + service.request_terminal_handoff.assert_not_awaited() diff --git a/tests/ai/test_error_cancellation_contract.py b/tests/ai/test_error_cancellation_contract.py index 2eca921071..7c1c9c3757 100644 --- a/tests/ai/test_error_cancellation_contract.py +++ b/tests/ai/test_error_cancellation_contract.py @@ -303,10 +303,13 @@ async def test_immediate_terminal_execution_waits_for_task_dependency_hold(execu release_started = asyncio.Event() release_finished = asyncio.Event() - async def release_terminal(execution_id: str, *, tenant_id: str) -> None: - del execution_id, tenant_id + async def release_terminal(execution_id: str, *, tenant_id: str, terminal_commit: bool = False) -> bool: + assert terminal_commit is True + current = await execution_storage.execution.executions.get(execution_id, tenant_id=tenant_id) + assert current is not None and not current.dependency_hold_ids release_started.set() await release_finished.wait() + return True service._release_terminal = release_terminal await service.acquire_dependency_hold( @@ -327,9 +330,11 @@ async def release_terminal(execution_id: str, *, tenant_id: str) -> None: hold_id="task:graph:node", ) ) - await release_started.wait() - release_finished.set() - await asyncio.gather(cleanup, release_hold) + try: + await asyncio.wait_for(release_started.wait(), 5) + finally: + release_finished.set() + await asyncio.wait_for(asyncio.gather(cleanup, release_hold), 5) @pytest.mark.asyncio diff --git a/tests/ai/test_error_contract_matrix.py b/tests/ai/test_error_contract_matrix.py index 31aecd6248..51d9d15e84 100644 --- a/tests/ai/test_error_contract_matrix.py +++ b/tests/ai/test_error_contract_matrix.py @@ -215,6 +215,7 @@ async def load_authorized( @pytest.mark.asyncio async def test_execution_wait_timeout_has_stable_code() -> None: service = object.__new__(DefaultExecutionService) + service._retained_publication = None async def inspect(execution_id: str, *, principal: Principal) -> object: del execution_id, principal diff --git a/tests/ai/test_history_exact_read_io.py b/tests/ai/test_history_exact_read_io.py new file mode 100644 index 0000000000..905bf680d3 --- /dev/null +++ b/tests/ai/test_history_exact_read_io.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Exact History reads through the public Runtime boundary.""" + +import asyncio +import json +import os +from collections import Counter +from contextvars import ContextVar +from dataclasses import replace +from pathlib import Path + +import pytest +from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart, ToolReturnPart +from sqlalchemy import event +from sqlalchemy.engine import Engine + +from linktools.ai.capability import AgentContext, ExecutionCapability, StepPersistenceCapability, TraceCapability +from linktools.ai.core import CommitPolicy, ExecutionStatus, Principal +from linktools.ai.errors import AIError, ErrorCode +from linktools.ai.history import History +from linktools.ai.persistence import Persistence, PersistenceStorage +from linktools.ai.runtime import Runtime, RuntimeFeatures +from linktools.ai.storage import InMemoryCacheBackend + +from ._runtime_test_helpers import _Models, _UsageFunctionModel, model_features +from .test_runtime_storage_io_session import _listen_sql, _record_selects +from .test_transcript_sparse_windows import _history + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hint", ("present", "missing", "invalid")) +async def test_nonzero_transcript_window_batches_head_and_seek(tmp_path: Path, hint: str) -> None: + state = Persistence(PersistenceStorage.sqlite(tmp_path / "window.db")) + try: + repository = await _history(state, 4) + if hint != "present": + async def change_hint(transaction): + key = repository._seek_key("history", 0) + current = await transaction.get_record(key) + assert current is not None + if hint == "missing": + await transaction.delete_record(key) + else: + assert await transaction.replace_record( + replace(current, data={}, storage_revision=current.storage_revision + 1), + expected_storage_revision=current.storage_revision, + ) + await repository._store.mutate(change_hint) + engine = repository._store.storage_group.context.engine + statements = [] + listener = _listen_sql(engine, statements) + try: + values = [message async for message in repository.iter_message_range("history", start=1, end=3)] + finally: + event.remove(engine.sync_engine, "before_cursor_execute", listener) + assert [message.parts[0].content for message in values] == ["1", "2"] + metadata = [item for item in _record_selects(statements) if repository._store.store_key in item[1]] + assert len(metadata) == 1 + finally: + await state.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cached", (False, True)) +@pytest.mark.parametrize("include_content", (False, True)) +async def test_exact_tool_history_keeps_authorized_projection( + tmp_path: Path, cached: bool, include_content: bool, +) -> None: + model_calls = 0 + effects = [] + + def model(messages, info): + nonlocal model_calls + model_calls += 1 + returned = [part for message in messages for part in message.parts if isinstance(part, ToolReturnPart)] + if not returned: + return ModelResponse(parts=[ToolCallPart("lookup", {"index": index}, tool_call_id=f"call-{index}") + for index in range(12)]) + assert len(returned) == 12 + return ModelResponse(parts=[TextPart("done")]) + + async def lookup(ctx: AgentContext[None], index: int) -> str: + """Read one deterministic synthetic value.""" + effects.append(index) + return "item" * 100 + str(index) + + features = RuntimeFeatures("history-read") + features.agent("default", model="default", allow_tools=("lookup",)) + features.tool(lookup, effect_policy="none") + features.capability(ExecutionCapability(retain_history=True)) + features.capability(StepPersistenceCapability()) + features.capability(TraceCapability(commit=CommitPolicy.sync())) + owner: ContextVar[asyncio.Task | None] = ContextVar("history_read_owner", default=None) + counts = Counter() + measurements = [] + output = os.environ.get("HISTORY_IO_OUTPUT") + + def active(): + return owner.get() is not None + + def sql(connection, cursor, statement, parameters, context, executemany): + if active(): + counts[statement.split()[0].upper()] += 1 + + def commit(connection): + if active(): + counts["commits"] += 1 + + def rollback(connection): + if active(): + counts["rollbacks"] += 1 + + def storage(): + return PersistenceStorage.from_root(tmp_path / "runtime", cache=( + InMemoryCacheBackend(max_bytes=16 * 1024 * 1024) if cached else None)) + + async def exact(reader, execution_id, principal, phase): + counts.clear() + token = owner.set(asyncio.current_task()) + try: + page = await reader.history( + execution_id, principal=principal, tool_call_id="call-11", + agent_run_seq=1, model_request_seq=1, include_content=include_content, limit=2, + ) + finally: + owner.reset(token) + measurements.append({"phase": phase, "counts": dict(counts)}) + return page + + if output: + event.listen(Engine, "before_cursor_execute", sql) + event.listen(Engine, "commit", commit) + event.listen(Engine, "rollback", rollback) + try: + async with Runtime.open("history-read", features=model_features(_Models(_UsageFunctionModel(model)), features), + storage=storage()) as runtime: + execution = await runtime.agents.get().start("read synthetic items") + assert (await execution.wait(timeout_seconds=15)).result.status is ExecutionStatus.SUCCEEDED + principal = runtime.default_principal + full = await runtime.executions.history(execution.execution_id, principal=principal, + include_content=include_content, limit=100) + expected = tuple(item for item in full.items if item.tool_call_id == "call-11") + selected = await exact(runtime.executions, execution.execution_id, principal, "live-warm") + assert selected.items == expected + assert len(selected.items) == 2 and selected.next_cursor is None + assert model_calls == 2 and sorted(effects) == list(range(12)) + async with Persistence.open("history-read", storage=storage(), read_only=True) as persistence, History.open( + persistence=persistence, + ) as history: + cold = await exact(history, execution.execution_id, principal, "reopened-cold") + assert cold.items == expected and cold.next_cursor is None + warm = await exact(history, execution.execution_id, principal, "reopened-warm") + assert warm.items == expected and warm.next_cursor is None + items, cursor = [], None + while True: + page = await history.history(execution.execution_id, principal=principal, tool_call_id="call-11", + agent_run_seq=1, model_request_seq=1, include_content=include_content, limit=1, cursor=cursor) + items.extend(page.items) + cursor = page.next_cursor + if cursor is None: + break + assert tuple(items) == expected + assert not (await history.history(execution.execution_id, principal=principal, + tool_call_id="missing", limit=2)).items + for identity, actor in (("missing-execution", principal), + (execution.execution_id, Principal(principal.principal_id, "other-tenant"))): + with pytest.raises(AIError) as error: + await history.history(identity, principal=actor, tool_call_id="call-11", limit=2) + assert error.value.code is ErrorCode.AUTHORIZATION_DENIED + finally: + if output: + event.remove(Engine, "before_cursor_execute", sql) + event.remove(Engine, "commit", commit) + event.remove(Engine, "rollback", rollback) + if output: + with Path(output).open("a", encoding="utf-8") as stream: + stream.write(json.dumps({"label": os.environ.get("HISTORY_IO_LABEL", "current"), "cached": cached, + "include_content": include_content, "measurements": measurements}) + "\n") diff --git a/tests/ai/test_joint_model_history_boundary.py b/tests/ai/test_joint_model_history_boundary.py index 8a90672c86..db89f8d837 100644 --- a/tests/ai/test_joint_model_history_boundary.py +++ b/tests/ai/test_joint_model_history_boundary.py @@ -149,6 +149,32 @@ async def observe_caller(cls, model): monkeypatch.setattr(native, "write_prepared_step", write_step) monkeypatch.setattr(StateStepArchive, "sync_prepared_projection", sync_projection) + if fault == "cancel_commit": + from linktools.ai.runtime._local import LocalExecutionBackend + + cleanup_entered, cleanup_release = asyncio.Event(), asyncio.Event() + cancel = LocalExecutionBackend.cancel + finish_revoked = LocalExecutionBackend._finish_revoked_worker + request_cancel = LocalExecutionBackend._request_worker_cancel + + async def cancel_during_cleanup(self, execution, **kwargs): + await asyncio.wait_for(cleanup_entered.wait(), 5) + return await cancel(self, execution, **kwargs) + + async def pause_cleanup(self, *args, **kwargs): + if injected and asyncio.current_task() is caller: + cleanup_entered.set() + await asyncio.wait_for(cleanup_release.wait(), 5) + return await finish_revoked(self, *args, **kwargs) + + def release_cleanup(self, execution_id, task): + request_cancel(self, execution_id, task) + cleanup_release.set() + + monkeypatch.setattr(LocalExecutionBackend, "cancel", cancel_during_cleanup) + monkeypatch.setattr(LocalExecutionBackend, "_finish_revoked_worker", pause_cleanup) + monkeypatch.setattr(LocalExecutionBackend, "_request_worker_cancel", release_cleanup) + async def provider(messages, info): provider_calls.append(True) return ModelResponse(parts=[TextPart("done")]) diff --git a/tests/ai/test_runtime_boundary_contract.py b/tests/ai/test_runtime_boundary_contract.py index c41a9ef8ea..61854dc27d 100644 --- a/tests/ai/test_runtime_boundary_contract.py +++ b/tests/ai/test_runtime_boundary_contract.py @@ -251,6 +251,7 @@ async def test_second_stream_close_does_not_close_first_subscription() -> None: @pytest.mark.asyncio async def test_wait_authorizes_before_abandoning_stream() -> None: service = object.__new__(DefaultExecutionService) + service._retained_publication = None service._runtime_bridge = _ExecutionRuntimeBridge() async def denied( @@ -412,6 +413,7 @@ def worker_installed(self, execution_id: str) -> bool: @pytest.mark.asyncio async def test_wait_timeout_includes_initial_authorized_read() -> None: service = object.__new__(DefaultExecutionService) + service._retained_publication = None service._runtime_bridge = _ExecutionRuntimeBridge() async def slow_authorized( diff --git a/tests/ai/test_task_failed_effect_resolution.py b/tests/ai/test_task_failed_effect_resolution.py new file mode 100644 index 0000000000..4f4fc1a048 --- /dev/null +++ b/tests/ai/test_task_failed_effect_resolution.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Explicit failure acceptance preserves already performed Task effects.""" +import asyncio +from pathlib import Path + +import pytest + +from linktools.ai.core import Principal, TaskStatus +from linktools.ai.errors import AIError, ErrorCode +from linktools.ai.persistence import Persistence, PersistenceStorage +from linktools.ai.task import Task, TaskEffectResolution, TaskGraph, TaskNode +from linktools.ai.task.engine import TaskEngine + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sql", (False, True)) +@pytest.mark.parametrize("lost_stage", ("graph", "execution")) +async def test_failed_effect_resolution_retains_partial_effect_and_reopens(tmp_path: Path, sql: bool, lost_stage: str, monkeypatch: pytest.MonkeyPatch) -> None: + effects: list[str] = [] + principal = Principal("owner", "default") + + + async def body(ctx): + effects.append("applied") + raise RuntimeError("failed after partial effect") + + async def dependent(ctx): + effects.append("dependent") + return "unexpected" + + task = Task("partial", body, effect_policy="non_replay_safe") + child = Task("dependent", dependent, effect_policy="none") + storage = PersistenceStorage.sqlite(tmp_path / "task.sqlite") if sql else PersistenceStorage.from_root(tmp_path / "local", exclusive_writer=True) + for first in (True, False): + async with Persistence.open("partial", storage=storage) as persistence, TaskEngine.open(task, child, persistence=persistence) as engine: + if first: + run = await engine.start(TaskGraph("partial", (TaskNode("n", task=task), TaskNode("dependent", task=child, dependencies=("n",)))), principal=principal, idempotency_key="start") + assert (await run.wait(timeout_seconds=5)).result.status is TaskStatus.RECOVERY_REQUIRED + node = next(n for n in (await run.state()).node_states if n.node_id == "n") + fence = node.fence + assert node.error_code == ErrorCode.TASK_EFFECT_UNKNOWN.value + assert effects == ["applied"] + resolution = TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED, safe_details={"phase": "after_partial_effect"}) + await run.resolve_effect("n", fence, TaskEffectResolution("unknown"), idempotency_key="unknown") + assert (await run.state()).status is TaskStatus.RECOVERY_REQUIRED + with pytest.raises(AIError) as stale: + await run.resolve_effect("n", fence + 1, resolution, idempotency_key="stale") + assert stale.value.code in {ErrorCode.TASK_NOT_READY, ErrorCode.TASK_FENCE_STALE} + with pytest.raises(AIError) as denied: + outsider = await engine.get("partial", principal=Principal("outsider", persistence.tenant_id)) + await outsider.resolve_effect("n", fence, resolution, idempotency_key="other") + assert denied.value.code is ErrorCode.AUTHORIZATION_DENIED + owner = persistence.task.tasks if lost_stage == "graph" else persistence.execution.executions + method = "fail" if lost_stage == "graph" else "commit_terminal" + original_fail = getattr(owner, method) + async def lost_response(*args, **kwargs): + await original_fail(*args, **kwargs) + monkeypatch.setattr(owner, method, original_fail) + raise AIError(ErrorCode.STORAGE_COMMIT_UNKNOWN) + monkeypatch.setattr(owner, method, lost_response) + with pytest.raises(AIError) as lost: + await run.resolve_effect("n", fence, resolution, idempotency_key="failure") + assert lost.value.code is ErrorCode.STORAGE_COMMIT_UNKNOWN + else: + run = await engine.get("partial", principal=principal) + await run.resolve_effect("n", fence, resolution, idempotency_key="failure") + state = await run.state() + assert state.status is TaskStatus.FAILED + assert next(n for n in state.node_states if n.node_id == "n").status is TaskStatus.FAILED + assert next(n for n in state.node_states if n.node_id == "n").error_code == ErrorCode.TASK_NODE_FAILED.value + assert effects == ["applied"] + with pytest.raises(AIError) as conflict: + await run.resolve_effect("n", fence, TaskEffectResolution("failed", error_code=ErrorCode.EXECUTION_WAIT_TIMEOUT), idempotency_key="failure") + assert conflict.value.code is ErrorCode.IDEMPOTENCY_CONFLICT + + +@pytest.mark.asyncio +async def test_failed_effect_resolution_does_not_enable_automatic_failure_acceptance(tmp_path: Path) -> None: + calls = 0 + + async def body(ctx): + nonlocal calls + calls += 1 + raise RuntimeError("partial effect") + + async def reconcile(ctx): + return TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED) + + task = Task("partial", body, effect_policy="non_replay_safe", reconcile=reconcile) + async with Persistence.open("partial", storage=PersistenceStorage.sqlite(tmp_path / "state.db")) as persistence, TaskEngine.open(task, persistence=persistence) as engine: + run = await engine.start(TaskGraph("graph", (TaskNode("n", task=task),)), idempotency_key="start") + assert (await run.wait(timeout_seconds=5)).result.status is TaskStatus.RECOVERY_REQUIRED + await run.recover(idempotency_key="recover") + assert (await run.wait(timeout_seconds=5)).result.status is TaskStatus.RECOVERY_REQUIRED + assert calls == 1 + + +@pytest.mark.asyncio +async def test_failed_effect_resolution_parallel_confirmations(tmp_path: Path) -> None: + calls = 0 + + async def body(ctx): + nonlocal calls + calls += 1 + raise RuntimeError("partial effect") + + task = Task("partial", body, effect_policy="non_replay_safe") + async with Persistence.open("parallel", storage=PersistenceStorage.sqlite(tmp_path / "state.db")) as persistence, TaskEngine.open(task, persistence=persistence) as engine: + run = await engine.start(TaskGraph("parallel", (TaskNode("n", task=task),)), idempotency_key="start") + assert (await run.wait(timeout_seconds=5)).result.status is TaskStatus.RECOVERY_REQUIRED + fence = (await run.state()).node_states[0].fence + resolution = TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED) + results = await asyncio.gather(*(run.resolve_effect("n", fence, resolution, idempotency_key="accept") for _ in range(2)), return_exceptions=True) + assert any(not isinstance(result, BaseException) for result in results) + for result in results: + if isinstance(result, BaseException): + assert isinstance(result, AIError) + assert result.code in {ErrorCode.STORAGE_CONFLICT, ErrorCode.TASK_FENCE_STALE, ErrorCode.EXECUTION_RESULT_CONFLICT} + await run.resolve_effect("n", fence, resolution, idempotency_key="accept") + assert (await run.state()).status is TaskStatus.FAILED + assert calls == 1 + + +@pytest.mark.asyncio +async def test_failed_effect_resolution_waits_for_sibling_drain(tmp_path: Path) -> None: + entered, cleanup, release, drained = (asyncio.Event() for _ in range(4)) + + async def partial(ctx): + await entered.wait() + raise RuntimeError("partial effect") + + async def sibling(ctx): + entered.set() + try: + await asyncio.Event().wait() + finally: + cleanup.set() + await release.wait() + drained.set() + + task = Task("partial", partial, effect_policy="non_replay_safe") + other = Task("sibling", sibling, effect_policy="none") + async with Persistence.open("drain", storage=PersistenceStorage.sqlite(tmp_path / "state.db")) as persistence, TaskEngine.open(task, other, persistence=persistence) as engine: + run = await engine.start(TaskGraph("graph", (TaskNode("n", task=task), TaskNode("sibling", task=other))), idempotency_key="start") + await asyncio.wait_for(cleanup.wait(), 5) + state = await run.state() + node = next(n for n in state.node_states if n.node_id == "n") + async with Persistence.open("drain", storage=PersistenceStorage.sqlite(tmp_path / "state.db")) as second_state, TaskEngine.open(task, other, persistence=second_state) as second_engine: + second_run = await second_engine.get("graph") + with pytest.raises(AIError) as foreign: + await second_run.resolve_effect("n", node.fence, TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED), idempotency_key="second-owner") + assert foreign.value.code is ErrorCode.TASK_NOT_READY + assert not drained.is_set() + assert (await second_run.state()).status is TaskStatus.RECOVERY_REQUIRED + acceptance = asyncio.create_task(run.resolve_effect("n", node.fence, TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED), idempotency_key="accept")) + try: + await asyncio.sleep(0.02) + assert not acceptance.done() + finally: + release.set() + with pytest.raises(AIError) as unsettled: + await asyncio.wait_for(acceptance, 5) + assert unsettled.value.code is ErrorCode.TASK_NOT_READY + assert drained.is_set() + assert (await run.state()).status is TaskStatus.RECOVERY_REQUIRED + await run.cancel(idempotency_key="cancel-siblings") + settled = await run.state() + assert settled.status is TaskStatus.RECOVERY_REQUIRED + assert next(n for n in settled.node_states if n.node_id == "n").status is TaskStatus.RECOVERY_REQUIRED + assert next(n for n in settled.node_states if n.node_id == "sibling").status is TaskStatus.CANCELLED + await run.resolve_effect("n", node.fence, TaskEffectResolution("failed", error_code=ErrorCode.TASK_NODE_FAILED), idempotency_key="accept") + assert (await run.state()).status is TaskStatus.FAILED diff --git a/tests/ai/test_transcript_sparse_windows.py b/tests/ai/test_transcript_sparse_windows.py index eaa77f589e..8f2a194c2f 100644 --- a/tests/ai/test_transcript_sparse_windows.py +++ b/tests/ai/test_transcript_sparse_windows.py @@ -82,18 +82,11 @@ async def remove_hint(transaction): await transaction.replace_record(replace(stored, data={}, storage_revision=stored.storage_revision + 1), expected_storage_revision=stored.storage_revision) await repository._store.mutate(remove_hint) - seeks = [] - original_seek = repository._seek_fact_sequence - async def track_seek(transaction, owner_id, index): - seeks.append(index) - return await original_seek(transaction, owner_id, index) - monkeypatch.setattr(repository, "_seek_fact_sequence", track_seek) windows = ((10, 12), (0, 0), (1, 3), (11, 13), (13, 13)) values = await repository.load_message_spans("history", windows) assert [[message.parts[0].content for message in span] for span in values] == [ ["10", "11"], [], ["1", "2"], ["11", "12"], [], ] - assert seeks == [1, 10] with pytest.raises(AIError) as error: await repository.load_message_spans("history", ((0, 1), (14, 14))) assert error.value.code is ErrorCode.STORAGE_INTEGRITY_ERROR diff --git a/tests/ai/test_upstream_runtime_measurement.py b/tests/ai/test_upstream_runtime_measurement.py index 8c85dda88b..3b130bc436 100644 --- a/tests/ai/test_upstream_runtime_measurement.py +++ b/tests/ai/test_upstream_runtime_measurement.py @@ -5,6 +5,9 @@ import asyncio import json import os +import re +from contextvars import ContextVar +from functools import wraps from collections import Counter from collections.abc import AsyncIterator from pathlib import Path @@ -15,11 +18,16 @@ from sqlalchemy.engine import Engine from linktools.ai.capability import AgentContext, ExecutionCapability +from linktools.ai.capability.steps import AtomicStepPersistence +from linktools.ai.persistence._runtime_commands import StateCommands +from linktools.ai.persistence._retention import RuntimeRetentionController from pydantic_ai_harness.step_persistence import StepPersistence from sqlalchemy.ext.asyncio import create_async_engine from linktools.ai.persistence import SqlStepStore from linktools.ai.migrate import provision_database from linktools.ai.core import CommitPolicy, ExecutionStatus +from linktools.ai.execution import ExecutionManager +from linktools.ai.runtime._conversation_service import ConversationService from linktools.ai.persistence import Persistence, PersistenceStorage from linktools.ai.persistence import _sql from linktools.ai.persistence.runs import RuntimeAgentRunStore @@ -57,6 +65,8 @@ async def test_two_turn_storage_measurement( assert Path(local_module.__file__).resolve().is_relative_to(root) assert Path(_sql.__file__).resolve().is_relative_to(root) payload = "x" * int(os.environ.get("TURN_IO_TOOL_BYTES", "0")) or "done" + tool_effect_policy = os.environ.get("TURN_IO_TOOL_POLICY", "replay_safe") + assert tool_effect_policy in {"none", "replay_safe"} counts: Counter[str] = Counter() active = False @@ -90,6 +100,53 @@ async def invoke(*args, **kwargs): for name in names: monkeypatch.setattr(cls, name, traced(getattr(cls, name), prefix, name)) + # Labels follow the real async owner, including its owned child tasks. Each + # SQL statement retains its own path: nested owners can share a transaction. + boundaries: ContextVar[tuple[str, ...]] = ContextVar("measurement_boundaries", default=()) + + def owned(original, label): + @wraps(original) + async def invoke(*args, **kwargs): + selected = label(*args, **kwargs) if callable(label) else label + token = boundaries.set((*boundaries.get(), selected)) + try: + return await original(*args, **kwargs) + finally: + boundaries.reset(token) + return invoke + + for cls, prefix, names in ( + (StateCommands, "execution", ( + "read_session_start", "reserve_start_checkpoint", "commit_start_checkpoint", + "commit_agent_attempt_checkpoint", "commit_prepared_agent_attempt_checkpoint", + "commit_start_attempt_checkpoint", "commit_tool_admission_boundary", + "commit_tool_terminal_boundary", "commit_tool_admission", "commit_tool_terminal", + "commit_terminal_checkpoint", "commit_pending_turn", "commit_cancel_checkpoint", + "commit_recovery_handoff", + )), + (ExecutionManager, "execution", ("load_result_authorized",)), + (RuntimeRetentionController, "retention", ("release_execution_handoff", "release_session")), + (local_module.LocalExecutionBackend, "runtime", ( + "_commit_success", "release_runtime_execution", "_release_session_execution", + )), + ): + for name in names: + monkeypatch.setattr(cls, name, owned(getattr(cls, name), f"{prefix}.{name}")) + + monkeypatch.setattr(ConversationService, "load_session_context", classmethod(owned( + ConversationService.load_session_context.__func__, "conversation.load_session_context", + ))) + + def native_boundary(self, **kwargs): + kinds = [value.kind for value in (*self._atomic_pending_events, *kwargs.get("events", ()))] + snapshot = kwargs.get("snapshot") + return "native.commit[" + ",".join(kinds or ["no_events"]) + ";snapshot=" + ( + "none" if snapshot is None else snapshot.state + ) + "]" + + monkeypatch.setattr(AtomicStepPersistence, "_atomic_commit", + owned(AtomicStepPersistence._atomic_commit, native_boundary)) + original_record = _sql._record_from_row def record_from_row(row): result = original_record(row) @@ -123,20 +180,56 @@ async def fresh(self, agent_run_id): return await original_require(self, agent_run_id) monkeypatch.setattr(RuntimeAgentRunStore, "require_active_producer", fresh) - mutations: dict[int, bool] = {} + transactions: list[dict[str, object]] = [] + open_transactions: dict[int, dict[str, object]] = {} + transaction_rounds = [] + table_pattern = re.compile(r'\b(?:FROM|JOIN|INTO|UPDATE)\s+["`\[]?(ai_[A-Za-z0-9_]+)', re.IGNORECASE) + + def new_transaction(connection, *, began_in_window): + transaction = {"id": len(transactions) + 1, "begin_boundary": list(boundaries.get()), + "began_in_window": began_in_window, "statements": [], "outcome": "open"} + transactions.append(transaction) + open_transactions[id(connection)] = transaction + return transaction + def sql(connection, cursor, statement, parameters, context, many): + if not active: + return word = statement.lstrip().split()[0].upper() count("sql." + word) count("sql.bind_text_blob_bytes", _bind_bytes(parameters)) - if word in {"INSERT", "UPDATE", "DELETE"}: - mutations[id(connection)] = True + transaction = open_transactions.get(id(connection)) + if transaction is None: + transaction = new_transaction(connection, began_in_window=False) + transaction["statements"].append({ + "verb": word, "tables": sorted(set(table_pattern.findall(statement))), + "access": "write" if word in {"INSERT", "UPDATE", "DELETE"} else "read" if word == "SELECT" else "control", + "boundary": list(boundaries.get()), "bound_text_blob_bytes": _bind_bytes(parameters), + }) + def begin(connection): - count("sql.transactions") - mutations[id(connection)] = False + if active: + count("sql.transactions") + new_transaction(connection, began_in_window=True) + + def finish(connection, outcome): + transaction = open_transactions.pop(id(connection), None) + if active: + count("sql." + ("commits" if outcome == "commit" else "rollbacks")) + if transaction is not None: + transaction["outcome"] = outcome + transaction["end_boundary"] = list(boundaries.get()) + transaction["access"] = "write" if any( + item["access"] == "write" for item in transaction["statements"] + ) else "read_only" + if active and outcome == "commit" and transaction["access"] == "write": + count("sql.write_commits") + def commit(connection): - count("sql.commits") - if mutations.pop(id(connection), False): - count("sql.write_commits") + finish(connection, "commit") + + def rollback(connection): + finish(connection, "rollback") calls: Counter[str] = Counter() effects: list[str] = [] @@ -160,7 +253,7 @@ async def model(messages, info): features.capability(ExecutionCapability(retain_history=True, commit=CommitPolicy.terminal() if terminal else CommitPolicy.sync()), scope="runtime") if checkpointed: features.capability(StepPersistence(store=steps, capture_frontier=os.environ.get("TURN_IO_FRONTIER", "1") == "1"), scope="runtime") - features.tool(work, effect_policy="replay_safe") + features.tool(work, effect_policy=tool_effect_policy) features.agent("default", model="default", allow_tools=("work",)) cache = InMemoryCacheBackend(max_bytes=16 * 1024 * 1024) storage = PersistenceStorage.sql(engine, cache=cache, coordination=cache if coordinated else None) @@ -173,6 +266,7 @@ async def table_counts(): event.listen(Engine, "before_cursor_execute", sql) event.listen(Engine, "begin", begin) event.listen(Engine, "commit", commit) + event.listen(Engine, "rollback", rollback) try: async with Persistence.open("turn-io", storage=storage) as persistence, Runtime.attach( persistence=persistence, features=model_features(_Models(_UsageFunctionModel(model)), features), @@ -183,6 +277,8 @@ async def table_counts(): backend = runtime._execution_service.runtime_backend() for turn in range(2): counts.clear() + transactions = [] + open_transactions.clear() active = True execution = await session.start(f"question-{turn}", idempotency_key=f"round-{turn}") result = await execution.wait(timeout_seconds=15) @@ -197,7 +293,18 @@ async def table_counts(): if owner is not None: await asyncio.wait_for(owner.closed.wait(), 10) active = False + assert sum(tx["outcome"] == "commit" for tx in transactions) == counts["sql.commits"] + assert sum(tx["outcome"] == "rollback" for tx in transactions) == counts["sql.rollbacks"] + assert sum(tx["outcome"] == "commit" and tx["access"] == "write" for tx in transactions) == counts["sql.write_commits"] rows.append(dict(counts)) + transaction_rounds.append({ + "turn": turn, "transactions": transactions, + "unattributed_statements": sum(not item["boundary"] for tx in transactions for item in tx["statements"]), + "unattributed_transactions": sum(not any(item["boundary"] for item in tx["statements"]) for tx in transactions), + "unattributed_write_statements": sum(item["access"] == "write" and not item["boundary"] + for tx in transactions for item in tx["statements"]), + "open_transactions": len(open_transactions), + }) table_rows_after = await table_counts() if inspect_tables else {} assert calls == {"question-0": 4, "question-1": 4} assert len(effects) == 6 @@ -217,12 +324,14 @@ async def table_counts(): event.remove(Engine, "before_cursor_execute", sql) event.remove(Engine, "begin", begin) event.remove(Engine, "commit", commit) + event.remove(Engine, "rollback", rollback) await steps.close() await engine.dispose() output = Path(os.environ.get("TURN_IO_OUTPUT", ".dev/turn-io-measurement.jsonl")) with output.open("a", encoding="utf-8") as stream: stream.write(json.dumps({"label": os.environ.get("TURN_IO_LABEL", "current"), "source_root": str(root), "tool_payload_bytes": len(payload), "checkpointed": checkpointed, "capture_frontier": os.environ.get("TURN_IO_FRONTIER", "1") == "1", - "coordinated": coordinated, "terminal": terminal, "rounds": rows, + "coordinated": coordinated, "terminal": terminal, "tool_effect_policy": tool_effect_policy, "rounds": rows, + "transaction_attribution": transaction_rounds, "table_row_delta": {name: table_rows_after[name] - table_rows_before.get(name, 0) for name in table_rows_after}, "assertions": {"tools": len(effects), "model_calls": sum(calls.values()), "history_messages": len(page.items), "history_tool_returns": 6,