From b62a608a0a4f93442e8f6af821e36dd788ce249f Mon Sep 17 00:00:00 2001 From: kl3inIT Date: Tue, 4 Aug 2026 23:21:03 +0700 Subject: [PATCH 1/5] docs: open the retrieval admission control increment (ADR 0020 Phase 1) Consolidates the two-model architecture debate into decision 0020 and opens the Phase 1 increment: fair 4-permit retrieval admission, API pool 8/8, topK reverted to the upstream default, and assistant-layer TTFT attribution. Co-Authored-By: Claude Opus 5 (1M context) --- ...0-authorized-multi-snapshot-query-plane.md | 105 ++++++++++++++++ .../design.md | 112 ++++++++++++++++++ .../plan.md | 50 ++++++++ docs/roadmap.md | 1 + 4 files changed, 268 insertions(+) create mode 100644 docs/decisions/0020-authorized-multi-snapshot-query-plane.md create mode 100644 docs/increments/active/2026-08-05-retrieval-admission-phase1/design.md create mode 100644 docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md diff --git a/docs/decisions/0020-authorized-multi-snapshot-query-plane.md b/docs/decisions/0020-authorized-multi-snapshot-query-plane.md new file mode 100644 index 000000000..8f9a350f9 --- /dev/null +++ b/docs/decisions/0020-authorized-multi-snapshot-query-plane.md @@ -0,0 +1,105 @@ +# 0020 — Authorized Multi-Snapshot Query Plane As Gated Retrieval Target + +## Status + +Accepted on 2026-08-04 via the required independent architecture challenge, +run as a two-model debate (`orgmemory-agent-debate`): Claude Fable 5 defended +the per-space status quo with mechanical fixes, GPT-5.6-sol ultra defended a +single authorized query plane, and a no-tools judge committed from the record +alone. Debate record and verdict: `tmp/retrieval-perf-debate-2026-08-04-*.md` +(untracked reference space; the substance is consolidated here). + +## Context + +Production assistant TTFT averages 7.4 s, of which the retrieval pipeline is +nearly all: keyword planning 2.1 s (cache hit rate ~4%), embedding 1.2 s, and +per-knowledge-space snapshot queries executed in barrier-synchronized batches +of `maximumConcurrentSpaces=4` against a 12-connection pool. ~18% of turns +die at the 120 s turn timeout with a bimodal latency distribution consistent +with pool exhaustion (4 connections × 3 concurrent turns = the whole pool). +Per-space top-k results are merged without score normalization, over-retrieve +by construction (~433 contributions dropped per turn), and multi-space rerank +fails closed. + +Facts established during the challenge that changed the analysis: + +- Isolation is already predicate-based, not physically structural: vector + rows share one table and every query applies organization, authorized-asset, + and pinned-batch predicates before distance ordering + (`PostgresVectorIndex`). +- The `RETRIEVE` telemetry stage is inclusive (it wraps authorize, prepare, + embed, and snapshot execution), and ~2.4 s of TTFT is above the retrieval + service, unattributed by current telemetry. +- Embedding consumes the keyword plan (`LightRagQueryEngine.prepare`), so + keyword planning cannot simply be parallelized away. +- No atomic cross-space snapshot instant exists today; each space is pinned + sequentially. +- ADR 0013 already requires an explicit multi-snapshot contract before + cross-namespace retrieval is exposed. + +## Decision + +The target retrieval architecture is one typed `AuthorizedMultiSnapshotQuery` +store operation: the application issues one JDBC call whose immutable input +carries the organization, authorization model, a sorted vector of per-space +`(spaceId, batch, generation, manifest, ACL generation, authorized asset ids)` +tuples, and a canonical scope fingerprint; the store returns one globally +scored candidate band; adapter and core independently validate per-row +space/batch/asset membership and reject the entire result on any mismatch. +Per-space namespaces remain the publication and rollback unit for writes. +Existing scope re-resolution, OpenFGA BatchCheck, and canonical +evidence-closure recheck before egress are unchanged. + +Cutover is gated, and the gates are binding: + +1. **Phase 1 ships first, independent of the target**: Hikari fixed pool + 8/8 (down from 12/2, per the HikariCP sizing formula on the 4-vCPU shared + host); one fair host-wide retrieval semaphore of 4 permits acquired before + connection checkout (replicas divide the permits); batch barriers replaced + by continuous admission; `topK` reverted 60 → 40 (upstream LightRAG + default); retrieval-stage tracing added, including the unattributed ~2.4 s + above the retrieval service. +2. **Latency gate**: a shadow prototype of the compound query on identical + hardware, embeddings, corpus, authorization scopes, and query set, per the + production-hardening runbook (1/7/20 spaces; narrow and broad grants; + current/10×/100× projection sizes; concurrency 1 and 4 under + shared-Postgres load; ≥5 repetitions; `EXPLAIN (ANALYZE, BUFFERS)`). + Predeclared thresholds: compound-query p95 ≤ 500 ms; cold-keyword-miss + median TTFT ≈ 2 s. Failure means no cutover. +3. **Recall gate**: an evaluation set must show that the cache-miss keyword + bypass (raw-query embedding seeds plus deterministic lexical terms) does + not regress recall@40 before it becomes the default miss path. +4. **Isolation parity gate**: negative isolation, revocation-during-query, + stale-snapshot, and poisoned-cache tests pass; no unscoped read overload + exists anywhere in the new port. +5. **Cache discipline**: the cold path is authoritative; the composite + retrieval-result cache (sorted snapshot vector + authorization fingerprint + + query semantics + model route) is optional and is disabled rather than + weakened if its hit rate proves useless. + +Until every gate passes, the per-space structural path remains production and +is the rollback. + +## Rejected Alternative + +Keeping the per-space query plane as the terminal architecture (semaphore +admission, post-merge rerank, pool and topK tuning only). Its strongest +argument — the compound query's performance is unproven, and a mis-planned +multi-tuple join ahead of vector distance ordering can be slower than seven +indexed point queries — is preserved as gates 2–3 rather than as grounds to +retain a design whose cross-space ranking is unsound by construction +(incomparable per-space scores, per-space truncation before merge) and whose +~2 s TTFT target is unreachable while the keyword LLM call sits on a ~96% +miss path. + +## Consequences + +- Phase 1 changes production configuration immediately and is expected to + eliminate the pool-exhaustion timeout signature; if the bimodal timeouts + survive Phase 1, the pool hypothesis is falsified and tracing decides next. +- The compound storage port, multi-snapshot cache contract, shadow-compare + harness, and retrieval eval set are new work items and belong to a future + increment; this decision does not authorize skipping its design/plan cycle. +- The debate corrected earlier telemetry interpretation; stage dashboards + should distinguish inclusive from exclusive stage timers before further + latency conclusions are drawn. diff --git a/docs/increments/active/2026-08-05-retrieval-admission-phase1/design.md b/docs/increments/active/2026-08-05-retrieval-admission-phase1/design.md new file mode 100644 index 000000000..d0b30cb89 --- /dev/null +++ b/docs/increments/active/2026-08-05-retrieval-admission-phase1/design.md @@ -0,0 +1,112 @@ +# Retrieval admission control and pool right-sizing (ADR 0020 Phase 1) + +Date: 2026-08-05 + +## Outcome + +Ship Phase 1 of [ADR 0020](../../../decisions/0020-authorized-multi-snapshot-query-plane.md): +remove the connection-pool exhaustion failure mode behind the bimodal +turn-timeout signature, replace batch-barrier snapshot scheduling with fair +continuous admission, revert `topK` to the upstream LightRAG default, and make +the currently unattributed ~2.4 s of assistant TTFT traceable — all without +touching the per-space query plane, authorization sequence, or cache identity. + +The compound `AuthorizedMultiSnapshotQuery` target, its shadow prototype, and +the keyword-bypass recall evaluation are **out of scope**; they are Phase 2 +work gated by ADR 0020 conditions 2–3 and require their own increment. + +## Production evidence (Prometheus, 7-day window ending 2026-08-04) + +- 45 assistant turns; mean TTFT 7.4 s; 8/45 turns (~18%) die at the 120 s + `turnTimeout` with nothing observed between 15.7 s and 120 s. +- Per-request snapshot fan-out holds up to `maximumConcurrentSpaces = 4` + connections; the API pool is 12 (`application-prod.yml`); three concurrent + turns can exhaust the pool. The bimodal fast-or-timeout distribution is the + expected signature of admission starvation, not slow retrieval. +- Stage means: retrieve 4.7 s (inclusive timer — wraps authorize, prepare, + embed, snapshots), prepare_query 2.1 s at ~4% keyword-cache hit rate, + embed 1.2 s, retrieve_snapshot 299 ms × ~7 calls/turn. +- Pipeline accounts for ≈ 5.0 s of the 7.4 s TTFT; ~2.4 s sits above + `GraphRagKnowledgeRetrievalService` and is unattributed by current + telemetry. +- Host: 4 vCPU EPYC slice, SSD, PostgreSQL container shared with other + stacks. HikariCP About-Pool-Sizing formula: `(4 × 2) + 0 ≈ 8`. + +## Decisions + +All four were settled by the two-model architecture debate consolidated in +ADR 0020; this design binds them to code. + +### Fixed pool 8/8 + +`maximum-pool-size` 12 → 8 and `minimum-idle` 2 → 8 for the API service. +The HikariCP guidance is a small saturated fixed pool; the shared-host caveat +argues against rounding up. The worker service pool is unchanged — its +workload (extraction) is not part of this failure mode. + +### Host-wide fair retrieval semaphore, 4 permits, acquired before checkout + +One JVM-wide fair `Semaphore(4)` at the snapshot-execution boundary in +`GraphRagKnowledgeRetrievalService`. A permit is acquired **before** any JDBC +connection is requested and released after the snapshot result is +materialized. This caps retrieval's total pool draw at 4 of 8 connections, +leaving 4 for scope resolution, BatchCheck support queries, conversation +writes, and unrelated API traffic. Replicas would divide, never multiply, +this budget; the current deployment is a single API replica, and the permit +count is a property so a future replica count change is a config edit plus +this documented rule. + +Rejected variant (debate A-R1): pool 10 with a 6-permit semaphore — rejected +because it contradicts the sizing formula on a shared host and lets one +seven-space turn monopolize permits ahead of a second turn's first query. + +### Barrier to continuous admission + +The current loop executes spaces in batches of `maximumConcurrentSpaces` +with a `future.get()` barrier per batch; a straggler in batch one delays +batch two, and the second batch of a seven-space turn runs three-wide under a +four-slot budget. Replace with submit-all + global-semaphore admission: +every space task is submitted to the virtual-thread executor immediately and +blocks on the fair semaphore, preserving result ordering at consolidation. +`maximumConcurrentSpaces` stops governing scheduling; it is superseded by the +semaphore permits (validation retained so existing configuration does not +break). + +The fail-closed behavior is unchanged: any snapshot failure cancels +outstanding work and fails the whole retrieval, exactly as the barrier loop +does today. + +### topK 60 → 40 and TTFT attribution + +`ORGMEMORY_GRAPH_QUERY_TOP_K` default reverts to the upstream LightRAG +v1.5.4 default of 40 (the 60 has no recorded rationale; the debate found +none). Graph expansion ceilings follow as `topK * 4` = 160. + +Add stage timing above the retrieval service so the ~2.4 s gap becomes +attributable: emit assistant-layer durations for grounding-to-prompt +assembly, conversation-history loading, and the delay between retrieval +completion and the first model token. Payload-free attributes only, matching +the existing OpenTelemetry event-sink discipline. + +## Safety argument + +- No authorization boundary moves: the sequence documented in the + `2026-07-28-lightrag-query-latency` design (organization check → scoped + ListObjects → filtered GraphRAG → scope comparison → final BatchCheck → + canonical recheck) is untouched. +- Admission control is above the storage port; per-space statements, cache + keys, and snapshot pinning are byte-identical. +- The semaphore cannot deadlock the pool: each admitted task holds at most + one connection (`C_m = 1`), so the HikariCP deadlock floor is 1 and the + 4-permit budget is far above it. +- Risk: a fair semaphore serializes admission order under contention; at 45 + turns/week the contention window is small, and fairness is exactly what + removes the starvation mode. + +## Relationship to open increments + +`2026-07-28-lightrag-query-latency` is merged (PR #102) with only its live +before/after production proof pending. This increment builds on its +prepare-once and stage-telemetry work. The production verification step here +supersedes that increment's pending timing capture: one deployment, one +before/after measurement window serves both. diff --git a/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md b/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md new file mode 100644 index 000000000..bf4abfeb9 --- /dev/null +++ b/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md @@ -0,0 +1,50 @@ +# Retrieval admission control and pool right-sizing plan + +Design: [design.md](design.md). Decision: ADR 0020 conditions 1 and 7. + +## 1. Lock current behavior + +- Add a test that fails if a snapshot task requests a JDBC connection before + holding an admission permit (admission-before-checkout invariant). +- Add a test proving a snapshot failure under continuous admission cancels + outstanding tasks and fails the retrieval (fail-closed parity with the + barrier loop). +- Add a property test for the superseded-but-validated + `maximumConcurrentSpaces` configuration. + +## 2. Admission control + +- Introduce the JVM-wide fair 4-permit retrieval semaphore behind a + configuration property, acquired before connection checkout, released after + snapshot materialization. +- Replace the batched `future.get()` barrier in + `GraphRagKnowledgeRetrievalService` with submit-all + semaphore admission, + preserving consolidation order and cancellation semantics. + +## 3. Configuration + +- API Hikari `maximum-pool-size` 12 → 8, `minimum-idle` 2 → 8 in + `application-prod.yml`; leave the worker pool unchanged. +- `topK` default 60 → 40 in `GraphRagQueryRuntimeProperties` and + `application.yml`; confirm the `topK * 4` ceilings follow. + +## 4. TTFT attribution + +- Emit payload-free assistant-layer stage durations: grounding-to-prompt + assembly, history load, retrieval-completion-to-first-token gap. +- Extend the OpenTelemetry sink test for the new closed attribute set. + +## 5. Verify and release + +- Terminating `clean test` gate; focused core retrieval, API property, and + OpenTelemetry tests. +- Deploy to production; capture a before/after window in Prometheus: + turn-latency histogram (expect the 120 s bucket population to disappear), + TTFT mean, and the new assistant-layer stages accounting for the ~2.4 s + gap. This measurement also closes the pending live-proof gate of + `2026-07-28-lightrag-query-latency`. +- If the bimodal timeout signature survives the pool/admission change, record + that the pool hypothesis is falsified in this plan and open tracing of the + surviving path before any further latency work. +- Consolidate: spec/test matrix refresh for the retrieval domain, roadmap + update, move to completed. diff --git a/docs/roadmap.md b/docs/roadmap.md index d32fd30e4..e61737907 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -54,6 +54,7 @@ The table is a delivery index, not a second description of current behavior. | Increment | Status | Remaining gate | | --- | --- | --- | +| [Retrieval admission control and pool right-sizing](increments/active/2026-08-05-retrieval-admission-phase1/plan.md) | active | ADR 0020 Phase 1: fair 4-permit admission, pool 8/8, topK 40, TTFT attribution; production before/after window also closes the LightRAG-latency live proof | | [Assistant composer and conversation model picker](increments/completed/2026-08-04-assistant-composer-model-picker/plan.md) | shipped | delivered administrator-bound model authority, conversation selection, and composer polish | | [Assistant interaction foundation](increments/completed/2026-08-04-assistant-interaction-foundation/plan.md) | shipped | delivered server-owned starters, scoped drafts, answer feedback, fresh retry, and interaction recovery | | [Apache AGE published-batch backfill](increments/completed/2026-08-03-apache-age-published-batch-backfill/verification.md) | shipped | challenged one-shot repair, 49/49 production reconciliation, Graph explorer, and cited Assistant proof | From 0c5bb4ea92c3aa835a015de987795161c365aef5 Mon Sep 17 00:00:00 2001 From: kl3inIT Date: Tue, 4 Aug 2026 23:47:18 +0700 Subject: [PATCH 2/5] perf(retrieval): add fair snapshot admission --- .../GraphRagQueryRuntimeProperties.java | 8 +- .../src/main/resources/application-prod.yml | 4 +- apps/api/src/main/resources/application.yml | 3 +- .../GraphRagQueryRuntimePropertiesTests.java | 35 +++ ...aultGraphRagKnowledgeRetrievalService.java | 135 +++++++---- ...aphRagKnowledgeRetrievalConfiguration.java | 11 +- .../retrieval/GraphRagRetrievalPolicy.java | 5 +- .../retrieval/RetrievalAdmissionControl.java | 39 +++ ...raphRagKnowledgeRetrievalServiceTests.java | 222 +++++++++++++++++- 9 files changed, 413 insertions(+), 49 deletions(-) create mode 100644 core/src/main/java/com/orgmemory/core/knowledge/retrieval/RetrievalAdmissionControl.java diff --git a/apps/api/src/main/java/com/orgmemory/api/assistant/GraphRagQueryRuntimeProperties.java b/apps/api/src/main/java/com/orgmemory/api/assistant/GraphRagQueryRuntimeProperties.java index 0d2c2ed06..895750a92 100644 --- a/apps/api/src/main/java/com/orgmemory/api/assistant/GraphRagQueryRuntimeProperties.java +++ b/apps/api/src/main/java/com/orgmemory/api/assistant/GraphRagQueryRuntimeProperties.java @@ -12,6 +12,7 @@ record GraphRagQueryRuntimeProperties( Integer maximumEmbeddingBatchSize, Integer maximumKnowledgeSpaces, Integer maximumConcurrentSpaces, + Integer retrievalAdmissionPermits, Integer topK, Integer chunkTopK, Integer relatedChunkNumber, @@ -45,7 +46,11 @@ record GraphRagQueryRuntimeProperties( throw new IllegalArgumentException( "maximumConcurrentSpaces must not exceed maximumKnowledgeSpaces"); } - topK = positive(topK, 60, "topK"); + retrievalAdmissionPermits = positive( + retrievalAdmissionPermits, + 4, + "retrievalAdmissionPermits"); + topK = positive(topK, 40, "topK"); chunkTopK = positive(chunkTopK, 20, "chunkTopK"); relatedChunkNumber = positive( relatedChunkNumber, @@ -98,6 +103,7 @@ GraphRagRetrievalPolicy toPolicy() { return new GraphRagRetrievalPolicy( maximumKnowledgeSpaces, maximumConcurrentSpaces, + retrievalAdmissionPermits, topK, chunkTopK, relatedChunkNumber, diff --git a/apps/api/src/main/resources/application-prod.yml b/apps/api/src/main/resources/application-prod.yml index ddf4266c2..583576fb4 100644 --- a/apps/api/src/main/resources/application-prod.yml +++ b/apps/api/src/main/resources/application-prod.yml @@ -7,8 +7,8 @@ spring: password: ${ORGMEMORY_DB_PASSWORD} hikari: pool-name: orgmemory-api - maximum-pool-size: ${ORGMEMORY_API_DB_POOL_MAXIMUM_SIZE:12} - minimum-idle: ${ORGMEMORY_API_DB_POOL_MINIMUM_IDLE:2} + maximum-pool-size: ${ORGMEMORY_API_DB_POOL_MAXIMUM_SIZE:8} + minimum-idle: ${ORGMEMORY_API_DB_POOL_MINIMUM_IDLE:8} connection-timeout: ${ORGMEMORY_DB_CONNECTION_TIMEOUT_MS:10000} validation-timeout: ${ORGMEMORY_DB_VALIDATION_TIMEOUT_MS:5000} max-lifetime: ${ORGMEMORY_DB_MAX_LIFETIME_MS:1800000} diff --git a/apps/api/src/main/resources/application.yml b/apps/api/src/main/resources/application.yml index 165975f05..0c87b0d80 100644 --- a/apps/api/src/main/resources/application.yml +++ b/apps/api/src/main/resources/application.yml @@ -150,7 +150,8 @@ orgmemory: keyword-cache-ttl: ${ORGMEMORY_GRAPH_QUERY_KEYWORD_CACHE_TTL:24h} maximum-knowledge-spaces: ${ORGMEMORY_GRAPH_QUERY_MAXIMUM_SPACES:20} maximum-concurrent-spaces: ${ORGMEMORY_GRAPH_QUERY_MAXIMUM_CONCURRENT_SPACES:4} - top-k: ${ORGMEMORY_GRAPH_QUERY_TOP_K:60} + retrieval-admission-permits: ${ORGMEMORY_GRAPH_QUERY_ADMISSION_PERMITS:4} + top-k: ${ORGMEMORY_GRAPH_QUERY_TOP_K:40} chunk-top-k: ${ORGMEMORY_GRAPH_QUERY_CHUNK_TOP_K:20} related-chunk-number: ${ORGMEMORY_GRAPH_QUERY_RELATED_CHUNKS:5} maximum-graph-depth: ${ORGMEMORY_GRAPH_QUERY_MAXIMUM_DEPTH:1} diff --git a/apps/api/src/test/java/com/orgmemory/api/assistant/GraphRagQueryRuntimePropertiesTests.java b/apps/api/src/test/java/com/orgmemory/api/assistant/GraphRagQueryRuntimePropertiesTests.java index c52200b09..db860d279 100644 --- a/apps/api/src/test/java/com/orgmemory/api/assistant/GraphRagQueryRuntimePropertiesTests.java +++ b/apps/api/src/test/java/com/orgmemory/api/assistant/GraphRagQueryRuntimePropertiesTests.java @@ -23,9 +23,21 @@ void defaultsKeepRerankingDisabledAndBoundTheEvidenceClosure() { assertEquals("none", policy.rerank().provider()); assertEquals(2_000, policy.maximumEvidenceClosure()); assertEquals(4, policy.maximumConcurrentSpaces()); + assertEquals(4, policy.retrievalAdmissionPermits()); + assertEquals(40, policy.topK()); assertEquals(Duration.ofHours(24), properties.keywordCacheTtl()); } + @Test + void supersededMaximumConcurrentSpacesRemainsValidated() { + assertThrows( + IllegalArgumentException.class, + () -> propertiesWithConcurrency(0)); + assertThrows( + IllegalArgumentException.class, + () -> propertiesWithConcurrency(21)); + } + @Test void enabledRerankingRequiresAnExplicitProvider() { assertThrows( @@ -65,9 +77,32 @@ private static GraphRagQueryRuntimeProperties properties( null, null, null, + null, rerankEnabled, rerankProvider, minimumRerankScore, null); } + + private static GraphRagQueryRuntimeProperties propertiesWithConcurrency( + Integer maximumConcurrentSpaces) { + return new GraphRagQueryRuntimeProperties( + null, + null, + null, + 20, + maximumConcurrentSpaces, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null); + } } diff --git a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/DefaultGraphRagKnowledgeRetrievalService.java b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/DefaultGraphRagKnowledgeRetrievalService.java index 118035e28..600a5d42c 100644 --- a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/DefaultGraphRagKnowledgeRetrievalService.java +++ b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/DefaultGraphRagKnowledgeRetrievalService.java @@ -43,6 +43,7 @@ import java.util.Set; import java.util.UUID; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -87,6 +88,7 @@ class DefaultGraphRagKnowledgeRetrievalService private final KnowledgeRetrievalProperties retrievalProperties; private final GraphRagEventSink events; private final GraphRagTaskDecorator tasks; + private final RetrievalAdmissionControl admission; DefaultGraphRagKnowledgeRetrievalService( KnowledgeSearchAuthorizationService searchAuthorization, @@ -114,7 +116,9 @@ class DefaultGraphRagKnowledgeRetrievalService audit, retrievalProperties, events, - GraphRagTaskDecorator.NONE); + GraphRagTaskDecorator.NONE, + new RetrievalAdmissionControl( + policy.retrievalAdmissionPermits())); } DefaultGraphRagKnowledgeRetrievalService( @@ -131,6 +135,39 @@ class DefaultGraphRagKnowledgeRetrievalService KnowledgeRetrievalProperties retrievalProperties, GraphRagEventSink events, GraphRagTaskDecorator tasks) { + this( + searchAuthorization, + evidenceScopes, + authorization, + canonicalEvidence, + embeddingProfiles, + embedding, + publications, + engine, + policy, + audit, + retrievalProperties, + events, + tasks, + new RetrievalAdmissionControl( + policy.retrievalAdmissionPermits())); + } + + DefaultGraphRagKnowledgeRetrievalService( + KnowledgeSearchAuthorizationService searchAuthorization, + KnowledgeEvidenceScopeResolver evidenceScopes, + RelationshipAuthorizationSetPort authorization, + SecureKnowledgeRetrievalStore canonicalEvidence, + EmbeddingProfileRegistry embeddingProfiles, + KnowledgeEmbeddingProperties embedding, + ProjectionPublicationStore publications, + LightRagQueryEngine engine, + GraphRagRetrievalPolicy policy, + PermissionAuditService audit, + KnowledgeRetrievalProperties retrievalProperties, + GraphRagEventSink events, + GraphRagTaskDecorator tasks, + RetrievalAdmissionControl admission) { this.searchAuthorization = searchAuthorization; this.evidenceScopes = evidenceScopes; this.batchRecheck = new OpenFgaBatchRecheck(authorization); @@ -144,6 +181,7 @@ class DefaultGraphRagKnowledgeRetrievalService this.retrievalProperties = retrievalProperties; this.events = Objects.requireNonNull(events, "events"); this.tasks = Objects.requireNonNull(tasks, "tasks"); + this.admission = Objects.requireNonNull(admission, "admission"); } @Override @@ -499,45 +537,52 @@ private List queryPublishedSpaces( List groundings = new ArrayList<>(); try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) { - for (int offset = 0; - offset < requests.size(); - offset += policy.maximumConcurrentSpaces()) { - int end = Math.min( - requests.size(), - offset + policy.maximumConcurrentSpaces()); - List> futures = - requests.subList(offset, end) - .stream() - .map(request -> executor.submit(tasks.decorate(() -> - queryPublishedSpace(request, prepared)))) - .toList(); - try { - for (Future future : futures) { - SnapshotQueryResult snapshotResult = future.get(); - emitSnapshotStage( - operationId, - scope.organizationId(), - snapshotResult.duration(), - snapshotResult.inputCount(), - snapshotResult.result() - .grounding() - .chunks() - .size(), - snapshotResult.namespace()); - emitRerank( - operationId, - scope.organizationId(), - snapshotResult.result()); - LightRagGrounding grounding = - snapshotResult.result().grounding(); - if (!grounding.empty()) { - groundings.add(grounding); - } - } - } catch (ExecutionException | InterruptedException - | RuntimeException failure) { - futures.forEach(future -> future.cancel(true)); - throw retrievalFailure(failure); + var completed = new ExecutorCompletionService( + executor); + List> futures = + new ArrayList<>(requests.size()); + for (int index = 0; index < requests.size(); index++) { + int resultIndex = index; + LightRagQueryRequest request = requests.get(index); + futures.add(completed.submit(tasks.decorate(() -> + new IndexedSnapshotQueryResult( + resultIndex, + queryPublishedSpace(request, prepared))))); + } + SnapshotQueryResult[] ordered = + new SnapshotQueryResult[requests.size()]; + try { + for (int completedCount = 0; + completedCount < requests.size(); + completedCount++) { + IndexedSnapshotQueryResult result = + completed.take().get(); + ordered[result.index()] = result.result(); + } + } catch (ExecutionException | InterruptedException + | RuntimeException failure) { + futures.forEach(future -> future.cancel(true)); + throw retrievalFailure(failure); + } + for (SnapshotQueryResult snapshotResult : ordered) { + emitSnapshotStage( + operationId, + scope.organizationId(), + snapshotResult.duration(), + snapshotResult.inputCount(), + snapshotResult.result() + .grounding() + .chunks() + .size(), + snapshotResult.namespace()); + emitRerank( + operationId, + scope.organizationId(), + snapshotResult.result()); + LightRagGrounding grounding = + snapshotResult.result().grounding(); + if (!grounding.empty()) { + groundings.add(grounding); } } } @@ -546,10 +591,11 @@ private List queryPublishedSpaces( private SnapshotQueryResult queryPublishedSpace( LightRagQueryRequest request, - LightRagPreparedQuery prepared) { + LightRagPreparedQuery prepared) throws Exception { long startedAt = System.nanoTime(); LightRagQueryResult result = - engine.executePrepared(request, prepared); + admission.execute(() -> + engine.executePrepared(request, prepared)); return new SnapshotQueryResult( result, Duration.ofNanos(Math.max( @@ -559,6 +605,11 @@ private SnapshotQueryResult queryPublishedSpace( request.snapshot().namespace()); } + private record IndexedSnapshotQueryResult( + int index, + SnapshotQueryResult result) { + } + private static RuntimeException retrievalFailure(Exception failure) { if (failure instanceof InterruptedException) { Thread.currentThread().interrupt(); diff --git a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalConfiguration.java b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalConfiguration.java index 5dcf3eb37..36a0a7d3f 100644 --- a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalConfiguration.java +++ b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalConfiguration.java @@ -15,6 +15,13 @@ @ConditionalOnBean(LightRagQueryEngine.class) class GraphRagKnowledgeRetrievalConfiguration { + @Bean + RetrievalAdmissionControl graphRagRetrievalAdmissionControl( + GraphRagRetrievalPolicy policy) { + return new RetrievalAdmissionControl( + policy.retrievalAdmissionPermits()); + } + @Bean GraphRagKnowledgeRetrievalService graphRagKnowledgeRetrievalService( KnowledgeSearchAuthorizationService searchAuthorization, @@ -28,6 +35,7 @@ GraphRagKnowledgeRetrievalService graphRagKnowledgeRetrievalService( GraphRagRetrievalPolicy policy, PermissionAuditService audit, KnowledgeRetrievalProperties retrievalProperties, + RetrievalAdmissionControl admission, ObjectProvider eventSinks, ObjectProvider taskDecorators) { return new DefaultGraphRagKnowledgeRetrievalService( @@ -44,6 +52,7 @@ GraphRagKnowledgeRetrievalService graphRagKnowledgeRetrievalService( retrievalProperties, GraphRagEventSink.failureTolerant( GraphRagEventSink.composite(eventSinks.orderedStream().toList())), - taskDecorators.getIfAvailable(() -> GraphRagTaskDecorator.NONE)); + taskDecorators.getIfAvailable(() -> GraphRagTaskDecorator.NONE), + admission); } } diff --git a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagRetrievalPolicy.java b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagRetrievalPolicy.java index 2462db8d7..364504ce1 100644 --- a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagRetrievalPolicy.java +++ b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/GraphRagRetrievalPolicy.java @@ -11,6 +11,7 @@ public record GraphRagRetrievalPolicy( int maximumKnowledgeSpaces, int maximumConcurrentSpaces, + int retrievalAdmissionPermits, int topK, int chunkTopK, int relatedChunkNumber, @@ -25,6 +26,7 @@ public record GraphRagRetrievalPolicy( if (maximumKnowledgeSpaces <= 0 || maximumConcurrentSpaces <= 0 || maximumConcurrentSpaces > maximumKnowledgeSpaces + || retrievalAdmissionPermits <= 0 || topK <= 0 || chunkTopK <= 0 || relatedChunkNumber <= 0 @@ -69,7 +71,8 @@ public static GraphRagRetrievalPolicy defaults() { return new GraphRagRetrievalPolicy( 20, 4, - 60, + 4, + 40, 20, 5, 1, diff --git a/core/src/main/java/com/orgmemory/core/knowledge/retrieval/RetrievalAdmissionControl.java b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/RetrievalAdmissionControl.java new file mode 100644 index 000000000..0cc8e65b6 --- /dev/null +++ b/core/src/main/java/com/orgmemory/core/knowledge/retrieval/RetrievalAdmissionControl.java @@ -0,0 +1,39 @@ +package com.orgmemory.core.knowledge.retrieval; + +import java.util.concurrent.Callable; +import java.util.concurrent.Semaphore; + +/** + * JVM-local admission boundary for snapshot queries that may check out a JDBC + * connection. The Spring runtime owns one instance, shared by every retrieval + * turn in the process. + */ +final class RetrievalAdmissionControl { + + private final Semaphore permits; + + RetrievalAdmissionControl(int maximumConcurrentQueries) { + if (maximumConcurrentQueries <= 0) { + throw new IllegalArgumentException( + "maximumConcurrentQueries must be positive"); + } + this.permits = new Semaphore(maximumConcurrentQueries, true); + } + + T execute(Callable query) throws Exception { + permits.acquire(); + try { + return query.call(); + } finally { + permits.release(); + } + } + + boolean fair() { + return permits.isFair(); + } + + int availablePermits() { + return permits.availablePermits(); + } +} diff --git a/core/src/test/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalServiceTests.java b/core/src/test/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalServiceTests.java index ddc06d387..fcafbc150 100644 --- a/core/src/test/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalServiceTests.java +++ b/core/src/test/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalServiceTests.java @@ -3,6 +3,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.atLeastOnce; @@ -28,6 +29,7 @@ import com.orgmemory.core.permission.PermissionAuditService; import com.orgmemory.graphrag.model.EvidenceReference; import com.orgmemory.graphrag.observability.GraphRagEventSink; +import com.orgmemory.graphrag.observability.GraphRagTaskDecorator; import com.orgmemory.graphrag.query.ContextTokenUsage; import com.orgmemory.graphrag.query.KeywordPlan; import com.orgmemory.graphrag.query.LightRagGrounding; @@ -35,6 +37,7 @@ import com.orgmemory.graphrag.query.LightRagPreparedQuery; import com.orgmemory.graphrag.query.LightRagQueryEngine; import com.orgmemory.graphrag.query.LightRagQueryMode; +import com.orgmemory.graphrag.query.LightRagQueryRequest; import com.orgmemory.graphrag.query.LightRagQueryResult; import com.orgmemory.graphrag.storage.ProjectionKind; import com.orgmemory.graphrag.storage.ProjectionNamespace; @@ -79,6 +82,10 @@ class GraphRagKnowledgeRetrievalServiceTests { UUID.fromString("40000000-0000-0000-0000-000000000011"); private static final UUID SECOND_ASSET_ID = UUID.fromString("40000000-0000-0000-0000-000000000012"); + private static final UUID THIRD_SPACE_ID = + UUID.fromString("40000000-0000-0000-0000-000000000013"); + private static final UUID THIRD_ASSET_ID = + UUID.fromString("40000000-0000-0000-0000-000000000014"); private static final String MODEL_ID = "model-v1"; private static final Instant NOW = Instant.parse("2026-07-24T00:00:00Z"); @@ -168,6 +175,157 @@ void multipleSpacesPrepareOneLogicalQueryBeforeSnapshotRetrieval() { .allMatch(event -> event.scopeFingerprint() != null)); } + @Test + void acquiresAdmissionPermitBeforeExecutingTheSnapshotStoreQuery() { + CurrentActor actor = actor(); + PermissionAuditService audit = mock(PermissionAuditService.class); + KnowledgeEvidenceScopeResolver scopes = + mock(KnowledgeEvidenceScopeResolver.class); + when(scopes.resolve(actor, MODEL_ID)) + .thenReturn(scope(Set.of(ASSET_ID), 1L)); + RetrievalAdmissionControl admission = + new RetrievalAdmissionControl(1); + LightRagQueryEngine engine = mock(LightRagQueryEngine.class); + LightRagPreparedQuery prepared = preparedQueryPlan(); + when(engine.prepare(any())).thenReturn(prepared); + when(engine.executePrepared(any(), any())).thenAnswer(invocation -> { + assertEquals( + 0, + admission.availablePermits(), + "the storage-query boundary must not be entered before admission"); + return noResults(); + }); + + service( + scopes, + mock(RelationshipAuthorizationSetPort.class), + mock(SecureKnowledgeRetrievalStore.class), + engine, + policy(4, 1), + audit, + mock(GraphRagEventSink.class), + admission) + .search( + actor, + "What is the leave policy?", + 10, + "request-admission-before-checkout"); + + verify(engine).executePrepared(any(), any()); + assertEquals(1, admission.availablePermits()); + } + + @Test + void continuousAdmissionDoesNotWaitForAnEarlierSnapshotBatch() + throws InterruptedException { + CurrentActor actor = actor(); + PermissionAuditService audit = mock(PermissionAuditService.class); + KnowledgeEvidenceScopeResolver scopes = + mock(KnowledgeEvidenceScopeResolver.class); + when(scopes.resolve(actor, MODEL_ID)) + .thenReturn(threeSpaceScope()); + LightRagQueryEngine engine = mock(LightRagQueryEngine.class); + LightRagPreparedQuery prepared = preparedQueryPlan(); + when(engine.prepare(any())).thenReturn(prepared); + CountDownLatch firstStarted = new CountDownLatch(1); + CountDownLatch releaseFirst = new CountDownLatch(1); + CountDownLatch thirdStarted = new CountDownLatch(1); + when(engine.executePrepared(any(), any())).thenAnswer(invocation -> { + LightRagQueryRequest request = invocation.getArgument(0); + Set assets = request.scope().authorizedAssetIds(); + if (assets.contains(ASSET_ID)) { + firstStarted.countDown(); + assertTrue(releaseFirst.await(2, TimeUnit.SECONDS)); + } else if (assets.contains(THIRD_ASSET_ID)) { + thirdStarted.countDown(); + } + return noResults(); + }); + GraphRagKnowledgeRetrievalService service = service( + scopes, + mock(RelationshipAuthorizationSetPort.class), + mock(SecureKnowledgeRetrievalStore.class), + engine, + policy(2, 2), + audit, + mock(GraphRagEventSink.class)); + + var result = java.util.concurrent.CompletableFuture.supplyAsync(() -> + service.search( + actor, + "What is the leave policy?", + 10, + "request-continuous-admission")); + + try { + assertTrue(firstStarted.await(2, TimeUnit.SECONDS)); + assertTrue( + thirdStarted.await(2, TimeUnit.SECONDS), + "the next space should enter as soon as a permit is free"); + } finally { + releaseFirst.countDown(); + } + assertEquals(List.of(), result.join().evidence()); + } + + @Test + void snapshotFailureCancelsOutstandingContinuouslyAdmittedTasks() + throws InterruptedException { + CurrentActor actor = actor(); + PermissionAuditService audit = mock(PermissionAuditService.class); + KnowledgeEvidenceScopeResolver scopes = + mock(KnowledgeEvidenceScopeResolver.class); + when(scopes.resolve(actor, MODEL_ID)) + .thenReturn(threeSpaceScope()); + LightRagQueryEngine engine = mock(LightRagQueryEngine.class); + LightRagPreparedQuery prepared = preparedQueryPlan(); + when(engine.prepare(any())).thenReturn(prepared); + CountDownLatch blockersStarted = new CountDownLatch(2); + CountDownLatch neverReleased = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(2); + when(engine.executePrepared(any(), any())).thenAnswer(invocation -> { + LightRagQueryRequest request = invocation.getArgument(0); + if (request.scope().authorizedAssetIds().contains(SECOND_ASSET_ID)) { + assertTrue(blockersStarted.await(2, TimeUnit.SECONDS)); + throw new IllegalStateException("snapshot failed"); + } + blockersStarted.countDown(); + try { + neverReleased.await(); + return noResults(); + } catch (InterruptedException cancelled) { + interrupted.countDown(); + throw cancelled; + } + }); + GraphRagKnowledgeRetrievalService service = service( + scopes, + mock(RelationshipAuthorizationSetPort.class), + mock(SecureKnowledgeRetrievalStore.class), + engine, + policy(2, 3), + audit, + mock(GraphRagEventSink.class)); + + assertTimeoutPreemptively( + Duration.ofSeconds(3), + () -> assertThrows( + IllegalStateException.class, + () -> service.search( + actor, + "What is the leave policy?", + 10, + "request-fail-closed-cancellation"))); + assertTrue( + interrupted.await(2, TimeUnit.SECONDS), + "every outstanding snapshot task must be interrupted"); + } + + @Test + void retrievalAdmissionControlUsesFairQueueing() { + assertTrue(new RetrievalAdmissionControl(4).fair()); + } + @Test void emptyAuthorizedScopeReturnsAllowedEmptyWithoutFinalBatchCheck() { CurrentActor actor = actor(); @@ -901,6 +1059,7 @@ private static GraphRagRetrievalPolicy rerankPolicy() { return new GraphRagRetrievalPolicy( defaults.maximumKnowledgeSpaces(), defaults.maximumConcurrentSpaces(), + defaults.retrievalAdmissionPermits(), defaults.topK(), defaults.chunkTopK(), defaults.relatedChunkNumber(), @@ -923,6 +1082,27 @@ private static GraphRagKnowledgeRetrievalService service( GraphRagRetrievalPolicy policy, PermissionAuditService audit, GraphRagEventSink events) { + return service( + scopes, + finalAuthorization, + canonical, + engine, + policy, + audit, + events, + new RetrievalAdmissionControl( + policy.retrievalAdmissionPermits())); + } + + private static GraphRagKnowledgeRetrievalService service( + KnowledgeEvidenceScopeResolver scopes, + RelationshipAuthorizationSetPort finalAuthorization, + SecureKnowledgeRetrievalStore canonical, + LightRagQueryEngine engine, + GraphRagRetrievalPolicy policy, + PermissionAuditService audit, + GraphRagEventSink events, + RetrievalAdmissionControl admission) { RelationshipAuthorizationPort entry = mock(RelationshipAuthorizationPort.class); when(entry.check(any())) @@ -961,7 +1141,47 @@ private static GraphRagKnowledgeRetrievalService service( 5, 5_000, 1_000), - events); + events, + GraphRagTaskDecorator.NONE, + admission); + } + + private static GraphRagRetrievalPolicy policy( + int maximumConcurrentSpaces, + int admissionPermits) { + GraphRagRetrievalPolicy defaults = + GraphRagRetrievalPolicy.defaults(); + return new GraphRagRetrievalPolicy( + defaults.maximumKnowledgeSpaces(), + maximumConcurrentSpaces, + admissionPermits, + defaults.topK(), + defaults.chunkTopK(), + defaults.relatedChunkNumber(), + defaults.maximumGraphDepth(), + defaults.maximumEvidenceClosure(), + defaults.minimumVectorSimilarity(), + defaults.includeHeadings(), + defaults.rerank(), + defaults.contextBudget()); + } + + private static ResolvedKnowledgeEvidenceScope threeSpaceScope() { + return new ResolvedKnowledgeEvidenceScope( + ORGANIZATION_ID, + USER_ID, + null, + false, + MODEL_ID, + NOW, + Map.of( + SPACE_ID, Set.of(ASSET_ID), + SECOND_SPACE_ID, Set.of(SECOND_ASSET_ID), + THIRD_SPACE_ID, Set.of(THIRD_ASSET_ID)), + Map.of( + SPACE_ID, 1L, + SECOND_SPACE_ID, 1L, + THIRD_SPACE_ID, 1L)); } private static LightRagGrounding grounding() { From d74711fad5510bc472bdda46397396ce37b012d0 Mon Sep 17 00:00:00 2001 From: kl3inIT Date: Tue, 4 Aug 2026 23:47:29 +0700 Subject: [PATCH 3/5] feat(assistant): attribute post-retrieval latency --- apps/api/build.gradle.kts | 1 + .../api/assistant/AssistantConfiguration.java | 35 +++++- .../MicrometerAssistantStageEventSink.java | 37 +++++++ .../api/assistant/ObservedChatMemory.java | 72 +++++++++++++ .../OpenTelemetryAssistantStageEventSink.java | 79 ++++++++++++++ .../api/MetricsDistributionTests.java | 2 + ...icrometerAssistantStageEventSinkTests.java | 41 +++++++ .../assistant/ObservedChatMemoryTests.java | 79 ++++++++++++++ ...TelemetryAssistantStageEventSinkTests.java | 87 +++++++++++++++ .../core/assistant/AssistantService.java | 101 +++++++++++++++--- .../AssistantStageEventSink.java | 86 +++++++++++++++ .../AssistantTurnObservationTests.java | 35 +++++- .../AssistantTurnEventTests.java | 13 +++ 13 files changed, 646 insertions(+), 22 deletions(-) create mode 100644 apps/api/src/main/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSink.java create mode 100644 apps/api/src/main/java/com/orgmemory/api/assistant/ObservedChatMemory.java create mode 100644 apps/api/src/main/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSink.java create mode 100644 apps/api/src/test/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSinkTests.java create mode 100644 apps/api/src/test/java/com/orgmemory/api/assistant/ObservedChatMemoryTests.java create mode 100644 apps/api/src/test/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSinkTests.java create mode 100644 core/src/main/java/com/orgmemory/core/assistant/observability/AssistantStageEventSink.java diff --git a/apps/api/build.gradle.kts b/apps/api/build.gradle.kts index 2ab2cc517..ff7d5b096 100644 --- a/apps/api/build.gradle.kts +++ b/apps/api/build.gradle.kts @@ -42,6 +42,7 @@ dependencies { // not have. testImplementation("io.micrometer:micrometer-registry-otlp") testImplementation("io.opentelemetry:opentelemetry-sdk") + testImplementation("io.opentelemetry:opentelemetry-sdk-testing") testRuntimeOnly("org.junit.platform:junit-platform-launcher") } diff --git a/apps/api/src/main/java/com/orgmemory/api/assistant/AssistantConfiguration.java b/apps/api/src/main/java/com/orgmemory/api/assistant/AssistantConfiguration.java index 113d46485..d1e354e21 100644 --- a/apps/api/src/main/java/com/orgmemory/api/assistant/AssistantConfiguration.java +++ b/apps/api/src/main/java/com/orgmemory/api/assistant/AssistantConfiguration.java @@ -4,6 +4,7 @@ import com.orgmemory.core.assistant.AssistantAssetToolService; import com.orgmemory.core.assistant.AssistantAssetTraceRecorder; import com.orgmemory.core.assistant.AssistantService; +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; import com.orgmemory.core.assistant.observability.AssistantTurnEvent; import com.orgmemory.core.assistant.observability.AssistantTurnMeterObservationHandler; import com.orgmemory.core.assetregistry.AssetRegistryService; @@ -15,7 +16,9 @@ import com.orgmemory.core.knowledge.search.PermissionAwareKnowledgeSearch; import io.micrometer.core.instrument.MeterRegistry; import io.micrometer.observation.ObservationRegistry; +import io.opentelemetry.api.OpenTelemetry; import java.time.Clock; +import java.util.List; import org.springframework.ai.chat.memory.ChatMemory; import org.springframework.ai.chat.memory.ChatMemoryRepository; import org.springframework.ai.chat.memory.MessageWindowChatMemory; @@ -35,11 +38,18 @@ Clock assistantClock() { } @Bean - ChatMemory assistantChatMemory(ChatMemoryRepository repository) { - return MessageWindowChatMemory.builder() + ChatMemory assistantChatMemory( + ChatMemoryRepository repository, + AssistantStageEventSink stages, + AssistantProperties properties) { + ChatMemory memory = MessageWindowChatMemory.builder() .chatMemoryRepository(repository) .maxMessages(20) .build(); + return new ObservedChatMemory( + memory, + stages, + observedEngine(properties)); } @Bean @@ -62,9 +72,26 @@ AssistantService assistantService( PermissionAwareKnowledgeSearch retrieval, ChatModelPort chat, ObservationRegistry observations, - AssistantProperties properties) { + AssistantProperties properties, + AssistantStageEventSink stages) { return new AssistantService( - retrieval, chat, observations, observedEngine(properties)); + retrieval, + chat, + observations, + observedEngine(properties), + stages); + } + + @Bean + AssistantStageEventSink assistantStageEventSink( + OpenTelemetry openTelemetry, + MeterRegistry meters) { + return AssistantStageEventSink.composite(List.of( + AssistantStageEventSink.failureTolerant( + new OpenTelemetryAssistantStageEventSink( + openTelemetry)), + AssistantStageEventSink.failureTolerant( + new MicrometerAssistantStageEventSink(meters)))); } /** diff --git a/apps/api/src/main/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSink.java b/apps/api/src/main/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSink.java new file mode 100644 index 000000000..9ed11b1d3 --- /dev/null +++ b/apps/api/src/main/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSink.java @@ -0,0 +1,37 @@ +package com.orgmemory.api.assistant; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.Timer; +import java.util.Locale; +import java.util.Objects; + +/** Bounded-cardinality timer for assistant latency attribution stages. */ +final class MicrometerAssistantStageEventSink + implements AssistantStageEventSink { + + static final String STAGE_TIMER = "orgmemory.assistant.stage"; + + private final MeterRegistry registry; + + MicrometerAssistantStageEventSink(MeterRegistry registry) { + this.registry = Objects.requireNonNull(registry, "registry"); + } + + @Override + public void emit(AssistantStageEvent event) { + Objects.requireNonNull(event, "event"); + Timer.builder(STAGE_TIMER) + .description( + "Assistant latency stages above permission-aware retrieval") + .tag("engine", value(event.engine())) + .tag("stage", value(event.stage())) + .tag("outcome", value(event.outcome())) + .register(registry) + .record(event.duration()); + } + + private static String value(Enum value) { + return value.name().toLowerCase(Locale.ROOT); + } +} diff --git a/apps/api/src/main/java/com/orgmemory/api/assistant/ObservedChatMemory.java b/apps/api/src/main/java/com/orgmemory/api/assistant/ObservedChatMemory.java new file mode 100644 index 000000000..e0954b0f8 --- /dev/null +++ b/apps/api/src/main/java/com/orgmemory/api/assistant/ObservedChatMemory.java @@ -0,0 +1,72 @@ +package com.orgmemory.api.assistant; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import com.orgmemory.core.assistant.observability.AssistantStageEventSink.AssistantStageEvent; +import com.orgmemory.core.assistant.observability.AssistantTurnEvent; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Objects; +import org.springframework.ai.chat.memory.ChatMemory; +import org.springframework.ai.chat.messages.Message; + +/** Measures the history read without exposing its conversation or messages. */ +final class ObservedChatMemory implements ChatMemory { + + private final ChatMemory delegate; + private final AssistantStageEventSink events; + private final AssistantTurnEvent.RetrievalEngine engine; + + ObservedChatMemory( + ChatMemory delegate, + AssistantStageEventSink events, + AssistantTurnEvent.RetrievalEngine engine) { + this.delegate = Objects.requireNonNull(delegate, "delegate"); + this.events = Objects.requireNonNull(events, "events"); + this.engine = Objects.requireNonNull(engine, "engine"); + } + + @Override + public void add(String conversationId, List messages) { + delegate.add(conversationId, messages); + } + + @Override + public List get(String conversationId) { + long startedAt = System.nanoTime(); + try { + List messages = delegate.get(conversationId); + emit( + AssistantStageEventSink.Outcome.SUCCEEDED, + startedAt, + null); + return messages; + } catch (RuntimeException | Error failure) { + emit( + AssistantStageEventSink.Outcome.FAILED, + startedAt, + "history_load_failed"); + throw failure; + } + } + + @Override + public void clear(String conversationId) { + delegate.clear(conversationId); + } + + private void emit( + AssistantStageEventSink.Outcome outcome, + long startedAt, + String failureCode) { + events.emit(new AssistantStageEvent( + engine, + AssistantStageEventSink.Stage.CONVERSATION_HISTORY_LOAD, + outcome, + Duration.ofNanos(Math.max( + 0L, + System.nanoTime() - startedAt)), + failureCode, + Instant.now())); + } +} diff --git a/apps/api/src/main/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSink.java b/apps/api/src/main/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSink.java new file mode 100644 index 000000000..4ee9e2814 --- /dev/null +++ b/apps/api/src/main/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSink.java @@ -0,0 +1,79 @@ +package com.orgmemory.api.assistant; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import java.time.Instant; +import java.util.Locale; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** Payload-free OpenTelemetry adapter for assistant latency attribution. */ +final class OpenTelemetryAssistantStageEventSink + implements AssistantStageEventSink { + + static final String INSTRUMENTATION_SCOPE = + "com.orgmemory.assistant"; + static final AttributeKey ENGINE = + AttributeKey.stringKey("orgmemory.assistant.engine"); + static final AttributeKey STAGE = + AttributeKey.stringKey("orgmemory.assistant.stage"); + static final AttributeKey OUTCOME = + AttributeKey.stringKey("orgmemory.assistant.outcome"); + static final AttributeKey DURATION_NANOS = + AttributeKey.longKey("orgmemory.assistant.duration_nanos"); + static final AttributeKey FAILURE_CODE = + AttributeKey.stringKey("orgmemory.assistant.failure_code"); + + private final Tracer tracer; + + OpenTelemetryAssistantStageEventSink(OpenTelemetry openTelemetry) { + tracer = Objects.requireNonNull(openTelemetry, "openTelemetry") + .getTracer(INSTRUMENTATION_SCOPE); + } + + @Override + public void emit(AssistantStageEvent event) { + Objects.requireNonNull(event, "event"); + long endEpochNanos = epochNanos(event.occurredAt()); + long startEpochNanos = Math.subtractExact( + endEpochNanos, + event.duration().toNanos()); + Span span = tracer.spanBuilder( + "orgmemory.assistant." + value(event.stage())) + .setSpanKind(SpanKind.INTERNAL) + .setStartTimestamp( + startEpochNanos, + TimeUnit.NANOSECONDS) + .startSpan(); + span.setAttribute(ENGINE, value(event.engine())); + span.setAttribute(STAGE, value(event.stage())); + span.setAttribute(OUTCOME, value(event.outcome())); + span.setAttribute( + DURATION_NANOS, + event.duration().toNanos()); + if (event.failureCode() != null) { + span.setAttribute(FAILURE_CODE, event.failureCode()); + } + if (event.outcome() == Outcome.FAILED) { + span.setStatus(StatusCode.ERROR); + } + span.end(endEpochNanos, TimeUnit.NANOSECONDS); + } + + private static String value(Enum value) { + return value.name().toLowerCase(Locale.ROOT); + } + + private static long epochNanos(Instant instant) { + return Math.addExact( + Math.multiplyExact( + instant.getEpochSecond(), + 1_000_000_000L), + instant.getNano()); + } +} diff --git a/apps/api/src/test/java/com/orgmemory/api/MetricsDistributionTests.java b/apps/api/src/test/java/com/orgmemory/api/MetricsDistributionTests.java index 6b7aa5260..0b8920469 100644 --- a/apps/api/src/test/java/com/orgmemory/api/MetricsDistributionTests.java +++ b/apps/api/src/test/java/com/orgmemory/api/MetricsDistributionTests.java @@ -45,6 +45,7 @@ class MetricsDistributionTests { "jvm.gc.pause", "orgmemory.assistant.turn", "orgmemory.assistant.time_to_first_token", + "orgmemory.assistant.stage", "orgmemory.graph_rag.stage", "gen_ai.client.operation" }) @@ -63,6 +64,7 @@ void meterChartedAsAQuantilePublishesAHistogram(String name) { strings = { "http.server.requests", "orgmemory.assistant.turn", + "orgmemory.assistant.stage", "orgmemory.graph_rag.stage", "gen_ai.client.operation" }) diff --git a/apps/api/src/test/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSinkTests.java b/apps/api/src/test/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSinkTests.java new file mode 100644 index 000000000..0ed04d0af --- /dev/null +++ b/apps/api/src/test/java/com/orgmemory/api/assistant/MicrometerAssistantStageEventSinkTests.java @@ -0,0 +1,41 @@ +package com.orgmemory.api.assistant; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import com.orgmemory.core.assistant.observability.AssistantTurnEvent; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.time.Duration; +import java.time.Instant; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; + +class MicrometerAssistantStageEventSinkTests { + + @Test + void recordsOnlyBoundedStageDimensions() { + SimpleMeterRegistry registry = new SimpleMeterRegistry(); + var sink = new MicrometerAssistantStageEventSink(registry); + + sink.emit(new AssistantStageEventSink.AssistantStageEvent( + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG, + AssistantStageEventSink.Stage.GROUNDING_TO_PROMPT, + AssistantStageEventSink.Outcome.SUCCEEDED, + Duration.ofMillis(25), + null, + Instant.parse("2026-08-05T01:02:03Z"))); + + Timer timer = registry.get( + MicrometerAssistantStageEventSink.STAGE_TIMER) + .timer(); + assertEquals(25.0, timer.totalTime(TimeUnit.MILLISECONDS)); + assertEquals( + Set.of("engine", "stage", "outcome"), + timer.getId().getTags().stream() + .map(tag -> tag.getKey()) + .collect(Collectors.toSet())); + } +} diff --git a/apps/api/src/test/java/com/orgmemory/api/assistant/ObservedChatMemoryTests.java b/apps/api/src/test/java/com/orgmemory/api/assistant/ObservedChatMemoryTests.java new file mode 100644 index 000000000..db308d874 --- /dev/null +++ b/apps/api/src/test/java/com/orgmemory/api/assistant/ObservedChatMemoryTests.java @@ -0,0 +1,79 @@ +package com.orgmemory.api.assistant; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import com.orgmemory.core.assistant.observability.AssistantTurnEvent; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.ai.chat.memory.ChatMemory; +import org.springframework.ai.chat.messages.Message; + +class ObservedChatMemoryTests { + + @Test + void measuresHistoryLoadingWithoutPublishingConversationOrMessages() { + ChatMemory delegate = mock(ChatMemory.class); + AssistantStageEventSink events = + mock(AssistantStageEventSink.class); + Message message = mock(Message.class); + when(delegate.get("conversation-secret")) + .thenReturn(List.of(message)); + var memory = new ObservedChatMemory( + delegate, + events, + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG); + + assertEquals( + List.of(message), + memory.get("conversation-secret")); + + ArgumentCaptor captured = + ArgumentCaptor.forClass( + AssistantStageEventSink.AssistantStageEvent.class); + verify(events).emit(captured.capture()); + assertEquals( + AssistantStageEventSink.Stage.CONVERSATION_HISTORY_LOAD, + captured.getValue().stage()); + assertEquals( + AssistantStageEventSink.Outcome.SUCCEEDED, + captured.getValue().outcome()); + assertEquals(null, captured.getValue().failureCode()); + } + + @Test + void reportsAClosedFailureCodeAndPreservesTheMemoryFailure() { + ChatMemory delegate = mock(ChatMemory.class); + AssistantStageEventSink events = + mock(AssistantStageEventSink.class); + IllegalStateException failure = + new IllegalStateException("private database detail"); + when(delegate.get("conversation-secret")).thenThrow(failure); + var memory = new ObservedChatMemory( + delegate, + events, + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG); + + assertEquals( + failure, + assertThrows( + IllegalStateException.class, + () -> memory.get("conversation-secret"))); + + ArgumentCaptor captured = + ArgumentCaptor.forClass( + AssistantStageEventSink.AssistantStageEvent.class); + verify(events).emit(captured.capture()); + assertEquals( + AssistantStageEventSink.Outcome.FAILED, + captured.getValue().outcome()); + assertEquals( + "history_load_failed", + captured.getValue().failureCode()); + } +} diff --git a/apps/api/src/test/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSinkTests.java b/apps/api/src/test/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSinkTests.java new file mode 100644 index 000000000..2e5e66db9 --- /dev/null +++ b/apps/api/src/test/java/com/orgmemory/api/assistant/OpenTelemetryAssistantStageEventSinkTests.java @@ -0,0 +1,87 @@ +package com.orgmemory.api.assistant; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import com.orgmemory.core.assistant.observability.AssistantTurnEvent; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import java.time.Duration; +import java.time.Instant; +import java.util.Set; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; + +class OpenTelemetryAssistantStageEventSinkTests { + + @Test + void exportsOnlyTheClosedPayloadFreeAttributeSetWithOriginalTiming() { + InMemorySpanExporter exporter = InMemorySpanExporter.create(); + try (SdkTracerProvider provider = SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(exporter)) + .build()) { + OpenTelemetrySdk telemetry = OpenTelemetrySdk.builder() + .setTracerProvider(provider) + .build(); + var sink = new OpenTelemetryAssistantStageEventSink(telemetry); + Instant endedAt = + Instant.parse("2026-08-05T01:02:03.123456789Z"); + Duration duration = Duration.ofMillis(125); + + sink.emit(new AssistantStageEventSink.AssistantStageEvent( + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG, + AssistantStageEventSink.Stage + .CONVERSATION_HISTORY_LOAD, + AssistantStageEventSink.Outcome.FAILED, + duration, + "history_load_failed", + endedAt)); + + var span = exporter.getFinishedSpanItems().getFirst(); + assertEquals( + "orgmemory.assistant.conversation_history_load", + span.getName()); + assertEquals( + StatusCode.ERROR, + span.getStatus().getStatusCode()); + assertEquals( + epochNanos(endedAt), + span.getEndEpochNanos()); + assertEquals( + duration.toNanos(), + span.getEndEpochNanos() + - span.getStartEpochNanos()); + Set keys = span.getAttributes().asMap().keySet() + .stream() + .map(key -> key.getKey()) + .collect(Collectors.toSet()); + assertEquals( + Set.of( + "orgmemory.assistant.engine", + "orgmemory.assistant.stage", + "orgmemory.assistant.outcome", + "orgmemory.assistant.duration_nanos", + "orgmemory.assistant.failure_code"), + keys); + assertFalse(keys.stream().anyMatch(key -> + key.contains("query") + || key.contains("prompt") + || key.contains("evidence") + || key.contains("user") + || key.contains("conversation") + || key.contains("exception"))); + } + } + + private static long epochNanos(Instant instant) { + return Math.addExact( + Math.multiplyExact( + instant.getEpochSecond(), + 1_000_000_000L), + instant.getNano()); + } +} diff --git a/core/src/main/java/com/orgmemory/core/assistant/AssistantService.java b/core/src/main/java/com/orgmemory/core/assistant/AssistantService.java index 20215c6f9..c81fdbaaf 100644 --- a/core/src/main/java/com/orgmemory/core/assistant/AssistantService.java +++ b/core/src/main/java/com/orgmemory/core/assistant/AssistantService.java @@ -4,6 +4,8 @@ import com.orgmemory.core.ai.AiWorkload; import com.orgmemory.core.ai.ChatModelPort; import com.orgmemory.core.ai.AssistantModelRouteAuthority; +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; +import com.orgmemory.core.assistant.observability.AssistantStageEventSink.AssistantStageEvent; import com.orgmemory.core.assistant.observability.AssistantTurnEvent; import com.orgmemory.core.assistant.observability.AssistantTurnObservationContext; import com.orgmemory.core.assistant.observability.AssistantTurnObservationDocumentation; @@ -13,6 +15,8 @@ import com.orgmemory.core.organization.CurrentActor; import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; +import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.List; import reactor.core.publisher.Flux; @@ -29,6 +33,7 @@ public class AssistantService { private final ChatModelPort chat; private final ObservationRegistry observations; private final AssistantTurnEvent.RetrievalEngine engine; + private final AssistantStageEventSink stages; /** * @param engine which retrieval implementation was wired in. Passed rather than asked of @@ -41,10 +46,25 @@ public AssistantService( ChatModelPort chat, ObservationRegistry observations, AssistantTurnEvent.RetrievalEngine engine) { + this( + retrieval, + chat, + observations, + engine, + AssistantStageEventSink.NO_OP); + } + + public AssistantService( + PermissionAwareKnowledgeSearch retrieval, + ChatModelPort chat, + ObservationRegistry observations, + AssistantTurnEvent.RetrievalEngine engine, + AssistantStageEventSink stages) { this.retrieval = retrieval; this.chat = chat; this.observations = observations; this.engine = engine; + this.stages = stages; } /** @@ -93,22 +113,38 @@ public AssistantTurn startTurn( return new AssistantTurn(search.requestId(), List.of(), Flux.just(NO_ACCESSIBLE_EVIDENCE)); } - PreparedTurn prepared = search.grounding() - .map(grounding -> new PreparedTurn( - AssistantPromptFactory.addUserContext( - grounding.generationRequest(), - actor), - numbered(grounding.citations()))) - .orElseGet(() -> { - AssistantPromptFactory.PreparedPrompt prompt = - AssistantPromptFactory.create( - question, - search.evidence(), - actor); - return new PreparedTurn( - prompt.request(), - prompt.citations()); - }); + long retrievalCompletedAt = System.nanoTime(); + PreparedTurn prepared; + try { + prepared = search.grounding() + .map(grounding -> new PreparedTurn( + AssistantPromptFactory.addUserContext( + grounding.generationRequest(), + actor), + numbered(grounding.citations()))) + .orElseGet(() -> { + AssistantPromptFactory.PreparedPrompt prompt = + AssistantPromptFactory.create( + question, + search.evidence(), + actor); + return new PreparedTurn( + prompt.request(), + prompt.citations()); + }); + emitStage( + AssistantStageEventSink.Stage.GROUNDING_TO_PROMPT, + AssistantStageEventSink.Outcome.SUCCEEDED, + retrievalCompletedAt, + null); + } catch (RuntimeException failure) { + emitStage( + AssistantStageEventSink.Stage.GROUNDING_TO_PROMPT, + AssistantStageEventSink.Outcome.FAILED, + retrievalCompletedAt, + "prompt_assembly_failed"); + throw failure; + } int evidenceCount = search.evidence().size(); int citationCount = prepared.citations().size(); @@ -116,6 +152,8 @@ public AssistantTurn startTurn( // stop is not idempotent, so the second terminal signal must not reach it. java.util.concurrent.atomic.AtomicBoolean stopped = new java.util.concurrent.atomic.AtomicBoolean(); + java.util.concurrent.atomic.AtomicBoolean firstTokenStageEmitted = + new java.util.concurrent.atomic.AtomicBoolean(); Flux generated = routeAuthority == null ? chat.stream( actor.organizationId(), @@ -137,7 +175,20 @@ public AssistantTurn startTurn( error -> new AssistantUnavailableException("The assistant is unavailable", error)) // Before the error mapping's terminal signal, so the moment the caller // could first see something is recorded even on a stream that later fails. - .doOnNext(token -> context.firstTokenAt(System.nanoTime())) + .doOnNext(token -> { + long firstTokenAt = System.nanoTime(); + context.firstTokenAt(firstTokenAt); + if (firstTokenStageEmitted.compareAndSet( + false, + true)) { + emitStage( + AssistantStageEventSink.Stage + .RETRIEVAL_TO_FIRST_TOKEN, + AssistantStageEventSink.Outcome.SUCCEEDED, + retrievalCompletedAt, + null); + } + }) .doOnError(error -> { context.unavailable(System.nanoTime(), "assistant_stream_failed"); observation.error(error); @@ -161,6 +212,22 @@ public AssistantTurn startTurn( } } + private void emitStage( + AssistantStageEventSink.Stage stage, + AssistantStageEventSink.Outcome outcome, + long startedAt, + String failureCode) { + stages.emit(new AssistantStageEvent( + engine, + stage, + outcome, + Duration.ofNanos(Math.max( + 0L, + System.nanoTime() - startedAt)), + failureCode, + Instant.now())); + } + /** * Ends the observation for a turn that failed before it could return a stream, so a * failure raised during retrieval is not left as an observation nothing ever stops. diff --git a/core/src/main/java/com/orgmemory/core/assistant/observability/AssistantStageEventSink.java b/core/src/main/java/com/orgmemory/core/assistant/observability/AssistantStageEventSink.java new file mode 100644 index 000000000..7e474f7f9 --- /dev/null +++ b/core/src/main/java/com/orgmemory/core/assistant/observability/AssistantStageEventSink.java @@ -0,0 +1,86 @@ +package com.orgmemory.core.assistant.observability; + +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Objects; + +/** + * Payload-free boundary for assistant latency stages above retrieval. + */ +@FunctionalInterface +public interface AssistantStageEventSink { + + AssistantStageEventSink NO_OP = event -> { + }; + + enum Stage { + GROUNDING_TO_PROMPT, + CONVERSATION_HISTORY_LOAD, + RETRIEVAL_TO_FIRST_TOKEN + } + + enum Outcome { + SUCCEEDED, + FAILED + } + + record AssistantStageEvent( + AssistantTurnEvent.RetrievalEngine engine, + Stage stage, + Outcome outcome, + Duration duration, + String failureCode, + Instant occurredAt) { + + public AssistantStageEvent { + Objects.requireNonNull(engine, "engine"); + Objects.requireNonNull(stage, "stage"); + Objects.requireNonNull(outcome, "outcome"); + Objects.requireNonNull(duration, "duration"); + Objects.requireNonNull(occurredAt, "occurredAt"); + if (duration.isNegative()) { + throw new IllegalArgumentException( + "duration must not be negative"); + } + failureCode = failureCode == null || failureCode.isBlank() + ? null + : failureCode.strip(); + if (outcome == Outcome.FAILED && failureCode == null) { + throw new IllegalArgumentException( + "failureCode is required for a failed stage"); + } + if (failureCode != null + && !AssistantTurnEvent.FAILURE_CODE + .matcher(failureCode) + .matches()) { + throw new IllegalArgumentException( + "failureCode must be a bounded machine code"); + } + } + } + + void emit(AssistantStageEvent event); + + static AssistantStageEventSink composite( + List sinks) { + List delegates = List.copyOf( + Objects.requireNonNull(sinks, "sinks")); + if (delegates.isEmpty()) { + return NO_OP; + } + return event -> delegates.forEach(sink -> sink.emit(event)); + } + + static AssistantStageEventSink failureTolerant( + AssistantStageEventSink delegate) { + Objects.requireNonNull(delegate, "delegate"); + return event -> { + try { + delegate.emit(event); + } catch (RuntimeException ignored) { + // Telemetry cannot make an otherwise valid assistant turn fail. + } + }; + } +} diff --git a/core/src/test/java/com/orgmemory/core/assistant/AssistantTurnObservationTests.java b/core/src/test/java/com/orgmemory/core/assistant/AssistantTurnObservationTests.java index 07880f927..caf3dadcb 100644 --- a/core/src/test/java/com/orgmemory/core/assistant/AssistantTurnObservationTests.java +++ b/core/src/test/java/com/orgmemory/core/assistant/AssistantTurnObservationTests.java @@ -8,6 +8,8 @@ import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.orgmemory.core.ai.AiGatewayUnavailableException; @@ -15,6 +17,7 @@ import com.orgmemory.core.ai.ChatGenerationRequest; import com.orgmemory.core.ai.ChatModelPort; import com.orgmemory.core.assistant.observability.AssistantTurnEvent; +import com.orgmemory.core.assistant.observability.AssistantStageEventSink; import com.orgmemory.core.assistant.observability.AssistantTurnMeterObservationHandler; import com.orgmemory.core.knowledge.retrieval.CanonicalHybridKnowledgeSearch; import com.orgmemory.core.knowledge.search.RetrievedKnowledgeEvidence; @@ -32,6 +35,7 @@ import java.util.UUID; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import reactor.core.publisher.Flux; /** @@ -53,6 +57,8 @@ class AssistantTurnObservationTests { private final ChatModelPort chat = mock(ChatModelPort.class); private final SimpleMeterRegistry meters = new SimpleMeterRegistry(); private final ObservationRegistry observations = observationRegistry(); + private final AssistantStageEventSink stages = + mock(AssistantStageEventSink.class); private final CurrentActor actor = new CurrentActor( UUID.randomUUID(), UUID.randomUUID(), @@ -169,13 +175,40 @@ void tagsTheEngineSoTwoRetrievalImplementationsStaySeparable() { "comparing engines is the reason the tag exists"); } + @Test + void attributesPromptAssemblyAndRetrievalToFirstTokenAboveRetrieval() { + whenSearchTakes(Duration.ZERO); + whenModelStreams(Flux.just("Sixty days.")); + + drain(startTurn()); + + ArgumentCaptor captured = + ArgumentCaptor.forClass( + AssistantStageEventSink.AssistantStageEvent.class); + verify(stages, times(2)).emit(captured.capture()); + assertEquals( + Set.of( + AssistantStageEventSink.Stage.GROUNDING_TO_PROMPT, + AssistantStageEventSink.Stage.RETRIEVAL_TO_FIRST_TOKEN), + captured.getAllValues().stream() + .map(AssistantStageEventSink.AssistantStageEvent::stage) + .collect(java.util.stream.Collectors.toSet())); + assertTrue(captured.getAllValues().stream().allMatch(event -> + !event.duration().isNegative() + && event.failureCode() == null)); + } + private AssistantTurn startTurn() { return service().startTurn(actor, "What is the probation policy?", 5, "request-1", CONVERSATION_ID); } private AssistantService service() { return new AssistantService( - retrieval, chat, observations, AssistantTurnEvent.RetrievalEngine.GRAPH_RAG); + retrieval, + chat, + observations, + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG, + stages); } private static void drain(AssistantTurn turn) { diff --git a/core/src/test/java/com/orgmemory/core/assistant/observability/AssistantTurnEventTests.java b/core/src/test/java/com/orgmemory/core/assistant/observability/AssistantTurnEventTests.java index 70151b441..51fcbbd49 100644 --- a/core/src/test/java/com/orgmemory/core/assistant/observability/AssistantTurnEventTests.java +++ b/core/src/test/java/com/orgmemory/core/assistant/observability/AssistantTurnEventTests.java @@ -19,6 +19,19 @@ */ class AssistantTurnEventTests { + @Test + void assistantStageEventRejectsFreeTextFailureCodes() { + assertThrows( + IllegalArgumentException.class, + () -> new AssistantStageEventSink.AssistantStageEvent( + AssistantTurnEvent.RetrievalEngine.GRAPH_RAG, + AssistantStageEventSink.Stage.GROUNDING_TO_PROMPT, + AssistantStageEventSink.Outcome.FAILED, + Duration.ofMillis(1), + "private prompt text is forbidden", + java.time.Instant.EPOCH)); + } + private static final UUID ORGANIZATION = UUID.randomUUID(); @Test From 080246b85760f67bcc4af6f399cd39177213f903 Mon Sep 17 00:00:00 2001 From: kl3inIT Date: Wed, 5 Aug 2026 00:04:31 +0700 Subject: [PATCH 4/5] docs(release): describe retrieval admission control Co-Authored-By: Claude Opus 5 (1M context) --- .tegami/2026-08-05-retrieval-admission-control.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .tegami/2026-08-05-retrieval-admission-control.md diff --git a/.tegami/2026-08-05-retrieval-admission-control.md b/.tegami/2026-08-05-retrieval-admission-control.md new file mode 100644 index 000000000..cc7082ee1 --- /dev/null +++ b/.tegami/2026-08-05-retrieval-admission-control.md @@ -0,0 +1,15 @@ +--- +packages: + orgmemory: patch +subject: Faster, fairer assistant retrieval under load +--- + +## Improvements + +Assistant knowledge retrieval now admits snapshot queries through one fair +process-wide limit instead of per-request batches, so concurrent +conversations can no longer exhaust the database connection pool and stall at +the turn timeout. The API connection pool is right-sized for the production +host, retrieval breadth returns to the upstream LightRAG default, and new +payload-free timing stages make the previously unattributed portion of +time-to-first-token observable. From ccce26390f402018ef31680f36f5820b2d2655bb Mon Sep 17 00:00:00 2001 From: kl3inIT Date: Wed, 5 Aug 2026 00:14:20 +0700 Subject: [PATCH 5/5] docs: record the deferred cancellation finding from PR review Co-Authored-By: Claude Opus 5 (1M context) --- .../2026-08-05-retrieval-admission-phase1/plan.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md b/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md index bf4abfeb9..793c536e6 100644 --- a/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md +++ b/docs/increments/active/2026-08-05-retrieval-admission-phase1/plan.md @@ -48,3 +48,13 @@ Design: [design.md](design.md). Decision: ADR 0020 conditions 1 and 7. surviving path before any further latency work. - Consolidate: spec/test matrix refresh for the retrieval domain, roadmap update, move to completed. + +## Deferred review finding (PR #292, CodeRabbit) + +The turn timeout does not interrupt a turn blocked in admission or in an +in-flight snapshot query; an abandoned turn consumes its permit and one +storage query after the timeout fires (bounded zombie work; permits always +release). Pre-existing in part — the timeout never interrupted the +synchronous search path. Deferred to the Phase 2 compound-query port, whose +design must include deadline-aware admission and cooperative cancellation +between the turn stream and the retrieval future.