diff --git a/proxy/account_failover.go b/proxy/account_failover.go index 02624a49..8a6ef5c0 100644 --- a/proxy/account_failover.go +++ b/proxy/account_failover.go @@ -1,6 +1,7 @@ package proxy import ( + "errors" "kiro-go/config" "kiro-go/logger" "strings" @@ -8,6 +9,80 @@ import ( const maxAccountRetryAttempts = 3 +// maxSameAccountStreamRetries bounds same-account recovery of a truncated +// stream. Kiro IDE caps its truncation retry at one (Dt3 = 1 in extension.js), +// but that budget was set for a single-credential client. This proxy also +// rotates accounts, and the two recover different failures: a same-account +// retry helps when the upstream hiccupped, rotation helps when that account's +// backend is unhealthy. Two is therefore not a copy of the IDE. +// +// Cost of the extra attempt is bounded and small. Truncation returns nil from +// CallKiroAPIContext, so it never reaches the endpoint fallback or +// maxStreamAttemptsPerEndpoint - those fire only on transport errors. Worst +// case across maxAccountRetryAttempts accounts is 3*(1+budget) requests: 9 at +// two retries versus 6 at one. +// +// The payoff lands mostly on the three fully buffered paths, whose canRetry is +// nil: a non-stream client gets a 500 with nothing usable, so one more chance +// is worth more there than on a stream that has already flushed partial text. +const maxSameAccountStreamRetries = 2 + +// errUpstreamTruncatedResponse is a soft failure raised when a transport-clean +// stream carried content but never a terminal signal. It is retryable on the +// same account and must not mark the account unhealthy. +// +// There is deliberately no empty-response error here. A stream that produced no +// output at all is already caught one layer down: parseEventStreamTracked +// returns errEmptyKiroStream when !sawOutput (proxy/kiro.go), and +// CallKiroAPIContext retries it internally. Since sawOutput is set by exactly +// the three signals classifyStreamIntegrity measures (content, reasoning, +// toolUse), an all-zero measurement can never reach this layer with a nil +// error. +var errUpstreamTruncatedResponse = errors.New("upstream truncated response without stop reason") + +// classifyStreamIntegrity decides whether an upstream stream that returned no +// transport error is actually complete. parseEventStream reports success on a +// clean EOF, so a stream that died mid-answer is otherwise indistinguishable +// from a finished one. +// +// Complete when a stopReason arrived, or when a tool call was delivered. Both +// match Kiro IDE, whose empty and truncation predicates each require +// toolCallCount === 0. +// +// Truncated when content arrived without any terminal signal. +// +// Reasoning-only with no answer is STRICTER THAN THE IDE, deliberately. The +// IDE's truncation predicate ends in (contentChars > 0 || !reasoningSeen), so +// reasoning with no answer and no stopReason is treated as complete there and +// is never retried. That is the exact shape of the production symptom this +// proxy exists to fix: thinking streams in full, then the turn dies before the +// answer or the tool call. Handing a client reasoning with no answer as a +// successful turn is what made the failure invisible, so it is classified as +// truncated here. +func classifyStreamIntegrity(contentChars, toolCallCount int, stopReason string, sawReasoning bool) error { + if strings.TrimSpace(stopReason) != "" { + return nil + } + if toolCallCount > 0 { + return nil + } + if contentChars > 0 || sawReasoning { + return errUpstreamTruncatedResponse + } + // No content, no reasoning, no tools: unreachable through the wired paths + // (errEmptyKiroStream fires first, see above). Treated as truncated rather + // than complete so a future caller that bypasses that guard still cannot + // ship an empty turn as a success. + return errUpstreamTruncatedResponse +} + +// isStreamIntegrityError reports whether err is a soft integrity failure. +// Callers may rotate accounts on these, but must not run them through +// handleAccountFailure: an upstream blip should not mark an account unhealthy. +func isStreamIntegrityError(err error) bool { + return errors.Is(err, errUpstreamTruncatedResponse) +} + func isQuotaErrorMessage(msg string) bool { msg = strings.ToLower(msg) return strings.Contains(msg, "429") || strings.Contains(msg, "quota") diff --git a/proxy/handler.go b/proxy/handler.go index 89c75174..8bcb83ed 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -1277,14 +1277,44 @@ func (h *Handler) handleClaudeStream(ctx context.Context, w http.ResponseWriter, }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return rawContentBuilder.Len(), len(toolUses), upstreamStopReason, rawThinkingBuilder.Len() > 0 + } + + // A same-account retry only happens while nothing has been flushed, so + // SSE block indices are still at their initial values and need no + // rollback. What must be cleared is every accumulator plus the thinking + // tag parser state: processClaudeText buffers up to 50 runes before + // flushing, so a short truncated attempt can leave a partial tag behind + // that would otherwise be prefixed onto the retry's first chunk. + reset := func() { + rawContentBuilder.Reset() + rawThinkingBuilder.Reset() + toolUses = nil + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + textBuffer = "" + inThinkingBlock = false + dropTagThinking = false + thinkingSource = thinkingSourceUnknown + thinkingStarted = false + eventThinkingOpen = false + } + + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, + func() bool { return !messageStarted }) if err != nil { if ctx.Err() != nil { return } lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } if !messageStarted { continue } @@ -1555,14 +1585,33 @@ func (h *Handler) handleClaudeNonStream(ctx context.Context, w http.ResponseWrit }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return len(content), len(toolUses), upstreamStopReason, thinkingContent != "" + } + + reset := func() { + content = "" + thinkingContent = "" + toolUses = nil + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + } + + // Fully buffered: nothing reaches the client until the response is + // encoded, so a retry can never duplicate output. + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, nil) if err != nil { if ctx.Err() != nil { return } lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } continue } @@ -2010,14 +2059,44 @@ func (h *Handler) handleOpenAIStream(ctx context.Context, w http.ResponseWriter, }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return rawContentBuilder.Len(), len(toolCalls), upstreamStopReason, rawReasoningBuilder.Len() > 0 + } + + // Retries only happen before anything is flushed, so chunk indices stay + // valid. The thinking tag parser state must be cleared too: processText + // holds back up to 50 runes, so a short truncated attempt would + // otherwise prepend its leftovers to the retry's first chunk. + reset := func() { + rawContentBuilder.Reset() + rawReasoningBuilder.Reset() + toolCalls = nil + toolCallIndex = 0 + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + textBuffer = "" + inThinkingBlock = false + dropTagThinking = false + thinkingSource = thinkingSourceUnknown + thinkingStarted = false + eventThinkingOpen = false + } + + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, + func() bool { return !responseStarted }) if err != nil { if ctx.Err() != nil { return } lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + // Integrity failures are upstream hiccups, not account faults. + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } if !responseStarted { continue } @@ -2133,14 +2212,34 @@ func (h *Handler) handleOpenAINonStream(ctx context.Context, w http.ResponseWrit }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return len(content), len(toolUses), upstreamStopReason, reasoningContent != "" + } + + reset := func() { + content = "" + reasoningContent = "" + toolUses = nil + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + } + + // Fully buffered: nothing reaches the client until the response is + // encoded, so a retry can never duplicate output. + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, nil) if err != nil { if ctx.Err() != nil { return } lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + // Integrity failures are upstream hiccups, not account faults. + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } continue } diff --git a/proxy/handler_test.go b/proxy/handler_test.go index 6489c4a4..f97a266f 100644 --- a/proxy/handler_test.go +++ b/proxy/handler_test.go @@ -70,6 +70,9 @@ func TestClaudeNonStreamRetriesNextAccountAfterPreResponseFailure(t *testing.T) _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ "content": "retried successfully", })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) })) defer server.Close() diff --git a/proxy/responses_handler.go b/proxy/responses_handler.go index 3d33666a..dadd9d26 100644 --- a/proxy/responses_handler.go +++ b/proxy/responses_handler.go @@ -171,14 +171,34 @@ func (h *Handler) handleResponsesNonStream( }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return len(content), len(toolUses), upstreamStopReason, reasoningContent != "" + } + + reset := func() { + content = "" + reasoningContent = "" + toolUses = nil + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + } + + // Fully buffered path: nothing reaches the client until the response is + // encoded, so a retry can never duplicate output. + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, nil) if err != nil { if ctx.Err() != nil { return } lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + // Integrity failures are upstream hiccups, not account faults. + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } continue } @@ -500,7 +520,27 @@ func (h *Handler) handleResponsesStream( }, } - err := CallKiroAPIContext(ctx, account, payload, callback) + measure := func() (int, int, string, bool) { + return fullText.Len(), len(toolUses), upstreamStopReason, reasoningText.Len() > 0 + } + + // Retries only run while responseStarted is false, i.e. before any + // content or function-call item has been sent, so the output_index / + // content_index cursors are still untouched. Only the accumulators need + // clearing. + reset := func() { + fullText.Reset() + reasoningText.Reset() + toolUses = nil + inputTokens = 0 + outputTokens = 0 + credits = 0 + realInputTokens = 0 + upstreamStopReason = "" + } + + err := runKiroWithIntegrityRetry(ctx, account, payload, callback, measure, reset, + func() bool { return !responseStarted }) if err != nil { if ctx.Err() != nil { return @@ -508,7 +548,10 @@ func (h *Handler) handleResponsesStream( if !responseStarted { lastErr = err excluded[account.ID] = true - h.handleAccountFailure(account, err) + // Integrity failures are upstream hiccups, not account faults. + if !isStreamIntegrityError(err) { + h.handleAccountFailure(account, err) + } continue } send("response.failed", map[string]interface{}{ diff --git a/proxy/responses_handler_test.go b/proxy/responses_handler_test.go index a89abe86..387c3012 100644 --- a/proxy/responses_handler_test.go +++ b/proxy/responses_handler_test.go @@ -303,6 +303,9 @@ func TestResponsesContinuationKeepsNewInstructions(t *testing.T) { _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ "content": "second reply", })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) })) defer server.Close() defer swapKiroEndpointsForTest(t, server)() @@ -381,6 +384,9 @@ func TestResponsesNonStreamRoundTrip(t *testing.T) { _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ "content": "responses non-stream OK", })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) })) defer server.Close() defer swapKiroEndpointsForTest(t, server)() @@ -467,6 +473,9 @@ func TestResponsesStreamSSE(t *testing.T) { _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ "content": "stream chunk", })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) })) defer server.Close() defer swapKiroEndpointsForTest(t, server)() diff --git a/proxy/stream_integrity.go b/proxy/stream_integrity.go new file mode 100644 index 00000000..247a9059 --- /dev/null +++ b/proxy/stream_integrity.go @@ -0,0 +1,88 @@ +package proxy + +import ( + "context" + "kiro-go/config" + "kiro-go/logger" +) + +// runKiroWithIntegrityRetry calls Kiro and recovers a truncated upstream stream +// the way Kiro IDE does: retry the same request on the same account within a +// bounded budget before surfacing the failure. +// +// callback is reused across attempts. Both CallKiroAPIContext and +// parseEventStreamTracked copy the struct before wrapping any field, so a retry +// cannot double-wrap it; per-attempt state is cleared by reset instead. +// measure reports the integrity inputs after a transport-successful call. +// reset clears per-attempt state before a same-account retry; may be nil. +// canRetry reports whether a retry is still safe (for streaming: nothing has +// been flushed to the client yet). nil means always retryable. +// +// Return contract: +// - nil: complete success only +// - transport error from CallKiroAPIContext: caller should rotate/ban as usual +// - integrity error while still retryable: retries exhausted; caller should +// rotate account without treating it as an auth/quota failure +// - integrity error after client flush: caller must surface failure to the +// client (do not fake end_turn / normal completion). Retry is unsafe. +func runKiroWithIntegrityRetry( + ctx context.Context, + account *config.Account, + payload *KiroPayload, + callback *KiroStreamCallback, + measure func() (contentChars, toolCount int, stopReason string, sawReasoning bool), + reset func(), + canRetry func() bool, +) error { + label := accountEmailForLog(account) + retryable := func() bool { + if canRetry == nil { + return true + } + return canRetry() + } + + for attempt := 0; attempt <= maxSameAccountStreamRetries; attempt++ { + if attempt > 0 && reset != nil { + reset() + } + + err := CallKiroAPIContext(ctx, account, payload, callback) + if err != nil { + return err + } + + contentChars, toolCount, stopReason, sawReasoning := measure() + integrityErr := classifyStreamIntegrity(contentChars, toolCount, stopReason, sawReasoning) + if integrityErr == nil { + return nil + } + + // A canceled client is not an integrity failure: the turn is over and + // reissuing it would only burn upstream quota. + if ctx != nil && ctx.Err() != nil { + return ctx.Err() + } + + if retryable() && attempt < maxSameAccountStreamRetries { + logger.Warnf("[StreamIntegrity] %v on %s; retrying same account (%d/%d)", + integrityErr, label, attempt+1, maxSameAccountStreamRetries) + continue + } + + if !retryable() { + // Bytes already reached the client; reissuing would duplicate output. + // Return the integrity error so callers emit an error event instead of + // finishing with a forged end_turn/tool_use success. + logger.Warnf("[StreamIntegrity] %v after client flush; signaling error (no retry)", integrityErr) + return integrityErr + } + + logger.Warnf("[StreamIntegrity] giving up after retries: %v", integrityErr) + return integrityErr + } + + // Unreachable: every branch inside the loop returns or continues, and the + // final iteration cannot continue. + return errUpstreamTruncatedResponse +} diff --git a/proxy/stream_integrity_paths_test.go b/proxy/stream_integrity_paths_test.go new file mode 100644 index 00000000..90c5a874 --- /dev/null +++ b/proxy/stream_integrity_paths_test.go @@ -0,0 +1,280 @@ +package proxy + +import ( + "encoding/json" + "kiro-go/config" + accountpool "kiro-go/pool" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +// setupIntegrityPathTest installs a single-account config plus a fake upstream +// and returns a handler wired to the reloaded pool. +func setupIntegrityPathTest(t *testing.T, server *httptest.Server) *Handler { + t.Helper() + if err := config.Init(t.TempDir() + "/config.json"); err != nil { + t.Fatalf("config.Init: %v", err) + } + if err := config.AddAccount(config.Account{ + ID: "test-account", + Enabled: true, + AccessToken: "token-test", + ProfileArn: "arn:aws:codewhisperer:profile/test", + }); err != nil { + t.Fatalf("add account: %v", err) + } + if err := config.UpdatePreferredEndpoint("kiro"); err != nil { + t.Fatalf("set preferred endpoint: %v", err) + } + if err := config.UpdateEndpointFallback(false); err != nil { + t.Fatalf("disable endpoint fallback: %v", err) + } + t.Cleanup(swapKiroEndpointsForTest(t, server)) + + p := accountpool.GetPool() + p.Reload() + return &Handler{ + pool: p, + promptCache: newPromptCacheTracker(defaultPromptCacheTTL), + } +} + +// truncatedUpstream serves content long enough to be flushed to the client but +// never sends a metadataEvent, i.e. a transport-successful truncated stream. +func truncatedUpstream(t *testing.T, hits *atomic.Int32) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if hits != nil { + hits.Add(1) + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": strings.Repeat("partial answer ", 8), + })) + })) +} + +// A truncated stream whose content already reached the client must end with an +// SSE error, never with a forged end_turn that tells the client it is done. +func TestClaudeStreamEmitsErrorOnTruncatedStream(t *testing.T) { + var hits atomic.Int32 + server := truncatedUpstream(t, &hits) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "max_tokens":100, + "messages":[{"role":"user","content":"hello"}], + "stream":true + }`)) + rec := httptest.NewRecorder() + h.handleClaudeMessages(rec, req) + + body := rec.Body.String() + if !strings.Contains(body, "partial answer") { + t.Fatalf("expected flushed content, got %s", body) + } + if strings.Contains(body, `"stop_reason":"end_turn"`) { + t.Fatalf("truncated stream must not be reported as end_turn, got %s", body) + } + if !strings.Contains(body, `"type":"error"`) { + t.Fatalf("expected SSE error event, got %s", body) + } + // Content was already flushed, so reissuing would duplicate output. + if hits.Load() != 1 { + t.Fatalf("must not retry after client flush, hits=%d", hits.Load()) + } +} + +// Non-stream buffers everything, so a truncated first attempt is safe to retry +// on the same account. The client must receive the recovered answer only. +func TestClaudeNonStreamRetriesTruncatedStream(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := hits.Add(1) + w.WriteHeader(http.StatusOK) + if n == 1 { + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "truncated attempt", + })) + return + } + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "recovered answer", + })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) + })) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "max_tokens":100, + "messages":[{"role":"user","content":"hello"}] + }`)) + rec := httptest.NewRecorder() + h.handleClaudeMessages(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200 after recovery, got %d body=%s", rec.Code, rec.Body.String()) + } + var resp ClaudeResponse + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v body=%s", err, rec.Body.String()) + } + if len(resp.Content) == 0 { + t.Fatalf("expected content, got %+v", resp) + } + text := resp.Content[0].Text + if strings.Contains(text, "truncated attempt") { + t.Fatalf("first attempt must be discarded on retry, got %q", text) + } + if !strings.Contains(text, "recovered answer") { + t.Fatalf("expected recovered answer, got %q", text) + } + if hits.Load() < 2 { + t.Fatalf("expected same-account retry, hits=%d", hits.Load()) + } +} + +// A soft integrity failure means the upstream blipped, not that the credential +// is bad. The account must stay enabled and unbanned. +func TestIntegrityFailureDoesNotBanAccount(t *testing.T) { + server := truncatedUpstream(t, nil) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "max_tokens":100, + "messages":[{"role":"user","content":"hello"}], + "stream":true + }`)) + h.handleClaudeMessages(httptest.NewRecorder(), req) + + var account *config.Account + for _, acc := range config.GetAccounts() { + if acc.ID == "test-account" { + found := acc + account = &found + break + } + } + if account == nil { + t.Fatal("account disappeared") + } + if !account.Enabled { + t.Fatal("integrity failure must not disable the account") + } + if account.BanStatus == "BANNED" || account.BanStatus == "DISABLED" { + t.Fatalf("integrity failure must not ban the account, status=%q reason=%q", + account.BanStatus, account.BanReason) + } +} + +// Responses streaming must surface response.failed instead of closing the turn +// with response.completed when the upstream truncated. +func TestResponsesStreamEmitsFailedOnTruncatedStream(t *testing.T) { + server := truncatedUpstream(t, nil) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "input":"hello", + "stream":true, + "store":false + }`)) + rec := httptest.NewRecorder() + h.handleOpenAIResponses(rec, req) + + body := rec.Body.String() + if !strings.Contains(body, "response.failed") { + t.Fatalf("expected response.failed, got %s", body) + } + if strings.Contains(body, "response.completed") { + t.Fatalf("truncated stream must not report response.completed, got %s", body) + } +} + +// OpenAI streaming must not close a truncated turn with a normal finish_reason. +func TestOpenAIStreamEmitsErrorOnTruncatedStream(t *testing.T) { + server := truncatedUpstream(t, nil) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "messages":[{"role":"user","content":"hello"}], + "stream":true + }`)) + rec := httptest.NewRecorder() + h.handleOpenAIChat(rec, req) + + body := rec.Body.String() + if !strings.Contains(body, "partial answer") { + t.Fatalf("expected flushed content, got %s", body) + } + if strings.Contains(body, `"finish_reason":"stop"`) { + t.Fatalf("truncated stream must not finish with stop, got %s", body) + } + if strings.Contains(body, "[DONE]") { + t.Fatalf("truncated stream must not report completion, got %s", body) + } + if !strings.Contains(body, `"error"`) { + t.Fatalf("expected an SSE error, got %s", body) + } +} + +// Regression for a defect the reference implementation does not cover: a short +// unflushed chunk stays inside processClaudeText's tag buffer, so a retry that +// does not clear it concatenates the previous attempt's text onto the new one. +func TestClaudeStreamRetryDoesNotLeakPreviousAttemptText(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := hits.Add(1) + w.WriteHeader(http.StatusOK) + if n == 1 { + // Short enough to stay in the buffer: never flushed to the client, + // so the stream is still retryable. + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "LEAK", + })) + return + } + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "clean answer", + })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) + })) + defer server.Close() + h := setupIntegrityPathTest(t, server) + + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{ + "model":"claude-sonnet-4.5", + "max_tokens":100, + "messages":[{"role":"user","content":"hello"}], + "stream":true + }`)) + rec := httptest.NewRecorder() + h.handleClaudeMessages(rec, req) + + body := rec.Body.String() + if hits.Load() < 2 { + t.Fatalf("expected retry for unflushed truncated stream, hits=%d", hits.Load()) + } + if strings.Contains(body, "LEAK") { + t.Fatalf("discarded attempt leaked into the retry output: %s", body) + } + if !strings.Contains(body, "clean answer") { + t.Fatalf("expected recovered answer, got %s", body) + } +} diff --git a/proxy/stream_integrity_test.go b/proxy/stream_integrity_test.go new file mode 100644 index 00000000..7604bab3 --- /dev/null +++ b/proxy/stream_integrity_test.go @@ -0,0 +1,271 @@ +package proxy + +import ( + "context" + "errors" + "kiro-go/config" + "net/http" + "net/http/httptest" + "path/filepath" + "sync/atomic" + "testing" + "time" +) + +// classifyStreamIntegrity is the completeness rule: a stream that returned no +// transport error is still incomplete when it carries no terminal signal. +// A stopReason of any value, or a delivered tool call, means complete. +// +// The reasoning-only case deliberately differs from Kiro IDE, which treats it +// as complete. See classifyStreamIntegrity's doc comment for why this proxy is +// stricter. +func TestClassifyStreamIntegrity(t *testing.T) { + for _, tc := range []struct { + name string + content int + tools int + stopReason string + sawReasoning bool + wantErr error + }{ + {"complete with stop", 12, 0, "end_turn", false, nil}, + {"complete with tools", 0, 1, "", false, nil}, + {"complete with tools despite content", 12, 1, "", false, nil}, + {"truncated content", 8, 0, "", false, errUpstreamTruncatedResponse}, + {"reasoning only stricter than ide", 0, 0, "", true, errUpstreamTruncatedResponse}, + {"no signal at all", 0, 0, "", false, errUpstreamTruncatedResponse}, + } { + t.Run(tc.name, func(t *testing.T) { + got := classifyStreamIntegrity(tc.content, tc.tools, tc.stopReason, tc.sawReasoning) + if tc.wantErr == nil { + if got != nil { + t.Fatalf("got %v, want nil", got) + } + return + } + if got == nil || got.Error() != tc.wantErr.Error() { + t.Fatalf("got %v, want %v", got, tc.wantErr) + } + if !isStreamIntegrityError(got) { + t.Fatalf("%v must be recognized as a stream integrity error", got) + } + }) + } +} + +// setupIntegrityTestUpstream points the Kiro endpoint list at a fake upstream +// and returns a restore func. Mirrors the fixture used by handler tests. +func setupIntegrityTestUpstream(t *testing.T, server *httptest.Server) func() { + t.Helper() + if err := config.Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("config.Init: %v", err) + } + if err := config.UpdatePreferredEndpoint("kiro"); err != nil { + t.Fatalf("set endpoint: %v", err) + } + if err := config.UpdateEndpointFallback(false); err != nil { + t.Fatalf("disable fallback: %v", err) + } + + oldEndpoints := kiroEndpoints + kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}} + oldClient := kiroHttpStore.Load() + kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}}) + return func() { + kiroEndpoints = oldEndpoints + kiroHttpStore.Store(oldClient) + } +} + +func integrityTestAccount() *config.Account { + return &config.Account{ + ID: "acc", + Email: "acc@test", + AccessToken: "token", + ProfileArn: "arn:aws:codewhisperer:profile/test", + } +} + +func integrityTestPayload() *KiroPayload { + payload := &KiroPayload{} + payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{ + Content: "hi", + Origin: "AI_EDITOR", + } + return payload +} + +// A stream that delivered content but no stopReason is truncated, not failed: +// the transport layer sees success and returns nil. The helper must catch that, +// reset the caller's accumulators, and retry on the same account. +// +// Note the fully-empty case is deliberately NOT used here: parseEventStreamTracked +// already returns errEmptyKiroStream when nothing was output, and +// CallKiroAPIContext retries it internally, so it never surfaces as a +// transport-successful call. +func TestRunKiroWithIntegrityRetryRecoversTruncatedThenComplete(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := hits.Add(1) + w.WriteHeader(http.StatusOK) + if n == 1 { + // Content but no metadataEvent => transport-successful truncation. + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "partial", + })) + return + } + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "recovered", + })) + _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{ + "stopReason": "end_turn", + })) + })) + defer server.Close() + defer setupIntegrityTestUpstream(t, server)() + + var content string + var stopReason string + var resets int + err := runKiroWithIntegrityRetry(context.Background(), integrityTestAccount(), integrityTestPayload(), + &KiroStreamCallback{ + OnText: func(s string, _ bool) { content += s }, + OnStopReason: func(r string) { stopReason = r }, + }, + func() (int, int, string, bool) { + return len(content), 0, stopReason, false + }, + func() { + resets++ + content = "" + stopReason = "" + }, + nil, + ) + if err != nil { + t.Fatalf("expected recovery, got %v", err) + } + if got := hits.Load(); got != 2 { + t.Fatalf("expected exactly one retry (2 upstream hits), got %d", got) + } + if resets < 1 { + t.Fatalf("expected reset before retry, got %d", resets) + } + // "partial" from the first attempt must not survive into the final result. + if content != "recovered" || stopReason != "end_turn" { + t.Fatalf("content=%q stopReason=%q", content, stopReason) + } +} + +// Once the client has already been flushed, an incomplete stream must not be +// retried (would duplicate output). Helper returns the integrity error so the +// caller can emit an error event instead of forging a normal completion. +func TestRunKiroWithIntegrityRetrySkipsRetryAfterClientFlush(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "partial", + })) + // no metadataEvent/stopReason => truncated + })) + defer server.Close() + defer setupIntegrityTestUpstream(t, server)() + + var content string + flushed := true + err := runKiroWithIntegrityRetry(context.Background(), integrityTestAccount(), integrityTestPayload(), + &KiroStreamCallback{ + OnText: func(s string, _ bool) { content += s }, + }, + func() (int, int, string, bool) { return len(content), 0, "", false }, + func() { content = "" }, + func() bool { return !flushed }, + ) + if !isStreamIntegrityError(err) { + t.Fatalf("expected integrity error after flush, got %v", err) + } + if hits.Load() != 1 { + t.Fatalf("must not retry after flush, hits=%d", hits.Load()) + } + if content != "partial" { + t.Fatalf("content=%q", content) + } +} + +// A canceled client context must not drive an integrity retry. The turn is over, +// so reissuing it would only burn upstream quota. +// +// The upstream here returns content with no stopReason, which classifies as +// truncated and would otherwise be retried; cancellation must suppress that. +func TestRunKiroWithIntegrityRetryStopsOnCanceledContext(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "partial", + })) + })) + defer server.Close() + defer setupIntegrityTestUpstream(t, server)() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var content string + var resets int + err := runKiroWithIntegrityRetry(ctx, integrityTestAccount(), integrityTestPayload(), + &KiroStreamCallback{OnText: func(s string, _ bool) { content += s }}, + func() (int, int, string, bool) { return len(content), 0, "", false }, + func() { resets++ }, + nil, + ) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if isStreamIntegrityError(err) { + t.Fatalf("cancellation must not be reported as an integrity failure: %v", err) + } + if resets != 0 { + t.Fatalf("cancellation must not trigger a retry reset, got %d", resets) + } + if got := hits.Load(); got > 1 { + t.Fatalf("canceled context must not drive retries, hits=%d", got) + } +} + +// The truncation retry budget must stay bounded and must be spent, not silently +// widened. Truncation never reaches the endpoint fallback (CallKiroAPIContext +// returns nil for it), so the only multiplier is account rotation; see +// maxSameAccountStreamRetries for the cost arithmetic. +func TestRunKiroWithIntegrityRetryStopsAfterBudgetExhausted(t *testing.T) { + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusOK) + // Always truncated: content with no terminal signal. + _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{ + "content": "partial", + })) + })) + defer server.Close() + defer setupIntegrityTestUpstream(t, server)() + + var content string + err := runKiroWithIntegrityRetry(context.Background(), integrityTestAccount(), integrityTestPayload(), + &KiroStreamCallback{OnText: func(s string, _ bool) { content += s }}, + func() (int, int, string, bool) { return len(content), 0, "", false }, + func() { content = "" }, + nil, + ) + if !isStreamIntegrityError(err) { + t.Fatalf("expected integrity error once the budget is spent, got %v", err) + } + if got := hits.Load(); got != int32(maxSameAccountStreamRetries+1) { + t.Fatalf("upstream hits=%d, want %d (initial attempt plus %d retry)", + got, maxSameAccountStreamRetries+1, maxSameAccountStreamRetries) + } +}