diff --git a/CHANGELOG.md b/CHANGELOG.md index 6614123a..defbb0ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1624,3 +1624,99 @@ 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. 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 + 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. + **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/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/DrivineExtractionRunStore.kt b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt new file mode 100644 index 00000000..3a4aac50 --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/DrivineExtractionRunStore.kt @@ -0,0 +1,1007 @@ +/* + * 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.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 +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.TransactionSynchronization +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. + * + * ## 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 @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) + + 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 announce(endRun(key, transition)) + } + + var lostRace: RuntimeException? = null + val result = try { + txTemplate.execute { endRun(key, transition) } + } catch (e: RuntimeException) { + // 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 + } + if (result != null) return announce(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 + // 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, + row["fingerprint"]?.toString(), + row["status"]?.toString(), + ) + } + + /** + * 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) + 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 Neo4j's own deadlock detector aborting a transaction, anywhere in the + * cause chain — the failure [DrivineExtractionRunStore.retryingTransientConflict] retries. + * + * 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 + * 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/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. 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..66a53db2 --- /dev/null +++ b/dice-storage/src/main/kotlin/com/embabel/dice/storage/ExtractionRunRowMappers.kt @@ -0,0 +1,668 @@ +/* + * 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.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 +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 holding eight fields: the + * classified code, the stage, the provider's status, the two halves of one measure, when it + * happened, and the two halves of the attempt it names. + * + * **Every field is a code, a number or a timestamp.** [ExtractionFailure] carries no `String` and no + * `Throwable`, so there is nothing free-text to write, and this writer names its fields one at a + * time, so a field added to the failure record cannot arrive here on its own. A text field could + * only get into a stored row through an edit to the list below, and + * `DrivineExtractionRunStoreIntegrationTest` round-trips a run carrying failures and compares the + * stored object's keys against an allowlist it states itself, so such an edit fails the build. + * + * The shape is flat. A failure's measure and the attempt it names are two fields each, flattened + * into the same object, which puts every key a stored failure can carry at one level where that + * exact-set comparison can see all of them at once. 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. Both halves of each pair are + * written together or not at all, and [deserializeFailures] refuses a stored object holding half of + * one. + * + * **A reader tolerates a key it does not recognise.** [deserializeFailures] asks the stored object + * for the fields it knows about and ignores anything else sitting beside them, so a node written by + * a later build, or corrupted by something outside DICE, still reads back as the fields this build + * understands. Refusing a whole run's failure list over one unexpected key would make an audit + * record unreadable at the moment someone needs to read it. Keeping DICE's own writer from putting a + * stray key there is the allowlist test's job; this is what a reader does when it meets one anyway, + * and the two are separate concerns. + */ +private fun serializeFailures(failures: List): String = + objectMapper.writeValueAsString( + failures.map { failure -> + linkedMapOf( + "code" to failure.code.name, + "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(), + "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"), + 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 = invocationIndex?.let { ExtractionInvocationId(it, 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<*, *>.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" + ) +} + +/** 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() + 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/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/DrivineExtractionRunStoreContractIntegrationTest.kt b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt new file mode 100644 index 00000000..9cf37bc0 --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreContractIntegrationTest.kt @@ -0,0 +1,66 @@ +/* + * 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.common.DiceEventListener +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. + * + * 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() { + + @Autowired + private lateinit var graphStore: DrivineExtractionRunStore + + @Autowired + private lateinit var eventListener: RedirectableEventListener + + @Autowired + private lateinit var persistenceManager: PersistenceManager + + 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 new file mode 100644 index 00000000..6ddbe3b7 --- /dev/null +++ b/dice-storage/src/test/kotlin/com/embabel/dice/storage/DrivineExtractionRunStoreIntegrationTest.kt @@ -0,0 +1,933 @@ +/* + * 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.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 +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 +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 com.fasterxml.jackson.databind.ObjectMapper +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 + + @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") + private val finishedAt: Instant = Instant.parse("2026-08-31T10:16:00Z") + + @AfterEach + fun cleanUp() { + eventListener.redirectTo(DiceEventListener.DEV_NULL) + 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( + 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) + + 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. + 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( + code = ExtractionFailureCode.DECODE_FAILED, + stage = ExtractionFailureStage.RESPONSE_DECODE, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.CHARACTER_COUNT, 4096), + at = 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") + } + + @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 ---- + + @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, + stage = ExtractionFailureStage.MODEL_CALL, + providerStatus = 504, + measure = ExtractionFailureMeasure(ExtractionFailureQuantity.ELAPSED_MILLIS, 90_000), + 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. */ + /** + * 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), + ).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/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)) + } +} 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..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 @@ -160,6 +185,23 @@ open class TestApplication { persistenceManager: PersistenceManager, ): DrivineCollectorTraceStore = DrivineCollectorTraceStore(persistenceManager) + @Bean + open fun extractionRunSchema(): SchemaCatalog = SchemaCatalog.of(ExtractionRunSchema.specs()) + + @Bean + open fun extractionRunEventListener(): RedirectableEventListener = RedirectableEventListener() + + @Bean + open fun extractionRunStore( + persistenceManager: PersistenceManager, + transactionManager: PlatformTransactionManager, + eventListener: RedirectableEventListener, + ): DrivineExtractionRunStore = DrivineExtractionRunStore( + persistenceManager = persistenceManager, + transactionManager = transactionManager, + listener = eventListener, + ) + @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..00ae9803 100644 --- a/docs/design/extraction-runs.md +++ b/docs/design/extraction-runs.md @@ -750,6 +750,163 @@ 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. + +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 +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. + ## OpenTelemetry GenAI naming, not adopted OTel's GenAI semantic conventions cover the same ground — `gen_ai.request.*`, `gen_ai.response.*`, @@ -837,9 +994,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