Merge PR #2474

# Conflicts:
#	internal/api/modules/amp/response_rewriter.go
#	internal/api/modules/amp/response_rewriter_test.go

internal/api/modules/amp/fallback_handlers.go

@@ -253,6 +253,7 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
 			log.Debugf("amp model mapping: request %s -> %s", normalizedModel, resolvedModel)
 			logAmpRouting(RouteTypeModelMapping, modelName, resolvedModel, providerName, requestPath)
 			rewriter := NewResponseRewriter(c.Writer, modelName)
+			rewriter.suppressThinking = true
 			c.Writer = rewriter
 			// Filter Anthropic-Beta header only for local handling paths
 			filterAntropicBetaHeader(c)
@@ -267,6 +268,7 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
 			// proxies (e.g. NewAPI) may return a different model name and lack
 			// Amp-required fields like thinking.signature.
 			rewriter := NewResponseRewriter(c.Writer, modelName)
+			rewriter.suppressThinking = providerName != "claude"
 			c.Writer = rewriter
 			// Filter Anthropic-Beta header only for local handling paths
 			filterAntropicBetaHeader(c)

internal/api/modules/amp/response_rewriter.go

@@ -17,19 +17,18 @@ import (
 // and to keep Amp-compatible response shapes.
 type ResponseRewriter struct {
 	gin.ResponseWriter
-	body                   *bytes.Buffer
-	originalModel          string
-	isStreaming            bool
-	suppressedContentBlock map[int]struct{}
+	body             *bytes.Buffer
+	originalModel    string
+	isStreaming      bool
+	suppressThinking bool
 }
 
 // NewResponseRewriter creates a new response rewriter for model name substitution.
 func NewResponseRewriter(w gin.ResponseWriter, originalModel string) *ResponseRewriter {
 	return &ResponseRewriter{
-		ResponseWriter:         w,
-		body:                   &bytes.Buffer{},
-		originalModel:          originalModel,
-		suppressedContentBlock: make(map[int]struct{}),
+		ResponseWriter: w,
+		body:           &bytes.Buffer{},
+		originalModel:  originalModel,
 	}
 }
 
@@ -91,7 +90,8 @@ func (rw *ResponseRewriter) Write(data []byte) (int, error) {
 	}
 
 	if rw.isStreaming {
-		n, err := rw.ResponseWriter.Write(rw.rewriteStreamChunk(data))
+		rewritten := rw.rewriteStreamChunk(data)
+		n, err := rw.ResponseWriter.Write(rewritten)
 		if err == nil {
 			if flusher, ok := rw.ResponseWriter.(http.Flusher); ok {
 				flusher.Flush()
@@ -154,19 +154,10 @@ func ensureAmpSignature(data []byte) []byte {
 	return data
 }
 
-func (rw *ResponseRewriter) markSuppressedContentBlock(index int) {
-	if rw.suppressedContentBlock == nil {
-		rw.suppressedContentBlock = make(map[int]struct{})
-	}
-	rw.suppressedContentBlock[index] = struct{}{}
-}
-
-func (rw *ResponseRewriter) isSuppressedContentBlock(index int) bool {
-	_, ok := rw.suppressedContentBlock[index]
-	return ok
-}
-
 func (rw *ResponseRewriter) suppressAmpThinking(data []byte) []byte {
+	if !rw.suppressThinking {
+		return data
+	}
 	if gjson.GetBytes(data, `content.#(type=="tool_use")`).Exists() {
 		filtered := gjson.GetBytes(data, `content.#(type!="thinking")#`)
 		if filtered.Exists() {
@@ -177,33 +168,11 @@ func (rw *ResponseRewriter) suppressAmpThinking(data []byte) []byte {
 				data, err = sjson.SetBytes(data, "content", filtered.Value())
 				if err != nil {
 					log.Warnf("Amp ResponseRewriter: failed to suppress thinking blocks: %v", err)
-				} else {
-					log.Debugf("Amp ResponseRewriter: Suppressed %d thinking blocks due to tool usage", originalCount-filteredCount)
 				}
 			}
 		}
 	}
 
-	eventType := gjson.GetBytes(data, "type").String()
-	indexResult := gjson.GetBytes(data, "index")
-	if eventType == "content_block_start" && gjson.GetBytes(data, "content_block.type").String() == "thinking" && indexResult.Exists() {
-		rw.markSuppressedContentBlock(int(indexResult.Int()))
-		return nil
-	}
-	if gjson.GetBytes(data, "delta.type").String() == "thinking_delta" {
-		if indexResult.Exists() {
-			rw.markSuppressedContentBlock(int(indexResult.Int()))
-		}
-		return nil
-	}
-	if eventType == "content_block_stop" && indexResult.Exists() {
-		index := int(indexResult.Int())
-		if rw.isSuppressedContentBlock(index) {
-			delete(rw.suppressedContentBlock, index)
-			return nil
-		}
-	}
-
 	return data
 }
 
@@ -254,6 +223,10 @@ func (rw *ResponseRewriter) rewriteStreamChunk(chunk []byte) []byte {
 				jsonData := bytes.TrimPrefix(bytes.TrimSpace(lines[dataIdx]), []byte("data: "))
 				if len(jsonData) > 0 && jsonData[0] == '{' {
 					rewritten := rw.rewriteStreamEvent(jsonData)
+					if rewritten == nil {
+						i = dataIdx + 1
+						continue
+					}
 					// Emit event line
 					out = append(out, line)
 					// Emit blank lines between event and data
@@ -279,7 +252,9 @@ func (rw *ResponseRewriter) rewriteStreamChunk(chunk []byte) []byte {
 			jsonData := bytes.TrimPrefix(trimmed, []byte("data: "))
 			if len(jsonData) > 0 && jsonData[0] == '{' {
 				rewritten := rw.rewriteStreamEvent(jsonData)
-				out = append(out, append([]byte("data: "), rewritten...))
+				if rewritten != nil {
+					out = append(out, append([]byte("data: "), rewritten...))
+				}
 				i++
 				continue
 			}

internal/api/modules/amp/response_rewriter_test.go

@@ -102,7 +102,7 @@ func TestRewriteStreamChunk_MessageModel(t *testing.T) {
 }
 
 func TestRewriteStreamChunk_PreservesThinkingWithSignatureInjection(t *testing.T) {
-	rw := &ResponseRewriter{suppressedContentBlock: make(map[int]struct{})}
+	rw := &ResponseRewriter{}
 
 	chunk := []byte("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"abc\"}}\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"tool_use\",\"name\":\"bash\",\"input\":{}}}\n\n")
 	result := rw.rewriteStreamChunk(chunk)