diff --git a/proxy/handler.go b/proxy/handler.go index 8bcb83ed..e7e01fb4 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -1310,6 +1310,18 @@ func (h *Handler) handleClaudeStream(ctx context.Context, w http.ResponseWriter, if ctx.Err() != nil { return } + if isKiroInputTooLongError(err) { + h.recordFailureWithDetails("claude", model, account.ID, err) + if !messageStarted { + h.sendClaudeError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + h.sendSSE(w, flusher, "error", map[string]interface{}{ + "type": "error", + "error": map[string]string{"type": "invalid_request_error", "message": err.Error()}, + }) + return + } lastErr = err excluded[account.ID] = true if !isStreamIntegrityError(err) { @@ -1607,6 +1619,11 @@ func (h *Handler) handleClaudeNonStream(ctx context.Context, w http.ResponseWrit if ctx.Err() != nil { return } + if isKiroInputTooLongError(err) { + h.recordFailureWithDetails("claude", model, account.ID, err) + h.sendClaudeError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } lastErr = err excluded[account.ID] = true if !isStreamIntegrityError(err) { @@ -2091,6 +2108,19 @@ func (h *Handler) handleOpenAIStream(ctx context.Context, w http.ResponseWriter, if ctx.Err() != nil { return } + if isKiroInputTooLongError(err) { + h.recordFailureWithDetails("openai", model, account.ID, err) + if !responseStarted { + h.sendOpenAIError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } + data, _ := json.Marshal(map[string]interface{}{ + "error": map[string]string{"message": err.Error(), "type": "invalid_request_error"}, + }) + fmt.Fprintf(w, "data: %s\n\n", data) + flusher.Flush() + return + } lastErr = err excluded[account.ID] = true // Integrity failures are upstream hiccups, not account faults. @@ -2234,6 +2264,11 @@ func (h *Handler) handleOpenAINonStream(ctx context.Context, w http.ResponseWrit if ctx.Err() != nil { return } + if isKiroInputTooLongError(err) { + h.recordFailureWithDetails("openai", model, account.ID, err) + h.sendOpenAIError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } lastErr = err excluded[account.ID] = true // Integrity failures are upstream hiccups, not account faults. diff --git a/proxy/input_too_long_test.go b/proxy/input_too_long_test.go new file mode 100644 index 00000000..248a6d08 --- /dev/null +++ b/proxy/input_too_long_test.go @@ -0,0 +1,299 @@ +package proxy + +import ( + "context" + "encoding/json" + "fmt" + "kiro-go/config" + accountpool "kiro-go/pool" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" +) + +const replayElisionMarker = "[Earlier turns of this conversation were omitted to fit the model's input limit.]" + +func replayFixtureText(tag string, approxTokens int) string { + unit := tag + " lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod. " + return strings.Repeat(unit, approxTokens/8+1) +} + +func storeOversizedReplayChain(t *testing.T, turns, tokensPerSide int) *ResponsesObject { + t.Helper() + if err := config.Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("config.Init: %v", err) + } + + var newest *ResponsesObject + previousID := "" + for i := 0; i < turns; i++ { + input, err := json.Marshal(fmt.Sprintf("turn %d user: %s", i, replayFixtureText(fmt.Sprintf("t%d", i), tokensPerSide))) + if err != nil { + t.Fatalf("marshal turn %d: %v", i, err) + } + response := &ResponsesObject{ + ID: fmt.Sprintf("resp_oversized_%d", i), + Object: "response", + Status: "completed", + Model: "gpt-5.6-terra", + StoredInput: input, + PreviousResponseID: previousID, + Output: []ResponseOutputItem{{ + Type: "message", + Role: "assistant", + Content: []ResponseContentPart{{ + Type: "output_text", + Text: fmt.Sprintf("turn %d assistant: %s", i, replayFixtureText("answer", tokensPerSide)), + }}, + }}, + } + if err := saveResponse(response); err != nil { + t.Fatalf("save turn %d: %v", i, err) + } + previousID = response.ID + newest = response + } + return newest +} + +func replayTokenEstimateForTest(messages []OpenAIMessage) int { + total := 0 + for _, message := range messages { + total += estimateOpenAIContentTokens(message.Content) + total += estimateApproxTokens(message.ToolCallID) + for _, call := range message.ToolCalls { + total += estimateApproxTokens(call.Function.Name) + total += estimateApproxTokens(call.Function.Arguments) + } + } + return total +} + +// A compacting Codex client keeps previous_response_id while replacing its +// current input with a short summary. Replaying the full pre-compaction chain +// defeats compaction and eventually triggers CONTENT_LENGTH_EXCEEDS_THRESHOLD. +func TestResponsesReplayHistoryLeavesRoomForCompactedInput(t *testing.T) { + newest := storeOversizedReplayChain(t, 32, 5_000) + replayed := expandPreviousResponseHistory(newest, "gpt-5.6-terra") + budget := getContextWindowSize("gpt-5.6-terra") / 2 + got := replayTokenEstimateForTest(replayed) + if got > budget { + t.Fatalf("replayed history uses %d tokens, over %d-token budget", got, budget) + } + + var joined strings.Builder + for _, message := range replayed { + if text, ok := message.Content.(string); ok { + joined.WriteString(text) + joined.WriteByte('\n') + } + } + transcript := joined.String() + if !strings.Contains(transcript, "turn 31 user") { + t.Fatal("newest turn was dropped") + } + if strings.Contains(transcript, "turn 0 user") { + t.Fatal("oldest turn must be dropped when replay exceeds its budget") + } + if !strings.Contains(transcript, replayElisionMarker) { + t.Fatal("history elision must be disclosed to the model") + } +} + +// A deterministic request-size rejection cannot succeed on another endpoint. +// Retrying it only multiplies latency and upstream load. +func TestCallKiroAPIDoesNotFallbackOnContentLengthRejection(t *testing.T) { + if err := config.Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("config.Init: %v", err) + } + var hits atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"message":"Input is too long.","reason":"CONTENT_LENGTH_EXCEEDS_THRESHOLD"}`)) + })) + defer server.Close() + + oldResolver := resolveKiroEndpoints + resolveKiroEndpoints = func(*config.Account) []kiroEndpoint { + return []kiroEndpoint{ + {URL: server.URL, Origin: "AI_EDITOR", Name: "Kiro IDE"}, + {URL: server.URL, Origin: "AI_EDITOR", Name: "CodeWhisperer"}, + {URL: server.URL, Origin: "AI_EDITOR", Name: "AmazonQ"}, + } + } + t.Cleanup(func() { resolveKiroEndpoints = oldResolver }) + + oldClient := kiroHttpStore.Load() + kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}}) + t.Cleanup(func() { kiroHttpStore.Store(oldClient) }) + + account := &config.Account{ + ID: "content-length-account", + AccessToken: "token", + ProfileArn: "arn:aws:codewhisperer:profile/test", + } + err := CallKiroAPIContext(context.Background(), account, integrityTestPayload(), &KiroStreamCallback{}) + if err == nil || !strings.Contains(err.Error(), "CONTENT_LENGTH_EXCEEDS_THRESHOLD") { + t.Fatalf("expected content-length error, got %v", err) + } + if hits.Load() != 1 { + t.Fatalf("content-length rejection must not fallback, hits=%d", hits.Load()) + } +} + +// The outer account loop must also stop after the first deterministic 400. +func TestResponsesDoesNotRotateAccountsOnContentLengthRejection(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.StatusBadRequest) + _, _ = w.Write([]byte(`{"message":"Input is too long.","reason":"CONTENT_LENGTH_EXCEEDS_THRESHOLD"}`)) + })) + defer server.Close() + + if err := config.Init(filepath.Join(t.TempDir(), "config.json")); err != nil { + t.Fatalf("config.Init: %v", err) + } + for i := 0; i < maxAccountRetryAttempts; i++ { + if err := config.AddAccount(config.Account{ + ID: fmt.Sprintf("oversize-account-%d", i), + Enabled: true, + AccessToken: "token", + ProfileArn: "arn:aws:codewhisperer:profile/test", + }); err != nil { + t.Fatalf("add account %d: %v", i, 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 fallback: %v", err) + } + t.Cleanup(swapKiroEndpointsForTest(t, server)) + + pool := accountpool.GetPool() + pool.Reload() + handler := &Handler{pool: pool, promptCache: newPromptCacheTracker(defaultPromptCacheTTL)} + req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{ + "model":"gpt-5.6-terra", + "input":"compacted summary", + "store":false + }`)) + rec := httptest.NewRecorder() + handler.handleOpenAIResponses(rec, req) + + if hits.Load() != 1 { + t.Fatalf("content-length rejection must not rotate accounts, hits=%d", hits.Load()) + } + if rec.Code != http.StatusBadRequest { + t.Fatalf("status=%d, want 400; body=%s", rec.Code, rec.Body.String()) + } + + streamReq := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{ + "model":"gpt-5.6-terra", + "input":"compacted summary", + "stream":true, + "store":false + }`)) + streamRec := httptest.NewRecorder() + handler.handleOpenAIResponses(streamRec, streamReq) + + if hits.Load() != 2 { + t.Fatalf("streaming rejection must use one additional attempt, total hits=%d", hits.Load()) + } + streamBody := streamRec.Body.String() + if !strings.Contains(streamBody, "response.failed") || !strings.Contains(streamBody, `"type":"invalid_request_error"`) { + t.Fatalf("stream must report a non-retryable invalid request, body=%s", streamBody) + } + if strings.Count(streamBody, "event: response.in_progress") != 1 { + t.Fatalf("stream must not rotate accounts, body=%s", streamBody) + } +} + +// Codex custom tool outputs can exceed a megabyte. They become structured +// current-message toolResults, which the byte-level payload limiter must trim; +// shortening only currentMessage.content leaves the oversized copy untouched. +func TestOpenAIToKiroTruncatesOversizedCurrentCustomToolResult(t *testing.T) { + call := ToolCall{ID: "call_large", Type: "function"} + call.Function.Name = "exec" + call.Function.Arguments = `{"input":"read a large command output"}` + payload := OpenAIToKiro(&OpenAIRequest{ + Model: "gpt-5.6-terra", + Messages: []OpenAIMessage{ + {Role: "assistant", Content: "", ToolCalls: []ToolCall{call}}, + {Role: "tool", ToolCallID: "call_large", Content: strings.Repeat("large tool output line\n", 80_000)}, + }, + }, false) + + encoded, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + if len(encoded) > maxPayloadBytes { + t.Fatalf("payload remains oversized after truncation: %d > %d", len(encoded), maxPayloadBytes) + } + ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext + if ctx == nil || len(ctx.ToolResults) != 1 { + t.Fatalf("expected one structured tool result, got %#v", ctx) + } + result := ctx.ToolResults[0] + if result.ToolUseID != "call_large" { + t.Fatalf("tool result pairing changed: %q", result.ToolUseID) + } + if len(result.Content) != 1 || !strings.Contains(result.Content[0].Text, "Tool output truncated") { + t.Fatalf("tool output truncation was not disclosed: %#v", result.Content) + } +} + +// Once a tool result is followed by another user message it moves into Kiro +// history and sanitizeKiroHistory flattens it into ordinary text. The minimum +// recent-history count must not make one multi-megabyte entry untrimmable. +func TestOpenAIToKiroTruncatesOversizedToolResultMovedIntoHistory(t *testing.T) { + call := ToolCall{ID: "call_historical_large", Type: "function"} + call.Function.Name = "exec" + call.Function.Arguments = `{"input":"read a large command output"}` + payload := OpenAIToKiro(&OpenAIRequest{ + Model: "gpt-5.6-sol", + Messages: []OpenAIMessage{ + {Role: "assistant", Content: "", ToolCalls: []ToolCall{call}}, + {Role: "tool", ToolCallID: "call_historical_large", Content: strings.Repeat("historical tool output line\n", 120_000)}, + {Role: "user", Content: "continue"}, + }, + }, false) + + encoded, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + if len(encoded) > maxPayloadBytes { + t.Fatalf("payload remains oversized after history truncation: %d > %d", len(encoded), maxPayloadBytes) + } + historyTokenLimit := int(float64(getContextWindowSize("gpt-5.6-sol")) * payloadHistoryWindowShare) + if got := estimateKiroHistoryTokens(payload.ConversationState.History); got > historyTokenLimit { + t.Fatalf("history remains over token budget: %d > %d", got, historyTokenLimit) + } + if len(payload.ConversationState.History) == 0 { + t.Fatal("history structure must remain present") + } + var historyText strings.Builder + for _, entry := range payload.ConversationState.History { + if entry.UserInputMessage != nil { + historyText.WriteString(entry.UserInputMessage.Content) + } + if entry.AssistantResponseMessage != nil { + historyText.WriteString(entry.AssistantResponseMessage.Content) + } + } + if !strings.Contains(historyText.String(), "Earlier message content truncated") { + t.Fatal("history content truncation must be disclosed to the model") + } + if payload.ConversationState.CurrentMessage.UserInputMessage.Content != "continue" { + t.Fatalf("current user instruction was changed: %q", payload.ConversationState.CurrentMessage.UserInputMessage.Content) + } +} diff --git a/proxy/kiro.go b/proxy/kiro.go index 7f263e2a..14ed50dc 100644 --- a/proxy/kiro.go +++ b/proxy/kiro.go @@ -479,8 +479,14 @@ endpointLoop: if resp.StatusCode != 200 { errBody, _ := io.ReadAll(resp.Body) resp.Body.Close() - lastErr = fmt.Errorf("HTTP %d from %s: %s", resp.StatusCode, ep.Name, string(errBody)) - // Authentication errors and payment errors are not retried across endpoints. + lastErr = newKiroHTTPError(resp.StatusCode, ep.Name, errBody) + // Authentication, payment, and deterministic request-shape errors + // cannot succeed on another endpoint. In particular, retrying an + // oversized payload only multiplies latency and upstream load. + if isKiroInputTooLongError(lastErr) { + logger.Warnf("[KiroAPI] Endpoint %s rejected oversized input; not retrying: %v", ep.Name, lastErr) + return lastErr + } if resp.StatusCode == 401 || resp.StatusCode == 403 || resp.StatusCode == 402 { return lastErr } diff --git a/proxy/kiro_errors.go b/proxy/kiro_errors.go new file mode 100644 index 00000000..0f9b687e --- /dev/null +++ b/proxy/kiro_errors.go @@ -0,0 +1,50 @@ +package proxy + +import ( + "encoding/json" + "errors" + "fmt" + "strings" +) + +const kiroContentLengthReason = "CONTENT_LENGTH_EXCEEDS_THRESHOLD" + +// kiroHTTPError preserves structured upstream status and reason fields so +// deterministic request errors can stop both endpoint and account retries. +type kiroHTTPError struct { + StatusCode int + Endpoint string + Message string + Reason string + Body string +} + +func newKiroHTTPError(statusCode int, endpoint string, body []byte) *kiroHTTPError { + parsed := struct { + Message string `json:"message"` + Reason string `json:"reason"` + }{} + _ = json.Unmarshal(body, &parsed) + return &kiroHTTPError{ + StatusCode: statusCode, + Endpoint: endpoint, + Message: parsed.Message, + Reason: parsed.Reason, + Body: string(body), + } +} + +func (e *kiroHTTPError) Error() string { + return fmt.Sprintf("HTTP %d from %s: %s", e.StatusCode, e.Endpoint, e.Body) +} + +func isKiroInputTooLongError(err error) bool { + if err == nil { + return false + } + var upstreamErr *kiroHTTPError + if errors.As(err, &upstreamErr) { + return upstreamErr.StatusCode == 400 && strings.EqualFold(strings.TrimSpace(upstreamErr.Reason), kiroContentLengthReason) + } + return strings.Contains(strings.ToUpper(err.Error()), kiroContentLengthReason) +} diff --git a/proxy/responses_handler.go b/proxy/responses_handler.go index dadd9d26..1403ffaf 100644 --- a/proxy/responses_handler.go +++ b/proxy/responses_handler.go @@ -42,6 +42,12 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) storeResponse = *req.Store } + // Resolve the model before replay expansion because the history budget is + // derived from the actual model's context window. A thinking suffix is a + // client-side convention and must not force the smallest fallback budget. + thinkingCfg := config.GetThinkingConfig() + actualModel, thinking := ParseModelAndThinking(req.Model, thinkingCfg.Suffix) + var historyMessages []OpenAIMessage if req.PreviousResponseID != "" { prev, loadErr := loadResponse(req.PreviousResponseID) @@ -50,7 +56,7 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) fmt.Sprintf("previous_response_id not found: %v", loadErr)) return } - historyMessages = expandPreviousResponseHistory(prev) + historyMessages = expandPreviousResponseHistory(prev, actualModel) } inputMessages, err := parseResponsesInput(req.Input) @@ -104,8 +110,6 @@ func (h *Handler) handleOpenAIResponses(w http.ResponseWriter, r *http.Request) openaiReq.MaxTokens = *req.MaxOutputTokens } - thinkingCfg := config.GetThinkingConfig() - actualModel, thinking := ParseModelAndThinking(req.Model, thinkingCfg.Suffix) openaiReq.Model = actualModel estimatedInputTokens := estimateOpenAIRequestInputTokens(openaiReq) @@ -193,6 +197,11 @@ func (h *Handler) handleResponsesNonStream( if ctx.Err() != nil { return } + if isKiroInputTooLongError(err) { + h.recordFailureWithDetails("responses", model, account.ID, err) + h.sendOpenAIError(w, http.StatusBadRequest, "invalid_request_error", err.Error()) + return + } lastErr = err excluded[account.ID] = true // Integrity failures are upstream hiccups, not account faults. @@ -545,7 +554,8 @@ func (h *Handler) handleResponsesStream( if ctx.Err() != nil { return } - if !responseStarted { + inputTooLong := isKiroInputTooLongError(err) + if !responseStarted && !inputTooLong { lastErr = err excluded[account.ID] = true // Integrity failures are upstream hiccups, not account faults. @@ -554,13 +564,17 @@ func (h *Handler) handleResponsesStream( } continue } + errorType := "server_error" + if inputTooLong { + errorType = "invalid_request_error" + } send("response.failed", map[string]interface{}{ "type": "response.failed", "response": map[string]interface{}{ "id": respID, "status": "failed", "error": map[string]string{ - "type": "server_error", + "type": errorType, "message": err.Error(), }, }, diff --git a/proxy/responses_handler_test.go b/proxy/responses_handler_test.go index 387c3012..6c977bc8 100644 --- a/proxy/responses_handler_test.go +++ b/proxy/responses_handler_test.go @@ -196,7 +196,7 @@ func TestResponsesPreviousResponseIDExpands(t *testing.T) { }, } - expanded := expandPreviousResponseHistory(prev) + expanded := expandPreviousResponseHistory(prev, prev.Model) if len(expanded) != 3 { t.Fatalf("expected 3 messages from history, got %d (%+v)", len(expanded), expanded) } @@ -256,7 +256,7 @@ func TestResponsesPreviousResponseIDExpandsFullChain(t *testing.T) { t.Fatalf("save b: %v", err) } - expanded := expandPreviousResponseHistory(b) + expanded := expandPreviousResponseHistory(b, b.Model) var transcript []string for _, m := range expanded { diff --git a/proxy/responses_history.go b/proxy/responses_history.go index 5525d1dc..5b2f85db 100644 --- a/proxy/responses_history.go +++ b/proxy/responses_history.go @@ -1,47 +1,129 @@ package proxy +import "strings" + // maxResponsesHistoryDepth caps how far back we walk the previous_response_id // chain when expanding history. The cap prevents pathological loops in // corrupted/cyclic stores from running forever; legitimate chains rarely go // this deep within the 30-day TTL. const maxResponsesHistoryDepth = 64 -// expandPreviousResponseHistory rebuilds the conversation history that led up -// to prev. It walks the previous_response_id chain backwards (oldest → newest) -// and emits OpenAI messages for both stored inputs and stored outputs of every -// ancestor, so a multi-turn /v1/responses session preserves full context. +// replayHistoryWindowShare bounds replayed ancestor history as a fraction of the +// model's context window, reserving the remainder for the client's current input +// and its tool schemas. // -// If a link in the chain is missing on disk (e.g. expired past TTL or the -// referenced ID was deleted), expansion stops at the deepest reachable -// ancestor instead of failing — the most recent context is still useful. -func expandPreviousResponseHistory(prev *ResponsesObject) []OpenAIMessage { +// A depth cap alone is not a size cap: 64 ancestors of full input and output can +// exceed any window on their own. This matters when a client compacts while its +// new request still points at the pre-compaction response. Replaying every turn +// the client discarded would refill the context and trigger compaction again. +const replayHistoryWindowShare = 0.5 + +// replayTruncationNote marks where replayed ancestor turns were elided. +const replayTruncationNote = "[Earlier turns of this conversation were omitted to fit the model's input limit.]" + +func maxReplayHistoryTokens(model string) int { + return int(float64(getContextWindowSize(model)) * replayHistoryWindowShare) +} + +// expandPreviousResponseHistory rebuilds the conversation history that led up +// to prev. It walks the previous_response_id chain backwards and emits messages +// oldest to newest. Replayed history is bounded by a share of the model window; +// the current request input is authoritative and is never trimmed here. +func expandPreviousResponseHistory(prev *ResponsesObject, model string) []OpenAIMessage { if prev == nil { return nil } chain := collectAncestorChain(prev) - - messages := make([]OpenAIMessage, 0) + groups := make([][]OpenAIMessage, 0, len(chain)) for _, node := range chain { - // Inject the instructions stored on the ancestor as a system message - // so it remains in scope for downstream turns. Without this, an early - // system prompt set on response A would be lost the moment a new - // turn omits it. + group := make([]OpenAIMessage, 0, 4) if node.Instructions != "" { - messages = append(messages, OpenAIMessage{ + group = append(group, OpenAIMessage{ Role: "system", Content: node.Instructions, }) } if prior, err := parseResponsesInput(node.StoredInput); err == nil { - messages = append(messages, prior...) + group = append(group, prior...) + } + group = append(group, outputToMessages(node.Output)...) + groups = append(groups, group) + } + + return flattenReplayGroups(groups, maxReplayHistoryTokens(model)) +} + +// flattenReplayGroups drops complete oldest turns until the newest suffix fits +// within budget. The newest turn is retained even if it alone exceeds budget; +// downstream payload truncation remains the final byte-level safety net. +func flattenReplayGroups(groups [][]OpenAIMessage, budget int) []OpenAIMessage { + keepFrom := 0 + if budget > 0 { + total := 0 + keepFrom = len(groups) + for i := len(groups) - 1; i >= 0; i-- { + total += estimateReplayGroupTokens(groups[i]) + if total > budget && i < len(groups)-1 { + break + } + keepFrom = i } - messages = append(messages, outputToMessages(node.Output)...) } + messages := make([]OpenAIMessage, 0) + if keepFrom > 0 { + messages = append(messages, OpenAIMessage{ + Role: "user", + Content: replayTruncationNote, + }) + } + for _, group := range groups[keepFrom:] { + messages = append(messages, group...) + } + if keepFrom > 0 { + messages = dropOrphanedLeadingToolMessages(messages) + } return messages } +// dropOrphanedLeadingToolMessages removes tool results whose assistant call was +// in an older turn removed by replay trimming. Kiro rejects unpaired results. +func dropOrphanedLeadingToolMessages(messages []OpenAIMessage) []OpenAIMessage { + seenCalls := make(map[string]bool) + out := make([]OpenAIMessage, 0, len(messages)) + sawAssistantCall := false + + for _, message := range messages { + if message.Role == "assistant" && len(message.ToolCalls) > 0 { + sawAssistantCall = true + for _, call := range message.ToolCalls { + if call.ID != "" { + seenCalls[call.ID] = true + } + } + } + if message.Role == "tool" && !sawAssistantCall && !seenCalls[message.ToolCallID] { + continue + } + out = append(out, message) + } + return out +} + +func estimateReplayGroupTokens(group []OpenAIMessage) int { + total := 0 + for _, message := range group { + total += estimateOpenAIContentTokens(message.Content) + total += estimateApproxTokens(message.ToolCallID) + for _, call := range message.ToolCalls { + total += estimateApproxTokens(call.Function.Name) + total += estimateApproxTokens(call.Function.Arguments) + } + } + return total +} + // collectAncestorChain walks previous_response_id backwards, returning the // chain in oldest-first order: [root, ..., parent, prev]. The walker is // bounded by maxResponsesHistoryDepth and a visited-set to short-circuit @@ -67,7 +149,6 @@ func collectAncestorChain(prev *ResponsesObject) []*ResponsesObject { cursor = ancestor } - // Reverse to oldest-first. for i, j := 0, len(stack)-1; i < j; i, j = i+1, j-1 { stack[i], stack[j] = stack[j], stack[i] } @@ -112,11 +193,11 @@ func joinTextParts(parts []ResponseContentPart) string { if len(parts) == 0 { return "" } - out := "" - for _, p := range parts { - if p.Type == "output_text" || p.Type == "text" || p.Type == "input_text" { - out += p.Text + var out strings.Builder + for _, part := range parts { + if part.Type == "output_text" || part.Type == "text" || part.Type == "input_text" { + out.WriteString(part.Text) } } - return out + return out.String() } diff --git a/proxy/translator.go b/proxy/translator.go index a86824dc..32011920 100644 --- a/proxy/translator.go +++ b/proxy/translator.go @@ -8,6 +8,7 @@ import ( "regexp" "strings" "time" + "unicode/utf8" "github.com/google/uuid" ) @@ -56,10 +57,24 @@ const toolResultImagePlaceholder = "[Tool returned an image; the image is attach // leave room for headers and minor serialization overhead. const maxPayloadBytes = 900 * 1024 +// payloadHistoryWindowShare reserves half of the model window for the current +// instruction, tool schemas/results, and output headroom. This applies to Kiro +// history produced directly from a client request, not only stored Responses +// replay, because Codex can resend prior tool output inline. +const payloadHistoryWindowShare = 0.5 + // truncationPlaceholder is inserted in history where older turns were dropped to // fit within maxPayloadBytes. const truncationPlaceholder = "[Earlier conversation history was truncated to fit the model's input limit. Older messages and tool activity have been omitted.]" +// toolResultTruncationMarker makes lossy tool-output trimming explicit while +// preserving the result's head, tail, call ID, and structured pairing. +const toolResultTruncationMarker = "\n...[Tool output truncated to fit the model input limit]...\n" + +// historyContentTruncationMarker is used only when a retained recent history +// entry is individually too large to fit, so its position and role stay intact. +const historyContentTruncationMarker = "\n...[Earlier message content truncated to fit the model input limit]...\n" + // minRecentHistoryTurns is the number of most-recent history entries always kept // (in addition to system priming and the active tool turn) when truncating. const minRecentHistoryTurns = 4 @@ -1730,6 +1745,13 @@ func truncatePayloadToLimit(payload *KiroPayload, hasPriming bool) { if payload == nil { return } + + // Byte size alone is not a reliable context limit: source code, JSON, and + // command output can tokenize densely. Cap inline history by the model window + // even when the serialized request is below maxPayloadBytes. + historyTokenLimit := int(float64(getContextWindowSize(currentMessageModelID(payload))) * payloadHistoryWindowShare) + truncateHistoryContentToTokenLimit(payload, historyTokenLimit) + if payloadByteSize(payload) <= maxPayloadBytes { return } @@ -1789,7 +1811,22 @@ func truncatePayloadToLimit(payload *KiroPayload, hasPriming bool) { rebuilt = append(rebuilt, tail...) payload.ConversationState.History = rebuilt - // If still too large (current message or retained tail alone exceeds the + // The minimum recent-turn count preserves structure, but an individual + // retained entry can itself be multiple megabytes (for example, a prior + // Codex exec result flattened into history). Shorten retained history text + // before touching the current turn. + if payloadByteSize(payload) > maxPayloadBytes { + truncateHistoryContentToLimit(payload, maxPayloadBytes) + } + + // Structured current tool results are outside current-message content and + // can dominate the payload (Codex exec output commonly exceeds 1 MB). Trim + // their text first so the user's current instruction remains intact. + if payloadByteSize(payload) > maxPayloadBytes { + truncateCurrentToolResults(payload) + } + + // If still too large (current message or tool schemas alone exceed the // limit), shrink the current message content as a last resort. if payloadByteSize(payload) > maxPayloadBytes { truncateCurrentMessage(payload) @@ -1828,6 +1865,227 @@ func currentMessageModelID(payload *KiroPayload) string { return payload.ConversationState.CurrentMessage.UserInputMessage.ModelID } +func estimateKiroHistoryTokens(history []KiroHistoryMessage) int { + total := 0 + for _, entry := range history { + if entry.UserInputMessage != nil { + total += estimateApproxTokens(entry.UserInputMessage.Content) + if ctx := entry.UserInputMessage.UserInputMessageContext; ctx != nil { + for _, result := range ctx.ToolResults { + for _, content := range result.Content { + total += estimateApproxTokens(content.Text) + } + } + for _, tool := range ctx.Tools { + total += estimateApproxTokens(tool.ToolSpecification.Name) + total += estimateApproxTokens(tool.ToolSpecification.Description) + total += estimateJSONTokens(tool.ToolSpecification.InputSchema.JSON) + } + } + } + if entry.AssistantResponseMessage != nil { + total += estimateApproxTokens(entry.AssistantResponseMessage.Content) + for _, toolUse := range entry.AssistantResponseMessage.ToolUses { + total += estimateApproxTokens(toolUse.Name) + total += estimateJSONTokens(toolUse.Input) + } + } + } + return total +} + +// truncateHistoryContentToTokenLimit handles dense text that can exceed Kiro's +// context threshold while remaining below the byte cap. It keeps message +// structure and trims only the largest retained text fields. +func truncateHistoryContentToTokenLimit(payload *KiroPayload, tokenLimit int) { + if payload == nil || tokenLimit <= 0 { + return + } + + for { + currentTokens := estimateKiroHistoryTokens(payload.ConversationState.History) + if currentTokens <= tokenLimit { + return + } + + entryIndex, largestTokens, largestBytes, isUser := -1, 0, 0, false + for i := range payload.ConversationState.History { + entry := &payload.ConversationState.History[i] + if entry.UserInputMessage != nil { + tokens := estimateApproxTokens(entry.UserInputMessage.Content) + if tokens > largestTokens { + entryIndex, largestTokens, largestBytes, isUser = i, tokens, len(entry.UserInputMessage.Content), true + } + } + if entry.AssistantResponseMessage != nil { + tokens := estimateApproxTokens(entry.AssistantResponseMessage.Content) + if tokens > largestTokens { + entryIndex, largestTokens, largestBytes, isUser = i, tokens, len(entry.AssistantResponseMessage.Content), false + } + } + } + if entryIndex < 0 || largestBytes <= len(historyContentTruncationMarker) { + return + } + + targetTokens := largestTokens - (currentTokens - tokenLimit) - 64 + if targetTokens < 1 { + targetTokens = 1 + } + targetBytes := largestBytes * targetTokens / max(1, largestTokens) + if targetBytes < len(historyContentTruncationMarker) { + targetBytes = len(historyContentTruncationMarker) + } + + entry := &payload.ConversationState.History[entryIndex] + var content *string + if isUser { + content = &entry.UserInputMessage.Content + } else { + content = &entry.AssistantResponseMessage.Content + } + trimmed := truncateTextAroundMarker(*content, targetBytes, historyContentTruncationMarker) + if len(trimmed) >= len(*content) { + trimmed = historyContentTruncationMarker + } + *content = trimmed + } +} + +// truncateHistoryContentToLimit reduces the largest retained history text +// fields until the complete serialized payload fits. Message roles, ordering, +// tool-call structures, and the current user instruction remain unchanged. +func truncateHistoryContentToLimit(payload *KiroPayload, limit int) { + if payload == nil || limit <= 0 { + return + } + + for { + currentSize := payloadByteSize(payload) + if currentSize <= limit { + return + } + + entryIndex, largestBytes, isUser := -1, 0, false + for i := range payload.ConversationState.History { + entry := &payload.ConversationState.History[i] + if entry.UserInputMessage != nil && len(entry.UserInputMessage.Content) > largestBytes { + entryIndex, largestBytes, isUser = i, len(entry.UserInputMessage.Content), true + } + if entry.AssistantResponseMessage != nil && len(entry.AssistantResponseMessage.Content) > largestBytes { + entryIndex, largestBytes, isUser = i, len(entry.AssistantResponseMessage.Content), false + } + } + if entryIndex < 0 || largestBytes <= len(historyContentTruncationMarker) { + return + } + + excess := currentSize - limit + targetBytes := largestBytes - excess - 256 + if targetBytes < len(historyContentTruncationMarker) { + targetBytes = len(historyContentTruncationMarker) + } + if targetBytes >= largestBytes { + targetBytes = largestBytes - 1 + } + + entry := &payload.ConversationState.History[entryIndex] + var content *string + if isUser { + content = &entry.UserInputMessage.Content + } else { + content = &entry.AssistantResponseMessage.Content + } + trimmed := truncateTextAroundMarker(*content, targetBytes, historyContentTruncationMarker) + if len(trimmed) >= len(*content) { + trimmed = historyContentTruncationMarker + } + *content = trimmed + } +} + +// truncateCurrentToolResults reduces the largest structured tool outputs until +// the serialized request fits. Re-measuring the complete JSON after each cut +// accounts for escaping overhead and keeps this guard correct for arbitrary +// tool output. IDs and result ordering are never changed. +func truncateCurrentToolResults(payload *KiroPayload) { + if payload == nil { + return + } + ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext + if ctx == nil || len(ctx.ToolResults) == 0 { + return + } + + for { + currentSize := payloadByteSize(payload) + if currentSize <= maxPayloadBytes { + return + } + + resultIndex, contentIndex, largestBytes := -1, -1, 0 + for i := range ctx.ToolResults { + for j := range ctx.ToolResults[i].Content { + if size := len(ctx.ToolResults[i].Content[j].Text); size > largestBytes { + resultIndex, contentIndex, largestBytes = i, j, size + } + } + } + if resultIndex < 0 || largestBytes <= len(toolResultTruncationMarker) { + return + } + + excess := currentSize - maxPayloadBytes + targetBytes := largestBytes - excess - 256 + if targetBytes < len(toolResultTruncationMarker) { + targetBytes = len(toolResultTruncationMarker) + } + if targetBytes >= largestBytes { + targetBytes = largestBytes - 1 + } + + content := &ctx.ToolResults[resultIndex].Content[contentIndex] + trimmed := truncateToolResultText(content.Text, targetBytes) + if len(trimmed) >= len(content.Text) { + trimmed = toolResultTruncationMarker + } + content.Text = trimmed + } +} + +// truncateToolResultText preserves both ends of tool output because command +// failures and summaries often appear at the tail. +func truncateToolResultText(text string, maxBytes int) string { + return truncateTextAroundMarker(text, maxBytes, toolResultTruncationMarker) +} + +// truncateTextAroundMarker preserves both ends of text and always cuts on +// UTF-8 rune boundaries. +func truncateTextAroundMarker(text string, maxBytes int, marker string) string { + if len(text) <= maxBytes { + return text + } + if maxBytes <= len(marker) { + end := min(maxBytes, len(marker)) + for end > 0 && !utf8.ValidString(marker[:end]) { + end-- + } + return marker[:end] + } + + remaining := maxBytes - len(marker) + headBytes := remaining * 2 / 3 + for headBytes > 0 && !utf8.ValidString(text[:headBytes]) { + headBytes-- + } + tailBytes := remaining - headBytes + tailStart := len(text) - tailBytes + for tailStart < len(text) && !utf8.RuneStart(text[tailStart]) { + tailStart++ + } + return text[:headBytes] + marker + text[tailStart:] +} + // truncateCurrentMessage hard-truncates the current message content as a last // resort when even the minimal retained history plus current message exceeds the // limit.