feat(client): grounded judge context and per-judge diagnostics

AGENTS.md

@@ -203,7 +203,28 @@ The value returned to callers of `config().invoke()`. A `dataclass` with the fol
 | `usage` | `UsageDict` | Normalized token counts (`input`, `output`, `total`). |
 | `track_data` | `TrackData` | Tracking payload from this invocation (run ID, config key, etc.). Carried inside each `JudgeTask` so background judge results are attributed to the originating request. |
 | `judge_results` | `dict[str, JudgeResult]?` | Results from inline judge evaluations. Present when `skip_judges=False` (default) and judges ran. |
-| `judge_tasks` | `list[JudgeTask]?` | Pre-packaged judge tasks. Present (as a list) when `skip_judges=True`. Each task is a serialisable dataclass ready to pass to a background thread running `run_judge(task, handlers)`. `None` when `skip_judges=False`. |
+| `judge_tasks` | `list[JudgeTask]?` | Pre-packaged judge tasks. Present (as a list) when `skip_judges=True`. Each task is a serialisable dataclass ready to pass to a background thread running `run_judge(task, handlers)`, and carries the resolved `judge_context` so the worker injects the identical evidence block. `None` when `skip_judges=False`. |
+| `judge_context` | `JsonValue?` | The value the `judge_context` callback returned, unchanged. `None` when no callback was given or it failed validation. |
+| `judge_diagnostics` | `list[JudgeDiagnostic]?` | Why judges were skipped or failed. `None` when nothing went wrong (never `[]`). |
+
+#### `JsonValue`
+
+`None | bool | int | float | str | list[JsonValue] | dict[str, JsonValue]` — a value that survives a `json.dumps` / `json.loads` round trip unchanged.
+
+#### `JudgeDiagnostic`
+
+One reason a judge produced no result, or a partial one. The strings are shared with the TypeScript SDK on the wire; never rename them. A diagnostic carries no raw exception text.
+
+| Field | Type | Description |
+|---|---|---|
+| `status` | `"skipped" \| "failed"` | Skipped means the judge never ran. |
+| `stage` | `"context" \| "config" \| "provider" \| "parse" \| "track" \| "timeout"` | Where it went wrong. |
+| `code` | `"context_callback_failed" \| "context_invalid_json" \| "context_too_large" \| "judge_duplicate_key" \| "judge_config_failed" \| "judge_provider_failed" \| "judge_response_invalid" \| "judge_tracking_failed" \| "judge_timed_out"` | The specific reason. |
+| `judge_key` | `str?` | The judge this concerns. Absent for context-stage diagnostics. |
+
+#### Stream `done` event
+
+The final event of `config().stream()`: `{"type": "done", "response": str, "usage": UsageDict, "judge_context": JsonValue \| None, "judge_results": dict[str, JudgeResult] \| None, "judge_diagnostics": list[JudgeDiagnostic] \| None}`. Exactly one is yielded, after every `chunk`. Empty results and diagnostics are `None`.
 
 #### `ProviderGraphResponse`
 
@@ -214,6 +235,7 @@ The value returned by `graph().invoke()`. A dataclass with attribute access.
 | `response` | `str` | The final text output (from the last node executed). |
 | `usage` | `UsageDict` | Aggregate token counts across all nodes. |
 | `judge_results` | `dict[str, JudgeResult]?` | Results from a graph-level judge, if configured. |
+| `judge_diagnostics` | `list[JudgeDiagnostic]?` | Diagnostics from the graph-level judge. Graph nodes do not receive a caller judge context in v1. |
 
 #### `ConfigArgs`
 
@@ -225,6 +247,8 @@ Arguments accepted by `config()`.
 | `handler` | `ProviderHandler \| list[ProviderHandler]`? | One handler or an ordered list of handlers. Routing selects the match by provider + mode. |
 | `tool_handlers` | `dict[str, Callable \| NativeTool]?` | Map of tool name → implementation function (or `NativeTool` sentinel). |
 | `registry` | `Registry?` | Registry to source handlers and tools from. Local `handler`/`tool_handlers` take precedence. |
+| `judge_context` | `Callable[[], JsonValue \| Awaitable[JsonValue]]?` | Lazily resolves JSON-safe evidence for the judges. Lazy on purpose: the value does not exist when `config()` is called, and the caller's tools fill it while the primary handler runs. |
+| `judge_timeout_ms` | `int` | Per-judge budget covering config lookup, provider call and parse. Default: `30_000`. |
 | `skip_judges` | `bool`? | When `True`, `invoke()` does not run judges inline. Instead it returns `judge_tasks: list[JudgeTask]` — pre-packaged tasks ready for background thread execution via `run_judge(task, handlers)`. Default: `False`. |
 
 #### `TrackData`
