feat(graph-rag): publish what one answer costs and what the budget refused

components/graph-rag-core/src/main/java/com/orgmemory/graphrag/observability/GraphRagEventSink.java

@@ -82,8 +82,44 @@ record GraphRagEvent(
             String scopeFingerprint,
             CacheStatus cacheStatus,
             String failureCode,
+            TokenUsage tokenUsage,
             Instant occurredAt) {
 
+        /**
+         * Describes a stage that measures no tokens, which is every stage but
+         * context assembly. The overload exists so that adding a measurement one
+         * stage can take does not oblige the other stages to pass {@code null},
+         * where a reader would have to count commas to learn what was omitted.
+         */
+        public GraphRagEvent(
+                UUID operationId,
+                UUID organizationId,
+                Stage stage,
+                Outcome outcome,
+                Duration duration,
+                int inputCount,
+                int outputCount,
+                String modelRouteFingerprint,
+                String scopeFingerprint,
+                CacheStatus cacheStatus,
+                String failureCode,
+                Instant occurredAt) {
+            this(
+                    operationId,
+                    organizationId,
+                    stage,
+                    outcome,
+                    duration,
+                    inputCount,
+                    outputCount,
+                    modelRouteFingerprint,
+                    scopeFingerprint,
+                    cacheStatus,
+                    failureCode,
+                    null,
+                    occurredAt);
+        }
+
         public GraphRagEvent {
             Objects.requireNonNull(operationId, "operationId");
             Objects.requireNonNull(organizationId, "organizationId");
@@ -122,6 +158,64 @@ record GraphRagEvent(
         }
     }
 
+    /**
+     * Token cost of one assembled generation context.
+     *
+     * <p>Counts, not content. A token count cannot reconstruct the text it
+     * measures, which is why this record is allowed through a boundary that
+     * refuses queries, prompts and evidence.
+     *
+     * <p>Deliberately a projection of the query model's own token record rather
+     * than a reference to it. That record exists to enforce a retrieval budget
+     * and is free to change shape for retrieval reasons; a telemetry contract
+     * that moved with it would rewrite dashboards for a decision that had
+     * nothing to do with them. The two overlap today and may diverge.
+     *
+     * @param promptTokens         what the model is actually charged for, which
+     *                             exceeds the sum of the channels below by the
+     *                             cost of rendering them into one prompt
+     * @param budgetTokens         the ceiling {@code promptTokens} was fitted to,
+     *                             so a collector can express headroom as a ratio
+     *                             without knowing this deployment's configuration
+     * @param droppedContributions entities, relations and chunks evicted to make
+     *                             the context fit. This is the input-side
+     *                             truncation signal: zero means the retrieved
+     *                             context reached the model whole, and a rising
+     *                             count means answers are degrading for a reason
+     *                             no latency or error metric reports.
+     */
+    record TokenUsage(
+            int promptTokens,
+            int systemPromptTokens,
+            int queryTokens,
+            int entityTokens,
+            int relationTokens,
+            int chunkTokens,
+            int budgetTokens,
+            int droppedContributions) {
+
+        public TokenUsage {
+            if (promptTokens < 0
+                    || systemPromptTokens < 0
+                    || queryTokens < 0
+                    || entityTokens < 0
+                    || relationTokens < 0
+                    || chunkTokens < 0
+                    || droppedContributions < 0) {
+                throw new IllegalArgumentException(
+                        "token counts must be non-negative");
+            }
+            if (budgetTokens <= 0) {
+                throw new IllegalArgumentException(
+                        "budgetTokens must be positive");
+            }
+        }
+
+        public boolean truncated() {
+            return droppedContributions > 0;
+        }
+    }
+
     enum Stage {
         PARSE,
         CHUNK,

components/graph-rag-core/src/main/java/com/orgmemory/graphrag/query/LightRagGroundingAssembler.java

@@ -78,7 +78,11 @@ public PreparedGrounding prepare(
                 distinctScopes(scopes),
                 usage,
                 chunkAllocation.usedTokens());
-        return fitAndRender(query, options, grounding);
+        // Counted after merging and before any budget runs, so the reported
+        // eviction is what the budget cost the answer and not what deduplication
+        // saved it.
+        int selected = entities.size() + relations.size() + chunks.size();
+        return fitAndRender(query, options, grounding, selected);
     }
 
     public PreparedGrounding consolidate(
@@ -111,8 +115,7 @@ public PreparedGrounding render(
         Objects.requireNonNull(options, "options");
         Objects.requireNonNull(grounding, "grounding");
         PreparedGrounding prepared = renderUnchecked(query, options, grounding);
-        int maximumInputTokens = options.contextBudget().maxTotalTokens()
-                - options.contextBudget().safetyBufferTokens();
+        int maximumInputTokens = options.contextBudget().maximumInputTokens();
         if (prepared.inputTokens() > maximumInputTokens) {
             throw new IllegalArgumentException(
                     "verified grounding exceeds the configured input budget");
@@ -123,10 +126,10 @@ public PreparedGrounding render(
     private PreparedGrounding fitAndRender(
             String query,
             LightRagQueryRequest.Options options,
-            LightRagGrounding grounding) {
+            LightRagGrounding grounding,
+            int selectedContributions) {
         SecureContextBudget budget = options.contextBudget();
-        int maximumInputTokens =
-                budget.maxTotalTokens() - budget.safetyBufferTokens();
+        int maximumInputTokens = budget.maximumInputTokens();
         TokenMeasurements measurements =
                 measureTokens(query, options, grounding);
         LightRagGrounding candidate =
@@ -155,7 +158,14 @@ private PreparedGrounding fitAndRender(
             throw new IllegalArgumentException(
                     "query and system instructions exceed the configured input budget");
         }
-        return prepared;
+        return prepared.withDroppedContributions(
+                selectedContributions - contributionCount(candidate));
+    }
+
+    private static int contributionCount(LightRagGrounding grounding) {
+        return grounding.entities().size()
+                + grounding.relations().size()
+                + grounding.chunks().size();
     }
 
     private TokenMeasurements measureTokens(
@@ -254,7 +264,8 @@ private PreparedGrounding renderUnchecked(
                 systemPrompt,
                 prompt,
                 references,
-                tokenizer.count(prompt));
+                tokenizer.count(prompt),
+                0);
     }
 
     private static LightRagGrounding removeLowestPriority(
@@ -577,13 +588,22 @@ private static String escape(String value) {
                 .replace("\n", "\\n");
     }
 
+    /**
+     * @param droppedContributions how many merged entities, relations and chunks
+     *                             the budget evicted before this prompt fit. The
+     *                             assembler is the only thing that knows the
+     *                             difference between what retrieval selected and
+     *                             what the model was shown, so it is the only
+     *                             thing that can report it.
+     */
     public record PreparedGrounding(
             LightRagGrounding grounding,
             String context,
             String systemPrompt,
             String prompt,
             List<LightRagQueryResult.Reference> references,
-            int inputTokens) {
+            int inputTokens,
+            int droppedContributions) {
 
         public PreparedGrounding {
             Objects.requireNonNull(grounding, "grounding");
@@ -596,6 +616,21 @@ public record PreparedGrounding(
                 throw new IllegalArgumentException(
                         "inputTokens must be non-negative");
             }
+            if (droppedContributions < 0) {
+                throw new IllegalArgumentException(
+                        "droppedContributions must be non-negative");
+            }
+        }
+
+        PreparedGrounding withDroppedContributions(int dropped) {
+            return new PreparedGrounding(
+                    grounding,
+                    context,
+                    systemPrompt,
+                    prompt,
+                    references,
+                    inputTokens,
+                    dropped);
         }
     }
 

components/graph-rag-core/src/main/java/com/orgmemory/graphrag/query/SecureContextBudget.java

@@ -22,6 +22,14 @@ public static SecureContextBudget lightRagCompatibleDefaults() {
         return new SecureContextBudget(6_000, 8_000, 30_000, 200);
     }
 
+    /**
+     * The largest prompt this budget will render: the total less the buffer held
+     * back for the model's own reply.
+     */
+    public int maximumInputTokens() {
+        return maxTotalTokens - safetyBufferTokens;
+    }
+
     public int availableChunkTokens(ContextTokenUsage usage) {
         if (usage.entityTokens() > maxEntityTokens) {
             throw new IllegalArgumentException("entity context exceeds its budget");

components/graph-rag-core/src/test/java/com/orgmemory/graphrag/observability/GraphRagEventSinkTests.java

@@ -1,8 +1,11 @@
 package com.orgmemory.graphrag.observability;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -29,6 +32,34 @@ void rejectsUnboundedFailureDiagnostics() {
                 "provider failed: raw payload"));
     }
 
+    @Test
+    void describesAStageThatMeasuredNoTokensAsAbsentRatherThanAsZero() {
+        assertNull(
+                event(GraphRagEventSink.Outcome.SUCCEEDED, null, null).tokenUsage(),
+                "a zero breakdown would claim a measurement the stage never took");
+    }
+
+    @Test
+    void rejectsABudgetThatCannotBoundAnything() {
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> new GraphRagEventSink.TokenUsage(10, 1, 1, 1, 1, 6, 0, 0),
+                "headroom expressed against a zero budget is a division by zero on every dashboard");
+    }
+
+    @Test
+    void rejectsNegativeTokenCounts() {
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> new GraphRagEventSink.TokenUsage(10, 1, 1, 1, 1, 6, 100, -1));
+    }
+
+    @Test
+    void treatsAnyEvictionAsTruncationBecauseTheAnswerIsAlreadyIncomplete() {
+        assertFalse(new GraphRagEventSink.TokenUsage(10, 1, 1, 1, 1, 6, 100, 0).truncated());
+        assertTrue(new GraphRagEventSink.TokenUsage(10, 1, 1, 1, 1, 6, 100, 1).truncated());
+    }
+
     @Test
     void deliversOneEventToEveryConfiguredBackend() {
         var otel = new RecordingSink();