From eaf121896546bde5d759472ba9ea69f227ca4269 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:23 +0800
Subject: [PATCH 1/6] fix(protocol): align default Kiro IDE fingerprint to
1.0.212
Default User-Agent / x-amz-user-agent still advertised 0.11.107. Bump the
fallback to the current desktop build and keep the per-endpoint header
builders explicit so OAuth and API Key paths stay distinguishable.
---
config/config.go | 6 +++-
proxy/kiro_headers.go | 38 +++++++++++++++++++---
proxy/kiro_headers_test.go | 66 +++++++++++++++++++++++++++++++-------
3 files changed, 92 insertions(+), 18 deletions(-)
diff --git a/config/config.go b/config/config.go
index be31fba1..1fbda693 100644
--- a/config/config.go
+++ b/config/config.go
@@ -270,7 +270,9 @@ var (
// Init initializes the configuration system with the specified file path.
// If the file doesn't exist, a default configuration is created.
func Init(path string) error {
+ cfgLock.Lock()
cfgPath = path
+ cfgLock.Unlock()
return Load()
}
@@ -1206,7 +1208,9 @@ func GetKiroClientConfig() KiroClientConfig {
cfgLock.RLock()
defer cfgLock.RUnlock()
- kiroVersion := "0.11.107"
+ // OAuth/SSO requests retain the author's IDE compatibility fingerprint;
+ // advertise the current official desktop version by default.
+ kiroVersion := "1.0.212"
if cfg != nil && cfg.KiroVersion != "" {
kiroVersion = cfg.KiroVersion
}
diff --git a/proxy/kiro_headers.go b/proxy/kiro_headers.go
index e704b3eb..13ab0945 100644
--- a/proxy/kiro_headers.go
+++ b/proxy/kiro_headers.go
@@ -8,8 +8,10 @@ import (
)
const (
- kiroStreamingSDKVersion = "1.0.34"
- kiroRuntimeSDKVersion = "1.0.0"
+ kiroLegacyStreamingSDKVersion = "1.0.39"
+ kiroRuntimeSDKVersion = "1.0.0"
+ kiroCLIVersion = "2.14.2"
+ kiroCLIServiceClientVersion = "0.1.17975"
)
type kiroHeaderValues struct {
@@ -18,14 +20,36 @@ type kiroHeaderValues struct {
Host string
}
-func buildStreamingHeaderValues(account *config.Account, host string) kiroHeaderValues {
- return buildKiroHeaderValues(account, host, "codewhispererstreaming", kiroStreamingSDKVersion, "m/E")
+func buildKiroRuntimeHeaderValues(account *config.Account, host string) kiroHeaderValues {
+ return buildKiroHeaderValues(account, host, "kiroruntime", kiroRuntimeSDKVersion, "m/E")
+}
+
+func buildLegacyStreamingHeaderValues(account *config.Account, host string) kiroHeaderValues {
+ return buildKiroHeaderValues(account, host, "codewhispererstreaming", kiroLegacyStreamingSDKVersion, "m/E")
}
func buildRuntimeHeaderValues(account *config.Account, host string) kiroHeaderValues {
+ if config.IsAPIKeyAccount(account) {
+ return buildKiroCLIHeaderValues(host, "codewhispererruntime")
+ }
return buildKiroHeaderValues(account, host, "codewhispererruntime", kiroRuntimeSDKVersion, "m/N,E")
}
+func buildKiroCLIHeaderValues(host, apiName string) kiroHeaderValues {
+ userAgent := fmt.Sprintf(
+ "KiroCLI/%s md/appVersion-%s app/AmazonQ-For-CLI api/%s#%s",
+ kiroCLIVersion,
+ kiroCLIVersion,
+ apiName,
+ kiroCLIServiceClientVersion,
+ )
+ return kiroHeaderValues{
+ UserAgent: userAgent,
+ AmzUserAgent: userAgent,
+ Host: host,
+ }
+}
+
func buildKiroHeaderValues(account *config.Account, host, apiName, sdkVersion, mode string) kiroHeaderValues {
clientCfg := config.GetKiroClientConfig()
machineID := ""
@@ -76,7 +100,11 @@ func applyKiroBaseHeaders(req *http.Request, account *config.Account, values kir
}
req.Header.Set("User-Agent", values.UserAgent)
req.Header.Set("x-amz-user-agent", values.AmzUserAgent)
- req.Header.Set("x-amzn-codewhisperer-optout", "true")
+ if account != nil && config.IsAPIKeyAccount(account) {
+ req.Header.Set("x-amzn-codewhisperer-optout", "false")
+ } else {
+ req.Header.Set("x-amzn-codewhisperer-optout", "true")
+ }
if values.Host != "" {
req.Host = values.Host
}
diff --git a/proxy/kiro_headers_test.go b/proxy/kiro_headers_test.go
index 74590265..ca1eff63 100644
--- a/proxy/kiro_headers_test.go
+++ b/proxy/kiro_headers_test.go
@@ -7,24 +7,63 @@ import (
"testing"
)
-func TestBuildStreamingHeaderValuesAlignsWithKiroIDEFormat(t *testing.T) {
+func TestBuildStreamingHeaderValuesAlignsWithKiroRuntimeIDE(t *testing.T) {
account := &config.Account{MachineId: "machine-123"}
- values := buildStreamingHeaderValues(account, "q.us-east-1.amazonaws.com")
+ values := buildKiroRuntimeHeaderValues(account, "runtime.us-east-1.kiro.dev")
- if values.Host != "q.us-east-1.amazonaws.com" {
+ if values.Host != "runtime.us-east-1.kiro.dev" {
t.Fatalf("expected host to be preserved, got %q", values.Host)
}
- if !strings.Contains(values.UserAgent, "aws-sdk-js/1.0.34") {
- t.Fatalf("expected streaming sdk version in user agent, got %q", values.UserAgent)
+ for _, want := range []string{
+ "aws-sdk-js/1.0.0",
+ "api/kiroruntime#1.0.0",
+ "KiroIDE-1.0.212-machine-123",
+ } {
+ if !strings.Contains(values.UserAgent, want) {
+ t.Fatalf("official KiroRuntime user agent %q missing %q", values.UserAgent, want)
+ }
}
- if !strings.Contains(values.UserAgent, "api/codewhispererstreaming#1.0.34") {
- t.Fatalf("expected streaming API marker in user agent, got %q", values.UserAgent)
+ if !strings.Contains(values.AmzUserAgent, "aws-sdk-js/1.0.0 KiroIDE-1.0.212-machine-123") {
+ t.Fatalf("x-amz-user-agent does not match KiroRuntime IDE format: %q", values.AmzUserAgent)
}
- if !strings.Contains(values.UserAgent, "KiroIDE-0.11.107-machine-123") {
- t.Fatalf("expected kiro version and machine id in user agent, got %q", values.UserAgent)
+}
+
+func TestBuildLegacyStreamingHeaderValuesRetainsEndpointSDK(t *testing.T) {
+ account := &config.Account{MachineId: "machine-123"}
+ values := buildLegacyStreamingHeaderValues(account, "q.us-east-1.amazonaws.com")
+
+ for _, want := range []string{
+ "aws-sdk-js/1.0.39",
+ "api/codewhispererstreaming#1.0.39",
+ "KiroIDE-1.0.212-machine-123",
+ } {
+ if !strings.Contains(values.UserAgent, want) {
+ t.Fatalf("legacy fallback user agent %q missing %q", values.UserAgent, want)
+ }
+ }
+}
+
+func TestResolveKiroCLIEndpointUsesCLIIdentity(t *testing.T) {
+ account := &config.Account{AuthMethod: "api_key", KiroApiKey: "ksk_test", Region: "eu-central-1"}
+ resolved, err := resolveKiroEndpoint(kiroCLIEndpoint, account, "")
+ if err != nil {
+ t.Fatal(err)
}
- if !strings.Contains(values.AmzUserAgent, "aws-sdk-js/1.0.34 KiroIDE-0.11.107-machine-123") {
- t.Fatalf("expected x-amz-user-agent to include version and machine id, got %q", values.AmzUserAgent)
+ if resolved.URL != "https://runtime.eu-central-1.kiro.dev/" || resolved.ContentType != "application/x-amz-json-1.0" || resolved.AgentMode {
+ t.Fatalf("resolved CLI endpoint = %+v", resolved)
+ }
+ for _, want := range []string{
+ "KiroCLI/2.14.2",
+ "md/appVersion-2.14.2",
+ "app/AmazonQ-For-CLI",
+ "api/codewhispererstreaming#0.1.17975",
+ } {
+ if !strings.Contains(resolved.HeaderValues.UserAgent, want) {
+ t.Fatalf("CLI user agent %q missing %q", resolved.HeaderValues.UserAgent, want)
+ }
+ }
+ if strings.Contains(resolved.HeaderValues.UserAgent, "KiroIDE") || strings.Contains(resolved.HeaderValues.UserAgent, "aws-sdk-js") {
+ t.Fatalf("API key path must not mix IDE/JS identity into CLI UA: %q", resolved.HeaderValues.UserAgent)
}
}
@@ -89,7 +128,7 @@ func TestApplyKiroBaseHeadersMarksAPIKeyCredentials(t *testing.T) {
AuthMethod: "api_key",
}
- applyKiroBaseHeaders(req, account, buildStreamingHeaderValues(account, req.URL.Host))
+ applyKiroBaseHeaders(req, account, buildKiroCLIHeaderValues(req.URL.Host, "codewhispererstreaming"))
if got := req.Header.Get("Authorization"); got != "Bearer ksk_test_key" {
t.Fatalf("expected API key bearer, got %q", got)
@@ -101,4 +140,7 @@ func TestApplyKiroBaseHeadersMarksAPIKeyCredentials(t *testing.T) {
if got := req.Header.Get("TokenType"); got != "API_KEY" {
t.Fatalf("expected TokenType/tokentype API_KEY, got %q", got)
}
+ if got := req.Header.Get("x-amzn-codewhisperer-optout"); got != "false" {
+ t.Fatalf("expected CLI optout=false, got %q", got)
+ }
}
From ae6d9ea9126747bb69e851873cfded26a1502c35 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:24 +0800
Subject: [PATCH 2/6] fix(stream): CRC-check event frames and buffer parallel
tools
Parse AWS event-stream messages with prelude/message CRC validation, keep
official optional events, and accumulate parallel tool_use frames by id.
Flush only after the full tool set validates so a malformed sibling cannot
leak a partial tool call to the client.
---
proxy/eventstream.go | 183 +++++++++++++
proxy/kiro.go | 512 ++++++++++++++++++++++-------------
proxy/kiro_test.go | 197 ++++++++++++--
proxy/stream_drop_test.go | 271 ++++++++++++++++++
proxy/tool_use_order_test.go | 33 +++
5 files changed, 976 insertions(+), 220 deletions(-)
create mode 100644 proxy/eventstream.go
create mode 100644 proxy/stream_drop_test.go
create mode 100644 proxy/tool_use_order_test.go
diff --git a/proxy/eventstream.go b/proxy/eventstream.go
new file mode 100644
index 00000000..1459c9da
--- /dev/null
+++ b/proxy/eventstream.go
@@ -0,0 +1,183 @@
+package proxy
+
+import (
+ "encoding/binary"
+ "fmt"
+ "hash/crc32"
+ "io"
+ "strings"
+)
+
+const (
+ eventStreamPreludeLength = 12
+ eventStreamMessageCRCLength = 4
+ eventStreamMaxHeadersLength = 128 * 1024
+ eventStreamMaxPayloadLength = 16 * 1024 * 1024
+)
+
+type eventStreamMessage struct {
+ headers map[string]string
+ payload []byte
+}
+
+type upstreamEventStreamError struct {
+ messageType string
+ code string
+ message string
+ reason string
+ retryAfterMilliseconds string
+}
+
+func (e *upstreamEventStreamError) Error() string {
+ parts := []string{"upstream event stream " + e.messageType}
+ if e.code != "" {
+ parts = append(parts, e.code)
+ }
+ if e.message != "" {
+ parts = append(parts, e.message)
+ }
+ if e.reason != "" {
+ parts = append(parts, "reason="+e.reason)
+ }
+ if e.retryAfterMilliseconds != "" {
+ parts = append(parts, "retryAfterMilliseconds="+e.retryAfterMilliseconds)
+ }
+ return strings.Join(parts, ": ")
+}
+
+func newUpstreamEventStreamError(messageType, code, headerMessage string, payload map[string]interface{}) error {
+ message := strings.TrimSpace(headerMessage)
+ if message == "" {
+ message = firstStringField(payload, "message", "errorMessage", "error")
+ }
+
+ retryAfter := ""
+ if value, ok := payload["retryAfterMilliseconds"]; ok && value != nil {
+ retryAfter = fmt.Sprint(value)
+ }
+
+ return &upstreamEventStreamError{
+ messageType: messageType,
+ code: strings.TrimSpace(code),
+ message: message,
+ reason: firstStringField(payload, "reason"),
+ retryAfterMilliseconds: retryAfter,
+ }
+}
+
+// readEventStreamMessage validates and decodes one AWS EventStream envelope.
+// A nil message and nil error means a clean EOF between messages.
+func readEventStreamMessage(body io.Reader) (*eventStreamMessage, error) {
+ prelude := make([]byte, eventStreamPreludeLength)
+ if _, err := io.ReadFull(body, prelude); err != nil {
+ if err == io.EOF {
+ return nil, nil
+ }
+ return nil, fmt.Errorf("read event stream prelude: %w", err)
+ }
+
+ wantPreludeCRC := binary.BigEndian.Uint32(prelude[8:12])
+ gotPreludeCRC := crc32.ChecksumIEEE(prelude[:8])
+ if gotPreludeCRC != wantPreludeCRC {
+ return nil, fmt.Errorf("event stream prelude CRC mismatch: got %08x, want %08x", gotPreludeCRC, wantPreludeCRC)
+ }
+
+ totalLength := binary.BigEndian.Uint32(prelude[0:4])
+ headersLength := binary.BigEndian.Uint32(prelude[4:8])
+ if totalLength < eventStreamPreludeLength+eventStreamMessageCRCLength {
+ return nil, fmt.Errorf("invalid event stream total length %d", totalLength)
+ }
+ if headersLength > eventStreamMaxHeadersLength {
+ return nil, fmt.Errorf("event stream headers length %d exceeds limit %d", headersLength, eventStreamMaxHeadersLength)
+ }
+ if uint64(headersLength)+eventStreamPreludeLength+eventStreamMessageCRCLength > uint64(totalLength) {
+ return nil, fmt.Errorf("invalid event stream headers length %d for total length %d", headersLength, totalLength)
+ }
+ payloadLength := totalLength - headersLength - eventStreamPreludeLength - eventStreamMessageCRCLength
+ if payloadLength > eventStreamMaxPayloadLength {
+ return nil, fmt.Errorf("event stream payload length %d exceeds limit %d", payloadLength, eventStreamMaxPayloadLength)
+ }
+
+ remaining := make([]byte, int(totalLength)-eventStreamPreludeLength)
+ if _, err := io.ReadFull(body, remaining); err != nil {
+ return nil, fmt.Errorf("read event stream message: %w", err)
+ }
+
+ messageCRCOffset := len(remaining) - eventStreamMessageCRCLength
+ wantMessageCRC := binary.BigEndian.Uint32(remaining[messageCRCOffset:])
+ checksum := crc32.NewIEEE()
+ _, _ = checksum.Write(prelude)
+ _, _ = checksum.Write(remaining[:messageCRCOffset])
+ gotMessageCRC := checksum.Sum32()
+ if gotMessageCRC != wantMessageCRC {
+ return nil, fmt.Errorf("event stream message CRC mismatch: got %08x, want %08x", gotMessageCRC, wantMessageCRC)
+ }
+
+ headersEnd := int(headersLength)
+ headers, err := parseEventStreamHeaders(remaining[:headersEnd])
+ if err != nil {
+ return nil, err
+ }
+ payload := remaining[headersEnd:messageCRCOffset]
+ return &eventStreamMessage{headers: headers, payload: payload}, nil
+}
+
+func parseEventStreamHeaders(encoded []byte) (map[string]string, error) {
+ headers := make(map[string]string)
+ for offset := 0; offset < len(encoded); {
+ nameLength := int(encoded[offset])
+ offset++
+ if nameLength == 0 || offset+nameLength+1 > len(encoded) {
+ return nil, fmt.Errorf("malformed event stream header name at offset %d", offset-1)
+ }
+ name := string(encoded[offset : offset+nameLength])
+ offset += nameLength
+ valueType := encoded[offset]
+ offset++
+
+ switch valueType {
+ case 0, 1: // boolean true / false: no value bytes
+ case 2:
+ if offset+1 > len(encoded) {
+ return nil, fmt.Errorf("malformed byte header %q", name)
+ }
+ offset++
+ case 3:
+ if offset+2 > len(encoded) {
+ return nil, fmt.Errorf("malformed int16 header %q", name)
+ }
+ offset += 2
+ case 4:
+ if offset+4 > len(encoded) {
+ return nil, fmt.Errorf("malformed int32 header %q", name)
+ }
+ offset += 4
+ case 5, 8:
+ if offset+8 > len(encoded) {
+ return nil, fmt.Errorf("malformed int64 header %q", name)
+ }
+ offset += 8
+ case 6, 7:
+ if offset+2 > len(encoded) {
+ return nil, fmt.Errorf("malformed variable-length header %q", name)
+ }
+ valueLength := int(binary.BigEndian.Uint16(encoded[offset : offset+2]))
+ offset += 2
+ if offset+valueLength > len(encoded) {
+ return nil, fmt.Errorf("malformed variable-length header %q", name)
+ }
+ if valueType == 7 {
+ headers[name] = string(encoded[offset : offset+valueLength])
+ }
+ offset += valueLength
+ case 9:
+ if offset+16 > len(encoded) {
+ return nil, fmt.Errorf("malformed UUID header %q", name)
+ }
+ offset += 16
+ default:
+ return nil, fmt.Errorf("unsupported event stream header type %d for %q", valueType, name)
+ }
+ }
+ return headers, nil
+}
diff --git a/proxy/kiro.go b/proxy/kiro.go
index 1786ba0f..755758c3 100644
--- a/proxy/kiro.go
+++ b/proxy/kiro.go
@@ -22,19 +22,35 @@ import (
)
// Endpoint configuration (auto-fallback on quota exhaustion).
+type kiroEndpointProtocol uint8
+
+const (
+ kiroProtocolLegacyStreaming kiroEndpointProtocol = iota
+ kiroProtocolRuntime
+ kiroProtocolCLI
+)
+
type kiroEndpoint struct {
URL string
Origin string
AmzTarget string
Name string
+ Protocol kiroEndpointProtocol
+}
+
+type resolvedKiroEndpoint struct {
+ URL string
+ HeaderValues kiroHeaderValues
+ ContentType string
+ AgentMode bool
}
var kiroEndpoints = []kiroEndpoint{
{
- URL: "https://q.us-east-1.amazonaws.com/generateAssistantResponse",
- Origin: "AI_EDITOR",
- AmzTarget: "",
- Name: "Kiro IDE",
+ URL: "https://runtime.us-east-1.kiro.dev/generateAssistantResponse",
+ Origin: "AI_EDITOR",
+ Name: "Kiro Runtime",
+ Protocol: kiroProtocolRuntime,
},
{
URL: "https://codewhisperer.us-east-1.amazonaws.com/generateAssistantResponse",
@@ -57,6 +73,7 @@ var kiroCLIEndpoint = kiroEndpoint{
Origin: "KIRO_CLI",
AmzTarget: "AmazonCodeWhispererStreamingService.GenerateAssistantResponse",
Name: "Kiro CLI",
+ Protocol: kiroProtocolCLI,
}
// Global HTTP clients, swappable at runtime to apply proxy reconfiguration without restart.
@@ -179,6 +196,7 @@ type KiroUserInputMessage struct {
ModelID string `json:"modelId,omitempty"`
Origin string `json:"origin"`
Images []KiroImage `json:"images,omitempty"`
+ Documents []KiroDocument `json:"documents,omitempty"`
UserInputMessageContext *UserInputMessageContext `json:"userInputMessageContext,omitempty"`
}
@@ -216,14 +234,34 @@ type KiroImage struct {
} `json:"source"`
}
+// KiroDocument is a non-image attachment (PDF, office, text, ...).
+type KiroDocument struct {
+ Name string `json:"name"`
+ Format string `json:"format"`
+ Source struct {
+ Bytes string `json:"bytes"`
+ } `json:"source"`
+}
+
+// KiroReasoningContent is the signed thinking stamp upstream expects on later
+// turns. Either open text+signature or redacted content.
+type KiroReasoningContent struct {
+ ReasoningText *struct {
+ Text string `json:"text"`
+ Signature string `json:"signature"`
+ } `json:"reasoningText,omitempty"`
+ RedactedContent []byte `json:"redactedContent,omitempty"`
+}
+
type KiroHistoryMessage struct {
UserInputMessage *KiroUserInputMessage `json:"userInputMessage,omitempty"`
AssistantResponseMessage *KiroAssistantResponseMessage `json:"assistantResponseMessage,omitempty"`
}
type KiroAssistantResponseMessage struct {
- Content string `json:"content"`
- ToolUses []KiroToolUse `json:"toolUses,omitempty"`
+ Content string `json:"content"`
+ ToolUses []KiroToolUse `json:"toolUses,omitempty"`
+ ReasoningContent *KiroReasoningContent `json:"reasoningContent,omitempty"`
}
type KiroToolUse struct {
@@ -248,6 +286,13 @@ type KiroStreamCallback struct {
OnError func(err error)
OnCredits func(credits float64)
OnContextUsage func(percentage float64)
+ // OnStopReason is fired when a metadataEvent carries stopReason.
+ // Absence after content is how callers detect a truncated stream
+ // (Kiro IDE retries those; see _streamResponseChunks).
+ OnStopReason func(reason string)
+ // OnReasoningMeta carries only the validation stamp for a thinking
+ // segment. Visible text still arrives via OnText(..., true).
+ OnReasoningMeta func(signature, redactedBase64 string)
}
// ==================== API Call ====================
@@ -281,15 +326,33 @@ func endpointsForAccount(account *config.Account) []kiroEndpoint {
return getSortedEndpoints(config.GetPreferredEndpoint())
}
-// cliRuntimeURL builds the regional Kiro CLI runtime URL.
-func cliRuntimeURL(account *config.Account) string {
- region := "us-east-1"
- if account != nil {
- if r := strings.TrimSpace(account.Region); r != "" {
- region = r
- }
+func resolveKiroEndpoint(endpoint kiroEndpoint, account *config.Account, profileArn string) (resolvedKiroEndpoint, error) {
+ region := kiroRegionForProfile(account, profileArn)
+ resolvedURL := endpoint.URL
+ if endpoint.Protocol == kiroProtocolRuntime || endpoint.Protocol == kiroProtocolCLI {
+ resolvedURL = strings.Replace(endpoint.URL, "runtime.us-east-1.kiro.dev", "runtime."+region+".kiro.dev", 1)
+ } else if endpoint.Protocol == kiroProtocolLegacyStreaming {
+ resolvedURL = regionalizeURLForProfile(endpoint.URL, account, profileArn)
+ } else {
+ return resolvedKiroEndpoint{}, fmt.Errorf("unsupported Kiro endpoint protocol %d", endpoint.Protocol)
}
- return fmt.Sprintf("https://runtime.%s.kiro.dev/", region)
+ parsed, err := url.Parse(resolvedURL)
+ if err != nil || parsed.Host == "" {
+ return resolvedKiroEndpoint{}, fmt.Errorf("invalid Kiro endpoint URL %q", resolvedURL)
+ }
+
+ resolved := resolvedKiroEndpoint{URL: resolvedURL, ContentType: "application/json", AgentMode: true}
+ switch endpoint.Protocol {
+ case kiroProtocolRuntime:
+ resolved.HeaderValues = buildKiroRuntimeHeaderValues(account, parsed.Host)
+ case kiroProtocolCLI:
+ resolved.HeaderValues = buildKiroCLIHeaderValues(parsed.Host, "codewhispererstreaming")
+ resolved.ContentType = "application/x-amz-json-1.0"
+ resolved.AgentMode = false
+ case kiroProtocolLegacyStreaming:
+ resolved.HeaderValues = buildLegacyStreamingHeaderValues(account, parsed.Host)
+ }
+ return resolved, nil
}
// getSortedEndpoints returns endpoints ordered by user preference, with optional fallback.
@@ -370,52 +433,33 @@ func CallKiroAPI(account *config.Account, payload *KiroPayload, callback *KiroSt
// Build endpoint list ordered by configuration / credential type.
endpoints := endpointsForAccount(account)
- isAPIKey := config.IsAPIKeyAccount(account)
var lastErr error
for _, ep := range endpoints {
// Update the origin field for the selected endpoint.
payload.ConversationState.CurrentMessage.UserInputMessage.Origin = ep.Origin
- // Target the profile's data-plane region; endpoint URLs are declared for us-east-1.
- // API Key accounts use the CLI runtime host instead of IDE/Q hosts.
- epURL := regionalizeURLForProfile(ep.URL, account, payload.ProfileArn)
- if isAPIKey {
- epURL = cliRuntimeURL(account)
+ resolved, err := resolveKiroEndpoint(ep, account, payload.ProfileArn)
+ if err != nil {
+ lastErr = err
+ continue
}
reqBody, _ := json.Marshal(payload)
- req, err := http.NewRequest("POST", epURL, bytes.NewReader(reqBody))
+ req, err := http.NewRequest("POST", resolved.URL, bytes.NewReader(reqBody))
if err != nil {
lastErr = err
continue
}
-
- host := ""
- if parsedURL, parseErr := url.Parse(epURL); parseErr == nil {
- host = parsedURL.Host
- }
- headerValues := buildStreamingHeaderValues(account, host)
-
- if isAPIKey {
- req.Header.Set("Content-Type", "application/x-amz-json-1.0")
- } else {
- req.Header.Set("Content-Type", "application/json")
- }
+ req.Header.Set("Content-Type", resolved.ContentType)
req.Header.Set("Accept", "*/*")
if ep.AmzTarget != "" {
req.Header.Set("X-Amz-Target", ep.AmzTarget)
}
- applyKiroBaseHeaders(req, account, headerValues)
- if !isAPIKey {
+ applyKiroBaseHeaders(req, account, resolved.HeaderValues)
+ if resolved.AgentMode {
req.Header.Set("x-amzn-kiro-agent-mode", "vibe")
}
- // CLI captures use optout=false; IDE path keeps true.
- if isAPIKey {
- req.Header.Set("x-amzn-codewhisperer-optout", "false")
- } else {
- req.Header.Set("x-amzn-codewhisperer-optout", "true")
- }
req.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
req.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())
@@ -471,105 +515,116 @@ func parseEventStream(body io.Reader, callback *KiroStreamCallback) error {
callback = &KiroStreamCallback{}
}
- // Read directly without bufio to avoid buffering latency in streaming responses.
var inputTokens, outputTokens int
var totalCredits float64
- var currentToolUse *toolUseState
+ pending := &pendingToolUses{}
for {
- // Prelude: 12 bytes (total_len + headers_len + crc)
- prelude := make([]byte, 12)
- _, err := io.ReadFull(body, prelude)
- if err == io.EOF {
- break
- }
+ message, err := readEventStreamMessage(body)
if err != nil {
return err
}
-
- totalLength := int(prelude[0])<<24 | int(prelude[1])<<16 | int(prelude[2])<<8 | int(prelude[3])
- headersLength := int(prelude[4])<<24 | int(prelude[5])<<16 | int(prelude[6])<<8 | int(prelude[7])
-
- if totalLength < 16 {
- continue
+ if message == nil {
+ break
}
- // Read the remaining message bytes.
- remaining := totalLength - 12
- msgBuf := make([]byte, remaining)
- _, err = io.ReadFull(body, msgBuf)
- if err != nil {
- return err
+ messageType := strings.TrimSpace(message.headers[":message-type"])
+ if messageType == "" {
+ return fmt.Errorf("event stream message missing :message-type header")
}
- if headersLength > len(msgBuf)-4 {
- continue
+ if messageType == "error" || messageType == "exception" {
+ var event map[string]interface{}
+ headerMessage := message.headers[":error-message"]
+ if len(message.payload) > 0 && string(message.payload) != "null" {
+ if err := json.Unmarshal(message.payload, &event); err != nil && headerMessage == "" {
+ headerMessage = strings.TrimSpace(string(message.payload))
+ }
+ }
+ if event == nil {
+ event = map[string]interface{}{}
+ }
+ code := message.headers[":error-code"]
+ if messageType == "exception" {
+ code = message.headers[":exception-type"]
+ }
+ return newUpstreamEventStreamError(messageType, code, headerMessage, event)
+ }
+ if messageType != "event" {
+ return fmt.Errorf("unsupported event stream message type %q", messageType)
}
- eventType := extractEventType(msgBuf[0:headersLength])
- payloadBytes := msgBuf[headersLength : len(msgBuf)-4]
- if len(payloadBytes) == 0 {
- continue
+ eventType := strings.TrimSpace(message.headers[":event-type"])
+ if eventType == "" {
+ return fmt.Errorf("event stream event missing :event-type header")
}
var event map[string]interface{}
- if err := json.Unmarshal(payloadBytes, &event); err != nil {
- continue
+ if len(message.payload) > 0 {
+ if err := json.Unmarshal(message.payload, &event); err != nil {
+ return fmt.Errorf("decode event stream %s payload: %w", eventType, err)
+ }
+ }
+ if event == nil {
+ event = map[string]interface{}{}
}
inputTokens, outputTokens = updateTokensFromEvent(event, inputTokens, outputTokens)
- // Dispatch by event type.
switch eventType {
- // Both text streams are passed through verbatim. Kiro sends
- // assistantResponseEvent and reasoningContentEvent as pure incremental deltas
- // (verified against real upstream traffic), never as cumulative snapshots, and
- // never replays a chunk: ordering and at-most-once delivery are already
- // guaranteed by TCP, and a dropped stream is retried as a whole new request
- // rather than resumed.
- //
- // Do NOT reintroduce content-based de-duplication here. At the string level a
- // replayed chunk is indistinguishable from text that simply repeats itself, and
- // the wire protocol carries no sequence number or message id to tell them apart
- // (the AWS event-stream base spec defines none), so such a heuristic can only
- // guess -- and when it guesses wrong it silently eats real output. The previous
- // implementation turned "6666666666" into "666", "abababab" into "abab" and
- // "1833" into "183", on both streams.
case "assistantResponseEvent":
- if content, ok := event["content"].(string); ok && content != "" {
- if callback.OnText != nil {
- callback.OnText(content, false)
- }
+ if content, ok := event["content"].(string); ok && content != "" && callback.OnText != nil {
+ callback.OnText(content, false)
}
case "reasoningContentEvent":
- if text, ok := event["text"].(string); ok && text != "" {
- if callback.OnText != nil {
- callback.OnText(text, true)
+ if text, ok := event["text"].(string); ok && text != "" && callback.OnText != nil {
+ callback.OnText(text, true)
+ }
+ if callback.OnReasoningMeta != nil {
+ sig, _ := event["signature"].(string)
+ redacted, _ := event["redactedContent"].(string)
+ if sig != "" || redacted != "" {
+ callback.OnReasoningMeta(sig, redacted)
}
}
case "toolUseEvent":
- currentToolUse = handleToolUseEvent(event, currentToolUse, callback)
+ if err := handleToolUseEvent(event, pending); err != nil {
+ return err
+ }
case "meteringEvent":
if usage, ok := event["usage"].(float64); ok {
totalCredits += usage
}
case "contextUsageEvent":
- if pct, ok := event["contextUsagePercentage"].(float64); ok {
- if callback.OnContextUsage != nil {
- callback.OnContextUsage(pct)
- }
+ if pct, ok := event["contextUsagePercentage"].(float64); ok && callback.OnContextUsage != nil {
+ callback.OnContextUsage(pct)
}
- }
- }
-
- if currentToolUse != nil {
- finishToolUse(currentToolUse, callback)
+ case "metadataEvent":
+ if reason := firstStringField(event, "stopReason", "stop_reason"); reason != "" && callback.OnStopReason != nil {
+ callback.OnStopReason(reason)
+ }
+ case "error", "throttlingError", "validationError", "serviceUnavailableError", "invalidStateEvent":
+ return newUpstreamEventStreamError("event", eventType, "", event)
+ case "codeReferenceEvent", "documentCitationEvent", "toolResultEvent",
+ "supplementaryWebLinksEvent", "messageMetadataEvent", "interactionComponentsEvent",
+ "intentsEvent", "followupPromptEvent", "citationEvent", "codeEvent", "dryRunSucceedEvent":
+ // These official metadata/presentation events have no portable
+ // representation in the OpenAI/Anthropic compatibility APIs.
+ logger.Debugf("[EventStream] ignoring known optional event %s", eventType)
+ default:
+ return fmt.Errorf("unsupported upstream event stream event %q", eventType)
+ }
+ }
+
+ // Flush tools that omitted an explicit stop frame. Any invalid argument
+ // buffer makes the entire parallel tool turn fail.
+ if err := pending.flushAll(callback); err != nil {
+ return err
}
if callback.OnCredits != nil && totalCredits > 0 {
callback.OnCredits(totalCredits)
}
-
if callback.OnComplete != nil {
callback.OnComplete(inputTokens, outputTokens)
}
@@ -701,7 +756,6 @@ func collectUsageMaps(v interface{}, out *[]map[string]interface{}) {
}
}
-
func readTokenNumber(m map[string]interface{}, keys ...string) (int, bool) {
for _, k := range keys {
v, ok := m[k]
@@ -737,70 +791,198 @@ type toolUseState struct {
ToolUseID string
Name string
InputBuffer strings.Builder
+ SawInput bool
GeneratedID bool
}
-func handleToolUseEvent(event map[string]interface{}, current *toolUseState, callback *KiroStreamCallback) *toolUseState {
- toolUseID := firstStringField(event, "toolUseId", "toolUseID", "tool_use_id", "id")
- name := firstStringField(event, "name", "toolName", "tool_name")
- isStop := firstBoolField(event, "stop", "isStop", "done")
+// pendingToolUses tracks in-flight tool calls for one stream. Entries are keyed
+// by toolUseId so interleaved parallel frames accumulate independently, and
+// `order` preserves arrival sequence: tool call order is semantic for clients,
+// so a bare map range (randomised in Go) must never decide emit order.
+type pendingToolUses struct {
+ byID map[string]*toolUseState
+ completed map[string]KiroToolUse
+ order []string
+ lastID string
+}
+
+func (p *pendingToolUses) get(id string) *toolUseState {
+ if p.byID == nil {
+ return nil
+ }
+ return p.byID[id]
+}
- if toolUseID != "" && name != "" {
- if current == nil {
- current = &toolUseState{ToolUseID: toolUseID, Name: name}
- } else if current.ToolUseID != toolUseID {
- if current.GeneratedID && current.Name == name {
- current.ToolUseID = toolUseID
- current.GeneratedID = false
- } else {
- finishToolUse(current, callback)
- current = &toolUseState{ToolUseID: toolUseID, Name: name}
+func (p *pendingToolUses) add(state *toolUseState) {
+ if p.byID == nil {
+ p.byID = map[string]*toolUseState{}
+ }
+ p.byID[state.ToolUseID] = state
+ p.order = append(p.order, state.ToolUseID)
+ p.lastID = state.ToolUseID
+}
+
+func (p *pendingToolUses) isCompleted(id string) bool {
+ _, ok := p.completed[id]
+ return ok
+}
+
+// rekey moves an entry from a generated id to the real id from upstream,
+// keeping its position in the arrival order.
+func (p *pendingToolUses) rekey(state *toolUseState, newID string) {
+ oldID := state.ToolUseID
+ delete(p.byID, oldID)
+ for i, id := range p.order {
+ if id == oldID {
+ p.order[i] = newID
+ break
+ }
+ }
+ state.ToolUseID = newID
+ state.GeneratedID = false
+ p.byID[newID] = state
+ if p.lastID == oldID {
+ p.lastID = newID
+ }
+}
+
+func (p *pendingToolUses) complete(state *toolUseState) error {
+ toolUse, err := decodeToolUse(state)
+ if err != nil {
+ return err
+ }
+ if p.completed == nil {
+ p.completed = make(map[string]KiroToolUse)
+ }
+ p.completed[state.ToolUseID] = toolUse
+ delete(p.byID, state.ToolUseID)
+ if p.lastID == state.ToolUseID {
+ p.lastID = ""
+ for i := len(p.order) - 1; i >= 0; i-- {
+ if p.byID[p.order[i]] != nil {
+ p.lastID = p.order[i]
+ break
}
}
- } else if name != "" && current == nil {
- current = &toolUseState{ToolUseID: "toolu_" + uuid.New().String(), Name: name, GeneratedID: true}
- } else if name != "" && current != nil && current.Name != name {
- finishToolUse(current, callback)
- current = &toolUseState{ToolUseID: "toolu_" + uuid.New().String(), Name: name, GeneratedID: true}
}
+ return nil
+}
- if current != nil {
- if input, ok := event["input"].(string); ok {
- current.InputBuffer.WriteString(input)
- } else if inputObj, ok := event["input"].(map[string]interface{}); ok {
- data, _ := json.Marshal(inputObj)
- current.InputBuffer.Reset()
- current.InputBuffer.Write(data)
+// flushAll validates the complete parallel set before exposing any tool to the
+// caller. A malformed sibling therefore cannot leak an otherwise valid tool.
+func (p *pendingToolUses) flushAll(callback *KiroStreamCallback) error {
+ for _, id := range p.order {
+ if p.isCompleted(id) {
+ continue
+ }
+ if state := p.byID[id]; state != nil {
+ if err := p.complete(state); err != nil {
+ return err
+ }
+ }
+ }
+ if callback != nil && callback.OnToolUse != nil {
+ for _, id := range p.order {
+ if toolUse, ok := p.completed[id]; ok {
+ callback.OnToolUse(toolUse)
+ }
}
}
+ p.byID = nil
+ p.completed = nil
+ p.order = nil
+ p.lastID = ""
+ return nil
+}
- if isStop && current != nil {
- finishToolUse(current, callback)
+// handleToolUseEvent accumulates a toolUseEvent into pending. Frames for
+// different IDs remain independent; completed tools stay buffered until EOF.
+func handleToolUseEvent(event map[string]interface{}, pending *pendingToolUses) error {
+ toolUseID := firstStringField(event, "toolUseId", "toolUseID", "tool_use_id", "id")
+ name := firstStringField(event, "name", "toolName", "tool_name")
+ isStop := firstBoolField(event, "stop", "isStop", "done")
+
+ var state *toolUseState
+ switch {
+ case toolUseID != "":
+ if pending.isCompleted(toolUseID) {
+ return nil
+ }
+ state = pending.get(toolUseID)
+ if state == nil && pending.lastID != "" {
+ if previous := pending.get(pending.lastID); previous != nil && previous.GeneratedID && (name == "" || previous.Name == name) {
+ pending.rekey(previous, toolUseID)
+ state = previous
+ }
+ }
+ if state == nil {
+ if name == "" {
+ return nil
+ }
+ state = &toolUseState{ToolUseID: toolUseID, Name: name}
+ pending.add(state)
+ } else {
+ if name != "" && state.Name == "" {
+ state.Name = name
+ }
+ pending.lastID = state.ToolUseID
+ }
+ case pending.lastID != "" && pending.get(pending.lastID) != nil:
+ state = pending.get(pending.lastID)
+ if name != "" && state.Name != name {
+ if err := pending.complete(state); err != nil {
+ return err
+ }
+ state = &toolUseState{ToolUseID: "toolu_" + uuid.New().String(), Name: name, GeneratedID: true}
+ pending.add(state)
+ }
+ case name != "":
+ state = &toolUseState{ToolUseID: "toolu_" + uuid.New().String(), Name: name, GeneratedID: true}
+ pending.add(state)
+ default:
return nil
}
- return current
+ if input, ok := event["input"].(string); ok {
+ state.SawInput = true
+ state.InputBuffer.WriteString(input)
+ } else if inputObj, ok := event["input"].(map[string]interface{}); ok {
+ state.SawInput = true
+ data, _ := json.Marshal(inputObj)
+ state.InputBuffer.Reset()
+ state.InputBuffer.Write(data)
+ }
+
+ if isStop {
+ return pending.complete(state)
+ }
+ return nil
}
-func finishToolUse(state *toolUseState, callback *KiroStreamCallback) {
- if state == nil || state.Name == "" || callback == nil || callback.OnToolUse == nil {
- return
+func decodeToolUse(state *toolUseState) (KiroToolUse, error) {
+ if state == nil || state.Name == "" {
+ return KiroToolUse{}, fmt.Errorf("%w: missing tool identity", errIncompleteToolUse)
}
if state.ToolUseID == "" {
state.ToolUseID = "toolu_" + uuid.New().String()
}
+ if !state.SawInput {
+ return KiroToolUse{}, fmt.Errorf("%w: %s (%s): missing input", errIncompleteToolUse, state.Name, state.ToolUseID)
+ }
var input map[string]interface{}
if state.InputBuffer.Len() > 0 {
- json.Unmarshal([]byte(state.InputBuffer.String()), &input)
+ if err := json.Unmarshal([]byte(state.InputBuffer.String()), &input); err != nil {
+ return KiroToolUse{}, fmt.Errorf("%w: %s (%s): %v", errIncompleteToolUse, state.Name, state.ToolUseID, err)
+ }
}
if input == nil {
input = make(map[string]interface{})
}
- callback.OnToolUse(KiroToolUse{
+ return KiroToolUse{
ToolUseID: state.ToolUseID,
Name: state.Name,
Input: input,
- })
+ }, nil
}
func firstStringField(m map[string]interface{}, keys ...string) string {
@@ -820,57 +1002,3 @@ func firstBoolField(m map[string]interface{}, keys ...string) bool {
}
return false
}
-
-// extractEventType extracts the event type string from AWS Event Stream message headers.
-func extractEventType(headers []byte) string {
- offset := 0
- for offset < len(headers) {
- if offset >= len(headers) {
- break
- }
- nameLen := int(headers[offset])
- offset++
- if offset+nameLen > len(headers) {
- break
- }
- name := string(headers[offset : offset+nameLen])
- offset += nameLen
- if offset >= len(headers) {
- break
- }
- valueType := headers[offset]
- offset++
-
- if valueType == 7 { // String
- if offset+2 > len(headers) {
- break
- }
- valueLen := int(headers[offset])<<8 | int(headers[offset+1])
- offset += 2
- if offset+valueLen > len(headers) {
- break
- }
- value := string(headers[offset : offset+valueLen])
- offset += valueLen
- if name == ":event-type" {
- return value
- }
- continue
- }
-
- // Skip other value types by their fixed byte widths.
- skipSizes := map[byte]int{0: 0, 1: 0, 2: 1, 3: 2, 4: 4, 5: 8, 8: 8, 9: 16}
- if valueType == 6 {
- if offset+2 > len(headers) {
- break
- }
- l := int(headers[offset])<<8 | int(headers[offset+1])
- offset += 2 + l
- } else if skip, ok := skipSizes[valueType]; ok {
- offset += skip
- } else {
- break
- }
- }
- return ""
-}
diff --git a/proxy/kiro_test.go b/proxy/kiro_test.go
index 14af12f5..0fcaaa37 100644
--- a/proxy/kiro_test.go
+++ b/proxy/kiro_test.go
@@ -4,9 +4,11 @@ import (
"bytes"
"encoding/binary"
"encoding/json"
+ "hash/crc32"
"kiro-go/config"
"net/http"
"net/url"
+ "strings"
"testing"
"time"
)
@@ -139,20 +141,116 @@ func TestParseEventStreamNilCallbackFieldsAreNoOp(t *testing.T) {
}
}
+func TestParseEventStreamRejectsPreludeCRCMismatch(t *testing.T) {
+ frame := awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "hello"})
+ frame[8] ^= 0xff
+
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{})
+ if err == nil || !strings.Contains(err.Error(), "prelude CRC") {
+ t.Fatalf("parse error = %v, want prelude CRC mismatch", err)
+ }
+}
+
+func TestParseEventStreamRejectsMessageCRCMismatchBeforeCallbacks(t *testing.T) {
+ frame := awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "hello"})
+ frame[len(frame)-1] ^= 0xff
+
+ called := false
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{
+ OnText: func(string, bool) { called = true },
+ })
+ if err == nil || !strings.Contains(err.Error(), "message CRC") {
+ t.Fatalf("parse error = %v, want message CRC mismatch", err)
+ }
+ if called {
+ t.Fatal("corrupt frame reached callback")
+ }
+}
+
+func TestParseEventStreamReturnsModeledException(t *testing.T) {
+ frame := awsEventStreamMessage(t, map[string]string{
+ ":message-type": "exception",
+ ":exception-type": "throttlingError",
+ }, map[string]interface{}{
+ "message": "rate limited",
+ "reason": "INSUFFICIENT_MODEL_CAPACITY",
+ "retryAfterMilliseconds": 1250,
+ })
+
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{})
+ if err == nil {
+ t.Fatal("expected modeled stream exception")
+ }
+ for _, want := range []string{"throttlingError", "rate limited", "INSUFFICIENT_MODEL_CAPACITY", "1250"} {
+ if !strings.Contains(err.Error(), want) {
+ t.Fatalf("stream exception %q missing %q", err, want)
+ }
+ }
+}
+
+func TestParseEventStreamReturnsUnmodeledError(t *testing.T) {
+ frame := awsEventStreamMessage(t, map[string]string{
+ ":message-type": "error",
+ ":error-code": "UpstreamUnavailable",
+ ":error-message": "try later",
+ }, nil)
+
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{})
+ if err == nil || !strings.Contains(err.Error(), "UpstreamUnavailable") || !strings.Contains(err.Error(), "try later") {
+ t.Fatalf("parse error = %v, want upstream error details", err)
+ }
+}
+
+func TestParseEventStreamRejectsUnknownEvent(t *testing.T) {
+ frame := awsEventStreamFrame(t, "futureEvent", map[string]interface{}{"value": true})
+
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{})
+ if err == nil || !strings.Contains(err.Error(), "futureEvent") {
+ t.Fatalf("parse error = %v, want unsupported event", err)
+ }
+}
+
+func TestParseEventStreamIgnoresKnownOptionalEvents(t *testing.T) {
+ stream := bytes.NewReader(bytes.Join([][]byte{
+ awsEventStreamFrame(t, "codeReferenceEvent", map[string]interface{}{"references": []interface{}{}}),
+ awsEventStreamFrame(t, "documentCitationEvent", map[string]interface{}{"title": "source"}),
+ awsEventStreamFrame(t, "toolResultEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "supplementaryWebLinksEvent", map[string]interface{}{"links": []interface{}{}}),
+ awsEventStreamFrame(t, "messageMetadataEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "interactionComponentsEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "intentsEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "followupPromptEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "citationEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "codeEvent", map[string]interface{}{}),
+ awsEventStreamFrame(t, "dryRunSucceedEvent", map[string]interface{}{}),
+ }, nil))
+
+ if err := parseEventStream(stream, &KiroStreamCallback{}); err != nil {
+ t.Fatalf("known optional events must remain compatible: %v", err)
+ }
+}
+
func TestHandleToolUseEventGeneratesMissingToolUseID(t *testing.T) {
var toolUses []KiroToolUse
- current := handleToolUseEvent(map[string]interface{}{
+ pending := &pendingToolUses{}
+ if err := handleToolUseEvent(map[string]interface{}{
"name": "mcpIdaProMcpStatus",
"input": `{"server":"ida-pro-mcp"}`,
"stop": true,
- }, nil, &KiroStreamCallback{
- OnToolUse: func(toolUse KiroToolUse) {
- toolUses = append(toolUses, toolUse)
- },
- })
+ }, pending); err != nil {
+ t.Fatalf("handle tool event: %v", err)
+ }
+ if len(toolUses) != 0 {
+ t.Fatalf("tool must remain buffered until full-set validation: %+v", toolUses)
+ }
+ if err := pending.flushAll(&KiroStreamCallback{OnToolUse: func(toolUse KiroToolUse) {
+ toolUses = append(toolUses, toolUse)
+ }}); err != nil {
+ t.Fatalf("flush tools: %v", err)
+ }
- if current != nil {
- t.Fatalf("expected stopped tool use to clear current state")
+ if len(pending.order) != 0 {
+ t.Fatalf("expected stopped tool use to clear pending state, got %d", len(pending.order))
}
if len(toolUses) != 1 {
t.Fatalf("expected one tool use, got %d", len(toolUses))
@@ -172,20 +270,31 @@ func TestHandleToolUseEventReplacesGeneratedIDWhenRealIDArrives(t *testing.T) {
toolUses = append(toolUses, toolUse)
},
}
+ pending := &pendingToolUses{}
- current := handleToolUseEvent(map[string]interface{}{
+ if err := handleToolUseEvent(map[string]interface{}{
"name": "mcpIdaProMcpStatus",
"input": `{"server":`,
- }, nil, callback)
- current = handleToolUseEvent(map[string]interface{}{
+ }, pending); err != nil {
+ t.Fatalf("first tool fragment: %v", err)
+ }
+ if err := handleToolUseEvent(map[string]interface{}{
"toolUseId": "toolu_real",
"name": "mcpIdaProMcpStatus",
"input": `"ida-pro-mcp"}`,
"stop": true,
- }, current, callback)
+ }, pending); err != nil {
+ t.Fatalf("final tool fragment: %v", err)
+ }
+ if len(toolUses) != 0 {
+ t.Fatalf("tool must remain buffered until full-set validation: %+v", toolUses)
+ }
+ if err := pending.flushAll(callback); err != nil {
+ t.Fatalf("flush tools: %v", err)
+ }
- if current != nil {
- t.Fatalf("expected stopped tool use to clear current state")
+ if len(pending.order) != 0 {
+ t.Fatalf("expected stopped tool use to clear pending state, got %d", len(pending.order))
}
if len(toolUses) != 1 {
t.Fatalf("expected one completed tool use, got %d", len(toolUses))
@@ -269,6 +378,23 @@ func TestSetPayloadProfileArnForAccountClearsAPIKeyProfile(t *testing.T) {
}
}
+func TestEndpointsForAccountPrefersKiroRuntimeForOAuth(t *testing.T) {
+ if err := config.Init(t.TempDir() + "/config.json"); err != nil {
+ t.Fatalf("init config: %v", err)
+ }
+ account := &config.Account{AccessToken: "token", Region: "eu-central-1"}
+ eps := endpointsForAccount(account)
+ if len(eps) == 0 {
+ t.Fatal("expected at least one endpoint")
+ }
+ if eps[0].Name != "Kiro Runtime" || eps[0].Origin != "AI_EDITOR" {
+ t.Fatalf("unexpected primary endpoint: %+v", eps[0])
+ }
+ if resolved, err := resolveKiroEndpoint(eps[0], account, ""); err != nil || resolved.URL != "https://runtime.eu-central-1.kiro.dev/generateAssistantResponse" {
+ t.Fatalf("OAuth Kiro endpoint = %+v, err=%v", resolved, err)
+ }
+}
+
func TestEndpointsForAccountUsesCLIForAPIKey(t *testing.T) {
eps := endpointsForAccount(&config.Account{AuthMethod: "api_key", KiroApiKey: "ksk_x"})
if len(eps) != 1 || eps[0].Name != "Kiro CLI" {
@@ -277,8 +403,9 @@ func TestEndpointsForAccountUsesCLIForAPIKey(t *testing.T) {
if eps[0].Origin != "KIRO_CLI" {
t.Fatalf("origin = %q", eps[0].Origin)
}
- if got := cliRuntimeURL(&config.Account{Region: "eu-central-1"}); got != "https://runtime.eu-central-1.kiro.dev/" {
- t.Fatalf("cli url = %q", got)
+ resolved, err := resolveKiroEndpoint(eps[0], &config.Account{Region: "eu-central-1", AuthMethod: "api_key", KiroApiKey: "ksk_x"}, "")
+ if err != nil || resolved.URL != "https://runtime.eu-central-1.kiro.dev/" {
+ t.Fatalf("cli endpoint = %+v, err=%v", resolved, err)
}
}
@@ -301,7 +428,7 @@ func assertProxyURL(t *testing.T, got *url.URL, want string) {
}
}
-func awsEventStreamFrame(t *testing.T, eventType string, payload map[string]interface{}) []byte {
+func awsEventStreamMessage(t *testing.T, fields map[string]string, payload interface{}) []byte {
t.Helper()
payloadBytes, err := json.Marshal(payload)
@@ -309,20 +436,34 @@ func awsEventStreamFrame(t *testing.T, eventType string, payload map[string]inte
t.Fatalf("marshal payload: %v", err)
}
- headerValue := []byte(eventType)
- headers := make([]byte, 0, 1+len(":event-type")+1+2+len(headerValue))
- headers = append(headers, byte(len(":event-type")))
- headers = append(headers, []byte(":event-type")...)
- headers = append(headers, byte(7))
- headers = append(headers, byte(len(headerValue)>>8), byte(len(headerValue)))
- headers = append(headers, headerValue...)
+ var headers []byte
+ for _, name := range []string{":message-type", ":event-type", ":exception-type", ":error-code", ":error-message"} {
+ value, ok := fields[name]
+ if !ok {
+ continue
+ }
+ headers = append(headers, byte(len(name)))
+ headers = append(headers, name...)
+ headers = append(headers, byte(7))
+ headers = append(headers, byte(len(value)>>8), byte(len(value)))
+ headers = append(headers, value...)
+ }
totalLength := 12 + len(headers) + len(payloadBytes) + 4
- frame := make([]byte, 12, totalLength)
+ frame := make([]byte, totalLength)
binary.BigEndian.PutUint32(frame[0:4], uint32(totalLength))
binary.BigEndian.PutUint32(frame[4:8], uint32(len(headers)))
- frame = append(frame, headers...)
- frame = append(frame, payloadBytes...)
- frame = append(frame, 0, 0, 0, 0)
+ binary.BigEndian.PutUint32(frame[8:12], crc32.ChecksumIEEE(frame[:8]))
+ copy(frame[12:], headers)
+ copy(frame[12+len(headers):], payloadBytes)
+ binary.BigEndian.PutUint32(frame[totalLength-4:], crc32.ChecksumIEEE(frame[:totalLength-4]))
return frame
}
+
+func awsEventStreamFrame(t *testing.T, eventType string, payload map[string]interface{}) []byte {
+ t.Helper()
+ return awsEventStreamMessage(t, map[string]string{
+ ":message-type": "event",
+ ":event-type": eventType,
+ }, payload)
+}
diff --git a/proxy/stream_drop_test.go b/proxy/stream_drop_test.go
new file mode 100644
index 00000000..e595ba75
--- /dev/null
+++ b/proxy/stream_drop_test.go
@@ -0,0 +1,271 @@
+package proxy
+
+// Regression coverage for the three stream-drop failure modes.
+// Each case asserts the user-facing symptom so the suite goes red on a
+// regression and green once the stream parser is correct.
+
+import (
+ "bytes"
+ "errors"
+ "testing"
+)
+
+// Truncation is undetectable when metadataEvent.stopReason is ignored:
+// content arrives, the body ends, parseEventStream returns nil and fires
+// OnComplete. Callers cannot tell this from a complete response.
+// Kiro IDE retries on "content>0 && stopReason === undefined".
+func TestParseEventStreamReportsStopReasonFromMetadata(t *testing.T) {
+ var stream bytes.Buffer
+ stream.Write(awsEventStreamFrame(t, "assistantResponseEvent",
+ map[string]interface{}{"content": "Let me check that file"}))
+ stream.Write(awsEventStreamFrame(t, "metadataEvent",
+ map[string]interface{}{"stopReason": "end_turn"}))
+
+ var text, reason string
+ completed := false
+ err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnText: func(s string, _ bool) { text += s },
+ OnStopReason: func(r string) { reason = r },
+ OnComplete: func(_, _ int) { completed = true },
+ })
+ if err != nil {
+ t.Fatalf("unexpected parse error: %v", err)
+ }
+ if text != "Let me check that file" {
+ t.Fatalf("text = %q", text)
+ }
+ if reason != "end_turn" {
+ t.Fatalf("stopReason = %q, want end_turn", reason)
+ }
+ if !completed {
+ t.Fatal("expected OnComplete")
+ }
+}
+
+// A truncated stream (content, no stopReason) must leave OnStopReason unset so
+// callers can detect truncation. parseEventStream itself still returns nil —
+// the signal is the missing callback, not an error.
+func TestParseEventStreamTruncationLeavesStopReasonEmpty(t *testing.T) {
+ var stream bytes.Buffer
+ stream.Write(awsEventStreamFrame(t, "assistantResponseEvent",
+ map[string]interface{}{"content": "Let me check that file"}))
+
+ var reason string
+ sawStop := false
+ err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnText: func(s string, _ bool) {},
+ OnStopReason: func(r string) {
+ sawStop = true
+ reason = r
+ },
+ })
+ if err != nil {
+ t.Fatalf("unexpected parse error: %v", err)
+ }
+ if sawStop {
+ t.Fatalf("truncated stream must not fire OnStopReason, got %q", reason)
+ }
+}
+
+// Interleaved parallel tool-call frames must each emit exactly once with
+// correct args. The previous single-slot state machine force-finished the
+// previous tool on every id change, so the client executed the same tool
+// multiple times.
+func TestParseEventStreamInterleavedParallelToolCallsEmitOnce(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"toolUseId": "call_A", "name": "read_file"})
+ frame(map[string]interface{}{"toolUseId": "call_B", "name": "list_dir"})
+ frame(map[string]interface{}{"toolUseId": "call_A", "name": "read_file", "input": `{"path":"a.go"}`})
+ frame(map[string]interface{}{"toolUseId": "call_B", "name": "list_dir", "input": `{"path":"/tmp"}`})
+ frame(map[string]interface{}{"toolUseId": "call_B", "name": "list_dir", "stop": true})
+
+ got := map[string]map[string]interface{}{}
+ var order []string
+ if err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) {
+ got[tu.ToolUseID] = tu.Input
+ order = append(order, tu.ToolUseID)
+ },
+ }); err != nil {
+ t.Fatalf("parse: %v", err)
+ }
+
+ if len(order) != 2 {
+ t.Fatalf("expected each tool emitted once (2 total), got %d: %v", len(order), order)
+ }
+ if p, _ := got["call_A"]["path"].(string); p != "a.go" {
+ t.Fatalf("call_A input=%v, want path=a.go", got["call_A"])
+ }
+ if p, _ := got["call_B"]["path"].(string); p != "/tmp" {
+ t.Fatalf("call_B input=%v, want path=/tmp", got["call_B"])
+ }
+}
+
+// Incomplete tool argument JSON must not be emitted as empty input.
+// finishToolUse used to discard the Unmarshal error and substitute {}, so the
+// client executed a parameterless tool call.
+func TestParseEventStreamDropsIncompleteToolArgs(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"toolUseId": "call_X", "name": "write_file"})
+ frame(map[string]interface{}{"toolUseId": "call_X", "name": "write_file", "input": `{"path":"main.go",`})
+ frame(map[string]interface{}{"toolUseId": "call_X", "name": "write_file", "input": `"content":"pack`})
+ // EOF: no stop frame, incomplete JSON.
+
+ var emitted []KiroToolUse
+ err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { emitted = append(emitted, tu) },
+ })
+ if !errors.Is(err, errIncompleteToolUse) {
+ t.Fatalf("parse error = %v, want incomplete tool error", err)
+ }
+ if len(emitted) != 0 {
+ t.Fatalf("expected incomplete tool args to be dropped, got %d: %+v", len(emitted), emitted)
+ }
+}
+
+func TestParseEventStreamRejectsToolWithoutInput(t *testing.T) {
+ frame := awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_missing", "name": "write_file", "stop": true,
+ })
+ var emitted []KiroToolUse
+ err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { emitted = append(emitted, tu) },
+ })
+ if !errors.Is(err, errIncompleteToolUse) {
+ t.Fatalf("parse error = %v, want missing-input tool error", err)
+ }
+ if len(emitted) != 0 {
+ t.Fatalf("tool without input must not be emitted: %+v", emitted)
+ }
+}
+
+func TestParseEventStreamAcceptsExplicitEmptyToolInput(t *testing.T) {
+ frame := awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_empty", "name": "list_items", "input": `{}`, "stop": true,
+ })
+ var emitted []KiroToolUse
+ if err := parseEventStream(bytes.NewReader(frame), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { emitted = append(emitted, tu) },
+ }); err != nil {
+ t.Fatalf("parse explicit empty input: %v", err)
+ }
+ if len(emitted) != 1 || len(emitted[0].Input) != 0 {
+ t.Fatalf("explicit empty input = %+v", emitted)
+ }
+}
+
+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},
+ {"empty", 0, 0, "", false, errUpstreamEmptyResponse},
+ {"truncated content", 8, 0, "", false, errUpstreamTruncatedResponse},
+ {"reasoning only truncated", 0, 0, "", true, 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)
+ }
+ })
+ }
+}
+
+// Two tools left pending at EOF must be emitted in arrival order. The earlier
+// map-based flush relied on Go's randomized map iteration, so a parallel tool
+// turn cut short by EOF delivered the calls in a nondeterministic order.
+func TestParseEventStreamFlushesPendingToolsInArrivalOrder(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"toolUseId": "call_1", "name": "read_file", "input": `{"path":"a.go"}`})
+ frame(map[string]interface{}{"toolUseId": "call_2", "name": "list_dir", "input": `{"path":"/tmp"}`})
+ frame(map[string]interface{}{"toolUseId": "call_3", "name": "grep", "input": `{"q":"x"}`})
+ body := stream.Bytes()
+
+ // Repeat: a single pass can match by luck under map iteration.
+ for i := 0; i < 20; i++ {
+ var order []string
+ if err := parseEventStream(bytes.NewReader(body), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { order = append(order, tu.ToolUseID) },
+ }); err != nil {
+ t.Fatalf("parse: %v", err)
+ }
+ want := []string{"call_1", "call_2", "call_3"}
+ if len(order) != len(want) {
+ t.Fatalf("expected %d tool calls, got %v", len(want), order)
+ }
+ for j := range want {
+ if order[j] != want[j] {
+ t.Fatalf("tool flush order = %v, want %v (run %d)", order, want, i)
+ }
+ }
+ }
+}
+
+func TestParseEventStreamReportsMixedCompleteAndIncompleteTools(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"toolUseId": "call_ok", "name": "read_file", "input": `{"path":"a.go"}`, "stop": true})
+ frame(map[string]interface{}{"toolUseId": "call_cut", "name": "write_file", "input": `{"path":"b.go",`})
+ stream.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "END_TURN"}))
+
+ var emitted []KiroToolUse
+ err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { emitted = append(emitted, tu) },
+ })
+ if !errors.Is(err, errIncompleteToolUse) {
+ t.Fatalf("mixed tool set error = %v, want incomplete tool error", err)
+ }
+ if len(emitted) != 0 {
+ t.Fatalf("invalid parallel tool set leaked tools before validation: %+v", emitted)
+ }
+}
+
+func TestParseEventStreamKeepsPendingTargetAfterSiblingStops(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"toolUseId": "call_a", "name": "read_file"})
+ frame(map[string]interface{}{"toolUseId": "call_b", "name": "list_dir", "input": `{"path":"/tmp"}`, "stop": true})
+ frame(map[string]interface{}{"input": `{"path":"a.go"}`, "stop": true})
+
+ var tools []KiroToolUse
+ if err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnToolUse: func(tool KiroToolUse) { tools = append(tools, tool) },
+ }); err != nil {
+ t.Fatalf("parse: %v", err)
+ }
+ if len(tools) != 2 {
+ t.Fatalf("tools = %+v, want two", tools)
+ }
+ if tools[0].ToolUseID != "call_a" || tools[0].Input["path"] != "a.go" {
+ t.Fatalf("ID-less continuation did not attach to call_a: %+v", tools[0])
+ }
+ if tools[1].ToolUseID != "call_b" || tools[1].Input["path"] != "/tmp" {
+ t.Fatalf("call_b changed: %+v", tools[1])
+ }
+}
diff --git a/proxy/tool_use_order_test.go b/proxy/tool_use_order_test.go
new file mode 100644
index 00000000..954146cc
--- /dev/null
+++ b/proxy/tool_use_order_test.go
@@ -0,0 +1,33 @@
+package proxy
+
+import (
+ "bytes"
+ "testing"
+)
+
+// A fragment without toolUseId opens a synthetic entry; when the real id arrives
+// it must adopt that entry in place rather than appending a second one at the
+// end of the arrival order.
+func TestParseEventStreamRekeyKeepsArrivalPosition(t *testing.T) {
+ var stream bytes.Buffer
+ frame := func(payload map[string]interface{}) {
+ stream.Write(awsEventStreamFrame(t, "toolUseEvent", payload))
+ }
+ frame(map[string]interface{}{"name": "first_tool", "input": `{"k":`})
+ frame(map[string]interface{}{"toolUseId": "real_first", "name": "first_tool", "input": `1}`})
+ frame(map[string]interface{}{"toolUseId": "second", "name": "second_tool", "input": `{"k":2}`})
+
+ var order []string
+ if err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { order = append(order, tu.ToolUseID) },
+ }); err != nil {
+ t.Fatalf("unexpected parse error: %v", err)
+ }
+
+ if len(order) != 2 {
+ t.Fatalf("expected 2 tool calls, got %d: %v", len(order), order)
+ }
+ if order[0] != "real_first" || order[1] != "second" {
+ t.Fatalf("rekey lost arrival position: got %v, want [real_first second]", order)
+ }
+}
From 06f3c203a4215b45440e3147c8960c92b9dd4df1 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:24 +0800
Subject: [PATCH 3/6] fix(stream): classify empty and truncated upstream turns
Add same-account integrity retry for empty/truncated streams and reject
incomplete tool arguments instead of forging {}. Successful-turn sealing
is gated on terminal outcomes that are safe to continue from.
---
proxy/account_failover.go | 58 ++++++
proxy/stream_integrity.go | 124 +++++++++++++
proxy/stream_integrity_test.go | 316 +++++++++++++++++++++++++++++++++
3 files changed, 498 insertions(+)
create mode 100644 proxy/stream_integrity.go
create mode 100644 proxy/stream_integrity_test.go
diff --git a/proxy/account_failover.go b/proxy/account_failover.go
index 02624a49..4af06769 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,63 @@ import (
const maxAccountRetryAttempts = 3
+// maxSameAccountStreamRetries mirrors Kiro IDE's empty/truncated response
+// recovery: retry the same request a few times before giving up or rotating.
+const maxSameAccountStreamRetries = 2
+
+// errUpstreamEmptyResponse and errUpstreamTruncatedResponse are soft failures
+// raised when the stream ended without a usable completion signal. They are
+// retryable on the same account and must not mark the account unhealthy.
+var (
+ errUpstreamEmptyResponse = errors.New("upstream returned empty response without stop reason")
+ errUpstreamTruncatedResponse = errors.New("upstream truncated response without stop reason")
+ errIncompleteToolUse = errors.New("upstream returned incomplete tool arguments")
+)
+
+// classifyStreamIntegrity decides whether an upstream stream that returned no
+// transport error is actually complete. Mirrors Kiro IDE:
+// - empty: no content, no tools, no stopReason
+// - truncated: some content, no tools, no stopReason
+// - complete: stopReason present, or tools present, or both
+//
+// A stopReason of any non-empty value counts as complete.
+func classifyStreamIntegrity(contentChars, toolCallCount int, stopReason string, sawReasoning bool) error {
+ if strings.TrimSpace(stopReason) != "" {
+ return nil
+ }
+ if toolCallCount > 0 {
+ // Tool turns often complete without a separate stopReason frame; treat
+ // a delivered tool call as a terminal signal.
+ return nil
+ }
+ if contentChars == 0 && !sawReasoning {
+ return errUpstreamEmptyResponse
+ }
+ if contentChars > 0 {
+ return errUpstreamTruncatedResponse
+ }
+ // reasoning-only with no stopReason is still truncated for clients that
+ // expected a final answer.
+ return errUpstreamTruncatedResponse
+}
+
+func isSuccessfulKiroTurn(stopReason string, toolCallCount int) bool {
+ if strings.TrimSpace(stopReason) == "" {
+ // Complete tool turns can legitimately omit metadataEvent.stopReason.
+ return toolCallCount > 0
+ }
+ switch classifyKiroStopReason(stopReason) {
+ case kiroStopLength, kiroStopContextLimit, kiroStopFiltered:
+ return false
+ default:
+ return true
+ }
+}
+
+func isStreamIntegrityError(err error) bool {
+ return errors.Is(err, errUpstreamEmptyResponse) || 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/stream_integrity.go b/proxy/stream_integrity.go
new file mode 100644
index 00000000..1c079dfe
--- /dev/null
+++ b/proxy/stream_integrity.go
@@ -0,0 +1,124 @@
+package proxy
+
+import (
+ "errors"
+ "fmt"
+ "io"
+ "kiro-go/config"
+ "kiro-go/logger"
+)
+
+type streamIntegrityState struct {
+ ContentChars int
+ ToolCount int
+ StopReason string
+ SawReasoning bool
+}
+
+func (s streamIntegrityState) classify() error {
+ return classifyStreamIntegrity(s.ContentChars, s.ToolCount, s.StopReason, s.SawReasoning)
+}
+
+func (s *streamIntegrityState) observeText(text string, isReasoning bool) {
+ if text == "" {
+ return
+ }
+ if isReasoning {
+ s.SawReasoning = true
+ return
+ }
+ s.ContentChars += len(text)
+}
+
+func (s *streamIntegrityState) observeToolUse() {
+ s.ToolCount++
+}
+
+func (s *streamIntegrityState) observeStopReason(reason string) {
+ s.StopReason = reason
+}
+
+func (s *streamIntegrityState) reset() {
+ *s = streamIntegrityState{}
+}
+
+// runKiroWithIntegrityRetry calls CallKiroAPI and recovers empty/truncated
+// upstream streams the way Kiro IDE does: retry the same request on the same
+// account a few times before surfacing the failure.
+//
+// build is invoked at the start of every attempt (including the first) so the
+// caller can close over fresh accumulators. state owns the integrity snapshot
+// and is reset automatically between attempts. reset clears the caller's other
+// per-attempt accumulators; it 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 CallKiroAPI: 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(
+ account *config.Account,
+ payload *KiroPayload,
+ state *streamIntegrityState,
+ build func() *KiroStreamCallback,
+ reset func(),
+ canRetry func() bool,
+) error {
+ label := accountEmailForLog(account)
+ if state == nil {
+ state = &streamIntegrityState{}
+ }
+ retryable := func() bool {
+ if canRetry == nil {
+ return true
+ }
+ return canRetry()
+ }
+ for attempt := 0; attempt <= maxSameAccountStreamRetries; attempt++ {
+ if attempt > 0 {
+ if reset != nil {
+ reset()
+ }
+ state.reset()
+ }
+ callback := build()
+ err := CallKiroAPI(account, payload, callback)
+
+ var integrityErr error
+ if err != nil {
+ if !errors.Is(err, errIncompleteToolUse) && !errors.Is(err, io.ErrUnexpectedEOF) {
+ return err
+ }
+ integrityErr = fmt.Errorf("%w: %v", errUpstreamTruncatedResponse, err)
+ } else {
+ integrityErr = state.classify()
+ }
+
+ if integrityErr == nil {
+ return nil
+ }
+
+ 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 (kiro2cc-proxy #13).
+ 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
+ }
+
+ return errUpstreamTruncatedResponse
+}
diff --git a/proxy/stream_integrity_test.go b/proxy/stream_integrity_test.go
new file mode 100644
index 00000000..490e24ab
--- /dev/null
+++ b/proxy/stream_integrity_test.go
@@ -0,0 +1,316 @@
+package proxy
+
+import (
+ "kiro-go/config"
+ "net/http"
+ "net/http/httptest"
+ "path/filepath"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+// First upstream response is empty (no content, no stopReason). The helper must
+// retry on the same account and succeed once a complete frame pair arrives.
+func TestRunKiroWithIntegrityRetryRecoversEmptyThenComplete(t *testing.T) {
+ 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)
+ }
+ 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 {
+ // Transport OK but empty body => empty-response integrity failure.
+ 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()
+
+ oldEndpoints := kiroEndpoints
+ kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}}
+ defer func() { kiroEndpoints = oldEndpoints }()
+
+ oldClient := kiroHttpStore.Load()
+ kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}})
+ defer kiroHttpStore.Store(oldClient)
+
+ account := &config.Account{
+ ID: "acc",
+ Email: "acc@test",
+ AccessToken: "token",
+ ProfileArn: "arn:aws:codewhisperer:profile/test",
+ }
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
+ Content: "hi",
+ Origin: "AI_EDITOR",
+ }
+
+ var content string
+ var resets int
+ state := streamIntegrityState{}
+ err := runKiroWithIntegrityRetry(account, payload,
+ &state,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(s string, reasoning bool) { state.observeText(s, reasoning); content += s },
+ OnStopReason: state.observeStopReason,
+ }
+ },
+ func() {
+ resets++
+ content = ""
+ },
+ nil,
+ )
+ if err != nil {
+ t.Fatalf("expected recovery, got %v", err)
+ }
+ if hits.Load() != 2 {
+ t.Fatalf("upstream hits = %d, want exactly one retry", hits.Load())
+ }
+ if resets != 1 {
+ t.Fatalf("resets = %d, want exactly one reset", resets)
+ }
+ if content != "recovered" || state.StopReason != "end_turn" {
+ t.Fatalf("content=%q stopReason=%q", content, state.StopReason)
+ }
+}
+
+func TestRunKiroWithIntegrityRetryRecoversTruncatedFrame(t *testing.T) {
+ 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)
+ }
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if hits.Add(1) == 1 {
+ frame := awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "cut"})
+ _, _ = w.Write(frame[:len(frame)-3])
+ 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 swapKiroEndpointsForTest(t, server)()
+
+ account := &config.Account{ID: "acc", AccessToken: "token", ProfileArn: "arn:aws:codewhisperer:profile/test"}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{Content: "hi", Origin: "AI_EDITOR"}
+ state := streamIntegrityState{}
+ var content string
+ err := runKiroWithIntegrityRetry(account, payload, &state,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(s string, reasoning bool) { state.observeText(s, reasoning); content += s },
+ OnStopReason: state.observeStopReason,
+ }
+ },
+ func() { content = "" },
+ nil,
+ )
+ if err != nil {
+ t.Fatalf("expected partial-frame recovery, got %v", err)
+ }
+ if hits.Load() != 2 || content != "recovered" || state.StopReason != "END_TURN" {
+ t.Fatalf("hits=%d content=%q stopReason=%q", hits.Load(), content, state.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) {
+ 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)
+ }
+ 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()
+
+ oldEndpoints := kiroEndpoints
+ kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}}
+ defer func() { kiroEndpoints = oldEndpoints }()
+
+ oldClient := kiroHttpStore.Load()
+ kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}})
+ defer kiroHttpStore.Store(oldClient)
+
+ account := &config.Account{ID: "acc", AccessToken: "token", ProfileArn: "arn:aws:codewhisperer:profile/test"}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{Content: "hi", Origin: "AI_EDITOR"}
+
+ var content string
+ flushed := false
+ state := streamIntegrityState{}
+ err := runKiroWithIntegrityRetry(account, payload,
+ &state,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(s string, reasoning bool) {
+ state.observeText(s, reasoning)
+ flushed = true
+ content += s
+ },
+ }
+ },
+ 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)
+ }
+}
+
+func TestSuccessfulKiroTurnSealGate(t *testing.T) {
+ for _, reason := range []string{"MAX_TOKENS", "MAX_OUTPUT_TOKENS", "LENGTH", "MODEL_CONTEXT_WINDOW_EXCEEDED", "CONTENT_FILTERED", "REFUSAL"} {
+ if isSuccessfulKiroTurn(reason, 0) {
+ t.Fatalf("%s must not qualify for successful sealing", reason)
+ }
+ }
+ if !isSuccessfulKiroTurn("END_TURN", 0) {
+ t.Fatal("END_TURN must qualify for sealing")
+ }
+ if !isSuccessfulKiroTurn("", 1) {
+ t.Fatal("tool turn without metadata stop must qualify for sealing")
+ }
+ if isSuccessfulKiroTurn("", 0) {
+ t.Fatal("empty terminal signal must not qualify for sealing")
+ }
+}
+
+func TestProtocolStopReasonMappings(t *testing.T) {
+ if got := mapClaudeStopReason("END_TURN", 0); got != "end_turn" {
+ t.Fatalf("Claude END_TURN = %q", got)
+ }
+ if got := mapClaudeStopReason("MAX_TOKENS", 0); got != "max_tokens" {
+ t.Fatalf("Claude MAX_TOKENS = %q", got)
+ }
+ if got := mapClaudeStopReason("END_TURN", 1); got != "tool_use" {
+ t.Fatalf("Claude tool turn = %q", got)
+ }
+ if got := mapOpenAIFinishReason("MAX_TOKENS", 0); got != "length" {
+ t.Fatalf("OpenAI MAX_TOKENS = %q", got)
+ }
+ if got := mapOpenAIFinishReason("CONTENT_FILTER", 0); got != "content_filter" {
+ t.Fatalf("OpenAI CONTENT_FILTER = %q", got)
+ }
+ if got := mapClaudeStopReason("CONTENT_FILTERED", 0); got != "refusal" {
+ t.Fatalf("Claude CONTENT_FILTERED = %q", got)
+ }
+ if got := mapOpenAIFinishReason("CONTENT_FILTERED", 0); got != "content_filter" {
+ t.Fatalf("OpenAI CONTENT_FILTERED = %q", got)
+ }
+ if got := mapOpenAIFinishReason("END_TURN", 1); got != "tool_calls" {
+ t.Fatalf("OpenAI tool turn = %q", got)
+ }
+ status, reason := mapResponsesCompletion("MAX_TOKENS")
+ if status != "incomplete" || reason != "max_output_tokens" {
+ t.Fatalf("Responses MAX_TOKENS = status %q reason %q", status, reason)
+ }
+ status, reason = mapResponsesCompletion("CONTENT_FILTERED")
+ if status != "incomplete" || reason != "content_filter" {
+ t.Fatalf("Responses CONTENT_FILTERED = status %q reason %q", status, reason)
+ }
+ status, reason = mapResponsesCompletion("END_TURN")
+ if status != "completed" || reason != "" {
+ t.Fatalf("Responses END_TURN = status %q reason %q", status, reason)
+ }
+}
+
+func TestRunKiroWithIntegrityRetryRetriesMixedToolSet(t *testing.T) {
+ 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)
+ }
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if hits.Add(1) == 1 {
+ _, _ = w.Write(awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_ok", "name": "read_file", "input": `{"path":"a.go"}`, "stop": true,
+ }))
+ _, _ = w.Write(awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_cut", "name": "write_file", "input": `{"path":"b.go",`, "stop": true,
+ }))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "END_TURN"}))
+ return
+ }
+ _, _ = w.Write(awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_fresh", "name": "read_file", "input": `{"path":"fresh.go"}`, "stop": true,
+ }))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "TOOL_USE"}))
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(t, server)()
+
+ account := &config.Account{ID: "acc", AccessToken: "token", ProfileArn: "arn:aws:codewhisperer:profile/test"}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{Content: "hi", Origin: "AI_EDITOR"}
+ var tools []KiroToolUse
+ state := streamIntegrityState{}
+ err := runKiroWithIntegrityRetry(account, payload, &state,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnToolUse: func(tu KiroToolUse) { state.observeToolUse(); tools = append(tools, tu) },
+ OnStopReason: state.observeStopReason,
+ }
+ },
+ func() { tools = nil },
+ nil,
+ )
+ if err != nil {
+ t.Fatalf("expected retry recovery, got %v", err)
+ }
+ if hits.Load() != 2 {
+ t.Fatalf("upstream hits = %d, want retry after mixed tool set", hits.Load())
+ }
+ if len(tools) != 1 || tools[0].ToolUseID != "call_fresh" {
+ t.Fatalf("partial first attempt leaked into recovered tool set: %+v", tools)
+ }
+}
From 121051824efed3bc14e2202e139b30804c9109b9 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:24 +0800
Subject: [PATCH 4/6] feat(proxy): exact session continuity and reasoning stamp
reattach
Mint conversation/continuation ids once and reuse them only on an exact
client-transcript match. Store thinking signatures by exact turn key and
reattach on the next request from the handler path, without fuzzy matching.
---
proxy/reasoning_store.go | 240 +++++++++++++++++++++
proxy/reasoning_store_test.go | 392 ++++++++++++++++++++++++++++++++++
proxy/session_store.go | 336 +++++++++++++++++++++++++++++
proxy/session_store_test.go | 375 ++++++++++++++++++++++++++++++++
4 files changed, 1343 insertions(+)
create mode 100644 proxy/reasoning_store.go
create mode 100644 proxy/reasoning_store_test.go
create mode 100644 proxy/session_store.go
create mode 100644 proxy/session_store_test.go
diff --git a/proxy/reasoning_store.go b/proxy/reasoning_store.go
new file mode 100644
index 00000000..70319a6b
--- /dev/null
+++ b/proxy/reasoning_store.go
@@ -0,0 +1,240 @@
+package proxy
+
+import (
+ "crypto/sha1"
+ "encoding/base64"
+ "encoding/hex"
+ "sort"
+ "strings"
+ "sync"
+ "time"
+)
+
+// sealedReasoning is one successful assistant turn's upstream thinking stamp.
+// Official client keeps this in process memory (lastSealedReasoning) and pastes
+// it onto that assistant message on the next request. As a stateless proxy we
+// re-derive the same binding with a turn fingerprint so we never attach a stamp
+// to the wrong conversation.
+type sealedReasoning struct {
+ Text string
+ Signature string
+ RedactedContent string
+ ModelID string
+ StoredAt time.Time
+}
+
+const (
+ reasoningStoreTTL = 2 * time.Hour
+ reasoningStoreMaxEntries = 4096
+)
+
+type reasoningStore struct {
+ mu sync.Mutex
+ // key -> stamp. key = turnKey(...)
+ byKey map[string]sealedReasoning
+}
+
+func newReasoningStore() *reasoningStore {
+ return &reasoningStore{byKey: make(map[string]sealedReasoning)}
+}
+
+var globalReasoningStore = newReasoningStore()
+
+// turnKey builds the exact-match address for one assistant turn.
+//
+// apiKey | model | conversationBucket | hash(parentUser, assistantText, ordered tool uses)
+//
+// No fuzzy matching. Same account + same model alone is never enough.
+func turnKey(apiKeyID, modelID, conversationID, parentUser, assistantText string, toolUses []KiroToolUse) string {
+ h := sha1.New()
+ h.Write([]byte(parentUser))
+ h.Write([]byte{0})
+ h.Write([]byte(assistantText))
+ h.Write([]byte{0})
+ for _, tu := range toolUses {
+ h.Write([]byte(tu.ToolUseID))
+ h.Write([]byte{0})
+ h.Write([]byte(tu.Name))
+ h.Write([]byte{0})
+ }
+ fp := hex.EncodeToString(h.Sum(nil))
+ return strings.Join([]string{
+ strings.TrimSpace(apiKeyID),
+ strings.TrimSpace(modelID),
+ strings.TrimSpace(conversationID),
+ fp,
+ }, "|")
+}
+
+// put seals a successful turn. Unsigned reasoning (no signature and no
+// redacted blob) is dropped — same rule as the official client.
+func (s *reasoningStore) put(apiKeyID, modelID, conversationID, parentUser, assistantText string, toolUses []KiroToolUse, text, signature, redacted string) {
+ if signature == "" && redacted == "" {
+ return
+ }
+ if text == "" && redacted == "" {
+ return
+ }
+ key := turnKey(apiKeyID, modelID, conversationID, parentUser, assistantText, toolUses)
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.purgeLocked(time.Now())
+ if len(s.byKey) >= reasoningStoreMaxEntries {
+ // Drop a quarter of entries when at capacity.
+ s.dropOldestLocked(reasoningStoreMaxEntries / 4)
+ }
+ s.byKey[key] = sealedReasoning{
+ Text: text,
+ Signature: signature,
+ RedactedContent: redacted,
+ ModelID: modelID,
+ StoredAt: time.Now(),
+ }
+}
+
+// lookup returns a stamp only on exact turn-key hit + model match + not expired.
+// Miss → nil. Never guesses.
+func (s *reasoningStore) lookup(apiKeyID, modelID, conversationID, parentUser, assistantText string, toolUses []KiroToolUse) *sealedReasoning {
+ key := turnKey(apiKeyID, modelID, conversationID, parentUser, assistantText, toolUses)
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.purgeLocked(time.Now())
+ e, ok := s.byKey[key]
+ if !ok {
+ return nil
+ }
+ if modelID != "" && e.ModelID != "" && e.ModelID != modelID {
+ return nil
+ }
+ cp := e
+ return &cp
+}
+
+func (s *reasoningStore) purgeLocked(now time.Time) {
+ for k, e := range s.byKey {
+ if now.Sub(e.StoredAt) > reasoningStoreTTL {
+ delete(s.byKey, k)
+ }
+ }
+}
+
+func (s *reasoningStore) dropOldestLocked(n int) {
+ type pair struct {
+ k string
+ t time.Time
+ }
+ all := make([]pair, 0, len(s.byKey))
+ for k, e := range s.byKey {
+ all = append(all, pair{k, e.StoredAt})
+ }
+ sort.Slice(all, func(i, j int) bool { return all[i].t.Before(all[j].t) })
+ if n > len(all) {
+ n = len(all)
+ }
+ for i := 0; i < n; i++ {
+ delete(s.byKey, all[i].k)
+ }
+}
+
+// toKiroReasoning converts a sealed stamp into the wire reasoningContent object.
+func toKiroReasoning(e *sealedReasoning) *KiroReasoningContent {
+ if e == nil {
+ return nil
+ }
+ if e.RedactedContent != "" {
+ raw, err := base64.StdEncoding.DecodeString(e.RedactedContent)
+ if err != nil || len(raw) == 0 {
+ return nil
+ }
+ return &KiroReasoningContent{RedactedContent: raw}
+ }
+ if e.Signature == "" {
+ return nil
+ }
+ rt := struct {
+ Text string `json:"text"`
+ Signature string `json:"signature"`
+ }{Text: e.Text, Signature: e.Signature}
+ return &KiroReasoningContent{ReasoningText: &rt}
+}
+
+// lastUserTextBeforeAssistant returns the nearest preceding user text for an
+// assistant history index. Used only for turn-key parent binding.
+func lastUserTextBeforeAssistant(history []KiroHistoryMessage, assistantIdx int) string {
+ for i := assistantIdx - 1; i >= 0; i-- {
+ if history[i].UserInputMessage != nil {
+ return history[i].UserInputMessage.Content
+ }
+ }
+ return ""
+}
+
+// attachStoredReasoning walks history and pastes exact-match stamps onto
+// assistant turns. No fuzzy fill. Safe to call every request.
+func attachStoredReasoning(history []KiroHistoryMessage, apiKeyID, modelID, conversationID string) {
+ if len(history) == 0 {
+ return
+ }
+ for i := range history {
+ as := history[i].AssistantResponseMessage
+ if as == nil || as.ReasoningContent != nil {
+ continue
+ }
+ parent := lastUserTextBeforeAssistant(history, i)
+ if e := globalReasoningStore.lookup(apiKeyID, modelID, conversationID, parent, as.Content, as.ToolUses); e != nil {
+ as.ReasoningContent = toKiroReasoning(e)
+ }
+ }
+}
+
+// sealSuccessfulTurn stores the stamp for this completed turn.
+// currentTurnParentUser is the user text of the turn being answered: always
+// payload.CurrentMessage, never the last user already in History.
+func currentTurnParentUser(payload *KiroPayload) string {
+ if payload == nil {
+ return ""
+ }
+ return payload.ConversationState.CurrentMessage.UserInputMessage.Content
+}
+
+func sealSuccessfulTurn(apiKeyID, modelID, conversationID, parentUser, assistantText string, toolUses []KiroToolUse, thinkingText, signature, redacted string) {
+ globalReasoningStore.put(apiKeyID, modelID, conversationID, parentUser, assistantText, toolUses, thinkingText, signature, redacted)
+}
+
+// debugTurnKey is for package tests.
+func debugTurnKey(apiKeyID, modelID, conversationID, parentUser, assistantText string, toolUses []KiroToolUse) string {
+ return turnKey(apiKeyID, modelID, conversationID, parentUser, assistantText, toolUses)
+}
+
+// reasoningCapture accumulates OnReasoningMeta for one upstream attempt.
+type reasoningCapture struct {
+ sig, red string
+}
+
+func (c *reasoningCapture) meta() func(string, string) {
+ return func(sig, red string) {
+ if c == nil {
+ return
+ }
+ if sig != "" {
+ c.sig = sig
+ }
+ if red != "" {
+ c.red = red
+ }
+ }
+}
+
+func (c *reasoningCapture) reset() {
+ if c == nil {
+ return
+ }
+ c.sig, c.red = "", ""
+}
+
+func (c *reasoningCapture) sealIfPresent(apiKeyID, model, convID, parent, assistant string, tools []KiroToolUse, thinkingText string) {
+ if c == nil || (c.sig == "" && c.red == "") {
+ return
+ }
+ sealSuccessfulTurn(apiKeyID, model, convID, parent, assistant, tools, thinkingText, c.sig, c.red)
+}
diff --git a/proxy/reasoning_store_test.go b/proxy/reasoning_store_test.go
new file mode 100644
index 00000000..06cb369d
--- /dev/null
+++ b/proxy/reasoning_store_test.go
@@ -0,0 +1,392 @@
+package proxy
+
+import (
+ "bytes"
+ "encoding/json"
+ "testing"
+)
+
+// Seam: turnKey — same turn content must hash equal; different assistant text must not.
+func TestTurnKeyStableAndIsolatesAssistantText(t *testing.T) {
+ a := debugTurnKey("k1", "claude-sonnet-4.5", "conv-a", "hello", "answer one", nil)
+ b := debugTurnKey("k1", "claude-sonnet-4.5", "conv-a", "hello", "answer one", nil)
+ if a != b {
+ t.Fatalf("same turn must produce same key")
+ }
+ c := debugTurnKey("k1", "claude-sonnet-4.5", "conv-a", "hello", "answer two", nil)
+ if a == c {
+ t.Fatalf("different assistant text must not share a key")
+ }
+}
+
+// Seam: store — two parallel chats with the same opener must not share stamps.
+func TestReasoningStoreDoesNotCrossConversations(t *testing.T) {
+ s := newReasoningStore()
+ tools := []KiroToolUse{}
+ s.put("keyA", "m1", "conv1", "hi", "reply-A", tools, "think-A", "sig-A", "")
+ s.put("keyA", "m1", "conv1", "hi", "reply-B", tools, "think-B", "sig-B", "")
+
+ // Lookup for reply-A must not return B's stamp.
+ got := s.lookup("keyA", "m1", "conv1", "hi", "reply-A", tools)
+ if got == nil || got.Signature != "sig-A" {
+ t.Fatalf("expected sig-A for reply-A, got %+v", got)
+ }
+ // Different conversation bucket, same text — miss.
+ if s.lookup("keyA", "m1", "conv-OTHER", "hi", "reply-A", tools) != nil {
+ t.Fatal("must not hit across conversation buckets")
+ }
+ // Different API key — miss.
+ if s.lookup("keyB", "m1", "conv1", "hi", "reply-A", tools) != nil {
+ t.Fatal("must not hit across API keys")
+ }
+ // Different model — miss.
+ if s.lookup("keyA", "other-model", "conv1", "hi", "reply-A", tools) != nil {
+ t.Fatal("must not hit across models")
+ }
+}
+
+// Seam: store — unsigned reasoning is never stored (official drops it).
+func TestReasoningStoreDropsUnsigned(t *testing.T) {
+ s := newReasoningStore()
+ s.put("k", "m", "c", "u", "a", nil, "thinking text", "", "")
+ if s.lookup("k", "m", "c", "u", "a", nil) != nil {
+ t.Fatal("unsigned reasoning must not be stored")
+ }
+}
+
+// Seam: parseEventStream — signature on reasoningContentEvent must reach OnReasoningMeta.
+func TestParseEventStreamFiresReasoningMetaWithSignature(t *testing.T) {
+ var stream bytes.Buffer
+ stream.Write(awsEventStreamFrame(t, "reasoningContentEvent", map[string]interface{}{
+ "text": "I should check the file",
+ "signature": "sig_abc123",
+ }))
+ stream.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{
+ "content": "done",
+ }))
+ stream.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{
+ "stopReason": "end_turn",
+ }))
+
+ var thinking, sig string
+ err := parseEventStream(bytes.NewReader(stream.Bytes()), &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ if isThinking {
+ thinking += text
+ }
+ },
+ OnReasoningMeta: func(signature, _ string) { sig = signature },
+ })
+ if err != nil {
+ t.Fatalf("parse: %v", err)
+ }
+ if thinking != "I should check the file" {
+ t.Fatalf("thinking text = %q", thinking)
+ }
+ if sig != "sig_abc123" {
+ t.Fatalf("signature = %q, want sig_abc123", sig)
+ }
+}
+
+// Seam: attachStoredReasoning — only exact-matching assistant turns get stamps.
+func TestAttachStoredReasoningExactMatchOnly(t *testing.T) {
+ // Reset global store by sealing via public API then attaching.
+ // Use unique conversation id to avoid pollution from other tests.
+ conv := "conv-attach-test"
+ apiKey := "ak-attach"
+ model := "claude-sonnet-4.5"
+ parent := "what is 1+1?"
+ assistant := "2"
+ sealSuccessfulTurn(apiKey, model, conv, parent, assistant, nil, "add the numbers", "sig-42", "")
+
+ history := []KiroHistoryMessage{
+ {UserInputMessage: &KiroUserInputMessage{Content: parent, Origin: "AI_EDITOR"}},
+ {AssistantResponseMessage: &KiroAssistantResponseMessage{Content: assistant}},
+ {UserInputMessage: &KiroUserInputMessage{Content: "and 2+2?", Origin: "AI_EDITOR"}},
+ // Different assistant text — must stay unstamped.
+ {AssistantResponseMessage: &KiroAssistantResponseMessage{Content: "4"}},
+ }
+ attachStoredReasoning(history, apiKey, model, conv)
+
+ if history[1].AssistantResponseMessage.ReasoningContent == nil {
+ t.Fatal("matching turn must receive reasoningContent")
+ }
+ rt := history[1].AssistantResponseMessage.ReasoningContent.ReasoningText
+ if rt == nil || rt.Signature != "sig-42" || rt.Text != "add the numbers" {
+ t.Fatalf("bad stamp on matching turn: %+v", history[1].AssistantResponseMessage.ReasoningContent)
+ }
+ if history[3].AssistantResponseMessage.ReasoningContent != nil {
+ t.Fatal("non-matching turn must not receive a stamp")
+ }
+}
+
+// Seam: ClaudeToKiro — after sealing turn N, the next request's history for that
+// assistant message must carry reasoningContent for upstream.
+func TestClaudeToKiroAttachesSealedReasoningOnNextTurn(t *testing.T) {
+ convSeedUser := "unique-opener-for-reasoning-tdd-xyz"
+ model := "claude-sonnet-4.5"
+ modelID := MapModel(model)
+
+ // Turn 1: mint session ids the same way the handler/translator will.
+ turn1 := []ClaudeMessage{{Role: "user", Content: convSeedUser}}
+ p1 := ClaudeToKiro(&ClaudeRequest{Model: model, Messages: turn1}, true)
+ convID := p1.ConversationState.ConversationID
+ contID := p1.ConversationState.AgentContinuationId
+
+ // After successful stream: seal reasoning + session chain (handler job).
+ sealSuccessfulTurn("", modelID, convID, convSeedUser, "first answer", nil, "think first", "sig-turn1", "")
+ sealSessionAfterSuccess("", modelID, convID, contID, []ClaudeMessage{
+ {Role: "user", Content: convSeedUser},
+ {Role: "assistant", Content: "first answer"},
+ })
+
+ // Turn 2: continuation must reuse ids; attach must stamp first answer.
+ req := &ClaudeRequest{
+ Model: model,
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: convSeedUser},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": "first answer"},
+ }},
+ {Role: "user", Content: "follow up"},
+ },
+ }
+ payload := ClaudeToKiro(req, true)
+ // Handler binds session ids with apiKey (empty here).
+ c, k, _ := resolveSessionIDs("", modelID, req.Messages)
+ payload.ConversationState.ConversationID = c
+ payload.ConversationState.AgentContinuationId = k
+ if payload.ConversationState.ConversationID != convID {
+ t.Fatalf("conversationId = %q, want reused %q", payload.ConversationState.ConversationID, convID)
+ }
+ attachStoredReasoning(payload.ConversationState.History, "", modelID, payload.ConversationState.ConversationID)
+ var found bool
+ for _, h := range payload.ConversationState.History {
+ if h.AssistantResponseMessage == nil || h.AssistantResponseMessage.Content != "first answer" {
+ continue
+ }
+ found = true
+ rc := h.AssistantResponseMessage.ReasoningContent
+ if rc == nil || rc.ReasoningText == nil || rc.ReasoningText.Signature != "sig-turn1" {
+ t.Fatalf("assistant history missing reasoningContent: %+v", rc)
+ }
+ }
+ if !found {
+ t.Fatalf("assistant turn not in history: %s", mustJSON(payload.ConversationState.History))
+ }
+}
+
+func mustJSON(v interface{}) string {
+ b, _ := json.Marshal(v)
+ return string(b)
+}
+
+// Handler must seal with CurrentMessage as parent (the user just answered),
+// not the last user already sitting in History (previous turn).
+func TestSealUsesCurrentMessageParentNotHistory(t *testing.T) {
+ model := "claude-sonnet-4.5"
+ opener := "seal-parent-opener-unique-aaa"
+ follow := "seal-parent-follow-up-bbb"
+ modelID := MapModel(model)
+
+ p1 := ClaudeToKiro(&ClaudeRequest{Model: model, Messages: []ClaudeMessage{{Role: "user", Content: opener}}}, true)
+ convID := p1.ConversationState.ConversationID
+ contID := p1.ConversationState.AgentContinuationId
+
+ // Seal as handler: parent = CurrentMessage (opener), then session chain.
+ sealSuccessfulTurn("", modelID, convID, opener, "answer-1", nil, "think-1", "sig-1", "")
+ sealSessionAfterSuccess("", modelID, convID, contID, []ClaudeMessage{
+ {Role: "user", Content: opener},
+ {Role: "assistant", Content: "answer-1"},
+ })
+
+ req := &ClaudeRequest{
+ Model: model,
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: opener},
+ {Role: "assistant", Content: []interface{}{map[string]interface{}{"type": "text", "text": "answer-1"}}},
+ {Role: "user", Content: follow},
+ },
+ }
+ payload := ClaudeToKiro(req, true)
+ // Handler binds session ids with apiKey (empty here).
+ c, k, _ := resolveSessionIDs("", modelID, req.Messages)
+ payload.ConversationState.ConversationID = c
+ payload.ConversationState.AgentContinuationId = k
+ if payload.ConversationState.ConversationID != convID {
+ t.Fatalf("want reused conv %q, got %q", convID, payload.ConversationState.ConversationID)
+ }
+ attachStoredReasoning(payload.ConversationState.History, "", modelID, payload.ConversationState.ConversationID)
+ var stamped bool
+ for _, h := range payload.ConversationState.History {
+ if h.AssistantResponseMessage != nil && h.AssistantResponseMessage.Content == "answer-1" {
+ rc := h.AssistantResponseMessage.ReasoningContent
+ if rc == nil || rc.ReasoningText == nil || rc.ReasoningText.Signature != "sig-1" {
+ t.Fatalf("expected sig-1 on answer-1, got %+v", rc)
+ }
+ stamped = true
+ }
+ }
+ if !stamped {
+ t.Fatal("answer-1 not found/stamped in history")
+ }
+ if globalReasoningStore.lookup("", modelID, convID, follow, "answer-1", nil) != nil {
+ t.Fatal("lookup with wrong parent must miss")
+ }
+}
+
+func TestRedactedReasoningRoundTripsBase64Once(t *testing.T) {
+ rc := toKiroReasoning(&sealedReasoning{RedactedContent: "AQID"})
+ if rc == nil {
+ t.Fatal("expected redacted reasoning content")
+ }
+ wire, err := json.Marshal(rc)
+ if err != nil {
+ t.Fatalf("marshal redacted reasoning: %v", err)
+ }
+ if got, want := string(wire), `{"redactedContent":"AQID"}`; got != want {
+ t.Fatalf("wire reasoning = %s, want %s", got, want)
+ }
+}
+
+func TestInvalidRedactedReasoningFailsClosed(t *testing.T) {
+ if rc := toKiroReasoning(&sealedReasoning{RedactedContent: "not base64"}); rc != nil {
+ t.Fatalf("invalid redacted content must not reach upstream: %+v", rc)
+ }
+}
+
+func TestReasoningTurnKeyUsesExactTextAndToolIdentity(t *testing.T) {
+ if debugTurnKey("k", "m", "c", "a b", "answer", nil) == debugTurnKey("k", "m", "c", "a b", "answer", nil) {
+ t.Fatal("reasoning key must preserve exact parent whitespace")
+ }
+ toolsA := []KiroToolUse{
+ {ToolUseID: "a", Name: "read", Input: map[string]interface{}{"path": "a.go"}},
+ {ToolUseID: "b", Name: "list", Input: map[string]interface{}{"path": "/tmp"}},
+ }
+ toolsB := []KiroToolUse{
+ {ToolUseID: "a", Name: "read", Input: map[string]interface{}{"path": "b.go"}},
+ {ToolUseID: "b", Name: "list", Input: map[string]interface{}{"path": "elsewhere"}},
+ }
+ if debugTurnKey("k", "m", "c", "u", "a", toolsA) != debugTurnKey("k", "m", "c", "u", "a", toolsB) {
+ t.Fatal("reasoning key contract binds tool IDs/names, not arguments")
+ }
+ reversed := []KiroToolUse{toolsA[1], toolsA[0]}
+ if debugTurnKey("k", "m", "c", "u", "a", toolsA) == debugTurnKey("k", "m", "c", "u", "a", reversed) {
+ t.Fatal("reasoning key must preserve tool order")
+ }
+}
+
+func TestReasoningReattachesAcrossRawToolResultParent(t *testing.T) {
+ globalReasoningStore = newReasoningStore()
+ globalSessionStore = newSessionStore()
+ const api = "key-tool-parent"
+ const model = "claude-sonnet-4.5"
+ const conv = "conv-tool-parent"
+ const cont = "cont-tool-parent"
+
+ inbound := []ClaudeMessage{
+ {Role: "user", Content: "read it"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "tool_use", "id": "tu_1", "name": "read_file", "input": map[string]interface{}{"path": "a.go"}},
+ }},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "tu_1", "content": "file body"},
+ }},
+ }
+ // Production seals against the translated current-message parent, then
+ // reattaches via attachStoredReasoning on the next request.
+ // For pure tool_result turns, ClaudeToKiro folds the parent into:
+ // "Tool results:\n\n[read_file] file body".
+ stamp := &reasoningCapture{sig: "sig-after-tool"}
+ parent := "Tool results:\n\n[read_file] file body"
+ stamp.sealIfPresent(api, MapModel(model), conv, parent, "final answer", nil, "reason after tool")
+ sealClaudeSession(api, model, conv, cont, inbound, "final answer", "", nil, "thinking", false)
+
+ req := &ClaudeRequest{Model: model, Messages: append(append([]ClaudeMessage(nil), inbound...),
+ ClaudeMessage{Role: "assistant", Content: "final answer"},
+ ClaudeMessage{Role: "user", Content: "next"},
+ )}
+ payload := ClaudeToKiro(req, true)
+ payload.ConversationState.ConversationID = conv
+ payload.ConversationState.AgentContinuationId = cont
+ attachStoredReasoning(payload.ConversationState.History, api, MapModel(model), conv)
+ for _, message := range payload.ConversationState.History {
+ asst := message.AssistantResponseMessage
+ if asst == nil || asst.Content != "final answer" {
+ continue
+ }
+ if asst.ReasoningContent == nil || asst.ReasoningContent.ReasoningText == nil || asst.ReasoningContent.ReasoningText.Signature != "sig-after-tool" {
+ t.Fatalf("post-tool assistant missing exact reasoning stamp: %+v", asst.ReasoningContent)
+ }
+ return
+ }
+ t.Fatal("post-tool assistant not found in translated history")
+}
+
+func TestReasoningSealUsesClientVisibleThinkFormat(t *testing.T) {
+ globalReasoningStore = newReasoningStore()
+ const api = "key-think-key"
+ const model = "claude-sonnet-4.5"
+ const conv = "conv-think-key"
+ inbound := []ClaudeMessage{{Role: "user", Content: "question"}}
+ stamp := &reasoningCapture{sig: "sig-think-key"}
+ // Production seals with the client-visible assistant text, which for
+ // "think" format is already folded by claudeAssistantTurn.
+ visible := "hidden thoughtanswer"
+ stamp.sealIfPresent(api, MapModel(model), conv, "question", visible, nil, "hidden thought")
+ sealClaudeSession(api, model, conv, "cont-think-key", inbound, "answer", "hidden thought", nil, "think", false)
+ req := &ClaudeRequest{Model: model, Messages: []ClaudeMessage{
+ {Role: "user", Content: "question"},
+ {Role: "assistant", Content: visible},
+ {Role: "user", Content: "next"},
+ }}
+ payload := ClaudeToKiro(req, true)
+ payload.ConversationState.ConversationID = conv
+ attachStoredReasoning(payload.ConversationState.History, api, MapModel(model), conv)
+ for _, message := range payload.ConversationState.History {
+ asst := message.AssistantResponseMessage
+ if asst != nil && asst.Content == visible {
+ if asst.ReasoningContent == nil || asst.ReasoningContent.ReasoningText == nil || asst.ReasoningContent.ReasoningText.Signature != "sig-think-key" {
+ t.Fatalf("think-format key did not reattach: %+v", asst.ReasoningContent)
+ }
+ return
+ }
+ }
+ t.Fatal("think-format assistant not found")
+}
+
+func TestHiddenReasoningStillSealsAndIncompleteTurnDoesNot(t *testing.T) {
+ globalReasoningStore = newReasoningStore()
+ globalSessionStore = newSessionStore()
+ const api = "key-hidden-reasoning"
+ const model = "claude-sonnet-4.5"
+ const conv = "conv-hidden-reasoning"
+ inbound := []ClaudeMessage{{Role: "user", Content: "question"}}
+ stamp := &reasoningCapture{sig: "sig-hidden"}
+ if !isSuccessfulKiroTurn("END_TURN", 0) {
+ t.Fatal("END_TURN must qualify for sealing")
+ }
+ stamp.sealIfPresent(api, MapModel(model), conv, "question", "answer", nil, "upstream hidden thought")
+ sealClaudeSession(api, model, conv, "cont-hidden", inbound, "answer", "upstream hidden thought", nil, "thinking", true)
+ if got := globalReasoningStore.lookup(api, MapModel(model), conv, "question", "answer", nil); got == nil || got.Signature != "sig-hidden" {
+ t.Fatalf("hidden reasoning was not stored: %+v", got)
+ }
+
+ globalReasoningStore = newReasoningStore()
+ globalSessionStore = newSessionStore()
+ if isSuccessfulKiroTurn("MAX_TOKENS", 0) {
+ t.Fatal("MAX_TOKENS must not qualify for sealing")
+ }
+ // Production only seals after isSuccessfulKiroTurn; incomplete turns skip both stores.
+ if got := globalReasoningStore.lookup(api, MapModel(model), conv, "question", "answer", nil); got != nil {
+ t.Fatalf("incomplete reasoning leaked into store: %+v", got)
+ }
+ _, _, reused := resolveSessionIDs(api, MapModel(model), []ClaudeMessage{
+ {Role: "user", Content: "question"},
+ {Role: "assistant", Content: "answer"},
+ {Role: "user", Content: "next"},
+ })
+ if reused {
+ t.Fatal("MAX_TOKENS turn must not seal session IDs")
+ }
+}
diff --git a/proxy/session_store.go b/proxy/session_store.go
new file mode 100644
index 00000000..e0479df9
--- /dev/null
+++ b/proxy/session_store.go
@@ -0,0 +1,336 @@
+package proxy
+
+import (
+ "crypto/sha1"
+ "encoding/hex"
+ "encoding/json"
+ "sort"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/google/uuid"
+)
+
+// sessionIDs are the two upstream continuity handles for one agent run.
+type sessionIDs struct {
+ ConversationID string
+ ContinuationID string
+ Ambiguous bool
+ StoredAt time.Time
+}
+
+const sessionStoreTTL = 2 * time.Hour
+const sessionStoreMax = 4096
+
+type sessionStore struct {
+ mu sync.Mutex
+ byKey map[string]sessionIDs
+}
+
+func newSessionStore() *sessionStore {
+ return &sessionStore{byKey: make(map[string]sessionIDs)}
+}
+
+var globalSessionStore = newSessionStore()
+
+// chainFingerprint hashes a normalized Claude message prefix.
+// Used as the exact-match address for "this transcript ends here".
+func chainFingerprintClaude(messages []ClaudeMessage) string {
+ h := sha1.New()
+ for _, msg := range messages {
+ h.Write([]byte(strings.ToLower(strings.TrimSpace(msg.Role))))
+ h.Write([]byte{0})
+ h.Write([]byte(normalizeClaudeMessageContent(msg.Content)))
+ h.Write([]byte{0})
+ }
+ return hex.EncodeToString(h.Sum(nil))
+}
+
+func normalizeClaudeMessageContent(content interface{}) string {
+ if content == nil {
+ return ""
+ }
+ if text, ok := content.(string); ok {
+ return stringifyLoose([]map[string]interface{}{{
+ "type": "text",
+ "text": text,
+ }})
+ }
+
+ blocks := contentBlocksAsMaps(content)
+ if blocks == nil {
+ return stringifyLoose(content)
+ }
+ normalized := make([]map[string]interface{}, 0, len(blocks))
+ for _, block := range blocks {
+ typ, _ := block["type"].(string)
+ switch typ {
+ case "text", "input_text":
+ text, _ := block["text"].(string)
+ normalized = append(normalized, map[string]interface{}{
+ "type": "text",
+ "text": text,
+ })
+ case "thinking":
+ thinking, _ := block["thinking"].(string)
+ normalized = append(normalized, map[string]interface{}{
+ "type": "thinking",
+ "thinking": thinking,
+ })
+ case "tool_use":
+ input := block["input"]
+ if input == nil {
+ input = map[string]interface{}{}
+ }
+ normalized = append(normalized, map[string]interface{}{
+ "type": "tool_use",
+ "id": block["id"],
+ "name": block["name"],
+ "input": input,
+ })
+ case "tool_result":
+ result := map[string]interface{}{
+ "type": "tool_result",
+ "tool_use_id": block["tool_use_id"],
+ "content": block["content"],
+ }
+ if isError, ok := block["is_error"]; ok {
+ result["is_error"] = isError
+ }
+ normalized = append(normalized, result)
+ default:
+ normalized = append(normalized, block)
+ }
+ }
+ return stringifyLoose(normalized)
+}
+
+// stringifyLoose renders a structured value deterministically for
+// fingerprinting. encoding/json sorts map keys, so equivalent in-memory map
+// order does not change the chain address. Marshal failures fail closed.
+func stringifyLoose(v interface{}) string {
+ b, err := json.Marshal(v)
+ if err != nil {
+ return ""
+ }
+ return string(b)
+}
+
+func sessionKey(apiKeyID, modelID, chainFP string) string {
+ return strings.Join([]string{
+ strings.TrimSpace(apiKeyID),
+ strings.TrimSpace(modelID),
+ strings.TrimSpace(chainFP),
+ }, "|")
+}
+
+// resolveSessionIDs returns reused IDs when the inbound transcript prefix exactly
+// matches a previously successful seal; otherwise mints fresh UUIDs.
+//
+// prefix = messages without the trailing user turn being answered now.
+// Only reuses when prefix is non-empty and ends with an assistant message
+// (i.e. this is a true continuation, not a brand-new chat).
+func resolveSessionIDs(apiKeyID, modelID string, messages []ClaudeMessage) (convID, contID string, reused bool) {
+ prefix := claudePrefixBeforeCurrentUser(messages)
+ if len(prefix) == 0 || !claudeEndsWithAssistant(prefix) {
+ return uuid.New().String(), uuid.New().String(), false
+ }
+ fp := chainFingerprintClaude(prefix)
+ if ids := globalSessionStore.lookup(apiKeyID, modelID, fp); ids != nil {
+ return ids.ConversationID, ids.ContinuationID, true
+ }
+ return uuid.New().String(), uuid.New().String(), false
+}
+
+func claudePrefixBeforeCurrentUser(messages []ClaudeMessage) []ClaudeMessage {
+ if len(messages) == 0 {
+ return nil
+ }
+ last := messages[len(messages)-1]
+ if strings.EqualFold(last.Role, "user") {
+ return messages[:len(messages)-1]
+ }
+ // Already ends mid-assistant (unusual); treat full list as prefix.
+ return messages
+}
+
+func claudeEndsWithAssistant(messages []ClaudeMessage) bool {
+ if len(messages) == 0 {
+ return false
+ }
+ return strings.EqualFold(messages[len(messages)-1].Role, "assistant")
+}
+
+// sealSessionAfterSuccess records continuity for the transcript the client will
+// hold after this successful assistant turn: inbound messages with the final
+// user turn answered, i.e. prefix+assistant — callers pass the full chain that
+// ends with the assistant content just produced.
+func sealSessionAfterSuccess(apiKeyID, modelID, convID, contID string, chainEndingWithAssistant []ClaudeMessage) {
+ if convID == "" || contID == "" || len(chainEndingWithAssistant) == 0 {
+ return
+ }
+ if !claudeEndsWithAssistant(chainEndingWithAssistant) {
+ return
+ }
+ fp := chainFingerprintClaude(chainEndingWithAssistant)
+ globalSessionStore.put(apiKeyID, modelID, fp, convID, contID)
+}
+
+func (s *sessionStore) put(apiKeyID, modelID, chainFP, convID, contID string) {
+ key := sessionKey(apiKeyID, modelID, chainFP)
+ now := time.Now()
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.purgeLocked(now)
+
+ if existing, ok := s.byKey[key]; ok {
+ if existing.Ambiguous || existing.ConversationID != convID || existing.ContinuationID != contID {
+ s.byKey[key] = sessionIDs{Ambiguous: true, StoredAt: now}
+ return
+ }
+ existing.StoredAt = now
+ s.byKey[key] = existing
+ return
+ }
+
+ if len(s.byKey) >= sessionStoreMax {
+ s.dropOldestLocked(sessionStoreMax / 4)
+ }
+ s.byKey[key] = sessionIDs{
+ ConversationID: convID,
+ ContinuationID: contID,
+ StoredAt: now,
+ }
+}
+
+func (s *sessionStore) lookup(apiKeyID, modelID, chainFP string) *sessionIDs {
+ key := sessionKey(apiKeyID, modelID, chainFP)
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.purgeLocked(time.Now())
+ e, ok := s.byKey[key]
+ if !ok || e.Ambiguous || e.ConversationID == "" || e.ContinuationID == "" {
+ return nil
+ }
+ cp := e
+ return &cp
+}
+
+func (s *sessionStore) purgeLocked(now time.Time) {
+ for k, e := range s.byKey {
+ if now.Sub(e.StoredAt) > sessionStoreTTL {
+ delete(s.byKey, k)
+ }
+ }
+}
+
+func (s *sessionStore) dropOldestLocked(n int) {
+ type pair struct {
+ key string
+ storedAt time.Time
+ }
+ all := make([]pair, 0, len(s.byKey))
+ for k, e := range s.byKey {
+ all = append(all, pair{k, e.StoredAt})
+ }
+ sort.Slice(all, func(i, j int) bool { return all[i].storedAt.Before(all[j].storedAt) })
+ if n > len(all) {
+ n = len(all)
+ }
+ for _, p := range all[:n] {
+ delete(s.byKey, p.key)
+ }
+}
+
+// sealClaudeSession records continuity after a successful Claude turn.
+//
+// The chain MUST be fingerprinted from the client transcript, never from the
+// wire payload: ClaudeToKiro rewrites history (system priming pair, flattened
+// tool turns, truncation), so a payload-derived chain can never match the
+// lookup key that resolveSessionIDs computes from the next inbound request.
+//
+// assistantText must be the VISIBLE assistant text the client stores — not the
+// integrity-check builder that may have tool name/input padding glued on.
+// thinkingText + format must mirror how the handler exposed reasoning to the
+// client, otherwise the next-turn lookup (which fingerprints thinking blocks
+// and folded text the same way the client replays them) will miss.
+func sealClaudeSession(apiKeyID, model, convID, contID string, inbound []ClaudeMessage, assistantText, thinkingText string, toolUses []KiroToolUse, format string, omitDisplay bool) {
+ if convID == "" || contID == "" || len(inbound) == 0 {
+ return
+ }
+ chain := make([]ClaudeMessage, 0, len(inbound)+1)
+ chain = append(chain, inbound...)
+ chain = append(chain, claudeAssistantTurn(assistantText, thinkingText, toolUses, format, omitDisplay))
+ sealSessionAfterSuccess(apiKeyID, MapModel(model), convID, contID, chain)
+}
+
+// sealClaudeSessionFromContent records continuity from the exact assistant
+// content the client will store and replay. Mixed web_search responses include
+// server_tool_use / web_search_tool_result blocks that are not raw Kiro tool
+// uses, so fingerprinting must use that client-visible shape.
+func sealClaudeSessionFromContent(apiKeyID, model, convID, contID string, inbound []ClaudeMessage, content []map[string]interface{}) {
+ if convID == "" || contID == "" || len(inbound) == 0 || len(content) == 0 {
+ return
+ }
+ blocks := make([]interface{}, 0, len(content))
+ for _, block := range content {
+ blocks = append(blocks, block)
+ }
+ chain := make([]ClaudeMessage, 0, len(inbound)+1)
+ chain = append(chain, inbound...)
+ chain = append(chain, ClaudeMessage{Role: "assistant", Content: blocks})
+ sealSessionAfterSuccess(apiKeyID, MapModel(model), convID, contID, chain)
+}
+
+// claudeAssistantTurn renders this turn's answer in the same block shape the
+// client will replay it as, so seal and lookup fingerprints agree.
+func claudeAssistantTurn(text, thinkingText string, toolUses []KiroToolUse, format string, omitDisplay bool) ClaudeMessage {
+ includeThinkingBlock := thinkingText != ""
+ // Formats that fold reasoning into the text stream: seal the folded text the
+ // client actually saw, with no separate thinking block.
+ switch format {
+ case "think":
+ if thinkingText != "" {
+ text = "" + thinkingText + "" + text
+ thinkingText = ""
+ includeThinkingBlock = false
+ }
+ case "reasoning_content":
+ if thinkingText != "" {
+ text = thinkingText + text
+ thinkingText = ""
+ includeThinkingBlock = false
+ }
+ default:
+ // Native thinking with omitted display emits an explicit empty shell.
+ if omitDisplay && includeThinkingBlock {
+ thinkingText = ""
+ }
+ }
+
+ needBlocks := includeThinkingBlock || len(toolUses) > 0
+ if !needBlocks {
+ return ClaudeMessage{Role: "assistant", Content: text}
+ }
+ blocks := make([]interface{}, 0, len(toolUses)+2)
+ if includeThinkingBlock {
+ blocks = append(blocks, map[string]interface{}{"type": "thinking", "thinking": thinkingText})
+ }
+ if text != "" {
+ blocks = append(blocks, map[string]interface{}{"type": "text", "text": text})
+ }
+ for _, tu := range toolUses {
+ input := tu.Input
+ if input == nil {
+ input = map[string]interface{}{}
+ }
+ blocks = append(blocks, map[string]interface{}{
+ "type": "tool_use",
+ "id": tu.ToolUseID,
+ "name": tu.Name,
+ "input": input,
+ })
+ }
+ return ClaudeMessage{Role: "assistant", Content: blocks}
+}
diff --git a/proxy/session_store_test.go b/proxy/session_store_test.go
new file mode 100644
index 00000000..ae0fb041
--- /dev/null
+++ b/proxy/session_store_test.go
@@ -0,0 +1,375 @@
+package proxy
+
+import (
+ "testing"
+)
+
+func TestResolveSessionIDsMintsFreshOnFirstTurn(t *testing.T) {
+ msgs := []ClaudeMessage{
+ {Role: "user", Content: "hello first turn only"},
+ }
+ c1, k1, reused := resolveSessionIDs("ak", "claude-sonnet-4.5", msgs)
+ if reused {
+ t.Fatal("first turn must not reuse")
+ }
+ if c1 == "" || k1 == "" {
+ t.Fatal("must mint both ids")
+ }
+ c2, k2, reused2 := resolveSessionIDs("ak", "claude-sonnet-4.5", msgs)
+ if reused2 {
+ t.Fatal("unsealed first turn must not reuse")
+ }
+ if c1 == c2 || k1 == k2 {
+ t.Fatalf("unsealed first turns must mint distinct ids: %s/%s vs %s/%s", c1, k1, c2, k2)
+ }
+}
+
+func TestSessionSealThenResolveReusesExactChain(t *testing.T) {
+ model := "claude-sonnet-4.5"
+ api := "ak-session-1"
+ user1 := "session-continuity-opener-unique-111"
+ asst1 := "session-continuity-answer-unique-111"
+
+ // After turn 1 success, client holds user1+asst1; we seal that chain.
+ chain1 := []ClaudeMessage{
+ {Role: "user", Content: user1},
+ {Role: "assistant", Content: asst1},
+ }
+ conv, cont := "conv-fixed-aaa", "cont-fixed-bbb"
+ sealSessionAfterSuccess(api, model, conv, cont, chain1)
+
+ // Turn 2: client sends user1+asst1+user2. Prefix before current user = chain1.
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: user1},
+ {Role: "assistant", Content: asst1},
+ {Role: "user", Content: "follow up please"},
+ }
+ gotConv, gotCont, reused := resolveSessionIDs(api, model, turn2)
+ if !reused {
+ t.Fatal("expected reuse on exact continuation chain")
+ }
+ if gotConv != conv || gotCont != cont {
+ t.Fatalf("reused ids = %s/%s, want %s/%s", gotConv, gotCont, conv, cont)
+ }
+}
+
+func TestSessionDoesNotReuseAcrossAPIKeysOrModels(t *testing.T) {
+ model := "claude-sonnet-4.5"
+ chain := []ClaudeMessage{
+ {Role: "user", Content: "iso-user"},
+ {Role: "assistant", Content: "iso-asst"},
+ }
+ sealSessionAfterSuccess("key-A", model, "c-A", "k-A", chain)
+
+ turn := []ClaudeMessage{
+ {Role: "user", Content: "iso-user"},
+ {Role: "assistant", Content: "iso-asst"},
+ {Role: "user", Content: "next"},
+ }
+ _, _, reused := resolveSessionIDs("key-B", model, turn)
+ if reused {
+ t.Fatal("must not reuse across API keys")
+ }
+ _, _, reused = resolveSessionIDs("key-A", "other-model", turn)
+ if reused {
+ t.Fatal("must not reuse across models")
+ }
+}
+
+func TestSessionDoesNotReuseWhenAssistantTextDiverges(t *testing.T) {
+ model := "m"
+ api := "ak"
+ sealSessionAfterSuccess(api, model, "c1", "k1", []ClaudeMessage{
+ {Role: "user", Content: "q"},
+ {Role: "assistant", Content: "answer-A"},
+ })
+ // Client rewrote history assistant text — chain fingerprint changes.
+ _, _, reused := resolveSessionIDs(api, model, []ClaudeMessage{
+ {Role: "user", Content: "q"},
+ {Role: "assistant", Content: "answer-TAMPERED"},
+ {Role: "user", Content: "next"},
+ })
+ if reused {
+ t.Fatal("tampered assistant history must not reuse session ids")
+ }
+}
+
+func TestSessionDoesNotReuseWhenWhitespaceDiverges(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ sealSessionAfterSuccess("key-space", "model", "conv-space", "cont-space", []ClaudeMessage{
+ {Role: "user", Content: "a b"},
+ {Role: "assistant", Content: "answer"},
+ })
+ _, _, reused := resolveSessionIDs("key-space", "model", []ClaudeMessage{
+ {Role: "user", Content: "a b"},
+ {Role: "assistant", Content: "answer"},
+ {Role: "user", Content: "next"},
+ })
+ if reused {
+ t.Fatal("different whitespace is different client history and must not reuse")
+ }
+}
+
+func TestClaudeToKiroReusesSessionIDsOnContinuation(t *testing.T) {
+ model := "claude-sonnet-4.5"
+ api := ""
+ user1 := "ctk-session-opener-zzz"
+ asst1 := "ctk-session-answer-zzz"
+
+ // Seal as handler would after turn 1.
+ sealSessionAfterSuccess(api, MapModel(model), "conv-ctk-1", "cont-ctk-1", []ClaudeMessage{
+ {Role: "user", Content: user1},
+ {Role: "assistant", Content: asst1},
+ })
+
+ req := &ClaudeRequest{
+ Model: model,
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: user1},
+ {Role: "assistant", Content: asst1},
+ {Role: "user", Content: "second question"},
+ },
+ }
+ // Handler-side resolve (ClaudeToKiro stays pure).
+ gotConv, gotCont, reused := resolveSessionIDs(api, MapModel(model), req.Messages)
+ if !reused || gotConv != "conv-ctk-1" || gotCont != "cont-ctk-1" {
+ t.Fatalf("resolve = %s/%s reused=%v", gotConv, gotCont, reused)
+ }
+ // Brand-new chat must not reuse.
+ _, _, reusedNew := resolveSessionIDs(api, MapModel(model), []ClaudeMessage{{Role: "user", Content: "totally new chat opener"}})
+ if reusedNew {
+ t.Fatal("new chat must not reuse")
+ }
+}
+
+// Regression: the seal must fingerprint the CLIENT transcript, not the wire
+// payload. ClaudeToKiro prepends a system-priming user/assistant pair whenever
+// a system prompt exists (always true for Claude Code / Pi), so a
+// payload-derived chain could never match the next request's lookup key and
+// continuity silently never happened in production.
+func TestSealClaudeSessionMatchesLookupWithSystemPrompt(t *testing.T) {
+ globalSessionStore = newSessionStore()
+
+ const api = "key-sys"
+ const model = "claude-sonnet-4.5"
+ system := "you are a coding agent with a long system prompt"
+
+ turn1 := []ClaudeMessage{{Role: "user", Content: "first question"}}
+ req1 := &ClaudeRequest{Model: model, System: system, Messages: turn1}
+ payload1 := ClaudeToKiro(req1, false)
+
+ conv1, cont1, reused := resolveSessionIDs(api, MapModel(model), turn1)
+ if reused {
+ t.Fatalf("first turn must not reuse")
+ }
+ // Sanity: the payload really does carry the priming pair that used to break this.
+ if len(payload1.ConversationState.History) == 0 {
+ t.Fatalf("expected system priming history in payload")
+ }
+
+ // Seal exactly as the handler does, from the inbound transcript.
+ sealClaudeSession(api, model, conv1, cont1, turn1, "first answer", "", nil, "thinking", false)
+
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: "first question"},
+ {Role: "assistant", Content: "first answer"},
+ {Role: "user", Content: "second question"},
+ }
+ conv2, cont2, reused2 := resolveSessionIDs(api, MapModel(model), turn2)
+ if !reused2 {
+ t.Fatalf("continuation must reuse after a real seal (system prompt present)")
+ }
+ if conv2 != conv1 || cont2 != cont1 {
+ t.Fatalf("ids not reused: conv %q vs %q, cont %q vs %q", conv2, conv1, cont2, cont1)
+ }
+}
+
+// Regression: a tool loop must keep the chain alive. The seal renders this
+// turn's tool_use blocks the same way the client replays them.
+func TestSealClaudeSessionSurvivesToolLoop(t *testing.T) {
+ globalSessionStore = newSessionStore()
+
+ const api = "key-tool"
+ const model = "claude-sonnet-4.5"
+
+ turn1 := []ClaudeMessage{{Role: "user", Content: "read the file"}}
+ conv, cont, _ := resolveSessionIDs(api, MapModel(model), turn1)
+ sealClaudeSession(api, model, conv, cont, turn1, "let me look", "",
+ []KiroToolUse{{ToolUseID: "tu_1", Name: "read_file", Input: map[string]interface{}{"path": "a.go"}}}, "thinking", false)
+
+ // Client replays assistant text + tool_use, then sends the tool_result.
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: "read the file"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": "let me look"},
+ map[string]interface{}{"type": "tool_use", "id": "tu_1", "name": "read_file", "input": map[string]interface{}{"path": "a.go"}},
+ }},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "tu_1", "content": "file body"},
+ }},
+ }
+ got, gotCont, reused := resolveSessionIDs(api, MapModel(model), turn2)
+ if !reused {
+ t.Fatalf("tool-result turn must continue the chain")
+ }
+ if got != conv || gotCont != cont {
+ t.Fatalf("tool loop lost ids: %q/%q want %q/%q", got, gotCont, conv, cont)
+ }
+}
+
+// Regression: with default "thinking" format the client stores a thinking
+// content block. Seal must include that block or the next-turn lookup misses.
+func TestSealClaudeSessionMatchesLookupWithThinkingBlock(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ const api = "key-think"
+ const model = "claude-sonnet-4.5"
+
+ turn1 := []ClaudeMessage{{Role: "user", Content: "first question"}}
+ conv, cont, _ := resolveSessionIDs(api, MapModel(model), turn1)
+ sealClaudeSession(api, model, conv, cont, turn1, "first answer", "let me reason", nil, "thinking", false)
+
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: "first question"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "thinking", "thinking": "let me reason"},
+ map[string]interface{}{"type": "text", "text": "first answer"},
+ }},
+ {Role: "user", Content: "second question"},
+ }
+ got, gotCont, reused := resolveSessionIDs(api, MapModel(model), turn2)
+ if !reused {
+ t.Fatal("thinking-format continuation must reuse")
+ }
+ if got != conv || gotCont != cont {
+ t.Fatalf("ids lost under thinking format: %q/%q want %q/%q", got, gotCont, conv, cont)
+ }
+}
+
+// Regression: "think" format folds reasoning into the text the client stores.
+// Seal must fingerprint the folded text, not the bare answer.
+func TestSealClaudeSessionMatchesLookupWithThinkTags(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ const api = "key-think-tags"
+ const model = "claude-sonnet-4.5"
+
+ turn1 := []ClaudeMessage{{Role: "user", Content: "first question"}}
+ conv, cont, _ := resolveSessionIDs(api, MapModel(model), turn1)
+ sealClaudeSession(api, model, conv, cont, turn1, "first answer", "let me reason", nil, "think", false)
+
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: "first question"},
+ {Role: "assistant", Content: "let me reasonfirst answer"},
+ {Role: "user", Content: "second question"},
+ }
+ got, gotCont, reused := resolveSessionIDs(api, MapModel(model), turn2)
+ if !reused {
+ t.Fatal("think-tag format continuation must reuse")
+ }
+ if got != conv || gotCont != cont {
+ t.Fatalf("ids lost under think tags: %q/%q want %q/%q", got, gotCont, conv, cont)
+ }
+}
+
+func TestSealClaudeSessionMatchesOmittedThinkingBlock(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ inbound := []ClaudeMessage{{Role: "user", Content: "question"}}
+ sealClaudeSession("key-omit", "model", "conv-omit", "cont-omit", inbound, "answer", "hidden", nil, "thinking", true)
+ _, _, reused := resolveSessionIDs("key-omit", MapModel("model"), []ClaudeMessage{
+ {Role: "user", Content: "question"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "thinking", "thinking": ""},
+ map[string]interface{}{"type": "text", "text": "answer"},
+ }},
+ {Role: "user", Content: "next"},
+ })
+ if !reused {
+ t.Fatal("empty client-visible thinking shell must remain in sealed shape")
+ }
+}
+
+// Regression: seal text must be the VISIBLE answer, never the integrity-check
+// padding that used to glue tool name+input onto rawContentBuilder.
+func TestSealClaudeSessionDoesNotFingerprintToolPadding(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ const api = "key-pad"
+ const model = "claude-sonnet-4.5"
+
+ turn1 := []ClaudeMessage{{Role: "user", Content: "read the file"}}
+ conv, cont, _ := resolveSessionIDs(api, MapModel(model), turn1)
+
+ // What the client actually stores.
+ clean := "let me look"
+ // What rawContentBuilder used to contain after OnToolUse padding.
+ polluted := clean + `read_file{"path":"a.go"}`
+
+ // Seal the clean shape (as the fixed handler now does).
+ sealClaudeSession(api, model, conv, cont, turn1, clean, "",
+ []KiroToolUse{{ToolUseID: "tu_1", Name: "read_file"}}, "thinking", false)
+
+ // Next request replays clean text + tool_use — must hit.
+ turn2 := []ClaudeMessage{
+ {Role: "user", Content: "read the file"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": clean},
+ map[string]interface{}{"type": "tool_use", "id": "tu_1", "name": "read_file"},
+ }},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "tu_1", "content": "body"},
+ }},
+ }
+ if _, _, reused := resolveSessionIDs(api, MapModel(model), turn2); !reused {
+ t.Fatal("clean seal must match clean client replay")
+ }
+
+ // A polluted seal must NOT match the clean client replay (proves the bug
+ // class: if we still sealed polluted text, continuity would die here).
+ globalSessionStore = newSessionStore()
+ conv2, cont2, _ := resolveSessionIDs(api, MapModel(model), turn1)
+ sealClaudeSession(api, model, conv2, cont2, turn1, polluted, "",
+ []KiroToolUse{{ToolUseID: "tu_1", Name: "read_file"}}, "thinking", false)
+ if _, _, reused := resolveSessionIDs(api, MapModel(model), turn2); reused {
+ t.Fatal("polluted seal must NOT match clean client replay — that was the production bug")
+ }
+}
+
+func TestSessionStoreCollisionFailsClosed(t *testing.T) {
+ s := newSessionStore()
+ s.put("key", "model", "same-chain", "conv-a", "cont-a")
+ s.put("key", "model", "same-chain", "conv-b", "cont-b")
+ if got := s.lookup("key", "model", "same-chain"); got != nil {
+ t.Fatalf("ambiguous transcript must not reuse either session: %+v", got)
+ }
+}
+
+func TestSessionFingerprintIncludesToolInput(t *testing.T) {
+ chain := func(path string) []ClaudeMessage {
+ return []ClaudeMessage{
+ {Role: "user", Content: "read a file"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{
+ "type": "tool_use", "id": "tool-1", "name": "read_file",
+ "input": map[string]interface{}{"path": path},
+ },
+ }},
+ }
+ }
+ if chainFingerprintClaude(chain("a.go")) == chainFingerprintClaude(chain("b.go")) {
+ t.Fatal("different tool inputs must produce different chain fingerprints")
+ }
+}
+
+func TestSessionFingerprintIncludesStructuredToolResult(t *testing.T) {
+ chain := func(result string) []ClaudeMessage {
+ return []ClaudeMessage{
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{
+ "type": "tool_result", "tool_use_id": "tool-1",
+ "content": []interface{}{map[string]interface{}{"type": "text", "text": result}},
+ },
+ }},
+ }
+ }
+ if chainFingerprintClaude(chain("first")) == chainFingerprintClaude(chain("second")) {
+ t.Fatal("different structured tool results must produce different chain fingerprints")
+ }
+}
From f4982990e26852f51e361c3be672dc8fda9adc57 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:24 +0800
Subject: [PATCH 5/6] feat(proxy): map Claude documents and centralize
stop-reason classes
Translate document/file blocks into Kiro documents (including raw text),
validate MIME/magic/name/limit rules, and classify upstream stop reasons
once for Claude, OpenAI, and Responses mappings.
---
proxy/documents_test.go | 260 +++++++++++++++++++++++++++
proxy/translator.go | 371 ++++++++++++++++++++++++++++++++++++---
proxy/translator_test.go | 41 +++--
3 files changed, 631 insertions(+), 41 deletions(-)
create mode 100644 proxy/documents_test.go
diff --git a/proxy/documents_test.go b/proxy/documents_test.go
new file mode 100644
index 00000000..04cf2a30
--- /dev/null
+++ b/proxy/documents_test.go
@@ -0,0 +1,260 @@
+package proxy
+
+import (
+ "encoding/base64"
+ "strings"
+ "testing"
+)
+
+// Seam: extractClaudeUserContent must split images vs documents.
+// PDF/base64 document blocks become KiroDocument; image/* stay images.
+func TestExtractClaudeUserContentDocumentsVsImages(t *testing.T) {
+ pdfB64 := base64.StdEncoding.EncodeToString([]byte("%PDF-1.4 fake"))
+ imgB64 := base64.StdEncoding.EncodeToString([]byte{0x89, 0x50, 0x4e, 0x47}) // PNG magic-ish
+
+ content := []interface{}{
+ map[string]interface{}{"type": "text", "text": "see attached"},
+ map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "type": "base64",
+ "media_type": "application/pdf",
+ "data": pdfB64,
+ },
+ "title": "report.pdf",
+ },
+ map[string]interface{}{
+ "type": "image",
+ "source": map[string]interface{}{
+ "type": "base64",
+ "media_type": "image/png",
+ "data": imgB64,
+ },
+ },
+ }
+
+ text, images, docs, tools, errMsg := extractClaudeUserContent(content)
+ if errMsg != "" {
+ t.Fatalf("unexpected document error: %s", errMsg)
+ }
+ if text != "see attached" {
+ t.Fatalf("text=%q", text)
+ }
+ if len(tools) != 0 {
+ t.Fatalf("tools=%d", len(tools))
+ }
+ if len(images) != 1 || images[0].Format != "png" {
+ t.Fatalf("images=%+v", images)
+ }
+ if len(docs) != 1 {
+ t.Fatalf("docs=%d %+v", len(docs), docs)
+ }
+ if docs[0].Format != "pdf" {
+ t.Fatalf("doc format=%q", docs[0].Format)
+ }
+ if docs[0].Name == "" {
+ t.Fatal("doc name empty")
+ }
+ if docs[0].Source.Bytes != pdfB64 {
+ t.Fatal("doc bytes mismatch")
+ }
+}
+
+// Seam: ClaudeToKiro must put documents on current user message.
+func TestClaudeToKiroWiresDocumentsOnCurrentMessage(t *testing.T) {
+ pdfB64 := base64.StdEncoding.EncodeToString([]byte("%PDF-1.4 x"))
+ req := &ClaudeRequest{
+ Model: "claude-sonnet-4.5",
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": "read this"},
+ map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "type": "base64",
+ "media_type": "application/pdf",
+ "data": pdfB64,
+ },
+ "title": "a.pdf",
+ },
+ }},
+ },
+ }
+ payload := ClaudeToKiro(req, false)
+ docs := payload.ConversationState.CurrentMessage.UserInputMessage.Documents
+ if len(docs) != 1 || docs[0].Format != "pdf" {
+ t.Fatalf("current documents=%+v", docs)
+ }
+ if payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" {
+ t.Fatal("content should remain non-empty with documents")
+ }
+}
+
+// Unsupported document mime types are skipped, not turned into images.
+func TestExtractClaudeUserContentSkipsUnknownDocumentMime(t *testing.T) {
+ b64 := base64.StdEncoding.EncodeToString([]byte("zzz"))
+ content := []interface{}{
+ map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "media_type": "application/x-unknown",
+ "data": b64,
+ },
+ "title": "x.bin",
+ },
+ }
+ _, images, docs, _, errMsg := extractClaudeUserContent(content)
+ if errMsg != "" {
+ t.Fatalf("unknown mime must skip without error: %s", errMsg)
+ }
+ if len(docs) != 0 || len(images) != 0 {
+ t.Fatalf("unknown mime must skip: docs=%+v images=%+v", docs, images)
+ }
+}
+
+func claudeDocumentBlock(name, mediaType string, data []byte) map[string]interface{} {
+ return map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "type": "base64",
+ "media_type": mediaType,
+ "data": base64.StdEncoding.EncodeToString(data),
+ },
+ "title": name,
+ }
+}
+
+func TestClaudeDocumentValidationRejectsBadBase64AndMagic(t *testing.T) {
+ for _, tc := range []struct {
+ name string
+ block map[string]interface{}
+ wantError string
+ }{
+ {
+ name: "bad base64",
+ block: map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "media_type": "application/pdf",
+ "data": "AQI",
+ },
+ "title": "bad.pdf",
+ },
+ wantError: "base64",
+ },
+ {
+ name: "bad magic",
+ block: claudeDocumentBlock("fake.pdf", "application/pdf", []byte("not a pdf")),
+ wantError: "does not match",
+ },
+ } {
+ t.Run(tc.name, func(t *testing.T) {
+ req := &ClaudeRequest{Messages: []ClaudeMessage{{Role: "user", Content: []interface{}{tc.block}}}}
+ if got := validateClaudeRequestShape(req); !strings.Contains(strings.ToLower(got), tc.wantError) {
+ t.Fatalf("validation error = %q, want %q", got, tc.wantError)
+ }
+ })
+ }
+}
+
+func TestClaudeDocumentsSanitizeAndDeduplicateNames(t *testing.T) {
+ pdf := []byte("%PDF-1.7 valid")
+ content := []interface{}{
+ claudeDocumentBlock("Quarter:Report.pdf", "application/pdf", pdf),
+ claudeDocumentBlock("Quarter?Report.pdf", "application/pdf", pdf),
+ }
+ _, _, docs, _, errMsg := extractClaudeUserContent(content)
+ if errMsg != "" {
+ t.Fatalf("unexpected document error: %s", errMsg)
+ }
+ if len(docs) != 1 {
+ t.Fatalf("normalized duplicate documents = %+v, want one", docs)
+ }
+ if docs[0].Name != "Quarter-Report" {
+ t.Fatalf("sanitized document name = %q", docs[0].Name)
+ }
+}
+
+func TestClaudeDocumentsAcceptOfficeMagicBytes(t *testing.T) {
+ content := []interface{}{
+ claudeDocumentBlock("report.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", []byte{0x50, 0x4b, 0x03, 0x04, 0x01}),
+ claudeDocumentBlock("legacy.xls", "application/vnd.ms-excel", []byte{0xd0, 0xcf, 0x11, 0xe0, 0x01}),
+ }
+ _, _, docs, _, errMsg := extractClaudeUserContent(content)
+ if errMsg != "" {
+ t.Fatalf("unexpected document error: %s", errMsg)
+ }
+ if len(docs) != 2 || docs[0].Format != "docx" || docs[1].Format != "xls" {
+ t.Fatalf("Office documents = %+v", docs)
+ }
+}
+
+func TestClaudeDocumentLimitAcrossConversation(t *testing.T) {
+ pdf := []byte("%PDF-1.7 valid")
+ messages := make([]ClaudeMessage, 0, 6)
+ for i := 0; i < 6; i++ {
+ messages = append(messages, ClaudeMessage{Role: "user", Content: []interface{}{
+ claudeDocumentBlock(string(rune('a'+i))+".pdf", "application/pdf", pdf),
+ }})
+ }
+ got := validateClaudeRequestShape(&ClaudeRequest{Messages: messages})
+ if !strings.Contains(strings.ToLower(got), "maximum is 5") {
+ t.Fatalf("document limit validation = %q", got)
+ }
+}
+
+func TestClaudeToKiroDropsDuplicateDocumentsAcrossConversation(t *testing.T) {
+ pdf := []byte("%PDF-1.7 valid")
+ req := &ClaudeRequest{Model: "claude-sonnet-4.5", Messages: []ClaudeMessage{
+ {Role: "user", Content: []interface{}{claudeDocumentBlock("same.pdf", "application/pdf", pdf)}},
+ {Role: "assistant", Content: "seen"},
+ {Role: "user", Content: []interface{}{claudeDocumentBlock("same.pdf", "application/pdf", pdf)}},
+ }}
+ payload := ClaudeToKiro(req, false)
+ total := len(payload.ConversationState.CurrentMessage.UserInputMessage.Documents)
+ for _, message := range payload.ConversationState.History {
+ if message.UserInputMessage != nil {
+ total += len(message.UserInputMessage.Documents)
+ }
+ }
+ if total != 1 {
+ t.Fatalf("duplicate documents across conversation = %d, want one", total)
+ }
+}
+
+func TestClaudeRequestWithOnlyDocumentHasUserContext(t *testing.T) {
+ req := &ClaudeRequest{Messages: []ClaudeMessage{{
+ Role: "user",
+ Content: []interface{}{
+ claudeDocumentBlock("only.pdf", "application/pdf", []byte("%PDF-1.7 valid")),
+ },
+ }}}
+ if got := validateClaudeRequestShape(req); got != "" {
+ t.Fatalf("document-only request rejected: %s", got)
+ }
+}
+
+func TestClaudeRawTextDocumentMapsToKiroTXT(t *testing.T) {
+ block := map[string]interface{}{
+ "type": "document",
+ "source": map[string]interface{}{
+ "type": "text",
+ "media_type": "text/plain",
+ "data": "hello world\nexact text",
+ },
+ "title": "notes.txt",
+ }
+ req := &ClaudeRequest{Model: "claude-sonnet-4.5", Messages: []ClaudeMessage{{Role: "user", Content: []interface{}{block}}}}
+ if got := validateClaudeRequestShape(req); got != "" {
+ t.Fatalf("valid raw-text document rejected: %s", got)
+ }
+ payload := ClaudeToKiro(req, false)
+ documents := payload.ConversationState.CurrentMessage.UserInputMessage.Documents
+ if len(documents) != 1 || documents[0].Format != "txt" {
+ t.Fatalf("raw-text documents = %+v", documents)
+ }
+ raw, err := base64.StdEncoding.DecodeString(documents[0].Source.Bytes)
+ if err != nil || string(raw) != "hello world\nexact text" {
+ t.Fatalf("raw-text bytes = %q, err=%v", raw, err)
+ }
+}
diff --git a/proxy/translator.go b/proxy/translator.go
index 0f77ddd9..7969a546 100644
--- a/proxy/translator.go
+++ b/proxy/translator.go
@@ -1,6 +1,7 @@
package proxy
import (
+ "bytes"
"encoding/base64"
"encoding/json"
"fmt"
@@ -212,18 +213,25 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload {
history := make([]KiroHistoryMessage, 0)
var currentContent string
var currentImages []KiroImage
+ var currentDocuments []KiroDocument
var currentToolResults []KiroToolResult
for i, msg := range req.Messages {
isLast := i == len(req.Messages)-1
if msg.Role == "user" {
- content, images, toolResults := extractClaudeUserContent(msg.Content)
- content = normalizeUserContent(content, len(images) > 0)
+ content, images, documents, toolResults, docErr := extractClaudeUserContent(msg.Content)
+ if docErr != "" {
+ // Request-shape validation should have rejected this already.
+ // Keep translation defensive for direct callers.
+ continue
+ }
+ content = normalizeUserContent(content, len(images) > 0 || len(documents) > 0)
if isLast {
currentContent = content
currentImages = images
+ currentDocuments = documents
currentToolResults = toolResults
} else {
userMsg := KiroUserInputMessage{
@@ -234,6 +242,9 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload {
if len(images) > 0 {
userMsg.Images = images
}
+ if len(documents) > 0 {
+ userMsg.Documents = documents
+ }
if len(toolResults) > 0 {
userMsg.UserInputMessageContext = &UserInputMessageContext{
ToolResults: toolResults,
@@ -245,12 +256,11 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload {
}
} else if msg.Role == "assistant" {
content, toolUses := extractClaudeAssistantContent(msg.Content)
- history = append(history, KiroHistoryMessage{
- AssistantResponseMessage: &KiroAssistantResponseMessage{
- Content: content,
- ToolUses: toolUses,
- },
- })
+ assistant := &KiroAssistantResponseMessage{
+ Content: content,
+ ToolUses: toolUses,
+ }
+ history = append(history, KiroHistoryMessage{AssistantResponseMessage: assistant})
}
}
@@ -289,12 +299,13 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload {
} else {
history = sanitizeKiroHistory(history, nil)
}
+ currentDocuments = deduplicateConversationDocuments(history, currentDocuments)
// 构建最终内容
finalContent := ""
if currentContent != "" {
finalContent = currentContent
- } else if len(currentImages) > 0 {
+ } else if len(currentImages) > 0 || len(currentDocuments) > 0 {
finalContent = normalizeUserContent("", true)
} else if len(currentToolResults) > 0 {
finalContent = buildToolResultsContinuation(currentToolResults)
@@ -310,13 +321,16 @@ func ClaudeToKiro(req *ClaudeRequest, thinking bool) *KiroPayload {
payload.ToolNameMap = toolNameMap
payload.ConversationState.ChatTriggerType = "MANUAL"
payload.ConversationState.AgentTaskType = "vibe"
+ // Provisional IDs; handler overwrites via resolveSessionIDs(apiKey, ...) so
+ // multi-tenant isolation is not lost inside this pure translator.
payload.ConversationState.AgentContinuationId = uuid.New().String()
- payload.ConversationState.ConversationID = buildConversationID(modelID, systemPrompt, firstClaudeConversationAnchor(req.Messages))
+ payload.ConversationState.ConversationID = uuid.New().String()
payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
- Content: finalContent,
- ModelID: modelID,
- Origin: origin,
- Images: currentImages,
+ Content: finalContent,
+ ModelID: modelID,
+ Origin: origin,
+ Images: currentImages,
+ Documents: currentDocuments,
}
// Only attach structured tool results when they answer the last history
@@ -622,13 +636,15 @@ func extractSystemPrompt(system interface{}) string {
return ""
}
-func extractClaudeUserContent(content interface{}) (string, []KiroImage, []KiroToolResult) {
+func extractClaudeUserContent(content interface{}) (string, []KiroImage, []KiroDocument, []KiroToolResult, string) {
var text string
var images []KiroImage
+ var documents []KiroDocument
var toolResults []KiroToolResult
+ documentNames := make(map[string]bool)
if s, ok := content.(string); ok {
- return s, nil, nil
+ return s, nil, nil, nil, ""
}
// Accept both JSON-decoded []interface{} and in-memory []map[string]interface{}
@@ -644,6 +660,15 @@ func extractClaudeUserContent(content interface{}) (string, []KiroImage, []KiroT
if img := extractImageFromClaudeBlock(block); img != nil {
images = append(images, *img)
}
+ case "document", "file", "input_document":
+ doc, validationError := parseClaudeDocumentBlock(block)
+ if validationError != "" {
+ return "", nil, nil, nil, validationError
+ }
+ if doc != nil && !documentNames[doc.Name] {
+ documentNames[doc.Name] = true
+ documents = append(documents, *doc)
+ }
case "tool_result":
toolUseID, _ := block["tool_use_id"].(string)
resultContent, resultImages := extractToolResultContent(block["content"])
@@ -661,7 +686,7 @@ func extractClaudeUserContent(content interface{}) (string, []KiroImage, []KiroT
}
}
- return text, images, toolResults
+ return text, images, documents, toolResults, ""
}
// contentBlocksAsMaps normalizes Claude content arrays for extraction.
@@ -685,6 +710,171 @@ func contentBlocksAsMaps(content interface{}) []map[string]interface{} {
}
}
+const maxDocumentsPerConversation = 5
+
+// documentMIMEToFormat mirrors the official client's supported document formats.
+var documentMIMEToFormat = map[string]string{
+ "application/pdf": "pdf",
+ "text/csv": "csv",
+ "application/msword": "doc",
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
+ "application/vnd.ms-excel": "xls",
+ "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
+ "text/html": "html",
+ "text/plain": "txt",
+ "text/markdown": "md",
+}
+
+var documentMagicBytes = map[string][]byte{
+ "application/pdf": {0x25, 0x50, 0x44, 0x46},
+ "application/msword": {0xd0, 0xcf, 0x11, 0xe0},
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document": {0x50, 0x4b, 0x03, 0x04},
+ "application/vnd.ms-excel": {0xd0, 0xcf, 0x11, 0xe0},
+ "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": {0x50, 0x4b, 0x03, 0x04},
+}
+
+var (
+ documentBase64Pattern = regexp.MustCompile(`^[A-Za-z0-9+/]*={0,2}$`)
+ documentNameInvalid = regexp.MustCompile(`[^a-zA-Z0-9\s\-()\[\]]`)
+ documentRepeatedHyphens = regexp.MustCompile(`-{2,}`)
+)
+
+func documentBlockName(block map[string]interface{}, format string) string {
+ name, _ := block["title"].(string)
+ if name == "" {
+ name, _ = block["name"].(string)
+ }
+ if name == "" {
+ name = "document." + format
+ }
+ return name
+}
+
+func sanitizeDocumentName(name string) string {
+ if dot := strings.LastIndex(name, "."); dot >= 0 {
+ name = name[:dot]
+ }
+ name = documentNameInvalid.ReplaceAllString(name, "-")
+ name = documentRepeatedHyphens.ReplaceAllString(name, "-")
+ name = strings.Join(strings.Fields(name), " ")
+ runes := []rune(strings.TrimSpace(name))
+ if len(runes) > 200 {
+ runes = runes[:200]
+ }
+ name = string(runes)
+ if name == "" {
+ return "document"
+ }
+ return name
+}
+
+// parseClaudeDocumentBlock returns nil without error for unsupported or empty
+// attachments, matching the official client. Supported malformed documents
+// return a user-facing validation message.
+func parseClaudeDocumentBlock(block map[string]interface{}) (*KiroDocument, string) {
+ source, _ := block["source"].(map[string]interface{})
+ if source == nil {
+ return nil, ""
+ }
+ mediaType, _ := source["media_type"].(string)
+ if mediaType == "" {
+ mediaType, _ = source["mediaType"].(string)
+ }
+ if mediaType == "" {
+ mediaType, _ = source["mime_type"].(string)
+ }
+ mediaType = strings.ToLower(strings.TrimSpace(mediaType))
+ format, ok := documentMIMEToFormat[mediaType]
+ if !ok {
+ return nil, ""
+ }
+
+ sourceType, _ := source["type"].(string)
+ sourceType = strings.ToLower(strings.TrimSpace(sourceType))
+ if sourceType != "" && sourceType != "base64" && sourceType != "text" {
+ return nil, ""
+ }
+ data, _ := source["data"].(string)
+ if data == "" {
+ return nil, ""
+ }
+
+ var raw []byte
+ if sourceType == "text" {
+ if mediaType != "text/plain" {
+ return nil, ""
+ }
+ raw = []byte(data)
+ } else {
+ if strings.HasPrefix(data, "data:") {
+ idx := strings.Index(data, "base64,")
+ if idx < 0 {
+ return nil, fmt.Sprintf("invalid base64 for document %q", documentBlockName(block, format))
+ }
+ data = data[idx+7:]
+ }
+ if len(data) == 0 || len(data)%4 != 0 || !documentBase64Pattern.MatchString(data) {
+ return nil, fmt.Sprintf("invalid base64 for document %q", documentBlockName(block, format))
+ }
+ var err error
+ raw, err = base64.StdEncoding.DecodeString(data)
+ if err != nil || len(raw) == 0 {
+ return nil, fmt.Sprintf("invalid base64 for document %q", documentBlockName(block, format))
+ }
+ if magic := documentMagicBytes[mediaType]; len(magic) > 0 && !bytes.HasPrefix(raw, magic) {
+ return nil, fmt.Sprintf("document %q content does not match declared type %s", documentBlockName(block, format), mediaType)
+ }
+ }
+
+ doc := &KiroDocument{Name: sanitizeDocumentName(documentBlockName(block, format)), Format: format}
+ doc.Source.Bytes = base64.StdEncoding.EncodeToString(raw)
+ return doc, ""
+}
+
+func validateClaudeDocuments(messages []ClaudeMessage) string {
+ seen := make(map[string]bool)
+ for _, message := range messages {
+ if !strings.EqualFold(strings.TrimSpace(message.Role), "user") {
+ continue
+ }
+ _, _, documents, _, validationError := extractClaudeUserContent(message.Content)
+ if validationError != "" {
+ return validationError
+ }
+ for _, doc := range documents {
+ if seen[doc.Name] {
+ continue
+ }
+ seen[doc.Name] = true
+ if len(seen) > maxDocumentsPerConversation {
+ return fmt.Sprintf("too many documents attached (%d); maximum is %d per conversation", len(seen), maxDocumentsPerConversation)
+ }
+ }
+ }
+ return ""
+}
+
+func deduplicateConversationDocuments(history []KiroHistoryMessage, current []KiroDocument) []KiroDocument {
+ seen := make(map[string]bool)
+ filter := func(documents []KiroDocument) []KiroDocument {
+ filtered := documents[:0]
+ for _, document := range documents {
+ if document.Name == "" || seen[document.Name] || len(seen) >= maxDocumentsPerConversation {
+ continue
+ }
+ seen[document.Name] = true
+ filtered = append(filtered, document)
+ }
+ return filtered
+ }
+ for i := range history {
+ if history[i].UserInputMessage != nil {
+ history[i].UserInputMessage.Documents = filter(history[i].UserInputMessage.Documents)
+ }
+ }
+ return filter(current)
+}
+
func extractImageFromClaudeBlock(block map[string]interface{}) *KiroImage {
if source, ok := block["source"].(map[string]interface{}); ok {
if data, ok := source["data"].(string); ok {
@@ -759,6 +949,7 @@ func extractToolResultContent(content interface{}) (string, []KiroImage) {
func extractClaudeAssistantContent(content interface{}) (string, []KiroToolUse) {
var text string
var toolUses []KiroToolUse
+ var pendingSearchQuery string
if s, ok := content.(string); ok {
return s, nil
@@ -766,6 +957,9 @@ func extractClaudeAssistantContent(content interface{}) (string, []KiroToolUse)
// Same dual-shape support as extractClaudeUserContent (JSON []interface{}
// and in-memory []map[string]interface{} from agentic loop feedback).
+ // Native Anthropic server web_search blocks are not valid Kiro history tool
+ // uses; fold their results into assistant text so the next turn still sees
+ // the search evidence the client already displayed.
for _, block := range contentBlocksAsMaps(content) {
blockType, _ := block["type"].(string)
switch blockType {
@@ -785,12 +979,55 @@ func extractClaudeAssistantContent(content interface{}) (string, []KiroToolUse)
Name: name,
Input: input,
})
+ case "server_tool_use":
+ name, _ := block["name"].(string)
+ if name != webSearchToolName {
+ continue
+ }
+ if input, ok := block["input"].(map[string]interface{}); ok {
+ pendingSearchQuery = toolUseQuery(input)
+ } else {
+ pendingSearchQuery = ""
+ }
+ case "web_search_tool_result":
+ summary := summarizeClaudeWebSearchResult(pendingSearchQuery, block["content"])
+ if summary != "" {
+ if text != "" && !strings.HasSuffix(text, "\n") {
+ text += "\n\n"
+ }
+ text += summary
+ }
+ pendingSearchQuery = ""
}
}
return text, toolUses
}
+func summarizeClaudeWebSearchResult(query string, content interface{}) string {
+ results := &WebSearchResults{}
+ for _, block := range contentBlocksAsMaps(content) {
+ if typ, _ := block["type"].(string); typ != "" && typ != "web_search_result" {
+ continue
+ }
+ title, _ := block["title"].(string)
+ url, _ := block["url"].(string)
+ snippet, _ := block["encrypted_content"].(string)
+ if snippet == "" {
+ snippet, _ = block["snippet"].(string)
+ }
+ item := WebSearchResult{Title: title, URL: url}
+ if snippet != "" {
+ item.Snippet = &snippet
+ }
+ results.Results = append(results.Results, item)
+ }
+ // Empty result arrays are valid MCP successes; keep the same
+ // generateSearchSummary("No results found.") representation so a later
+ // client-tool continuation still retains that the search happened.
+ return generateSearchSummary(query, results)
+}
+
func convertClaudeTools(tools []ClaudeTool) ([]KiroToolWrapper, map[string]string) {
if len(tools) == 0 {
return nil, nil
@@ -995,7 +1232,85 @@ func shortenToolName(name string) string {
// ==================== Kiro -> Claude 转换 ====================
-func KiroToClaudeResponse(content, thinkingContent string, includeEmptyThinkingBlock bool, toolUses []KiroToolUse, inputTokens, outputTokens int, model string) *ClaudeResponse {
+func normalizeKiroStopReason(reason string) string {
+ reason = strings.ToLower(strings.TrimSpace(reason))
+ return strings.NewReplacer("-", "_", " ", "_").Replace(reason)
+}
+
+type kiroStopClass uint8
+
+const (
+ kiroStopComplete kiroStopClass = iota
+ kiroStopLength
+ kiroStopContextLimit
+ kiroStopFiltered
+ kiroStopSequence
+ kiroStopPaused
+)
+
+func classifyKiroStopReason(reason string) kiroStopClass {
+ switch normalizeKiroStopReason(reason) {
+ case "max_tokens", "max_output_tokens", "length":
+ return kiroStopLength
+ case "model_context_window_exceeded", "context_window_exceeded":
+ return kiroStopContextLimit
+ case "refusal", "content_filter", "content_filtered", "guardrail_intervened":
+ return kiroStopFiltered
+ case "stop_sequence":
+ return kiroStopSequence
+ case "pause_turn":
+ return kiroStopPaused
+ default:
+ return kiroStopComplete
+ }
+}
+
+func mapClaudeStopReason(reason string, toolCount int) string {
+ if toolCount > 0 {
+ return "tool_use"
+ }
+ switch classifyKiroStopReason(reason) {
+ case kiroStopLength:
+ return "max_tokens"
+ case kiroStopContextLimit:
+ return "model_context_window_exceeded"
+ case kiroStopFiltered:
+ return "refusal"
+ case kiroStopSequence:
+ return "stop_sequence"
+ case kiroStopPaused:
+ return "pause_turn"
+ default:
+ return "end_turn"
+ }
+}
+
+func mapOpenAIFinishReason(reason string, toolCount int) string {
+ if toolCount > 0 {
+ return "tool_calls"
+ }
+ switch classifyKiroStopReason(reason) {
+ case kiroStopLength, kiroStopContextLimit:
+ return "length"
+ case kiroStopFiltered:
+ return "content_filter"
+ default:
+ return "stop"
+ }
+}
+
+func mapResponsesCompletion(reason string) (status, incompleteReason string) {
+ switch classifyKiroStopReason(reason) {
+ case kiroStopLength, kiroStopContextLimit:
+ return "incomplete", "max_output_tokens"
+ case kiroStopFiltered:
+ return "incomplete", "content_filter"
+ default:
+ return "completed", ""
+ }
+}
+
+func KiroToClaudeResponse(content, thinkingContent string, includeEmptyThinkingBlock bool, toolUses []KiroToolUse, inputTokens, outputTokens int, model, stopReason string) *ClaudeResponse {
blocks := make([]ClaudeContentBlock, 0)
if thinkingContent != "" || includeEmptyThinkingBlock {
@@ -1021,10 +1336,7 @@ func KiroToClaudeResponse(content, thinkingContent string, includeEmptyThinkingB
})
}
- stopReason := "end_turn"
- if len(toolUses) > 0 {
- stopReason = "tool_use"
- }
+ clientStopReason := mapClaudeStopReason(stopReason, len(toolUses))
return &ClaudeResponse{
ID: "msg_" + uuid.New().String(),
@@ -1032,7 +1344,7 @@ func KiroToClaudeResponse(content, thinkingContent string, includeEmptyThinkingB
Role: "assistant",
Content: blocks,
Model: model,
- StopReason: stopReason,
+ StopReason: clientStopReason,
Usage: ClaudeUsage{
InputTokens: inputTokens,
OutputTokens: outputTokens,
@@ -1314,6 +1626,7 @@ func OpenAIToKiro(req *OpenAIRequest, thinking bool) *KiroPayload {
// 构建 payload
payload := &KiroPayload{}
payload.ConversationState.ChatTriggerType = "MANUAL"
+ // The OpenAI path has no chain seal yet, so it keeps the content-derived id.
payload.ConversationState.ConversationID = buildConversationID(modelID, systemPrompt, firstOpenAIConversationAnchor(nonSystemMessages))
payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
Content: finalContent,
@@ -1635,7 +1948,7 @@ func sanitizeKiroHistory(history []KiroHistoryMessage, currentToolResultIDs map[
// placeholder like ".": replayed across a long history that produces
// dozens of "." assistant turns, which the model then imitates by
// replying ".". Mark such turns for removal instead.
- if msg.UserInputMessage != nil && strings.TrimSpace(msg.UserInputMessage.Content) == "" && len(msg.UserInputMessage.Images) == 0 {
+ if msg.UserInputMessage != nil && strings.TrimSpace(msg.UserInputMessage.Content) == "" && len(msg.UserInputMessage.Images) == 0 && len(msg.UserInputMessage.Documents) == 0 {
msg.UserInputMessage.Content = minimalFallbackUserContent
}
}
@@ -1665,7 +1978,8 @@ func sanitizeKiroHistory(history []KiroHistoryMessage, currentToolResultIDs map[
if last.UserInputMessage != nil &&
strings.TrimSpace(last.UserInputMessage.Content) == strings.TrimSpace(msg.UserInputMessage.Content) &&
strings.TrimSpace(msg.UserInputMessage.Content) != "" &&
- len(msg.UserInputMessage.Images) == 0 {
+ len(last.UserInputMessage.Images) == 0 && len(msg.UserInputMessage.Images) == 0 &&
+ len(last.UserInputMessage.Documents) == 0 && len(msg.UserInputMessage.Documents) == 0 {
continue // skip duplicate consecutive user turn
}
}
@@ -1855,7 +2169,7 @@ func firstClaudeConversationAnchor(messages []ClaudeMessage) string {
if msg.Role != "user" {
continue
}
- text, _, toolResults := extractClaudeUserContent(msg.Content)
+ text, _, _, toolResults, _ := extractClaudeUserContent(msg.Content)
if strings.TrimSpace(text) != "" {
return strings.TrimSpace(text)
}
@@ -2157,8 +2471,8 @@ func extractThinkingFromContent(content string) (string, string) {
}
// KiroToOpenAIResponseWithReasoning 带 reasoning_content 的 OpenAI 响应
-func KiroToOpenAIResponseWithReasoning(content, reasoningContent string, toolUses []KiroToolUse, inputTokens, outputTokens int, model, thinkingFormat string) map[string]interface{} {
- finishReason := "stop"
+func KiroToOpenAIResponseWithReasoning(content, reasoningContent string, toolUses []KiroToolUse, inputTokens, outputTokens int, model, thinkingFormat, stopReason string) map[string]interface{} {
+ finishReason := mapOpenAIFinishReason(stopReason, len(toolUses))
message := map[string]interface{}{
"role": "assistant",
@@ -2179,7 +2493,6 @@ func KiroToOpenAIResponseWithReasoning(content, reasoningContent string, toolUse
}
}
message["tool_calls"] = toolCalls
- finishReason = "tool_calls"
} else {
// 根据配置格式化 thinking 输出
if reasoningContent != "" {
diff --git a/proxy/translator_test.go b/proxy/translator_test.go
index ad6ef174..fef77374 100644
--- a/proxy/translator_test.go
+++ b/proxy/translator_test.go
@@ -203,6 +203,8 @@ func TestOpenAIConversationIDStableFromAnchor(t *testing.T) {
}
func TestClaudeConversationIDStableFromAnchor(t *testing.T) {
+ // Continuity requires a successful seal of turn 1; bare content hash is no
+ // longer used (it collided across chats that shared an opener).
reqA := &ClaudeRequest{
Model: "claude-sonnet-4.5",
System: "sys",
@@ -210,6 +212,18 @@ func TestClaudeConversationIDStableFromAnchor(t *testing.T) {
{Role: "user", Content: "hello"},
},
}
+ payloadA := ClaudeToKiro(reqA, false)
+ if payloadA.ConversationState.ConversationID == "" {
+ t.Fatal("expected non-empty conversation ID")
+ }
+ sealSessionAfterSuccess("", MapModel("claude-sonnet-4.5"),
+ payloadA.ConversationState.ConversationID,
+ payloadA.ConversationState.AgentContinuationId,
+ []ClaudeMessage{
+ {Role: "user", Content: "hello"},
+ {Role: "assistant", Content: "ok"},
+ })
+
reqB := &ClaudeRequest{
Model: "claude-sonnet-4.5",
System: "sys",
@@ -219,15 +233,12 @@ func TestClaudeConversationIDStableFromAnchor(t *testing.T) {
{Role: "user", Content: "next"},
},
}
-
- payloadA := ClaudeToKiro(reqA, false)
- payloadB := ClaudeToKiro(reqB, false)
-
- if payloadA.ConversationState.ConversationID == "" || payloadB.ConversationState.ConversationID == "" {
- t.Fatalf("expected non-empty conversation IDs")
+ gotConv, gotCont, reused := resolveSessionIDs("", MapModel("claude-sonnet-4.5"), reqB.Messages)
+ if !reused {
+ t.Fatal("expected reuse after seal")
}
- if payloadA.ConversationState.ConversationID != payloadB.ConversationState.ConversationID {
- t.Fatalf("expected stable conversation ID across turns, got %q vs %q", payloadA.ConversationState.ConversationID, payloadB.ConversationState.ConversationID)
+ if gotConv != payloadA.ConversationState.ConversationID || gotCont != payloadA.ConversationState.AgentContinuationId {
+ t.Fatalf("stable ids after seal: got %s/%s want %s/%s", gotConv, gotCont, payloadA.ConversationState.ConversationID, payloadA.ConversationState.AgentContinuationId)
}
}
@@ -268,7 +279,7 @@ func TestClaudeToKiroDropsLeadingAssistantHistory(t *testing.T) {
}
func TestKiroToClaudeResponseCanEmitEmptyThinkingBlock(t *testing.T) {
- resp := KiroToClaudeResponse("final answer", "", true, nil, 10, 20, "claude-sonnet-4.6")
+ resp := KiroToClaudeResponse("final answer", "", true, nil, 10, 20, "claude-sonnet-4.6", "END_TURN")
if len(resp.Content) != 2 {
t.Fatalf("expected empty thinking block plus text block, got %d blocks", len(resp.Content))
@@ -526,6 +537,13 @@ func TestClaudeToolResultMixedTextAndImage(t *testing.T) {
req := &ClaudeRequest{
Model: "claude-opus-4.8",
Messages: []ClaudeMessage{
+ {Role: "user", Content: "read this image"},
+ {
+ Role: "assistant",
+ Content: []interface{}{
+ map[string]interface{}{"type": "tool_use", "id": "tool_2", "name": "read", "input": map[string]interface{}{"path": "a.png"}},
+ },
+ },
{
Role: "user",
Content: []interface{}{
@@ -626,13 +644,12 @@ func TestOpenAIToolResultImageCarriedWhenFollowedByUser(t *testing.T) {
var toolHistImages int
for _, h := range payload.ConversationState.History {
- if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil &&
- len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 {
+ if h.UserInputMessage != nil && strings.Contains(h.UserInputMessage.Content, toolResultsContinuationPrefix) {
toolHistImages += len(h.UserInputMessage.Images)
}
}
if toolHistImages != 1 {
- t.Fatalf("expected tool image carried on the flushed tool-result history entry, got %d", toolHistImages)
+ t.Fatalf("expected tool image carried on the narrated tool-result history entry, got %d", toolHistImages)
}
cur := payload.ConversationState.CurrentMessage.UserInputMessage
From b784418a301b5cf953055c6a7d5c56d615a67a28 Mon Sep 17 00:00:00 2001
From: asuan-dev <116167305+asuan-dev@users.noreply.github.com>
Date: Sun, 26 Jul 2026 13:16:25 +0800
Subject: [PATCH 6/6] fix(proxy): wire integrity, continuity, and web-search
seals on author paths
Route Claude/OpenAI/Responses through integrity retry, keep ClaudeToKiro
pure with handler-side attachStoredReasoning, and restore web-search
lifecycle seals: intermediate rounds seal the raw upstream tool shape so
the next loop reuses ids; final flush seals client-visible content.
---
proxy/handler.go | 498 ++++++++++++++++++++------------
proxy/handler_test.go | 214 +++++++++++++-
proxy/responses_handler.go | 299 +++++++++++--------
proxy/responses_handler_test.go | 110 +++++++
proxy/responses_types.go | 33 ++-
proxy/websearch.go | 2 +-
proxy/websearch_loop.go | 155 ++++++----
proxy/websearch_test.go | 241 +++++++++++++++-
8 files changed, 1177 insertions(+), 375 deletions(-)
diff --git a/proxy/handler.go b/proxy/handler.go
index dfec8535..4def3f3b 100644
--- a/proxy/handler.go
+++ b/proxy/handler.go
@@ -135,6 +135,9 @@ func validateClaudeRequestShape(req *ClaudeRequest) string {
if len(req.Messages) == 0 {
return "messages must not be empty"
}
+ if msg := validateClaudeDocuments(req.Messages); msg != "" {
+ return msg
+ }
if msg := validateClaudeThinkingConfig(req.Thinking, req.MaxTokens); msg != "" {
return msg
}
@@ -151,8 +154,8 @@ func validateClaudeRequestShape(req *ClaudeRequest) string {
continue
}
- text, images, toolResults := extractClaudeUserContent(msg.Content)
- if normalizeUserContent(text, len(images) > 0) != "" || len(toolResults) > 0 {
+ text, images, documents, toolResults, _ := extractClaudeUserContent(msg.Content)
+ if normalizeUserContent(text, len(images) > 0 || len(documents) > 0) != "" || len(toolResults) > 0 {
hasUserContext = true
}
}
@@ -876,33 +879,37 @@ func (h *Handler) handleClaudeMessagesInternal(w http.ResponseWriter, r *http.Re
apiKeyID := apiKeyIDFromContext(r.Context())
- // Pure native web_search: relay via Kiro MCP (generateAssistantResponse does not run it).
+ // Pure native web_search is handled by Kiro MCP because
+ // generateAssistantResponse does not execute it.
if hasWebSearchTool(&req) {
h.handleWebSearchRequest(w, &req, estimatedInputTokens, apiKeyID)
return
}
- // Mixed tools including native web_search: agentic loop digests web_search internally
- // and returns client tool_use blocks as-is.
+ // Mixed native web_search and client tools use the existing local agentic
+ // loop; client tools remain visible to the caller.
if hasWebSearchAmongTools(&req) {
logger.Infof("[WebSearch] Mixed tools with native web_search, entering agentic loop")
h.runWebSearchLoop(w, &req, thinking, estimatedInputTokens, apiKeyID)
return
}
- // 转换请求
+ // 转换请求（纯映射）；会话 ID / 思考戳在 handler 侧按 apiKey 绑定。
kiroPayload := ClaudeToKiro(&req, thinking)
-
- // Stream or non-stream
+ modelID := MapModel(req.Model)
+ convID, contID, _ := resolveSessionIDs(apiKeyID, modelID, req.Messages)
+ kiroPayload.ConversationState.ConversationID = convID
+ kiroPayload.ConversationState.AgentContinuationId = contID
+ attachStoredReasoning(kiroPayload.ConversationState.History, apiKeyID, modelID, convID)
if req.Stream {
- h.handleClaudeStream(w, kiroPayload, req.Model, thinking, thinkingResponseOpts, estimatedInputTokens, cacheProfile, apiKeyID)
+ h.handleClaudeStream(w, kiroPayload, req.Model, thinking, thinkingResponseOpts, estimatedInputTokens, cacheProfile, apiKeyID, req.Messages)
} else {
- h.handleClaudeNonStream(w, kiroPayload, req.Model, thinking, thinkingResponseOpts, estimatedInputTokens, cacheProfile, apiKeyID)
+ h.handleClaudeNonStream(w, kiroPayload, req.Model, thinking, thinkingResponseOpts, estimatedInputTokens, cacheProfile, apiKeyID, req.Messages)
}
}
// handleClaudeStream Claude 流式响应
-func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload, model string, thinking bool, thinkingOpts claudeThinkingResponseOptions, estimatedInputTokens int, cacheProfile *promptCacheProfile, apiKeyID string) {
+func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload, model string, thinking bool, thinkingOpts claudeThinkingResponseOptions, estimatedInputTokens int, cacheProfile *promptCacheProfile, apiKeyID string, inbound []ClaudeMessage) {
w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
@@ -1209,73 +1216,113 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload
}
}
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if text == "" {
- return
- }
- if isThinking {
- rawThinkingBuilder.WriteString(text)
- } else {
- rawContentBuilder.WriteString(text)
- }
- processClaudeText(text, isThinking, false)
- },
- OnToolUse: func(tu KiroToolUse) {
- processClaudeText("", false, true)
- rawContentBuilder.WriteString(tu.Name)
- if b, err := json.Marshal(tu.Input); err == nil {
- rawContentBuilder.Write(b)
- }
+ var integrity streamIntegrityState
+ var stamp reasoningCapture
- toolUses = append(toolUses, tu)
- ensureMessageStart()
- closeActiveBlock()
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ if text == "" {
+ return
+ }
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ rawThinkingBuilder.WriteString(text)
+ } else {
+ rawContentBuilder.WriteString(text)
+ }
+ processClaudeText(text, isThinking, false)
+ },
+ OnReasoningMeta: stamp.meta(),
+ OnToolUse: func(tu KiroToolUse) {
+ processClaudeText("", false, true)
+ // Do NOT append tool name/input into rawContentBuilder.
+ // That padding was only for integrity length, but toolCount is
+ // measured separately; padding polluted the sealed assistant
+ // text so the next turn's fingerprint never matched.
+ toolUses = append(toolUses, tu)
+ integrity.observeToolUse()
+ ensureMessageStart()
+ closeActiveBlock()
- idx := nextContentIndex
- nextContentIndex++
+ idx := nextContentIndex
+ nextContentIndex++
- h.sendSSE(w, flusher, "content_block_start", map[string]interface{}{
- "type": "content_block_start",
- "index": idx,
- "content_block": map[string]interface{}{
- "type": "tool_use",
- "id": tu.ToolUseID,
- "name": tu.Name,
- "input": map[string]interface{}{},
- },
- })
+ h.sendSSE(w, flusher, "content_block_start", map[string]interface{}{
+ "type": "content_block_start",
+ "index": idx,
+ "content_block": map[string]interface{}{
+ "type": "tool_use",
+ "id": tu.ToolUseID,
+ "name": tu.Name,
+ "input": map[string]interface{}{},
+ },
+ })
- inputJSON, _ := json.Marshal(tu.Input)
- h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{
- "type": "content_block_delta",
- "index": idx,
- "delta": map[string]interface{}{
- "type": "input_json_delta",
- "partial_json": string(inputJSON),
- },
- })
+ inputJSON, _ := json.Marshal(tu.Input)
+ h.sendSSE(w, flusher, "content_block_delta", map[string]interface{}{
+ "type": "content_block_delta",
+ "index": idx,
+ "delta": map[string]interface{}{
+ "type": "input_json_delta",
+ "partial_json": string(inputJSON),
+ },
+ })
- h.sendSSE(w, flusher, "content_block_stop", map[string]interface{}{
- "type": "content_block_stop",
- "index": idx,
- })
- },
- OnComplete: func(inTok, outTok int) {
- inputTokens = inTok
- outputTokens = outTok
- },
- OnCredits: func(c float64) {
- credits = c
+ h.sendSSE(w, flusher, "content_block_stop", map[string]interface{}{
+ "type": "content_block_stop",
+ "index": idx,
+ })
+ },
+ OnComplete: func(inTok, outTok int) {
+ inputTokens = inTok
+ outputTokens = outTok
+ },
+ OnCredits: func(c float64) {
+ credits = c
+ },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
+ }
},
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ func() {
+ rawContentBuilder.Reset()
+ rawThinkingBuilder.Reset()
+ toolUses = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
+ stamp.reset()
+ nextContentIndex = 0
+ textBuffer = ""
+ inThinkingBlock = false
+ dropTagThinking = false
+ thinkingSource = thinkingSourceUnknown
+ thinkingStarted = false
+ eventThinkingOpen = false
+ activeBlockIndex = -1
+ activeBlockType = ""
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ func() bool { return !messageStarted },
+ )
if err != nil {
lastErr = err
+ if isStreamIntegrityError(err) {
+ if !messageStarted {
+ // Still safe to rotate: nothing flushed yet.
+ excluded[account.ID] = true
+ continue
+ }
+ // Partial content already reached the client — do not forge end_turn.
+ h.recordFailureWithDetails("claude", model, account.ID, err)
+ h.sendSSE(w, flusher, "error", map[string]interface{}{
+ "type": "error",
+ "error": map[string]string{"type": "api_error", "message": err.Error()},
+ })
+ return
+ }
excluded[account.ID] = true
h.handleAccountFailure(account, err)
if !messageStarted {
@@ -1301,10 +1348,11 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload
inputTokens = estimatedInputTokens
}
outputContent, extractedReasoning := extractThinkingFromContent(rawContentBuilder.String())
- thinkingOutput := rawThinkingBuilder.String()
- if thinking && thinkingOutput == "" && extractedReasoning != "" {
- thinkingOutput = extractedReasoning
+ rawThinkingForSeal := rawThinkingBuilder.String()
+ if rawThinkingForSeal == "" && extractedReasoning != "" {
+ rawThinkingForSeal = extractedReasoning
}
+ thinkingOutput := rawThinkingForSeal
if !thinking {
thinkingOutput = ""
}
@@ -1316,16 +1364,21 @@ func (h *Handler) handleClaudeStream(w http.ResponseWriter, payload *KiroPayload
h.promptCache.Update(account.ID, cacheProfile)
h.recordSuccessLog("claude", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
- stopReason := "end_turn"
- if len(toolUses) > 0 {
- stopReason = "tool_use"
+ // Seal thinking stamp for this turn only after a successful complete
+ // response. Keyed by parent user text + assistant text so the next
+ // request can reattach without trusting the client to echo signature.
+ if isSuccessfulKiroTurn(integrity.StopReason, len(toolUses)) {
+ stamp.sealIfPresent(apiKeyID, MapModel(model), payload.ConversationState.ConversationID, currentTurnParentUser(payload), outputContent, toolUses, thinkingOutput)
+ sealClaudeSession(apiKeyID, model, payload.ConversationState.ConversationID, payload.ConversationState.AgentContinuationId, inbound, outputContent, thinkingOutput, toolUses, thinkingFormat, thinkingOpts.OmitDisplay)
}
+ clientStopReason := mapClaudeStopReason(integrity.StopReason, len(toolUses))
+
ensureMessageStart()
h.sendSSE(w, flusher, "message_delta", map[string]interface{}{
"type": "message_delta",
"delta": map[string]interface{}{
- "stop_reason": stopReason,
+ "stop_reason": clientStopReason,
},
"usage": buildClaudeUsageMap(inputTokens, outputTokens, cacheUsage, cacheProfile != nil),
})
@@ -1499,7 +1552,7 @@ func (h *Handler) getRequestLogs() []RequestLog {
}
// handleClaudeNonStream Claude 非流式响应
-func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayload, model string, thinking bool, thinkingOpts claudeThinkingResponseOptions, estimatedInputTokens int, cacheProfile *promptCacheProfile, apiKeyID string) {
+func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayload, model string, thinking bool, thinkingOpts claudeThinkingResponseOptions, estimatedInputTokens int, cacheProfile *promptCacheProfile, apiKeyID string, inbound []ClaudeMessage) {
excluded := make(map[string]bool)
var lastErr error
reqStart := time.Now()
@@ -1524,43 +1577,63 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl
var credits float64
var realInputTokens int
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if isThinking {
- thinkingContent += text
- } else {
- content += text
+ var integrity streamIntegrityState
+ var stamp reasoningCapture
+
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ thinkingContent += text
+ } else {
+ content += text
+ }
+ },
+ OnReasoningMeta: stamp.meta(),
+ OnToolUse: func(tu KiroToolUse) {
+ toolUses = append(toolUses, tu)
+ integrity.observeToolUse()
+ },
+ OnComplete: func(inTok, outTok int) {
+ inputTokens = inTok
+ outputTokens = outTok
+ },
+ OnCredits: func(c float64) {
+ credits = c
+ },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
}
},
- OnToolUse: func(tu KiroToolUse) {
- toolUses = append(toolUses, tu)
- },
- OnComplete: func(inTok, outTok int) {
- inputTokens = inTok
- outputTokens = outTok
- },
- OnCredits: func(c float64) {
- credits = c
+ func() {
+ content, thinkingContent = "", ""
+ toolUses = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
+ stamp.reset()
},
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
- },
- }
-
- err := CallKiroAPI(account, payload, callback)
+ nil,
+ )
if err != nil {
lastErr = err
excluded[account.ID] = true
- h.handleAccountFailure(account, err)
+ if !isStreamIntegrityError(err) {
+ h.handleAccountFailure(account, err)
+ }
continue
}
thinkingFormat := thinkingOpts.Format
finalContent, extractedReasoning := extractThinkingFromContent(content)
- rawThinkingContent := thinkingContent
- if thinking && rawThinkingContent == "" && extractedReasoning != "" {
- rawThinkingContent = extractedReasoning
+ rawThinkingForSeal := thinkingContent
+ if rawThinkingForSeal == "" && extractedReasoning != "" {
+ rawThinkingForSeal = extractedReasoning
}
+ rawThinkingContent := rawThinkingForSeal
if !thinking {
rawThinkingContent = ""
}
@@ -1578,6 +1651,14 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl
h.promptCache.Update(account.ID, cacheProfile)
h.recordSuccessLog("claude", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
+ if isSuccessfulKiroTurn(integrity.StopReason, len(toolUses)) {
+ stamp.sealIfPresent(apiKeyID, MapModel(model), payload.ConversationState.ConversationID, currentTurnParentUser(payload), finalContent, toolUses, rawThinkingContent)
+ // Seal BEFORE folding reasoning into text for response formatting, but
+ // pass the raw thinking + format so claudeAssistantTurn can render the
+ // same shape the client will replay.
+ sealClaudeSession(apiKeyID, model, payload.ConversationState.ConversationID, payload.ConversationState.AgentContinuationId, inbound, finalContent, rawThinkingContent, toolUses, thinkingFormat, thinkingOpts.OmitDisplay)
+ }
+
responseThinkingContent := rawThinkingContent
includeEmptyThinkingBlock := thinking && thinkingOpts.OmitDisplay && rawThinkingContent != ""
if includeEmptyThinkingBlock {
@@ -1596,7 +1677,7 @@ func (h *Handler) handleClaudeNonStream(w http.ResponseWriter, payload *KiroPayl
}
}
- resp := KiroToClaudeResponse(finalContent, responseThinkingContent, includeEmptyThinkingBlock, toolUses, inputTokens, outputTokens, model)
+ resp := KiroToClaudeResponse(finalContent, responseThinkingContent, includeEmptyThinkingBlock, toolUses, inputTokens, outputTokens, model, integrity.StopReason)
resp.Usage.InputTokens = billedClaudeInputTokens(inputTokens, cacheUsage)
resp.Usage.CacheCreationInputTokens = cacheUsage.CacheCreationInputTokens
resp.Usage.CacheReadInputTokens = cacheUsage.CacheReadInputTokens
@@ -1922,76 +2003,119 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload
}
}
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if text == "" {
- return
- }
- if isThinking {
- rawReasoningBuilder.WriteString(text)
- } else {
- rawContentBuilder.WriteString(text)
- }
- processText(text, isThinking, false)
- },
- OnToolUse: func(tu KiroToolUse) {
- processText("", false, true)
-
- args, _ := json.Marshal(tu.Input)
- rawContentBuilder.WriteString(tu.Name)
- rawContentBuilder.Write(args)
- tc := ToolCall{ID: tu.ToolUseID, Type: "function"}
- tc.Function.Name = tu.Name
- tc.Function.Arguments = string(args)
- toolCalls = append(toolCalls, tc)
-
- chunk := map[string]interface{}{
- "id": chatID,
- "object": "chat.completion.chunk",
- "created": time.Now().Unix(),
- "model": model,
- "choices": []map[string]interface{}{{
- "index": 0,
- "delta": map[string]interface{}{
- "tool_calls": []map[string]interface{}{{
- "index": toolCallIndex,
- "id": tu.ToolUseID,
- "type": "function",
- "function": map[string]string{
- "name": tu.Name,
- "arguments": string(args),
+ sendStreamError := func(err error) {
+ payload := map[string]interface{}{
+ "error": map[string]string{
+ "message": err.Error(),
+ "type": "server_error",
+ },
+ }
+ data, _ := json.Marshal(payload)
+ fmt.Fprintf(w, "data: %s\n\n", data)
+ flusher.Flush()
+ }
+
+ var integrity streamIntegrityState
+
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ if text == "" {
+ return
+ }
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ rawReasoningBuilder.WriteString(text)
+ } else {
+ rawContentBuilder.WriteString(text)
+ }
+ processText(text, isThinking, false)
+ },
+ OnToolUse: func(tu KiroToolUse) {
+ processText("", false, true)
+
+ args, _ := json.Marshal(tu.Input)
+ tc := ToolCall{ID: tu.ToolUseID, Type: "function"}
+ tc.Function.Name = tu.Name
+ tc.Function.Arguments = string(args)
+ toolCalls = append(toolCalls, tc)
+ integrity.observeToolUse()
+
+ chunk := map[string]interface{}{
+ "id": chatID,
+ "object": "chat.completion.chunk",
+ "created": time.Now().Unix(),
+ "model": model,
+ "choices": []map[string]interface{}{{
+ "index": 0,
+ "delta": map[string]interface{}{
+ "tool_calls": []map[string]interface{}{{
+ "index": toolCallIndex,
+ "id": tu.ToolUseID,
+ "type": "function",
+ "function": map[string]string{
+ "name": tu.Name,
+ "arguments": string(args),
+ },
+ }},
},
+ "finish_reason": nil,
}},
- },
- "finish_reason": nil,
- }},
+ }
+ toolCallIndex++
+ data, _ := json.Marshal(chunk)
+ fmt.Fprintf(w, "data: %s\n\n", string(data))
+ flusher.Flush()
+ responseStarted = true
+ },
+ OnComplete: func(inTok, outTok int) {
+ inputTokens = inTok
+ outputTokens = outTok
+ },
+ OnCredits: func(c float64) {
+ credits = c
+ },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
}
- toolCallIndex++
- data, _ := json.Marshal(chunk)
- fmt.Fprintf(w, "data: %s\n\n", string(data))
- flusher.Flush()
- responseStarted = true
},
- OnComplete: func(inTok, outTok int) {
- inputTokens = inTok
- outputTokens = outTok
- },
- OnCredits: func(c float64) {
- credits = c
- },
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ func() {
+ rawContentBuilder.Reset()
+ rawReasoningBuilder.Reset()
+ toolCalls = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
+ toolCallIndex = 0
+ textBuffer = ""
+ inThinkingBlock = false
+ dropTagThinking = false
+ thinkingSource = thinkingSourceUnknown
+ thinkingStarted = false
+ eventThinkingOpen = false
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ func() bool { return !responseStarted },
+ )
if err != nil {
lastErr = err
+ if isStreamIntegrityError(err) {
+ if !responseStarted {
+ excluded[account.ID] = true
+ continue
+ }
+ // Partial content already reached the client — signal failure, no success terminator.
+ sendStreamError(err)
+ h.recordFailureWithDetails("openai", model, account.ID, err)
+ return
+ }
excluded[account.ID] = true
h.handleAccountFailure(account, err)
if !responseStarted {
continue
}
+ sendStreamError(err)
h.recordFailureWithDetails("openai", model, account.ID, err)
return
}
@@ -2025,10 +2149,7 @@ func (h *Handler) handleOpenAIStream(w http.ResponseWriter, payload *KiroPayload
h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits)
h.recordSuccessLog("openai", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
- finishReason := "stop"
- if len(toolCalls) > 0 {
- finishReason = "tool_calls"
- }
+ finishReason := mapOpenAIFinishReason(integrity.StopReason, len(toolCalls))
chunk := map[string]interface{}{
"id": chatID,
@@ -2087,27 +2208,42 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl
var credits float64
var realInputTokens int
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if isThinking {
- reasoningContent += text
- } else {
- content += text
+ var integrity streamIntegrityState
+
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ reasoningContent += text
+ } else {
+ content += text
+ }
+ },
+ OnToolUse: func(tu KiroToolUse) { integrity.observeToolUse(); toolUses = append(toolUses, tu) },
+ OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
+ OnCredits: func(c float64) { credits = c },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
}
},
- OnToolUse: func(tu KiroToolUse) { toolUses = append(toolUses, tu) },
- OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
- OnCredits: func(c float64) { credits = c },
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ func() {
+ content, reasoningContent = "", ""
+ toolUses = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ nil,
+ )
if err != nil {
lastErr = err
excluded[account.ID] = true
- h.handleAccountFailure(account, err)
+ if !isStreamIntegrityError(err) {
+ h.handleAccountFailure(account, err)
+ }
continue
}
@@ -2131,7 +2267,7 @@ func (h *Handler) handleOpenAINonStream(w http.ResponseWriter, payload *KiroPayl
h.recordSuccessLog("openai", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
thinkingFormat := config.GetThinkingConfig().OpenAIFormat
- resp := KiroToOpenAIResponseWithReasoning(finalContent, reasoningContent, toolUses, inputTokens, outputTokens, model, thinkingFormat)
+ resp := KiroToOpenAIResponseWithReasoning(finalContent, reasoningContent, toolUses, inputTokens, outputTokens, model, thinkingFormat, integrity.StopReason)
w.Header().Set("Content-Type", "application/json; charset=utf-8")
json.NewEncoder(w).Encode(resp)
return
diff --git a/proxy/handler_test.go b/proxy/handler_test.go
index 728f88f7..1b4a01a7 100644
--- a/proxy/handler_test.go
+++ b/proxy/handler_test.go
@@ -7,6 +7,7 @@ import (
"net/http"
"net/http/httptest"
"strings"
+ "sync/atomic"
"testing"
"time"
)
@@ -68,6 +69,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()
@@ -98,7 +102,7 @@ func TestClaudeNonStreamRetriesNextAccountAfterPreResponseFailure(t *testing.T)
}
rec := httptest.NewRecorder()
- h.handleClaudeNonStream(rec, payload, "claude-sonnet-4.5", false, claudeThinkingResponseOptions{}, 1, nil, "")
+ h.handleClaudeNonStream(rec, payload, "claude-sonnet-4.5", false, claudeThinkingResponseOptions{}, 1, nil, "", nil)
if rec.Code != http.StatusOK {
t.Fatalf("expected retry to succeed, status=%d body=%s", rec.Code, rec.Body.String())
@@ -480,3 +484,211 @@ func TestBuildAnthropicModelsResponseGeneratesThinkingVariants(t *testing.T) {
t.Fatalf("expected image capability to be preserved, got %#v", models[0]["supports_image"])
}
}
+
+func TestClaudeStreamRetryDiscardsBufferedPartialText(t *testing.T) {
+ cfgFile := t.TempDir() + "/config.json"
+ if err := config.Init(cfgFile); err != nil {
+ t.Fatalf("config.Init: %v", err)
+ }
+ if err := config.AddAccount(config.Account{
+ ID: "only",
+ Enabled: true,
+ AccessToken: "token",
+ ProfileArn: "arn:aws:codewhisperer:profile/only",
+ }); err != nil {
+ t.Fatalf("add account: %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)
+ }
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if hits.Add(1) == 1 {
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "stale"}))
+ return
+ }
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "fresh"}))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "MAX_TOKENS"}))
+ }))
+ defer server.Close()
+
+ oldEndpoints := kiroEndpoints
+ kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}}
+ defer func() { kiroEndpoints = oldEndpoints }()
+ oldClient := kiroHttpStore.Load()
+ kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}})
+ defer kiroHttpStore.Store(oldClient)
+
+ p := accountpool.GetPool()
+ p.Reload()
+ h := &Handler{pool: p, promptCache: newPromptCacheTracker(defaultPromptCacheTTL)}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
+ Content: "hello", ModelID: "claude-sonnet-4.5", Origin: "AI_EDITOR",
+ }
+
+ rec := httptest.NewRecorder()
+ h.handleClaudeStream(rec, payload, "claude-sonnet-4.5", false, claudeThinkingResponseOptions{}, 1, nil, "", []ClaudeMessage{{Role: "user", Content: "hello"}})
+ body := rec.Body.String()
+ if hits.Load() != 2 {
+ t.Fatalf("upstream hits = %d, want one retry", hits.Load())
+ }
+ if strings.Contains(body, "stale") {
+ t.Fatalf("retry leaked first attempt text into SSE: %s", body)
+ }
+ if !strings.Contains(body, "fresh") {
+ t.Fatalf("retry response missing fresh text: %s", body)
+ }
+ if !strings.Contains(body, `"stop_reason":"max_tokens"`) {
+ t.Fatalf("Claude stop reason was not normalized: %s", body)
+ }
+}
+
+func TestOpenAIStreamRetryDiscardsBufferedPartialText(t *testing.T) {
+ cfgFile := t.TempDir() + "/config.json"
+ if err := config.Init(cfgFile); err != nil {
+ t.Fatalf("config.Init: %v", err)
+ }
+ if err := config.AddAccount(config.Account{
+ ID: "only",
+ Enabled: true,
+ AccessToken: "token",
+ ProfileArn: "arn:aws:codewhisperer:profile/only",
+ }); err != nil {
+ t.Fatalf("add account: %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)
+ }
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if hits.Add(1) == 1 {
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "stale"}))
+ return
+ }
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "fresh"}))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "MAX_TOKENS"}))
+ }))
+ defer server.Close()
+
+ oldEndpoints := kiroEndpoints
+ kiroEndpoints = []kiroEndpoint{{URL: server.URL, Origin: "AI_EDITOR", Name: "test"}}
+ defer func() { kiroEndpoints = oldEndpoints }()
+ oldClient := kiroHttpStore.Load()
+ kiroHttpStore.Store(&http.Client{Timeout: time.Second, Transport: &http.Transport{}})
+ defer kiroHttpStore.Store(oldClient)
+
+ p := accountpool.GetPool()
+ p.Reload()
+ h := &Handler{pool: p, promptCache: newPromptCacheTracker(defaultPromptCacheTTL)}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
+ Content: "hello", ModelID: "claude-sonnet-4.5", Origin: "AI_EDITOR",
+ }
+
+ rec := httptest.NewRecorder()
+ h.handleOpenAIStream(rec, payload, "claude-sonnet-4.5", false, 1, "")
+ body := rec.Body.String()
+ if hits.Load() != 2 {
+ t.Fatalf("upstream hits = %d, want one retry", hits.Load())
+ }
+ if strings.Contains(body, "stale") {
+ t.Fatalf("retry leaked first attempt text into SSE: %s", body)
+ }
+ if !strings.Contains(body, "fresh") {
+ t.Fatalf("retry response missing fresh text: %s", body)
+ }
+ if !strings.Contains(body, `"finish_reason":"length"`) {
+ t.Fatalf("OpenAI finish reason was not mapped: %s", body)
+ }
+}
+
+func TestOpenAIStreamSignalsErrorAfterPartialFlush(t *testing.T) {
+ if err := config.Init(t.TempDir() + "/config.json"); err != nil {
+ t.Fatalf("config.Init: %v", err)
+ }
+ if err := config.AddAccount(config.Account{
+ ID: "only", Enabled: true, AccessToken: "token",
+ ProfileArn: "arn:aws:codewhisperer:profile/only",
+ }); err != nil {
+ t.Fatalf("add account: %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)
+ }
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": strings.Repeat("partial ", 8)}))
+ // No metadataEvent: output has already flushed, so retry is unsafe.
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(t, server)()
+
+ p := accountpool.GetPool()
+ p.Reload()
+ h := &Handler{pool: p, promptCache: newPromptCacheTracker(defaultPromptCacheTTL)}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
+ Content: "hello", ModelID: "claude-sonnet-4.5", Origin: "AI_EDITOR",
+ }
+ rec := httptest.NewRecorder()
+ h.handleOpenAIStream(rec, payload, "claude-sonnet-4.5", false, 1, "")
+ body := rec.Body.String()
+ if !strings.Contains(body, `"error":{"message":`) || !strings.Contains(body, `"type":"server_error"`) {
+ t.Fatalf("post-flush stream missing explicit error event: %s", body)
+ }
+ if strings.Contains(body, "data: [DONE]") || strings.Contains(body, `"finish_reason":"stop"`) {
+ t.Fatalf("failed stream must not emit a success terminator: %s", body)
+ }
+}
+
+func TestClaudeNonStreamIntegrityExhaustionExcludesAccount(t *testing.T) {
+ if err := config.Init(t.TempDir() + "/config.json"); err != nil {
+ t.Fatalf("config.Init: %v", err)
+ }
+ if err := config.AddAccount(config.Account{
+ ID: "only", Enabled: true, AccessToken: "token",
+ ProfileArn: "arn:aws:codewhisperer:profile/only",
+ }); err != nil {
+ t.Fatalf("add account: %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)
+ }
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ hits.Add(1)
+ w.WriteHeader(http.StatusOK)
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(t, server)()
+
+ p := accountpool.GetPool()
+ p.Reload()
+ h := &Handler{pool: p, promptCache: newPromptCacheTracker(defaultPromptCacheTTL)}
+ payload := &KiroPayload{}
+ payload.ConversationState.CurrentMessage.UserInputMessage = KiroUserInputMessage{
+ Content: "hello", ModelID: "claude-sonnet-4.5", Origin: "AI_EDITOR",
+ }
+ h.handleClaudeNonStream(httptest.NewRecorder(), payload, "claude-sonnet-4.5", false, claudeThinkingResponseOptions{}, 1, nil, "", nil)
+
+ want := int32(maxSameAccountStreamRetries + 1)
+ if hits.Load() != want {
+ t.Fatalf("empty account called %d times, want one same-account retry budget (%d)", hits.Load(), want)
+ }
+}
diff --git a/proxy/responses_handler.go b/proxy/responses_handler.go
index a58106f6..7dbd9daf 100644
--- a/proxy/responses_handler.go
+++ b/proxy/responses_handler.go
@@ -150,27 +150,42 @@ func (h *Handler) handleResponsesNonStream(
var credits float64
var realInputTokens int
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if isThinking {
- reasoningContent += text
- } else {
- content += text
+ var integrity streamIntegrityState
+
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ reasoningContent += text
+ } else {
+ content += text
+ }
+ },
+ OnToolUse: func(tu KiroToolUse) { integrity.observeToolUse(); toolUses = append(toolUses, tu) },
+ OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
+ OnCredits: func(c float64) { credits = c },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
}
},
- OnToolUse: func(tu KiroToolUse) { toolUses = append(toolUses, tu) },
- OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
- OnCredits: func(c float64) { credits = c },
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ func() {
+ content, reasoningContent = "", ""
+ toolUses = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ nil,
+ )
if err != nil {
lastErr = err
excluded[account.ID] = true
- h.handleAccountFailure(account, err)
+ if !isStreamIntegrityError(err) {
+ h.handleAccountFailure(account, err)
+ }
continue
}
@@ -191,7 +206,7 @@ func (h *Handler) handleResponsesNonStream(
h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits)
h.recordSuccessLog("responses", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
- respObj := buildResponsesObject(respID, model, finalContent, toolUses, inputTokens, outputTokens, req)
+ respObj := buildResponsesObject(respID, model, finalContent, toolUses, inputTokens, outputTokens, integrity.StopReason, req)
respObj.StoredInput = storedInput
respObj.Instructions = req.Instructions
@@ -216,16 +231,17 @@ func (h *Handler) handleResponsesNonStream(
func buildResponsesObject(
id, model, content string, toolUses []KiroToolUse,
- inputTokens, outputTokens int, req *ResponsesRequest,
+ inputTokens, outputTokens int, stopReason string, req *ResponsesRequest,
) *ResponsesObject {
output := make([]ResponseOutputItem, 0, 1+len(toolUses))
+ status, incompleteReason := mapResponsesCompletion(stopReason)
if strings.TrimSpace(content) != "" {
output = append(output, ResponseOutputItem{
ID: generateOutputItemID("msg"),
Type: "message",
Role: "assistant",
- Status: "completed",
+ Status: status,
Content: []ResponseContentPart{{
Type: "output_text",
Text: content,
@@ -250,7 +266,7 @@ func buildResponsesObject(
ID: generateOutputItemID("msg"),
Type: "message",
Role: "assistant",
- Status: "completed",
+ Status: status,
Content: []ResponseContentPart{{
Type: "output_text",
Text: "",
@@ -258,16 +274,22 @@ func buildResponsesObject(
})
}
+ var incompleteDetails *ResponsesIncompleteDetails
+ if incompleteReason != "" {
+ incompleteDetails = &ResponsesIncompleteDetails{Reason: incompleteReason}
+ }
+
return &ResponsesObject{
ID: id,
Object: "response",
CreatedAt: time.Now().Unix(),
- Status: "completed",
+ Status: status,
Model: model,
Output: output,
Usage: ResponsesUsage{InputTokens: inputTokens, OutputTokens: outputTokens, TotalTokens: inputTokens + outputTokens},
PreviousResponseID: req.PreviousResponseID,
Metadata: req.Metadata,
+ IncompleteDetails: incompleteDetails,
}
}
@@ -316,6 +338,10 @@ func (h *Handler) handleResponsesStream(
var lastErr error
responseStarted := false
reqStart := time.Now()
+ send("response.in_progress", map[string]interface{}{
+ "type": "response.in_progress",
+ "response": initial,
+ })
for attempt := 0; attempt < maxAccountRetryAttempts; attempt++ {
account := h.pool.GetNextForModelExcluding(model, excluded)
@@ -329,11 +355,6 @@ func (h *Handler) handleResponsesStream(
continue
}
- send("response.in_progress", map[string]interface{}{
- "type": "response.in_progress",
- "response": initial,
- })
-
var (
fullText strings.Builder
reasoningText strings.Builder
@@ -377,101 +398,138 @@ func (h *Handler) handleResponsesStream(
})
}
- callback := &KiroStreamCallback{
- OnText: func(text string, isThinking bool) {
- if text == "" {
- return
- }
- if isThinking {
- reasoningText.WriteString(text)
- return
- }
- fullText.WriteString(text)
- ensureMessageStarted()
- send("response.output_text.delta", map[string]interface{}{
- "type": "response.output_text.delta",
- "item_id": messageItemID,
- "output_index": outputIndex,
- "content_index": contentIndex,
- "delta": text,
- })
- responseStarted = true
- },
- OnToolUse: func(tu KiroToolUse) {
- if messageStarted {
- send("response.content_part.done", map[string]interface{}{
- "type": "response.content_part.done",
- "item_id": messageItemID,
- "output_index": outputIndex,
- "content_index": contentIndex,
- "part": map[string]interface{}{
- "type": "output_text",
- "text": fullText.String(),
- },
- })
- send("response.output_item.done", map[string]interface{}{
- "type": "response.output_item.done",
- "output_index": outputIndex,
- "item": map[string]interface{}{
- "id": messageItemID,
- "type": "message",
- "role": "assistant",
- "status": "completed",
- "content": []map[string]interface{}{{
- "type": "output_text",
- "text": fullText.String(),
- }},
- },
- })
- messageStarted = false
- outputIndex++
- }
-
- toolUses = append(toolUses, tu)
- args, _ := json.Marshal(tu.Input)
- fcID := generateOutputItemID("fc")
- send("response.output_item.added", map[string]interface{}{
- "type": "response.output_item.added",
- "output_index": outputIndex,
- "item": map[string]interface{}{
- "id": fcID,
- "type": "function_call",
- "status": "in_progress",
- "call_id": tu.ToolUseID,
- "name": tu.Name,
- "arguments": "",
+ var integrity streamIntegrityState
+
+ err := runKiroWithIntegrityRetry(account, payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ return &KiroStreamCallback{
+ OnText: func(text string, isThinking bool) {
+ if text == "" {
+ return
+ }
+ integrity.observeText(text, isThinking)
+ if isThinking {
+ reasoningText.WriteString(text)
+ return
+ }
+ fullText.WriteString(text)
+ ensureMessageStarted()
+ send("response.output_text.delta", map[string]interface{}{
+ "type": "response.output_text.delta",
+ "item_id": messageItemID,
+ "output_index": outputIndex,
+ "content_index": contentIndex,
+ "delta": text,
+ })
+ responseStarted = true
},
- })
- send("response.function_call_arguments.delta", map[string]interface{}{
- "type": "response.function_call_arguments.delta",
- "item_id": fcID,
- "output_index": outputIndex,
- "delta": string(args),
- })
- send("response.output_item.done", map[string]interface{}{
- "type": "response.output_item.done",
- "output_index": outputIndex,
- "item": map[string]interface{}{
- "id": fcID,
- "type": "function_call",
- "status": "completed",
- "call_id": tu.ToolUseID,
- "name": tu.Name,
- "arguments": string(args),
+ OnToolUse: func(tu KiroToolUse) {
+ integrity.observeToolUse()
+ if messageStarted {
+ messageStatus, _ := mapResponsesCompletion(integrity.StopReason)
+ send("response.content_part.done", map[string]interface{}{
+ "type": "response.content_part.done",
+ "item_id": messageItemID,
+ "output_index": outputIndex,
+ "content_index": contentIndex,
+ "part": map[string]interface{}{
+ "type": "output_text",
+ "text": fullText.String(),
+ },
+ })
+ send("response.output_item.done", map[string]interface{}{
+ "type": "response.output_item.done",
+ "output_index": outputIndex,
+ "item": map[string]interface{}{
+ "id": messageItemID,
+ "type": "message",
+ "role": "assistant",
+ "status": messageStatus,
+ "content": []map[string]interface{}{{
+ "type": "output_text",
+ "text": fullText.String(),
+ }},
+ },
+ })
+ messageStarted = false
+ outputIndex++
+ }
+
+ toolUses = append(toolUses, tu)
+ args, _ := json.Marshal(tu.Input)
+ fcID := generateOutputItemID("fc")
+ send("response.output_item.added", map[string]interface{}{
+ "type": "response.output_item.added",
+ "output_index": outputIndex,
+ "item": map[string]interface{}{
+ "id": fcID,
+ "type": "function_call",
+ "status": "in_progress",
+ "call_id": tu.ToolUseID,
+ "name": tu.Name,
+ "arguments": "",
+ },
+ })
+ send("response.function_call_arguments.delta", map[string]interface{}{
+ "type": "response.function_call_arguments.delta",
+ "item_id": fcID,
+ "output_index": outputIndex,
+ "delta": string(args),
+ })
+ send("response.output_item.done", map[string]interface{}{
+ "type": "response.output_item.done",
+ "output_index": outputIndex,
+ "item": map[string]interface{}{
+ "id": fcID,
+ "type": "function_call",
+ "status": "completed",
+ "call_id": tu.ToolUseID,
+ "name": tu.Name,
+ "arguments": string(args),
+ },
+ })
+ outputIndex++
+ responseStarted = true
},
- })
- outputIndex++
- responseStarted = true
+ OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
+ OnCredits: func(c float64) { credits = c },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ },
+ OnStopReason: func(reason string) { integrity.observeStopReason(reason) },
+ }
},
- OnComplete: func(inTok, outTok int) { inputTokens = inTok; outputTokens = outTok },
- OnCredits: func(c float64) { credits = c },
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(model)) / 100.0)
+ func() {
+ fullText.Reset()
+ reasoningText.Reset()
+ toolUses = nil
+ inputTokens, outputTokens, credits, realInputTokens = 0, 0, 0, 0
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ func() bool { return !responseStarted },
+ )
if err != nil {
+ if isStreamIntegrityError(err) {
+ if !responseStarted {
+ lastErr = err
+ excluded[account.ID] = true
+ continue
+ }
+ // Partial content already reached the client — do not forge completed.
+ send("response.failed", map[string]interface{}{
+ "type": "response.failed",
+ "response": map[string]interface{}{
+ "id": respID,
+ "status": "failed",
+ "error": map[string]string{
+ "type": "server_error",
+ "message": err.Error(),
+ },
+ },
+ })
+ h.recordFailureWithDetails("responses", model, account.ID, err)
+ return
+ }
if !responseStarted {
lastErr = err
excluded[account.ID] = true
@@ -499,6 +557,8 @@ func (h *Handler) handleResponsesStream(
reasoning = ""
}
+ responseStatus, _ := mapResponsesCompletion(integrity.StopReason)
+
if messageStarted {
send("response.content_part.done", map[string]interface{}{
"type": "response.content_part.done",
@@ -517,7 +577,7 @@ func (h *Handler) handleResponsesStream(
"id": messageItemID,
"type": "message",
"role": "assistant",
- "status": "completed",
+ "status": responseStatus,
"content": []map[string]interface{}{{
"type": "output_text",
"text": finalContent,
@@ -538,7 +598,7 @@ func (h *Handler) handleResponsesStream(
h.pool.UpdateStats(account.ID, inputTokens+outputTokens, credits)
h.recordSuccessLog("responses", model, account.ID, inputTokens+outputTokens, credits, time.Since(reqStart).Milliseconds())
- respObj := buildResponsesObject(respID, model, finalContent, toolUses, inputTokens, outputTokens, req)
+ respObj := buildResponsesObject(respID, model, finalContent, toolUses, inputTokens, outputTokens, integrity.StopReason, req)
respObj.CreatedAt = createdAt
respObj.StoredInput = storedInput
respObj.Instructions = req.Instructions
@@ -549,8 +609,9 @@ func (h *Handler) handleResponsesStream(
}
}
- send("response.completed", map[string]interface{}{
- "type": "response.completed",
+ completionEvent := "response." + respObj.Status
+ send(completionEvent, map[string]interface{}{
+ "type": completionEvent,
"response": respObj,
})
fmt.Fprintf(w, "data: [DONE]\n\n")
diff --git a/proxy/responses_handler_test.go b/proxy/responses_handler_test.go
index 72003eb9..f4a13071 100644
--- a/proxy/responses_handler_test.go
+++ b/proxy/responses_handler_test.go
@@ -10,6 +10,7 @@ import (
"net/http/httptest"
"path/filepath"
"strings"
+ "sync/atomic"
"testing"
"time"
)
@@ -234,6 +235,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)()
@@ -312,6 +316,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)()
@@ -365,6 +372,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)()
@@ -391,3 +401,103 @@ func TestResponsesStreamSSE(t *testing.T) {
t.Fatalf("expected stream content delta, got:\n%s", bodyStr)
}
}
+
+func TestResponsesStreamMaxTokensIsIncomplete(t *testing.T) {
+ h, cleanup := setupResponsesTestHandler(t)
+ defer cleanup()
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "partial output"}))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "MAX_TOKENS"}))
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(t, server)()
+
+ req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(
+ `{"model":"claude-sonnet-4.5","input":"write a lot","stream":true,"store":false}`,
+ ))
+ rec := httptest.NewRecorder()
+ h.handleOpenAIResponses(rec, req)
+ body := rec.Body.String()
+
+ if !strings.Contains(body, "event: response.incomplete") {
+ t.Fatalf("missing response.incomplete event:\n%s", body)
+ }
+ if !strings.Contains(body, `"status":"incomplete"`) ||
+ !strings.Contains(body, `"incomplete_details":{"reason":"max_output_tokens"}`) {
+ t.Fatalf("missing incomplete status details:\n%s", body)
+ }
+}
+
+func TestResponsesStreamMaxTokensMarksPreToolMessageIncomplete(t *testing.T) {
+ h, cleanup := setupResponsesTestHandler(t)
+ defer cleanup()
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ _, _ = w.Write(awsEventStreamFrame(t, "assistantResponseEvent", map[string]interface{}{"content": "partial text"}))
+ _, _ = w.Write(awsEventStreamFrame(t, "toolUseEvent", map[string]interface{}{
+ "toolUseId": "call_1", "name": "read_file", "input": `{"path":"a.go"}`, "stop": true,
+ }))
+ _, _ = w.Write(awsEventStreamFrame(t, "metadataEvent", map[string]interface{}{"stopReason": "MAX_TOKENS"}))
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(t, server)()
+
+ req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(
+ `{"model":"claude-sonnet-4.5","input":"write a lot","stream":true,"store":false}`,
+ ))
+ rec := httptest.NewRecorder()
+ h.handleOpenAIResponses(rec, req)
+
+ foundMessage := false
+ for _, line := range strings.Split(rec.Body.String(), "\n") {
+ if !strings.HasPrefix(line, "data: {") {
+ continue
+ }
+ var event map[string]interface{}
+ if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &event); err != nil {
+ continue
+ }
+ if event["type"] != "response.output_item.done" {
+ continue
+ }
+ item, _ := event["item"].(map[string]interface{})
+ if item["type"] != "message" {
+ continue
+ }
+ foundMessage = true
+ if item["status"] != "incomplete" {
+ t.Fatalf("pre-tool message status = %v, body:\n%s", item["status"], rec.Body.String())
+ }
+ }
+ if !foundMessage {
+ t.Fatalf("missing pre-tool message completion event:\n%s", rec.Body.String())
+ }
+}
+
+func TestResponsesIntegrityExhaustionRotatesOnce(t *testing.T) {
+ h, cleanup := setupResponsesTestHandler(t)
+ defer cleanup()
+
+ var hits atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ hits.Add(1)
+ w.WriteHeader(http.StatusOK)
+ }))
+ defer server.Close()
+ defer swapKiroEndpointsForTest(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)
+
+ want := int32(maxSameAccountStreamRetries + 1)
+ if hits.Load() != want {
+ t.Fatalf("empty account called %d times, want one same-account retry budget (%d)", hits.Load(), want)
+ }
+ if count := strings.Count(rec.Body.String(), "event: response.in_progress"); count != 1 {
+ t.Fatalf("response.in_progress emitted %d times:\n%s", count, rec.Body.String())
+ }
+}
diff --git a/proxy/responses_types.go b/proxy/responses_types.go
index c36413a0..4c2dd688 100644
--- a/proxy/responses_types.go
+++ b/proxy/responses_types.go
@@ -17,20 +17,21 @@ type ResponsesRequest struct {
}
type ResponsesObject struct {
- ID string `json:"id"`
- Object string `json:"object"`
- CreatedAt int64 `json:"created_at"`
- Status string `json:"status"`
- Model string `json:"model"`
- Output []ResponseOutputItem `json:"output"`
- Usage ResponsesUsage `json:"usage"`
- PreviousResponseID string `json:"previous_response_id,omitempty"`
- Metadata map[string]string `json:"metadata,omitempty"`
- Error *ResponsesError `json:"error,omitempty"`
- Instructions string `json:"instructions,omitempty"`
- StoredInput json.RawMessage `json:"-"`
- StoredInstr string `json:"-"`
- StoredAt int64 `json:"stored_at,omitempty"`
+ ID string `json:"id"`
+ Object string `json:"object"`
+ CreatedAt int64 `json:"created_at"`
+ Status string `json:"status"`
+ Model string `json:"model"`
+ Output []ResponseOutputItem `json:"output"`
+ Usage ResponsesUsage `json:"usage"`
+ PreviousResponseID string `json:"previous_response_id,omitempty"`
+ Metadata map[string]string `json:"metadata,omitempty"`
+ Error *ResponsesError `json:"error,omitempty"`
+ IncompleteDetails *ResponsesIncompleteDetails `json:"incomplete_details,omitempty"`
+ Instructions string `json:"instructions,omitempty"`
+ StoredInput json.RawMessage `json:"-"`
+ StoredInstr string `json:"-"`
+ StoredAt int64 `json:"stored_at,omitempty"`
}
type ResponseOutputItem struct {
@@ -60,3 +61,7 @@ type ResponsesError struct {
Code string `json:"code,omitempty"`
Message string `json:"message"`
}
+
+type ResponsesIncompleteDetails struct {
+ Reason string `json:"reason"`
+}
diff --git a/proxy/websearch.go b/proxy/websearch.go
index 2996e459..4eb26f55 100644
--- a/proxy/websearch.go
+++ b/proxy/websearch.go
@@ -281,7 +281,7 @@ func callMcpAPI(account *config.Account, mcpReq *McpRequest) (*McpResponse, erro
if parsed, perr := url.Parse(endpoint); perr == nil {
host = parsed.Host
}
- headerValues := buildStreamingHeaderValues(account, host)
+ headerValues := buildLegacyStreamingHeaderValues(account, host)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "*/*")
diff --git a/proxy/websearch_loop.go b/proxy/websearch_loop.go
index 352485f0..e2603c3a 100644
--- a/proxy/websearch_loop.go
+++ b/proxy/websearch_loop.go
@@ -25,11 +25,13 @@ const maxWebSearchRounds = 5
// webSearchRoundOutcome is one buffered generateAssistantResponse round.
type webSearchRoundOutcome struct {
- text string
- toolUses []KiroToolUse
- inputTokens int
- credits float64
- stopReasonOverride string
+ text string
+ toolUses []KiroToolUse
+ inputTokens int
+ credits float64
+ stopReason string
+ convID string
+ contID string
}
// runWebSearchLoop is the mixed-tools entry point.
@@ -50,7 +52,7 @@ func (h *Handler) runWebSearchLoop(w http.ResponseWriter, req *ClaudeRequest, th
// Allow one extra iteration so a terminal flush can run after the last
// search-only round (same pattern as 0..=MAX_WEB_SEARCH_ROUNDS in kiro-rs).
for roundIdx := 0; roundIdx <= maxUses; roundIdx++ {
- round, account, err := h.callUpstreamForWebSearch(&working, thinking, fallbackInput)
+ round, account, err := h.callUpstreamForWebSearch(&working, thinking, fallbackInput, apiKeyID)
if err != nil {
logger.Warnf("[WebSearchLoop] upstream round %d failed: %v", roundIdx, err)
accountID := ""
@@ -115,7 +117,7 @@ func (h *Handler) runWebSearchLoop(w http.ResponseWriter, req *ClaudeRequest, th
}
content := buildFlushContent(presentation, round.text, round.toolUses, searched)
- stopReason := resolveFlushStopReason(round.stopReasonOverride, round.toolUses, content)
+ stopReason := resolveFlushStopReason(round.stopReason, round.toolUses, content)
outputTokens := estimateContentBlocksTokens(content)
inputTokens := round.inputTokens
if inputTokens <= 0 {
@@ -128,6 +130,11 @@ func (h *Handler) runWebSearchLoop(w http.ResponseWriter, req *ClaudeRequest, th
}
h.recordSuccessForApiKey(apiKeyID, inputTokens, outputTokens, totalCredits)
h.recordSuccessLog("claude", req.Model, lastAccountID, inputTokens+outputTokens, totalCredits, time.Since(reqStart).Milliseconds())
+ // Seal the exact assistant content the client receives so the next
+ // client-side tool_result continuation reuses the same IDs.
+ if round.convID != "" && round.contID != "" {
+ sealClaudeSessionFromContent(apiKeyID, req.Model, round.convID, round.contID, req.Messages, content)
+ }
if req.Stream {
h.renderWebSearchLoopSSE(w, req.Model, content, stopReason, inputTokens, outputTokens)
@@ -141,8 +148,14 @@ func (h *Handler) runWebSearchLoop(w http.ResponseWriter, req *ClaudeRequest, th
}
// callUpstreamForWebSearch converts the Claude request and buffers one Kiro stream.
-func (h *Handler) callUpstreamForWebSearch(req *ClaudeRequest, thinking bool, estimatedInputTokens int) (*webSearchRoundOutcome, *config.Account, error) {
+func (h *Handler) callUpstreamForWebSearch(req *ClaudeRequest, thinking bool, estimatedInputTokens int, apiKeyID string) (*webSearchRoundOutcome, *config.Account, error) {
payload := ClaudeToKiro(req, thinking)
+ modelID := MapModel(req.Model)
+ convID, contID, _ := resolveSessionIDs(apiKeyID, modelID, req.Messages)
+ payload.ConversationState.ConversationID = convID
+ payload.ConversationState.AgentContinuationId = contID
+ attachStoredReasoning(payload.ConversationState.History, apiKeyID, modelID, convID)
+
excluded := make(map[string]bool)
var lastErr error
@@ -158,48 +171,71 @@ func (h *Handler) callUpstreamForWebSearch(req *ClaudeRequest, thinking bool, es
continue
}
- var text string
+ var text, thinkingText string
var toolUses []KiroToolUse
var inputTokens int
var credits float64
var realInputTokens int
- var stopOverride string
-
- callback := &KiroStreamCallback{
- OnText: func(t string, isThinking bool) {
- if isThinking {
- return
+ var integrity streamIntegrityState
+ var stamp reasoningCapture
+
+ err := runKiroWithIntegrityRetry(
+ account,
+ payload,
+ &integrity,
+ func() *KiroStreamCallback {
+ onReasoningMeta := stamp.meta()
+ return &KiroStreamCallback{
+ OnText: func(t string, isThinking bool) {
+ integrity.observeText(t, isThinking)
+ if isThinking {
+ thinkingText += t
+ return
+ }
+ text += t
+ },
+ OnToolUse: func(tu KiroToolUse) {
+ toolUses = append(toolUses, tu)
+ integrity.observeToolUse()
+ },
+ OnComplete: func(inTok, _ int) {
+ inputTokens = inTok
+ },
+ OnCredits: func(c float64) {
+ credits = c
+ },
+ OnContextUsage: func(pct float64) {
+ realInputTokens = int(pct * float64(getContextWindowSize(req.Model)) / 100.0)
+ if pct >= 100.0 {
+ integrity.observeStopReason("model_context_window_exceeded")
+ }
+ },
+ OnStopReason: func(reason string) {
+ if integrity.StopReason == "" {
+ integrity.observeStopReason(reason)
+ }
+ },
+ OnReasoningMeta: func(signature, redacted string) {
+ integrity.SawReasoning = true
+ onReasoningMeta(signature, redacted)
+ },
}
- text += t
- },
- OnToolUse: func(tu KiroToolUse) {
- toolUses = append(toolUses, tu)
- },
- OnComplete: func(inTok, outTok int) {
- inputTokens = inTok
- _ = outTok
},
- OnCredits: func(c float64) {
- credits = c
- },
- OnContextUsage: func(pct float64) {
- realInputTokens = int(pct * float64(getContextWindowSize(req.Model)) / 100.0)
- if pct >= 100.0 {
- stopOverride = "model_context_window_exceeded"
- }
- },
- OnError: func(err error) {
- if err != nil {
- lastErr = err
- }
+ func() {
+ text, thinkingText = "", ""
+ toolUses = nil
+ inputTokens, realInputTokens = 0, 0
+ credits = 0
+ stamp.reset()
},
- }
-
- err := CallKiroAPI(account, payload, callback)
+ nil,
+ )
if err != nil {
lastErr = err
excluded[account.ID] = true
- h.handleAccountFailure(account, err)
+ if !isStreamIntegrityError(err) {
+ h.handleAccountFailure(account, err)
+ }
continue
}
@@ -209,12 +245,20 @@ func (h *Handler) callUpstreamForWebSearch(req *ClaudeRequest, thinking bool, es
inputTokens = estimatedInputTokens
}
+ // Intermediate search-only rounds still seal the raw upstream shape so
+ // the next loop iteration reuses the same IDs. The final client-facing
+ // flush seals the visible server_tool_use content separately.
+ stamp.sealIfPresent(apiKeyID, modelID, convID, currentTurnParentUser(payload), text, toolUses, thinkingText)
+ sealClaudeSession(apiKeyID, req.Model, convID, contID, req.Messages, text, "", toolUses, "", true)
+
return &webSearchRoundOutcome{
- text: text,
- toolUses: toolUses,
- inputTokens: inputTokens,
- credits: credits,
- stopReasonOverride: stopOverride,
+ text: text,
+ toolUses: toolUses,
+ inputTokens: inputTokens,
+ credits: credits,
+ stopReason: integrity.StopReason,
+ convID: convID,
+ contID: contID,
}, account, nil
}
@@ -386,23 +430,20 @@ func buildFlushContent(
// resolveFlushStopReason picks stop_reason for the flushed response.
// web_search-only rounds end as end_turn; client tool_use yields tool_use.
-func resolveFlushStopReason(override string, toolUses []KiroToolUse, content []map[string]interface{}) string {
- if override != "" {
- return override
- }
- for _, c := range content {
- if c["type"] == "tool_use" {
- if name, _ := c["name"].(string); name != webSearchToolName {
- return "tool_use"
+func resolveFlushStopReason(reason string, toolUses []KiroToolUse, content []map[string]interface{}) string {
+ for _, block := range content {
+ if block["type"] == "tool_use" {
+ if name, _ := block["name"].(string); name != webSearchToolName {
+ return mapClaudeStopReason(reason, 1)
}
}
}
- for _, tu := range toolUses {
- if tu.Name != webSearchToolName {
- return "tool_use"
+ for _, toolUse := range toolUses {
+ if toolUse.Name != webSearchToolName {
+ return mapClaudeStopReason(reason, 1)
}
}
- return "end_turn"
+ return mapClaudeStopReason(reason, 0)
}
func estimateContentBlocksTokens(content []map[string]interface{}) int {
diff --git a/proxy/websearch_test.go b/proxy/websearch_test.go
index d93a1ee1..11976573 100644
--- a/proxy/websearch_test.go
+++ b/proxy/websearch_test.go
@@ -490,7 +490,10 @@ func TestAppendSearchRound_ContentSurvivesClaudeToKiro(t *testing.T) {
if aText != "Looking it up." || len(aUses) != 1 || aUses[0].ToolUseID != "toolu_ws_1" {
t.Fatalf("assistant extract failed: text=%q uses=%+v", aText, aUses)
}
- _, _, uResults := extractClaudeUserContent(req.Messages[2].Content)
+ _, _, _, uResults, errMsg := extractClaudeUserContent(req.Messages[2].Content)
+ if errMsg != "" {
+ t.Fatalf("unexpected document error: %s", errMsg)
+ }
if len(uResults) != 1 || uResults[0].ToolUseID != "toolu_ws_1" {
t.Fatalf("user tool_result extract failed: %+v", uResults)
}
@@ -529,7 +532,10 @@ func TestExtractClaudeContent_MapSliceShape(t *testing.T) {
user := []map[string]interface{}{
{"type": "tool_result", "tool_use_id": "id1", "content": "ok"},
}
- _, _, results := extractClaudeUserContent(user)
+ _, _, _, results, errMsg := extractClaudeUserContent(user)
+ if errMsg != "" {
+ t.Fatalf("unexpected document error: %s", errMsg)
+ }
if len(results) != 1 || results[0].ToolUseID != "id1" {
t.Fatalf("map-slice user content: %+v", results)
}
@@ -559,3 +565,234 @@ func TestParseSearchResults_EmptyResultsArrayIsValid(t *testing.T) {
t.Fatalf("results len = %d", len(results.Results))
}
}
+
+func TestClaudeToKiroFoldsServerWebSearchResultsIntoHistory(t *testing.T) {
+ snippet := "Go 1.18 introduced generics"
+ req := &ClaudeRequest{
+ Model: "claude-sonnet-4.5",
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: "search and then run bash"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": "I searched first."},
+ map[string]interface{}{
+ "type": "server_tool_use",
+ "id": "srvtoolu_1",
+ "name": "web_search",
+ "input": map[string]interface{}{"query": "golang generics"},
+ },
+ map[string]interface{}{
+ "type": "web_search_tool_result",
+ "content": []interface{}{
+ map[string]interface{}{
+ "type": "web_search_result",
+ "title": "Generics",
+ "url": "https://go.dev",
+ "encrypted_content": snippet,
+ },
+ },
+ },
+ map[string]interface{}{
+ "type": "tool_use",
+ "id": "toolu_bash",
+ "name": "Bash",
+ "input": map[string]interface{}{"command": "ls"},
+ },
+ }},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "toolu_bash", "content": "a.go"},
+ }},
+ },
+ }
+
+ payload := ClaudeToKiro(req, false)
+ if payload == nil {
+ t.Fatal("nil payload")
+ }
+ history := payload.ConversationState.History
+ if len(history) == 0 {
+ t.Fatal("expected history")
+ }
+ last := history[len(history)-1]
+ if last.AssistantResponseMessage == nil {
+ t.Fatal("expected assistant history entry")
+ }
+ if !strings.Contains(last.AssistantResponseMessage.Content, "golang generics") {
+ t.Fatalf("assistant history missing folded search summary: %q", last.AssistantResponseMessage.Content)
+ }
+ if !strings.Contains(last.AssistantResponseMessage.Content, "https://go.dev") {
+ t.Fatalf("assistant history missing search source: %q", last.AssistantResponseMessage.Content)
+ }
+ if len(last.AssistantResponseMessage.ToolUses) != 1 || last.AssistantResponseMessage.ToolUses[0].ToolUseID != "toolu_bash" {
+ t.Fatalf("client tool uses = %+v, want only Bash", last.AssistantResponseMessage.ToolUses)
+ }
+}
+
+func TestClaudeToKiroFoldsEmptyServerWebSearchResults(t *testing.T) {
+ req := &ClaudeRequest{
+ Model: "claude-sonnet-4.5",
+ Messages: []ClaudeMessage{
+ {Role: "user", Content: "search and then run bash"},
+ {Role: "assistant", Content: []interface{}{
+ map[string]interface{}{"type": "text", "text": "I searched first."},
+ map[string]interface{}{
+ "type": "server_tool_use",
+ "id": "srvtoolu_empty",
+ "name": "web_search",
+ "input": map[string]interface{}{"query": "unlikely query"},
+ },
+ map[string]interface{}{
+ "type": "web_search_tool_result",
+ "content": []interface{}{},
+ },
+ map[string]interface{}{
+ "type": "tool_use",
+ "id": "toolu_bash",
+ "name": "Bash",
+ "input": map[string]interface{}{"command": "ls"},
+ },
+ }},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "toolu_bash", "content": "a.go"},
+ }},
+ },
+ }
+
+ payload := ClaudeToKiro(req, false)
+ last := payload.ConversationState.History[len(payload.ConversationState.History)-1]
+ if last.AssistantResponseMessage == nil {
+ t.Fatal("expected assistant history entry")
+ }
+ if !strings.Contains(last.AssistantResponseMessage.Content, "unlikely query") {
+ t.Fatalf("empty search missing query in folded summary: %q", last.AssistantResponseMessage.Content)
+ }
+ if !strings.Contains(last.AssistantResponseMessage.Content, "No results found.") {
+ t.Fatalf("empty search missing No results found text: %q", last.AssistantResponseMessage.Content)
+ }
+ if len(last.AssistantResponseMessage.ToolUses) != 1 || last.AssistantResponseMessage.ToolUses[0].ToolUseID != "toolu_bash" {
+ t.Fatalf("client tool uses = %+v, want only Bash", last.AssistantResponseMessage.ToolUses)
+ }
+}
+
+func TestSealClientVisibleWebSearchContentMatchesContinuation(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ const api = "key-ws-mixed"
+ const model = "claude-sonnet-4.5"
+ inbound := []ClaudeMessage{{Role: "user", Content: "search and then run bash"}}
+ conv, cont, _ := resolveSessionIDs(api, MapModel(model), inbound)
+
+ content := []map[string]interface{}{
+ {"type": "text", "text": "I searched first."},
+ {
+ "type": "server_tool_use",
+ "id": "srvtoolu_1",
+ "name": "web_search",
+ "input": map[string]interface{}{"query": "golang generics"},
+ },
+ {
+ "type": "web_search_tool_result",
+ "content": []interface{}{
+ map[string]interface{}{
+ "type": "web_search_result",
+ "title": "Generics",
+ "url": "https://go.dev",
+ "encrypted_content": "Go 1.18 introduced generics",
+ },
+ },
+ },
+ {
+ "type": "tool_use",
+ "id": "toolu_bash",
+ "name": "Bash",
+ "input": map[string]interface{}{"command": "ls"},
+ },
+ }
+ // Old path: seal the raw upstream tool_use shape. That must not match the
+ // client-visible continuation that stores server_tool_use blocks.
+ sealClaudeSession(api, model, conv, cont, inbound, "I searched first.", "",
+ []KiroToolUse{
+ {ToolUseID: "toolu_ws", Name: "web_search", Input: map[string]interface{}{"query": "golang generics"}},
+ {ToolUseID: "toolu_bash", Name: "Bash", Input: map[string]interface{}{"command": "ls"}},
+ }, "", true)
+ turn2 := []ClaudeMessage{
+ inbound[0],
+ {Role: "assistant", Content: content},
+ {Role: "user", Content: []interface{}{
+ map[string]interface{}{"type": "tool_result", "tool_use_id": "toolu_bash", "content": "a.go"},
+ }},
+ }
+ if _, _, reused := resolveSessionIDs(api, MapModel(model), turn2); reused {
+ t.Fatal("raw upstream tool shape must not match client-visible server_tool_use continuation")
+ }
+
+ // Target path: seal the exact client-visible content. Continuation must reuse.
+ globalSessionStore = newSessionStore()
+ conv, cont, _ = resolveSessionIDs(api, MapModel(model), inbound)
+ sealClaudeSessionFromContent(api, model, conv, cont, inbound, content)
+ gotConv, gotCont, reused := resolveSessionIDs(api, MapModel(model), turn2)
+ if !reused {
+ t.Fatal("client-visible flush content must seal the continuation chain")
+ }
+ if gotConv != conv || gotCont != cont {
+ t.Fatalf("ids = %s/%s, want %s/%s", gotConv, gotCont, conv, cont)
+ }
+}
+
+// Intermediate search-only rounds seal the raw upstream tool shape so the next
+// loop iteration can reuse the same conversation / continuation IDs. The final
+// client-facing flush still seals server_tool_use content separately.
+func TestIntermediateWebSearchRoundSealReusesNextLoopIDs(t *testing.T) {
+ globalSessionStore = newSessionStore()
+ globalReasoningStore = newReasoningStore()
+
+ const api = "key-ws-intermediate"
+ const model = "claude-sonnet-4.5"
+ inbound := []ClaudeMessage{{Role: "user", Content: "search golang generics"}}
+ conv, cont, reused := resolveSessionIDs(api, MapModel(model), inbound)
+ if reused {
+ t.Fatal("first web-search round must mint fresh ids")
+ }
+
+ round := &webSearchRoundOutcome{
+ text: "Looking it up.",
+ toolUses: []KiroToolUse{{
+ ToolUseID: "toolu_ws_1",
+ Name: webSearchToolName,
+ Input: map[string]interface{}{"query": "golang generics"},
+ }},
+ convID: conv,
+ contID: cont,
+ }
+ // Same seal the production intermediate path performs after a successful
+ // upstream round, before appendSearchRound mutates the working transcript.
+ stamp := &reasoningCapture{sig: "sig-ws-1"}
+ stamp.sealIfPresent(api, MapModel(model), conv, "search golang generics", round.text, round.toolUses, "hidden search plan")
+ sealClaudeSession(api, model, conv, cont, inbound, round.text, "", round.toolUses, "", true)
+
+ working := &ClaudeRequest{Model: model, Messages: append([]ClaudeMessage(nil), inbound...)}
+ presentation := make([]map[string]interface{}, 0)
+ snippet := "Go 1.18 introduced generics"
+ searched := []*WebSearchResults{{
+ Results: []WebSearchResult{{Title: "Generics", URL: "https://go.dev", Snippet: &snippet}},
+ }}
+ appendSearchRound(working, round, searched, &presentation)
+
+ gotConv, gotCont, reusedNext := resolveSessionIDs(api, MapModel(model), working.Messages)
+ if !reusedNext {
+ t.Fatal("next loop iteration must reuse ids after intermediate raw-tool seal")
+ }
+ if gotConv != conv || gotCont != cont {
+ t.Fatalf("next-loop ids = %s/%s, want %s/%s", gotConv, gotCont, conv, cont)
+ }
+
+ // Final client-visible flush still needs its own seal; raw intermediate seal
+ // must not accidentally match a server_tool_use continuation from the client.
+ content := buildFlushContent(presentation, "done", nil, nil)
+ clientTurn := []ClaudeMessage{
+ inbound[0],
+ {Role: "assistant", Content: content},
+ {Role: "user", Content: "thanks"},
+ }
+ if _, _, reusedClient := resolveSessionIDs(api, MapModel(model), clientTurn); reusedClient {
+ t.Fatal("intermediate raw-tool seal must not match client-visible server_tool_use chain")
+ }
+}