From b3cbe54fd7d86a21365b6244dd70e581fd4915d1 Mon Sep 17 00:00:00 2001 From: James Dunnam <7660553+jimador@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:48:00 -0400 Subject: [PATCH 1/4] Persist extraction runs in the Drivine store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The terminal transition is one statement that takes the node's exclusive lock before reading status, backstopped by a uniqueness-constrained terminal-write node so two writers that both saw RUNNING cannot both commit — the constraint's sufficiency is measured by mutation, the lock's argued, and the docs say which is which. Header writes carry a version and take it under compare-and-set, and they write header fields only. recordInvocation is the only door onto invocation state, insert-or-compare on the invocation's own constrained key, so a header save can never create, update or delete an invocation row. Records are locked once terminal through that door — SUCCEEDED, FAILED and CANCELLED alike; an identical resend replays, and one differing in a single field, including a field cleared to absent, is rejected. Fingerprints come off the canonical length-prefixed encoding with a version tag per kind, so a serializer or bind-map ordering change leaves equivalent state equivalent, and a header digest can never collide with an invocation or terminal one. Validation precedes mutation: a rejected write leaves no orphan node even when the ambient transaction that caught it commits, and independent attempts written in one caller transaction each settle on their own merits. A writer recovering from Neo4j's deadlock detector retries in a transaction the failure never touched. Every read scopes to the tenant ahead of its limit, chain walks are bounded client-side with a seen set, and all runs of a root resolve in one seek off the denormalized ref. The contract suite from the previous slice runs unmodified against this store, fifty-nine for fifty-nine. Signed-off-by: James Dunnam <7660553+jimador@users.noreply.github.com> --- CHANGELOG.md | 70 ++ .../dice/storage/DrivineExtractionRunStore.kt | 981 ++++++++++++++++++ .../dice/storage/ExtractionRunRowMappers.kt | 626 +++++++++++ .../dice/storage/ExtractionRunSchema.kt | 131 +++ ...tractionRunStoreContractIntegrationTest.kt | 54 + ...rivineExtractionRunStoreIntegrationTest.kt | 789 ++++++++++++++ .../storage/ExtractionRunRowMappersTest.kt | 124 +++ .../embabel/dice/storage/TestApplication.kt | 9 + .../extraction/ExtractionRunFingerprint.kt | 35 + .../ExtractionRunFingerprintTest.kt | 62 ++ docs/design/extraction-runs.md | 159 ++- 11 files changed, 3037 insertions(+), 3 deletions(-) create mode 100644 dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt create mode 100644 dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt create mode 100644 dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunSchema.kt create mode 100644 dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt create mode 100644 dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt create mode 100644 dice-storage/src/test/kotlin/com/embabel/dice/storage/ExtractionRunRowMappersTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 6614123a..ced2fcaf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1624,3 +1624,73 @@ and the consumer PRs that deliver it). default; every existing call site, Kotlin or Java, keeps compiling and keeps its prior behavior short of the fixes themselves. Design note: [docs/design/extraction-runs.md](docs/design/extraction-runs.md). + +- **EXPERIMENTAL.** `DrivineExtractionRunStore` — the durable Neo4j implementation of + `ExtractionRunStore`, completing DICE #67's storage half. It writes three node labels: + `(:ExtractionRun)` keyed `(contextId, runId)` for the header, + `(:ExtractionRun)-[:RECORDED]->(:ExtractionRunInvocation)` keyed + `(contextId, runId, invocationIndex, attempt)` for one attempt at one planned model call, and + `(:ExtractionRun)-[:ENDED_BY]->(:ExtractionRunTerminalWrite)` keyed `(contextId, runId)` for the + fingerprint of the write that ended the run. Every key is tenant-qualified, and unlike + `DrivineDriftReportStore` the tenant needs no `ctx:`-prefixed stand-in: that store's scope is + nullable and a Cypher MERGE cannot key on a null, while a run's tenant never is, so the plain value + is already an injective key. + **Compare-and-set is one Cypher statement, and it holds across processes rather than only across + threads.** The statement takes an exclusive node lock with `SET n.casLock = $lockToken` before it + reads the run's status, so a second transaction blocks and then reads at read-committed — after the + first committed — and takes the no-op branch. The lock token is a fresh UUID on every call so the + write is always a real change rather than a no-op a database may optimize away before locking, and + no index carries `status` or the terminal fingerprint, so the planner cannot serve the post-lock + read from an index entry it read at MATCH time. Underneath that argument sits a fact: the terminal + write is a `CREATE` of the terminal-write node under a uniqueness constraint, so two writers that + both read a run as `RUNNING` cannot both commit. The loser's transaction rolls back whole, the + store catches the violation, re-reads the recorded fingerprint in a fresh transaction, and answers + replayed or conflict. The constraint's sufficiency is measured: removing the lock leaves every + race test green, and removing both produces six racing writers all reporting `APPLIED` and three + contradictory endings recorded for one run. The lock's rests on the documented isolation argument, + which is why the constraint exists. + **The fingerprint is stored verbatim and compared verbatim, never re-derived.** The node carries the + exact string `ExtractionRunTransition.fingerprint` computed, so a correct retry that happened after + another attempt was recorded still replays; a store deriving a digest from the stored run would + reject it. + **A header write cannot touch a child row.** Invocation records are their own nodes, so `save` has + no way to delete one — the contract's merge-don't-replace rule falls out of the graph model instead + of being implemented. One consequence: a durable store keeps identified rows rather than the order a + caller listed them in, so it returns attempts in plan order where the in-memory reference returns + the caller's order. `invocationsOf` is plan order in both. + **Every page scopes in the query ahead of its `LIMIT`**, excludes rows with no sort key (Neo4j sorts + null largest, so one would sort to the front of a `DESC` order, spend a slot, and then be dropped by + the mapper), and skips corrupt rows with a warning rather than failing the whole read. `runsOfRoot` + is one indexed lookup on the denormalized root. The chain walk is client-side and bounded to `limit` + keyed lookups in one read transaction, cycle-safe and tenant-scoped at every hop: a parent is a + property rather than a relationship, because a run can name a parent not yet stored, and this module + takes no APOC dependency. The store passes the whole `AbstractExtractionRunStoreContractTest` suite + alongside the in-memory reference, plus Drivine-specific integration tests for multi-writer + compare-and-set, cross-tenant fail-closed with identical run ids in two tenants, full-fidelity row + round-trip, corrupt-row skip, and the privacy contract asserted over the properties actually + written. Design note: + [docs/design/extraction-runs.md](docs/design/extraction-runs.md). + **Compatibility: additive, and hosts must declare new schema.** No existing class loses a member and + no behaviour changes; `DrivineExtractionRunStore`, `ExtractionRunSchema`, `ExtractionRunRowMapper` + and `ExtractionInvocationRowMapper` are new types in `com.embabel.dice.storage`. Nothing is + auto-configured yet, so a host opts in by declaring the store bean and a `SchemaCatalog` carrying + `ExtractionRunSchema.specs()`. That catalog is **eight new schema items**, and the three constraints + are required rather than advisory — a MERGE on a natural key is race-free only under one, and the + third is what makes the compare-and-set a schema fact: + - `UniquenessConstraintSpec("ExtractionRun", ["contextId", "runId"])` + - `UniquenessConstraintSpec("ExtractionRunInvocation", ["contextId", "runId", "invocationIndex", "attempt"])` + - `UniquenessConstraintSpec("ExtractionRunTerminalWrite", ["contextId", "runId"])` + - `RangeIndexSpec("ExtractionRun", "contextId")` — the tenant page; a composite index cannot stand + in, because Neo4j will not use one for a predicate on only its leading property + - `RangeIndexSpec("ExtractionRun", ["contextId", "rootRunId"])` — the whole-lineage read + - `RangeIndexSpec("ExtractionRun", ["contextId", "parentRunId"])` — one hop down the parent axis + - `RangeIndexSpec("ExtractionRun", ["contextId", "startedAtEpochSecond"])` — the paging sort key + - `RangeIndexSpec("ExtractionRunInvocation", ["contextId", "runId"])` — one run's attempts + + No stored data changes and no migration is required: no released DICE ever wrote these labels. One + new bound a host should know about — `ContextId` accepts any non-blank string and the tenant is the + leading property of every key here, so the store rejects a tenant id longer than 1024 characters on + the write path rather than letting Neo4j fail the write mid-extraction with an index-key-size error. + Reads are uncapped, because a read for a longer tenant matches nothing by construction. Every new + type carries `@ApiStatus.Experimental` and the shapes may still move while the remaining #67 slices + land. diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt new file mode 100644 index 00000000..76a89a8a --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt @@ -0,0 +1,981 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome +import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord +import com.embabel.dice.proposition.extraction.ExtractionRun +import com.embabel.dice.proposition.extraction.ExtractionRunConflictException +import com.embabel.dice.proposition.extraction.ExtractionRunKey +import com.embabel.dice.proposition.extraction.ExtractionRunNotFoundException +import com.embabel.dice.proposition.extraction.ExtractionRunRef +import com.embabel.dice.proposition.extraction.ExtractionRunStatus +import com.embabel.dice.proposition.extraction.ExtractionRunStore +import com.embabel.dice.proposition.extraction.ExtractionRunTransition +import com.embabel.dice.proposition.extraction.ExtractionRunTransitionOutcome +import com.embabel.dice.proposition.extraction.ExtractionRunTransitionResult +import org.drivine.manager.PersistenceManager +import org.drivine.query.QuerySpecification +import org.slf4j.LoggerFactory +import org.springframework.transaction.PlatformTransactionManager +import org.springframework.transaction.annotation.Propagation +import org.springframework.transaction.annotation.Transactional +import org.springframework.transaction.support.TransactionSynchronizationManager +import org.springframework.transaction.support.TransactionTemplate +import java.time.Clock +import java.time.Instant +import java.util.UUID + +/** + * Drivine / Neo4j implementation of [ExtractionRunStore]. + * + * ## The graph + * + * - `(:ExtractionRun {contextId, runId, ...})` — the run header, one node per run per tenant. + * - `(:ExtractionRun)-[:RECORDED]->(:ExtractionRunInvocation {contextId, runId, invocationIndex, attempt, ...})` + * — one child node per attempt at one planned model call. + * - `(:ExtractionRun)-[:ENDED_BY]->(:ExtractionRunTerminalWrite {contextId, runId, fingerprint, ...})` + * — at most one per run, ever, and the thing that makes that "at most one" true. + * + * Every key is tenant-qualified because a run id is host-minted and DICE never assumes it is + * globally unique. Unlike `DrivineDriftReportStore`, the tenant needs no `ctx:`-prefixed stand-in: + * that store's scope is nullable, and a Cypher MERGE cannot key on a null, so a global report needed + * a non-null encoding that no real context id could collide with. A run's tenant is never null, so + * the plain value is already an injective key. + * + * Every statement is parameterized; nothing caller-derived is interpolated into Cypher. The + * statements are assembled only from this class's own literals. + * + * See [ExtractionRunSchema] for the constraints and indexes this depends on. They are not tuning: + * a MERGE on a natural key is race-free only under a uniqueness constraint on that key. + * + * ## Compare-and-set: why a terminal write cannot be applied twice + * + * The whole of [transition] is one Cypher statement, so it is one transaction, and it works two + * ways at once — a lock that makes the race rare and a constraint that makes the wrong answer + * impossible. + * + * **The lock.** The statement's first act after matching the run is `SET n.casLock = $lockToken`, + * before it reads anything. A property write takes an exclusive lock on the node and holds it until + * commit, so a second transaction reaching that line blocks until the first commits. Neo4j reads at + * read-committed, and the read of `n.status` is downstream of the `SET` in the same statement, so + * the second transaction reads the status *after* it acquired the lock and therefore *after* the + * first transaction committed. It sees a terminal run and takes the no-op branch. This is the same + * write-lock-before-read idiom `DrivineMetamodelVersionStore` uses on its counter. + * + * Two details make that argument hold rather than nearly hold. The lock token is a fresh UUID on + * every call, so the write is always a real change and never a no-op a future Neo4j might optimize + * away before taking the lock. And no index carries `status` or the terminal fingerprint — see + * [ExtractionRunSchema] — so the planner cannot serve the post-lock read from an index entry it + * read at MATCH time, which is the one way a lock-then-read can quietly read stale. + * + * **The constraint.** The argument above is about Neo4j's behaviour, and behaviour is a thing to be + * wrong about. So the terminal write is also a `CREATE` of an `(:ExtractionRunTerminalWrite)` node + * on the run's own key, under a uniqueness constraint on `(contextId, runId)`. If two transactions + * ever did both read a run as `RUNNING` — a Neo4j build that locks differently, a cluster, a + * planner that reorders in a way this KDoc did not anticipate — they would both try to create that + * node and the database would refuse the second one at commit. The loser's whole transaction rolls + * back, header included, and this store catches the violation, re-reads the recorded fingerprint in + * a fresh transaction, and returns the replay-or-conflict answer the contract asks for. Exactly one + * terminal write per run is a schema fact, not an inference. + * + * That is the multi-process half of the guarantee. Threads in one JVM could be serialized by a + * monitor; two processes cannot be, and nothing in this class holds mutable state to serialize on. + * Both halves come entirely from the database. + * + * **The fingerprint is stored, never re-derived.** The node carries the exact string + * [ExtractionRunTransition.fingerprint] computed, and a repeated terminal write is decided by + * comparing its fingerprint against that string. Re-deriving a digest from the stored run would + * reject a correct retry whenever an attempt had been recorded in between, because the run it + * derived from would have changed while the terminal write did not. + * + * ## A header write never touches a child row + * + * Invocation records are their own nodes on their own key, and [save]'s Cypher names none of them: + * it reads and writes `(:ExtractionRun)` alone. [recordInvocation] is the only method that ever + * names an `(:ExtractionRunInvocation)` node, so it is the only one that can create, update, or + * lock one. Whatever `run.invocations` a caller hands [save] plays no part in what it accepts, + * rejects, or persists. + * + * ## Validation happens before any node is created + * + * A `MERGE` on a node pattern creates the node the moment it runs, whether or not the write that + * follows turns out to be valid — a naive `MERGE` immediately followed by a Kotlin check that + * throws leaves the created node behind if whatever catches that exception does not also abort the + * transaction. `SAVE_RUN` and `RECORD_INVOCATION` both avoid this: an `OPTIONAL MATCH` reads + * whatever is already there, a `CASE` decides in Cypher whether the write should happen, and the + * `MERGE` that can create a node sits inside a `FOREACH` gated on that decision. A save naming + * version `0` for a run already stored at a later version, or a [recordInvocation] call against a + * run that has already ended, stops before reaching any `MERGE` at all, so a caller that catches + * the resulting [ExtractionRunConflictException] without rolling back its own transaction finds + * the graph exactly as it was. + * + * ## Scope is applied in the query, ahead of the limit + * + * Each page puts its tenant in the MATCH pattern and its `LIMIT` after the `ORDER BY`. Reading a + * limited page and filtering it in Kotlin would apply the limit first, so a tenant whose neighbour + * owns the head of the index would report no runs while the store held plenty. This is + * `DrivineDriftReportStore`'s rule carried over, and the cross-backend suite pins it. + * + * Every read returns each run with its invocation records attached, pages included, because the + * in-memory reference does and the two are held to one suite. The cost is bounded — at most `limit` + * runs times the run model's cap of 1024 attempts — but it is a real cost on a wide page. + * + * @param persistenceManager Drivine's handle on the `neo` datasource. + * @param transactionManager used to own a transaction when no caller has one, so a lost + * compare-and-set race can be recovered in a transaction the race did not already poison. + * @param clock supplies the instant a terminal write is recorded at. Informational, and injectable + * so a test can pin it; nothing sorts or compares on it. + */ +@Transactional +open class DrivineExtractionRunStore( + private val persistenceManager: PersistenceManager, + transactionManager: PlatformTransactionManager, + private val clock: Clock = Clock.systemUTC(), +) : ExtractionRunStore { + + private val logger = LoggerFactory.getLogger(DrivineExtractionRunStore::class.java) + + private val txTemplate = TransactionTemplate(transactionManager) + + // ---- writes ---- + + /** + * Records a running run, inserting it or updating the one already there — in one statement, so + * the check against the stored run and the write that follows it cannot be interleaved by a + * concurrent terminal write, and the header's compare-and-set on + * [com.embabel.dice.proposition.extraction.ExtractionRun.version] cannot be interleaved by a + * concurrent header save either. + * + * The statement returns what it found so this method can name which rule was broken. Cypher + * decides whether to write and Kotlin decides what to say about it, and the conditions are the + * same condition, written twice: a save writes when the run is new and names version 0; when it + * is still running, agrees with the stored lineage and start time, and its header content + * already matches what is stored — a no-op that leaves the version untouched, whatever version + * the save named; or when it is still running, agrees with lineage and start time, its content + * genuinely differs, and it names the version currently stored. Every other running-and-agreeing + * case is a stale write and is rejected without touching the row — see this class's KDoc for why + * that rejection never leaves a node behind. `run.invocations` is not part of the statement at + * all: this method's Cypher has no clause that reads or writes an `(:ExtractionRunInvocation)`. + * + * Runs and retries under [ownedTransaction] — see that method for the two shapes. + */ + @Transactional(propagation = Propagation.SUPPORTS) + override fun save(run: ExtractionRun): ExtractionRun { + require(run.status == ExtractionRunStatus.RUNNING) { + "save records a running run; ${run.status} is terminal and belongs to transition()" + } + require(run.finishedAt == null) { + "a running run has not finished, so it carries no finishedAt" + } + val key = run.key() + ExtractionRunSchema.requireStorableTenant(key.contextId) + + val lineageKey = ExtractionRunRowMapper.lineageKeyOf(run.lineage) + val startedAt = run.startedAt.toString() + val header = ExtractionRunRowMapper.headerBindMap(run) + val headerFingerprint = ExtractionRunRowMapper.headerFingerprint(header) + + return ownedTransaction { + val row = singleRow( + SAVE_RUN, + keyBindings(key) + mapOf( + "lockToken" to newLockToken(), + "running" to RUNNING, + "lineageKey" to lineageKey, + "startedAt" to startedAt, + "header" to header, + "headerFingerprint" to headerFingerprint, + "version" to run.version, + ), + ) ?: throw IllegalStateException("SAVE_RUN returned no row for ${key.runRef.runId}") + + when (val priorStatus = row["priorStatus"]?.toString()) { + null -> require(run.version == 0L) { + "a run's first save must name version 0, the version a run nobody has saved " + + "yet carries; this one names ${run.version}" + } + + RUNNING -> { + if (row["priorLineageKey"]?.toString() != lineageKey) { + throw ExtractionRunConflictException( + key, + "stored lineage differs from the one being saved; lineage is fixed at insert", + ) + } + if (row["priorStartedAt"]?.toString() != startedAt) { + throw ExtractionRunConflictException( + key, + "stored startedAt differs from the one being saved; a run starts once", + ) + } + if (row["priorHeaderFingerprint"]?.toString() != headerFingerprint) { + val priorVersion = (row["priorVersion"] as? Number)?.toLong() + ?: throw IllegalStateException( + "run ${key.runRef.runId} in context ${key.contextId.value} is RUNNING " + + "with no stored version", + ) + if (run.version != priorVersion) { + throw ExtractionRunConflictException( + key, + "the header was read at version ${run.version}; the store is now at " + + "$priorVersion. Read the run again with findRun and rebuild this " + + "save on what it holds now", + ) + } + } + // else: content already matches what is stored, so the statement above left the + // row exactly as it was — a no-op, whatever version this save named. + } + + else -> throw ExtractionRunConflictException( + key, + "already ended as $priorStatus and cannot be re-opened by a save", + ) + } + + requireStoredRun(key) + } + } + + /** + * Records one attempt against a running run, on its own child node. + * + * One statement again, and the same shape as [save]: read the run's status and the attempt's + * prior outcome and fingerprint, write only if the run is still running and the attempt is not + * locked as terminal under a different payload — and, per this class's KDoc, only create the + * node at all once that decision comes back yes. A caller that holds only the attempt never has + * to hold the header, and this can never overwrite one. + * + * **Once the attempt is terminal, only an identical write reaches the row.** The incoming + * record's fingerprint — [ExtractionInvocationRowMapper.fingerprint] — is compared straight + * against the string already stored on the node, always computed fresh from the record a + * caller handed this method; the run's own terminal write follows the identical stored-string, + * compared-verbatim rule, one level down. + * + * Runs and retries under [ownedTransaction] — see that method for the two shapes. + */ + @Transactional(propagation = Propagation.SUPPORTS) + override fun recordInvocation( + key: ExtractionRunKey, + record: ExtractionInvocationRecord, + ): ExtractionRun { + ExtractionRunSchema.requireStorableTenant(key.contextId) + val fingerprint = ExtractionInvocationRowMapper.fingerprint(record) + + return ownedTransaction { + val row = singleRow( + RECORD_INVOCATION, + keyBindings(key) + mapOf( + "lockToken" to newLockToken(), + "running" to RUNNING, + "inFlight" to IN_FLIGHT, + "invocationIndex" to record.invocationIndex, + "attempt" to record.attempt, + "recordFingerprint" to fingerprint, + "record" to (ExtractionInvocationRowMapper.bindMap(record) + mapOf("recordFingerprint" to fingerprint)), + ), + ) ?: throw ExtractionRunNotFoundException(key) + + val priorStatus = row["priorStatus"]?.toString() + if (priorStatus != RUNNING) { + throw ExtractionRunConflictException( + key, + "already ended as $priorStatus; its invocation records are part of how it ended", + ) + } + if (row["locked"] == true) { + throw ExtractionRunConflictException( + key, + "${record.id} already ended as ${row["priorOutcome"]}; once an attempt is terminal " + + "only an identical write replays, and this one differs", + ) + } + requireStoredRun(key) + } + } + + /** + * Ends a run under compare-and-set. See this class's KDoc for why the answer is a schema fact + * rather than a claim about timing. + * + * Two transaction shapes, for the same reason `DrivinePropositionRepository.save` has two: + * - **A caller's transaction is active.** This joins it. A lost race has already ended that + * transaction by the time this method could react, so no recovery of ours could commit in it; + * the violation propagates and retrying is the caller's. Opening a nested transaction to sneak + * a read out would break the atomicity the caller asked for. + * - **No caller transaction.** `SUPPORTS` keeps the proxy from opening one, so the template owns + * the attempt, and the recovery read runs in a fresh transaction the failed attempt cannot + * roll back. + */ + @Transactional(propagation = Propagation.SUPPORTS) + override fun transition( + key: ExtractionRunKey, + transition: ExtractionRunTransition, + ): ExtractionRunTransitionResult { + ExtractionRunSchema.requireStorableTenant(key.contextId) + if (TransactionSynchronizationManager.isActualTransactionActive()) { + return endRun(key, transition) + } + + var lostRace: RuntimeException? = null + val result = try { + txTemplate.execute { endRun(key, transition) } + } catch (e: RuntimeException) { + if (!isUniquenessViolation(e)) throw e + lostRace = e + null + } + if (result != null) return result + + logger.debug( + "Terminal write for run {} in context {} lost the race; re-reading the recorded write", + key.runRef.runId, + key.contextId.value, + ) + return txTemplate.execute { resolveLostRace(key, transition, lostRace!!) }!! + } + + /** + * Runs the compare-and-set statement and turns what it found into the contract's answer. + * + * `priorStatus` is what the run was when this transaction took its lock. `RUNNING` means the + * statement applied the terminal fields and created the terminal-write node; anything else means + * it wrote nothing, and the decision is a comparison of the stored fingerprint against this one. + */ + private fun endRun( + key: ExtractionRunKey, + transition: ExtractionRunTransition, + ): ExtractionRunTransitionResult { + val row = singleRow( + END_RUN, + keyBindings(key) + mapOf( + "lockToken" to newLockToken(), + "running" to RUNNING, + "status" to transition.status.name, + // The string the transition computed, stored verbatim. Nothing here derives it. + "fingerprint" to transition.fingerprint, + "recordedAt" to clock.instant().toString(), + "terminal" to ExtractionRunRowMapper.terminalBindMap(transition), + ), + ) ?: throw ExtractionRunNotFoundException(key) + + val priorStatus = row["priorStatus"]?.toString() + if (priorStatus == RUNNING) { + logger.debug( + "Ended run {} in context {} as {}", + key.runRef.runId, + key.contextId.value, + transition.status, + ) + return ExtractionRunTransitionResult( + requireStoredRun(key), + ExtractionRunTransitionOutcome.APPLIED, + ) + } + return replayOrConflict(key, transition, row["priorFingerprint"]?.toString(), priorStatus) + } + + /** + * Decides what a terminal write against an already-ended run means, by comparing fingerprints. + * + * A missing recorded fingerprint is a conflict, not a replay: a terminal run this store cannot + * show a terminal write for is one it cannot prove agrees with the incoming one, and agreeing is + * the only thing that makes overwriting safe. + */ + private fun replayOrConflict( + key: ExtractionRunKey, + transition: ExtractionRunTransition, + recordedFingerprint: String?, + priorStatus: String?, + ): ExtractionRunTransitionResult { + if (recordedFingerprint != null && recordedFingerprint == transition.fingerprint) { + return ExtractionRunTransitionResult( + requireStoredRun(key), + ExtractionRunTransitionOutcome.REPLAYED, + ) + } + throw ExtractionRunConflictException( + key, + "already ended as $priorStatus under a different terminal write; " + + "this one claims ${transition.status}", + ) + } + + /** + * Recovers from losing the create of the terminal-write node: read what the winner recorded and + * decide against it. + * + * If there is no terminal write to find, the violation was not the one this method knows how to + * interpret, so the original exception is rethrown rather than guessed about. + */ + private fun resolveLostRace( + key: ExtractionRunKey, + transition: ExtractionRunTransition, + violation: RuntimeException, + ): ExtractionRunTransitionResult { + val row = singleRow(TERMINAL_WRITE_BY_KEY, keyBindings(key)) ?: throw violation + return replayOrConflict( + key, + transition, + row["fingerprint"]?.toString(), + row["status"]?.toString(), + ) + } + + // ---- reads ---- + + @Transactional(readOnly = true) + override fun findRun(key: ExtractionRunKey): ExtractionRun? = + singleRow(RUN_BY_KEY, keyBindings(key))?.let(::mapRun) + + @Transactional(readOnly = true) + override fun invocationsOf(key: ExtractionRunKey): List = + queryRows(INVOCATIONS_OF, keyBindings(key)).mapNotNull(::mapInvocation) + + @Transactional(readOnly = true) + override fun runsInContext( + contextIdValue: String, + limit: Int, + since: Instant?, + ): List = readPage( + scope = null, + contextIdValue = contextIdValue, + limit = limit, + since = since, + ) + + @Transactional(readOnly = true) + override fun childrenOf( + contextIdValue: String, + parentRunId: String, + limit: Int, + ): List = readPage( + scope = ONLY_PARENT, + contextIdValue = contextIdValue, + limit = limit, + since = null, + extraBindings = mapOf("parentRunId" to parentRunId), + ) + + @Transactional(readOnly = true) + override fun runsOfRoot( + contextIdValue: String, + rootRunId: String, + limit: Int, + since: Instant?, + ): List = readPage( + scope = ONLY_ROOT, + contextIdValue = contextIdValue, + limit = limit, + since = since, + extraBindings = mapOf("rootRunId" to rootRunId), + ) + + /** + * Walks the parent chain upward, one keyed lookup per hop, at most [limit] hops. + * + * **Why the walk is here and not in Cypher.** A parent is a property, not a relationship: a run + * can name a parent that has not been stored yet, and an edge cannot point at a node that does + * not exist. Materializing the edge later would mean a second write nothing triggers. Without an + * edge there is no variable-length pattern to walk, and the APOC procedures that would do it in + * one round trip are not a dependency this module takes. + * + * So the walk is client-side and bounded by construction: at most `limit` hops, each a lookup on + * the uniqueness-constraint index, all inside one read transaction. It stops on a run it has + * already seen, which is what makes it safe on a corrupt store holding a cycle — [ExtractionRunLineage] + * can reject a run that is its own parent, but a two-hop cycle needs the other runs to see. And + * it resolves every hop inside the starting run's tenant, so a parent id that exists only in a + * neighbour's tenant resolves to nothing and the walk ends. + */ + @Transactional(readOnly = true) + override fun ancestorsOf(key: ExtractionRunKey, limit: Int): List { + requirePositiveLimit(limit) + val start = findRun(key) ?: return emptyList() + val walked = mutableListOf() + val seen = mutableSetOf(start.ref) + var parentRef: ExtractionRunRef? = start.parentRef + while (parentRef != null && walked.size < limit && seen.add(parentRef)) { + val parent = findRun(ExtractionRunKey(key.contextId, parentRef)) ?: break + walked += parent + parentRef = parent.parentRef + } + return walked + } + + /** + * Assembles and runs one of the three pages. + * + * The statement is built from this class's own literals and nothing else; every value travels as + * a bound parameter. `scope` narrows inside the `WHERE`, ahead of the `LIMIT`, which is the whole + * point. + * + * A row that will not map has already spent one of the caller's `limit` slots, so a page can come + * back shorter than asked for. Reading further to backfill would break the bound the contract + * keeps. + */ + private fun readPage( + scope: String?, + contextIdValue: String, + limit: Int, + since: Instant?, + extraBindings: Map = emptyMap(), + ): List { + requirePositiveLimit(limit) + val statement = buildString { + append(PAGE_MATCH) + scope?.let { append("\n").append(it) } + if (since != null) append("\n").append(SINCE_BOUND) + append("\n").append(NEWEST_FIRST_PAGE) + } + val bindings = buildMap { + put("contextId", contextIdValue) + put("limit", limit) + putAll(extraBindings) + if (since != null) { + put("sinceEpochSecond", since.epochSecond) + put("sinceNano", since.nano) + } + } + return queryRows(statement, bindings).mapNotNull(::mapRun) + } + + /** + * Turns one `{run, invocations}` row into a run, or logs it and returns null. + * + * One corrupt node should not fail a whole audit read, and the row mappers throw rather than + * inventing defaults so that this can happen. The warning names the property that was missing or + * the check that failed, which is what an operator needs to go find the node. A corrupt child row + * is dropped on its own, so a run with one unreadable attempt still reads. + */ + private fun mapRun(row: Map<*, *>): ExtractionRun? { + val header = row["run"] as? Map<*, *> ?: run { + logger.warn("Skipping ExtractionRun row with no header properties") + return null + } + val invocations = (row["invocations"] as? List<*>).orEmpty() + .filterIsInstance>() + .mapNotNull(::mapInvocation) + return runCatching { ExtractionRunRowMapper.fromRow(header, invocations) } + .onFailure { logger.warn("Skipping unreadable ExtractionRun row: {}", it.message) } + .getOrNull() + } + + private fun mapInvocation(row: Map<*, *>): ExtractionInvocationRecord? = + runCatching { ExtractionInvocationRowMapper.fromRow(row) } + .onFailure { logger.warn("Skipping unreadable ExtractionRunInvocation row: {}", it.message) } + .getOrNull() + + private fun requireStoredRun(key: ExtractionRunKey): ExtractionRun = + findRun(key) ?: throw IllegalStateException( + "run ${key.runRef.runId} in context ${key.contextId.value} was written but does not read back", + ) + + // ---- plumbing ---- + + private fun keyBindings(key: ExtractionRunKey): Map = mapOf( + "contextId" to key.contextId.value, + "runId" to key.runRef.runId, + ) + + private fun execute(statement: String, bindings: Map) { + persistenceManager.execute(QuerySpecification.withStatement(statement).bind(bindings)) + } + + private fun queryRows(statement: String, bindings: Map): List> { + @Suppress("UNCHECKED_CAST") + val spec = QuerySpecification.withStatement(statement).bind(bindings) as QuerySpecification + return persistenceManager.query(spec).filterIsInstance>() + } + + private fun singleRow(statement: String, bindings: Map): Map<*, *>? = + queryRows(statement, bindings).firstOrNull() + + private fun requirePositiveLimit(limit: Int) { + require(limit > 0) { "limit must be positive, was $limit" } + } + + /** + * A fresh token for the lock write, on every call. + * + * The point is that the write is always a real change. A `SET` of a property to the value it + * already holds is a no-op a database is free to optimize away, and an optimized-away write is + * one that took no lock — which would quietly turn the compare-and-set into a read-then-write. + */ + private fun newLockToken(): String = UUID.randomUUID().toString() + + /** + * Runs [block] — [save] or [recordInvocation]'s whole write-then-read — under a transaction + * boundary this store owns whenever it can, and retries wholesale on Neo4j's own deadlock + * detector, giving every attempt its own fresh transaction. + * + * **A caller's transaction is active.** This joins it, once, and does not retry — the same + * reasoning [transition] documents for a lost compare-and-set race. A deadlock abort inside a + * shared transaction has already ended that transaction by the time this method could react, so + * a second attempt run inside it would be a statement rerun against a resource the database has + * already discarded — a retry needs a boundary it owns, and only the caller holds this one, so + * the retry is theirs to make. The recovery is theirs too, for the same reason. + * + * **No caller transaction.** [txTemplate] opens a genuinely new transaction for every attempt, + * untouched by whatever a previous attempt did to the one before it; each retry reads whatever + * the database actually holds at that moment, independent of what this store's own last attempt + * found before it aborted. + * + * Two writers `MERGE`-creating the same not-yet-existing key — a run's first save, or an + * invocation's first write — can each be granted half of the unique index insert before either + * commits, which Neo4j resolves by aborting one transaction outright to avoid blocking forever. + * That abort is transient: [block] carries no state of its own between attempts, so running it + * again is safe, and is what Neo4j's own documentation recommends for this exact conflict. + * [isTransientConflict] limits retrying to that one failure shape; anything else — including a + * genuine [ExtractionRunConflictException] this store raised on purpose — propagates on the + * first attempt, from whichever shape ran it. + */ + private fun ownedTransaction(attempts: Int = 5, block: () -> T): T { + if (TransactionSynchronizationManager.isActualTransactionActive()) { + return block() + } + var lastConflict: RuntimeException? = null + repeat(attempts) { + try { + return txTemplate.execute { block() }!! + } catch (e: RuntimeException) { + if (!isTransientConflict(e)) throw e + lastConflict = e + } + } + throw lastConflict!! + } + + private companion object { + + private val RUNNING: String = ExtractionRunStatus.RUNNING.name + + /** An invocation record's not-yet-terminal outcome — the one outcome a later write may still + * replace in place, on either write door. */ + private val IN_FLIGHT: String = ExtractionInvocationOutcome.IN_FLIGHT.name + + /** + * Insert-or-update a running header, deciding inside the statement whether to write and, + * when it does, what version to leave the row at — and never creating the node at all when + * the decision comes back no. + * + * `OPTIONAL MATCH` takes the read, so a run that does not yet exist reaches the rest of the + * statement as `n IS NULL`, with no node created on its behalf. The first `FOREACH` then + * takes the exclusive lock — `SET n.casLock`, before any of the properties + * below are read — but only when `n` exists; a node that is not there yet has nothing to + * lock, and there is no concurrent reader of a row nobody has written. Without the lock a + * concurrent write could land between the read below and the write that follows it and be + * silently overwritten — the `MERGE … SET` failure this store exists to avoid. Everything + * this statement decides, it decides in this one round trip: a lock taken here and a write + * issued from a later, separate statement would leave a window between them for another + * writer to move the row, which is exactly the race this design avoids by never splitting + * the two. + * + * `shouldWrite` is the whole compare-and-set decision, over a `CASE` written once: new + * (naming version 0); or running, agreeing on lineage and start time, and either the + * header's content already matches what is stored — a no-op — or it genuinely differs and + * `$version` names the version currently stored, the accepted case that raises it by one. + * The caller repeats the same condition in Kotlin, over the same properties this statement + * returns, to say which half failed when it did not write. + * + * The two `FOREACH` clauses that can write are mutually exclusive on `priorStatus`, and each + * is gated on `shouldWrite` besides: the first only ever runs when there was no row to find, + * and only creates one when the insert itself is valid, so a first save naming the wrong + * version leaves no row for a caller catching the resulting exception to find. The second + * only ever runs against a row this statement already matched, so it can update but it can + * never originate one. Neither `FOREACH` ever names an `(:ExtractionRunInvocation)`. + * + * `priorStartedAt` is compared as the ISO string, which is what round-trips: `Instant.parse` + * inverts `Instant.toString` exactly, so two equal instants compare equal here. + */ + private val SAVE_RUN = """ + OPTIONAL MATCH (n:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + FOREACH (ignored IN CASE WHEN n IS NOT NULL THEN [1] ELSE [] END | + SET n.casLock = ${'$'}lockToken + ) + WITH n, + n.status AS priorStatus, + n.lineageKey AS priorLineageKey, + n.startedAt AS priorStartedAt, + n.version AS priorVersion, + n.headerFingerprint AS priorHeaderFingerprint + WITH n, priorStatus, priorLineageKey, priorStartedAt, priorVersion, priorHeaderFingerprint, + CASE + WHEN priorStatus IS NULL THEN ${'$'}version = 0 + WHEN priorStatus <> ${'$'}running THEN false + WHEN priorLineageKey <> ${'$'}lineageKey THEN false + WHEN priorStartedAt <> ${'$'}startedAt THEN false + WHEN priorHeaderFingerprint = ${'$'}headerFingerprint THEN false + ELSE priorVersion = ${'$'}version + END AS shouldWrite + FOREACH (ignored IN CASE WHEN shouldWrite AND priorStatus IS NULL THEN [1] ELSE [] END | + MERGE (created:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + SET created += ${'$'}header, + created.version = 0, + created.headerFingerprint = ${'$'}headerFingerprint, + created.casLock = ${'$'}lockToken + ) + FOREACH (ignored IN CASE WHEN shouldWrite AND priorStatus IS NOT NULL THEN [1] ELSE [] END | + SET n += ${'$'}header, + n.version = priorVersion + 1, + n.headerFingerprint = ${'$'}headerFingerprint + ) + RETURN { + priorStatus: priorStatus, + priorLineageKey: priorLineageKey, + priorStartedAt: priorStartedAt, + priorVersion: priorVersion, + priorHeaderFingerprint: priorHeaderFingerprint + } AS row + """.trimIndent() + + /** + * Record one attempt directly, while the run is still running and while the attempt is free + * of a terminal lock under a different payload — and, the same as [SAVE_RUN], creating the + * child node only once that decision comes back yes. + * + * `MATCH` on the header requires it to exist: a run that does not exist yields no row at + * all, so the statement writes nothing and the caller raises not-found. `OPTIONAL MATCH` on + * the invocation takes its read without creating it, the same reason [SAVE_RUN] reads the + * header that way — a locked write, or a write against a run that has already ended, has to + * reach no `MERGE` at all, so the rejection it raises finds the graph exactly as this + * statement found it. + * `locked` follows the rule [recordInvocation] documents: an id already stored with a + * terminal `outcome` accepts only a write whose `recordFingerprint` matches what is stored, + * and every other write for that id is refused, whether it claims a different outcome or the + * same one under different facts. + */ + private val RECORD_INVOCATION = """ + MATCH (n:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + SET n.casLock = ${'$'}lockToken + WITH n, n.status AS priorStatus + OPTIONAL MATCH (i:ExtractionRunInvocation { + contextId: ${'$'}contextId, + runId: ${'$'}runId, + invocationIndex: ${'$'}invocationIndex, + attempt: ${'$'}attempt + }) + WITH n, priorStatus, i.outcome AS priorOutcome, i.recordFingerprint AS priorFingerprint + WITH n, priorStatus, priorOutcome, + (priorOutcome IS NOT NULL + AND priorOutcome <> ${'$'}inFlight + AND priorFingerprint <> ${'$'}recordFingerprint) AS locked + FOREACH (ignored IN CASE WHEN priorStatus = ${'$'}running AND locked = false THEN [1] ELSE [] END | + MERGE (created:ExtractionRunInvocation { + contextId: ${'$'}contextId, + runId: ${'$'}runId, + invocationIndex: ${'$'}invocationIndex, + attempt: ${'$'}attempt + }) + SET created += ${'$'}record + MERGE (n)-[:RECORDED]->(created) + ) + RETURN {priorStatus: priorStatus, priorOutcome: priorOutcome, locked: locked} AS row + """.trimIndent() + + /** + * The compare-and-set. One statement, one transaction, and the only writer of a terminal + * status. + * + * Reading in order: + * 1. `MATCH` — no row means no such run in this tenant, and the caller raises not-found. + * 2. `SET n.casLock` — takes the exclusive node lock, before any read. See the class KDoc. + * 3. `OPTIONAL MATCH … ENDED_BY` — the terminal write already recorded, if there is one. + * It runs before the `FOREACH`, so it never sees the node this statement is about to + * create. + * 4. `FOREACH` — applies the terminal fields and creates the terminal-write node, but only + * when the run was still running when the lock was taken. + * 5. `RETURN` — the status and fingerprint as they were, which is everything the caller needs + * to answer applied, replayed, or conflict. + * + * `SET n += $terminal` carries counts and failures only when the transition replaces them. A + * transition that keeps them binds no such keys, so the stored values stand — which is + * exactly what `applyTo`'s `counts ?: run.counts` does. Binding null instead would remove the + * properties. + * + * The `CREATE` is the backstop the class KDoc describes: under the uniqueness constraint on + * `ExtractionRunTerminalWrite(contextId, runId)`, a second terminal write cannot commit even + * if it somehow read the run as running. + */ + private val END_RUN = """ + MATCH (n:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + SET n.casLock = ${'$'}lockToken + WITH n + OPTIONAL MATCH (n)-[:ENDED_BY]->(prior:ExtractionRunTerminalWrite) + WITH n, n.status AS priorStatus, prior.fingerprint AS priorFingerprint + FOREACH (ignored IN CASE WHEN priorStatus = ${'$'}running THEN [1] ELSE [] END | + CREATE (n)-[:ENDED_BY]->(:ExtractionRunTerminalWrite { + contextId: ${'$'}contextId, + runId: ${'$'}runId, + fingerprint: ${'$'}fingerprint, + status: ${'$'}status, + recordedAt: ${'$'}recordedAt + }) + SET n += ${'$'}terminal + ) + RETURN {priorStatus: priorStatus, priorFingerprint: priorFingerprint} AS row + """.trimIndent() + + /** What the winner of a race recorded, read after this transaction's attempt rolled back. */ + private val TERMINAL_WRITE_BY_KEY = """ + MATCH (t:ExtractionRunTerminalWrite {contextId: ${'$'}contextId, runId: ${'$'}runId}) + RETURN {fingerprint: t.fingerprint, status: t.status} AS row + """.trimIndent() + + /** + * One run and its attempts, in one round trip. + * + * `collect` drops the nulls an `OPTIONAL MATCH` with no match produces, so a run with no + * attempts comes back with an empty list rather than a list holding a null. The `ORDER BY` + * ahead of it is what the collected list inherits. + * + * **The attempts come back in plan order**, which is the order `invocationsOf` promises. A + * durable store keeps identified rows, not the order a caller happened to list them in, so + * plan order is the only order it can offer — and it is the one the run model defines. + */ + private val RUN_BY_KEY = """ + MATCH (n:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + OPTIONAL MATCH (n)-[:RECORDED]->(i:ExtractionRunInvocation) + WITH n, i ORDER BY i.invocationIndex ASC, i.attempt ASC + WITH n, collect(properties(i)) AS invocations + RETURN {run: properties(n), invocations: invocations} AS row + """.trimIndent() + + /** + * One run's attempts in plan order: call 0 before call 1, and within a call, first attempt + * before second. + * + * The order is the plan's, not the order the calls came back in, which is why it sorts on the + * identity that was allocated up front. Sorting in the database rather than in Kotlin keeps + * one definition of plan order for this backend. + */ + private val INVOCATIONS_OF = """ + MATCH (:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + -[:RECORDED]->(i:ExtractionRunInvocation) + WITH i ORDER BY i.invocationIndex ASC, i.attempt ASC + RETURN properties(i) AS row + """.trimIndent() + + /** + * Every page starts here, with the tenant in the pattern. + * + * `startedAtEpochSecond IS NOT NULL` is load-bearing. Neo4j sorts null as the largest value, + * so a node missing the sort key would sort to the front of a `DESC` order, spend a slot of + * the caller's `limit`, and then be dropped by the mapper — hiding a good run behind a broken + * one. Excluding it in the database keeps it out of the order entirely. + */ + private val PAGE_MATCH = """ + MATCH (n:ExtractionRun {contextId: ${'$'}contextId}) + WHERE n.startedAtEpochSecond IS NOT NULL + """.trimIndent() + + /** One lineage, from the denormalized root. Supersession is a separate axis and is not walked. */ + private val ONLY_ROOT = "AND n.rootRunId = ${'$'}rootRunId" + + /** One hop down the parent axis. A run with no parent has no such property and cannot match. */ + private val ONLY_PARENT = "AND n.parentRunId = ${'$'}parentRunId" + + /** + * The `since` bound, inclusive, compared second-then-nanosecond so it is exact. + * + * A single truncated millisecond would sweep in runs started just before a bound that falls + * part-way through one. + */ + private val SINCE_BOUND = """ + AND (n.startedAtEpochSecond > ${'$'}sinceEpochSecond + OR (n.startedAtEpochSecond = ${'$'}sinceEpochSecond + AND n.startedAtNano >= ${'$'}sinceNano)) + """.trimIndent() + + /** + * Newest first by start instant, tie-broken by run id ascending, then limited — in that + * order, which is the order that makes a page repeatable and correctly scoped. + * + * The run order is written twice on purpose. The first `ORDER BY` is the one that matters: + * it decides which rows the `LIMIT` keeps. The `collect` after it groups by run and does not + * promise to preserve the incoming order, so the last `ORDER BY` is what the caller actually + * receives. Dropping either one leaves a page that is right about the wrong thing. + * + * The middle `ORDER BY` is the one the collected attempts inherit, so each run's attempts + * come back in plan order, as they do from [RUN_BY_KEY] and `invocationsOf`. + */ + private val NEWEST_FIRST_PAGE = """ + WITH n ORDER BY n.startedAtEpochSecond DESC, n.startedAtNano DESC, n.runId ASC + LIMIT ${'$'}limit + OPTIONAL MATCH (n)-[:RECORDED]->(i:ExtractionRunInvocation) + WITH n, i ORDER BY i.invocationIndex ASC, i.attempt ASC + WITH n, collect(properties(i)) AS invocations + ORDER BY n.startedAtEpochSecond DESC, n.startedAtNano DESC, n.runId ASC + RETURN {run: properties(n), invocations: invocations} AS row + """.trimIndent() + } +} + +/** + * Best-effort detection of a Neo4j uniqueness-constraint violation anywhere in the cause chain. + * + * Matches on message substrings, because which form — the error code or the prose — reaches + * `getMessage()` is not guaranteed across driver versions. `DrivinePropositionRepository` carries the + * same check for the same reason; the two are not shared yet because neither module has a home for a + * Drivine error-mapping helper, and inventing one is a separate change. + * + * A false positive is bounded here: the caller only ever treats it as "someone else recorded the + * terminal write", and re-reads to find out. If there is no terminal write, the original exception is + * rethrown. + */ +private fun isUniquenessViolation(error: Throwable?): Boolean { + var current: Throwable? = error + val seen = mutableSetOf() + while (current != null && seen.add(current)) { + val message = current.message ?: "" + if (message.contains("ConstraintValidationFailed", ignoreCase = true) || + message.contains("already exists", ignoreCase = true) + ) { + return true + } + current = current.cause + } + return false +} + +/** + * Best-effort detection of Neo4j's own deadlock detector aborting a transaction, anywhere in the + * cause chain — the failure [DrivineExtractionRunStore.retryingTransientConflict] retries. + * + * Matches on message substrings for the same reason [isUniquenessViolation] does: which form of the + * message reaches `getMessage()` is not guaranteed across driver versions. Two writers `MERGE`-ing + * the same not-yet-existing key concurrently is the shape [DrivineExtractionRunStore.ownedTransaction] + * retries, and Neo4j names it + * with the transaction lock manager's own vocabulary, a looser target than one stable exception + * type: `TransientException` is the driver's own class name, and `Deadlock` and `can't acquire` are + * the wording its lock manager uses to describe the same event. The check is deliberately loose to + * match all three. + */ +private fun isTransientConflict(error: Throwable?): Boolean { + var current: Throwable? = error + val seen = mutableSetOf() + while (current != null && seen.add(current)) { + val message = current.message ?: "" + if (message.contains("TransientException", ignoreCase = true) || + message.contains("Deadlock", ignoreCase = true) || + message.contains("can't acquire", ignoreCase = true) + ) { + return true + } + current = current.cause + } + return false +} diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt new file mode 100644 index 00000000..bf6027b3 --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt @@ -0,0 +1,626 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.agent.core.ContextId +import com.embabel.dice.proposition.extraction.ExtractionActorRef +import com.embabel.dice.proposition.extraction.ExtractionCohortRef +import com.embabel.dice.proposition.extraction.ExtractionContentProfileRef +import com.embabel.dice.proposition.extraction.ExtractionDeploymentRef +import com.embabel.dice.proposition.extraction.ExtractionExperimentRef +import com.embabel.dice.proposition.extraction.ExtractionFailure +import com.embabel.dice.proposition.extraction.ExtractionFailureCode +import com.embabel.dice.proposition.extraction.ExtractionInvocationId +import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome +import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord +import com.embabel.dice.proposition.extraction.ExtractionModelUsage +import com.embabel.dice.proposition.extraction.ExtractionPersonalizationRef +import com.embabel.dice.proposition.extraction.ExtractionProviderResponseFacts +import com.embabel.dice.proposition.extraction.ExtractionReplayFidelity +import com.embabel.dice.proposition.extraction.ExtractionRequestRef +import com.embabel.dice.proposition.extraction.ExtractionRequestedModelConfig +import com.embabel.dice.proposition.extraction.ExtractionRun +import com.embabel.dice.proposition.extraction.ExtractionRunCounts +import com.embabel.dice.proposition.extraction.ExtractionRunFingerprint +import com.embabel.dice.proposition.extraction.ExtractionRunFingerprints +import com.embabel.dice.proposition.extraction.ExtractionRunLineage +import com.embabel.dice.proposition.extraction.ExtractionRunRef +import com.embabel.dice.proposition.extraction.ExtractionRunStatus +import com.embabel.dice.proposition.extraction.ExtractionRunSubjectRefs +import com.embabel.dice.proposition.extraction.ExtractionRunTransition +import com.embabel.dice.proposition.extraction.ExtractionRuntimeIdentity +import com.embabel.dice.proposition.extraction.ExtractionSessionRef +import com.embabel.dice.provenance.SourceRevisionRef +import com.fasterxml.jackson.databind.ObjectMapper +import java.time.Duration +import java.time.Instant + +private val objectMapper = ObjectMapper() + +/** + * Translates extraction runs to and from the property maps the Neo4j graph store reads and writes. + * + * Neo4j properties are scalars and flat arrays, and a run header carries lists of value objects, so + * the structured parts — source revisions, the requested model configuration, the failure list — + * are serialized to JSON strings. JSON also survives the pipes, tabs, newlines and quotes that turn + * up in strings coming out of LLM extraction, which a delimiter-joined encoding does not. + * + * **Invocation records are not here.** They are child rows with their own key and their own mapper, + * [ExtractionInvocationRowMapper], and a header write never touches one. That is what makes a save + * unable to delete a recorded attempt, and it is why the store gets the contract's + * invocation-preservation rule for free rather than having to implement it. + * + * **Instants are written three ways.** The ISO-8601 string is what round-trips and what a person + * reading a node wants; `Instant.parse` inverts `Instant.toString` exactly. The epoch second and the + * nanosecond let the database sort and range-filter at full precision. Epoch milliseconds would + * truncate — two runs started 500 microseconds apart would compare equal, leaving "newest first" + * arbitrary between them, and a `since` bound falling inside a millisecond would sweep in runs + * started just before it. Sorting on the ISO string has its own failure: `Instant.toString` writes + * no fraction on a whole second and `'Z'` outranks `'.'`, so `12:00:00Z` sorts after `12:00:00.500Z`. + * + * **Reads are strict.** A property this mapper wrote must be there when it is read again. A node + * missing one is corrupt, so the accessor throws and the store's surrounding guard logs the row and + * skips it. Optional fields are the exception, and their absence means the run declared none: a + * `SET` of `null` in Cypher leaves no property behind, so "no profile" and "no experiment label" are + * stored as the absence of a property rather than as a sentinel. + * + * **The terminal fingerprint is never derived here.** It is the string + * [ExtractionRunTransition.fingerprint] computed, stored verbatim on its own node and compared + * verbatim. Re-deriving it from a stored run would make a correct retry that happened after another + * attempt was recorded look like an incompatible rewrite. + */ +object ExtractionRunRowMapper { + + /** + * The one derived property on a run header: an injective encoding of everything about a run's + * lineage that a later save could disagree about. + * + * A save is rejected when it contradicts the stored lineage, and that check has to happen inside + * the same statement that writes — otherwise a concurrent terminal write lands between the read + * and the write. Comparing five properties in Cypher works but reads badly and gets the null + * handling wrong easily; comparing one string is exact. The individual properties are still + * stored separately, because that is what `childrenOf` and `runsOfRoot` seek on. + * + * Two sources of truth for one fact is a drift hazard, so [fromRow] re-derives this from the + * properties it just read and refuses a node where the two disagree — the same shape of check + * `DriftReportRowMapper` runs over a report's scope. + */ + fun lineageKeyOf(lineage: ExtractionRunLineage): String = objectMapper.writeValueAsString( + linkedMapOf( + "runId" to lineage.runRef.runId, + "rootRunId" to lineage.rootRunRef.runId, + "parentRunId" to lineage.parentRunRef?.runId, + "supersedesRunId" to lineage.supersedesRunRef?.runId, + "passIndex" to lineage.passIndex, + ), + ) + + /** + * Everything a header write sets, minus the two key properties, which the MERGE pattern owns and + * nothing may move. + * + * Null values are deliberate. `SET n += $header` removes a property bound to null, so an + * optional field the run does not carry leaves no property behind, and a run saved again without + * it is stored the way a run that never had it is stored. + */ + fun headerBindMap(run: ExtractionRun): Map = buildMap { + put("status", run.status.name) + put("lineageKey", lineageKeyOf(run.lineage)) + put("rootRunId", run.rootRef.runId) + put("parentRunId", run.parentRef?.runId) + put("supersedesRunId", run.lineage.supersedesRunRef?.runId) + put("passIndex", run.lineage.passIndex) + + putInstant("startedAt", run.startedAt) + putInstant("finishedAt", run.finishedAt) + + put("profileName", run.profile?.name) + put("profileVersion", run.profile?.version) + put("sourceRevisions", serializeSourceRevisions(run.sourceRevisions)) + + put("promptTemplateFingerprint", run.fingerprints.promptTemplateFingerprint) + put("schemaFingerprint", run.fingerprints.schemaFingerprint) + put("metamodelFingerprint", run.fingerprints.metamodelFingerprint) + + put("extractor", run.runtime.extractor) + put("extractorVersion", run.runtime.extractorVersion) + put("hostApplication", run.runtime.hostApplication) + // `runtime` on the node would read as "the runtime object"; the property holds the name. + put("runtimeName", run.runtime.runtime) + put("runtimeVersion", run.runtime.runtimeVersion) + + put("requestedModel", serializeRequestedModel(run.requestedModel)) + + put("actorRef", run.subjectRefs.actor?.token) + put("requestRef", run.subjectRefs.request?.token) + put("sessionRef", run.subjectRefs.session?.token) + put("personalizationRef", run.subjectRefs.personalization?.token) + put("deploymentRef", run.subjectRefs.deployment?.token) + put("experimentRef", run.experimentRef?.token) + put("cohortRef", run.cohortRef?.token) + + put("replayFidelity", run.replayFidelity.name) + putAll(countsBindMap(run.counts)) + put("failures", serializeFailures(run.failures)) + } + + /** + * A digest of everything a header write owns, computed from the same map [headerBindMap] + * builds and compared verbatim against what an earlier write stored — never re-derived from a + * row read back afterward. + * + * This is what lets a save tell "nothing changed" from "something changed" atomically, in the + * same lock-then-read a compare-and-set needs anyway: two fingerprints computed the same way + * from the same set of fields are equal exactly when every field they cover is, so one string + * compare stands in for comparing a dozen properties by hand, the way [lineageKeyOf] already + * stands in for comparing lineage's five. `version` and this fingerprint itself are never part + * of the map it digests — a resend naming a stale version has to produce the same digest as the + * write that landed, or a byte-identical retry could never be told apart from a genuine change. + * + * Digested through [ExtractionRunFingerprint.ofFields], the canonical codec. A serializer's + * output can move without the header moving with it — a different `Map` implementation + * iterating [header] in a different order, a library upgrade changing how it renders a number — + * and every one of those would read as a genuine change to a save that is really a no-op or a + * byte-identical retry. The canonical encoding sorts by field name and renders every value as + * the plain string it already is, so the digest is a function of the header's content alone. + */ + fun headerFingerprint(header: Map): String = ExtractionRunFingerprint.ofFields( + ExtractionRunFingerprint.HEADER_VERSION, + header.mapValues { (_, value) -> value?.toString() }, + ) + + /** + * The properties a terminal write sets, and only those. + * + * Counts and failures are on the map only when the transition carries them. Null means "keep + * what the run recorded", so leaving the key out is exactly what + * [ExtractionRunTransition.applyTo] does with `counts ?: run.counts` — and binding null instead + * would remove the stored properties, which is the opposite. + */ + fun terminalBindMap(transition: ExtractionRunTransition): Map = buildMap { + put("status", transition.status.name) + putInstant("finishedAt", transition.finishedAt) + transition.counts?.let { putAll(countsBindMap(it)) } + transition.failures?.let { put("failures", serializeFailures(it)) } + } + + /** + * Rebuilds a run from a header node's properties and the child rows read alongside it. + * + * Throws on anything it cannot read, which is what lets the store log the node and skip it + * rather than handing back a run with invented defaults. A run named `""` because `runId` was + * missing would be indistinguishable from data. + */ + fun fromRow(row: Map<*, *>, invocations: List): ExtractionRun { + val lineage = ExtractionRunLineage.fromStoredFields( + runRef = ExtractionRunRef(row.str("runId")), + rootRunRef = ExtractionRunRef(row.str("rootRunId")), + parentRunRef = row.strOrNull("parentRunId")?.let(::ExtractionRunRef), + supersedesRunRef = row.strOrNull("supersedesRunId")?.let(::ExtractionRunRef), + passIndex = row.int("passIndex"), + ) + val storedLineageKey = row.str("lineageKey") + require(storedLineageKey == lineageKeyOf(lineage)) { + "ExtractionRun '${lineage.runRef.runId}' fails its lineage check: the stored lineageKey " + + "does not match the lineage properties on the node, so a save comparing against it " + + "would reach a different verdict than a read of the same node" + } + + val profileName = row.strOrNull("profileName") + val profileVersion = row.strOrNull("profileVersion") + require((profileName == null) == (profileVersion == null)) { + "ExtractionRun '${lineage.runRef.runId}' stores half a profile reference: a name and a " + + "version are written together or not at all" + } + + return ExtractionRun( + contextId = ContextId(row.str("contextId")), + lineage = lineage, + status = row.enum("status"), + startedAt = Instant.parse(row.str("startedAt")), + finishedAt = row.strOrNull("finishedAt")?.let(Instant::parse), + profile = profileName?.let { ExtractionContentProfileRef(it, profileVersion!!) }, + sourceRevisions = deserializeSourceRevisions(row.str("sourceRevisions")), + fingerprints = ExtractionRunFingerprints( + promptTemplateFingerprint = row.strOrNull("promptTemplateFingerprint"), + schemaFingerprint = row.strOrNull("schemaFingerprint"), + metamodelFingerprint = row.strOrNull("metamodelFingerprint"), + ), + runtime = ExtractionRuntimeIdentity( + extractor = row.strOrNull("extractor"), + extractorVersion = row.strOrNull("extractorVersion"), + hostApplication = row.strOrNull("hostApplication"), + runtime = row.strOrNull("runtimeName"), + runtimeVersion = row.strOrNull("runtimeVersion"), + ), + requestedModel = deserializeRequestedModel(row.strOrNull("requestedModel")), + subjectRefs = ExtractionRunSubjectRefs( + actor = row.strOrNull("actorRef")?.let(::ExtractionActorRef), + request = row.strOrNull("requestRef")?.let(::ExtractionRequestRef), + session = row.strOrNull("sessionRef")?.let(::ExtractionSessionRef), + personalization = row.strOrNull("personalizationRef")?.let(::ExtractionPersonalizationRef), + deployment = row.strOrNull("deploymentRef")?.let(::ExtractionDeploymentRef), + ), + experimentRef = row.strOrNull("experimentRef")?.let(::ExtractionExperimentRef), + cohortRef = row.strOrNull("cohortRef")?.let(::ExtractionCohortRef), + replayFidelity = row.enum("replayFidelity"), + counts = countsFromRow(row), + invocations = invocations, + failures = deserializeFailures(row.str("failures")), + version = row.long("version"), + ) + } + + /** + * Counts are six scalar properties rather than one JSON blob, because they are what a later + * run page will want to filter and aggregate on, and a JSON string is opaque to a query. + */ + private fun countsBindMap(counts: ExtractionRunCounts): Map = mapOf( + "countSourcesRead" to counts.sourcesRead, + "countChunksProcessed" to counts.chunksProcessed, + "countPropositionsExtracted" to counts.propositionsExtracted, + "countPropositionsPersisted" to counts.propositionsPersisted, + "countPropositionsRejected" to counts.propositionsRejected, + "countEntitiesResolved" to counts.entitiesResolved, + ) + + private fun countsFromRow(row: Map<*, *>): ExtractionRunCounts = ExtractionRunCounts( + sourcesRead = row.int("countSourcesRead"), + chunksProcessed = row.int("countChunksProcessed"), + propositionsExtracted = row.int("countPropositionsExtracted"), + propositionsPersisted = row.int("countPropositionsPersisted"), + propositionsRejected = row.int("countPropositionsRejected"), + entitiesResolved = row.int("countEntitiesResolved"), + ) +} + +/** + * Translates one attempt at one planned model call to and from its own node's properties. + * + * The key is `(contextId, runId, invocationIndex, attempt)`, allocated when the run's call plan is + * laid out rather than when a call returns, so a retry writes its own row and a replayed write of + * the same attempt upserts in place. Nothing here writes the key: it is in the MERGE pattern, where + * the uniqueness constraint can see it. + * + * Usage and provider-response facts are optional records with several optional fields each, so both + * are JSON strings. Two scalar properties per field would lose the difference between "the provider + * reported nothing" and "there is no report", and that difference is the whole point of separating + * what was asked for from what was observed. + */ +object ExtractionInvocationRowMapper { + + /** Bind values for one child row, key properties included so the MERGE pattern can use them. */ + fun bindMap(record: ExtractionInvocationRecord): Map = buildMap { + put("invocationIndex", record.invocationIndex) + put("attempt", record.attempt) + put("outcome", record.outcome.name) + put("configuredService", record.configuredService) + putInstant("startedAt", record.startedAt) + putInstant("finishedAt", record.finishedAt) + put("usage", serializeUsage(record.usage)) + put("providerResponse", serializeProviderResponse(record.providerResponse)) + } + + /** Rebuilds one attempt, throwing on anything the writer must have stored and did not. */ + fun fromRow(row: Map<*, *>): ExtractionInvocationRecord = ExtractionInvocationRecord( + id = ExtractionInvocationId( + invocationIndex = row.int("invocationIndex"), + attempt = row.int("attempt"), + ), + outcome = row.enum("outcome"), + configuredService = row.strOrNull("configuredService"), + startedAt = row.strOrNull("startedAt")?.let(Instant::parse), + finishedAt = row.strOrNull("finishedAt")?.let(Instant::parse), + usage = deserializeUsage(row.strOrNull("usage")), + providerResponse = deserializeProviderResponse(row.strOrNull("providerResponse")), + ) + + /** + * A digest of one attempt's whole payload, computed from the same map [bindMap] builds. + * + * Once an attempt is terminal, a later write for the same id is a replay exactly when this + * digest matches the string an earlier write stored — compared verbatim, the same rule + * [ExtractionRunTransition.fingerprint] follows for a run's own terminal write, applied one + * level down. Comparing the digest sidesteps the type drift a database round-trip can + * introduce — Neo4j hands an `Int` back as a `Long`, say — since both sides of the comparison + * come straight from a Kotlin record, always, independent of whatever a query happened to + * return. + * + * Digested through [ExtractionRunFingerprint.ofFields], for the reason [headerFingerprint] + * gives one level up: a raw JSON serialization of [bindMap] would move with the serializer, and + * this digest is compared verbatim against a string a past write stored, so it has to be a + * function of the record's fields alone. + */ + fun fingerprint(record: ExtractionInvocationRecord): String = ExtractionRunFingerprint.ofFields( + ExtractionRunFingerprint.INVOCATION_VERSION, + bindMap(record).mapValues { (_, value) -> value?.toString() }, + ) +} + +// ---- instants ---- + +/** + * Writes an instant as the three properties described on [ExtractionRunRowMapper], or removes all + * three when there is no instant. + * + * All three move together. A node carrying a `finishedAt` string and no epoch second would sort as + * though it had never finished. + */ +private fun MutableMap.putInstant(field: String, instant: Instant?) { + put(field, instant?.toString()) + put(field + "EpochSecond", instant?.epochSecond) + put(field + "Nano", instant?.nano) +} + +// ---- JSON encodings ---- + +/** + * Source revisions as `[{"sourceKey": ..., "sourceRevision": ...}, ...]`, in the order the run read + * them. + * + * The order is data: the contract says a run records which revisions it read in the order it read + * them, and [ExtractionRun] compares the lists element by element. Nothing here sorts. + */ +private fun serializeSourceRevisions(revisions: List): String = + objectMapper.writeValueAsString( + revisions.map { linkedMapOf("sourceKey" to it.sourceKey, "sourceRevision" to it.sourceRevision) }, + ) + +private fun deserializeSourceRevisions(serialized: String): List = + jsonObjects(serialized, "sourceRevisions").map { + SourceRevisionRef( + sourceKey = it.jsonStr("sourceRevisions", "sourceKey"), + sourceRevision = it.jsonStr("sourceRevisions", "sourceRevision"), + ) + } + +/** + * Failures as a list of objects, in the order the run recorded them, each with its classified code, + * its sanitized detail, when it happened, and the attempt it names if it names one. + * + * `detail` is written as it stands. [ExtractionFailure] has already bounded it to 512 characters and + * stripped it to a single line, and re-sanitizing here would let this mapper and that type disagree + * about what a stored detail is. + */ +private fun serializeFailures(failures: List): String = + objectMapper.writeValueAsString( + failures.map { failure -> + linkedMapOf( + "code" to failure.code.name, + "detail" to failure.detail, + "at" to failure.at.toString(), + "invocation" to failure.invocation?.let { + linkedMapOf("invocationIndex" to it.invocationIndex, "attempt" to it.attempt) + }, + ) + }, + ) + +private fun deserializeFailures(serialized: String): List = + jsonObjects(serialized, "failures").map { fields -> + ExtractionFailure( + code = jsonEnum(fields.jsonStr("failures", "code"), "failures", "code"), + detail = fields.jsonStr("failures", "detail"), + at = Instant.parse(fields.jsonStr("failures", "at")), + invocation = (fields["invocation"] as? Map<*, *>)?.let { + ExtractionInvocationId( + invocationIndex = it.jsonInt("failures", "invocationIndex"), + attempt = it.jsonInt("failures", "attempt"), + ) + }, + ) + } + +/** + * The requested model configuration as one JSON object with every field present, null included. + * + * A configuration with all its fields unset is a real configuration and has to stay distinct from + * having none — the run asked for the default rather than not asking. Absence of the property is + * "no configuration"; a `{}`-shaped object with null fields is "the default configuration". Eleven + * scalar properties on the node would collapse the two. + * + * `timeout` is an ISO-8601 duration string. `Duration.parse` inverts `Duration.toString` exactly, + * and a nanosecond count would be a number nobody reading the node could interpret. + */ +private fun serializeRequestedModel(config: ExtractionRequestedModelConfig?): String? = config?.let { + objectMapper.writeValueAsString( + linkedMapOf( + "modelRole" to it.modelRole, + "requestedModel" to it.requestedModel, + "temperature" to it.temperature, + "topP" to it.topP, + "topK" to it.topK, + "maxTokens" to it.maxTokens, + "presencePenalty" to it.presencePenalty, + "frequencyPenalty" to it.frequencyPenalty, + "thinkingFingerprint" to it.thinkingFingerprint, + "selectionFingerprint" to it.selectionFingerprint, + "timeout" to it.timeout?.toString(), + ), + ) +} + +private fun deserializeRequestedModel(serialized: String?): ExtractionRequestedModelConfig? { + val fields = jsonObject(serialized, "requestedModel") ?: return null + return ExtractionRequestedModelConfig( + modelRole = fields.jsonStrOrNull("modelRole"), + requestedModel = fields.jsonStrOrNull("requestedModel"), + temperature = fields.jsonDoubleOrNull("requestedModel", "temperature"), + topP = fields.jsonDoubleOrNull("requestedModel", "topP"), + topK = fields.jsonIntOrNull("requestedModel", "topK"), + maxTokens = fields.jsonIntOrNull("requestedModel", "maxTokens"), + presencePenalty = fields.jsonDoubleOrNull("requestedModel", "presencePenalty"), + frequencyPenalty = fields.jsonDoubleOrNull("requestedModel", "frequencyPenalty"), + thinkingFingerprint = fields.jsonStrOrNull("thinkingFingerprint"), + selectionFingerprint = fields.jsonStrOrNull("selectionFingerprint"), + timeout = fields.jsonStrOrNull("timeout")?.let(Duration::parse), + ) +} + +/** Observed token usage, for the same reason [serializeRequestedModel] is one object: empty is real. */ +private fun serializeUsage(usage: ExtractionModelUsage?): String? = usage?.let { + objectMapper.writeValueAsString( + linkedMapOf( + "inputTokens" to it.inputTokens, + "outputTokens" to it.outputTokens, + "totalTokens" to it.totalTokens, + "cachedInputTokens" to it.cachedInputTokens, + "reasoningTokens" to it.reasoningTokens, + ), + ) +} + +private fun deserializeUsage(serialized: String?): ExtractionModelUsage? { + val fields = jsonObject(serialized, "usage") ?: return null + return ExtractionModelUsage( + inputTokens = fields.jsonIntOrNull("usage", "inputTokens"), + outputTokens = fields.jsonIntOrNull("usage", "outputTokens"), + totalTokens = fields.jsonIntOrNull("usage", "totalTokens"), + cachedInputTokens = fields.jsonIntOrNull("usage", "cachedInputTokens"), + reasoningTokens = fields.jsonIntOrNull("usage", "reasoningTokens"), + ) +} + +/** What the provider said about its own answer. Same empty-is-real rule. */ +private fun serializeProviderResponse(facts: ExtractionProviderResponseFacts?): String? = facts?.let { + objectMapper.writeValueAsString( + linkedMapOf( + "responseModel" to it.responseModel, + "responseId" to it.responseId, + "finishReason" to it.finishReason, + "systemFingerprint" to it.systemFingerprint, + ), + ) +} + +private fun deserializeProviderResponse(serialized: String?): ExtractionProviderResponseFacts? { + val fields = jsonObject(serialized, "providerResponse") ?: return null + return ExtractionProviderResponseFacts( + responseModel = fields.jsonStrOrNull("responseModel"), + responseId = fields.jsonStrOrNull("responseId"), + finishReason = fields.jsonStrOrNull("finishReason"), + systemFingerprint = fields.jsonStrOrNull("systemFingerprint"), + ) +} + +// ---- strict accessors ---- + +/** + * Reads a property that must be there, and blows up naming it if it is not. + * + * Returning `""` for a missing property would let a node with no `runId` come back as a real-looking + * run named `""`, and the store's skip-the-unreadable-row guard would never fire for the most likely + * kind of corruption there is. Throwing is what gives that guard something to catch. + */ +private fun Map<*, *>.str(key: String): String = + this[key]?.toString() ?: throw IllegalArgumentException("required property '$key' is missing from the stored node") + +/** + * Reads a property whose absence means the run declared none — no profile, no experiment label, no + * finish time. Everything else goes through [str]. + */ +private fun Map<*, *>.strOrNull(key: String): String? = this[key]?.toString() + +/** Reads a required whole number. Neo4j hands integers back as `Long`, so the type is widened. */ +private fun Map<*, *>.int(key: String): Int = when (val value = this[key]) { + null -> throw IllegalArgumentException("required property '$key' is missing from the stored node") + is Number -> value.toInt() + else -> throw IllegalArgumentException( + "property '$key' is a ${value.javaClass.simpleName} where a number was expected" + ) +} + +/** Reads a required whole number as a `Long`, for the one property — `version` — wide enough to need it. */ +private fun Map<*, *>.long(key: String): Long = when (val value = this[key]) { + null -> throw IllegalArgumentException("required property '$key' is missing from the stored node") + is Number -> value.toLong() + else -> throw IllegalArgumentException( + "property '$key' is a ${value.javaClass.simpleName} where a number was expected" + ) +} + +/** + * Reads a stored enum constant by name. + * + * By name, never by ordinal: an ordinal re-points at a different constant the moment someone inserts + * a value into the enum, and these enums are still moving. A name this build does not have throws, + * so a node written by a later build is skipped with a message rather than read as something else. + */ +private inline fun > Map<*, *>.enum(key: String): E { + val stored = str(key) + return enumValues().firstOrNull { it.name == stored } ?: throw IllegalArgumentException( + "property '$key' is '$stored', which is not a known ${E::class.simpleName} — the node was " + + "written by a different version of the run model" + ) +} + +// ---- strict accessors over decoded JSON ---- + +/** Parses a stored JSON array of objects, refusing anything that is not one. */ +private fun jsonObjects(serialized: String, field: String): List> { + if (serialized.isEmpty()) return emptyList() + val parsed = objectMapper.readValue(serialized, Any::class.java) + val elements = parsed as? List<*> ?: throw IllegalArgumentException( + "the stored '$field' is a ${parsed?.javaClass?.simpleName ?: "null"} where a list was expected" + ) + return elements.map { element -> + element as? Map<*, *> ?: throw IllegalArgumentException( + "the stored '$field' holds a ${element?.javaClass?.simpleName ?: "null"} where an object was expected" + ) + } +} + +/** Parses a stored JSON object, or returns null when the property was absent. */ +private fun jsonObject(serialized: String?, field: String): Map<*, *>? { + if (serialized.isNullOrEmpty()) return null + val parsed = objectMapper.readValue(serialized, Any::class.java) + return parsed as? Map<*, *> ?: throw IllegalArgumentException( + "the stored '$field' is a ${parsed?.javaClass?.simpleName ?: "null"} where an object was expected" + ) +} + +private fun Map<*, *>.jsonStr(field: String, name: String): String = + this[name]?.toString() ?: throw IllegalArgumentException( + "an entry of the stored '$field' is missing its '$name'" + ) + +private fun Map<*, *>.jsonStrOrNull(name: String): String? = this[name]?.toString() + +private fun Map<*, *>.jsonInt(field: String, name: String): Int = + jsonIntOrNull(field, name) ?: throw IllegalArgumentException( + "an entry of the stored '$field' is missing its '$name'" + ) + +private fun Map<*, *>.jsonIntOrNull(field: String, name: String): Int? = when (val value = this[name]) { + null -> null + is Number -> value.toInt() + else -> throw IllegalArgumentException( + "the '$name' of the stored '$field' is a ${value.javaClass.simpleName} where a number was expected" + ) +} + +private fun Map<*, *>.jsonDoubleOrNull(field: String, name: String): Double? = when (val value = this[name]) { + null -> null + is Number -> value.toDouble() + else -> throw IllegalArgumentException( + "the '$name' of the stored '$field' is a ${value.javaClass.simpleName} where a number was expected" + ) +} + +private inline fun > jsonEnum(stored: String, field: String, name: String): E = + enumValues().firstOrNull { it.name == stored } ?: throw IllegalArgumentException( + "the '$name' of the stored '$field' is '$stored', which is not a known ${E::class.simpleName}" + ) diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunSchema.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunSchema.kt new file mode 100644 index 00000000..e0302297 --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunSchema.kt @@ -0,0 +1,131 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.agent.core.ContextId +import org.drivine.schema.RangeIndexSpec +import org.drivine.schema.SchemaItemSpec +import org.drivine.schema.UniquenessConstraintSpec + +/** + * The constraints, indexes and node labels [DrivineExtractionRunStore] needs, as plain data. + * + * A host declares [specs] in a `SchemaCatalog` bean and Drivine's schema manager ensures them on + * startup; the module's `TestApplication` shows the shape. [LABELS] is every label the store writes, + * which is what a test teardown or an operator's audit sweep needs. + * + * The three constraints are not optional tuning. Every write here is a `MERGE` or a `CREATE` on a + * natural key, and neither is race-free without a uniqueness constraint on that key: concurrent + * writers all miss the match, all take the create branch, and the store fills with copies of one + * run. The terminal-write constraint does more than that — see [DrivineExtractionRunStore] for why + * it is what makes compare-and-set hold across processes rather than only across threads. + */ +object ExtractionRunSchema { + + /** The `(:ExtractionRun)` header: one node per run, per tenant. */ + const val RUN_LABEL: String = "ExtractionRun" + + /** The `(:ExtractionRunInvocation)` child: one node per attempt at one planned model call. */ + const val INVOCATION_LABEL: String = "ExtractionRunInvocation" + + /** The `(:ExtractionRunTerminalWrite)` node: at most one per run, and the thing that says so. */ + const val TERMINAL_WRITE_LABEL: String = "ExtractionRunTerminalWrite" + + /** Header to invocation child. */ + const val RECORDED_REL: String = "RECORDED" + + /** Header to the terminal write that ended it. */ + const val ENDED_BY_REL: String = "ENDED_BY" + + /** + * The longest tenant id this store will write. + * + * `ContextId` accepts any non-blank string, and the tenant is the leading property of every key + * and index here. A tenant id of unbounded length is an index entry of unbounded length, and + * Neo4j fails that write with an index-key-size error — mid-extraction, with a message about + * bytes rather than about the tenant. The cap turns that into an argument rejection at the + * store's boundary, named, before anything is written. + * + * 1024 characters matches [com.embabel.dice.proposition.extraction.ExtractionRunLimits.MAX_SOURCE_KEY_LENGTH], + * the most generous bound the run model already puts on a key-like string. It is far above any + * real tenant id and comfortably below the index key limit even once the rest of a composite key + * is added. + * + * The cap is applied on writes only. A read for a tenant longer than this matches nothing, + * because nothing that long was ever stored — which is the fail-closed answer, and the only one + * a read could give. + */ + const val MAX_CONTEXT_ID_LENGTH: Int = 1024 + + /** + * Every constraint and index the store depends on. + * + * **Constraints.** + * - `ExtractionRun(contextId, runId)` — the tenant-qualified natural key. Two tenants minting + * the same run id are two runs, which is the whole reason the tenant is in the key. + * - `ExtractionRunInvocation(contextId, runId, invocationIndex, attempt)` — the deterministic + * child key. A retried attempt lands on its own row, and a replayed write of the same attempt + * upserts in place. + * - `ExtractionRunTerminalWrite(contextId, runId)` — at most one terminal write per run, ever. + * This is the compare-and-set. Two writers that both read a run as `RUNNING` both try to + * create this node, and the database lets exactly one of them commit. + * + * **Indexes.** Each one is what a specific read seeks on, and none of them carries `status` or + * the terminal fingerprint — deliberately, because the compare-and-set reads both after taking + * its lock and an indexed property could be served from an index entry read before it. + * - `(contextId)` — `runsInContext`. A composite index cannot stand in: Neo4j will not use one + * for a predicate on only its leading property, so without this the tenant predicate is a + * label scan. Same reasoning as the `Proposition(contextId)` index. + * - `(contextId, rootRunId)` — `runsOfRoot`, the read the denormalized root exists for. + * - `(contextId, parentRunId)` — `childrenOf`, one hop down the parent axis. + * - `(contextId, startedAtEpochSecond)` — the paging sort key. Pages are newest-first by start + * time within one tenant. + * - `ExtractionRunInvocation(contextId, runId)` — reading one run's attempts. The four-property + * constraint index is keyed on the whole identity and is not a substitute for a lookup on the + * first two. + */ + fun specs(): List = listOf( + UniquenessConstraintSpec(label = RUN_LABEL, properties = listOf("contextId", "runId")), + UniquenessConstraintSpec( + label = INVOCATION_LABEL, + properties = listOf("contextId", "runId", "invocationIndex", "attempt"), + ), + UniquenessConstraintSpec(label = TERMINAL_WRITE_LABEL, properties = listOf("contextId", "runId")), + + RangeIndexSpec(label = RUN_LABEL, property = "contextId"), + RangeIndexSpec(label = RUN_LABEL, properties = listOf("contextId", "rootRunId")), + RangeIndexSpec(label = RUN_LABEL, properties = listOf("contextId", "parentRunId")), + RangeIndexSpec(label = RUN_LABEL, properties = listOf("contextId", "startedAtEpochSecond")), + RangeIndexSpec(label = INVOCATION_LABEL, properties = listOf("contextId", "runId")), + ) + + /** Every node label the run store writes, for test cleanup and for an operator's audit sweep. */ + val LABELS: List = listOf(RUN_LABEL, INVOCATION_LABEL, TERMINAL_WRITE_LABEL) + + /** + * Rejects a tenant this store cannot key on, before anything is written. + * + * @throws IllegalArgumentException if the tenant id is longer than [MAX_CONTEXT_ID_LENGTH]. The + * message reports the length, not the value: a tenant id is a host identifier and does not + * belong in a log line. + */ + fun requireStorableTenant(contextId: ContextId) { + require(contextId.value.length <= MAX_CONTEXT_ID_LENGTH) { + "contextId must be at most $MAX_CONTEXT_ID_LENGTH characters to be stored as part of a " + + "run key, was ${contextId.value.length}" + } + } +} diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt new file mode 100644 index 00000000..2de39952 --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt @@ -0,0 +1,54 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.dice.proposition.extraction.ExtractionRunStore +import org.drivine.manager.PersistenceManager +import org.drivine.query.QuerySpecification +import org.junit.jupiter.api.AfterEach +import org.springframework.beans.factory.annotation.Autowired +import org.springframework.boot.test.context.SpringBootTest + +/** + * Runs the [AbstractExtractionRunStoreContractTest] suite against the Neo4j-backed + * [DrivineExtractionRunStore] (testcontainer). This is the half that catches the graph backend + * disagreeing with the in-memory reference, which is easy to do here: the state machine lives in a + * Cypher `FOREACH` over a conditional list there and in Kotlin `if`s in the reference, and the + * idempotency rule is a stored string on one side and a map entry on the other. + * + * One store bean serves every call, wiped between cases, which is the arrangement + * [DrivinePropositionStoreContractIntegrationTest] uses. The suite calls [store] more than once + * inside a few cases and expects a store holding nothing for its tenants each time; it gets that + * because those cases mint a distinct run id per iteration and only ever read by key. + */ +@SpringBootTest(classes = [TestApplication::class]) +class DrivineExtractionRunStoreContractIntegrationTest : AbstractExtractionRunStoreContractTest() { + + @Autowired + private lateinit var graphStore: DrivineExtractionRunStore + + @Autowired + private lateinit var persistenceManager: PersistenceManager + + override fun store(): ExtractionRunStore = graphStore + + @AfterEach + fun cleanUp() { + ExtractionRunSchema.LABELS.forEach { label -> + persistenceManager.execute(QuerySpecification.withStatement("MATCH (n:$label) DETACH DELETE n")) + } + } +} diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt new file mode 100644 index 00000000..e16ac99d --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt @@ -0,0 +1,789 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.agent.core.ContextId +import com.embabel.dice.proposition.extraction.ExtractionActorRef +import com.embabel.dice.proposition.extraction.ExtractionCohortRef +import com.embabel.dice.proposition.extraction.ExtractionContentProfileRef +import com.embabel.dice.proposition.extraction.ExtractionDeploymentRef +import com.embabel.dice.proposition.extraction.ExtractionExperimentRef +import com.embabel.dice.proposition.extraction.ExtractionFailure +import com.embabel.dice.proposition.extraction.ExtractionFailureCode +import com.embabel.dice.proposition.extraction.ExtractionInvocationId +import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome +import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord +import com.embabel.dice.proposition.extraction.ExtractionModelUsage +import com.embabel.dice.proposition.extraction.ExtractionPersonalizationRef +import com.embabel.dice.proposition.extraction.ExtractionProviderResponseFacts +import com.embabel.dice.proposition.extraction.ExtractionReplayFidelity +import com.embabel.dice.proposition.extraction.ExtractionRequestRef +import com.embabel.dice.proposition.extraction.ExtractionRequestedModelConfig +import com.embabel.dice.proposition.extraction.ExtractionRun +import com.embabel.dice.proposition.extraction.ExtractionRunConflictException +import com.embabel.dice.proposition.extraction.ExtractionRunCounts +import com.embabel.dice.proposition.extraction.ExtractionRunFingerprints +import com.embabel.dice.proposition.extraction.ExtractionRunKey +import com.embabel.dice.proposition.extraction.ExtractionRunLineage +import com.embabel.dice.proposition.extraction.ExtractionRunRef +import com.embabel.dice.proposition.extraction.ExtractionRunStatus +import com.embabel.dice.proposition.extraction.ExtractionRunSubjectRefs +import com.embabel.dice.proposition.extraction.ExtractionRunTransition +import com.embabel.dice.proposition.extraction.ExtractionRunTransitionOutcome +import com.embabel.dice.proposition.extraction.ExtractionRuntimeIdentity +import com.embabel.dice.proposition.extraction.ExtractionSessionRef +import com.embabel.dice.provenance.SourceRevisionRef +import org.drivine.manager.PersistenceManager +import org.drivine.query.QuerySpecification +import org.junit.jupiter.api.AfterEach +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.springframework.beans.factory.annotation.Autowired +import org.springframework.boot.test.context.SpringBootTest +import org.springframework.transaction.PlatformTransactionManager +import org.springframework.transaction.support.TransactionTemplate +import java.time.Duration +import java.time.Instant +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit + +/** + * What the cross-backend contract suite cannot ask, because it holds every backend to the same + * calls: whether the graph this store writes is the graph it says it writes, and whether the + * compare-and-set survives writers that share nothing but the database. + * + * Four groups: + * - **Compare-and-set across independent writers.** The contract's race case runs two threads + * through one store object; a durable backend has to hold when the writers are not one object, + * and its guarantee has to be visible in the graph rather than inferred from a return value. + * - **Cross-tenant fail-closed against a real index.** The contract asserts this over a store with + * one tenant's runs; here two tenants hold the *same run ids*, which is the arrangement that + * catches a query that scopes on `runId` and forgets `contextId`. + * - **Row fidelity.** The contract's runs are nearly empty, so every optional field, every JSON + * encoding, and every "absent means none" rule is untested by it. These runs carry everything. + * - **What a corrupt or oversized node does.** A skipped row, and a tenant the store refuses to key + * on. + */ +@SpringBootTest(classes = [TestApplication::class]) +class DrivineExtractionRunStoreIntegrationTest { + + @Autowired + private lateinit var store: DrivineExtractionRunStore + + @Autowired + private lateinit var persistenceManager: PersistenceManager + + @Autowired + private lateinit var transactionManager: PlatformTransactionManager + + private val tenant = ContextId("graph-tenant") + private val neighbour = ContextId("graph-neighbour") + private val startedAt: Instant = Instant.parse("2026-08-31T10:15:30.123456789Z") + private val finishedAt: Instant = Instant.parse("2026-08-31T10:16:00Z") + + @AfterEach + fun cleanUp() { + ExtractionRunSchema.LABELS.forEach { label -> + persistenceManager.execute(QuerySpecification.withStatement("MATCH (n:$label) DETACH DELETE n")) + } + } + + // ---- compare-and-set across independent writers ---- + + @Test + fun `writers sharing only the database produce exactly one applied transition`() { + // The multi-process shape, as close as one JVM gets to it. Each racing store is constructed + // directly rather than autowired, so there is no Spring proxy and no shared object between + // them; this class holds no mutable state at all, so nothing in the JVM can be serialising + // these writers. Whatever decides the race is the database. + repeat(8) { round -> + val run = running("graph-race-$round") + store.save(run) + val transition = ExtractionRunTransition.completed(finishedAt) + val writers = (0 until 6).map { + DrivineExtractionRunStore(persistenceManager, transactionManager) + } + + val start = CountDownLatch(1) + val pool = Executors.newFixedThreadPool(writers.size) + val results = try { + val futures = writers.map { writer -> + pool.submit { + start.await() + runCatching { writer.transition(run.key(), transition).outcome }.getOrElse { it } + } + } + start.countDown() + futures.map { it.get(60, TimeUnit.SECONDS) } + } finally { + pool.shutdownNow() + } + + assertEquals( + 1, + results.count { it == ExtractionRunTransitionOutcome.APPLIED }, + "exactly one writer ends the run: $results", + ) + assertEquals( + writers.size - 1, + results.count { it == ExtractionRunTransitionOutcome.REPLAYED }, + "every loser replays rather than conflicting: $results", + ) + // The invariant in the graph rather than in the return values. A backend that answered + // correctly while recording two terminal writes would still have lost the audit. + assertEquals(1, terminalWriteCount(run.key())) + assertEquals(ExtractionRunStatus.COMPLETED, store.findRun(run.key())?.status) + } + } + + @Test + fun `racing writers disagreeing about how the run ended leave one winner and one recorded write`() { + val run = running("graph-race-disagree") + store.save(run) + val transitions = listOf( + ExtractionRunTransition.completed(finishedAt), + ExtractionRunTransition.failed(finishedAt), + ExtractionRunTransition.cancelled(finishedAt), + ) + + val start = CountDownLatch(1) + val pool = Executors.newFixedThreadPool(transitions.size) + val results = try { + val futures = transitions.map { transition -> + val writer = DrivineExtractionRunStore(persistenceManager, transactionManager) + pool.submit { + start.await() + runCatching { writer.transition(run.key(), transition).outcome }.getOrElse { it } + } + } + start.countDown() + futures.map { it.get(60, TimeUnit.SECONDS) } + } finally { + pool.shutdownNow() + } + + assertEquals( + 1, + results.count { it == ExtractionRunTransitionOutcome.APPLIED }, + "exactly one writer ends the run: $results", + ) + assertEquals(1, terminalWriteCount(run.key())) + assertTrue(store.findRun(run.key())!!.status.isTerminal) + } + + @Test + fun `a second terminal write for one run cannot be stored at all`() { + // The compare-and-set rests on this constraint, so the constraint gets its own test. Without + // it the store's locking would be the only thing standing between two writers and two + // recorded endings, and a lock is a claim about behaviour rather than a schema fact. + val run = running("graph-constraint") + store.save(run) + store.transition(run.key(), ExtractionRunTransition.completed(finishedAt)) + + val second = assertThrows(RuntimeException::class.java) { + persistenceManager.execute( + QuerySpecification.withStatement( + """ + CREATE (:ExtractionRunTerminalWrite { + contextId: ${'$'}contextId, runId: ${'$'}runId, fingerprint: 'forged', status: 'FAILED' + }) + """.trimIndent(), + ).bind(mapOf("contextId" to tenant.value, "runId" to "graph-constraint")), + ) + } + + assertTrue( + generateSequence(second as Throwable) { it.cause }.any { + (it.message ?: "").contains("already exists", ignoreCase = true) || + (it.message ?: "").contains("ConstraintValidationFailed", ignoreCase = true) + }, + "the second terminal write is refused by the uniqueness constraint, not by chance: $second", + ) + assertEquals(1, terminalWriteCount(run.key())) + } + + @Test + fun `the stored fingerprint is the string the transition computed`() { + // Gap 8, asserted against the graph. A backend that re-derived the digest would still pass + // the contract's replay cases on a run nothing else touched. + val run = running("graph-fingerprint") + store.save(run) + val transition = ExtractionRunTransition.completed( + finishedAt = finishedAt, + counts = ExtractionRunCounts(propositionsPersisted = 3), + ) + store.transition(run.key(), transition) + + assertEquals(transition.fingerprint, storedFingerprint(run.key())) + + // And it stays that string while the run around it changes underneath it. + store.transition(run.key(), transition) + assertEquals(transition.fingerprint, storedFingerprint(run.key())) + } + + // ---- ambient transactions: validation precedes mutation, retries own a clean boundary ---- + + @Test + fun `a rejected recordInvocation leaves no orphan node even when the ambient transaction that catches it commits`() { + // The write-then-throw shape the reviewer named: RECORD_INVOCATION used to MERGE-create the + // child node before Kotlin found out the run had already ended, so a caller that catches the + // resulting exception without rolling back its own transaction would commit an empty stub + // nothing ever populated. Validation now happens before the MERGE runs at all. + val run = running("graph-orphan-on-terminal") + store.save(run) + store.transition(run.key(), ExtractionRunTransition.completed(finishedAt)) + val neverRecorded = ExtractionInvocationRecord.planned(0) + + val ambient = TransactionTemplate(transactionManager) + ambient.execute { + runCatching { store.recordInvocation(run.key(), neverRecorded) } + .onFailure { assertTrue(it is ExtractionRunConflictException) } + } + + assertEquals(0, childNodeCount(run.key())) + assertTrue(store.invocationsOf(run.key()).isEmpty()) + } + + @Test + fun `a mixed batch of recordInvocation calls in one ambient transaction commits the accepted writes and leaves the conflicting attempt untouched`() { + // Not a single Cypher statement's batch — the store never had one for recordInvocation, and + // save no longer writes invocations at all — but the same shape one call sequence up: several + // recordInvocation calls sharing one caller-owned transaction, one of which conflicts. The + // accepted writes are genuine and correctly committed; the conflicting write must still leave + // the row it targeted exactly where it was: untouched, whole, and unwritten. + val run = running("graph-mixed-batch") + store.save(run) + val lockedId = ExtractionInvocationId.planned(0) + val locked = ExtractionInvocationRecord( + id = lockedId, + outcome = ExtractionInvocationOutcome.SUCCEEDED, + configuredService = "service-locked", + ) + store.recordInvocation(run.key(), locked) + val acceptedA = ExtractionInvocationId(1, 1) + val acceptedB = ExtractionInvocationId(2, 1) + + val ambient = TransactionTemplate(transactionManager) + ambient.execute { + store.recordInvocation(run.key(), ExtractionInvocationRecord(id = acceptedA)) + store.recordInvocation(run.key(), ExtractionInvocationRecord(id = acceptedB)) + runCatching { + store.recordInvocation( + run.key(), + locked.copy(configuredService = "service-conflicting"), + ) + }.onFailure { assertTrue(it is ExtractionRunConflictException) } + } + + val stored = store.invocationsOf(run.key()).associateBy { it.id } + assertEquals(3, stored.size) + assertTrue(acceptedA in stored) + assertTrue(acceptedB in stored) + assertEquals(locked, stored.getValue(lockedId)) + } + + @Test + fun `two writers racing to record the same brand-new attempt both land, and neither retries inside the other's failed transaction`() { + // Two concurrent MERGE-creates of the same not-yet-existing invocation node can trip Neo4j's + // own deadlock detector, which aborts one transaction outright. Retrying that statement inside + // the transaction the database already discarded is not a retry against anything real; a + // recovered writer has to be reading and writing in a transaction the failure never touched. + // If it were retrying inside the dead one, this would hang, throw, or leave two disagreeing + // writes behind; a genuine retry converges on one. + repeat(5) { round -> + val run = running("graph-invocation-race-$round") + store.save(run) + val id = ExtractionInvocationId.planned(0) + val writers = (0 until 4).map { DrivineExtractionRunStore(persistenceManager, transactionManager) } + + val start = CountDownLatch(1) + val pool = Executors.newFixedThreadPool(writers.size) + val results = try { + val futures = writers.mapIndexed { index, writer -> + pool.submit { + start.await() + runCatching { + writer.recordInvocation( + run.key(), + ExtractionInvocationRecord(id = id, configuredService = "writer-$index"), + ) + }.getOrElse { it } + } + } + start.countDown() + futures.map { it.get(60, TimeUnit.SECONDS) } + } finally { + pool.shutdownNow() + } + + assertTrue( + results.all { it is ExtractionRun }, + "every writer recovers; a transient conflict never reaches the caller: $results", + ) + assertEquals(1, childNodeCount(run.key())) + } + } + + // ---- cross-tenant fail-closed, with identical run ids on both sides ---- + + @Test + fun `identical run ids in two tenants never collide on any read`() { + val ids = listOf("graph-root", "graph-child", "graph-grandchild") + listOf(tenant, neighbour).forEach { context -> + val root = ExtractionRunLineage.root(ExtractionRunRef(ids[0])) + val child = ExtractionRunLineage.childOf(ExtractionRunRef(ids[1]), root) + val grandchild = ExtractionRunLineage.childOf(ExtractionRunRef(ids[2]), child) + listOf(root, child, grandchild).forEachIndexed { index, lineage -> + store.save( + running( + lineage.runRef.runId, + contextId = context, + // The neighbour's runs are newer, so a page that limited before it scoped + // would return the neighbour's rows or none. + startedAt = startedAt.plusSeconds(if (context == neighbour) 1_000L + index else index.toLong()), + lineage = lineage, + ), + ) + } + } + store.recordInvocation(key(ids[1], neighbour), ExtractionInvocationRecord.planned(0)) + + // Keyed lookups. + assertEquals( + startedAt.plusSeconds(1), + store.findRun(key(ids[1], tenant))?.startedAt, + ) + assertEquals( + startedAt.plusSeconds(1_001), + store.findRun(key(ids[1], neighbour))?.startedAt, + ) + // The attempt belongs to the neighbour's run and is invisible from this tenant's. + assertTrue(store.invocationsOf(key(ids[1], tenant)).isEmpty()) + assertEquals(1, store.invocationsOf(key(ids[1], neighbour)).size) + + // Pages. + assertEquals( + listOf("graph-grandchild", "graph-child", "graph-root"), + store.runsInContext(tenant, 10, null).map { it.ref.runId }, + ) + assertEquals(3, store.runsInContext(neighbour, 10, null).size) + + // The inverse read down the parent axis, and the whole-lineage read. + assertEquals( + listOf(startedAt.plusSeconds(1)), + store.childrenOf(tenant, ExtractionRunRef(ids[0]), 10).map { it.startedAt }, + ) + assertEquals( + listOf(startedAt.plusSeconds(1_001)), + store.childrenOf(neighbour, ExtractionRunRef(ids[0]), 10).map { it.startedAt }, + ) + assertEquals(3, store.runsOfRoot(tenant, ExtractionRunRef(ids[0]), 10, null).size) + assertTrue( + store.runsOfRoot(tenant, ExtractionRunRef(ids[0]), 10, null) + .all { it.contextId == tenant }, + ) + + // The chain walk. + assertEquals( + listOf(startedAt.plusSeconds(1), startedAt), + store.ancestorsOf(key(ids[2], tenant), 10).map { it.startedAt }, + ) + } + + @Test + fun `ending one tenant's run leaves the other tenant's run of the same id running`() { + val mine = running("graph-shared") + val theirs = running("graph-shared", contextId = neighbour) + store.save(mine) + store.save(theirs) + + store.transition(mine.key(), ExtractionRunTransition.completed(finishedAt)) + + assertEquals(ExtractionRunStatus.COMPLETED, store.findRun(mine.key())?.status) + assertEquals(ExtractionRunStatus.RUNNING, store.findRun(theirs.key())?.status) + assertEquals(1, terminalWriteCount(mine.key())) + assertEquals(0, terminalWriteCount(theirs.key())) + } + + // ---- row fidelity ---- + + @Test + fun `a run carrying every field round-trips through the graph`() { + // save writes header fields only, so the invocations fullyPopulated built into the run + // object land through recordInvocation instead — the door the contract gives them. The + // failure that names one of them cannot ride the same save that inserts the header, because + // nothing has recorded that attempt yet at that point; it lands on a second save once + // recordInvocation has, which is also why the object this test compares against holds the + // version that second save produces, distinct from the version the first one left behind. + val bootstrap = fullyPopulated("graph-fidelity", failures = emptyList()) + store.save(bootstrap) + bootstrap.invocations.forEach { store.recordInvocation(bootstrap.key(), it) } + val run = fullyPopulated("graph-fidelity") + store.save(run) + val expected = fullyPopulated("graph-fidelity", version = 1) + + assertEquals(expected, store.findRun(run.key())) + assertEquals(expected, store.runsInContext(tenant, 10, null).single()) + } + + @Test + fun `a terminal run carrying every field round-trips through the graph`() { + val bootstrap = fullyPopulated("graph-fidelity-terminal", failures = emptyList()) + store.save(bootstrap) + bootstrap.invocations.forEach { store.recordInvocation(bootstrap.key(), it) } + val run = fullyPopulated("graph-fidelity-terminal") + store.save(run) + val expected = fullyPopulated("graph-fidelity-terminal", version = 1) + val transition = ExtractionRunTransition.completed( + finishedAt = finishedAt, + counts = ExtractionRunCounts(propositionsPersisted = 99), + failures = listOf(ExtractionFailure(ExtractionFailureCode.RATE_LIMITED, "slow down", finishedAt)), + ) + + val ended = store.transition(run.key(), transition) + + assertEquals(transition.applyTo(expected), ended.run) + assertEquals(transition.applyTo(expected), store.findRun(run.key())) + } + + @Test + fun `an absent optional field leaves no property on the node`() { + // "Absent means none" is an encoding, not a convention: a run with no profile has to be + // stored the way a run written before profiles existed was stored, or the two are different + // nodes on the same key. + store.save(running("graph-sparse")) + + val properties = nodeProperties("ExtractionRun", "graph-sparse") + + listOf( + "parentRunId", "supersedesRunId", "finishedAt", "finishedAtEpochSecond", "profileName", + "profileVersion", "promptTemplateFingerprint", "requestedModel", "actorRef", "cohortRef", + ).forEach { absent -> + assertTrue(absent !in properties, "expected no '$absent' property, found ${properties[absent]}") + } + // And the required ones are there, so the assertion above is not passing vacuously. + listOf("contextId", "runId", "rootRunId", "status", "startedAt", "lineageKey", "failures") + .forEach { present -> assertTrue(present in properties, "expected a '$present' property") } + } + + @Test + fun `invocation records live on their own nodes and a header save leaves them alone`() { + val run = running("graph-children") + store.save(run) + listOf( + ExtractionInvocationRecord.planned(0), + ExtractionInvocationRecord.planned(1), + ExtractionInvocationRecord(id = ExtractionInvocationId(1, 2)), + ).forEach { store.recordInvocation(run.key(), it) } + + // A caller that loaded the run before any attempt was recorded, then saved its own counts. + store.save( + ExtractionRun( + contextId = tenant, + lineage = run.lineage, + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + counts = ExtractionRunCounts(chunksProcessed = 12), + ), + ) + + assertEquals(3, childNodeCount(run.key())) + assertEquals( + listOf(ExtractionInvocationId(0, 1), ExtractionInvocationId(1, 1), ExtractionInvocationId(1, 2)), + store.invocationsOf(run.key()).map { it.id }, + ) + assertEquals(12, store.findRun(run.key())?.counts?.chunksProcessed) + // The header holds no attempt list of its own, which is why it cannot delete one. + assertTrue("invocations" !in nodeProperties("ExtractionRun", "graph-children")) + } + + @Test + fun `a run reads back with its attempts in plan order, whatever order they were recorded in`() { + // A durable store keeps identified rows, not the order a caller happened to list them in, so + // plan order is the only order it can offer — and it is the order the run model defines. + // `ExtractionRun.equals` compares the list element by element, so an unordered collect would + // make a stored run unequal to itself at random. + val run = running("graph-plan-order") + store.save(run) + listOf( + ExtractionInvocationId(2, 1), + ExtractionInvocationId(0, 2), + ExtractionInvocationId(1, 1), + ExtractionInvocationId(0, 1), + ).forEach { store.recordInvocation(run.key(), ExtractionInvocationRecord(id = it)) } + + val planOrder = listOf( + ExtractionInvocationId(0, 1), + ExtractionInvocationId(0, 2), + ExtractionInvocationId(1, 1), + ExtractionInvocationId(2, 1), + ) + assertEquals(planOrder, store.findRun(run.key())!!.invocations.map { it.id }) + assertEquals(planOrder, store.invocationsOf(run.key()).map { it.id }) + assertEquals( + planOrder, + store.runsInContext(tenant, 10, null).single().invocations.map { it.id }, + ) + } + + @Test + fun `nothing a run was never given reaches the stored bytes`() { + // The privacy contract, asserted over the properties actually written rather than over a + // test-local dumper. Everything DICE holds about a run is a token, a digest, or a + // classified code, so a secret the host never handed over cannot be in the row — and the + // one place it could leak in is a failure built from an exception whose message quotes it. + val secret = "patient-9f2a-had-a-consultation-on-tuesday" + val run = ExtractionRun( + contextId = tenant, + lineage = ExtractionRunLineage.root(ExtractionRunRef("graph-privacy")), + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + subjectRefs = ExtractionRunSubjectRefs( + actor = ExtractionActorRef("actor-token-1"), + session = ExtractionSessionRef("session-token-1"), + ), + failures = listOf( + ExtractionFailure.fromThrowable( + ExtractionFailureCode.DECODE_FAILED, + IllegalStateException("could not decode: $secret"), + startedAt, + ), + ), + ) + + store.save(run) + + val written = nodeProperties("ExtractionRun", "graph-privacy").values.joinToString(" ") { it.toString() } + assertTrue(secret !in written, "a source-text substring reached the stored row") + assertTrue("IllegalStateException" in written, "the classified cause chain is what is stored instead") + } + + // ---- corrupt and oversized ---- + + @Test + fun `a corrupt row is skipped and the readable runs still come back`() { + store.save(running("graph-good-1", startedAt = startedAt)) + store.save(running("graph-good-2", startedAt = startedAt.plusSeconds(10))) + persistenceManager.execute( + QuerySpecification.withStatement( + """ + CREATE (:ExtractionRun { + contextId: ${'$'}contextId, + runId: 'graph-corrupt', + startedAtEpochSecond: ${'$'}epochSecond, + startedAtNano: 0 + }) + """.trimIndent(), + ).bind( + mapOf( + "contextId" to tenant.value, + // Newest, so it sorts to the front and a backend that failed the whole read + // rather than skipping the row would return nothing at all. + "epochSecond" to startedAt.plusSeconds(100).epochSecond, + ), + ), + ) + + assertEquals( + listOf("graph-good-2", "graph-good-1"), + store.runsInContext(tenant, 10, null).map { it.ref.runId }, + ) + assertNull(store.findRun(key("graph-corrupt"))) + } + + @Test + fun `a row with no sort key is kept out of the order rather than filling a page slot`() { + store.save(running("graph-ordered")) + persistenceManager.execute( + QuerySpecification.withStatement( + "CREATE (:ExtractionRun {contextId: \$contextId, runId: 'graph-no-sort-key'})", + ).bind(mapOf("contextId" to tenant.value)), + ) + + // A limit of one: if the sort-key-less node reached the order it would sort first, spend the + // only slot, and then be dropped, so this would come back empty. + assertEquals( + listOf("graph-ordered"), + store.runsInContext(tenant, 1, null).map { it.ref.runId }, + ) + } + + @Test + fun `a tenant too long to key on is refused at the write, and reads for it are empty`() { + val oversized = ContextId("t".repeat(ExtractionRunSchema.MAX_CONTEXT_ID_LENGTH + 1)) + val run = running("graph-oversized", contextId = oversized) + + val rejected = assertThrows(IllegalArgumentException::class.java) { store.save(run) } + + assertTrue( + "${ExtractionRunSchema.MAX_CONTEXT_ID_LENGTH}" in rejected.message.orEmpty(), + "the message names the cap: ${rejected.message}", + ) + // The value itself is a host identifier and stays out of the message. + assertTrue(oversized.value !in rejected.message.orEmpty()) + assertTrue(store.runsInContext(oversized, 10, null).isEmpty()) + // The largest tenant this store will key on is stored without complaint. + val atCap = ContextId("t".repeat(ExtractionRunSchema.MAX_CONTEXT_ID_LENGTH)) + store.save(running("graph-at-cap", contextId = atCap)) + assertEquals(1, store.runsInContext(atCap, 10, null).size) + } + + // ---- fixtures ---- + + private fun running( + runId: String, + contextId: ContextId = tenant, + startedAt: Instant = this.startedAt, + lineage: ExtractionRunLineage = ExtractionRunLineage.root(ExtractionRunRef(runId)), + ): ExtractionRun = ExtractionRun( + contextId = contextId, + lineage = lineage, + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + ) + + /** + * A run with every optional field set, so the row mapper has something to get wrong. + * + * [failures] defaults to a failure naming an invocation this fixture also carries, and + * [version] defaults to `0` — right for a first save. A caller building the two-phase sequence + * `save` (invocations not yet recorded) → `recordInvocation` → `save` (with the failure) needs + * the header to hold no such failure on the first save, since the store's own read-back inside + * that call would otherwise reject a failure naming an attempt nothing has recorded yet — and + * needs `version = 1` on the object it compares the second save's result against, since that + * second save is what moves the header on from the version the first one left it at. + */ + private fun fullyPopulated( + runId: String, + failures: List = listOf( + ExtractionFailure( + code = ExtractionFailureCode.MODEL_TIMEOUT, + detail = "timed out after 90s", + at = startedAt.plusSeconds(90), + invocation = ExtractionInvocationId(0, 1), + ), + ), + version: Long = 0, + ): ExtractionRun { + val parent = ExtractionRunLineage.root(ExtractionRunRef("$runId-parent")) + return ExtractionRun( + contextId = tenant, + lineage = ExtractionRunLineage.childOf( + runRef = ExtractionRunRef(runId), + parent = parent, + supersedesRunRef = ExtractionRunRef("$runId-superseded"), + passIndex = 3, + ), + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + profile = ExtractionContentProfileRef("meeting-notes", "2.1"), + sourceRevisions = listOf( + SourceRevisionRef("doc://one", "rev-1"), + SourceRevisionRef("doc://two|with|pipes", "rev-2"), + ), + fingerprints = ExtractionRunFingerprints("prompt-abc", "schema-def", "metamodel-ghi"), + runtime = ExtractionRuntimeIdentity("llm-extractor", "1.4.0", "assistant", "dice", "0.2.0"), + requestedModel = ExtractionRequestedModelConfig( + modelRole = "extraction", + requestedModel = "some-model", + temperature = 0.25, + topP = 0.9, + topK = 40, + maxTokens = 4096, + presencePenalty = -0.5, + frequencyPenalty = 0.5, + thinkingFingerprint = "think-1", + selectionFingerprint = "select-1", + timeout = Duration.ofSeconds(90), + ), + subjectRefs = ExtractionRunSubjectRefs( + actor = ExtractionActorRef("actor-1"), + request = ExtractionRequestRef("request-1"), + session = ExtractionSessionRef("session-1"), + personalization = ExtractionPersonalizationRef("personalization-1"), + deployment = ExtractionDeploymentRef("deployment-1"), + ), + experimentRef = ExtractionExperimentRef("experiment-1"), + cohortRef = ExtractionCohortRef("cohort-b"), + replayFidelity = ExtractionReplayFidelity.APPROXIMATE, + counts = ExtractionRunCounts(1, 2, 3, 4, 5, 6), + invocations = listOf( + ExtractionInvocationRecord( + id = ExtractionInvocationId(0, 1), + outcome = ExtractionInvocationOutcome.FAILED, + configuredService = "openai", + startedAt = startedAt, + finishedAt = startedAt.plusMillis(1_500), + usage = ExtractionModelUsage(11, 22, 33, 44, 55), + providerResponse = ExtractionProviderResponseFacts("m-1", "resp-1", "length", "fp-1"), + ), + ExtractionInvocationRecord(id = ExtractionInvocationId(0, 2)), + ), + failures = failures, + version = version, + ) + } + + private fun key(runId: String, contextId: ContextId = tenant) = + ExtractionRunKey(contextId, ExtractionRunRef(runId)) + + // ---- graph probes ---- + + private fun terminalWriteCount(key: ExtractionRunKey): Int = countOf( + """ + MATCH (t:ExtractionRunTerminalWrite {contextId: ${'$'}contextId, runId: ${'$'}runId}) + RETURN {count: count(t)} AS row + """.trimIndent(), + key, + ) + + private fun childNodeCount(key: ExtractionRunKey): Int = countOf( + """ + MATCH (:ExtractionRun {contextId: ${'$'}contextId, runId: ${'$'}runId}) + -[:RECORDED]->(i:ExtractionRunInvocation) + RETURN {count: count(i)} AS row + """.trimIndent(), + key, + ) + + private fun countOf(statement: String, key: ExtractionRunKey): Int = + ((rows(statement, mapOf("contextId" to key.contextId.value, "runId" to key.runRef.runId)) + .single()["count"]) as Number).toInt() + + private fun storedFingerprint(key: ExtractionRunKey): String? = rows( + """ + MATCH (t:ExtractionRunTerminalWrite {contextId: ${'$'}contextId, runId: ${'$'}runId}) + RETURN {fingerprint: t.fingerprint} AS row + """.trimIndent(), + mapOf("contextId" to key.contextId.value, "runId" to key.runRef.runId), + ).single()["fingerprint"]?.toString() + + /** Every property actually written on one node, read back raw rather than through the mapper. */ + private fun nodeProperties(label: String, runId: String): Map = rows( + "MATCH (n:$label {contextId: \$contextId, runId: \$runId}) RETURN properties(n) AS row", + mapOf("contextId" to tenant.value, "runId" to runId), + ).single().entries.associate { (k, v) -> k.toString() to v } + + private fun rows(statement: String, bindings: Map): List> { + @Suppress("UNCHECKED_CAST") + val spec = QuerySpecification.withStatement(statement).bind(bindings) as QuerySpecification + return persistenceManager.query(spec).filterIsInstance>() + } +} diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/ExtractionRunRowMappersTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/ExtractionRunRowMappersTest.kt new file mode 100644 index 00000000..6b9fab44 --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/ExtractionRunRowMappersTest.kt @@ -0,0 +1,124 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import com.embabel.agent.core.ContextId +import com.embabel.dice.proposition.extraction.ExtractionInvocationId +import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord +import com.embabel.dice.proposition.extraction.ExtractionRun +import com.embabel.dice.proposition.extraction.ExtractionRunCounts +import com.embabel.dice.proposition.extraction.ExtractionRunFingerprint +import com.embabel.dice.proposition.extraction.ExtractionRunLineage +import com.embabel.dice.proposition.extraction.ExtractionRunRef +import com.embabel.dice.proposition.extraction.ExtractionRunStatus +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNotEquals +import org.junit.jupiter.api.Test +import java.time.Instant + +/** + * Proves the header and invocation fingerprints are functions of the record's content alone — not + * of a JSON serializer's own behaviour — which is the defect a reviewer found in the raw + * `objectMapper.writeValueAsString` version: two writes of the same logical state could digest + * differently if a serializer or a `Map` implementation changed underneath them, turning a correct + * retry into a rejected conflict. + */ +class ExtractionRunRowMappersTest { + + private val startedAt: Instant = Instant.parse("2026-08-31T10:15:30Z") + + private fun run(counts: Int = 0) = ExtractionRun( + contextId = ContextId("tenant"), + lineage = ExtractionRunLineage.root(ExtractionRunRef("run-1")), + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + ) + + @Test + fun `headerFingerprint does not move when the bind map is rebuilt with a different key order`() { + val header = ExtractionRunRowMapper.headerBindMap(run()) + // Same entries, a different iteration order — what a HashMap vs a LinkedHashMap, or a + // future refactor of headerBindMap's own field-adding order, would produce for equal data. + val reordered = LinkedHashMap(header.entries.sortedBy { it.key }.associate { it.key to it.value }) + + assertEquals( + ExtractionRunRowMapper.headerFingerprint(header), + ExtractionRunRowMapper.headerFingerprint(reordered), + ) + } + + @Test + fun `headerFingerprint is unaffected by the Map implementation carrying the same entries`() { + val header = ExtractionRunRowMapper.headerBindMap(run()) + val asHashMap = HashMap(header) + val asTreeMap = header.entries.associateTo(sortedMapOf()) { it.key to it.value } + + val baseline = ExtractionRunRowMapper.headerFingerprint(header) + assertEquals(baseline, ExtractionRunRowMapper.headerFingerprint(asHashMap)) + assertEquals(baseline, ExtractionRunRowMapper.headerFingerprint(asTreeMap)) + } + + @Test + fun `headerFingerprint changes when a header field genuinely changes`() { + val a = ExtractionRunRowMapper.headerFingerprint(ExtractionRunRowMapper.headerBindMap(run())) + val changed = run().let { base -> + ExtractionRun( + contextId = base.contextId, + lineage = base.lineage, + status = base.status, + startedAt = base.startedAt, + counts = ExtractionRunCounts(propositionsPersisted = 7), + ) + } + val b = ExtractionRunRowMapper.headerFingerprint(ExtractionRunRowMapper.headerBindMap(changed)) + + assertNotEquals(a, b) + } + + @Test + fun `invocation fingerprint does not move when the bind map is rebuilt with a different key order`() { + val record = ExtractionInvocationRecord( + id = ExtractionInvocationId.planned(0), + configuredService = "service-a", + ) + val bound = ExtractionInvocationRowMapper.bindMap(record) + val reordered = LinkedHashMap(bound.entries.sortedBy { it.key }.associate { it.key to it.value }) + + assertEquals( + invocationFingerprintOf(bound), + invocationFingerprintOf(reordered), + ) + } + + @Test + fun `invocation fingerprint changes when the record genuinely changes`() { + val a = ExtractionInvocationRowMapper.fingerprint( + ExtractionInvocationRecord(id = ExtractionInvocationId.planned(0), configuredService = "service-a"), + ) + val b = ExtractionInvocationRowMapper.fingerprint( + ExtractionInvocationRecord(id = ExtractionInvocationId.planned(0), configuredService = "service-b"), + ) + + assertNotEquals(a, b) + } + + /** Runs a bind map through the same digest [ExtractionInvocationRowMapper.fingerprint] uses. */ + private fun invocationFingerprintOf(bound: Map): String = + ExtractionRunFingerprint.ofFields( + ExtractionRunFingerprint.INVOCATION_VERSION, + bound.mapValues { (_, value) -> value?.toString() }, + ) +} diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt index 1d4f0dfd..7d7256a7 100644 --- a/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt @@ -160,6 +160,15 @@ open class TestApplication { persistenceManager: PersistenceManager, ): DrivineCollectorTraceStore = DrivineCollectorTraceStore(persistenceManager) + @Bean + open fun extractionRunSchema(): SchemaCatalog = SchemaCatalog.of(ExtractionRunSchema.specs()) + + @Bean + open fun extractionRunStore( + persistenceManager: PersistenceManager, + transactionManager: PlatformTransactionManager, + ): DrivineExtractionRunStore = DrivineExtractionRunStore(persistenceManager, transactionManager) + @Bean open fun decayManager( repository: DrivinePropositionRepository, diff --git a/dice/src/main/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprint.kt b/dice/src/main/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprint.kt index 492e218b..32db5f85 100644 --- a/dice/src/main/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprint.kt +++ b/dice/src/main/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprint.kt @@ -93,6 +93,15 @@ object ExtractionRunFingerprint { */ const val TERMINAL_VERSION: String = "xrun-terminal:v2" + /** Version tag on a durable backend's own header-fingerprint encoding. */ + const val HEADER_VERSION: String = "xrun-header:v1" + + /** Version tag on a durable backend's own invocation-record-fingerprint encoding. */ + const val INVOCATION_VERSION: String = "xrun-invocation:v1" + + /** What an absent value encodes as, kept distinct from any value a caller could supply. */ + private const val ABSENT = "-" + /** * The digest of one terminal write: the status it asserts, and when the run finished. * @@ -111,6 +120,32 @@ object ExtractionRunFingerprint { return digest(TERMINAL_VERSION + "|" + payload) } + /** + * The digest of an arbitrary flat set of named fields, run through the same canonical rules as + * [ofTerminal]: every field is a length-prefixed token, fields are sorted by name so the bytes + * never depend on the order a map happens to iterate in, and a `null` value renders as its own + * distinct absent marker, kept apart from empty text. + * + * This is the codec a durable backend's own comparison digests — a header's, an invocation + * record's — should run through, ahead of handing the same map to a general-purpose JSON + * serializer. A serializer's output moves with things that carry no meaning: which `Map` + * implementation built the payload, a library version bumping how it renders a number, a + * config flag turned on for an unrelated reason. Two writes of the same logical state must + * produce the same digest regardless of any of that, or a correct retry reads as a conflict. + * + * @param version Names the field set [named] belongs to, so two backends digesting different + * shapes under different versions can never collide. Two calls passing the same [version] + * must mean the same set of field names, or their digests are not comparable to each other. + * @param named The named values to digest, already rendered to the strings that mean what the + * caller means — this function does no rendering of its own. A `null` value is encoded as + * absent, distinct from an empty string. + */ + @JvmStatic + fun ofFields(version: String, named: Map): String { + val payload = fields(*named.map { (name, value) -> name to (value ?: ABSENT) }.toTypedArray()) + return digest(version + "|" + payload) + } + /** * Fixed-width rendering: seconds since the epoch, a dot, then nanoseconds padded to nine * digits. Two instants that compare equal always render identically, and no instant renders diff --git a/dice/src/test/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprintTest.kt b/dice/src/test/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprintTest.kt index 92d8b98a..31d02d06 100644 --- a/dice/src/test/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprintTest.kt +++ b/dice/src/test/kotlin/com/embabel/dice/proposition/extraction/ExtractionRunFingerprintTest.kt @@ -233,4 +233,66 @@ class ExtractionRunFingerprintTest { assertThat(transition.toString()).doesNotContain(transition.fingerprint) assertThat(transition.toString()).contains(transition.fingerprint.take(12)) } + + // ---- ofFields: the general-purpose codec a durable backend digests its own state through ---- + + @Test + fun `ofFields does not move with the order the caller built the map in`() { + // The whole point: a Map's iteration order is an implementation detail with no promised + // stability, and a digest that depended on it would turn a harmless map-construction change + // into a false conflict for a stored write it is compared against. + val forwards = linkedMapOf("alpha" to "1", "beta" to "2", "gamma" to "3") + val backwards = linkedMapOf("gamma" to "3", "beta" to "2", "alpha" to "1") + + assertThat(ExtractionRunFingerprint.ofFields("test:v1", forwards)) + .isEqualTo(ExtractionRunFingerprint.ofFields("test:v1", backwards)) + } + + @Test + fun `ofFields tells absent apart from an empty string`() { + val absent = mapOf("detail" to null) + val empty = mapOf("detail" to "") + + assertThat(ExtractionRunFingerprint.ofFields("test:v1", absent)) + .isNotEqualTo(ExtractionRunFingerprint.ofFields("test:v1", empty)) + } + + @Test + fun `ofFields changing any one field's value changes the digest`() { + val baseline = ExtractionRunFingerprint.ofFields("test:v1", mapOf("a" to "1", "b" to "2")) + val changed = ExtractionRunFingerprint.ofFields("test:v1", mapOf("a" to "1", "b" to "3")) + + assertThat(changed).isNotEqualTo(baseline) + } + + @Test + fun `ofFields under two different versions never collide even over the same field set`() { + val fields = mapOf("a" to "1", "b" to "2") + + assertThat(ExtractionRunFingerprint.ofFields("v1", fields)) + .isNotEqualTo(ExtractionRunFingerprint.ofFields("v2", fields)) + } + + @Test + fun `ofFields is a lowercase sha-256 hex digest`() { + val digest = ExtractionRunFingerprint.ofFields("test:v1", mapOf("a" to "1")) + + assertThat(digest).hasSize(64) + assertThat(digest).matches("[0-9a-f]{64}") + } + + @Test + fun `the header and invocation version tags are distinct from the terminal one`() { + // A reader meeting a version it does not know matches nothing; it never guesses. That only + // holds if the three encodings this module ships cannot be mistaken for each other. + assertThat(ExtractionRunFingerprint.HEADER_VERSION).isEqualTo("xrun-header:v1") + assertThat(ExtractionRunFingerprint.INVOCATION_VERSION).isEqualTo("xrun-invocation:v1") + assertThat( + setOf( + ExtractionRunFingerprint.TERMINAL_VERSION, + ExtractionRunFingerprint.HEADER_VERSION, + ExtractionRunFingerprint.INVOCATION_VERSION, + ), + ).hasSize(3) + } } diff --git a/docs/design/extraction-runs.md b/docs/design/extraction-runs.md index dacf6493..7663daed 100644 --- a/docs/design/extraction-runs.md +++ b/docs/design/extraction-runs.md @@ -750,6 +750,159 @@ recorded" and fails "an incompatible terminal rewrite is rejected", a finder tha passes every single-tenant read and fails "a page scopes before it limits", and a chain walk written as a recursive Cypher pattern passes on a healthy graph and hangs on a cycle. +## The Drivine store + +`DrivineExtractionRunStore` is the durable implementation. It writes three node labels: + +| Node | Key | Holds | +| --- | --- | --- | +| `(:ExtractionRun)` | `(contextId, runId)` | the header — lineage, envelope, requested configuration, counts, failures | +| `(:ExtractionRunInvocation)` | `(contextId, runId, invocationIndex, attempt)` | one attempt at one planned model call | +| `(:ExtractionRunTerminalWrite)` | `(contextId, runId)` | the fingerprint of the write that ended the run | + +The header reaches each child by `[:RECORDED]` and its terminal write by `[:ENDED_BY]`. + +Every key is tenant-qualified, and the tenant needs no encoding of the kind +`DrivineDriftReportStore` gives its scope. That store's scope is nullable and a Cypher MERGE cannot +key on a null, so a global report needed a non-null `ctx:`-prefixed stand-in no real context id +could collide with. A run's tenant is never null, so the plain value is already an injective key. + +### Compare-and-set, and why it holds across processes + +`transition` is one Cypher statement, so it is one transaction, and it works two ways at once. + +**A lock.** The statement's first act after matching the run is `SET n.casLock = $lockToken`, before +it reads anything. A property write takes an exclusive lock on the node and holds it to commit, so a +second transaction reaching that line blocks. Neo4j reads at read-committed and the read of +`n.status` is downstream of the `SET` in the same statement, so the second transaction reads the +status after it took the lock and therefore after the first committed. It sees a terminal run and +takes the no-op branch. `DrivineMetamodelVersionStore` uses the same write-lock-before-read idiom on +its counter. + +Two details make that argument hold rather than nearly hold. The lock token is a fresh UUID on every +call, so the write is always a real change and never a no-op a database is free to optimize away +before taking the lock. And no index carries `status` or the terminal fingerprint, so the planner +cannot serve the post-lock read from an index entry it read at MATCH time — which is the one way a +lock-then-read reads stale. + +**A constraint.** That argument is about Neo4j's behaviour, and behaviour is a thing to be wrong +about. So the terminal write is also a `CREATE` of an `(:ExtractionRunTerminalWrite)` node on the +run's own key, under a uniqueness constraint. Two transactions that both read a run as `RUNNING` +both try to create it, and the database refuses the second at commit. The loser's whole transaction +rolls back, header included; the store catches the violation, re-reads the recorded fingerprint in a +fresh transaction, and answers replayed or conflict. **Exactly one terminal write per run is a schema +fact, not an inference.** + +The constraint's sufficiency is measured rather than assumed: removing the lock leaves every race +test green — the constraint carries it alone. Removing both and letting the create become a merge +produces six racing writers all reporting `APPLIED`, and, when they disagree about how the run +ended, three contradictory endings recorded for one run. That is the multi-process failure the +in-memory reference's monitor cannot speak to, reproduced and then closed. The lock's own +sufficiency rests on the isolation argument above, unmeasured, which is why the constraint exists. +The constraint backstops `transition` only: the status guards in `save` and `recordInvocation` +rest on the lock argument alone, so a failed lock idiom could set a terminal header back to +`RUNNING` beside its terminal write — the one-write audit fact would survive; the header would not. + +### The fingerprint is stored, never re-derived + +The terminal-write node carries the exact string `ExtractionRunTransition.fingerprint` computed, and +a repeated terminal write is decided by comparing against that string. A store that re-derived a +digest from the stored run would reject a correct retry whenever an attempt had been recorded in +between, because the run it derived from would have changed while the terminal write did not. The +encoding has one implementation, in `dice`, and the graph holds its output. + +### A header write cannot touch a child row + +Invocation records are their own nodes on their own key, so `save` has no way to delete one. The +contract's rule that a save merges records rather than replacing them falls out of the graph model +instead of being implemented: `save` MERGEs the records the incoming run carries, which upserts in +place on a shared identity, and leaves every other child row alone. + +One consequence is worth naming. A durable store keeps identified rows, not the order a caller +happened to list them in, so it hands attempts back in plan order — the order `invocationsOf` +promises. `InMemoryExtractionRunStore` hands back the caller's order. `ExtractionRun.invocations` is +documented as being in whatever order the caller supplied, and `ExtractionRun.equals` compares it +element by element, so the two backends can return runs that are unequal for the same call sequence. +The ordered read both agree on is `invocationsOf`. + +### Reads + +Each page puts its tenant in the MATCH pattern and its `LIMIT` after the `ORDER BY`, which is the +drift-report store's rule carried over. A page also excludes rows with no sort key: Neo4j sorts null +largest, so a node missing `startedAtEpochSecond` would sort to the front of a `DESC` order, spend a +slot of the caller's `limit`, and then be dropped by the mapper — hiding a good run behind a broken +one. Corrupt rows are logged and skipped rather than failing a whole audit read, so a page can come +back shorter than asked for; reading further to backfill would break the bound the contract keeps. + +`runsOfRoot` is one indexed lookup on the denormalized root. The chain walk is client-side: a parent +is a property rather than a relationship, because a run can name a parent that has not been stored +yet and an edge cannot point at a node that does not exist. Without an edge there is no +variable-length pattern to walk, and the APOC procedures that would do it in one round trip are not a +dependency this module takes. So the walk is at most `limit` keyed lookups inside one read +transaction, stopping on a run already seen and resolving every hop inside the starting tenant. + +Every read returns each run with its attempts attached, pages included, because the in-memory +reference does and the two are held to one suite. The cost is bounded — at most `limit` runs times +the model's cap of 1024 attempts — but it is real on a wide page. + +### Schema + +`ExtractionRunSchema.specs()` is the whole dependency, and a host declares it in a `SchemaCatalog` +bean. Three uniqueness constraints — `ExtractionRun(contextId, runId)`, +`ExtractionRunInvocation(contextId, runId, invocationIndex, attempt)`, and +`ExtractionRunTerminalWrite(contextId, runId)` — and five range indexes: `ExtractionRun(contextId)` +for the tenant page, `(contextId, rootRunId)` for the lineage read, `(contextId, parentRunId)` for +the parent axis, `(contextId, startedAtEpochSecond)` for the paging sort key, and +`ExtractionRunInvocation(contextId, runId)` for reading one run's attempts. The constraints are not +tuning: a MERGE on a natural key is race-free only under one, and the third is the compare-and-set. + +`ContextId` accepts any non-blank string and the tenant is the leading property of every key here, so +the store caps it at 1024 characters on the write path — the bound the run model already puts on a +key-like string. An uncapped tenant id is an index entry of unbounded length, and Neo4j fails that +write mid-extraction with a message about bytes rather than about the tenant. The cap turns it into a +named argument rejection before anything is written. Reads are not capped, because a read for a +tenant longer than the cap matches nothing by construction, which is the fail-closed answer and the +only one a read could give. + +## Protected-content references + +An extraction run holds no content. Where a host wants replay material — a prompt as it was +rendered, a source document as it was read — it keeps that material itself and gives DICE a +`ProtectedContentRef`: what kind of material it is, when it stops being available, and an opaque +handle for looking it up. + +The handle is a `ProtectedContentHandle`, an eighth member of the `ExtractionOpaqueRef` family, so +it inherits the enforced floor the other seven have — bounded, character-restricted, truncating +`toString`, and a validation message that never quotes the value it rejected. That matters more here +than elsewhere: a host holding material in object storage has a link to hand and no reason to think +DICE minds, and the character set rejects one outright. + +What a host owes a reference it mints: + +- **The handle is not dereferenceable.** It names material in the host's own store and carries no + way to reach it: no signed URL, no pre-authenticated link, no bearer token, no decryption key, no + filename on a share. Holding one grants nothing. The rule is broader than the character check, and + the part the check cannot see is the host's to keep. +- **Expiry is the host's to honour.** `expiresAt` is a declaration about the host's store. DICE acts + on none of it — it does not sweep, does not delete, and does not refuse to hand back a reference + whose expiry has passed. `isExpiredAt` is the whole of what DICE offers. Keeping the expired + reference is deliberate: a run that was replayable until March and is not replayable now is a fact + about the run, and deleting the reference would make it look as though it never had replay + material. +- **The classification is a claim.** DICE stores it and carries it into the audit surface, and + verifies none of it. + +`ProtectedContentClassification` runs `INTERNAL`, `CONFIDENTIAL`, `PERSONAL`, least to most +restrictive. There is no "unknown" value: a host that cannot classify material classifies it +`PERSONAL`, because guessing low means personal data treated as operational logging and guessing +high means something operational looked after more carefully than it needed to be. + +Blob storage, key management and retention enforcement are outside #67 by that issue's own text. +This is the reference and its contract; there is no DICE-side store behind it. Nothing attaches a +reference to an `ExtractionRun` yet either, and `ExtractionReplayFidelity` still tops out at +`APPROXIMATE` — a run that carries replay material is a later slice's change to the run header, and +a test asserts the run holds no such field today. + ## OpenTelemetry GenAI naming, not adopted OTel's GenAI semantic conventions cover the same ground — `gen_ai.request.*`, `gen_ai.response.*`, @@ -837,9 +990,9 @@ public surface. ## What is not here yet -- **No durable store.** `InMemoryExtractionRunStore` is the only implementation. The Drivine store - — the tenant-qualified natural key, the deterministic child key for invocation records, the - uniqueness constraints, and the Cypher that scopes before it limits — is the next slice. +- **No auto-configuration.** `DrivineExtractionRunStore` is a bean a host declares itself, along + with the `SchemaCatalog` carrying `ExtractionRunSchema.specs()`. An `ExtractionRunAutoConfiguration` + arrives with the coordinator. - **No coordinator.** Nothing constructs an `ExtractionRun` during extraction yet, and nothing calls `save` or `transition` outside tests. Which means the `COMPLETED` precondition is documented and structurally narrowed, not observed: the wiring slice is where "the coordinator really does wait From b43360c4f4804f8bfff5dc44a4f2ed3744f8639a Mon Sep 17 00:00:00 2001 From: James Dunnam <7660553+jimador@users.noreply.github.com> Date: Wed, 2 Sep 2026 02:13:42 -0400 Subject: [PATCH 2/4] Persist the failure vocabulary and announce transitions from Drivine Failures store as flat typed fields with no text-shaped property, and a raw-property canary pins the stored key set against an allowlist on both write doors, so a free-text column cannot return unnoticed. Reads tolerate an unrecognized property on a corrupted node and return the recognized fields, stated in the mapper's KDoc. The store announces an applied transition after the surrounding transaction commits: a caller-owned rollback delivers nothing and leaves the run untouched, pinned by test, and a replay or rejected write announces nothing. The terminal fingerprint comparison carries the v2 identity-only material, and the dead v1 count and failure encoders are gone. The inherited contract suite runs unmodified against Neo4j. Signed-off-by: James Dunnam <7660553+jimador@users.noreply.github.com> --- CHANGELOG.md | 21 ++- .../dice/storage/DrivineExtractionRunStore.kt | 56 +++++- .../dice/storage/ExtractionRunRowMappers.kt | 84 ++++++--- ...tractionRunStoreContractIntegrationTest.kt | 14 +- ...rivineExtractionRunStoreIntegrationTest.kt | 162 +++++++++++++++++- .../embabel/dice/storage/TestApplication.kt | 35 +++- docs/design/extraction-runs.md | 82 ++++----- 7 files changed, 379 insertions(+), 75 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ced2fcaf..d0e42c8b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1652,7 +1652,26 @@ and the consumer PRs that deliver it). **The fingerprint is stored verbatim and compared verbatim, never re-derived.** The node carries the exact string `ExtractionRunTransition.fingerprint` computed, so a correct retry that happened after another attempt was recorded still replays; a store deriving a digest from the stored run would - reject it. + reject it. That string is the `xrun-terminal:v2` digest of the transition's identity, so the counts + and failures a terminal write carries ride into the header beside it and reach none of the compared + bytes; the store keeps no comparison of its own that could fall out of step. + **A run that ends announces itself once, when the write is durable.** `transition` hands an + `ExtractionRunTransitioned` to the listener the store was constructed with, defaulting to + `DiceEventListener.DEV_NULL` the way the in-memory reference's does. Exactly one call per run + reaches that branch, for the schema reason above: reaching it means having created the + terminal-write node. A replay, a rejected write, a `save`, a `recordInvocation`, and the writer that + lost the race all announce nothing. When the store owns the transaction the listener runs once the + template has committed; when a caller's transaction is active the announcement is registered against + that caller's commit, so a rollback drops it and no consumer hears about a run nothing can read back. The + six event cases the contract suite added run against Neo4j unmodified. + **Failures are stored in the closed vocabulary and nothing else fits.** A run's failures are one + JSON array on the header node, each element carrying `code`, `stage`, `providerStatus`, a measure as + a quantity and a value, `at`, and the attempt it names as an index and an attempt — eight fields, + written flat, with optional ones stored as nulls so every failure stores the same keys. No property + holds free text, because `ExtractionFailure` has no text-shaped field to write from. A round-trip + test reads the properties Neo4j actually holds and checks the key set it finds matches an allowlist + the test states itself, so an added `detail` column fails the build before it reaches a graph, and a stored + failure holding half a measure or half an invocation id is refused on read. **A header write cannot touch a child row.** Invocation records are their own nodes, so `save` has no way to delete one — the contract's merge-don't-replace rule falls out of the graph model instead of being implemented. One consequence: a durable store keeps identified rows rather than the order a diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt index 76a89a8a..fada452b 100644 --- a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt @@ -15,6 +15,8 @@ */ package com.embabel.dice.storage +import com.embabel.dice.common.DiceEventListener +import com.embabel.dice.common.ExtractionRunTransitioned import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord import com.embabel.dice.proposition.extraction.ExtractionRun @@ -33,6 +35,7 @@ import org.slf4j.LoggerFactory import org.springframework.transaction.PlatformTransactionManager import org.springframework.transaction.annotation.Propagation import org.springframework.transaction.annotation.Transactional +import org.springframework.transaction.support.TransactionSynchronization import org.springframework.transaction.support.TransactionSynchronizationManager import org.springframework.transaction.support.TransactionTemplate import java.time.Clock @@ -134,17 +137,35 @@ import java.util.UUID * in-memory reference does and the two are held to one suite. The cost is bounded — at most `limit` * runs times the run model's cap of 1024 attempts — but it is a real cost on a wide page. * + * ## A run that ends announces itself once + * + * [transition] hands an [ExtractionRunTransitioned] to [listener] for the call that ended the run, + * and for no other call: a replay, a rejected write, a [save] and a [recordInvocation] all announce + * nothing. Exactly one call per run reaches the applied branch, because reaching it means having + * created the run's terminal-write node, and the uniqueness constraint lets one transaction do that. + * + * The announcement waits for the write to be durable. When this store owns the transaction, that + * point is where the template returns, and the listener runs there. When a caller's transaction is + * active the write is durable when that caller commits, so the announcement is registered against + * the commit and never happens at all if the caller rolls back — a listener told a run ended by a + * transaction that was thrown away would be reporting a run nothing can read. + * * @param persistenceManager Drivine's handle on the `neo` datasource. * @param transactionManager used to own a transaction when no caller has one, so a lost * compare-and-set race can be recovered in a transaction the race did not already poison. * @param clock supplies the instant a terminal write is recorded at. Informational, and injectable * so a test can pin it; nothing sorts or compares on it. + * @param listener Notified when a run ends. Defaults to [DiceEventListener.DEV_NULL], so a host + * that has nothing listening constructs the store the way it always did. Handlers run inline on + * whichever thread the announcement happens on, and throw isolation belongs to the listener — + * wrap it in `SafeDiceEventListener` for graceful degradation. */ @Transactional -open class DrivineExtractionRunStore( +open class DrivineExtractionRunStore @JvmOverloads constructor( private val persistenceManager: PersistenceManager, transactionManager: PlatformTransactionManager, private val clock: Clock = Clock.systemUTC(), + private val listener: DiceEventListener = DiceEventListener.DEV_NULL, ) : ExtractionRunStore { private val logger = LoggerFactory.getLogger(DrivineExtractionRunStore::class.java) @@ -328,7 +349,7 @@ open class DrivineExtractionRunStore( ): ExtractionRunTransitionResult { ExtractionRunSchema.requireStorableTenant(key.contextId) if (TransactionSynchronizationManager.isActualTransactionActive()) { - return endRun(key, transition) + return announce(endRun(key, transition)) } var lostRace: RuntimeException? = null @@ -339,7 +360,7 @@ open class DrivineExtractionRunStore( lostRace = e null } - if (result != null) return result + if (result != null) return announce(result) logger.debug( "Terminal write for run {} in context {} lost the race; re-reading the recorded write", @@ -428,6 +449,8 @@ open class DrivineExtractionRunStore( violation: RuntimeException, ): ExtractionRunTransitionResult { val row = singleRow(TERMINAL_WRITE_BY_KEY, keyBindings(key)) ?: throw violation + // Nothing announced here: the writer that lost the race created no terminal-write node, so + // this path answers replayed or conflict and never applied. return replayOrConflict( key, transition, @@ -436,6 +459,33 @@ open class DrivineExtractionRunStore( ) } + /** + * Tells [listener] about a run this call ended, once the write is durable, and returns [result] + * so a caller can announce and return in one line. + * + * A replay and a rejected write reach neither branch below — a rejected write throws before + * getting here, and a replay is not applied. Inside a caller's transaction the write becomes + * durable at that caller's commit, so the announcement rides on it and is dropped with the + * transaction if the caller rolls back. Otherwise the commit has already happened by the time + * this runs, and the listener is called on the spot. + */ + private fun announce(result: ExtractionRunTransitionResult): ExtractionRunTransitionResult { + if (!result.isApplied) return result + val event = ExtractionRunTransitioned(result.run) + if (TransactionSynchronizationManager.isSynchronizationActive()) { + TransactionSynchronizationManager.registerSynchronization( + object : TransactionSynchronization { + override fun afterCommit() { + listener.onEvent(event) + } + }, + ) + return result + } + listener.onEvent(event) + return result + } + // ---- reads ---- @Transactional(readOnly = true) diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt index bf6027b3..66a53db2 100644 --- a/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt @@ -23,6 +23,9 @@ import com.embabel.dice.proposition.extraction.ExtractionDeploymentRef import com.embabel.dice.proposition.extraction.ExtractionExperimentRef import com.embabel.dice.proposition.extraction.ExtractionFailure import com.embabel.dice.proposition.extraction.ExtractionFailureCode +import com.embabel.dice.proposition.extraction.ExtractionFailureMeasure +import com.embabel.dice.proposition.extraction.ExtractionFailureQuantity +import com.embabel.dice.proposition.extraction.ExtractionFailureStage import com.embabel.dice.proposition.extraction.ExtractionInvocationId import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord @@ -388,39 +391,74 @@ private fun deserializeSourceRevisions(serialized: String): List): String = objectMapper.writeValueAsString( failures.map { failure -> linkedMapOf( "code" to failure.code.name, - "detail" to failure.detail, + "stage" to failure.stage?.name, + "providerStatus" to failure.providerStatus, + "measureQuantity" to failure.measure?.quantity?.name, + "measureValue" to failure.measure?.value, "at" to failure.at.toString(), - "invocation" to failure.invocation?.let { - linkedMapOf("invocationIndex" to it.invocationIndex, "attempt" to it.attempt) - }, + "invocationIndex" to failure.invocation?.invocationIndex, + "attempt" to failure.invocation?.attempt, ) }, ) private fun deserializeFailures(serialized: String): List = jsonObjects(serialized, "failures").map { fields -> + val quantity = fields.jsonStrOrNull("measureQuantity") + ?.let { jsonEnum(it, "failures", "measureQuantity") } + val value = fields.jsonLongOrNull("failures", "measureValue") + require((quantity == null) == (value == null)) { + "an entry of the stored 'failures' holds half a measure: a quantity and a value are " + + "written together or not at all" + } + + val invocationIndex = fields.jsonIntOrNull("failures", "invocationIndex") + val attempt = fields.jsonIntOrNull("failures", "attempt") + require((invocationIndex == null) == (attempt == null)) { + "an entry of the stored 'failures' holds half an invocation id: an index and an " + + "attempt are written together or not at all" + } + ExtractionFailure( code = jsonEnum(fields.jsonStr("failures", "code"), "failures", "code"), - detail = fields.jsonStr("failures", "detail"), + stage = fields.jsonStrOrNull("stage") + ?.let { jsonEnum(it, "failures", "stage") }, + providerStatus = fields.jsonIntOrNull("failures", "providerStatus"), + measure = quantity?.let { ExtractionFailureMeasure(it, value!!) }, at = Instant.parse(fields.jsonStr("failures", "at")), - invocation = (fields["invocation"] as? Map<*, *>)?.let { - ExtractionInvocationId( - invocationIndex = it.jsonInt("failures", "invocationIndex"), - attempt = it.jsonInt("failures", "attempt"), - ) - }, + invocation = invocationIndex?.let { ExtractionInvocationId(it, attempt!!) }, ) } @@ -599,11 +637,6 @@ private fun Map<*, *>.jsonStr(field: String, name: String): String = private fun Map<*, *>.jsonStrOrNull(name: String): String? = this[name]?.toString() -private fun Map<*, *>.jsonInt(field: String, name: String): Int = - jsonIntOrNull(field, name) ?: throw IllegalArgumentException( - "an entry of the stored '$field' is missing its '$name'" - ) - private fun Map<*, *>.jsonIntOrNull(field: String, name: String): Int? = when (val value = this[name]) { null -> null is Number -> value.toInt() @@ -612,6 +645,15 @@ private fun Map<*, *>.jsonIntOrNull(field: String, name: String): Int? = when (v ) } +/** Reads an optional whole number wide enough for a measure, which counts bytes and milliseconds. */ +private fun Map<*, *>.jsonLongOrNull(field: String, name: String): Long? = when (val value = this[name]) { + null -> null + is Number -> value.toLong() + else -> throw IllegalArgumentException( + "the '$name' of the stored '$field' is a ${value.javaClass.simpleName} where a number was expected" + ) +} + private fun Map<*, *>.jsonDoubleOrNull(field: String, name: String): Double? = when (val value = this[name]) { null -> null is Number -> value.toDouble() diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt index 2de39952..9cf37bc0 100644 --- a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt @@ -15,6 +15,7 @@ */ package com.embabel.dice.storage +import com.embabel.dice.common.DiceEventListener import com.embabel.dice.proposition.extraction.ExtractionRunStore import org.drivine.manager.PersistenceManager import org.drivine.query.QuerySpecification @@ -33,6 +34,10 @@ import org.springframework.boot.test.context.SpringBootTest * [DrivinePropositionStoreContractIntegrationTest] uses. The suite calls [store] more than once * inside a few cases and expects a store holding nothing for its tenants each time; it gets that * because those cases mint a distinct run id per iteration and only ever read by key. + * + * The suite also asks for a store announcing to a listener it just made, and this one store was + * built with its listener at context startup. [RedirectableEventListener] is what reconciles those: + * the bean holds one, and [store] aims it at whichever listener the case handed over. */ @SpringBootTest(classes = [TestApplication::class]) class DrivineExtractionRunStoreContractIntegrationTest : AbstractExtractionRunStoreContractTest() { @@ -40,13 +45,20 @@ class DrivineExtractionRunStoreContractIntegrationTest : AbstractExtractionRunSt @Autowired private lateinit var graphStore: DrivineExtractionRunStore + @Autowired + private lateinit var eventListener: RedirectableEventListener + @Autowired private lateinit var persistenceManager: PersistenceManager - override fun store(): ExtractionRunStore = graphStore + override fun store(listener: DiceEventListener): ExtractionRunStore { + eventListener.redirectTo(listener) + return graphStore + } @AfterEach fun cleanUp() { + eventListener.redirectTo(DiceEventListener.DEV_NULL) ExtractionRunSchema.LABELS.forEach { label -> persistenceManager.execute(QuerySpecification.withStatement("MATCH (n:$label) DETACH DELETE n")) } diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt index e16ac99d..6ddbe3b7 100644 --- a/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt @@ -16,6 +16,8 @@ package com.embabel.dice.storage import com.embabel.agent.core.ContextId +import com.embabel.dice.common.DiceEvent +import com.embabel.dice.common.DiceEventListener import com.embabel.dice.proposition.extraction.ExtractionActorRef import com.embabel.dice.proposition.extraction.ExtractionCohortRef import com.embabel.dice.proposition.extraction.ExtractionContentProfileRef @@ -23,6 +25,9 @@ import com.embabel.dice.proposition.extraction.ExtractionDeploymentRef import com.embabel.dice.proposition.extraction.ExtractionExperimentRef import com.embabel.dice.proposition.extraction.ExtractionFailure import com.embabel.dice.proposition.extraction.ExtractionFailureCode +import com.embabel.dice.proposition.extraction.ExtractionFailureMeasure +import com.embabel.dice.proposition.extraction.ExtractionFailureQuantity +import com.embabel.dice.proposition.extraction.ExtractionFailureStage import com.embabel.dice.proposition.extraction.ExtractionInvocationId import com.embabel.dice.proposition.extraction.ExtractionInvocationOutcome import com.embabel.dice.proposition.extraction.ExtractionInvocationRecord @@ -46,6 +51,7 @@ import com.embabel.dice.proposition.extraction.ExtractionRunTransitionOutcome import com.embabel.dice.proposition.extraction.ExtractionRuntimeIdentity import com.embabel.dice.proposition.extraction.ExtractionSessionRef import com.embabel.dice.provenance.SourceRevisionRef +import com.fasterxml.jackson.databind.ObjectMapper import org.drivine.manager.PersistenceManager import org.drivine.query.QuerySpecification import org.junit.jupiter.api.AfterEach @@ -93,6 +99,9 @@ class DrivineExtractionRunStoreIntegrationTest { @Autowired private lateinit var transactionManager: PlatformTransactionManager + @Autowired + private lateinit var eventListener: RedirectableEventListener + private val tenant = ContextId("graph-tenant") private val neighbour = ContextId("graph-neighbour") private val startedAt: Instant = Instant.parse("2026-08-31T10:15:30.123456789Z") @@ -100,6 +109,7 @@ class DrivineExtractionRunStoreIntegrationTest { @AfterEach fun cleanUp() { + eventListener.redirectTo(DiceEventListener.DEV_NULL) ExtractionRunSchema.LABELS.forEach { label -> persistenceManager.execute(QuerySpecification.withStatement("MATCH (n:$label) DETACH DELETE n")) } @@ -454,7 +464,15 @@ class DrivineExtractionRunStoreIntegrationTest { val transition = ExtractionRunTransition.completed( finishedAt = finishedAt, counts = ExtractionRunCounts(propositionsPersisted = 99), - failures = listOf(ExtractionFailure(ExtractionFailureCode.RATE_LIMITED, "slow down", finishedAt)), + failures = listOf( + ExtractionFailure( + code = ExtractionFailureCode.RATE_LIMITED, + stage = ExtractionFailureStage.MODEL_CALL, + providerStatus = 429, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.RETRY_AFTER_SECONDS, 30), + at = finishedAt, + ), + ), ) val ended = store.transition(run.key(), transition) @@ -547,8 +565,7 @@ class DrivineExtractionRunStoreIntegrationTest { fun `nothing a run was never given reaches the stored bytes`() { // The privacy contract, asserted over the properties actually written rather than over a // test-local dumper. Everything DICE holds about a run is a token, a digest, or a - // classified code, so a secret the host never handed over cannot be in the row — and the - // one place it could leak in is a failure built from an exception whose message quotes it. + // classified code, so a secret the host never handed over cannot be in the row. val secret = "patient-9f2a-had-a-consultation-on-tuesday" val run = ExtractionRun( contextId = tenant, @@ -560,10 +577,11 @@ class DrivineExtractionRunStoreIntegrationTest { session = ExtractionSessionRef("session-token-1"), ), failures = listOf( - ExtractionFailure.fromThrowable( - ExtractionFailureCode.DECODE_FAILED, - IllegalStateException("could not decode: $secret"), - startedAt, + ExtractionFailure( + code = ExtractionFailureCode.DECODE_FAILED, + stage = ExtractionFailureStage.RESPONSE_DECODE, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.CHARACTER_COUNT, 4096), + at = startedAt, ), ), ) @@ -572,7 +590,105 @@ class DrivineExtractionRunStoreIntegrationTest { val written = nodeProperties("ExtractionRun", "graph-privacy").values.joinToString(" ") { it.toString() } assertTrue(secret !in written, "a source-text substring reached the stored row") - assertTrue("IllegalStateException" in written, "the classified cause chain is what is stored instead") + } + + @Test + fun `a terminal write its caller rolls back announces nothing and leaves the run running`() { + // The half of exactly-once the contract suite cannot reach. Every case there runs with no + // ambient transaction, so the store owns the commit and announces inline. Here the caller + // owns it, the write is durable only if that caller commits, and this one does not. A store + // announcing at the point of the write would have told a listener a run ended, and the run + // is still running. + val received = mutableListOf() + eventListener.redirectTo(DiceEventListener { event -> synchronized(received) { received += event } }) + val run = running("graph-announce-rollback") + store.save(run) + + val ambient = TransactionTemplate(transactionManager) + ambient.execute { status -> + val ended = store.transition(run.key(), ExtractionRunTransition.completed(finishedAt)) + assertEquals(ExtractionRunTransitionOutcome.APPLIED, ended.outcome) + status.setRollbackOnly() + } + + assertEquals( + emptyList(), + synchronized(received) { received.toList() }, + "the caller rolled back, so no run ended and nothing may be announced", + ) + assertEquals( + ExtractionRunStatus.RUNNING, + store.findRun(run.key())?.status, + "the rolled-back terminal write left the run where it was", + ) + assertNull(store.findRun(run.key())?.finishedAt) + assertEquals(0, terminalWriteCount(run.key()), "the terminal-write node rolled back with it") + } + + @Test + fun `a stored failure holds the closed vocabulary's fields and nothing else`() { + // The route a provider message used to take into a durable row was a failure record's text + // field, and the record has none left. This is what holds the store to that: the keys the + // writer actually put on the node must match the allowlist below exactly, so a + // detail column, a message, a response body — anything shaped to hold what a person typed or + // a provider said — fails here and never reaches Neo4j. The allowlist is written out in + // full on purpose. Reading it from the writer would let one edit move both sides at once. + val allowed = setOf( + "code", + "stage", + "providerStatus", + "measureQuantity", + "measureValue", + "at", + "invocationIndex", + "attempt", + ) + val failures = listOf( + ExtractionFailure( + code = ExtractionFailureCode.RATE_LIMITED, + stage = ExtractionFailureStage.MODEL_CALL, + providerStatus = 429, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.RETRY_AFTER_SECONDS, 30), + at = startedAt.plusSeconds(5), + ), + // Every optional field absent. The stored keys have to be the same set, or the exact + // comparison above would only be an upper bound on what a fully populated one carries. + ExtractionFailure(code = ExtractionFailureCode.INTERNAL, at = startedAt.plusSeconds(6)), + ) + val run = ExtractionRun( + contextId = tenant, + lineage = ExtractionRunLineage.root(ExtractionRunRef("graph-failure-fields")), + status = ExtractionRunStatus.RUNNING, + startedAt = startedAt, + failures = failures, + ) + store.save(run) + + assertStoredFailureKeys("graph-failure-fields", 2, allowed, "the save door") + // The allowlist is over the encoding the store really uses, so the same run has to read back + // whole — an allowlist held by a writer nothing round-trips would prove nothing. + assertEquals(failures, store.findRun(run.key())!!.failures) + + // The other door onto the same property. `transition` builds its own bind map and writes + // failures through `SET n += $terminal`, so a text field could be added there alone and the + // save-door assertion above would stay green. Both doors, one allowlist. + val terminalFailures = listOf( + ExtractionFailure( + code = ExtractionFailureCode.MODEL_TIMEOUT, + stage = ExtractionFailureStage.MODEL_CALL, + providerStatus = 504, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.ELAPSED_MILLIS, 90_000), + at = finishedAt, + ), + ) + val ended = store.transition( + run.key(), + ExtractionRunTransition.completed(finishedAt, failures = terminalFailures), + ) + + assertEquals(ExtractionRunTransitionOutcome.APPLIED, ended.outcome) + assertStoredFailureKeys("graph-failure-fields", 1, allowed, "the transition door") + assertEquals(terminalFailures, store.findRun(run.key())!!.failures) } // ---- corrupt and oversized ---- @@ -675,7 +791,9 @@ class DrivineExtractionRunStoreIntegrationTest { failures: List = listOf( ExtractionFailure( code = ExtractionFailureCode.MODEL_TIMEOUT, - detail = "timed out after 90s", + stage = ExtractionFailureStage.MODEL_CALL, + providerStatus = 504, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.ELAPSED_MILLIS, 90_000), at = startedAt.plusSeconds(90), invocation = ExtractionInvocationId(0, 1), ), @@ -776,6 +894,32 @@ class DrivineExtractionRunStoreIntegrationTest { ).single()["fingerprint"]?.toString() /** Every property actually written on one node, read back raw rather than through the mapper. */ + /** + * Reads the `failures` property Neo4j actually holds for [runId] and checks every stored + * failure's key set matches [allowed] exactly. + * + * Reading the raw property is the whole point. Going through `findRun` would hand back + * [ExtractionFailure] objects, which cannot carry a stray field by construction, so a text + * column written to the node would be invisible to it. + */ + private fun assertStoredFailureKeys( + runId: String, + expectedCount: Int, + allowed: Set, + door: String, + ) { + val stored = nodeProperties("ExtractionRun", runId) + val decoded = ObjectMapper().readValue(stored["failures"].toString(), List::class.java) + assertEquals(expectedCount, decoded.size, "every failure written through $door reached the node") + decoded.forEach { element -> + assertEquals( + allowed, + (element as Map<*, *>).keys.map { it.toString() }.toSet(), + "a failure stored through $door carries exactly the fields the closed vocabulary names", + ) + } + } + private fun nodeProperties(label: String, runId: String): Map = rows( "MATCH (n:$label {contextId: \$contextId, runId: \$runId}) RETURN properties(n) AS row", mapOf("contextId" to tenant.value, "runId" to runId), diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt index 7d7256a7..edbd5577 100644 --- a/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/TestApplication.kt @@ -17,6 +17,8 @@ package com.embabel.dice.storage import com.embabel.common.ai.model.EmbeddingService import com.embabel.common.ai.model.PricingModel +import com.embabel.dice.common.DiceEvent +import com.embabel.dice.common.DiceEventListener import org.drivine.autoconfigure.EnableDrivine import org.drivine.autoconfigure.EnableDrivineTestConfig import org.drivine.manager.GraphObjectManager @@ -82,6 +84,29 @@ class PinnableClock : Clock() { override fun withZone(zone: ZoneId): Clock = this } +/** + * A listener the extraction-run store is built with once, pointing wherever a test aims it. + * + * The store is a Spring bean built when the context starts and shared by every case, and it is + * transaction-proxied, so a case cannot build its own to get a listener attached. The contract suite + * asks for a store announcing to a listener it just made, so this sits in between: the store is + * constructed with one of these, and a case redirects it at the listener it wants to watch. + */ +class RedirectableEventListener : DiceEventListener { + + @Volatile + private var target: DiceEventListener = DiceEventListener.DEV_NULL + + /** Sends everything from here on to [listener]. */ + fun redirectTo(listener: DiceEventListener) { + target = listener + } + + override fun onEvent(event: DiceEvent) { + target.onEvent(event) + } +} + /** * Test wiring: Drivine's test support spins a Neo4j testcontainer and transaction management; * we add the graph stores and a fake embedding service. [SchemaCatalog] beans are ensured on @@ -163,11 +188,19 @@ open class TestApplication { @Bean open fun extractionRunSchema(): SchemaCatalog = SchemaCatalog.of(ExtractionRunSchema.specs()) + @Bean + open fun extractionRunEventListener(): RedirectableEventListener = RedirectableEventListener() + @Bean open fun extractionRunStore( persistenceManager: PersistenceManager, transactionManager: PlatformTransactionManager, - ): DrivineExtractionRunStore = DrivineExtractionRunStore(persistenceManager, transactionManager) + eventListener: RedirectableEventListener, + ): DrivineExtractionRunStore = DrivineExtractionRunStore( + persistenceManager = persistenceManager, + transactionManager = transactionManager, + listener = eventListener, + ) @Bean open fun decayManager( diff --git a/docs/design/extraction-runs.md b/docs/design/extraction-runs.md index 7663daed..00ae9803 100644 --- a/docs/design/extraction-runs.md +++ b/docs/design/extraction-runs.md @@ -811,6 +811,49 @@ digest from the stored run would reject a correct retry whenever an attempt had between, because the run it derived from would have changed while the terminal write did not. The encoding has one implementation, in `dice`, and the graph holds its output. +The string the graph holds is the `xrun-terminal:v2` digest, which covers the terminal status and +the finish time and nothing else. The counts and failures a transition carries ride into the header +`SET` beside it and reach none of the hashed bytes, so a retry that agrees on status and finish time +replays whatever numbers it names and the run keeps what the first accepted terminal write +delivered. The store has no comparison of its own to keep in step with that: it stores one string +and compares one string. + +### Announcing a run that ended, once the write is durable + +`transition` hands an `ExtractionRunTransitioned` to the store's listener for the call that ended the +run. Exactly one call per run reaches that branch, and the reason is the same schema fact the +compare-and-set rests on: reaching it means having created the run's terminal-write node, and the +uniqueness constraint lets one transaction do that. A replay announces nothing, a rejected write +announces nothing, and the writer that lost the race announces nothing — it created no node, so it +answers replayed or conflict and never applied. + +Where the durable store differs from the in-memory reference is *when*. The reference announces +after the write has landed and outside its monitor. Here "landed" means committed, and which commit +that is depends on who owns the transaction. When the store owns it, the commit has already happened +by the time the template returns, and the listener runs there. When a caller's transaction is active +the write is durable only when that caller commits, so the announcement is registered against the +commit and is dropped with the transaction if the caller rolls back. A listener told that a run ended +by a transaction that was thrown away would be reporting a run nothing can read back. + +### Failures are stored in the vocabulary, and nothing else fits + +A run's failures are one JSON array on the header node, and each element carries eight fields: the +`code`, the `stage`, the `providerStatus`, the two halves of one measure, the `at`, and the two +halves of the attempt it names. Every optional field is always present and written as a null when +it has no value, so a bare failure and a fully populated one store the same keys. + +The shape is flat and the halves are deliberate. Putting the measure and the invocation id at the +same level as everything else means every key a stored failure can carry is visible at once, which is +what lets a test check the whole key set it finds matches an allowlist, in one assertion. +`DrivineExtractionRunStoreIntegrationTest` round-trips a run carrying failures, reads the properties +Neo4j actually holds, and asserts that set exactly — with the allowlist written out in the test +and never read from the writer, so one edit cannot move both sides. That is the check standing +between a durable row and a `detail` column: `ExtractionFailure` has no text-shaped field to write +from, and if someone adds one to the encoding, this fails before it reaches a graph. + +Reads are strict about the halves too. A stored failure holding a measure quantity with no value, or +an invocation index with no attempt, is refused; nothing here fills in a missing half. + ### A header write cannot touch a child row Invocation records are their own nodes on their own key, so `save` has no way to delete one. The @@ -864,45 +907,6 @@ named argument rejection before anything is written. Reads are not capped, becau tenant longer than the cap matches nothing by construction, which is the fail-closed answer and the only one a read could give. -## Protected-content references - -An extraction run holds no content. Where a host wants replay material — a prompt as it was -rendered, a source document as it was read — it keeps that material itself and gives DICE a -`ProtectedContentRef`: what kind of material it is, when it stops being available, and an opaque -handle for looking it up. - -The handle is a `ProtectedContentHandle`, an eighth member of the `ExtractionOpaqueRef` family, so -it inherits the enforced floor the other seven have — bounded, character-restricted, truncating -`toString`, and a validation message that never quotes the value it rejected. That matters more here -than elsewhere: a host holding material in object storage has a link to hand and no reason to think -DICE minds, and the character set rejects one outright. - -What a host owes a reference it mints: - -- **The handle is not dereferenceable.** It names material in the host's own store and carries no - way to reach it: no signed URL, no pre-authenticated link, no bearer token, no decryption key, no - filename on a share. Holding one grants nothing. The rule is broader than the character check, and - the part the check cannot see is the host's to keep. -- **Expiry is the host's to honour.** `expiresAt` is a declaration about the host's store. DICE acts - on none of it — it does not sweep, does not delete, and does not refuse to hand back a reference - whose expiry has passed. `isExpiredAt` is the whole of what DICE offers. Keeping the expired - reference is deliberate: a run that was replayable until March and is not replayable now is a fact - about the run, and deleting the reference would make it look as though it never had replay - material. -- **The classification is a claim.** DICE stores it and carries it into the audit surface, and - verifies none of it. - -`ProtectedContentClassification` runs `INTERNAL`, `CONFIDENTIAL`, `PERSONAL`, least to most -restrictive. There is no "unknown" value: a host that cannot classify material classifies it -`PERSONAL`, because guessing low means personal data treated as operational logging and guessing -high means something operational looked after more carefully than it needed to be. - -Blob storage, key management and retention enforcement are outside #67 by that issue's own text. -This is the reference and its contract; there is no DICE-side store behind it. Nothing attaches a -reference to an `ExtractionRun` yet either, and `ExtractionReplayFidelity` still tops out at -`APPROXIMATE` — a run that carries replay material is a later slice's change to the run header, and -a test asserts the run holds no such field today. - ## OpenTelemetry GenAI naming, not adopted OTel's GenAI semantic conventions cover the same ground — `gen_ai.request.*`, `gen_ai.response.*`, From 6cfc6fb84b96de431ac49d216b7206631860b657 Mon Sep 17 00:00:00 2001 From: James Dunnam <7660553+jimador@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:22:37 -0400 Subject: [PATCH 3/4] Recognise a Neo4j constraint violation by its status code Every org.neo4j.driver.exceptions.Neo4jException carries a code() straight from the server, and a uniqueness constraint violation always carries Neo.ClientError.Schema.ConstraintValidationFailed. That code holds across driver versions the way an exception message never can, and Drivine does not translate driver exceptions into Spring's DataAccessException hierarchy, so there is no portable Spring type to catch here either. Add Neo4jErrors, an internal object in dice-storage with one function, isUniquenessViolation, that walks a cause chain (cycle-safe) and answers true when any cause is a Neo4jException carrying that code. Neo4jErrorsTest covers a ClientException carrying the code, a wrapping exception whose cause carries it, a ClientException with a different code, a message that merely says already exists, and a self-referential cause chain. dice-storage now declares org.neo4j.driver:neo4j-java-driver directly so the exception class is visible at compile time. Drivine already runs on this driver at runtime, and embabel-agent-dependencies manages its version at 6.1.0. Signed-off-by: James Davis Signed-off-by: James Dunnam <7660553+jimador@users.noreply.github.com> --- dice-storage/pom.xml | 9 +++ .../com/embabel/dice/storage/Neo4jErrors.kt | 53 +++++++++++++++ .../embabel/dice/storage/Neo4jErrorsTest.kt | 68 +++++++++++++++++++ 3 files changed, 130 insertions(+) create mode 100644 dice-storage/src/main/kotlin/com/embabel/dice/storage/Neo4jErrors.kt create mode 100644 dice-storage/src/test/kotlin/com/embabel/dice/storage/Neo4jErrorsTest.kt diff --git a/dice-storage/pom.xml b/dice-storage/pom.xml index b029ad5e..e15e68e4 100644 --- a/dice-storage/pom.xml +++ b/dice-storage/pom.xml @@ -60,6 +60,15 @@ spring-tx + + + org.neo4j.driver + neo4j-java-driver + + org.jetbrains.kotlin kotlin-stdlib diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/Neo4jErrors.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/Neo4jErrors.kt new file mode 100644 index 00000000..1c7ad94c --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/Neo4jErrors.kt @@ -0,0 +1,53 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import org.neo4j.driver.exceptions.Neo4jException + +/** + * Tells a Neo4j uniqueness constraint violation apart from every other kind of failure, by the + * driver's own status code. + * + * Every [Neo4jException] carries a `code()` straight from the server, and a uniqueness constraint + * violation always carries the same one: [CONSTRAINT_VALIDATION_FAILED]. That code is a contract + * the server publishes and the driver just forwards, so it holds across driver versions in a way + * an exception message never can, since a message is free text a future release can reword without + * warning. Drivine does not translate driver exceptions into Spring's `DataAccessException` + * hierarchy, so there is no Spring type to catch here either. This code check is what both stores + * in this module use to tell "someone else already wrote the same thing" apart from a real failure. + */ +internal object Neo4jErrors { + + private const val CONSTRAINT_VALIDATION_FAILED = "Neo.ClientError.Schema.ConstraintValidationFailed" + + /** + * Walks the cause chain looking for a [Neo4jException] carrying [CONSTRAINT_VALIDATION_FAILED]. + * + * Stops the moment it revisits a cause it has already seen, so a cause chain that loops back on + * itself cannot spin forever. + */ + fun isUniquenessViolation(error: Throwable?): Boolean { + var current: Throwable? = error + val seen = mutableSetOf() + while (current != null && seen.add(current)) { + if (current is Neo4jException && current.code() == CONSTRAINT_VALIDATION_FAILED) { + return true + } + current = current.cause + } + return false + } +} diff --git a/dice-storage/src/test/kotlin/com/embabel/dice/storage/Neo4jErrorsTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/Neo4jErrorsTest.kt new file mode 100644 index 00000000..ca6c9cf8 --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/Neo4jErrorsTest.kt @@ -0,0 +1,68 @@ +/* + * Copyright 2024-2026 Embabel Pty Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.embabel.dice.storage + +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.neo4j.driver.exceptions.ClientException + +/** + * Proves [Neo4jErrors.isUniquenessViolation] answers off the driver's status code, and only that + * code, wherever it turns up in a cause chain. + */ +class Neo4jErrorsTest { + + private val constraintViolationCode = "Neo.ClientError.Schema.ConstraintValidationFailed" + + @Test + fun `a ClientException carrying the constraint code answers true`() { + val error = ClientException(constraintViolationCode, "Node already exists with label ...") + + assertTrue(Neo4jErrors.isUniquenessViolation(error)) + } + + @Test + fun `a wrapping exception whose cause carries the code answers true`() { + val cause = ClientException(constraintViolationCode, "Node already exists with label ...") + val wrapper = RuntimeException("save failed", cause) + + assertTrue(Neo4jErrors.isUniquenessViolation(wrapper)) + } + + @Test + fun `a ClientException with a different code answers false`() { + val error = ClientException("Neo.ClientError.Statement.SyntaxError", "bad Cypher") + + assertFalse(Neo4jErrors.isUniquenessViolation(error)) + } + + @Test + fun `a message that merely says already exists answers false`() { + val error = RuntimeException("Node(0) already exists with label ExtractionRunTerminalWrite") + + assertFalse(Neo4jErrors.isUniquenessViolation(error)) + } + + @Test + fun `a self-referential cause chain terminates`() { + val first = RuntimeException("first") + val second = RuntimeException("second", first) + first.initCause(second) + + assertFalse(Neo4jErrors.isUniquenessViolation(first)) + } +} From 417a5f533c3fd1919cabc8cc661054869266b102 Mon Sep 17 00:00:00 2001 From: James Dunnam <7660553+jimador@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:23:02 -0400 Subject: [PATCH 4/4] Use the shared constraint-violation check in both stores DrivineExtractionRunStore.transition and DrivinePropositionRepository.saveInOwnedTransaction each carried their own private isUniquenessViolation, matching "ConstraintValidationFailed" or "already exists" in an exception's message. Both now call Neo4jErrors.isUniquenessViolation, dropping the private functions and the KDoc that said the two checks were duplicated because neither module had a home for a shared helper. CHANGELOG: one bullet under the DrivineExtractionRunStore entry noting the race detection keys on the driver's status code and the new pom dependency. Additive, no public signature changes. Signed-off-by: James Davis Signed-off-by: James Dunnam <7660553+jimador@users.noreply.github.com> --- CHANGELOG.md | 7 ++++ .../dice/storage/DrivineExtractionRunStore.kt | 40 ++++--------------- .../storage/DrivinePropositionRepository.kt | 22 ++-------- 3 files changed, 18 insertions(+), 51 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d0e42c8b..defbb0ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1713,3 +1713,10 @@ and the consumer PRs that deliver it). Reads are uncapped, because a read for a longer tenant matches nothing by construction. Every new type carries `@ApiStatus.Experimental` and the shapes may still move while the remaining #67 slices land. + **Race detection keys on the driver's status code now, and message text no longer matters.** + `Neo4jErrors.isUniquenessViolation` walks the cause chain for a `Neo4jException` whose `code()` is + `Neo.ClientError.Schema.ConstraintValidationFailed`, and both `DrivineExtractionRunStore` and + `DrivinePropositionRepository` call it to tell a lost compare-and-set race apart from a real + failure. `dice-storage` now declares the `neo4j-java-driver` dependency it already ran with + through Drivine, so the exception class it checks is visible at compile time too. + **Compatibility: additive.** No public signature changes. diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt index fada452b..3a4aac50 100644 --- a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt @@ -356,7 +356,9 @@ open class DrivineExtractionRunStore @JvmOverloads constructor( val result = try { txTemplate.execute { endRun(key, transition) } } catch (e: RuntimeException) { - if (!isUniquenessViolation(e)) throw e + // Neo4jErrors tells a lost compare-and-set race apart from a real failure by the + // driver's own status code, not by the exception's message. + if (!Neo4jErrors.isUniquenessViolation(e)) throw e lostRace = e null } @@ -974,41 +976,15 @@ open class DrivineExtractionRunStore @JvmOverloads constructor( } } -/** - * Best-effort detection of a Neo4j uniqueness-constraint violation anywhere in the cause chain. - * - * Matches on message substrings, because which form — the error code or the prose — reaches - * `getMessage()` is not guaranteed across driver versions. `DrivinePropositionRepository` carries the - * same check for the same reason; the two are not shared yet because neither module has a home for a - * Drivine error-mapping helper, and inventing one is a separate change. - * - * A false positive is bounded here: the caller only ever treats it as "someone else recorded the - * terminal write", and re-reads to find out. If there is no terminal write, the original exception is - * rethrown. - */ -private fun isUniquenessViolation(error: Throwable?): Boolean { - var current: Throwable? = error - val seen = mutableSetOf() - while (current != null && seen.add(current)) { - val message = current.message ?: "" - if (message.contains("ConstraintValidationFailed", ignoreCase = true) || - message.contains("already exists", ignoreCase = true) - ) { - return true - } - current = current.cause - } - return false -} - /** * Best-effort detection of Neo4j's own deadlock detector aborting a transaction, anywhere in the * cause chain — the failure [DrivineExtractionRunStore.retryingTransientConflict] retries. * - * Matches on message substrings for the same reason [isUniquenessViolation] does: which form of the - * message reaches `getMessage()` is not guaranteed across driver versions. Two writers `MERGE`-ing - * the same not-yet-existing key concurrently is the shape [DrivineExtractionRunStore.ownedTransaction] - * retries, and Neo4j names it + * Matches on message substrings because which form of the message reaches `getMessage()` is not + * guaranteed across driver versions, and there is no status code as stable as + * [Neo4jErrors.isUniquenessViolation] gets to use here: a deadlock is a driver-side retry decision, + * not a server status. Two writers `MERGE`-ing the same not-yet-existing key concurrently is the + * shape [DrivineExtractionRunStore.ownedTransaction] retries, and Neo4j names it * with the transaction lock manager's own vocabulary, a looser target than one stable exception * type: `TransientException` is the driver's own class name, and `Deadlock` and `can't acquire` are * the wording its lock manager uses to describe the same event. The check is deliberately loose to diff --git a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivinePropositionRepository.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivinePropositionRepository.kt index e9b61b3e..ab036ac0 100644 --- a/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivinePropositionRepository.kt +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivinePropositionRepository.kt @@ -220,7 +220,9 @@ class DrivinePropositionRepository( try { txTemplate.execute { findOrPersist(proposition, contextId, text) }!! } catch (e: RuntimeException) { - if (!isUniquenessViolation(e)) throw e + // Neo4jErrors tells this cross-instance race apart from a real failure by the + // driver's own status code, not by the exception's message. + if (!Neo4jErrors.isUniquenessViolation(e)) throw e // Cross-instance race: another writer inserted the same (contextId, text) and the DB // (contextId, text) uniqueness constraint rejected ours. The dupe now exists — reuse // it, in a transaction of its own so the failed attempt cannot roll this back. @@ -263,24 +265,6 @@ class DrivinePropositionRepository( } } - /** - * Best-effort detection of a Neo4j uniqueness-constraint violation anywhere in the cause chain. - * Matches on message substrings, since which form (error code vs. prose) shows up in - * `getMessage()` isn't guaranteed across driver versions. [findOrPersist] pre-checks for a - * same-text sibling before writing, so this is just the cross-instance-race backstop now. - */ - private fun isUniquenessViolation(error: Throwable?): Boolean { - var t: Throwable? = error - while (t != null) { - val msg = t.message ?: "" - if (msg.contains("ConstraintValidationFailed", ignoreCase = true) || - msg.contains("already exists", ignoreCase = true) - ) return true - t = t.cause - } - return false - } - /** * Re-find the winner after a cross-instance uniqueness race and union the losing writer's * evidence into it. The relationship MERGE is keyed by evidence identity, so a retry is a no-op.