🤖 feat: add provider response model mapping

Author: dreamhunter2333

internal/executor/middleware_dispatch.go

@@ -10,6 +10,7 @@ import (
 	"github.com/awsl-project/maxx/internal/converter"
 	"github.com/awsl-project/maxx/internal/cooldown"
 	"github.com/awsl-project/maxx/internal/domain"
+	"github.com/awsl-project/maxx/internal/executor/responsemodifier"
 	"github.com/awsl-project/maxx/internal/flow"
 	"github.com/awsl-project/maxx/internal/pricing"
 	"github.com/awsl-project/maxx/internal/usage"
@@ -158,24 +159,40 @@ func (e *Executor) dispatch(c *flow.Ctx) {
 			var responseWriter http.ResponseWriter
 			var convertingWriter *ConvertingResponseWriter
 			responseCapture := NewResponseCapture(c.Writer)
+			modifierWriter := responsemodifier.NewResponseModifierWriter(
+				responseCapture,
+				matchedRoute.Provider,
+				originalClientType,
+				state.requestModel,
+				mappedModel,
+				state.isStream,
+			)
+			if modifierWriter != nil {
+				responseWriter = modifierWriter
+			} else {
+				responseWriter = responseCapture
+			}
 			if needsConversion {
 				convertingWriter = NewConvertingResponseWriter(
-					responseCapture, e.converter, originalClientType, currentClientType, state.isStream, state.originalRequestBody)
+					responseWriter, e.converter, originalClientType, currentClientType, state.isStream, state.originalRequestBody)
 				responseWriter = convertingWriter
-			} else {
-				responseWriter = responseCapture
 			}
 
 			originalWriter := c.Writer
 			c.Writer = responseWriter
 			err := matchedRoute.ProviderAdapter.Execute(c, matchedRoute.Provider)
 			c.Writer = originalWriter
 
-			if needsConversion && convertingWriter != nil && !state.isStream {
+			if needsConversion && convertingWriter != nil && !state.isStream && (err == nil || modifierWriter == nil) {
 				if finalizeErr := convertingWriter.Finalize(); finalizeErr != nil {
 					log.Printf("[Executor] Response conversion finalize failed: %v", finalizeErr)
 				}
 			}
+			if err == nil && modifierWriter != nil {
+				if finalizeErr := modifierWriter.Finalize(); finalizeErr != nil {
+					log.Printf("[Executor] Response modifier finalize failed: %v", finalizeErr)
+				}
+			}
 
 			eventChan.Close()
 			<-eventDone

internal/executor/responsemodifier/claude_response_modifier.go

@@ -0,0 +1,99 @@
+package responsemodifier
+
+import (
+	"bytes"
+	"encoding/json"
+	"regexp"
+	"strings"
+
+	"github.com/awsl-project/maxx/internal/domain"
+)
+
+var claudeSSEDataLineRE = regexp.MustCompile(`(?m)^(\s*data:\s*)(.*?)(\r?\n?)$`)
+
+type claudeResponseModifier struct {
+	requestModel string
+	mappedModel  string
+}
+
+func newClaudeResponseModifier(config *domain.ProviderConfig, clientType domain.ClientType, requestModel string, mappedModel string) *claudeResponseModifier {
+	if clientType != domain.ClientTypeClaude {
+		return nil
+	}
+	if config == nil || config.Claude == nil {
+		return nil
+	}
+	requestModel = strings.TrimSpace(requestModel)
+	mappedModel = strings.TrimSpace(mappedModel)
+	if requestModel == "" || mappedModel == "" || requestModel == mappedModel {
+		return nil
+	}
+	return &claudeResponseModifier{requestModel: requestModel, mappedModel: mappedModel}
+}
+
+func (m *claudeResponseModifier) modifyResponse(body []byte) []byte {
+	return m.rewriteJSON(body)
+}
+
+func (m *claudeResponseModifier) modifyStreamEvent(body []byte) []byte {
+	return claudeSSEDataLineRE.ReplaceAllFunc(body, func(line []byte) []byte {
+		parts := claudeSSEDataLineRE.FindSubmatch(line)
+		if len(parts) != 4 || bytes.Equal(bytes.TrimSpace(parts[2]), []byte("[DONE]")) {
+			return line
+		}
+		payload := m.rewriteJSON(parts[2])
+		if bytes.Equal(payload, parts[2]) {
+			return line
+		}
+		out := make([]byte, 0, len(parts[1])+len(payload)+len(parts[3]))
+		out = append(out, parts[1]...)
+		out = append(out, payload...)
+		out = append(out, parts[3]...)
+		return out
+	})
+}
+
+func (m *claudeResponseModifier) rewriteJSON(body []byte) []byte {
+	if len(bytes.TrimSpace(body)) == 0 {
+		return body
+	}
+	var object map[string]any
+	if err := json.Unmarshal(body, &object); err != nil {
+		return body
+	}
+	changed := m.rewriteModel(object, "model")
+	if message, ok := object["message"].(map[string]any); ok {
+		changed = m.rewriteModel(message, "model") || changed
+	}
+	if !changed {
+		return body
+	}
+	var buf bytes.Buffer
+	encoder := json.NewEncoder(&buf)
+	encoder.SetEscapeHTML(false)
+	if err := encoder.Encode(object); err != nil {
+		return body
+	}
+	return bytes.TrimSuffix(buf.Bytes(), []byte("\n"))
+}
+
+func (m *claudeResponseModifier) rewriteModel(object map[string]any, key string) bool {
+	value, ok := object[key].(string)
+	if !ok {
+		return false
+	}
+	mapped := m.mapModel(value)
+	if mapped == value {
+		return false
+	}
+	object[key] = mapped
+	return true
+}
+
+func (m *claudeResponseModifier) mapModel(model string) string {
+	model = strings.TrimSpace(model)
+	if model != m.mappedModel {
+		return model
+	}
+	return m.requestModel
+}

internal/executor/responsemodifier/response_modifier_writer.go

@@ -0,0 +1,156 @@
+package responsemodifier
+
+import (
+	"bytes"
+	"net/http"
+	"strconv"
+
+	"github.com/awsl-project/maxx/internal/domain"
+)
+
+type responseModifier interface {
+	modifyResponse(body []byte) []byte
+	modifyStreamEvent(event []byte) []byte
+}
+
+// ResponseModifierWriter buffers a response and applies provider-specific response modifications before sending it.
+type ResponseModifierWriter struct {
+	underlying  http.ResponseWriter
+	modifier    responseModifier
+	isStream    bool
+	statusCode  int
+	buffer      bytes.Buffer
+	headersSent bool
+}
+
+func NewResponseModifierWriter(
+	w http.ResponseWriter,
+	provider *domain.Provider,
+	clientType domain.ClientType,
+	requestModel string,
+	mappedModel string,
+	isStream bool,
+) *ResponseModifierWriter {
+	modifier := newResponseModifier(provider, clientType, requestModel, mappedModel)
+	if modifier == nil {
+		return nil
+	}
+	return &ResponseModifierWriter{underlying: w, modifier: modifier, isStream: isStream, statusCode: http.StatusOK}
+}
+
+func newResponseModifier(provider *domain.Provider, clientType domain.ClientType, requestModel string, mappedModel string) responseModifier {
+	if provider == nil {
+		return nil
+	}
+	switch provider.Type {
+	case "claude":
+		modifier := newClaudeResponseModifier(provider.Config, clientType, requestModel, mappedModel)
+		if modifier == nil {
+			return nil
+		}
+		return modifier
+	default:
+		return nil
+	}
+}
+
+func (w *ResponseModifierWriter) Header() http.Header {
+	return w.underlying.Header()
+}
+
+func (w *ResponseModifierWriter) WriteHeader(code int) {
+	w.statusCode = code
+	if w.isStream {
+		w.writeHeaderIfNeeded()
+	}
+}
+
+func (w *ResponseModifierWriter) Write(b []byte) (int, error) {
+	if !w.isStream {
+		_, err := w.buffer.Write(b)
+		return len(b), err
+	}
+	if _, err := w.buffer.Write(b); err != nil {
+		return 0, err
+	}
+	if err := w.flushCompleteStreamEvents(false); err != nil {
+		return 0, err
+	}
+	return len(b), nil
+}
+
+func (w *ResponseModifierWriter) Flush() {
+	if w.isStream {
+		_ = w.flushCompleteStreamEvents(false)
+	}
+	if f, ok := w.underlying.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (w *ResponseModifierWriter) Finalize() error {
+	if w.isStream {
+		return w.flushCompleteStreamEvents(true)
+	}
+	body := w.modifier.modifyResponse(w.buffer.Bytes())
+	if w.underlying.Header().Get("Content-Length") != "" {
+		w.underlying.Header().Set("Content-Length", strconv.Itoa(len(body)))
+	}
+	w.writeHeaderIfNeeded()
+	_, err := w.underlying.Write(body)
+	return err
+}
+
+func (w *ResponseModifierWriter) flushCompleteStreamEvents(final bool) error {
+	for {
+		eventLen := completeSSEEventLen(w.buffer.Bytes())
+		if eventLen == 0 {
+			break
+		}
+		if err := w.writeStreamEvent(w.buffer.Next(eventLen)); err != nil {
+			return err
+		}
+	}
+	if final && w.buffer.Len() > 0 {
+		if err := w.writeStreamEvent(w.buffer.Next(w.buffer.Len())); err != nil {
+			return err
+		}
+	}
+	if final {
+		w.writeHeaderIfNeeded()
+	}
+	return nil
+}
+
+func (w *ResponseModifierWriter) writeStreamEvent(event []byte) error {
+	body := w.modifier.modifyStreamEvent(event)
+	w.writeHeaderIfNeeded()
+	_, err := w.underlying.Write(body)
+	if err != nil {
+		return err
+	}
+	if f, ok := w.underlying.(http.Flusher); ok {
+		f.Flush()
+	}
+	return nil
+}
+
+func (w *ResponseModifierWriter) writeHeaderIfNeeded() {
+	if w.headersSent {
+		return
+	}
+	w.underlying.WriteHeader(w.statusCode)
+	w.headersSent = true
+}
+
+func completeSSEEventLen(data []byte) int {
+	lf := bytes.Index(data, []byte("\n\n"))
+	crlf := bytes.Index(data, []byte("\r\n\r\n"))
+	if lf == -1 && crlf == -1 {
+		return 0
+	}
+	if lf == -1 || (crlf != -1 && crlf < lf) {
+		return crlf + len("\r\n\r\n")
+	}
+	return lf + len("\n\n")
+}