The previous wave-4 fix added `host == "::1"` for IPv6 loopback support,
but verification against `mcr.microsoft.com/dotnet/sdk:9.0` showed that
`new Uri("http://[::1]:8000/").Host` returns `"[::1]"` WITH brackets, not
`"::1"`. The string-equality check therefore never matched and IPv6
loopback detection was silently broken.
Switch to `Uri.IsLoopback`, which is the framework-provided helper that
robustly identifies all loopback variations: `localhost`, the entire
`127.0.0.0/8` range, `::1` in any bracketed form, and IPv4-mapped IPv6
loopback (`::ffff:127.0.0.1`). This sidesteps `Uri.Host`'s bracket-
formatting quirks entirely.
The explicit `0.0.0.0` (unspecified address, kept for back-compat) and
`aimock` (Docker-compose service name) checks are retained because they
are not loopback addresses.
ResolveApiKey and ResolveEndpoint previously cascaded through env var →
configuration sources using the `??` operator, which only short-circuits
on null. When `OPENAI_API_KEY=""` or `OPENAI_BASE_URL=""` was set in the
environment, the empty string was returned by the cascade WITHOUT
consulting the configuration fallbacks (`configuration["OPENAI_API_KEY"]`,
`configuration["GitHubToken"]`, `configuration["OPENAI_BASE_URL"]`),
masking the configured values entirely. The downstream
`IsNullOrWhiteSpace` check would then push the resolver into the
mock-or-throw path even when a valid key was configured.
Replace the `??` chains with a `FirstNonBlank` helper that treats both
null and whitespace-only candidates as absent, so empty env vars fall
through to the configuration fallbacks as intended.
Wave-1 widened the catch list in BeautifulChatAgent.GenerateA2ui from
specific upstream exceptions to a blanket catch (Exception ex). That
inadvertently swallowed InvalidOperationException thrown by
ApiKeyResolver.ResolveApiKey, which is the intentional fail-fast contract
for misconfigured deployments. With the blanket catch in place, a missing
or invalid API key produced a generic "unexpected_error" structured
response to end users while the operator received no clear startup-level
signal — exactly the silent-degradation failure mode the fail-fast was
designed to prevent.
Replace the blanket catch with specific catches for the two
runtime-recoverable exception types the secondary tool caller can throw
after the transport/SDK/cancellation handlers:
- JsonException — upstream returned malformed JSON
- KeyNotFoundException — upstream JSON missing required fields
(defensive; TryGetProperty guards most paths after wave-1)
Both map to the existing "upstream_malformed" structured-error taxonomy.
All other exceptions, including configuration errors like
ApiKeyResolver's InvalidOperationException, now propagate so a
misconfigured deployment fails loudly at the operator layer instead of
limping along behind a generic user-facing message.
ASP.NET Core hands control back to the middleware once the endpoint handler
completes via `await _next(context)`, but for streaming endpoints (AG-UI uses
`IAsyncEnumerable`/SSE) the response delegate continues writing to the response
body — and may issue downstream OpenAI calls — AFTER `_next` returns. The
previous `finally` clause reset `AimockHeaderContext` to an empty dictionary at
that point, which caused `AimockHeaderPolicy` to silently drop aimock headers
on those streaming-tail calls. That defeated the entire purpose of the
middleware for the exact codepath it was built to serve.
The `finally`-clear was also unnecessary for isolation: `AsyncLocal<T>` is
scoped to the current execution context, so each ASP.NET request gets its own
slot automatically. There is no cross-request leakage to defend against.
Remove the `try`/`finally` and just call `AimockHeaderContext.Set(headers)`
followed by `await _next(context)`. Add a comment documenting why no reset is
performed.
ApiKeyResolver.IsMockEndpoint compared Uri.Host against "[::1]", but
System.Uri.Host strips the surrounding brackets and returns "::1" for
inputs like http://[::1]:8000/. The bracketed comparison was therefore
dead code: developers running aimock on IPv6 loopback got the fail-fast
error instead of the intended mock-key fallback.
Change the literal to "::1" and update the inline + docstring comments
to reflect the bracket-less form.
- BeautifulChatStateSnapshotAgent.RunCoreAsync now coalesces a null
response.Messages to an empty array before constructing the new
List<ChatMessage>. Previously, an inner agent that returned no
messages would cause `new List<ChatMessage>(null)` to throw
ArgumentNullException, masking the real upstream behaviour.
- GenerateA2ui's empty-content branch now returns the canonical
BeautifulChatA2ui.StructuredError shape (error / message /
remediation / errorId) like every other error path in the method,
instead of an ad-hoc { error, errorId } object. The frontend
expects the structured shape with remediation guidance.
Program.cs#CreateOpenAiClient previously read the OpenAI endpoint solely from
the OPENAI_BASE_URL environment variable and fell back to
ApiKeyResolver.DefaultOpenAiEndpoint. The secondary tool-calling HTTP client
(A2uiSecondaryToolCaller) instead routes through ApiKeyResolver.ResolveEndpoint,
which checks env, then configuration[OPENAI_BASE_URL] (appsettings.json /
user-secrets), then the default.
When OPENAI_BASE_URL was supplied only via configuration (e.g. user-secrets in
local dev), the primary client silently dialed the public Azure-hosted models
endpoint while the secondary client dialed the configured aimock. The two
endpoints diverged with no visible signal, breaking aimock-routed flows and
masking misconfigurations.
Switch CreateOpenAiClient to call ApiKeyResolver.ResolveEndpoint so both
clients share a single source of truth. Logging is preserved and now reports
which source supplied the endpoint (env, configuration, or default).
When iterating IHeaderDictionary's underlying store yields case-variant
duplicates (e.g., a misbehaving proxy injecting both `X-Foo` and `x-foo`),
the previous GroupBy + ToDictionary path silently kept only the first value
and discarded the rest. For aimock context routing — where header semantics
drive fixture selection — silent drops are a debugging nightmare.
We still keep `.First()` because HTTP defines no canonical merge for
case-variant collisions across distinct keys (the comma-join rule only
applies when keys are ASCII-equal). Instead, when a group has more than
one entry, we emit a structured warning naming the key, the kept value,
and the number of dropped variants, so operators can spot the upstream
misbehavior in logs.
Constructor now takes `ILogger<AimockHeaderMiddleware>` — ASP.NET's
middleware activator injects it automatically via `UseMiddleware<T>()`.
Two surgical hardening fixes to ApiKeyResolver:
1. IsMockEndpoint: drop host.StartsWith("aimock.", ...) and
host.EndsWith(".aimock", ...). The StartsWith match is exploitable
— an attacker-registered domain like aimock.attacker.example.com
would be classified as a mock endpoint and bypass the fail-fast,
silently returning the mock key. The EndsWith match is unused
noise (no .aimock TLD exists). Only exact, well-known mock hosts
("localhost", "127.0.0.1", "0.0.0.0", "[::1]", "aimock") are
accepted. IPv6 loopback [::1] is added so dual-stack dev paths
are still covered.
2. ResolveApiKey: change !string.IsNullOrEmpty(apiKey) to
!string.IsNullOrWhiteSpace(apiKey). A whitespace-only key such as
" " would otherwise be accepted as valid, bypassing the mock-key
/fail-fast logic entirely and sending a bogus key upstream.
- RunCoreAsync: avoid mutating AgentResponse.Messages in place. Reassign
Messages to a fresh List<ChatMessage> so the path is safe regardless of
whether the inner agent returns a mutable list or an immutable
IReadOnlyList wrapper. Prevents NotSupportedException at runtime; all
other response metadata (AgentId, ResponseId, Usage, RawRepresentation,
...) is preserved via the property setter.
- RunCoreStreamingAsync: only emit the trailing todos snapshot if the
inner stream completed normally. Wrap the inner enumerator in
try/finally with a streamCompletedNormally flag; on early exit
(throw or cancellation) log a warning and skip the snapshot rather
than emitting potentially stale state to the frontend. The trailing
yield lives outside the try block (idiomatic C# pattern since `yield`
cannot live inside `try { } catch { }`).
The previous implementation used Contains() substring matching against the
full endpoint URL, which is exploitable. An attacker-controlled endpoint
such as https://attacker.example.com/aimock-decoy or
https://api.openai.com/?env=localhost would be classified as a mock and
bypass the fail-fast guard, silently returning the sk-mock-local key for
what is actually a non-mock destination.
Parse the URL and inspect only the host component, matching exact dev
hosts (localhost, 127.0.0.1, 0.0.0.0, aimock) plus aimock subdomains.
Path, query, and arbitrary subdomain segments containing "aimock" or
"localhost" no longer trigger the mock fallback.
The header propagation chain (middleware -> context -> policy) had two bugs
that combined to threaten D5/D6 header-forwarding:
1. AimockHeaderMiddleware.ToDictionary used the default ORDINAL case-sensitive
comparer. ASP.NET's IHeaderDictionary is case-insensitive, but iterating
the underlying store can yield case-variant duplicates (e.g., a misbehaving
proxy injecting both `X-Foo` and `x-foo`). Default-comparer ToDictionary
throws ArgumentException on duplicates and fails the request.
2. AimockHeaderContext.Set lowercased all keys via ToLowerInvariant.
AimockHeaderMiddleware captured original case; the context then mutated
the casing; AimockHeaderPolicy.TryGetValue then matched against whatever
case the OpenAI SDK happened to use. This is inconsistent and aimock
fixture matching can be case-sensitive depending on configuration.
Canonical strategy: preserve original header casing as captured by the
middleware, but compare case-insensitively throughout via
StringComparer.OrdinalIgnoreCase. The middleware now also groups variants
defensively so duplicate keys cannot blow up the dictionary build.
AimockHeaderPolicy's add-if-absent TryGetValue then works correctly under
case-insensitive comparison without any further changes.
- Replace silent `return null` paths in A2uiSecondaryToolCaller with
TryGetProperty guards that each emit a structured LogWarning naming
the exact missing/unexpected field (choices, message, tool_calls,
function, name mismatch, arguments). Callers can now tell why a
design-tool call produced no content.
- Add ILogger parameter to GetDesignToolArgumentsAsync and pass
BeautifulChatAgent._logger from the single call site so warnings
flow into the existing log stream.
- Log the response body (truncated to 1 KB) at LogWarning before
EnsureSuccessStatusCode throws, so upstream error payloads survive
the throw and reach operators.
- Extract the OPENAI_API_KEY/GitHubToken/sk-mock-local fallback chain
from Program.cs and A2uiSecondaryToolCaller.cs into a new
ApiKeyResolver helper. Both call sites now share one implementation.
- ApiKeyResolver fails fast with InvalidOperationException + LogCritical
when no real key is present and OPENAI_BASE_URL is not an
aimock/localhost endpoint, so misconfigured prod deploys cannot
silently send sk-mock-local to a real LLM provider. The silent
mock-key fallback is preserved for aimock/localhost dev endpoints.
- RunCoreAsync: append todos snapshot DataContent to AgentResponse so non-streaming
callers receive the same state mirror that RunCoreStreamingAsync already emits.
- ManageTodos: defensive-copy each incoming todo before storing (mirrors the
symmetry of GetTodosSnapshot) so callers cannot mutate our backing list by
retaining input references.
- ManageTodos: validate Status against the documented "pending" | "completed" set;
coerce out-of-range values to "pending" with a LogWarning instead of letting
arbitrary LLM-supplied strings into shared state.
- GenerateA2ui: add a final catch (Exception) returning a StructuredError
("unexpected_error", ...) matching the existing taxonomy, and log an info-level
entry when OperationCanceledException flows through (was previously silent).
AimockHeaderPolicy previously called message.Request.Headers.Set(...)
unconditionally for every key returned by AimockHeaderContext, which
silently clobbered any header already set by an earlier pipeline policy
or the SDK itself (e.g. x-request-id, x-correlation-id).
Switch both Process and ProcessAsync to add-if-absent semantics, using
PipelineRequestHeaders.TryGetValue to skip keys that already have a value
on the outbound request. New aimock x-* headers still propagate; existing
correlation/SDK headers are preserved.
Jordan Ritter