@@ -353,10 +377,11 @@ Returns a `ConfigInstance` with:
 2. Selects the handler by matching on `[config.provider.name, normalized mode]`. Selection priority: (a) exact provider match, (b) wildcard `['*', mode]` fallback for multi-provider adapters (e.g. LangChain). Raises if no matching handler is found.
 3. Invokes the selected handler with the config, user input, tool handlers, variables, and history. The `context` passed to `.invoke()` is automatically merged into `variables` under the key `ldContext`, so templates can reference `{{ldContext.key}}`, `{{ldContext.email}}`, etc. If `history` is provided, it is passed to the handler as the 5th positional argument — messages-mode handlers splice it into the messages array; agent-mode handlers append it to the system prompt.
 4. Emits LaunchDarkly telemetry events: duration (`$ld:ai:duration:total`), outcome (`$ld:ai:generation:success` / `$ld:ai:generation:error`), and token counts (`$ld:ai:tokens:*`).
-5. If `judgeConfiguration` is present:
-   - **Default (`skip_judges=False`):** runs each configured judge inline at its `samplingRate`. Results are returned in `ProviderResponse.judge_results`.
-   - **`skip_judges=True`:** builds serialisable `JudgeTask` objects for each judge (no AI calls). Returns them in `ProviderResponse.judge_tasks`. Pass each task to a background thread running `run_judge(task, handlers)`.
-6. Returns a `ProviderResponse` (always includes `response`, `usage`, and `track_data`).
+5. If a `judge_context` callback was given, resolves it exactly once — immediately after the handler succeeds, before output-format parsing, whether or not any judge is sampled. The value must be acyclic JSON of at most 64 KiB encoded; it is returned unchanged on `ProviderResponse.judge_context` and is never truncated or transformed. An invalid value produces a `context` diagnostic and skips every judge.
+6. If `judgeConfiguration` is present:
+   - **Default (`skip_judges=False`):** runs sampled judges sequentially, in configured order, only the first occurrence of each key. Results are returned in `ProviderResponse.judge_results`. Each judge is isolated: a failure adds one `JudgeDiagnostic` and never erases the primary result or another judge's result. Each judge is bounded by `judge_timeout_ms`, and its reasoning is capped at 4 KiB. The resolved context reaches a judge only through its `message_history` variable, between the lines `UNTRUSTED_ACTUATOR_EVIDENCE_BEGIN` and `UNTRUSTED_ACTUATOR_EVIDENCE_END`: never the primary model, never track data, never a span.
+   - **`skip_judges=True`:** builds serialisable `JudgeTask` objects for each judge (no AI calls). Returns them in `ProviderResponse.judge_tasks`, with the resolved context on each task and any build-step diagnostics on `ProviderResponse.judge_diagnostics`. Pass each task to a background thread running `run_judge(task, handlers)`.
+7. Returns a `ProviderResponse` (always includes `response`, `usage`, and `track_data`).
 
 ### `graph(key, **options)`
 

TELEMETRY-CONTRACT.md

@@ -238,6 +238,10 @@ attributes require `captureContent` / `capture_content`, a handler-factory optio
 not receive. The reasoning is still returned to the caller in `judgeResults` / `judge_results`;
 only the telemetry copy is withheld. Exporting it needs its own opt-in.
 
+The caller-supplied judge context (`judge_context` / `judgeContext`) is never recorded in
+telemetry either: it reaches the judge only through its prompt, and never becomes a span
+attribute, a span event, or track data.
+
 ---
 
 ## 5. Finish reasons

packages/client/agents.md

@@ -123,7 +123,7 @@ Handlers may return any of these — the client normalizes them before emitting
       - On success: emits `$ld:ai:generation:success` + token tracks
       - On error: emits `$ld:ai:generation:error` then re-raises
 3. If `judge_configuration.judges` is present, runs each judge handler (sampled by `sampling_rate`) against the primary response, tracks `evaluation_metric_key`, and emits a `gen_ai.evaluation.result` span event on the judge's `invoke_agent` span (`gen_ai.evaluation.name` / `.score.value` / `.explanation`).
