diff --git a/internal/server/api/chat.go b/internal/server/api/chat.go index 3a05d24d3..09e7dbed4 100644 --- a/internal/server/api/chat.go +++ b/internal/server/api/chat.go @@ -81,6 +81,8 @@ func (handlers *ChatCompletionHandlers) ChatCompletionWithRequest(c *gin.Context return } + writeForwardResponseHeaders(c, result) + if result.ChatCompletion != nil { resp := result.ChatCompletion @@ -115,6 +117,18 @@ func (handlers *ChatCompletionHandlers) ChatCompletionWithRequest(c *gin.Context } } +func writeForwardResponseHeaders(c *gin.Context, result orchestrator.ChatCompletionResult) { + var headers http.Header + if result.ChatCompletion != nil { + headers = result.ChatCompletion.Headers + } else { + headers = httpclient.GetResponseHeaders(result.ChatCompletionStream) + } + + // c.Writer.Header() is non-nil, so the merge updates it in place. + _ = httpclient.MergeForwardResponseHeaders(c.Writer.Header(), headers) +} + // StreamErrorFormatter formats a stream error into a JSON-serializable object for SSE error events. type StreamErrorFormatter func(ctx context.Context, err error) any diff --git a/internal/server/api/chat_test.go b/internal/server/api/chat_test.go index f3a70bc9b..be7b74253 100644 --- a/internal/server/api/chat_test.go +++ b/internal/server/api/chat_test.go @@ -135,6 +135,36 @@ func TestWriteSSEStream_Success(t *testing.T) { assert.Contains(t, body, `[DONE]`) } +func TestWriteForwardResponseHeaders(t *testing.T) { + headers := http.Header{ + httpclient.ReasoningIncludedHeader: []string{"true"}, + "Set-Cookie": []string{"secret=1"}, + } + tests := map[string]orchestrator.ChatCompletionResult{ + "non-streaming": { + ChatCompletion: &httpclient.Response{Headers: headers}, + }, + "streaming": { + ChatCompletionStream: httpclient.WithResponseHeaders( + streams.SliceStream([]*httpclient.StreamEvent{}), + headers, + ), + }, + } + + for name, result := range tests { + t.Run(name, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + writeForwardResponseHeaders(c, result) + + require.Equal(t, "true", w.Header().Get(httpclient.ReasoningIncludedHeader)) + require.Empty(t, w.Header().Get("Set-Cookie")) + }) + } +} + func TestWriteSSEStream_ErrorFormatsAsJSON(t *testing.T) { w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) diff --git a/llm/httpclient/client.go b/llm/httpclient/client.go index a595498e1..075307957 100644 --- a/llm/httpclient/client.go +++ b/llm/httpclient/client.go @@ -362,7 +362,7 @@ func (hc *HttpClient) DoStream(ctx context.Context, request *Request) (streams.S stream := decoderFactory(ctx, rawResp.Body) - return stream, nil + return WithResponseHeaders(stream, MergeForwardResponseHeaders(nil, rawResp.Header)), nil } // BuildHttpRequest builds an HTTP request from Request. diff --git a/llm/httpclient/client_test.go b/llm/httpclient/client_test.go index 28586cb16..52cb54339 100644 --- a/llm/httpclient/client_test.go +++ b/llm/httpclient/client_test.go @@ -166,7 +166,7 @@ func TestHttpClientImpl_DoStream(t *testing.T) { serverResponse func(w http.ResponseWriter, r *http.Request) wantErr bool wantErrContains string - validate func(stream any) bool + validate func(stream StreamDecoder) bool }{ { name: "successful streaming request", @@ -189,6 +189,8 @@ func TestHttpClientImpl_DoStream(t *testing.T) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") + w.Header().Set(ReasoningIncludedHeader, "true") + w.Header().Set("Set-Cookie", "secret=1") w.WriteHeader(http.StatusOK) // Write SSE events @@ -211,9 +213,9 @@ func TestHttpClientImpl_DoStream(t *testing.T) { } }, wantErr: false, - validate: func(stream any) bool { - // This is a basic validation - in a real test we'd iterate through the stream - return stream != nil + validate: func(stream StreamDecoder) bool { + headers := GetResponseHeaders(stream) + return headers.Get(ReasoningIncludedHeader) == "true" && headers.Get("Set-Cookie") == "" }, }, { @@ -230,7 +232,7 @@ func TestHttpClientImpl_DoStream(t *testing.T) { w.Write([]byte(`{"error": "unauthorized"}`)) }, wantErr: true, - validate: func(stream any) bool { + validate: func(stream StreamDecoder) bool { return stream == nil }, }, diff --git a/llm/httpclient/response_headers.go b/llm/httpclient/response_headers.go new file mode 100644 index 000000000..93f15bcff --- /dev/null +++ b/llm/httpclient/response_headers.go @@ -0,0 +1,85 @@ +package httpclient + +import ( + "net/http" + "strings" + + "github.com/looplj/axonhub/llm/streams" +) + +// ReasoningIncludedHeader tells Codex that upstream usage already includes +// previously generated reasoning tokens. +const ReasoningIncludedHeader = "X-Reasoning-Included" + +type responseHeadersProvider interface { + responseHeaders() http.Header +} + +type responseHeadersStream struct { + streams.Stream[*StreamEvent] + headers http.Header +} + +func (s *responseHeadersStream) responseHeaders() http.Header { + return s.headers +} + +// WithResponseHeaders attaches HTTP response metadata to a stream without +// widening the generic streams.Stream interface. +func WithResponseHeaders(stream streams.Stream[*StreamEvent], headers http.Header) streams.Stream[*StreamEvent] { + if stream == nil || len(headers) == 0 { + return stream + } + + return &responseHeadersStream{ + Stream: stream, + headers: headers.Clone(), + } +} + +// GetResponseHeaders returns a copy of response metadata attached to a stream. +func GetResponseHeaders(stream streams.Stream[*StreamEvent]) http.Header { + provider, ok := stream.(responseHeadersProvider) + if !ok { + return nil + } + + return provider.responseHeaders().Clone() +} + +// MergeForwardResponseHeaders copies the small, explicit set of upstream +// headers that are safe and meaningful at AxonHub's client boundary. +func MergeForwardResponseHeaders(dst, src http.Header) http.Header { + forward := hasOnlyTrueHeaderValues(src, ReasoningIncludedHeader) + if dst != nil { + dst.Del(ReasoningIncludedHeader) + } + if !forward { + return dst + } + if dst == nil { + dst = make(http.Header) + } + + dst[ReasoningIncludedHeader] = []string{"true"} + + return dst +} + +func hasOnlyTrueHeaderValues(headers http.Header, name string) bool { + found := false + for key, values := range headers { + if !strings.EqualFold(key, name) { + continue + } + + for _, value := range values { + found = true + if !strings.EqualFold(strings.TrimSpace(value), "true") { + return false + } + } + } + + return found +} diff --git a/llm/httpclient/response_headers_test.go b/llm/httpclient/response_headers_test.go new file mode 100644 index 000000000..a8fbab1eb --- /dev/null +++ b/llm/httpclient/response_headers_test.go @@ -0,0 +1,48 @@ +package httpclient + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/looplj/axonhub/llm/streams" +) + +func TestMergeForwardResponseHeaders(t *testing.T) { + tests := []struct { + name string + src http.Header + want string + }{ + {name: "true", src: http.Header{"x-reasoning-included": []string{" TRUE "}}, want: "true"}, + {name: "false", src: http.Header{ReasoningIncludedHeader: []string{"false"}}}, + {name: "missing", src: http.Header{"Set-Cookie": []string{"secret=1"}}}, + {name: "conflicting duplicates", src: http.Header{ + ReasoningIncludedHeader: []string{"true"}, + "x-reasoning-included": []string{"false"}, + }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dst := http.Header{ReasoningIncludedHeader: []string{"stale"}} + got := MergeForwardResponseHeaders(dst, tt.src) + + require.Equal(t, tt.want, got.Get(ReasoningIncludedHeader)) + require.Empty(t, got.Get("Set-Cookie")) + }) + } +} + +func TestResponseHeadersStreamCopiesHeaders(t *testing.T) { + headers := http.Header{ReasoningIncludedHeader: []string{"true"}} + stream := WithResponseHeaders(streams.SliceStream([]*StreamEvent{}), headers) + headers.Set(ReasoningIncludedHeader, "false") + + got := GetResponseHeaders(stream) + got.Set(ReasoningIncludedHeader, "false") + + require.Equal(t, "true", GetResponseHeaders(stream).Get(ReasoningIncludedHeader)) + require.Nil(t, GetResponseHeaders(streams.SliceStream([]*StreamEvent{}))) +} diff --git a/llm/pipeline/integration_test.go b/llm/pipeline/integration_test.go index c526c82bf..b17e08c10 100644 --- a/llm/pipeline/integration_test.go +++ b/llm/pipeline/integration_test.go @@ -92,7 +92,9 @@ func TestPipeline_OpenAI_to_OpenAI(t *testing.T) { return &httpclient.Response{ StatusCode: http.StatusOK, Headers: http.Header{ - "Content-Type": []string{"application/json"}, + "Content-Type": []string{"application/json"}, + httpclient.ReasoningIncludedHeader: []string{"true"}, + "Set-Cookie": []string{"secret=1"}, }, Body: responseBody, }, nil @@ -136,6 +138,8 @@ func TestPipeline_OpenAI_to_OpenAI(t *testing.T) { // Verify response require.Equal(t, http.StatusOK, result.Response.StatusCode) require.Equal(t, "application/json", result.Response.Headers.Get("Content-Type")) + require.Equal(t, "true", result.Response.Headers.Get(httpclient.ReasoningIncludedHeader)) + require.Empty(t, result.Response.Headers.Get("Set-Cookie")) var finalResponse llm.Response diff --git a/llm/pipeline/non_streaming.go b/llm/pipeline/non_streaming.go index f7b65f8df..c7bbebb85 100644 --- a/llm/pipeline/non_streaming.go +++ b/llm/pipeline/non_streaming.go @@ -36,6 +36,7 @@ func (p *pipeline) notStream( return nil, fmt.Errorf("failed to apply raw response middlewares: %w", err) } + responseHeaders := httpclient.MergeForwardResponseHeaders(nil, httpResp.Headers) llmResp, err := p.Outbound.TransformResponse(ctx, httpResp) if err != nil { @@ -74,6 +75,7 @@ func (p *pipeline) notStream( return nil, fmt.Errorf("failed to apply inbound raw response middlewares: %w", err) } + finalResp.Headers = httpclient.MergeForwardResponseHeaders(finalResp.Headers, responseHeaders) return finalResp, nil } @@ -88,6 +90,7 @@ func (p *pipeline) autoAggregateStream( return nil, err } defer inboundStream.Close() + responseHeaders := httpclient.GetResponseHeaders(inboundStream) chunks := make([]*httpclient.StreamEvent, 0, 8) for inboundStream.Next() { @@ -132,6 +135,7 @@ func (p *pipeline) autoAggregateStream( p.applyRawErrorResponseMiddlewares(ctx, err) return nil, fmt.Errorf("failed to apply inbound raw response middlewares: %w", err) } + resp.Headers = httpclient.MergeForwardResponseHeaders(resp.Headers, responseHeaders) return resp, nil } diff --git a/llm/pipeline/pipeline_retry_test.go b/llm/pipeline/pipeline_retry_test.go index baee750ab..c1457657e 100644 --- a/llm/pipeline/pipeline_retry_test.go +++ b/llm/pipeline/pipeline_retry_test.go @@ -470,7 +470,15 @@ func TestPipeline_Process_StreamRetriesPreCommitError(t *testing.T) { executor := &mockExecutor{ doStream: func(ctx context.Context, req *httpclient.Request) (streams.Stream[*httpclient.StreamEvent], error) { attempts++ - return streams.SliceStream([]*httpclient.StreamEvent{{Data: []byte("raw")}}), nil + reasoningIncluded := "false" + if attempts == 1 { + reasoningIncluded = "true" + } + + return httpclient.WithResponseHeaders( + streams.SliceStream([]*httpclient.StreamEvent{{Data: []byte("raw")}}), + http.Header{httpclient.ReasoningIncludedHeader: []string{reasoningIncluded}}, + ), nil }, } @@ -513,6 +521,7 @@ func TestPipeline_Process_StreamRetriesPreCommitError(t *testing.T) { require.True(t, res.Stream) require.Equal(t, 2, attempts) require.Equal(t, 1, prepareCalls) + require.Empty(t, httpclient.GetResponseHeaders(res.EventStream).Get(httpclient.ReasoningIncludedHeader)) events, err := streams.All(res.EventStream) require.NoError(t, err) diff --git a/llm/pipeline/stream.go b/llm/pipeline/stream.go index 529601f84..06028f19f 100644 --- a/llm/pipeline/stream.go +++ b/llm/pipeline/stream.go @@ -321,6 +321,7 @@ func (p *pipeline) stream( return nil, WrapUpstreamError(err) } + responseHeaders := httpclient.MergeForwardResponseHeaders(nil, httpclient.GetResponseHeaders(outboundStream)) // Apply raw stream middlewares rawStream := outboundStream @@ -445,5 +446,5 @@ func (p *pipeline) stream( } } - return inboundStream, nil + return httpclient.WithResponseHeaders(inboundStream, responseHeaders), nil } diff --git a/llm/pipeline/streaming_integration_test.go b/llm/pipeline/streaming_integration_test.go index de9cda28f..89754a078 100644 --- a/llm/pipeline/streaming_integration_test.go +++ b/llm/pipeline/streaming_integration_test.go @@ -72,7 +72,10 @@ func TestPipeline_Streaming_OpenAI_to_OpenAI(t *testing.T) { require.Equal(t, "Bearer test-api-key", request.Headers.Get("Authorization")) // Return mock stream - return streams.SliceStream(streamEvents), nil + return httpclient.WithResponseHeaders(streams.SliceStream(streamEvents), http.Header{ + httpclient.ReasoningIncludedHeader: []string{"true"}, + "Set-Cookie": []string{"secret=1"}, + }), nil }, } @@ -128,6 +131,8 @@ func TestPipeline_Streaming_OpenAI_to_OpenAI(t *testing.T) { require.NotNil(t, result) require.True(t, result.Stream) require.NotNil(t, result.EventStream) + require.Equal(t, "true", httpclient.GetResponseHeaders(result.EventStream).Get(httpclient.ReasoningIncludedHeader)) + require.Empty(t, httpclient.GetResponseHeaders(result.EventStream).Get("Set-Cookie")) // Collect all events from the stream var collectedEvents []*httpclient.StreamEvent @@ -461,7 +466,10 @@ func TestPipeline_NonStreaming_AutoAggregateUpgradedStream(t *testing.T) { require.NoError(t, err) require.Equal(t, true, reqBody["stream"]) - return streams.SliceStream(streamEvents), nil + return httpclient.WithResponseHeaders(streams.SliceStream(streamEvents), http.Header{ + httpclient.ReasoningIncludedHeader: []string{"true"}, + "Set-Cookie": []string{"secret=1"}, + }), nil }, } @@ -515,6 +523,8 @@ func TestPipeline_NonStreaming_AutoAggregateUpgradedStream(t *testing.T) { require.NotNil(t, result.Response) require.Equal(t, http.StatusOK, result.Response.StatusCode) require.Equal(t, "application/json", result.Response.Headers.Get("Content-Type")) + require.Equal(t, "true", result.Response.Headers.Get(httpclient.ReasoningIncludedHeader)) + require.Empty(t, result.Response.Headers.Get("Set-Cookie")) var finalResponse map[string]any err = json.Unmarshal(result.Response.Body, &finalResponse) diff --git a/llm/transformer/openai/codex/outbound.go b/llm/transformer/openai/codex/outbound.go index 8be2ab497..b2b27ab54 100644 --- a/llm/transformer/openai/codex/outbound.go +++ b/llm/transformer/openai/codex/outbound.go @@ -376,9 +376,9 @@ func (e *codexExecutor) Do(ctx context.Context, request *httpclient.Request) (*h return &httpclient.Response{ StatusCode: http.StatusOK, - Headers: http.Header{ + Headers: httpclient.MergeForwardResponseHeaders(http.Header{ "Content-Type": []string{"application/json"}, - }, + }, httpclient.GetResponseHeaders(stream)), Body: body, Request: request, }, nil diff --git a/llm/transformer/openai/codex/outbound_executor_test.go b/llm/transformer/openai/codex/outbound_executor_test.go index f8111597d..691100d0f 100644 --- a/llm/transformer/openai/codex/outbound_executor_test.go +++ b/llm/transformer/openai/codex/outbound_executor_test.go @@ -321,12 +321,18 @@ func TestCodexOutbound_CustomizeExecutorAggregatesNonStreamRequests(t *testing.T {Type: "response.output_item.done", Data: []byte(`{"type":"response.output_item.done","sequence_number":5,"output_index":0,"item":{"id":"msg_test_456","type":"message","status":"completed","role":"assistant"}}`)}, {Type: "response.completed", Data: []byte(`{"type":"response.completed","sequence_number":6,"response":{"id":"resp_test_123","object":"response","created_at":1700000000,"model":"gpt-5-codex","status":"completed","output":[]}}`)}, }, + streamHeaders: http.Header{ + httpclient.ReasoningIncludedHeader: []string{"true"}, + "Set-Cookie": []string{"secret=1"}, + }, }) response, err := executor.Do(ctx, request) require.NoError(t, err) require.Equal(t, http.StatusOK, response.StatusCode) require.Equal(t, "application/json", response.Headers.Get("Content-Type")) + require.Equal(t, "true", response.Headers.Get(httpclient.ReasoningIncludedHeader)) + require.Empty(t, response.Headers.Get("Set-Cookie")) var body map[string]any require.NoError(t, json.Unmarshal(response.Body, &body)) @@ -365,7 +371,8 @@ func TestCodexOutbound_DoReturnsWebSocketErrorEvents(t *testing.T) { var _ pipeline.ChannelCustomizedExecutor = (*OutboundTransformer)(nil) type mockCodexExecutor struct { - streamEvents []*httpclient.StreamEvent + streamEvents []*httpclient.StreamEvent + streamHeaders http.Header } func (m *mockCodexExecutor) Do(_ context.Context, _ *httpclient.Request) (*httpclient.Response, error) { @@ -373,7 +380,7 @@ func (m *mockCodexExecutor) Do(_ context.Context, _ *httpclient.Request) (*httpc } func (m *mockCodexExecutor) DoStream(_ context.Context, _ *httpclient.Request) (streams.Stream[*httpclient.StreamEvent], error) { - return streams.SliceStream(m.streamEvents), nil + return httpclient.WithResponseHeaders(streams.SliceStream(m.streamEvents), m.streamHeaders), nil } func TestCodexOutbound_DoesNotInjectCLIInstructions(t *testing.T) {