refactor(tracing): simplifies span generic impl (#1622)

**Description**

This simplifies the span generic implementation by removing unnecessary
indirection

**Related Issues/PRs (if applicable)**
Follow up on #1582

---------

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

internal/tracing/api/api.go

@@ -98,6 +98,7 @@ type (
 type (
 	// SpanRecorder standardizes recorder implementations for non-MCP tracers.
 	SpanRecorder[ReqT any, RespT any, RespChunkT any] interface {
+		SpanResponseRecorder[RespT, RespChunkT]
 		// StartParams returns the name and options to start the span with.
 		//
 		// Parameters:
@@ -108,6 +109,11 @@ type (
 		StartParams(req *ReqT, body []byte) (spanName string, opts []trace.SpanStartOption)
 		// RecordRequest records request attributes to the span.
 		RecordRequest(span trace.Span, req *ReqT, body []byte)
+	}
+	// SpanResponseRecorder standardizes response recording for non-MCP tracers.
+	//
+	// It is separated from SpanRecorder to allow reusing response recording logic.
+	SpanResponseRecorder[RespT, RespChunkT any] interface {
 		// RecordResponse records response attributes to the span.
 		RecordResponse(span trace.Span, resp *RespT)
 		// RecordResponseOnError ends recording the span with an error status.

internal/tracing/span.go

@@ -14,78 +14,41 @@ import (
 	tracing "github.com/envoyproxy/ai-gateway/internal/tracing/api"
 )
 
-type responseRecorder[RespT any] interface {
-	RecordResponse(trace.Span, *RespT)
-	RecordResponseOnError(trace.Span, int, []byte)
-}
-
-type streamResponseRecorder[RespT, ChunkT any] interface {
-	responseRecorder[RespT]
-	RecordResponseChunks(trace.Span, []*ChunkT)
-}
-
-type noopChunkRecorder[ChunkT any] struct{}
-
-func (noopChunkRecorder[ChunkT]) RecordResponseChunk(*ChunkT) {}
-
-type streamingSpan[RespT, ChunkT any] struct {
+type span[RespT, ChunkT any] struct {
 	span     trace.Span
-	recorder streamResponseRecorder[RespT, ChunkT]
+	recorder tracing.SpanResponseRecorder[RespT, ChunkT]
 	chunks   []*ChunkT
 }
 
-func (s *streamingSpan[RespT, ChunkT]) RecordResponseChunk(resp *ChunkT) {
+// RecordResponseChunk implements [tracing.Span.RecordResponseChunk]
+func (s *span[RespT, ChunkT]) RecordResponseChunk(resp *ChunkT) {
 	s.chunks = append(s.chunks, resp)
 }
 
-func (s *streamingSpan[RespT, ChunkT]) RecordResponse(resp *RespT) {
+// RecordResponse implements [tracing.Span.RecordResponse]
+func (s *span[RespT, ChunkT]) RecordResponse(resp *RespT) {
 	s.recorder.RecordResponse(s.span, resp)
 }
 
-func (s *streamingSpan[RespT, ChunkT]) EndSpan() {
+// EndSpan implements [tracing.Span.EndSpan]
+func (s *span[RespT, ChunkT]) EndSpan() {
 	if len(s.chunks) > 0 {
 		s.recorder.RecordResponseChunks(s.span, s.chunks)
 	}
 	s.span.End()
 }
 
-func (s *streamingSpan[RespT, ChunkT]) EndSpanOnError(statusCode int, body []byte) {
-	s.recorder.RecordResponseOnError(s.span, statusCode, body)
-	s.span.End()
-}
-
-type responseSpan[RespT, ChunkT any] struct {
-	noopChunkRecorder[ChunkT]
-	span     trace.Span
-	recorder streamResponseRecorder[RespT, ChunkT]
-}
-
-func (s *responseSpan[RespT, ChunkT]) RecordResponse(resp *RespT) {
-	s.recorder.RecordResponse(s.span, resp)
-}
-
-func (s *responseSpan[RespT, ChunkT]) EndSpan() {
-	s.span.End()
-}
-
-func (s *responseSpan[RespT, ChunkT]) EndSpanOnError(statusCode int, body []byte) {
+// EndSpanOnError implements [tracing.Span.EndSpanOnError]
+func (s *span[RespT, ChunkT]) EndSpanOnError(statusCode int, body []byte) {
 	s.recorder.RecordResponseOnError(s.span, statusCode, body)
 	s.span.End()
 }
 
 // Type aliases tying generic implementations to concrete recorder contracts.
 type (
-	chatCompletionSpan  = streamingSpan[openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]
-	completionSpan      = streamingSpan[openai.CompletionResponse, openai.CompletionResponse]
-	embeddingsSpan      = responseSpan[openai.EmbeddingResponse, struct{}]
-	imageGenerationSpan = responseSpan[openaisdk.ImagesResponse, struct{}]
-	rerankSpan          = responseSpan[cohereschema.RerankV2Response, struct{}]
-)
-
-var (
-	_ tracing.ChatCompletionSpan  = (*chatCompletionSpan)(nil)
-	_ tracing.CompletionSpan      = (*completionSpan)(nil)
-	_ tracing.EmbeddingsSpan      = (*embeddingsSpan)(nil)
-	_ tracing.ImageGenerationSpan = (*imageGenerationSpan)(nil)
-	_ tracing.RerankSpan          = (*rerankSpan)(nil)
+	chatCompletionSpan  = span[openai.ChatCompletionResponse, openai.ChatCompletionResponseChunk]
+	completionSpan      = span[openai.CompletionResponse, openai.CompletionResponse]
+	embeddingsSpan      = span[openai.EmbeddingResponse, struct{}]
+	imageGenerationSpan = span[openaisdk.ImagesResponse, struct{}]
+	rerankSpan          = span[cohereschema.RerankV2Response, struct{}]
 )