mirror of
https://github.com/CopilotKit/CopilotKit.git
synced 2026-09-14 16:26:20 +08:00
Jordan Ritter
4f66ba23ed
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.
61 lines
3.1 KiB
C#
61 lines
3.1 KiB
C#
// STOPGAP: This integration-level header propagation replaces once copilotkit-sdk-dotnet
|
|
// ships (Microsoft contribution, ETA mid-2026). When that SDK lands, delete this code
|
|
// and use the SDK's built-in header propagation.
|
|
// See: https://www.notion.so/copilotkit/3543aa3818528150b6acc5b872ad7fe5
|
|
|
|
// TODO(copilotkit-sdk-dotnet): migrate to SDK-level header propagation
|
|
public class AimockHeaderMiddleware
|
|
{
|
|
private readonly RequestDelegate _next;
|
|
private readonly ILogger<AimockHeaderMiddleware> _logger;
|
|
|
|
public AimockHeaderMiddleware(RequestDelegate next, ILogger<AimockHeaderMiddleware> logger)
|
|
{
|
|
_next = next;
|
|
_logger = logger;
|
|
}
|
|
|
|
public async Task InvokeAsync(HttpContext context)
|
|
{
|
|
// Use case-insensitive comparer because ASP.NET's IHeaderDictionary is itself
|
|
// case-insensitive, but iterating its underlying store can in rare cases yield
|
|
// case-variant duplicates (e.g., a misbehaving proxy injecting both `X-Foo`
|
|
// and `x-foo`). With the default ordinal comparer, ToDictionary would throw
|
|
// ArgumentException on duplicate keys and fail the request.
|
|
//
|
|
// When such a collision occurs, we keep the first value because HTTP has no
|
|
// canonical merge rule for case-variant headers across distinct keys (joining
|
|
// with comma would only be defined if the keys were ASCII-equal). We log a
|
|
// warning so operators can see that an upstream proxy is misbehaving and that
|
|
// downstream consumers may be observing only one of several values.
|
|
var groupedHeaders = context.Request.Headers
|
|
.Where(h => h.Key.StartsWith("x-", StringComparison.OrdinalIgnoreCase))
|
|
.GroupBy(h => h.Key, StringComparer.OrdinalIgnoreCase)
|
|
.ToList();
|
|
|
|
foreach (var group in groupedHeaders.Where(g => g.Count() > 1))
|
|
{
|
|
_logger.LogWarning(
|
|
"[aimock-header-middleware] header '{Key}' arrived with {Count} case-variant entries; keeping the first ('{Kept}'), dropping {DroppedCount} others",
|
|
group.Key, group.Count(), group.First().Value.ToString(), group.Count() - 1);
|
|
}
|
|
|
|
var headers = groupedHeaders.ToDictionary(
|
|
g => g.First().Key,
|
|
g => g.First().Value.ToString(),
|
|
StringComparer.OrdinalIgnoreCase);
|
|
// Do NOT reset AimockHeaderContext after _next returns. ASP.NET Core hands control
|
|
// back to this middleware once the endpoint handler completes, but for streaming
|
|
// endpoints (AG-UI uses IAsyncEnumerable/SSE) the response delegate continues
|
|
// writing to the response body — and may invoke downstream OpenAI calls — AFTER
|
|
// _next returns. Clearing the AsyncLocal here would cause AimockHeaderPolicy to
|
|
// silently drop aimock headers on those streaming-tail calls.
|
|
//
|
|
// AsyncLocal<T> is scoped to the current execution context (per-request), so the
|
|
// value does not leak across concurrent ASP.NET requests; the finally-clear was
|
|
// unnecessary and actively harmful.
|
|
AimockHeaderContext.Set(headers);
|
|
await _next(context);
|
|
}
|
|
}
|