Unify Agent Gateway request context and terminal observation - #7860

I reviewed this end to end. The architecture is right and the test suite is unusually strong for a change this size — but I found three things I would like resolved before this lands, plus a stale description.

What I verified as correct

• AgentTraceContext.headers is a real allow-list: only traceparent/tracestate, regex-validated, zero-trace-id and zero-span-id rejected, tracestate bounded to 512 chars / 32 entries / unique keys. Combined with Context.root().with(Span.wrap(...)) in both AiOutboundTraceContext and OpenTelemetryAgentMcpCallTracer.start, baggage and unrelated context genuinely cannot leak. GATEWAY_REACTOR_CONTEXT_KEY matches TracingConstants.REACTOR_CONTEXT_KEY in the tracing plugin, so the parent is actually found.
• WebClientAiProxyTransport strips traceparent/tracestate/baggage twice — on the mutated default headers and again after request.getHeaders().forEach(headers::put) — then sets only validated values. That is the correct order; a plain put would have let configured stale telemetry win.
• AiUpstreamUsageTap never subscribes; it is a doOnNext side tap on the same body, so the single-subscription guarantee holds. [DONE] stops parsing, doFinally clears, eventBytes bounds each SSE event, MAX_BYTES bounds the non-streaming body, and unavailable() marks UNSUPPORTED/LIMIT_EXCEEDED instead of guessing. Replacing the snapshot per event rather than summing cumulative SSE counters is right, and AiProxyObservationTest.recordsFourPrimaryAttemptsThenFallbackAndOnlyItsUsageAtLogicalEnd pins the "last attempt owns the terminal usage" semantic precisely.
• AgentGatewayPluginOrderTest.shouldNotSilentlyRepairAnUnmigratedCachedGateway is the best test in the PR: it proves a cached sort = 198 keeps the gateway after the limiter, which is exactly why 2.7.2-agent-gateway-order.sql is mandatory rather than optional. shouldKeepTheExistingInstallationMigrationExplicitAndDefaultOnly pinning the migration text (and asserting it never touches enabled/config) is a good guard against someone "helpfully" widening it later.
• 169 is genuinely free. I parsed all 38 plugin rows in db/init/mysql/schema.sql: 170 is already shared by aiPrompt and loggingRocketMQ, which is what the doc claims.

Blocking: the ordering flip changes precedence for more plugins than the doc admits

docs/agent-gateway-governance-order.md names RocketMQ logging and sensitiveWord as the plugins a matched MCP request now bypasses. It does not mention mcpServer (sort 180).

That is the most directly affected plugin. Before this change mcpServer at 180 ran before agentGateway at 198, so it shadowed the gateway for overlapping selectors. After this change the gateway shadows it instead. Either direction is "one wins", but the direction is now inverted, and shouldRetainExplicitCustomOrdersEvenWhenUnsafe confirms nothing auto-repairs a deployment that has both enabled. Anyone running mcpServer alongside agentGateway on the same path will silently start routing to a different implementation on upgrade.

Please:

1. Add mcpServer (and loggingKafka at 180, loggingElasticSearch at 190, aiTokenLimiter at 171) to the bypass list in the governance doc — the current wording understates the blast radius.
2. Add one test asserting a matched MCP request does not reach a later-ordered plugin, so the "gateway terminates MCP" contract is pinned rather than inferred.
3. Call the precedence flip out explicitly in the release notes. This is a routing change, not just a default value.

Blocking: outcome() and error() disagree, and one of them is unreachable

AiObservationSupport.outcome maps every ShenyuException to REJECTED and does not walk the cause chain. AiRequestObservation.error walks up to 16 causes looking for TimeoutException and maps it to TIMEOUT. Two classifiers for the same field, with different rules.

The consequence is concrete in AiProxyPlugin.doExecute:

executeRequest(exchange, selector)
        .doOnError(error -> observation.failure(AiObservationSupport.outcome(error)))
        .onErrorResume(error -> handleFailure(exchange.getResponse(), error))
        .doOnError(observation::error)      // never sees the primary error

onErrorResume consumes the error, so the second doOnError only fires when handleFailure itself fails. For a standalone AI Proxy request (no gateway owner, inherited == null) a TimeoutException wrapped in a ShenyuException is therefore reported as REJECTED — a server-side timeout classified as client fault, with the cause-chain logic that was written for exactly this case never running.

Suggested fix: move observation.error(error) before onErrorResume (and drop the redundant failure() call, or keep failure() for the layer-specific override and make it first-wins), and stop mapping bare ShenyuException to REJECTED — it is used broadly for internal failures. AiProxyException.Kind.REQUEST is the signal that actually means "client's fault"; ShenyuException alone is not.

