Unified · Split
| @@ -100,17 +100,14 @@ func TestRewriteStreamChunk_MessageModel(t *testing.T) { · 99 unmodified lines | |||
| 100 | } | 100 | } |
| 101 | } | 101 | } |
| 102 | 102 | ||
| 103 | func TestRewriteStreamChunk_SuppressesThinkingContentBlockFrames(t *testing.T) { | 103 | func TestRewriteStreamChunk_PassesThroughThinkingBlocks(t *testing.T) { |
| 104 | rw := &ResponseRewriter{suppressedContentBlock: make(map[int]struct{})} | 104 | rw := &ResponseRewriter{} |
| 105 | 105 | ||
| 106 | 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") | 106 | 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") |
| 107 | result := rw.rewriteStreamChunk(chunk) | 107 | result := rw.rewriteStreamChunk(chunk) |
| 108 | 108 | ||
| 109 | if contains(result, []byte("\"thinking\"")) || contains(result, []byte("\"thinking_delta\"")) { | 109 | if !contains(result, []byte("\"thinking_delta\"")) { |
| 110 | t.Fatalf("expected thinking content_block frames to be suppressed, got %s", string(result)) | 110 | t.Fatalf("expected thinking blocks to pass through in streaming, got %s", string(result)) |
| 111 | } | ||
| 112 | if contains(result, []byte("content_block_stop")) { | ||
| 113 | t.Fatalf("expected suppressed thinking content_block_stop to be removed, got %s", string(result)) | ||
| 114 | } | 111 | } |
| 115 | if !contains(result, []byte("\"tool_use\"")) { | 112 | if !contains(result, []byte("\"tool_use\"")) { |
| 116 | t.Fatalf("expected tool_use content_block frame to remain, got %s", string(result)) | 113 | t.Fatalf("expected tool_use content_block frame to remain, got %s", string(result)) |