-4. Returns `ProviderResponse`: `{ response: str, usage: UsageDict, track_data: TrackData, judge_results?: dict[str, JudgeResult], judge_tasks?: list[JudgeTask] }`. `judge_results` is populated when `skip_judges=False` (default) and judges ran; `judge_tasks` is populated when `skip_judges=True`.
+4. Returns `ProviderResponse`: `{ response: str, usage: UsageDict, track_data: TrackData, judge_context?: JsonValue, judge_diagnostics?: list[JudgeDiagnostic], judge_results?: dict[str, JudgeResult], judge_tasks?: list[JudgeTask] }`. `judge_context` is the caller callback's value, resolved once after the primary handler and injected only into each judge's `message_history`; `judge_diagnostics` says why a judge was skipped or failed. `judge_results` is populated when `skip_judges=False` (default) and judges ran; `judge_tasks` is populated when `skip_judges=True`.
 
 ---
 

packages/client/src/launchdarkly_ai_server/__init__.py

@@ -33,7 +33,15 @@
     init_evaluations,
 )
 from .graph import GraphInstance, graph, resolve_graph
-from .judges import build_judge_tasks, run_judge, run_judges
+from .judges import (
+    BuildJudgeTasksResult,
+    JudgeContextResolution,
+    RunJudgesResult,
+    build_judge_tasks,
+    resolve_judge_context,
+    run_judge,
+    run_judges,
+)
 from .lifecycle import (
     extract_variation,
     get_client,
@@ -65,6 +73,8 @@
     HandlerStreamEvent,
     InitClientOptions,
     InputTokenDetails,
+    JsonValue,
+    JudgeDiagnostic,
     JudgeResult,
     JudgeRunResult,
     JudgeTask,
@@ -125,6 +135,10 @@
     "HandlerResult",
     "HandlerStreamEvent",
     "InitClientOptions",
+    "BuildJudgeTasksResult",
+    "JsonValue",
+    "JudgeContextResolution",
+    "JudgeDiagnostic",
     "JudgeResult",
     "JudgeRunResult",
     "JudgeTask",
@@ -209,7 +223,9 @@
     # judges
     "build_judge_tasks",
     "run_judge",
+    "resolve_judge_context",
     "run_judges",
+    "RunJudgesResult",
     # client
     "config",
     "ConfigInstance",

packages/client/src/launchdarkly_ai_server/client.py

@@ -1,16 +1,23 @@
 from __future__ import annotations
 
 import json
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import Any
 
 from .conversation import bind_conversation_id
-from .judges import build_judge_tasks, run_judges
+from .judges import (
+    DEFAULT_JUDGE_TIMEOUT_MS,
+    build_judge_tasks,
+    resolve_judge_context,
+    run_judges,
+)
 from .lifecycle import extract_variation
 from .registry import resolve_handlers, resolve_tools
 from .tracking import execute_and_stream, execute_and_track
 from .types import (
     AiConfigRep,
+    JsonValue,
+    JudgeDiagnostic,
     LDContext,
     NativeTool,
     ProviderHandler,
@@ -52,12 +59,16 @@ def __init__(
         tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None,
         registry: Any,  # Registry | None
         skip_judges: bool = False,
+        judge_context: Callable[[], JsonValue | Awaitable[JsonValue]] | None = None,
+        judge_timeout_ms: int = DEFAULT_JUDGE_TIMEOUT_MS,
     ) -> None:
         self._key = key
         self._handler = handler
         self._tool_handlers = tool_handlers
         self._registry = registry
         self._skip_judges = skip_judges
+        self._judge_context = judge_context
+        self._judge_timeout_ms = judge_timeout_ms
 
     def _normalize_handlers(self) -> list[ProviderHandler] | None:
         if self._handler is None:
@@ -102,6 +113,10 @@ async def invoke(
         usage: dict[str, int] = result["usage"]
         track_data = result["track_data"]
 
+        # Freeze the caller's context the moment the primary handler succeeds, before output
+        # parsing. Sampling controls judge execution, never this boundary.
+        context_resolution = await resolve_judge_context(self._judge_context)
+
         parsed_response = _resolve_output_format_response(
             raw_response,
             config.get("outputFormat") if isinstance(config, dict) else None,
@@ -116,35 +131,54 @@ async def invoke(
         usage_obj = to_usage_dict(usage)
 
         if self._skip_judges:
-            judge_tasks = await build_judge_tasks(
+            build_result = await build_judge_tasks(
                 config=config,
                 user_context=context,
                 handler=handler,
                 handlers=resolved_handler_list,
                 llm_response=llm_str,
                 base_track_data=track_data,
+                context_resolution=context_resolution,
             )
             return ProviderResponse(
                 response=parsed_response,
                 usage=usage_obj,
-                judge_tasks=judge_tasks,
+                judge_context=build_result.judge_context,
+                judge_diagnostics=build_result.judge_diagnostics or None,
+                judge_tasks=build_result.judge_tasks,
                 track_data=track_data,
             )
 
-        judge_results = await run_judges(
-            config=config,
-            user_context=context,
-            handler=handler,
-            handlers=resolved_handler_list,
-            user_input=user_input,
-            llm_response=llm_str,
-            base_track_data=track_data,
-            tool_handlers=resolved_tools,
+        diagnostics: list[JudgeDiagnostic] = (
+            [context_resolution.diagnostic]
+            if context_resolution.diagnostic is not None
+            else []
         )
+        if not context_resolution.failed:
+            judge_run = await run_judges(
+                config=config,
+                user_context=context,
+                handler=handler,
+                handlers=resolved_handler_list,
+                user_input=user_input,
+                llm_response=llm_str,
+                base_track_data=track_data,
+                tool_handlers=resolved_tools,
+                judge_context=context_resolution.judge_context,
+                judge_context_json=context_resolution.serialized,
+                judge_timeout_ms=self._judge_timeout_ms,
+            )
+            diagnostics.extend(judge_run.judge_diagnostics)
+            judge_results = judge_run.judge_results
+        else:
+            judge_results = {}
+
         return ProviderResponse(
             response=parsed_response,
             usage=usage_obj,
-            judge_results=judge_results if judge_results else None,
+            judge_context=context_resolution.judge_context,
+            judge_diagnostics=diagnostics or None,
+            judge_results=judge_results or None,
             track_data=track_data,
         )
 
@@ -206,25 +240,39 @@ async def _stream_events(
 
         if done_event:
             track_data = done_event.get("track_data", {})
-            judge_results = (
-                {}
-                if self._skip_judges
-                else await run_judges(
-                    config=config,
-                    user_context=context,
-                    handler=handler,
-                    handlers=resolved_handler_list,
-                    user_input=user_input,
-                    llm_response=done_event.get("response", ""),
-                    base_track_data=track_data,
-                    tool_handlers=resolved_tools,
-                )
-            )
+            judge_results: dict[str, Any] = {}
+            diagnostics: list[JudgeDiagnostic] = []
+            judge_context: JsonValue | None = None
+
+            if not self._skip_judges:
+                context_resolution = await resolve_judge_context(self._judge_context)
+                judge_context = context_resolution.judge_context
+                if context_resolution.diagnostic is not None:
+                    diagnostics.append(context_resolution.diagnostic)
+                if not context_resolution.failed:
+                    judge_run = await run_judges(
+                        config=config,
+                        user_context=context,
+                        handler=handler,
+                        handlers=resolved_handler_list,
+                        user_input=user_input,
+                        llm_response=done_event.get("response", ""),
+                        base_track_data=track_data,
+                        tool_handlers=resolved_tools,
+                        judge_context=context_resolution.judge_context,
+                        judge_context_json=context_resolution.serialized,
+                        judge_timeout_ms=self._judge_timeout_ms,
+                    )
+                    judge_results = judge_run.judge_results
+                    diagnostics.extend(judge_run.judge_diagnostics)
+
             yield {
                 "type": "done",
                 "response": done_event.get("response", ""),
                 "usage": done_event.get("usage"),
-                "judge_results": judge_results if judge_results else None,
+                "judge_context": judge_context,
+                "judge_results": judge_results or None,
+                "judge_diagnostics": diagnostics or None,
             }
 
 
@@ -235,6 +283,8 @@ def config(
     tool_handlers: dict[str, Callable[..., Any] | NativeTool] | None = None,
     registry: Any = None,
     skip_judges: bool = False,
+    judge_context: Callable[[], JsonValue | Awaitable[JsonValue]] | None = None,
+    judge_timeout_ms: int = DEFAULT_JUDGE_TIMEOUT_MS,
 ) -> ConfigInstance:
     """
     Creates a ``ConfigInstance`` bound to *key*. Accepts a single handler or a
@@ -247,11 +297,19 @@ def config(
     ``.invoke()`` / ``.stream()``. When set, ``invoke()`` returns
     ``judge_tasks: list[JudgeTask]`` — pre-packaged tasks ready for a background
     thread calling ``run_judge(task, handlers)``.
+
+    ``judge_context`` is a callback returning JSON-safe evidence for the judges. It is lazy on
+    purpose: the value does not exist yet when ``config()`` is called, and the caller's tools
+    fill it while the primary handler runs. It resolves exactly once per request, right after
+    the primary handler succeeds, and reaches the judges only through their ``message_history``
+    variable. ``judge_timeout_ms`` bounds one judge's config lookup, provider call and parse.
     """
     return ConfigInstance(
         key=key,
         handler=handler,
         tool_handlers=tool_handlers,
         registry=registry,
         skip_judges=skip_judges,
+        judge_context=judge_context,
+        judge_timeout_ms=judge_timeout_ms,
     )