Blocking: AgentMcpTerminalObservation.failure is last-wins, which contradicts its own comment

In AgentMcpHttpHandler.handleFailure, the recovery-write branch does:

.onErrorResume(writeFailure -> {
    recordFailure(terminal, writeFailure);   // overwrites the original classification

failure(...) writes the single outcome field unconditionally. So a 403 authorization rejection whose error body fails to write ends up classified SERVER_ERROR — the authorization signal is erased by a transport detail. The inline comment ("Recovery writes can fail independently of the execution error classified above") reads as if both facts are retained, but only one can be.

Either make it first-wins (keep the first classification, since that is why the request failed) or carry the write failure in a separate field. And please extend recoveredErrorsKeepTheirOriginalClassification to cover the write-failure case — right now that test only exercises execution errors.

Blocking: the description is stale

It still says "This is a draft integration candidate", "Please do not merge this PR before #7827" (#7827 merged and is in master), and states verification against base 81a85f3a4 — while the head now carries three master merges (3c09a392, dd6d8786, 4a9544d9) plus a CI re-run. The diff itself looks correctly reconciled (no leftover #7827 commits), so this is just the text. Anyone landing on this PR today would think it is still blocked.

Nits (non-blocking)

• AiRequestObservation.finish freezes a still-active attempt with a hardcoded AiCallOutcome.ERROR. Every current path calls Attempt.finish via doOnSuccess/doOnError/doOnCancel, so it is defensive today — but if a future path forgets, a successful request emits an attempt event marked ERROR. terminalOutcome(active.status, failure) would degrade more honestly than a hardcoded constant.
• db/upgrade/2.7.2-agent-gateway-order.sql is dialect-generic while every other upgrade script is per-dialect. I checked it does work on all three (Oracle's namespace_plugin_rel (... sort ...) is unquoted and resolves to uppercase, SORT is not reserved in MySQL/PG/Oracle), so this is a consistency question, not a bug — but a maintainer will ask. A one-line "verified on MySQL/PostgreSQL/Oracle" comment would close it.
• UpstreamCircuitBreakerTest moving the wait window from 100ms to 100s removes the race the test was written to catch (the breaker closing mid-test while 32 threads hammer it). The de-flake is understandable, but the original intent is gone. TraceContinuityTest setting HTTP_TIME_OUT is the same category. Both are unrelated to this feature; worth a sentence in the description, or a separate PR.
• AiUpstreamUsageTap.attach reaches stopped = true through doFinally → clear(). Re-subscribing the same AiUpstreamResponse would then silently observe nothing. Unreachable today because each attempt rebuilds the response object, but it is a trap for the next refactor.
• AgentRequestContext has no equals/hashCode, so AiAttemptEvent/AiTerminalEvent record equality compares contexts by identity. Fine if you never rely on it, but assertNotEquals in the tests depends on that behaviour implicitly.

The recordFailure overwrite and the two-classifier split are the two I would actually hold the merge on; the rest is documentation and cleanup. Happy to re-review as soon as those are in — CI is fully green at 2324714a (55 checks, 0 failures) and the branch is mergeable, so nothing else is in the way.

Review — unify Agent Gateway request context and terminal observation

This is a large but well-structured change (62 files, +4858/-129). The core idea — one subscription-local AiRequestObservation owned by whoever creates it, with everything else attaching to it — is sound, and the defensive posture (bounded snapshots, no throwable/message retention, observer callbacks wrapped in try/catch, AgentTraceContext header validation) is better than what most observability PRs do.

I have one blocking item (the migration story) and a list of non-blocking items below.

Blocking: the released 2.7.1-upgrade-2.7.2-*.sql scripts must not be edited

v2.7.2 is already tagged and released, so db/upgrade/2.7.1-upgrade-2.7.2-{mysql,pg,oracle,ob,og}.sql are shipped, immutable migration artefacts. This PR changes the seeded sort from 198 to 169 inside all five of them:

-INSERT INTO `plugin` VALUES ('68', 'agentGateway', null, 'Ai', 198, 0, ...);
+INSERT INTO `plugin` VALUES ('68', 'agentGateway', null, 'Ai', 169, 0, ...);

Two problems:

1. It has no effect on anyone who already upgraded. Their plugin / namespace_plugin_rel rows still say 198, which is why the standalone 2.7.2-agent-gateway-order.sql exists — but then the edit is redundant for them and harmful for the next group.
2. It makes the 2.7.2 upgrade non-deterministic. Someone upgrading 2.7.1 → 2.7.2 before this merge gets 198; someone doing the same upgrade after this merge gets 169. Same release, same input, different schema state. That is the one thing migration scripts exist to prevent.

AgentGatewayPluginOrderTest.shouldKeepGatewayAndNamespaceSeedsInSync currently enforces the edit across those five files, so the test locks in the inconsistency.

Suggested fix: revert the five 2.7.1-upgrade-2.7.2-*.sql hunks, and put the 198 → 169 seed change in the 2.7.2-upgrade-2.7.3-{h2,mysql,ob,og,oracle,pg}.sql files that already exist for this release, following the same per-dialect pattern as every other migration in the repo. Keep the fresh-install seeds (db/init/*, h2 schema.sql) at 169 — those are correct as changed. Then narrow the test to the fresh-install scripts plus the new 2.7.2→2.7.3 scripts.

Related: db/upgrade/2.7.2-agent-gateway-order.sql is a single "common SQL" file with a non-standard name, while every other migration is X-upgrade-Y-<dialect>.sql. As written it will not be picked up by anything automated, and it silently assumes the four dialects accept the same DDL/DML subset. If it is meant to be an operator runbook, it belongs under docs/; if it is meant to be a migration, it belongs in the per-dialect 2.7.2→2.7.3 files. (The two UPDATEs themselves are fine and correctly refuse to touch enabled / config.)

Non-blocking

AiRequestObservation.error() clobbers a more specific classification

AiProxyPlugin and AiProxyEngine classify with the bounded mapper:

.doOnError(error -> observation.failure(AiObservationSupport.outcome(error)))

but AgentGatewayPlugin.handleLlm uses the coarser one, and it runs after (it wraps the whole downstream chain):

.doOnError(observation::error)

error(Throwable) sets failure = ERROR (or TIMEOUT) unconditionally, so a REJECTED outcome already recorded for a 400-class failure is overwritten with ERROR by the time the terminal event is built. Suggest either making error() only fill when failure == null, or having the gateway use the same AiObservationSupport.outcome(...) classifier so the two layers cannot disagree.

selectedModel() re-parses the request body on every request

private String selectedModel(final AiUpstreamRequest request) {
    JsonNode payload = OBJECT_MAPPER.readTree(request.getBody());
    ...
}

The body has just been serialized by prepareOutgoingRequest / provider adaptation, and now gets parsed again purely to fill a metric label. For a multimodal request (base64 images) that is a second full JSON parse of a multi-megabyte payload on the hot path. The model is already known inside prepare(...) — thread it through PreparedExchange instead of re-reading the bytes.

AgentRequestContext has no equals / hashCode / toString

It is an immutable value object used as a Reactor context value, an exchange attribute, and an event field. Without equals every test has to compare eleven getters, and any future de-duplication or assertion on context identity silently degrades to reference equality. Adding the three methods (with toString emitting only ids, not subject) would help; there is nothing secret in the current field set except subject/tenant, which you may want to elide.

Two sources of truth for the request context

AgentRequestContext.REACTOR_KEY is written once per scope (AiProxyPlugin, AiProxyEngine.execute, AgentGatewayPlugin), but AiRequestObservation.updateContext(...) keeps enriching its own copy afterwards via withIdentity, withTarget, withConfigurationVersion. A consumer that reads the Reactor key therefore sees a context without identity/target while the terminal event has them. Either make the Reactor key carry the observation (single source of truth) or document explicitly that the Reactor key is the initial snapshot and observation.getContext() is authoritative.

AiUpstreamUsageTap can downgrade REPORTED to UNKNOWN

AiUsage.Status quality = ... Objects.isNull(prompt) && Objects.isNull(completion) && Objects.isNull(total)
        ? AiUsage.Status.UNKNOWN : AiUsage.Status.REPORTED;
observation.usage(quality == INVALID ? ... : new AiUsage(prompt, completion, total, quality));

The "replace, never sum" rule is right for cumulative SSE counters, but it also means a later chunk carrying "usage": {} (several OpenAI-compatible servers emit exactly that on the final chunk) overwrites a previously captured REPORTED value with UNKNOWN. Consider ignoring an event whose usage object yields no counts at all once a REPORTED snapshot exists.

shenyu-common relies on a transitive reactor-core

AiRequestObservation imports reactor.core.publisher.SignalType, but shenyu-common/pom.xml only declares reactor-netty-core and reactor-netty-http. It compiles today via transitivity; please add an explicit reactor-core dependency (or move the class to a reactive module) so a future dependency cleanup doesn't break shenyu-common — which the admin also depends on.

Smaller notes

• AiProxyEngine now has five constructors, including (…, AiProxyRetryProperties) and (…, Duration); passing null as the fifth argument is now ambiguous. A builder or a single (…, AiProxyRetryProperties, AiCallObserver) plus static factories would read better.
• AiUpstreamUsageTap.readLine() allocates a fresh CharsetDecoder per line; on a long SSE stream that is a lot of short-lived garbage. A single reusable decoder, reset() per line, would do.
• AiOutboundTraceContext.headers injects the gateway's own span context rather than creating a CLIENT child span. That is a deliberate simplification and the javadoc says so, but the docs should state it, otherwise people will look for a client span that does not exist.
• AgentMcpTerminalObservation.finish calls observer.record(event) without a null check; the current caller guards it, but a defensive Objects.isNull(observer) avoids a swallowed NPE per request if that ever changes.
• RELEASE-NOTES.md — same note as on #8034: release notes are normally assembled by the release manager, please keep that hunk in the PR description instead.

What I checked and found correct

• AiRequestObservation.finish freezes under the monitor and calls the observer outside the lock; Attempt.finish does the same. No observer callback can run while holding the monitor, so a slow exporter cannot serialize requests or deadlock.
• Attempt.freeze is idempotent (if (completed) return null), so a retry, a fallback and a cancel racing on the same attempt produce exactly one event.
• startAttempt after completion returns an inert Attempt instead of throwing — the right choice for a metrics path.
• terminalOutcome maps 504 → TIMEOUT, 429 → RATE_LIMITED, 4xx → REJECTED, 5xx → ERROR, and SignalType.CANCEL → CANCELLED before any status mapping, so a cancelled write is never counted as delivered.
• AgentTraceContext.headers validates traceparent against the W3C v00 shape, rejects the all-zero trace-id and span-id, rejects uppercase hex, and bounds tracestate (512 chars, 32 members, 256-char values, unique keys, single =). Invalid telemetry is dropped rather than failing the request.
• AiOutboundTraceContext uses deferContextual and only explicit Reactor parents, never Context.current(), and it strips stale traceparent/tracestate/baggage from both defaultHeaders and the per-request headers before injecting. RequestScopedMcpClient correctly deletes the header key from the shared handshake and only forwards trace headers for tools/call.
• AiUpstreamUsageTap is a pure side tap (doOnNext + doFinally), never subscribes twice, bounds both per-event bytes and line length, enables STRICT_DUPLICATE_DETECTION and FAIL_ON_TRAILING_TOKENS, and handles CRLF split across chunk boundaries.
• AgentGatewayPlugin.handleLlm resolves identity through AgentMcpSecurityResolver and only recovers identity errors — a downstream write failure keeps its original error path and cannot trigger a second proxy request.
• AgentGatewayPluginOrderTest is a genuinely useful contract test (it is what makes the seed drift visible); the regex-based seed check handles the three different INSERT shapes across dialects.

Overall this is good work. The migration placement is the only thing I would not merge as-is.

Re-reviewed the latest changes against my previous change requests. The blocking items are addressed: released 2.7.1 -> 2.7.2 scripts are no longer edited; the order migration is in the six 2.7.2 -> 2.7.3 dialect scripts and preserves custom orders/config/enabled; the governance documentation explicitly covers the mcpServer precedence change; wrapped timeout classification and first-failure MCP terminal classification have regression coverage.

No remaining blocking finding in the reviewed change. 216 targeted tests passed across common context/migration contracts, AI observation/plugin behavior, MCP terminal/call observation, web plugin ordering, loopback HTTP governance, starter wiring, and tracing.

Important validation baseline: the isolated worktree contained PR head a34127b plus a clean, uncommitted merge of current master aab453f. The raw historical head fails Maven model loading because its inherited AI PII module parents still name 2.7.2-SNAPSHOT; the master merge fixes those baseline versions. No PR files were manually patched and no branch was pushed. Refresh the branch onto master before merging so these validation conditions are represented in CI.

Validation used isolated worktrees, the Maven wrapper, JDK 21 with the project compiler settings, and targeted tests. Checkstyle passed; RAT and JaCoCo were skipped. This does not claim a full reactor, e2e, or live deployment test run.

Reviewed 6b1207c.

Re-reviewed the current-master branch. The previously requested released-migration preservation and dialect-specific upgrade fixes are present, the mcpServer ordering impact is documented, and the observation/terminal ownership regressions pass. The source now includes the baseline merge that was only present locally during the prior review. No remaining blocker found in the reviewed change.

Targeted test report: 367 cases, 0 failures, 0 errors, 0 skipped.

Validation used ./mvnw, Amazon Corretto 21 and project target Java 17. Targeted Checkstyle, separately executed Apache RAT and git diff --check pass. No full-repository or Docker E2E pass is claimed; required CI must succeed before merge.