From ddf04f335f7ffccb4ae21ef96b41fe0f09f28b6f Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 15:32:33 +0800 Subject: [PATCH 01/20] feat(plugin): add Kiro provider --- .env.example | 8 +- config.example.yaml | 5 + docker-compose.yml | 1 + examples/plugin/Makefile | 9 +- examples/plugin/README.md | 1 + examples/plugin/kiro/README.md | 148 +++++ examples/plugin/kiro/THIRD_PARTY_NOTICES.md | 15 + examples/plugin/kiro/go/auth.go | 161 ++++++ examples/plugin/kiro/go/eventstream.go | 330 +++++++++++ examples/plugin/kiro/go/executor.go | 451 +++++++++++++++ examples/plugin/kiro/go/go.mod | 10 + examples/plugin/kiro/go/go.sum | 4 + examples/plugin/kiro/go/kiro_test.go | 575 ++++++++++++++++++++ examples/plugin/kiro/go/main.go | 422 ++++++++++++++ examples/plugin/kiro/go/model_discovery.go | 303 +++++++++++ examples/plugin/kiro/go/protocol.go | 165 ++++++ examples/plugin/kiro/go/response.go | 296 ++++++++++ examples/plugin/kiro/go/token_estimator.go | 137 +++++ examples/plugin/kiro/go/translate.go | 405 ++++++++++++++ 19 files changed, 3441 insertions(+), 5 deletions(-) create mode 100644 examples/plugin/kiro/README.md create mode 100644 examples/plugin/kiro/THIRD_PARTY_NOTICES.md create mode 100644 examples/plugin/kiro/go/auth.go create mode 100644 examples/plugin/kiro/go/eventstream.go create mode 100644 examples/plugin/kiro/go/executor.go create mode 100644 examples/plugin/kiro/go/go.mod create mode 100644 examples/plugin/kiro/go/go.sum create mode 100644 examples/plugin/kiro/go/kiro_test.go create mode 100644 examples/plugin/kiro/go/main.go create mode 100644 examples/plugin/kiro/go/model_discovery.go create mode 100644 examples/plugin/kiro/go/protocol.go create mode 100644 examples/plugin/kiro/go/response.go create mode 100644 examples/plugin/kiro/go/token_estimator.go create mode 100644 examples/plugin/kiro/go/translate.go diff --git a/.env.example b/.env.example index 5b0546f4c59..a8b03899d6b 100644 --- a/.env.example +++ b/.env.example @@ -1,8 +1,12 @@ # Example environment configuration for CLIProxyAPI. # Copy this file to `.env` and uncomment the variables you need. # -# NOTE: Environment variables are only required when using remote storage options. -# For local file-based storage (default), no environment variables need to be set. +# NOTE: Environment variables are required only for the features that use them. + +# ------------------------------------------------------------------------------ +# Kiro provider plugin (optional) +# ------------------------------------------------------------------------------ +# KIRO_API_KEY=replace-with-your-kiro-api-key # ------------------------------------------------------------------------------ # Management Web UI diff --git a/config.example.yaml b/config.example.yaml index acb689062a2..36c87b4e669 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -107,6 +107,11 @@ plugins: config2: "string" config3: 3 mode: "safe" # enum example: safe, fast + # kiro-go: + # enabled: true + # priority: 1 + # models: + # - "*" # expose every model discovered for each Kiro account # When true, disable high-overhead request logging and HTTP middleware features to reduce per-request memory usage under high concurrency. commercial-mode: false diff --git a/docker-compose.yml b/docker-compose.yml index 2205d30acac..d4c6cc27055 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -14,6 +14,7 @@ services: # - .env environment: DEPLOY: ${DEPLOY:-} + KIRO_API_KEY: ${KIRO_API_KEY:-} ports: - "8317:8317" - "8085:8085" diff --git a/examples/plugin/Makefile b/examples/plugin/Makefile index 78ff07a4f1f..201066a3803 100644 --- a/examples/plugin/Makefile +++ b/examples/plugin/Makefile @@ -1,4 +1,5 @@ EXAMPLES := simple model auth frontend-auth executor protocol-format request-translator request-normalizer response-translator response-normalizer thinking usage cli management-api host-callback host-callback-auth-files host-model-callback claude-web-search-router +GO_ONLY_EXAMPLES := kiro LANGUAGES := go c rust BIN_DIR := $(CURDIR)/bin BUILD_DIR := $(BIN_DIR)/build @@ -19,12 +20,14 @@ RUST_DYLIB_PREFIX := lib RUST_DYLIB_EXT := so endif -.PHONY: build list clean +.PHONY: build build-kiro list clean -build: $(foreach example,$(EXAMPLES),$(foreach lang,$(LANGUAGES),$(BIN_DIR)/$(example)-$(lang).$(PLUGIN_EXT))) +build: $(foreach example,$(EXAMPLES),$(foreach lang,$(LANGUAGES),$(BIN_DIR)/$(example)-$(lang).$(PLUGIN_EXT))) $(foreach example,$(GO_ONLY_EXAMPLES),$(BIN_DIR)/$(example)-go.$(PLUGIN_EXT)) + +build-kiro: $(BIN_DIR)/kiro-go.$(PLUGIN_EXT) list: - @$(foreach example,$(EXAMPLES),$(foreach lang,$(LANGUAGES),echo $(example)/$(lang);)) + @$(foreach example,$(EXAMPLES),$(foreach lang,$(LANGUAGES),echo $(example)/$(lang);)) $(foreach example,$(GO_ONLY_EXAMPLES),echo $(example)/go;) clean: rm -rf $(BIN_DIR) diff --git a/examples/plugin/README.md b/examples/plugin/README.md index 2e7b2de0c49..8a37c4e4026 100644 --- a/examples/plugin/README.md +++ b/examples/plugin/README.md @@ -17,6 +17,7 @@ This directory contains standard dynamic library plugin examples for the CLIProx - `request-lifecycle/`: Go-only request admission example with concurrency control, active HTTP termination, and terminal callbacks. - `scheduler/`: Go-only scheduler that can select a configured auth ID, delegate to a built-in scheduler, or deny picks. - `claude-web-search-router/`: ModelRouter + executor for Claude Code built-in `web_search` (antigravity / codex / xai / Tavily). See `claude-web-search-router/README.md`. +- `kiro/`: API-key-only Kiro provider with Claude-native request/response handling. See `kiro/README.md`. - `response-translator/`: response translation capability only. - `response-normalizer/`: response normalization capability only. - `thinking/`: thinking applier capability only. diff --git a/examples/plugin/kiro/README.md b/examples/plugin/kiro/README.md new file mode 100644 index 00000000000..2a0026296cf --- /dev/null +++ b/examples/plugin/kiro/README.md @@ -0,0 +1,148 @@ +# Kiro Provider Plugin + +This Go plugin adds an API-key-only Kiro provider to CLIProxyAPI. It uses the +Kiro CLI runtime protocol and exposes Claude as its native protocol. The host +translates OpenAI Chat Completions and Responses API requests to and from Claude +when those entry points are used. + +The implementation deliberately does not include Kiro OAuth, device login, +refresh tokens, account selection, or its own HTTP client. CLIProxyAPI owns +credential selection, session affinity, cooldowns, protocol translation, proxy +configuration, and outbound request logging. + +## Build + +Build for the current platform: + +```bash +make -C examples/plugin build-kiro +``` + +On Linux, the extension is `.so`. The plugin must be built for the operating +system and architecture where CLIProxyAPI runs. For a Podman deployment on an +Apple Silicon Mac, build the Linux ARM64 plugin inside the Podman VM: + +```bash +podman run --rm \ + -v "$PWD:/src" \ + -w /src/examples/plugin/kiro/go \ + docker.io/library/golang:1.26 \ + sh -c 'CGO_ENABLED=1 go build -buildmode=c-shared -o /src/plugins/linux/arm64/kiro-go.so . && rm -f /src/plugins/linux/arm64/kiro-go.h' +``` + +The image name uses an OCI registry reference; the command is executed by +Podman and does not require a Docker daemon. + +## Configure the plugin + +Dynamic plugins are disabled by default. Add this to `config.yaml`: + +```yaml +plugins: + enabled: true + dir: "plugins" + configs: + kiro-go: + enabled: true + priority: 1 + models: + - "*" +``` + +The plugin discovers models separately for every Kiro account. Use `"*"` to +expose every model returned for that account, or list specific model IDs to use +the configured list as an allow-list: + +```yaml + models: + - "claude-sonnet-4.5" + - "claude-haiku-4.5" +``` + +Successful results are cached for five minutes. If a refresh fails, the plugin +keeps the last successful result. On a cold-start failure it falls back to the +configured models; when `"*"` is configured, the cold-start fallback is Claude +Sonnet 4.5 and Claude Haiku 4.5. + +## Add an auth record + +Create `auths/kiro-pro.json`: + +```json +{ + "type": "kiro", + "label": "kiro-pro", + "region": "us-east-1", + "api_key_env": "KIRO_API_KEY" +} +``` + +Set restrictive permissions: + +```bash +chmod 600 auths/kiro-pro.json +``` + +Pass `KIRO_API_KEY` to the CLIProxyAPI process or Podman container. The plugin +resolves the environment variable only when it sends a request. The resolved +key is not written back to the auth JSON. + +For a temporary local test: + +```bash +export KIRO_API_KEY='replace-me' +``` + +An `api_key` field is also accepted for installations that already protect the +auth directory, but `api_key_env` or a Podman secret is preferred. + +## Request flow + +1. The auth parser recognizes an auth JSON whose `type` is `kiro`. +2. During auth registration, the plugin calls Kiro's model-list service and + returns the discovered, allow-listed models for that account. +3. CLIProxyAPI selects a compatible Kiro auth using its normal routing and + affinity logic. +4. The host translates the incoming request to Claude format when necessary. +5. The plugin converts Claude messages, tools, tool results, images, system + instructions, and inference settings into Kiro conversation-state JSON. +6. The plugin asks the host HTTP bridge to POST to + `https://runtime..kiro.dev/` with API-key headers. +7. The plugin validates and decodes the AWS EventStream response. +8. It returns Claude JSON or emits Claude SSE. CLIProxyAPI translates that back + to the original client protocol when necessary. + +## Initial test cases + +Run these after providing a real key: + +1. List models and confirm the account's discovered models appear. If an + explicit allow-list is configured, confirm only its intersection appears. +2. Send a non-streaming `/v1/messages` text request. +3. Send the same request with `stream: true` and verify Anthropic SSE ordering. +4. Call `/v1/chat/completions` to verify host-side OpenAI translation. +5. Call `/v1/responses` to verify Responses API translation. +6. Exercise one client tool call and return its `tool_result`. +7. Send a small base64 image if the selected account model supports images. +8. Use an invalid key and verify a 401/403 does not leak the key. +9. Use a model outside the account entitlement and verify the upstream error is + surfaced without disabling unrelated credentials. +10. Cancel a streaming request and verify the upstream stream closes. + +## Limitations + +- The Kiro subscription API key is documented, but the raw runtime protocol is + not a public model API and may change. +- Model discovery uses Kiro's internal `ListAvailableModels` service rather + than a documented public model API and may change. +- Discovery is refreshed when the host registers an auth. The five-minute + cache avoids repeated calls but does not run its own background refresh. +- The host model registry has no credit-multiplier field, so discovery maps + names, descriptions, token limits, and modalities but does not expose Kiro's + per-model credit multiplier. +- `/v1/messages/count_tokens` returns an explicitly marked local estimate. +- Web search server tools are not forwarded in the initial implementation. +- Kiro client compatibility versions are configurable because upstream header + expectations may change. + +See `THIRD_PARTY_NOTICES.md` for implementation provenance. diff --git a/examples/plugin/kiro/THIRD_PARTY_NOTICES.md b/examples/plugin/kiro/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000000..8438a898106 --- /dev/null +++ b/examples/plugin/kiro/THIRD_PARTY_NOTICES.md @@ -0,0 +1,15 @@ +# Third-party notices + +The Kiro runtime integration was informed by the Kiro-Go project: + +- Project: https://github.com/Quorinex/Kiro-Go +- Reviewed revision: `f8f6071c9298a4266ad3e0c7e483d4a2510cbcaf` +- License: MIT + +Kiro-Go's API-key request path, Kiro conversation-state shapes, required runtime +headers, and AWS EventStream handling were used as interoperability references. +The CLIProxyAPI plugin implementation is adapted to CLIProxyAPI's plugin ABI, +host HTTP transport, auth storage, routing, and streaming bridge. + +Copyright notices and license terms remain available in the upstream project: +https://github.com/Quorinex/Kiro-Go/blob/f8f6071c9298a4266ad3e0c7e483d4a2510cbcaf/LICENSE diff --git a/examples/plugin/kiro/go/auth.go b/examples/plugin/kiro/go/auth.go new file mode 100644 index 00000000000..64c6043d431 --- /dev/null +++ b/examples/plugin/kiro/go/auth.go @@ -0,0 +1,161 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +var regionPattern = regexp.MustCompile(`^[a-z]{2}(?:-gov)?-[a-z]+-\d+$`) + +type kiroCredential struct { + Type string `json:"type"` + APIKey string `json:"api_key,omitempty"` + APIKeyEnv string `json:"api_key_env,omitempty"` + Region string `json:"region,omitempty"` + Label string `json:"label,omitempty"` + Prefix string `json:"prefix,omitempty"` + ProxyURL string `json:"proxy_url,omitempty"` + Disabled bool `json:"disabled,omitempty"` +} + +func parseAuth(raw []byte) ([]byte, error) { + var request pluginapi.AuthParseRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, fmt.Errorf("decode auth parse request: %w", errUnmarshal) + } + credential, handled, errCredential := decodeCredential(request.RawJSON) + if errCredential != nil { + return errorEnvelope("invalid_auth", errCredential.Error(), 0, false), nil + } + if !handled { + return okEnvelope(pluginapi.AuthParseResponse{Handled: false}) + } + storage, errMarshal := json.Marshal(credential) + if errMarshal != nil { + return nil, errMarshal + } + label := firstNonEmpty(credential.Label, strings.TrimSuffix(request.FileName, filepath.Ext(request.FileName)), keyFingerprint(credential)) + authID := strings.TrimSpace(request.FileName) + if authID == "" { + authID = "kiro-" + keyFingerprint(credential) + } + return okEnvelope(pluginapi.AuthParseResponse{ + Handled: true, + Auth: pluginapi.AuthData{ + Provider: pluginIdentifier, + ID: authID, + FileName: request.FileName, + Label: label, + Prefix: credential.Prefix, + ProxyURL: credential.ProxyURL, + Disabled: credential.Disabled, + StorageJSON: storage, + Metadata: map[string]any{ + "type": pluginIdentifier, + "auth_kind": "api_key", + "region": credential.Region, + "key_source": keySource(credential), + "key_present": credential.APIKey != "" || os.Getenv(credential.APIKeyEnv) != "", + }, + Attributes: map[string]string{ + "auth_kind": "api_key", + "region": credential.Region, + }, + }, + }) +} + +func refreshAuth(raw []byte) ([]byte, error) { + var request pluginapi.AuthRefreshRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, fmt.Errorf("decode auth refresh request: %w", errUnmarshal) + } + credential, handled, errCredential := decodeCredential(request.StorageJSON) + if errCredential != nil { + return errorEnvelope("invalid_auth", errCredential.Error(), 0, false), nil + } + if !handled { + return errorEnvelope("invalid_auth", "Kiro auth storage is not recognized", 0, false), nil + } + return okEnvelope(pluginapi.AuthRefreshResponse{Auth: pluginapi.AuthData{ + Provider: pluginIdentifier, + ID: request.AuthID, + StorageJSON: request.StorageJSON, + Metadata: request.Metadata, + Attributes: request.Attributes, + Label: credential.Label, + Prefix: credential.Prefix, + ProxyURL: credential.ProxyURL, + Disabled: credential.Disabled, + }}) +} + +func decodeCredential(raw []byte) (kiroCredential, bool, error) { + var credential kiroCredential + if len(raw) == 0 { + return credential, false, nil + } + if errUnmarshal := json.Unmarshal(raw, &credential); errUnmarshal != nil { + return credential, false, nil + } + credential.Type = strings.ToLower(strings.TrimSpace(credential.Type)) + if credential.Type != pluginIdentifier { + return credential, false, nil + } + credential.APIKey = strings.TrimSpace(credential.APIKey) + credential.APIKeyEnv = strings.TrimSpace(credential.APIKeyEnv) + credential.Region = firstNonEmpty(strings.ToLower(strings.TrimSpace(credential.Region)), "us-east-1") + credential.Label = strings.TrimSpace(credential.Label) + credential.Prefix = strings.TrimSpace(credential.Prefix) + credential.ProxyURL = strings.TrimSpace(credential.ProxyURL) + if credential.APIKey == "" && credential.APIKeyEnv == "" { + return credential, true, fmt.Errorf("either api_key or api_key_env is required") + } + if credential.APIKeyEnv != "" && strings.ContainsRune(credential.APIKeyEnv, '=') { + return credential, true, fmt.Errorf("api_key_env must be an environment variable name") + } + if !regionPattern.MatchString(credential.Region) { + return credential, true, fmt.Errorf("invalid Kiro region %q", credential.Region) + } + return credential, true, nil +} + +func resolveAPIKey(credential kiroCredential) (string, error) { + if credential.APIKey != "" { + return credential.APIKey, nil + } + key := strings.TrimSpace(os.Getenv(credential.APIKeyEnv)) + if key == "" { + return "", fmt.Errorf("environment variable %s is empty", credential.APIKeyEnv) + } + return key, nil +} + +func keySource(credential kiroCredential) string { + if credential.APIKey != "" { + return "auth_file" + } + return "env:" + credential.APIKeyEnv +} + +func keyFingerprint(credential kiroCredential) string { + seed := credential.APIKey + if seed == "" { + seed = "env:" + credential.APIKeyEnv + } + sum := sha256.Sum256([]byte("KiroAPIKey/" + seed)) + return hex.EncodeToString(sum[:6]) +} + +func machineID(key string) string { + sum := sha256.Sum256([]byte("KiroAPIKey/" + key)) + return hex.EncodeToString(sum[:]) +} diff --git a/examples/plugin/kiro/go/eventstream.go b/examples/plugin/kiro/go/eventstream.go new file mode 100644 index 00000000000..6783a4ee7ba --- /dev/null +++ b/examples/plugin/kiro/go/eventstream.go @@ -0,0 +1,330 @@ +package main + +import ( + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "hash/crc32" + "strconv" + "strings" +) + +const maxEventStreamMessageSize = 16 * 1024 * 1024 + +type eventStreamDecoder struct { + buffer []byte +} + +func (d *eventStreamDecoder) Feed(chunk []byte) ([]kiroEvent, error) { + if len(chunk) > 0 { + d.buffer = append(d.buffer, chunk...) + } + var events []kiroEvent + for { + if len(d.buffer) < 12 { + return events, nil + } + totalLength := int(binary.BigEndian.Uint32(d.buffer[0:4])) + headersLength := int(binary.BigEndian.Uint32(d.buffer[4:8])) + if totalLength < 16 || totalLength > maxEventStreamMessageSize { + return nil, fmt.Errorf("invalid Kiro event frame length %d", totalLength) + } + if headersLength < 0 || headersLength > totalLength-16 { + return nil, fmt.Errorf("invalid Kiro event header length %d", headersLength) + } + if crc32.ChecksumIEEE(d.buffer[:8]) != binary.BigEndian.Uint32(d.buffer[8:12]) { + return nil, errors.New("Kiro event prelude CRC mismatch") + } + if len(d.buffer) < totalLength { + return events, nil + } + frame := d.buffer[:totalLength] + if crc32.ChecksumIEEE(frame[:totalLength-4]) != binary.BigEndian.Uint32(frame[totalLength-4:]) { + return nil, errors.New("Kiro event message CRC mismatch") + } + headers, errHeaders := parseEventHeaders(frame[12 : 12+headersLength]) + if errHeaders != nil { + return nil, fmt.Errorf("decode Kiro event headers: %w", errHeaders) + } + payloadRaw := frame[12+headersLength : totalLength-4] + payload := make(map[string]any) + if len(payloadRaw) > 0 { + if errUnmarshal := json.Unmarshal(payloadRaw, &payload); errUnmarshal != nil { + return nil, fmt.Errorf("decode Kiro event payload: %w", errUnmarshal) + } + } + messageType := headers[":message-type"] + if messageType == "error" || messageType == "exception" { + message := firstStringField(payload, "message", "errorMessage") + if message == "" { + message = messageType + } + return nil, fmt.Errorf("Kiro upstream stream error: %s", message) + } + events = append(events, kiroEvent{Type: headers[":event-type"], Payload: payload}) + d.buffer = append(d.buffer[:0], d.buffer[totalLength:]...) + } +} + +func (d *eventStreamDecoder) Finish() error { + if len(d.buffer) != 0 { + return fmt.Errorf("incomplete Kiro event frame: %d buffered bytes", len(d.buffer)) + } + return nil +} + +func parseEventHeaders(data []byte) (map[string]string, error) { + headers := make(map[string]string) + for offset := 0; offset < len(data); { + nameLength := int(data[offset]) + offset++ + if nameLength == 0 || offset+nameLength >= len(data) { + return nil, errors.New("malformed header name") + } + name := string(data[offset : offset+nameLength]) + offset += nameLength + valueType := data[offset] + offset++ + var valueLength int + switch valueType { + case 0, 1: + continue + case 2: + valueLength = 1 + case 3: + valueLength = 2 + case 4: + valueLength = 4 + case 5, 8: + valueLength = 8 + case 9: + valueLength = 16 + case 6, 7: + if offset+2 > len(data) { + return nil, errors.New("truncated variable header length") + } + valueLength = int(binary.BigEndian.Uint16(data[offset : offset+2])) + offset += 2 + default: + return nil, fmt.Errorf("unsupported header value type %d", valueType) + } + if offset+valueLength > len(data) { + return nil, errors.New("truncated header value") + } + if valueType == 7 { + headers[name] = string(data[offset : offset+valueLength]) + } + offset += valueLength + } + return headers, nil +} + +type pendingToolUse struct { + ID string + Name string + InputJSON strings.Builder + Generated bool +} + +type pendingToolUses struct { + byID map[string]*pendingToolUse + order []string + lastID string +} + +func (p *pendingToolUses) accept(event map[string]any) ([]kiroToolUse, error) { + id := firstStringField(event, "toolUseId", "toolUseID", "tool_use_id", "id") + name := firstStringField(event, "name", "toolName", "tool_name") + stop := firstBoolField(event, "stop", "isStop", "done") + if p.byID == nil { + p.byID = make(map[string]*pendingToolUse) + } + if id == "" { + id = p.lastID + } + generated := false + if id == "" && name != "" { + id = "toolu_" + randomUUID() + generated = true + } + if id == "" { + return nil, nil + } + state := p.byID[id] + if state == nil && !generated && p.lastID != "" { + previous := p.byID[p.lastID] + if previous != nil && previous.Generated && (name == "" || previous.Name == name) { + oldID := previous.ID + delete(p.byID, oldID) + previous.ID = id + previous.Generated = false + p.byID[id] = previous + for index, existing := range p.order { + if existing == oldID { + p.order[index] = id + break + } + } + state = previous + } + } + if state == nil { + state = &pendingToolUse{ID: id, Name: name, Generated: generated} + p.byID[id] = state + p.order = append(p.order, id) + } + if state.Name == "" { + state.Name = name + } + p.lastID = id + if fragment, ok := event["input"].(string); ok { + state.InputJSON.WriteString(fragment) + } else if input, ok := event["input"].(map[string]any); ok { + raw, _ := json.Marshal(input) + state.InputJSON.Reset() + state.InputJSON.Write(raw) + } + if !stop { + return nil, nil + } + tool, errFinish := finishPendingTool(state) + if errFinish != nil { + return nil, errFinish + } + p.remove(id) + return []kiroToolUse{tool}, nil +} + +func (p *pendingToolUses) flush() ([]kiroToolUse, error) { + tools := make([]kiroToolUse, 0, len(p.order)) + for _, id := range append([]string(nil), p.order...) { + state := p.byID[id] + if state == nil || state.Name == "" { + continue + } + tool, errFinish := finishPendingTool(state) + if errFinish != nil { + return nil, errFinish + } + tools = append(tools, tool) + } + p.byID = nil + p.order = nil + p.lastID = "" + return tools, nil +} + +func (p *pendingToolUses) remove(id string) { + delete(p.byID, id) + for index, existing := range p.order { + if existing == id { + p.order = append(p.order[:index], p.order[index+1:]...) + break + } + } + if p.lastID == id { + p.lastID = "" + } +} + +func finishPendingTool(state *pendingToolUse) (kiroToolUse, error) { + input := make(map[string]any) + if state.InputJSON.Len() > 0 { + if errUnmarshal := json.Unmarshal([]byte(state.InputJSON.String()), &input); errUnmarshal != nil { + return kiroToolUse{}, fmt.Errorf("decode Kiro tool input: %w", errUnmarshal) + } + } + return kiroToolUse{ToolUseID: state.ID, Name: state.Name, Input: input}, nil +} + +func firstStringField(values map[string]any, keys ...string) string { + for _, key := range keys { + if value, ok := values[key].(string); ok && value != "" { + return value + } + } + return "" +} + +func firstBoolField(values map[string]any, keys ...string) bool { + for _, key := range keys { + if value, ok := values[key].(bool); ok { + return value + } + } + return false +} + +func readNumber(values map[string]any, keys ...string) (int, bool) { + value, ok := readFloat(values, keys...) + return int(value), ok +} + +func readFloat(values map[string]any, keys ...string) (float64, bool) { + for _, key := range keys { + switch value := values[key].(type) { + case float64: + return value, true + case int: + return float64(value), true + case int64: + return float64(value), true + case json.Number: + parsed, errParse := value.Float64() + if errParse == nil { + return parsed, true + } + case string: + parsed, errParse := strconv.ParseFloat(value, 64) + if errParse == nil { + return parsed, true + } + } + } + return 0, false +} + +func updateUsage(event map[string]any, inputTokens, outputTokens int) (int, int) { + candidates := []map[string]any{event} + collectUsageMaps(event, &candidates) + for _, candidate := range candidates { + if value, ok := readNumber(candidate, "outputTokens", "completionTokens", "totalOutputTokens", "output_tokens", "completion_tokens", "total_output_tokens"); ok { + outputTokens = value + } + if value, ok := readNumber(candidate, "inputTokens", "promptTokens", "totalInputTokens", "input_tokens", "prompt_tokens", "total_input_tokens"); ok { + inputTokens = value + continue + } + uncached, _ := readNumber(candidate, "uncachedInputTokens", "uncached_input_tokens") + cacheRead, _ := readNumber(candidate, "cacheReadInputTokens", "cache_read_input_tokens") + cacheWrite, _ := readNumber(candidate, "cacheWriteInputTokens", "cache_write_input_tokens", "cacheCreationInputTokens", "cache_creation_input_tokens") + if uncached+cacheRead+cacheWrite > 0 { + inputTokens = uncached + cacheRead + cacheWrite + continue + } + if total, ok := readNumber(candidate, "totalTokens", "total_tokens"); ok && total > outputTokens { + inputTokens = total - outputTokens + } + } + return inputTokens, outputTokens +} + +func collectUsageMaps(value any, candidates *[]map[string]any) { + switch typed := value.(type) { + case map[string]any: + for key, child := range typed { + normalized := strings.ToLower(key) + if normalized == "usage" || normalized == "tokenusage" || normalized == "token_usage" { + if nested, ok := child.(map[string]any); ok { + *candidates = append(*candidates, nested) + } + } + collectUsageMaps(child, candidates) + } + case []any: + for _, child := range typed { + collectUsageMaps(child, candidates) + } + } +} diff --git a/examples/plugin/kiro/go/executor.go b/examples/plugin/kiro/go/executor.go new file mode 100644 index 00000000000..6ac028b1a2c --- /dev/null +++ b/examples/plugin/kiro/go/executor.go @@ -0,0 +1,451 @@ +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "unicode/utf8" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +const ( + kiroAmzTarget = "AmazonCodeWhispererStreamingService.GenerateAssistantResponse" + maxAttempts = 2 +) + +type rpcHostHTTPRequest struct { + HostCallbackID string `json:"host_callback_id,omitempty"` + Method string `json:"method"` + URL string `json:"url"` + Headers http.Header `json:"headers,omitempty"` + Body []byte `json:"body,omitempty"` +} + +type rpcHostHTTPStreamResponse struct { + StatusCode int `json:"status_code"` + Headers http.Header `json:"headers,omitempty"` + StreamID string `json:"stream_id"` +} + +type rpcHostHTTPStreamReadRequest struct { + StreamID string `json:"stream_id"` +} + +type rpcHostHTTPStreamReadResponse struct { + Payload []byte `json:"payload,omitempty"` + Error string `json:"error,omitempty"` + Done bool `json:"done,omitempty"` +} + +type rpcHostHTTPStreamCloseRequest struct { + StreamID string `json:"stream_id"` +} + +type rpcStreamEmitRequest struct { + StreamID string `json:"stream_id"` + Payload []byte `json:"payload,omitempty"` + Error string `json:"error,omitempty"` +} + +type rpcStreamCloseRequest struct { + StreamID string `json:"stream_id"` + Error string `json:"error,omitempty"` +} + +type upstreamRequest struct { + URL string + Headers http.Header + Body []byte +} + +func execute(raw []byte) ([]byte, error) { + var request rpcExecutorRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, fmt.Errorf("decode executor request: %w", errUnmarshal) + } + upstream, payload, errPrepare := prepareUpstreamRequest(request.ExecutorRequest) + if errPrepare != nil { + return errorEnvelope("invalid_request", errPrepare.Error(), http.StatusBadRequest, false), nil + } + + var lastErr error + for attempt := 1; attempt <= maxAttempts; attempt++ { + accumulator := newAccumulator(payload) + emitted, status, errRun := consumeUpstream(request.HostCallbackID, upstream, attempt, func(events []kiroEvent) error { + for _, event := range events { + if _, errAccept := accumulator.accept(event); errAccept != nil { + return errAccept + } + } + return nil + }) + if errRun == nil { + if errFinish := accumulator.finish(); errFinish != nil { + lastErr = errFinish + } else { + body, errMarshal := accumulator.responseJSON() + if errMarshal != nil { + return nil, errMarshal + } + return okEnvelope(pluginapi.ExecutorResponse{ + Payload: body, + Headers: http.Header{"Content-Type": []string{"application/json"}}, + Metadata: map[string]any{"kiro_credits": accumulator.Credits}, + }) + } + } else { + lastErr = errRun + } + if emitted || !retryableUpstream(status, lastErr) || attempt == maxAttempts { + return upstreamErrorEnvelope(status, lastErr), nil + } + } + return upstreamErrorEnvelope(0, lastErr), nil +} + +func executeStream(raw []byte) ([]byte, error) { + var request rpcExecutorRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, fmt.Errorf("decode streaming executor request: %w", errUnmarshal) + } + if strings.TrimSpace(request.StreamID) == "" { + return errorEnvelope("invalid_request", "stream_id is required", http.StatusBadRequest, false), nil + } + upstream, payload, errPrepare := prepareUpstreamRequest(request.ExecutorRequest) + if errPrepare != nil { + return errorEnvelope("invalid_request", errPrepare.Error(), http.StatusBadRequest, false), nil + } + go runStream(request, upstream, payload) + return okEnvelope(map[string]any{"headers": http.Header{"Content-Type": []string{"text/event-stream"}}}) +} + +func runStream(request rpcExecutorRequest, upstream upstreamRequest, payload *kiroPayload) { + var terminalErr error + defer func() { + if recovered := recover(); recovered != nil { + terminalErr = fmt.Errorf("Kiro stream panic: %v", recovered) + } + closePluginStream(request.StreamID, terminalErr) + }() + + for attempt := 1; attempt <= maxAttempts; attempt++ { + accumulator := newAccumulator(payload) + writer := newClaudeSSEWriter(accumulator) + clientEmitted := false + emitted, status, errRun := consumeUpstream(request.HostCallbackID, upstream, attempt, func(events []kiroEvent) error { + for _, event := range events { + blocks, errAccept := accumulator.accept(event) + if errAccept != nil { + return errAccept + } + if len(blocks) == 0 { + continue + } + if !clientEmitted { + frames, errStart := writer.start() + if errStart != nil { + return errStart + } + if errEmit := emitFrames(request.StreamID, frames); errEmit != nil { + return errEmit + } + clientEmitted = true + } + frames, errBlocks := writer.blocks(blocks) + if errBlocks != nil { + return errBlocks + } + if errEmit := emitFrames(request.StreamID, frames); errEmit != nil { + return errEmit + } + } + return nil + }) + if errRun == nil { + blockCount := len(accumulator.Blocks) + if errFinish := accumulator.finish(); errFinish != nil { + errRun = errFinish + } else { + if !clientEmitted { + frames, errStart := writer.start() + if errStart != nil { + terminalErr = errStart + return + } + if errEmit := emitFrames(request.StreamID, frames); errEmit != nil { + terminalErr = errEmit + return + } + clientEmitted = true + } + if blockCount < len(accumulator.Blocks) { + frames, errBlocks := writer.blocks(accumulator.Blocks[blockCount:]) + if errBlocks != nil { + terminalErr = errBlocks + return + } + if errEmit := emitFrames(request.StreamID, frames); errEmit != nil { + terminalErr = errEmit + return + } + } + frames, errFinishFrames := writer.finish() + if errFinishFrames != nil { + terminalErr = errFinishFrames + return + } + terminalErr = emitFrames(request.StreamID, frames) + return + } + } + terminalErr = errRun + if clientEmitted || emitted || !retryableUpstream(status, errRun) || attempt == maxAttempts { + return + } + } +} + +func prepareUpstreamRequest(request pluginapi.ExecutorRequest) (upstreamRequest, *kiroPayload, error) { + credential, handled, errCredential := decodeCredential(request.StorageJSON) + if errCredential != nil { + return upstreamRequest{}, nil, errCredential + } + if !handled { + return upstreamRequest{}, nil, fmt.Errorf("selected auth is not a Kiro API-key credential") + } + key, errKey := resolveAPIKey(credential) + if errKey != nil { + return upstreamRequest{}, nil, errKey + } + payload, _, errTranslate := claudeToKiro(request.Payload, request.Model) + if errTranslate != nil { + return upstreamRequest{}, nil, errTranslate + } + body, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return upstreamRequest{}, nil, fmt.Errorf("encode Kiro request: %w", errMarshal) + } + endpoint := "https://runtime." + credential.Region + ".kiro.dev/" + parsed, errParse := url.Parse(endpoint) + if errParse != nil || parsed.Scheme != "https" || parsed.Hostname() != "runtime."+credential.Region+".kiro.dev" { + return upstreamRequest{}, nil, fmt.Errorf("invalid Kiro endpoint") + } + cfg := loadedConfig() + machine := machineID(key) + userAgent := fmt.Sprintf("aws-sdk-js/1.0.34 ua/2.1 os/%s lang/js md/nodejs#%s api/codewhispererstreaming#1.0.34 m/E KiroIDE-%s-%s", cfg.SystemVersion, cfg.NodeVersion, cfg.KiroVersion, machine) + amzUserAgent := fmt.Sprintf("aws-sdk-js/1.0.34 KiroIDE-%s-%s", cfg.KiroVersion, machine) + headers := http.Header{ + "Accept": []string{"*/*"}, + "Authorization": []string{"Bearer " + key}, + "Content-Type": []string{"application/x-amz-json-1.0"}, + "Tokentype": []string{"API_KEY"}, + "User-Agent": []string{userAgent}, + "X-Amz-Target": []string{kiroAmzTarget}, + "X-Amz-User-Agent": []string{amzUserAgent}, + "X-Amzn-Codewhisperer-Optout": []string{"false"}, + } + return upstreamRequest{URL: endpoint, Headers: headers, Body: body}, payload, nil +} + +func consumeUpstream(hostCallbackID string, request upstreamRequest, attempt int, onEvents func([]kiroEvent) error) (bool, int, error) { + headers := request.Headers.Clone() + headers.Set("Amz-Sdk-Request", fmt.Sprintf("attempt=%d; max=%d", attempt, maxAttempts)) + headers.Set("Amz-Sdk-Invocation-Id", randomUUID()) + rawResponse, errCall := invokeHost(pluginabi.MethodHostHTTPDoStream, rpcHostHTTPRequest{ + HostCallbackID: hostCallbackID, + Method: http.MethodPost, + URL: request.URL, + Headers: headers, + Body: request.Body, + }) + if errCall != nil { + return false, 0, fmt.Errorf("start Kiro upstream stream: %w", errCall) + } + var response rpcHostHTTPStreamResponse + if errUnmarshal := json.Unmarshal(rawResponse, &response); errUnmarshal != nil { + return false, 0, fmt.Errorf("decode Kiro upstream response: %w", errUnmarshal) + } + if strings.TrimSpace(response.StreamID) == "" { + return false, response.StatusCode, fmt.Errorf("Kiro upstream returned no stream") + } + defer closeHostHTTPStream(response.StreamID) + if response.StatusCode != http.StatusOK { + body, _ := readErrorBody(response.StreamID, 4096) + return false, response.StatusCode, fmt.Errorf("Kiro upstream HTTP %d: %s", response.StatusCode, strings.TrimSpace(string(body))) + } + + decoder := &eventStreamDecoder{} + emitted := false + for { + readRaw, errReadCall := invokeHost(pluginabi.MethodHostHTTPStreamRead, rpcHostHTTPStreamReadRequest{StreamID: response.StreamID}) + if errReadCall != nil { + return emitted, response.StatusCode, fmt.Errorf("read Kiro upstream stream: %w", errReadCall) + } + var readResponse rpcHostHTTPStreamReadResponse + if errUnmarshal := json.Unmarshal(readRaw, &readResponse); errUnmarshal != nil { + return emitted, response.StatusCode, fmt.Errorf("decode Kiro stream read: %w", errUnmarshal) + } + if readResponse.Error != "" { + return emitted, response.StatusCode, fmt.Errorf("Kiro upstream stream: %s", readResponse.Error) + } + if len(readResponse.Payload) > 0 { + events, errFeed := decoder.Feed(readResponse.Payload) + if errFeed != nil { + return emitted, response.StatusCode, errFeed + } + if len(events) > 0 { + emitted = true + if errEvents := onEvents(events); errEvents != nil { + return emitted, response.StatusCode, errEvents + } + } + } + if readResponse.Done { + break + } + } + if errFinish := decoder.Finish(); errFinish != nil { + return emitted, response.StatusCode, errFinish + } + return emitted, response.StatusCode, nil +} + +func readErrorBody(streamID string, maxBytes int) ([]byte, error) { + var body bytes.Buffer + for body.Len() < maxBytes { + raw, errCall := invokeHost(pluginabi.MethodHostHTTPStreamRead, rpcHostHTTPStreamReadRequest{StreamID: streamID}) + if errCall != nil { + return body.Bytes(), errCall + } + var response rpcHostHTTPStreamReadResponse + if errUnmarshal := json.Unmarshal(raw, &response); errUnmarshal != nil { + return body.Bytes(), errUnmarshal + } + remaining := maxBytes - body.Len() + if len(response.Payload) > remaining { + response.Payload = response.Payload[:remaining] + } + body.Write(response.Payload) + if response.Done || response.Error != "" { + break + } + } + return body.Bytes(), nil +} + +func closeHostHTTPStream(streamID string) { + if strings.TrimSpace(streamID) == "" { + return + } + _, _ = invokeHost(pluginabi.MethodHostHTTPStreamClose, rpcHostHTTPStreamCloseRequest{StreamID: streamID}) +} + +func emitFrames(streamID string, frames [][]byte) error { + for _, frame := range frames { + if len(frame) == 0 { + continue + } + if _, errCall := invokeHost(pluginabi.MethodHostStreamEmit, rpcStreamEmitRequest{StreamID: streamID, Payload: frame}); errCall != nil { + return errCall + } + } + return nil +} + +func closePluginStream(streamID string, errValue error) { + message := "" + if errValue != nil { + message = errValue.Error() + } + _, _ = invokeHost(pluginabi.MethodHostStreamClose, rpcStreamCloseRequest{StreamID: streamID, Error: message}) +} + +func retryableUpstream(status int, errValue error) bool { + if errValue == nil { + return false + } + if status == http.StatusUnauthorized || status == http.StatusForbidden || status == http.StatusPaymentRequired || status == http.StatusBadRequest { + return false + } + return status == 0 || status == http.StatusTooManyRequests || status >= 500 +} + +func upstreamErrorEnvelope(status int, errValue error) []byte { + if errValue == nil { + errValue = fmt.Errorf("Kiro upstream request failed") + } + if status == 0 { + status = http.StatusBadGateway + } + code := "upstream_error" + if status == http.StatusUnauthorized || status == http.StatusForbidden { + code = "invalid_auth" + } else if status == http.StatusTooManyRequests { + code = "quota_exhausted" + } else if status == http.StatusPaymentRequired { + code = "billing_required" + } + return errorEnvelope(code, errValue.Error(), status, retryableUpstream(status, errValue)) +} + +func countTokens(raw []byte) ([]byte, error) { + var request rpcExecutorRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, errUnmarshal + } + var payload any + if errUnmarshal := json.Unmarshal(request.Payload, &payload); errUnmarshal != nil { + return errorEnvelope("invalid_request", "invalid Claude token-count request", http.StatusBadRequest, false), nil + } + serialized, _ := json.Marshal(payload) + characters := utf8.RuneCount(serialized) + estimate := (characters + 3) / 4 + return okEnvelope(pluginapi.ExecutorResponse{ + Payload: []byte(fmt.Sprintf(`{"input_tokens":%d}`, estimate)), + Headers: http.Header{"Content-Type": []string{"application/json"}}, + Metadata: map[string]any{"estimated": true}, + }) +} + +func executorHTTPRequest(raw []byte) ([]byte, error) { + var request rpcExecutorHTTPRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, errUnmarshal + } + credential, handled, errCredential := decodeCredential(request.StorageJSON) + if errCredential != nil || !handled { + return errorEnvelope("invalid_auth", "selected auth is not a Kiro credential", http.StatusUnauthorized, false), nil + } + endpoint := "https://runtime." + credential.Region + ".kiro.dev/" + parsed, errParse := url.Parse(request.URL) + if errParse != nil || parsed.Scheme != "https" || parsed.Hostname() != "runtime."+credential.Region+".kiro.dev" { + return errorEnvelope("invalid_request", "Kiro HTTP bridge only permits the configured runtime host", http.StatusBadRequest, false), nil + } + key, errKey := resolveAPIKey(credential) + if errKey != nil { + return errorEnvelope("invalid_auth", errKey.Error(), http.StatusUnauthorized, false), nil + } + headers := request.Headers.Clone() + headers.Set("Authorization", "Bearer "+key) + headers.Set("tokentype", "API_KEY") + resultRaw, errCall := invokeHost(pluginabi.MethodHostHTTPDo, rpcHostHTTPRequest{ + HostCallbackID: request.HostCallbackID, + Method: request.Method, + URL: firstNonEmpty(request.URL, endpoint), + Headers: headers, + Body: request.Body, + }) + if errCall != nil { + return upstreamErrorEnvelope(0, errCall), nil + } + var response pluginapi.ExecutorHTTPResponse + if errUnmarshal := json.Unmarshal(resultRaw, &response); errUnmarshal != nil { + return nil, errUnmarshal + } + return okEnvelope(response) +} diff --git a/examples/plugin/kiro/go/go.mod b/examples/plugin/kiro/go/go.mod new file mode 100644 index 00000000000..85e7c9ea747 --- /dev/null +++ b/examples/plugin/kiro/go/go.mod @@ -0,0 +1,10 @@ +module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/kiro/go + +go 1.26.0 + +require ( + github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 + gopkg.in/yaml.v3 v3.0.1 +) + +replace github.com/router-for-me/CLIProxyAPI/v7 => ../../../.. diff --git a/examples/plugin/kiro/go/go.sum b/examples/plugin/kiro/go/go.sum new file mode 100644 index 00000000000..a62c313c5b0 --- /dev/null +++ b/examples/plugin/kiro/go/go.sum @@ -0,0 +1,4 @@ +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go new file mode 100644 index 00000000000..0bcced94a13 --- /dev/null +++ b/examples/plugin/kiro/go/kiro_test.go @@ -0,0 +1,575 @@ +package main + +import ( + "encoding/binary" + "encoding/json" + "hash/crc32" + "net/http" + "net/url" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func TestParseAuthEnvironmentReferenceDoesNotPersistSecret(t *testing.T) { + t.Setenv("KIRO_TEST_KEY", "secret-test-value") + request, _ := json.Marshal(pluginapi.AuthParseRequest{ + FileName: "kiro-pro.json", + RawJSON: []byte(`{"type":"kiro","api_key_env":"KIRO_TEST_KEY","region":"us-east-1","label":"pro"}`), + }) + raw, errParse := parseAuth(request) + if errParse != nil { + t.Fatalf("parseAuth() error = %v", errParse) + } + var env envelope + if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if !env.OK { + t.Fatalf("parseAuth() envelope = %#v", env) + } + var response pluginapi.AuthParseResponse + if errUnmarshal := json.Unmarshal(env.Result, &response); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if !response.Handled || response.Auth.Provider != pluginIdentifier { + t.Fatalf("ParseAuth response = %#v", response) + } + if strings.Contains(string(response.Auth.StorageJSON), "secret-test-value") { + t.Fatal("StorageJSON contains resolved API key") + } + if response.Auth.Attributes["region"] != "us-east-1" { + t.Fatalf("region = %q", response.Auth.Attributes["region"]) + } +} + +func TestDecodeCredentialRejectsHostInjectionRegion(t *testing.T) { + _, handled, errCredential := decodeCredential([]byte(`{"type":"kiro","api_key_env":"KEY","region":"us-east-1.kiro.dev.evil"}`)) + if !handled || errCredential == nil { + t.Fatalf("decodeCredential() handled=%v error=%v", handled, errCredential) + } +} + +func TestModelsForAuthDiscoversPaginatesAndCaches(t *testing.T) { + configureKiroTest(t, "models:\n - '*'\n") + t.Setenv("KIRO_TEST_KEY", "test-discovery-key") + + var requests []rpcHostHTTPRequest + previous := invokeHost + invokeHost = func(method string, payload any) (json.RawMessage, error) { + if method != pluginabi.MethodHostHTTPDo { + return nil, os.ErrInvalid + } + request := payload.(rpcHostHTTPRequest) + requests = append(requests, request) + if request.Headers.Get("Authorization") != "Bearer test-discovery-key" || request.Headers.Get("TokenType") != "API_KEY" { + t.Fatalf("discovery headers = %#v", request.Headers) + } + parsed, errParse := url.Parse(request.URL) + if errParse != nil { + t.Fatal(errParse) + } + if parsed.Host != "codewhisperer.us-east-1.amazonaws.com" || parsed.Path != "/ListAvailableModels" { + t.Fatalf("discovery URL = %q", request.URL) + } + var body []byte + switch len(requests) { + case 1: + if parsed.Query().Get("nextToken") != "" { + t.Fatalf("first nextToken = %q", parsed.Query().Get("nextToken")) + } + body = []byte(`{ + "models":[ + {"modelId":"claude-opus-5","modelName":"Claude Opus 5","description":"Large context","supportedInputTypes":["TEXT","IMAGE"],"tokenLimits":{"maxInputTokens":1000000,"maxOutputTokens":128000}}, + {"modelId":"","modelName":"invalid"} + ], + "nextToken":"page +/=" + }`) + case 2: + if parsed.Query().Get("nextToken") != "page +/=" { + t.Fatalf("second nextToken = %q", parsed.Query().Get("nextToken")) + } + body = []byte(`{ + "models":[ + {"modelId":"claude-opus-5","modelName":"duplicate"}, + {"modelId":"glm-5","modelName":"GLM 5","supportedInputTypes":["TEXT"],"tokenLimits":{"maxInputTokens":200000,"maxOutputTokens":64000}} + ], + "nextToken":null + }`) + default: + t.Fatalf("unexpected discovery request %d", len(requests)) + } + return mustJSONRaw(pluginapi.HTTPResponse{StatusCode: http.StatusOK, Body: body}), nil + } + t.Cleanup(func() { invokeHost = previous }) + + response := callModelsForAuthTest(t, "kiro-pro", "callback-1") + if len(response.Models) != 2 { + t.Fatalf("discovered models = %#v", response.Models) + } + if response.Models[0].ID != "claude-opus-5" || response.Models[0].InputTokenLimit != 1_000_000 || response.Models[0].OutputTokenLimit != 128_000 { + t.Fatalf("first model = %#v", response.Models[0]) + } + if strings.Join(response.Models[0].SupportedInputModalities, ",") != "text,image" { + t.Fatalf("first model modalities = %#v", response.Models[0].SupportedInputModalities) + } + if response.Models[1].ID != "glm-5" || strings.Join(response.Models[1].SupportedInputModalities, ",") != "text" { + t.Fatalf("second model = %#v", response.Models[1]) + } + + cached := callModelsForAuthTest(t, "kiro-pro", "callback-2") + if len(cached.Models) != 2 || len(requests) != 2 { + t.Fatalf("cached models=%d HTTP requests=%d", len(cached.Models), len(requests)) + } +} + +func TestModelsForAuthAppliesConfiguredAllowList(t *testing.T) { + configureKiroTest(t, "models:\n - claude-sonnet-4.5\n") + t.Setenv("KIRO_TEST_KEY", "test-allow-list-key") + + previous := invokeHost + invokeHost = func(method string, payload any) (json.RawMessage, error) { + if method != pluginabi.MethodHostHTTPDo { + return nil, os.ErrInvalid + } + body := []byte(`{"models":[{"modelId":"claude-sonnet-4.5","modelName":"Sonnet"},{"modelId":"claude-opus-5","modelName":"Opus"}]}`) + return mustJSONRaw(pluginapi.HTTPResponse{StatusCode: http.StatusOK, Body: body}), nil + } + t.Cleanup(func() { invokeHost = previous }) + + response := callModelsForAuthTest(t, "kiro-allow-list", "callback-allow-list") + if len(response.Models) != 1 || response.Models[0].ID != "claude-sonnet-4.5" { + t.Fatalf("allow-listed models = %#v", response.Models) + } +} + +func TestModelsForAuthUsesConfiguredFallbackOnColdStartFailure(t *testing.T) { + configureKiroTest(t, "models:\n - '*'\n") + t.Setenv("KIRO_TEST_KEY", "test-fallback-key") + + previous := invokeHost + invokeHost = func(method string, payload any) (json.RawMessage, error) { + switch method { + case pluginabi.MethodHostHTTPDo: + return mustJSONRaw(pluginapi.HTTPResponse{StatusCode: http.StatusServiceUnavailable, Body: []byte(`{"message":"temporary"}`)}), nil + case pluginabi.MethodHostLog: + return json.RawMessage(`{}`), nil + default: + return nil, os.ErrInvalid + } + } + t.Cleanup(func() { invokeHost = previous }) + + response := callModelsForAuthTest(t, "kiro-fallback", "callback-fallback") + if len(response.Models) != 2 || response.Models[0].ID != "claude-sonnet-4.5" || response.Models[1].ID != "claude-haiku-4.5" { + t.Fatalf("fallback models = %#v", response.Models) + } + for _, model := range response.Models { + if model.ID == "*" { + t.Fatal("wildcard was registered as a model") + } + } +} + +func TestModelsForAuthUsesStaleCacheOnRefreshFailure(t *testing.T) { + configureKiroTest(t, "models:\n - '*'\n") + t.Setenv("KIRO_TEST_KEY", "test-stale-key") + cacheKey := modelCacheKey("kiro-stale", "us-east-1", "test-stale-key") + modelCacheMu.Lock() + modelCache[cacheKey] = modelCacheEntry{ + Models: []pluginapi.ModelInfo{{ID: "claude-opus-5", Name: "Claude Opus 5"}}, + FetchedAt: time.Now().Add(-modelCacheTTL - time.Second), + } + modelCacheMu.Unlock() + + previous := invokeHost + invokeHost = func(method string, payload any) (json.RawMessage, error) { + switch method { + case pluginabi.MethodHostHTTPDo: + return nil, os.ErrDeadlineExceeded + case pluginabi.MethodHostLog: + return json.RawMessage(`{}`), nil + default: + return nil, os.ErrInvalid + } + } + t.Cleanup(func() { invokeHost = previous }) + + response := callModelsForAuthTest(t, "kiro-stale", "callback-stale") + if len(response.Models) != 1 || response.Models[0].ID != "claude-opus-5" { + t.Fatalf("stale models = %#v", response.Models) + } +} + +func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { + raw := []byte(`{ + "model":"claude-sonnet-4-5", + "max_tokens":128, + "system":"Be concise.", + "tools":[{"name":"math.add/unsafe","description":"Add values","input_schema":{"type":"object","properties":{"a":{"type":"number"}}}}], + "messages":[ + {"role":"user","content":"Add 2 and 3"}, + {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"math.add/unsafe","input":{"a":2}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"5"}]} + ] + }`) + payload, request, errTranslate := claudeToKiro(raw, "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + current := payload.ConversationState.CurrentMessage.UserInputMessage + if current.ModelID != "claude-sonnet-4.5" { + t.Fatalf("model ID = %q", current.ModelID) + } + if current.Origin != "KIRO_CLI" { + t.Fatalf("origin = %q", current.Origin) + } + if request.MaxTokens != 128 || payload.InferenceConfig.MaxTokens != 128 { + t.Fatalf("max tokens were not preserved") + } + if len(payload.ConversationState.History) < 4 { + t.Fatalf("history = %#v", payload.ConversationState.History) + } + context := current.UserInputMessageContext + if context == nil || len(context.Tools) != 1 || len(context.ToolResults) != 1 { + t.Fatalf("current context = %#v", context) + } + toolName := context.Tools[0].ToolSpecification.Name + if toolName != "math_add_unsafe" || payload.ToolNameMap[toolName] != "math.add/unsafe" { + t.Fatalf("tool name=%q map=%#v", toolName, payload.ToolNameMap) + } + if !strings.Contains(current.Content, "Tool results:") { + t.Fatalf("current content = %q", current.Content) + } +} + +func TestNormalizeKiroModelDoesNotRewriteDatedSnapshotAsDecimal(t *testing.T) { + if got := normalizeKiroModel("claude-sonnet-4-20250514"); got != "claude-sonnet-4" { + t.Fatalf("normalizeKiroModel() = %q", got) + } + if got := normalizeKiroModel("claude-haiku-4-5"); got != "claude-haiku-4.5" { + t.Fatalf("normalizeKiroModel() = %q", got) + } +} + +func TestEventStreamDecoderHandlesSplitFramesAndCRC(t *testing.T) { + frame := testEventFrame(t, "assistantResponseEvent", map[string]any{"content": "hello"}) + decoder := &eventStreamDecoder{} + first, errFirst := decoder.Feed(frame[:7]) + if errFirst != nil || len(first) != 0 { + t.Fatalf("first Feed() events=%d error=%v", len(first), errFirst) + } + events, errSecond := decoder.Feed(frame[7:]) + if errSecond != nil || len(events) != 1 || events[0].Type != "assistantResponseEvent" { + t.Fatalf("second Feed() events=%#v error=%v", events, errSecond) + } + corrupt := append([]byte(nil), frame...) + corrupt[len(corrupt)-1] ^= 0xff + if _, errCorrupt := (&eventStreamDecoder{}).Feed(corrupt); errCorrupt == nil { + t.Fatal("corrupt frame was accepted") + } +} + +func TestPendingToolUseAdoptsLateUpstreamID(t *testing.T) { + pending := &pendingToolUses{} + tools, errFirst := pending.accept(map[string]any{"name": "lookup", "input": `{"q":`}) + if errFirst != nil || len(tools) != 0 { + t.Fatalf("first tool fragment = %#v, %v", tools, errFirst) + } + tools, errSecond := pending.accept(map[string]any{"toolUseId": "toolu_real", "name": "lookup", "input": `"x"}`, "stop": true}) + if errSecond != nil || len(tools) != 1 { + t.Fatalf("second tool fragment = %#v, %v", tools, errSecond) + } + if tools[0].ToolUseID != "toolu_real" || tools[0].Input["q"] != "x" { + t.Fatalf("tool = %#v", tools[0]) + } +} + +func TestExecuteUsesHostTransportAndReturnsClaudeResponse(t *testing.T) { + t.Setenv("KIRO_TEST_KEY", "test-key") + frames := append(testEventFrame(t, "assistantResponseEvent", map[string]any{"content": "391", "usage": map[string]any{"inputTokens": 12, "outputTokens": 1}}), testEventFrame(t, "metadataEvent", map[string]any{"stopReason": "end_turn"})...) + mock := newHostMock(frames) + previous := invokeHost + invokeHost = mock.call + t.Cleanup(func() { invokeHost = previous }) + + storage := []byte(`{"type":"kiro","api_key_env":"KIRO_TEST_KEY","region":"us-east-1"}`) + requestRaw, _ := json.Marshal(rpcExecutorRequest{ExecutorRequest: pluginapi.ExecutorRequest{ + Model: "claude-sonnet-4.5", + Payload: []byte(`{"model":"claude-sonnet-4.5","max_tokens":32,"messages":[{"role":"user","content":"Calculate 17 * 23. Return only the number."}]}`), + StorageJSON: storage, + }}) + raw, errExecute := execute(requestRaw) + if errExecute != nil { + t.Fatalf("execute() error = %v", errExecute) + } + var env envelope + _ = json.Unmarshal(raw, &env) + if !env.OK { + t.Fatalf("execute() envelope = %#v", env) + } + var executorResponse pluginapi.ExecutorResponse + _ = json.Unmarshal(env.Result, &executorResponse) + var response claudeResponse + if errUnmarshal := json.Unmarshal(executorResponse.Payload, &response); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if len(response.Content) != 1 || response.Content[0].Text != "391" || response.Usage.InputTokens != 12 || response.Usage.OutputTokens != 1 { + t.Fatalf("response = %#v", response) + } + if got := mock.request.Headers.Get("Authorization"); got != "Bearer test-key" { + t.Fatalf("Authorization = %q", got) + } + if strings.Contains(string(mock.request.Body), "test-key") { + t.Fatal("upstream body contains API key") + } +} + +func TestExecuteStreamEmitsAnthropicSSE(t *testing.T) { + t.Setenv("KIRO_TEST_KEY", "test-key") + frames := append(testEventFrame(t, "assistantResponseEvent", map[string]any{"content": "hello"}), testEventFrame(t, "assistantResponseEvent", map[string]any{"content": " world"})...) + frames = append(frames, testEventFrame(t, "metadataEvent", map[string]any{"stopReason": "end_turn"})...) + mock := newHostMock(frames) + previous := invokeHost + invokeHost = mock.call + t.Cleanup(func() { invokeHost = previous }) + + requestRaw, _ := json.Marshal(rpcExecutorRequest{ + ExecutorRequest: pluginapi.ExecutorRequest{ + Model: "claude-sonnet-4.5", + Payload: []byte(`{"model":"claude-sonnet-4.5","max_tokens":32,"messages":[{"role":"user","content":"hello"}],"stream":true}`), + StorageJSON: []byte(`{"type":"kiro","api_key_env":"KIRO_TEST_KEY","region":"us-east-1"}`), + }, + StreamID: "plugin-stream-1", + }) + if _, errStream := executeStream(requestRaw); errStream != nil { + t.Fatal(errStream) + } + select { + case <-mock.closed: + case <-time.After(3 * time.Second): + t.Fatal("stream did not close") + } + joined := string(mock.emittedBytes()) + for _, expected := range []string{"event: message_start", "event: content_block_delta", `"text":"hello"`, `"text":" world"`, `"input_tokens":`, `"output_tokens":3`, "event: message_stop"} { + if !strings.Contains(joined, expected) { + t.Fatalf("stream missing %q: %s", expected, joined) + } + } + if got := strings.Count(joined, `"type":"content_block_start"`); got != 1 { + t.Fatalf("content block starts = %d, want 1: %s", got, joined) + } + if got := strings.Count(joined, `"type":"content_block_stop"`); got != 1 { + t.Fatalf("content block stops = %d, want 1: %s", got, joined) + } +} + +func TestAccumulatorMergesAdjacentFragments(t *testing.T) { + payload := &kiroPayload{} + payload.ConversationState.CurrentMessage.UserInputMessage.ModelID = "claude-sonnet-4.5" + accumulator := newAccumulator(payload) + + events := []kiroEvent{ + {Type: "reasoningContentEvent", Payload: map[string]any{"text": "Think"}}, + {Type: "reasoningContentEvent", Payload: map[string]any{"text": " first."}}, + {Type: "assistantResponseEvent", Payload: map[string]any{"content": "Hello"}}, + {Type: "assistantResponseEvent", Payload: map[string]any{"content": " world."}}, + } + for _, event := range events { + if _, errAccept := accumulator.accept(event); errAccept != nil { + t.Fatal(errAccept) + } + } + if len(accumulator.Blocks) != 2 { + t.Fatalf("content blocks = %#v, want one thinking and one text block", accumulator.Blocks) + } + if accumulator.Blocks[0].Thinking != "Think first." || accumulator.Blocks[1].Text != "Hello world." { + t.Fatalf("merged content blocks = %#v", accumulator.Blocks) + } +} + +func TestAccumulatorEstimatesUsageWhenUpstreamOmitsTokens(t *testing.T) { + payload, _, errTranslate := claudeToKiro([]byte(`{ + "model":"claude-sonnet-4.5", + "max_tokens":32, + "messages":[{"role":"user","content":"Calculate 17 * 23. Return only the number."}] + }`), "") + if errTranslate != nil { + t.Fatal(errTranslate) + } + accumulator := newAccumulator(payload) + if _, errAccept := accumulator.accept(kiroEvent{Type: "assistantResponseEvent", Payload: map[string]any{"content": "391"}}); errAccept != nil { + t.Fatal(errAccept) + } + if errFinish := accumulator.finish(); errFinish != nil { + t.Fatal(errFinish) + } + if accumulator.InputTokens <= 0 || accumulator.OutputTokens <= 0 { + t.Fatalf("estimated usage = input:%d output:%d, want both positive", accumulator.InputTokens, accumulator.OutputTokens) + } +} + +func TestAccumulatorPrefersContextUsageOverInputEstimate(t *testing.T) { + payload := &kiroPayload{EstimatedInputTokens: 7} + payload.ConversationState.CurrentMessage.UserInputMessage.ModelID = "claude-sonnet-4.5" + accumulator := newAccumulator(payload) + _, _ = accumulator.accept(kiroEvent{Type: "contextUsageEvent", Payload: map[string]any{"contextUsagePercentage": 10.5}}) + _, _ = accumulator.accept(kiroEvent{Type: "assistantResponseEvent", Payload: map[string]any{"content": "ok"}}) + if errFinish := accumulator.finish(); errFinish != nil { + t.Fatal(errFinish) + } + if accumulator.InputTokens != 21_000 { + t.Fatalf("input tokens = %d, want 21000 from context usage", accumulator.InputTokens) + } +} + +func TestUpdateUsageReadsNestedCacheBuckets(t *testing.T) { + event := map[string]any{"metrics": map[string]any{"usage": map[string]any{ + "uncachedInputTokens": "5", + "cacheReadInputTokens": 4.0, + "cacheWriteInputTokens": 3.0, + "totalOutputTokens": 2.0, + }}} + inputTokens, outputTokens := updateUsage(event, 0, 0) + if inputTokens != 12 || outputTokens != 2 { + t.Fatalf("usage = input:%d output:%d, want input:12 output:2", inputTokens, outputTokens) + } +} + +func configureKiroTest(t *testing.T, configYAML string) { + t.Helper() + request, errMarshal := json.Marshal(lifecycleRequest{ConfigYAML: []byte(configYAML)}) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if errConfigure := configure(request); errConfigure != nil { + t.Fatal(errConfigure) + } + t.Cleanup(func() { + if errConfigure := configure(nil); errConfigure != nil { + t.Errorf("restore default configuration: %v", errConfigure) + } + }) +} + +func callModelsForAuthTest(t *testing.T, authID, hostCallbackID string) pluginapi.ModelResponse { + t.Helper() + rawRequest, errMarshal := json.Marshal(rpcAuthModelRequest{ + AuthModelRequest: pluginapi.AuthModelRequest{ + AuthID: authID, + StorageJSON: []byte(`{"type":"kiro","api_key_env":"KIRO_TEST_KEY","region":"us-east-1"}`), + }, + HostCallbackID: hostCallbackID, + }) + if errMarshal != nil { + t.Fatal(errMarshal) + } + rawResponse, errModels := modelsForAuth(rawRequest) + if errModels != nil { + t.Fatal(errModels) + } + var env envelope + if errUnmarshal := json.Unmarshal(rawResponse, &env); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if !env.OK { + t.Fatalf("modelsForAuth() envelope = %#v", env) + } + var response pluginapi.ModelResponse + if errUnmarshal := json.Unmarshal(env.Result, &response); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + return response +} + +type hostMock struct { + mu sync.Mutex + chunks [][]byte + request rpcHostHTTPRequest + emitted [][]byte + closed chan struct{} + once sync.Once +} + +func newHostMock(stream []byte) *hostMock { + middle := len(stream) / 2 + return &hostMock{chunks: [][]byte{stream[:middle], stream[middle:]}, closed: make(chan struct{})} +} + +func (m *hostMock) call(method string, payload any) (json.RawMessage, error) { + m.mu.Lock() + defer m.mu.Unlock() + switch method { + case pluginabi.MethodHostHTTPDoStream: + m.request = payload.(rpcHostHTTPRequest) + return mustJSONRaw(rpcHostHTTPStreamResponse{StatusCode: http.StatusOK, StreamID: "upstream-1"}), nil + case pluginabi.MethodHostHTTPStreamRead: + if len(m.chunks) == 0 { + return mustJSONRaw(rpcHostHTTPStreamReadResponse{Done: true}), nil + } + chunk := m.chunks[0] + m.chunks = m.chunks[1:] + return mustJSONRaw(rpcHostHTTPStreamReadResponse{Payload: chunk, Done: len(m.chunks) == 0}), nil + case pluginabi.MethodHostHTTPStreamClose: + return json.RawMessage(`{}`), nil + case pluginabi.MethodHostStreamEmit: + request := payload.(rpcStreamEmitRequest) + m.emitted = append(m.emitted, append([]byte(nil), request.Payload...)) + return json.RawMessage(`{}`), nil + case pluginabi.MethodHostStreamClose: + m.once.Do(func() { close(m.closed) }) + return json.RawMessage(`{}`), nil + default: + return nil, os.ErrInvalid + } +} + +func (m *hostMock) emittedBytes() []byte { + m.mu.Lock() + defer m.mu.Unlock() + return []byte(strings.Join(byteStrings(m.emitted), "")) +} + +func byteStrings(values [][]byte) []string { + out := make([]string, 0, len(values)) + for _, value := range values { + out = append(out, string(value)) + } + return out +} + +func mustJSONRaw(value any) json.RawMessage { + raw, _ := json.Marshal(value) + return raw +} + +func testEventFrame(t *testing.T, eventType string, payload any) []byte { + t.Helper() + headers := append(testEventHeader(":message-type", "event"), testEventHeader(":event-type", eventType)...) + body, errMarshal := json.Marshal(payload) + if errMarshal != nil { + t.Fatal(errMarshal) + } + totalLength := 16 + len(headers) + len(body) + frame := make([]byte, totalLength) + binary.BigEndian.PutUint32(frame[0:4], uint32(totalLength)) + binary.BigEndian.PutUint32(frame[4:8], uint32(len(headers))) + binary.BigEndian.PutUint32(frame[8:12], crc32.ChecksumIEEE(frame[:8])) + copy(frame[12:], headers) + copy(frame[12+len(headers):], body) + binary.BigEndian.PutUint32(frame[totalLength-4:], crc32.ChecksumIEEE(frame[:totalLength-4])) + return frame +} + +func testEventHeader(name, value string) []byte { + header := []byte{byte(len(name))} + header = append(header, name...) + header = append(header, 7, byte(len(value)>>8), byte(len(value))) + header = append(header, value...) + return header +} diff --git a/examples/plugin/kiro/go/main.go b/examples/plugin/kiro/go/main.go new file mode 100644 index 00000000000..6b9be637266 --- /dev/null +++ b/examples/plugin/kiro/go/main.go @@ -0,0 +1,422 @@ +package main + +/* +#include +#include + +typedef struct { + void* ptr; + size_t len; +} cliproxy_buffer; + +typedef int (*cliproxy_host_call_fn)(void*, const char*, const uint8_t*, size_t, cliproxy_buffer*); +typedef void (*cliproxy_host_free_fn)(void*, size_t); + +typedef struct { + uint32_t abi_version; + void* host_ctx; + cliproxy_host_call_fn call; + cliproxy_host_free_fn free_buffer; +} cliproxy_host_api; + +typedef int (*cliproxy_plugin_call_fn)(char*, uint8_t*, size_t, cliproxy_buffer*); +typedef void (*cliproxy_plugin_free_fn)(void*, size_t); +typedef void (*cliproxy_plugin_shutdown_fn)(void); + +typedef struct { + uint32_t abi_version; + cliproxy_plugin_call_fn call; + cliproxy_plugin_free_fn free_buffer; + cliproxy_plugin_shutdown_fn shutdown; +} cliproxy_plugin_api; + +extern int cliproxyPluginCall(char*, uint8_t*, size_t, cliproxy_buffer*); +extern void cliproxyPluginFree(void*, size_t); +extern void cliproxyPluginShutdown(void); + +static const cliproxy_host_api* stored_host; + +static void store_host_api(const cliproxy_host_api* host) { + stored_host = host; +} + +static int call_host_api(const char* method, const uint8_t* request, size_t request_len, cliproxy_buffer* response) { + if (stored_host == NULL || stored_host->call == NULL) { + return 1; + } + return stored_host->call(stored_host->host_ctx, method, request, request_len, response); +} + +static void free_host_buffer(void* ptr, size_t len) { + if (stored_host != NULL && stored_host->free_buffer != NULL && ptr != NULL) { + stored_host->free_buffer(ptr, len); + } +} +*/ +import "C" + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "sync/atomic" + "unsafe" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + "gopkg.in/yaml.v3" +) + +const pluginIdentifier = "kiro" + +type envelope struct { + OK bool `json:"ok"` + Result json.RawMessage `json:"result,omitempty"` + Error *pluginabi.Error `json:"error,omitempty"` +} + +type lifecycleRequest struct { + ConfigYAML []byte `json:"config_yaml"` +} + +type registration struct { + SchemaVersion uint32 `json:"schema_version"` + Metadata pluginapi.Metadata `json:"metadata"` + Capabilities registrationCapability `json:"capabilities"` +} + +type registrationCapability struct { + ModelRegistrar bool `json:"model_registrar"` + ModelProvider bool `json:"model_provider"` + AuthProvider bool `json:"auth_provider"` + Executor bool `json:"executor"` + ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"` + ExecutorInputFormats []string `json:"executor_input_formats"` + ExecutorOutputFormats []string `json:"executor_output_formats"` +} + +type pluginConfig struct { + Enabled bool `yaml:"enabled"` + Models []string `yaml:"models"` + KiroVersion string `yaml:"kiro_version"` + NodeVersion string `yaml:"node_version"` + SystemVersion string `yaml:"system_version"` +} + +type rpcExecutorRequest struct { + pluginapi.ExecutorRequest + StreamID string `json:"stream_id,omitempty"` + HostCallbackID string `json:"host_callback_id,omitempty"` +} + +type rpcExecutorHTTPRequest struct { + pluginapi.ExecutorHTTPRequest + HostCallbackID string `json:"host_callback_id,omitempty"` +} + +type rpcAuthModelRequest struct { + pluginapi.AuthModelRequest + HostCallbackID string `json:"host_callback_id,omitempty"` +} + +var currentConfig atomic.Value + +func main() {} + +//export cliproxy_plugin_init +func cliproxy_plugin_init(host *C.cliproxy_host_api, plugin *C.cliproxy_plugin_api) C.int { + if plugin == nil { + return 1 + } + C.store_host_api(host) + plugin.abi_version = C.uint32_t(pluginabi.ABIVersion) + plugin.call = C.cliproxy_plugin_call_fn(C.cliproxyPluginCall) + plugin.free_buffer = C.cliproxy_plugin_free_fn(C.cliproxyPluginFree) + plugin.shutdown = C.cliproxy_plugin_shutdown_fn(C.cliproxyPluginShutdown) + return 0 +} + +//export cliproxyPluginCall +func cliproxyPluginCall(method *C.char, request *C.uint8_t, requestLen C.size_t, response *C.cliproxy_buffer) C.int { + if response != nil { + response.ptr = nil + response.len = 0 + } + if method == nil { + writeResponse(response, errorEnvelope("invalid_method", "method is required", 0, false)) + return 1 + } + var requestBytes []byte + if request != nil && requestLen > 0 { + requestBytes = C.GoBytes(unsafe.Pointer(request), C.int(requestLen)) + } + raw, errHandle := handleMethod(C.GoString(method), requestBytes) + if errHandle != nil { + writeResponse(response, errorEnvelope("plugin_error", errHandle.Error(), 0, false)) + return 1 + } + writeResponse(response, raw) + return 0 +} + +//export cliproxyPluginFree +func cliproxyPluginFree(ptr unsafe.Pointer, _ C.size_t) { + if ptr != nil { + C.free(ptr) + } +} + +//export cliproxyPluginShutdown +func cliproxyPluginShutdown() {} + +func handleMethod(method string, request []byte) ([]byte, error) { + switch method { + case pluginabi.MethodPluginRegister, pluginabi.MethodPluginReconfigure: + if errConfigure := configure(request); errConfigure != nil { + return nil, errConfigure + } + return okEnvelope(pluginRegistration()) + case pluginabi.MethodModelRegister: + return okEnvelope(pluginapi.ModelRegistrationResponse{Provider: pluginIdentifier, Models: configuredModels()}) + case pluginabi.MethodModelStatic: + return okEnvelope(pluginapi.ModelResponse{Provider: pluginIdentifier, Models: configuredModels()}) + case pluginabi.MethodModelForAuth: + return modelsForAuth(request) + case pluginabi.MethodAuthIdentifier, pluginabi.MethodExecutorIdentifier: + return okEnvelope(map[string]string{"identifier": pluginIdentifier}) + case pluginabi.MethodAuthParse: + return parseAuth(request) + case pluginabi.MethodAuthLoginStart: + return errorEnvelope("unsupported_auth", "Kiro provider supports API-key authentication only", http.StatusBadRequest, false), nil + case pluginabi.MethodAuthLoginPoll: + return okEnvelope(pluginapi.AuthLoginPollResponse{Status: pluginapi.AuthLoginStatusError, Message: "Kiro provider supports API-key authentication only"}) + case pluginabi.MethodAuthRefresh: + return refreshAuth(request) + case pluginabi.MethodExecutorExecute: + return execute(request) + case pluginabi.MethodExecutorExecuteStream: + return executeStream(request) + case pluginabi.MethodExecutorCountTokens: + return countTokens(request) + case pluginabi.MethodExecutorHTTPRequest: + return executorHTTPRequest(request) + default: + return errorEnvelope("unknown_method", "unknown method: "+method, 0, false), nil + } +} + +func defaultPluginConfig() pluginConfig { + return pluginConfig{ + Enabled: true, + Models: []string{"claude-sonnet-4.5", "claude-haiku-4.5"}, + KiroVersion: "0.11.107", + NodeVersion: "22.22.0", + SystemVersion: "linux#6.6.87", + } +} + +func configure(raw []byte) error { + var request lifecycleRequest + if len(raw) > 0 { + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return fmt.Errorf("decode plugin configuration: %w", errUnmarshal) + } + } + cfg := defaultPluginConfig() + if len(request.ConfigYAML) > 0 { + if errUnmarshal := yaml.Unmarshal(request.ConfigYAML, &cfg); errUnmarshal != nil { + return fmt.Errorf("decode Kiro plugin YAML: %w", errUnmarshal) + } + } + cfg.KiroVersion = firstNonEmpty(strings.TrimSpace(cfg.KiroVersion), "0.11.107") + cfg.NodeVersion = firstNonEmpty(strings.TrimSpace(cfg.NodeVersion), "22.22.0") + cfg.SystemVersion = firstNonEmpty(strings.TrimSpace(cfg.SystemVersion), "linux#6.6.87") + cfg.Models = normalizeModels(cfg.Models) + if len(cfg.Models) == 0 { + cfg.Models = defaultPluginConfig().Models + } + currentConfig.Store(cfg) + resetModelCache() + return nil +} + +func loadedConfig() pluginConfig { + if raw := currentConfig.Load(); raw != nil { + if cfg, ok := raw.(pluginConfig); ok { + return cfg + } + } + return defaultPluginConfig() +} + +func pluginRegistration() registration { + return registration{ + SchemaVersion: pluginabi.SchemaVersion, + Metadata: pluginapi.Metadata{ + Name: "kiro", + Version: "0.1.0", + Author: "router-for-me", + GitHubRepository: "https://github.com/router-for-me/CLIProxyAPI", + ConfigFields: []pluginapi.ConfigField{ + {Name: "models", Type: pluginapi.ConfigFieldTypeArray, Description: "Allowed Kiro model IDs. Use * to expose every model discovered for an account."}, + {Name: "kiro_version", Type: pluginapi.ConfigFieldTypeString, Description: "Kiro client version sent in upstream compatibility headers."}, + {Name: "node_version", Type: pluginapi.ConfigFieldTypeString, Description: "Node.js version sent in upstream compatibility headers."}, + {Name: "system_version", Type: pluginapi.ConfigFieldTypeString, Description: "Operating-system version sent in upstream compatibility headers."}, + }, + }, + Capabilities: registrationCapability{ + ModelRegistrar: true, + ModelProvider: true, + AuthProvider: true, + Executor: true, + ExecutorModelScope: pluginapi.ExecutorModelScopeBoth, + ExecutorInputFormats: []string{"claude"}, + ExecutorOutputFormats: []string{"claude"}, + }, + } +} + +func configuredModels() []pluginapi.ModelInfo { + models := loadedConfig().Models + for _, model := range models { + if model == "*" { + models = defaultPluginConfig().Models + break + } + } + out := make([]pluginapi.ModelInfo, 0, len(models)) + for _, model := range models { + out = append(out, pluginapi.ModelInfo{ + ID: model, + Name: model, + DisplayName: model, + Object: "model", + OwnedBy: "kiro", + Type: "claude", + ContextLength: 200000, + InputTokenLimit: 200000, + OutputTokenLimit: 64000, + MaxCompletionTokens: 64000, + SupportedGenerationMethods: []string{"generateContent", "streamGenerateContent"}, + SupportedParameters: []string{"max_tokens", "temperature", "top_p", "tools"}, + SupportedInputModalities: []string{"text", "image"}, + SupportedOutputModalities: []string{"text"}, + }) + } + return out +} + +func normalizeModels(models []string) []string { + seen := make(map[string]struct{}, len(models)) + out := make([]string, 0, len(models)) + for _, model := range models { + model = strings.TrimSpace(model) + if model == "" { + continue + } + if _, ok := seen[model]; ok { + continue + } + seen[model] = struct{}{} + out = append(out, model) + } + return out +} + +func okEnvelope(value any) ([]byte, error) { + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + return nil, errMarshal + } + return json.Marshal(envelope{OK: true, Result: raw}) +} + +func errorEnvelope(code, message string, status int, retryable bool) []byte { + raw, _ := json.Marshal(envelope{OK: false, Error: &pluginabi.Error{ + Code: code, + Message: message, + HTTPStatus: status, + Retryable: retryable, + }}) + return raw +} + +func writeResponse(response *C.cliproxy_buffer, raw []byte) { + if response == nil || len(raw) == 0 { + return + } + ptr := C.CBytes(raw) + if ptr == nil { + return + } + response.ptr = ptr + response.len = C.size_t(len(raw)) +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + return "" +} + +type hostEnvelope struct { + OK bool `json:"ok"` + Result json.RawMessage `json:"result,omitempty"` + Error *hostEnvelopeErr `json:"error,omitempty"` +} + +type hostEnvelopeErr struct { + Code string `json:"code"` + Message string `json:"message"` +} + +var invokeHost = callHost + +func callHost(method string, payload any) (json.RawMessage, error) { + rawPayload, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return nil, fmt.Errorf("marshal host callback %s: %w", method, errMarshal) + } + cMethod := C.CString(method) + defer C.free(unsafe.Pointer(cMethod)) + + var response C.cliproxy_buffer + var requestPtr *C.uint8_t + if len(rawPayload) > 0 { + cPayload := C.CBytes(rawPayload) + if cPayload == nil { + return nil, fmt.Errorf("allocate host callback %s", method) + } + defer C.free(cPayload) + requestPtr = (*C.uint8_t)(cPayload) + } + code := C.call_host_api(cMethod, requestPtr, C.size_t(len(rawPayload)), &response) + var rawResponse []byte + if response.ptr != nil && response.len > 0 { + rawResponse = C.GoBytes(response.ptr, C.int(response.len)) + } + if response.ptr != nil { + C.free_host_buffer(response.ptr, response.len) + } + if len(rawResponse) == 0 { + return nil, fmt.Errorf("host callback %s returned no response, code=%d", method, int(code)) + } + var env hostEnvelope + if errUnmarshal := json.Unmarshal(rawResponse, &env); errUnmarshal != nil { + return nil, fmt.Errorf("decode host callback %s: %w", method, errUnmarshal) + } + if !env.OK { + if env.Error != nil { + return nil, fmt.Errorf("%s: %s", env.Error.Code, env.Error.Message) + } + return nil, fmt.Errorf("host callback %s failed", method) + } + if code != 0 { + return nil, fmt.Errorf("host callback %s returned code=%d", method, int(code)) + } + return append(json.RawMessage(nil), env.Result...), nil +} diff --git a/examples/plugin/kiro/go/model_discovery.go b/examples/plugin/kiro/go/model_discovery.go new file mode 100644 index 00000000000..8ef0e927d3b --- /dev/null +++ b/examples/plugin/kiro/go/model_discovery.go @@ -0,0 +1,303 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginabi" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +const ( + modelCacheTTL = 5 * time.Minute + modelDiscoveryPageSize = 50 + modelDiscoveryMaxPages = 10 +) + +type kiroModelListResponse struct { + Models []kiroAvailableModel `json:"models"` + NextToken string `json:"nextToken"` +} + +type kiroAvailableModel struct { + ModelID string `json:"modelId"` + ModelName string `json:"modelName"` + Description string `json:"description"` + SupportedInputTypes []string `json:"supportedInputTypes"` + TokenLimits struct { + MaxInputTokens int `json:"maxInputTokens"` + MaxOutputTokens int `json:"maxOutputTokens"` + } `json:"tokenLimits"` +} + +type modelCacheEntry struct { + Models []pluginapi.ModelInfo + FetchedAt time.Time +} + +type rpcHostLogRequest struct { + HostCallbackID string `json:"host_callback_id,omitempty"` + Level string `json:"level,omitempty"` + Message string `json:"message,omitempty"` + Fields map[string]any `json:"fields,omitempty"` +} + +var ( + modelCacheMu sync.Mutex + modelCache = make(map[string]modelCacheEntry) +) + +func modelsForAuth(raw []byte) ([]byte, error) { + var request rpcAuthModelRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, fmt.Errorf("decode Kiro model discovery request: %w", errUnmarshal) + } + + models, errDiscovery := accountModels(request) + if errDiscovery != nil { + logModelDiscoveryFallback(request, errDiscovery) + } + return okEnvelope(pluginapi.ModelResponse{Provider: pluginIdentifier, Models: models}) +} + +func accountModels(request rpcAuthModelRequest) ([]pluginapi.ModelInfo, error) { + credential, handled, errCredential := decodeCredential(request.StorageJSON) + if errCredential != nil { + return configuredModels(), fmt.Errorf("decode credential: %w", errCredential) + } + if !handled { + return configuredModels(), fmt.Errorf("selected auth is not a Kiro credential") + } + key, errKey := resolveAPIKey(credential) + if errKey != nil { + return configuredModels(), fmt.Errorf("resolve API key: %w", errKey) + } + + cacheKey := modelCacheKey(request.AuthID, credential.Region, key) + if models, okCache := cachedAccountModels(cacheKey, true); okCache { + return filterDiscoveredModels(models), nil + } + + models, errDiscovery := fetchAvailableModels(request.HostCallbackID, credential, key) + if errDiscovery != nil { + if stale, okStale := cachedAccountModels(cacheKey, false); okStale { + return filterDiscoveredModels(stale), fmt.Errorf("refresh models; using stale cache: %w", errDiscovery) + } + return configuredModels(), fmt.Errorf("discover models; using configured fallback: %w", errDiscovery) + } + storeAccountModels(cacheKey, models) + return filterDiscoveredModels(models), nil +} + +func fetchAvailableModels(hostCallbackID string, credential kiroCredential, key string) ([]pluginapi.ModelInfo, error) { + models := make([]pluginapi.ModelInfo, 0) + seenModels := make(map[string]struct{}) + seenTokens := make(map[string]struct{}) + nextToken := "" + + for page := 0; page < modelDiscoveryMaxPages; page++ { + endpoint, errEndpoint := modelDiscoveryURL(credential.Region, nextToken) + if errEndpoint != nil { + return nil, errEndpoint + } + rawResponse, errCall := invokeHost(pluginabi.MethodHostHTTPDo, rpcHostHTTPRequest{ + HostCallbackID: hostCallbackID, + Method: http.MethodGet, + URL: endpoint, + Headers: http.Header{ + "Accept": []string{"application/json"}, + "Authorization": []string{"Bearer " + key}, + "Tokentype": []string{"API_KEY"}, + "X-Amzn-Codewhisperer-Optout": []string{"false"}, + }, + }) + if errCall != nil { + return nil, fmt.Errorf("call model-list endpoint: %w", errCall) + } + var response pluginapi.HTTPResponse + if errUnmarshal := json.Unmarshal(rawResponse, &response); errUnmarshal != nil { + return nil, fmt.Errorf("decode host HTTP response: %w", errUnmarshal) + } + if response.StatusCode != http.StatusOK { + return nil, fmt.Errorf("model-list endpoint returned HTTP %d: %s", response.StatusCode, limitedErrorBody(response.Body)) + } + + var pageResponse kiroModelListResponse + if errUnmarshal := json.Unmarshal(response.Body, &pageResponse); errUnmarshal != nil { + return nil, fmt.Errorf("decode model-list response: %w", errUnmarshal) + } + for _, model := range pageResponse.Models { + modelID := strings.TrimSpace(model.ModelID) + if modelID == "" { + continue + } + keyModel := strings.ToLower(modelID) + if _, exists := seenModels[keyModel]; exists { + continue + } + seenModels[keyModel] = struct{}{} + models = append(models, discoveredModelInfo(model)) + } + + nextToken = strings.TrimSpace(pageResponse.NextToken) + if nextToken == "" { + return models, nil + } + if _, repeated := seenTokens[nextToken]; repeated { + return nil, fmt.Errorf("model-list pagination repeated a next token") + } + seenTokens[nextToken] = struct{}{} + } + + return models, nil +} + +func modelDiscoveryURL(region, nextToken string) (string, error) { + if !regionPattern.MatchString(region) { + return "", fmt.Errorf("invalid Kiro region %q", region) + } + query := url.Values{ + "origin": []string{"AI_EDITOR"}, + "maxResults": []string{fmt.Sprintf("%d", modelDiscoveryPageSize)}, + } + if nextToken != "" { + query.Set("nextToken", nextToken) + } + return "https://codewhisperer." + region + ".amazonaws.com/ListAvailableModels?" + query.Encode(), nil +} + +func discoveredModelInfo(model kiroAvailableModel) pluginapi.ModelInfo { + inputLimit := int64(model.TokenLimits.MaxInputTokens) + if inputLimit <= 0 { + inputLimit = 200000 + } + outputLimit := int64(model.TokenLimits.MaxOutputTokens) + if outputLimit <= 0 { + outputLimit = 64000 + } + displayName := firstNonEmpty(model.ModelName, model.ModelID) + return pluginapi.ModelInfo{ + ID: strings.TrimSpace(model.ModelID), + Name: displayName, + DisplayName: displayName, + Object: "model", + OwnedBy: "kiro", + Type: "claude", + Description: strings.TrimSpace(model.Description), + ContextLength: inputLimit, + InputTokenLimit: inputLimit, + OutputTokenLimit: outputLimit, + MaxCompletionTokens: outputLimit, + SupportedGenerationMethods: []string{"generateContent", "streamGenerateContent"}, + SupportedParameters: []string{"max_tokens", "temperature", "top_p", "tools"}, + SupportedInputModalities: discoveredInputModalities(model.SupportedInputTypes), + SupportedOutputModalities: []string{"text"}, + } +} + +func discoveredInputModalities(inputTypes []string) []string { + modalities := make([]string, 0, 2) + seen := make(map[string]struct{}, 2) + for _, inputType := range inputTypes { + var modality string + switch strings.ToUpper(strings.TrimSpace(inputType)) { + case "TEXT": + modality = "text" + case "IMAGE": + modality = "image" + default: + continue + } + if _, exists := seen[modality]; exists { + continue + } + seen[modality] = struct{}{} + modalities = append(modalities, modality) + } + if len(modalities) == 0 { + return []string{"text"} + } + return modalities +} + +func filterDiscoveredModels(models []pluginapi.ModelInfo) []pluginapi.ModelInfo { + configured := loadedConfig().Models + allowed := make(map[string]struct{}, len(configured)) + allowAll := false + for _, model := range configured { + model = strings.ToLower(strings.TrimSpace(model)) + if model == "*" { + allowAll = true + break + } + if model != "" { + allowed[model] = struct{}{} + } + } + + out := make([]pluginapi.ModelInfo, 0, len(models)) + for _, model := range models { + if !allowAll { + if _, okAllowed := allowed[strings.ToLower(strings.TrimSpace(model.ID))]; !okAllowed { + continue + } + } + out = append(out, model) + } + return out +} + +func modelCacheKey(authID, region, key string) string { + sum := sha256.Sum256([]byte(key)) + return strings.TrimSpace(authID) + "\x00" + region + "\x00" + hex.EncodeToString(sum[:8]) +} + +func cachedAccountModels(cacheKey string, freshOnly bool) ([]pluginapi.ModelInfo, bool) { + modelCacheMu.Lock() + defer modelCacheMu.Unlock() + entry, okCache := modelCache[cacheKey] + if !okCache || freshOnly && time.Since(entry.FetchedAt) >= modelCacheTTL { + return nil, false + } + return append([]pluginapi.ModelInfo(nil), entry.Models...), true +} + +func storeAccountModels(cacheKey string, models []pluginapi.ModelInfo) { + modelCacheMu.Lock() + modelCache[cacheKey] = modelCacheEntry{Models: append([]pluginapi.ModelInfo(nil), models...), FetchedAt: time.Now()} + modelCacheMu.Unlock() +} + +func resetModelCache() { + modelCacheMu.Lock() + modelCache = make(map[string]modelCacheEntry) + modelCacheMu.Unlock() +} + +func logModelDiscoveryFallback(request rpcAuthModelRequest, errDiscovery error) { + _, _ = invokeHost(pluginabi.MethodHostLog, rpcHostLogRequest{ + HostCallbackID: request.HostCallbackID, + Level: "warn", + Message: "Kiro model discovery failed; using cached or configured models", + Fields: map[string]any{ + "auth_id": request.AuthID, + "error": errDiscovery.Error(), + }, + }) +} + +func limitedErrorBody(body []byte) string { + const limit = 1024 + body = []byte(strings.TrimSpace(string(body))) + if len(body) > limit { + body = body[:limit] + } + return string(body) +} diff --git a/examples/plugin/kiro/go/protocol.go b/examples/plugin/kiro/go/protocol.go new file mode 100644 index 00000000000..479f943919e --- /dev/null +++ b/examples/plugin/kiro/go/protocol.go @@ -0,0 +1,165 @@ +package main + +import "encoding/json" + +type kiroPayload struct { + ConversationState struct { + AgentContinuationID string `json:"agentContinuationId,omitempty"` + AgentTaskType string `json:"agentTaskType,omitempty"` + ChatTriggerType string `json:"chatTriggerType"` + ConversationID string `json:"conversationId"` + CurrentMessage kiroCurrentMessage `json:"currentMessage"` + History []kiroHistoryMessage `json:"history,omitempty"` + } `json:"conversationState"` + InferenceConfig *kiroInferenceConfig `json:"inferenceConfig,omitempty"` + ToolNameMap map[string]string `json:"-"` + EstimatedInputTokens int `json:"-"` +} + +type kiroCurrentMessage struct { + UserInputMessage kiroUserInputMessage `json:"userInputMessage"` +} + +type kiroUserInputMessage struct { + Content string `json:"content"` + ModelID string `json:"modelId,omitempty"` + Origin string `json:"origin"` + Images []kiroImage `json:"images,omitempty"` + UserInputMessageContext *kiroUserMessageContext `json:"userInputMessageContext,omitempty"` +} + +type kiroUserMessageContext struct { + Tools []kiroToolWrapper `json:"tools,omitempty"` + ToolResults []kiroToolResult `json:"toolResults,omitempty"` +} + +type kiroToolWrapper struct { + ToolSpecification kiroToolSpecification `json:"toolSpecification"` +} + +type kiroToolSpecification struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema kiroInputSchema `json:"inputSchema"` +} + +type kiroInputSchema struct { + JSON any `json:"json"` +} + +type kiroToolResult struct { + ToolUseID string `json:"toolUseId"` + Content []kiroResultContent `json:"content"` + Status string `json:"status"` +} + +type kiroResultContent struct { + Text string `json:"text"` +} + +type kiroImage struct { + Format string `json:"format"` + Source kiroImageSource `json:"source"` +} + +type kiroImageSource struct { + Bytes string `json:"bytes"` +} + +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"` +} + +type kiroToolUse struct { + ToolUseID string `json:"toolUseId"` + Name string `json:"name"` + Input map[string]any `json:"input"` +} + +type kiroInferenceConfig struct { + MaxTokens int `json:"maxTokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` +} + +type claudeRequest struct { + Model string `json:"model"` + Messages []claudeMessage `json:"messages"` + MaxTokens int `json:"max_tokens"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"top_p,omitempty"` + Stream bool `json:"stream,omitempty"` + System any `json:"system,omitempty"` + Thinking *struct { + Type string `json:"type,omitempty"` + } `json:"thinking,omitempty"` + Tools []claudeTool `json:"tools,omitempty"` +} + +type claudeMessage struct { + Role string `json:"role"` + Content any `json:"content"` +} + +type claudeTool struct { + Type string `json:"type,omitempty"` + Name string `json:"name"` + Description string `json:"description"` + InputSchema any `json:"input_schema"` +} + +type claudeContentBlock struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + Thinking string `json:"thinking,omitempty"` + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Input map[string]any `json:"input,omitempty"` +} + +type claudeUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +type claudeResponse struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Content []claudeContentBlock `json:"content"` + Model string `json:"model"` + StopReason string `json:"stop_reason"` + StopSequence *string `json:"stop_sequence"` + Usage claudeUsage `json:"usage"` +} + +type kiroEvent struct { + Type string + Payload map[string]any +} + +type responseAccumulator struct { + Model string + ID string + Blocks []claudeContentBlock + InputTokens int + OutputTokens int + StopReason string + Credits float64 + EstimatedInputTokens int + ContextUsagePercentage float64 + ToolNames map[string]string + pendingTools pendingToolUses +} + +func decodeJSONMap(raw json.RawMessage) map[string]any { + var value map[string]any + _ = json.Unmarshal(raw, &value) + return value +} diff --git a/examples/plugin/kiro/go/response.go b/examples/plugin/kiro/go/response.go new file mode 100644 index 00000000000..d268b14842c --- /dev/null +++ b/examples/plugin/kiro/go/response.go @@ -0,0 +1,296 @@ +package main + +import ( + "encoding/json" + "fmt" + "strings" +) + +func newAccumulator(payload *kiroPayload) *responseAccumulator { + return &responseAccumulator{ + Model: payload.ConversationState.CurrentMessage.UserInputMessage.ModelID, + ID: "msg_" + strings.ReplaceAll(randomUUID(), "-", ""), + ToolNames: payload.ToolNameMap, + EstimatedInputTokens: payload.EstimatedInputTokens, + } +} + +func (a *responseAccumulator) accept(event kiroEvent) ([]claudeContentBlock, error) { + a.InputTokens, a.OutputTokens = updateUsage(event.Payload, a.InputTokens, a.OutputTokens) + switch event.Type { + case "assistantResponseEvent": + if text := firstStringField(event.Payload, "content", "text"); text != "" { + block := claudeContentBlock{Type: "text", Text: text} + a.appendFragment(block) + return []claudeContentBlock{block}, nil + } + case "reasoningContentEvent": + if text := firstStringField(event.Payload, "text", "content"); text != "" { + block := claudeContentBlock{Type: "thinking", Thinking: text} + a.appendFragment(block) + return []claudeContentBlock{block}, nil + } + case "toolUseEvent": + tools, errTools := a.pendingTools.accept(event.Payload) + if errTools != nil { + return nil, errTools + } + blocks := make([]claudeContentBlock, 0, len(tools)) + for _, tool := range tools { + name := tool.Name + if original := a.ToolNames[name]; original != "" { + name = original + } + block := claudeContentBlock{Type: "tool_use", ID: tool.ToolUseID, Name: name, Input: tool.Input} + a.Blocks = append(a.Blocks, block) + blocks = append(blocks, block) + } + return blocks, nil + case "metadataEvent": + if reason := firstStringField(event.Payload, "stopReason", "stop_reason"); reason != "" { + a.StopReason = mapStopReason(reason) + } + case "meteringEvent": + if credits, ok := event.Payload["usage"].(float64); ok { + a.Credits += credits + } + case "contextUsageEvent": + if percentage, ok := readFloat(event.Payload, "contextUsagePercentage", "context_usage_percentage"); ok && percentage > 0 { + a.ContextUsagePercentage = percentage + } + } + return nil, nil +} + +func (a *responseAccumulator) appendFragment(block claudeContentBlock) { + if len(a.Blocks) > 0 { + last := &a.Blocks[len(a.Blocks)-1] + if last.Type == block.Type { + switch block.Type { + case "text": + last.Text += block.Text + return + case "thinking": + last.Thinking += block.Thinking + return + } + } + } + a.Blocks = append(a.Blocks, block) +} + +func (a *responseAccumulator) finish() error { + tools, errTools := a.pendingTools.flush() + if errTools != nil { + return errTools + } + for _, tool := range tools { + name := tool.Name + if original := a.ToolNames[name]; original != "" { + name = original + } + a.Blocks = append(a.Blocks, claudeContentBlock{Type: "tool_use", ID: tool.ToolUseID, Name: name, Input: tool.Input}) + } + if len(a.Blocks) == 0 { + return fmt.Errorf("Kiro stream ended before producing output") + } + if a.StopReason == "" { + for _, block := range a.Blocks { + if block.Type == "tool_use" { + a.StopReason = "tool_use" + break + } + } + if a.StopReason == "" { + a.StopReason = "end_turn" + } + } + if a.InputTokens <= 0 { + a.InputTokens = a.currentInputTokens() + } + if a.OutputTokens <= 0 { + a.OutputTokens = estimateClaudeOutputTokens(a.Blocks) + } + return nil +} + +func (a *responseAccumulator) currentInputTokens() int { + if a.InputTokens > 0 { + return a.InputTokens + } + if a.ContextUsagePercentage > 0 { + return int(a.ContextUsagePercentage * float64(contextWindowTokens(a.Model)) / 100) + } + return a.EstimatedInputTokens +} + +func (a *responseAccumulator) responseJSON() ([]byte, error) { + return json.Marshal(claudeResponse{ + ID: a.ID, + Type: "message", + Role: "assistant", + Content: a.Blocks, + Model: a.Model, + StopReason: a.StopReason, + StopSequence: nil, + Usage: claudeUsage{InputTokens: a.InputTokens, OutputTokens: a.OutputTokens}, + }) +} + +func mapStopReason(reason string) string { + switch strings.ToLower(strings.TrimSpace(reason)) { + case "max_tokens", "max-tokens", "max tokens", "length", "max_tokens_reached": + return "max_tokens" + case "tool_use", "tool-use", "tool use", "tool_call": + return "tool_use" + case "stop_sequence", "stop-sequence": + return "stop_sequence" + default: + return "end_turn" + } +} + +type claudeSSEWriter struct { + accumulator *responseAccumulator + started bool + blockIndex int + openType string + openIndex int +} + +func newClaudeSSEWriter(accumulator *responseAccumulator) *claudeSSEWriter { + return &claudeSSEWriter{accumulator: accumulator} +} + +func (w *claudeSSEWriter) start() ([][]byte, error) { + if w.started { + return nil, nil + } + w.started = true + message := map[string]any{ + "id": w.accumulator.ID, "type": "message", "role": "assistant", "model": w.accumulator.Model, + "content": []any{}, "stop_reason": nil, "stop_sequence": nil, + "usage": map[string]int{"input_tokens": w.accumulator.currentInputTokens(), "output_tokens": 0}, + } + return marshalSSE("message_start", map[string]any{"type": "message_start", "message": message}) +} + +func (w *claudeSSEWriter) blocks(blocks []claudeContentBlock) ([][]byte, error) { + var frames [][]byte + for _, block := range blocks { + switch block.Type { + case "thinking": + blockFrames, errBlock := w.fragment("thinking", map[string]any{"type": "thinking", "thinking": ""}, map[string]any{"type": "thinking_delta", "thinking": block.Thinking}) + if errBlock != nil { + return nil, errBlock + } + frames = append(frames, blockFrames...) + case "tool_use": + closeFrames, errClose := w.closeOpenBlock() + if errClose != nil { + return nil, errClose + } + frames = append(frames, closeFrames...) + + index := w.blockIndex + w.blockIndex++ + contentBlock := map[string]any{"type": "tool_use", "id": block.ID, "name": block.Name, "input": map[string]any{}} + inputJSON, errMarshal := json.Marshal(block.Input) + if errMarshal != nil { + return nil, errMarshal + } + delta := map[string]any{"type": "input_json_delta", "partial_json": string(inputJSON)} + blockFrames, errBlock := marshalCompleteContentBlock(index, contentBlock, delta) + if errBlock != nil { + return nil, errBlock + } + frames = append(frames, blockFrames...) + default: + blockFrames, errBlock := w.fragment("text", map[string]any{"type": "text", "text": ""}, map[string]any{"type": "text_delta", "text": block.Text}) + if errBlock != nil { + return nil, errBlock + } + frames = append(frames, blockFrames...) + } + } + return frames, nil +} + +func (w *claudeSSEWriter) fragment(blockType string, contentBlock, delta any) ([][]byte, error) { + var frames [][]byte + if w.openType != blockType { + closeFrames, errClose := w.closeOpenBlock() + if errClose != nil { + return nil, errClose + } + frames = append(frames, closeFrames...) + w.openType = blockType + w.openIndex = w.blockIndex + w.blockIndex++ + startFrames, errStart := marshalSSE("content_block_start", map[string]any{"type": "content_block_start", "index": w.openIndex, "content_block": contentBlock}) + if errStart != nil { + return nil, errStart + } + frames = append(frames, startFrames...) + } + deltaFrames, errDelta := marshalSSE("content_block_delta", map[string]any{"type": "content_block_delta", "index": w.openIndex, "delta": delta}) + if errDelta != nil { + return nil, errDelta + } + return append(frames, deltaFrames...), nil +} + +func (w *claudeSSEWriter) closeOpenBlock() ([][]byte, error) { + if w.openType == "" { + return nil, nil + } + index := w.openIndex + w.openType = "" + return marshalSSE("content_block_stop", map[string]any{"type": "content_block_stop", "index": index}) +} + +func marshalCompleteContentBlock(index int, contentBlock, delta any) ([][]byte, error) { + startFrames, errStart := marshalSSE("content_block_start", map[string]any{"type": "content_block_start", "index": index, "content_block": contentBlock}) + if errStart != nil { + return nil, errStart + } + deltaFrames, errDelta := marshalSSE("content_block_delta", map[string]any{"type": "content_block_delta", "index": index, "delta": delta}) + if errDelta != nil { + return nil, errDelta + } + frames := append(startFrames, deltaFrames...) + stopFrames, errStop := marshalSSE("content_block_stop", map[string]any{"type": "content_block_stop", "index": index}) + if errStop != nil { + return nil, errStop + } + return append(frames, stopFrames...), nil +} + +func (w *claudeSSEWriter) finish() ([][]byte, error) { + closeFrames, errClose := w.closeOpenBlock() + if errClose != nil { + return nil, errClose + } + deltaFrames, errDelta := marshalSSE("message_delta", map[string]any{ + "type": "message_delta", + "delta": map[string]any{"stop_reason": w.accumulator.StopReason, "stop_sequence": nil}, + "usage": map[string]int{"input_tokens": w.accumulator.InputTokens, "output_tokens": w.accumulator.OutputTokens}, + }) + if errDelta != nil { + return nil, errDelta + } + stopFrames, errStop := marshalSSE("message_stop", map[string]any{"type": "message_stop"}) + if errStop != nil { + return nil, errStop + } + frames := append(closeFrames, deltaFrames...) + return append(frames, stopFrames...), nil +} + +func marshalSSE(event string, value any) ([][]byte, error) { + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + return nil, errMarshal + } + return [][]byte{[]byte("event: " + event + "\ndata: " + string(raw) + "\n\n")}, nil +} diff --git a/examples/plugin/kiro/go/token_estimator.go b/examples/plugin/kiro/go/token_estimator.go new file mode 100644 index 00000000000..837d89e4408 --- /dev/null +++ b/examples/plugin/kiro/go/token_estimator.go @@ -0,0 +1,137 @@ +package main + +import ( + "encoding/json" + "math" + "regexp" + "strconv" + "strings" +) + +var claudeVersionPattern = regexp.MustCompile(`claude-(?:opus|sonnet|haiku)-(\d+)(?:[.-](\d+))?`) + +func estimateApproxTokens(text string) int { + if text == "" { + return 0 + } + runes := []rune(text) + if len(runes) < 5 { + return max(1, int(math.Ceil(float64(len(runes))/3))) + } + var ascii, digits, symbols, nonASCII int + for _, value := range runes { + switch { + case value >= 0x80: + nonASCII++ + case value >= '0' && value <= '9': + digits++ + case (value >= '!' && value <= '/') || (value >= ':' && value <= '@') || + (value >= '[' && value <= '`') || (value >= '{' && value <= '~'): + symbols++ + default: + ascii++ + } + } + estimated := int(math.Ceil(float64(ascii)/4.5 + float64(digits)/2 + float64(symbols)/1.5 + float64(nonASCII)/1.5)) + if estimated < 1 { + return 1 + } + return estimated +} + +func estimateClaudeRequestInputTokens(request *claudeRequest) int { + if request == nil { + return 0 + } + total := estimateClaudeValueTokens(request.System) + for _, message := range request.Messages { + total += estimateClaudeValueTokens(message.Content) + } + for _, tool := range request.Tools { + total += estimateApproxTokens(tool.Name) + total += estimateApproxTokens(tool.Description) + total += estimateJSONTokens(tool.InputSchema) + } + return total +} + +func estimateClaudeOutputTokens(blocks []claudeContentBlock) int { + total := 0 + for _, block := range blocks { + switch block.Type { + case "text": + total += estimateApproxTokens(block.Text) + case "thinking": + total += estimateApproxTokens(block.Thinking) + case "tool_use": + total += estimateApproxTokens(block.Name) + total += estimateJSONTokens(block.Input) + } + } + return total +} + +func estimateClaudeValueTokens(value any) int { + switch typed := value.(type) { + case nil: + return 0 + case string: + return estimateApproxTokens(typed) + case []any: + total := 0 + for _, item := range typed { + total += estimateClaudeValueTokens(item) + } + return total + case map[string]any: + blockType, _ := typed["type"].(string) + switch blockType { + case "text", "input_text", "output_text": + if text, ok := typed["text"].(string); ok { + return estimateApproxTokens(text) + } + case "thinking": + if thinking, ok := typed["thinking"].(string); ok { + return estimateApproxTokens(thinking) + } + case "tool_use": + return estimateApproxTokens(stringValue(typed["name"])) + estimateJSONTokens(typed["input"]) + case "tool_result": + return estimateClaudeValueTokens(typed["content"]) + } + return estimateJSONTokens(typed) + default: + return estimateJSONTokens(typed) + } +} + +func estimateJSONTokens(value any) int { + if value == nil { + return 0 + } + raw, errMarshal := json.Marshal(value) + if errMarshal != nil { + return 0 + } + return estimateApproxTokens(string(raw)) +} + +func stringValue(value any) string { + text, _ := value.(string) + return text +} + +func contextWindowTokens(model string) int { + match := claudeVersionPattern.FindStringSubmatch(strings.ToLower(model)) + if len(match) == 3 { + major, errMajor := strconv.Atoi(match[1]) + minor := 0 + if match[2] != "" { + minor, _ = strconv.Atoi(match[2]) + } + if errMajor == nil && (major > 4 || major == 4 && minor >= 6) { + return 1_000_000 + } + } + return 200_000 +} diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go new file mode 100644 index 00000000000..f5931ffe5ff --- /dev/null +++ b/examples/plugin/kiro/go/translate.go @@ -0,0 +1,405 @@ +package main + +import ( + "crypto/rand" + "crypto/sha1" // #nosec G505 -- SHA-1 is used only for deterministic UUID-compatible identifiers. + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "regexp" + "strings" +) + +const ( + maxToolDescriptionLength = 10237 + maxKiroPayloadBytes = 900 * 1024 +) + +var ( + modelDashVersionPattern = regexp.MustCompile(`^claude-(opus|sonnet|haiku)-(\d+)-(\d{1,2})([^0-9].*)?$`) + invalidToolNamePattern = regexp.MustCompile(`[^A-Za-z0-9_-]+`) +) + +func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeRequest, error) { + var request claudeRequest + if errUnmarshal := json.Unmarshal(raw, &request); errUnmarshal != nil { + return nil, nil, fmt.Errorf("decode Claude request: %w", errUnmarshal) + } + model := normalizeKiroModel(firstNonEmpty(requestedModel, request.Model)) + if model == "" { + return nil, nil, fmt.Errorf("model is required") + } + if len(request.Messages) == 0 { + return nil, nil, fmt.Errorf("messages must not be empty") + } + + systemPrompt := extractSystemPrompt(request.System) + if request.Thinking != nil { + kind := strings.ToLower(strings.TrimSpace(request.Thinking.Type)) + if kind == "enabled" || kind == "adaptive" { + systemPrompt = strings.TrimSpace("enabled\n200000\n\n" + systemPrompt) + } + } + + payload := &kiroPayload{} + payload.ConversationState.ChatTriggerType = "MANUAL" + payload.ConversationState.AgentTaskType = "vibe" + payload.ConversationState.AgentContinuationID = randomUUID() + payload.ConversationState.ConversationID = conversationID(model, systemPrompt, firstUserAnchor(request.Messages)) + + history := make([]kiroHistoryMessage, 0, len(request.Messages)+2) + if systemPrompt != "" { + history = append(history, + kiroHistoryMessage{UserInputMessage: &kiroUserInputMessage{Content: systemPrompt, ModelID: model, Origin: "KIRO_CLI"}}, + kiroHistoryMessage{AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "I will follow these instructions."}}, + ) + } + + var currentText string + var currentImages []kiroImage + var currentToolResults []kiroToolResult + for index, message := range request.Messages { + last := index == len(request.Messages)-1 + switch strings.ToLower(strings.TrimSpace(message.Role)) { + case "user": + text, images, toolResults := extractUserContent(message.Content) + if last { + currentText, currentImages, currentToolResults = text, images, toolResults + continue + } + userMessage := &kiroUserInputMessage{Content: fallbackContent(text, len(images) > 0), ModelID: model, Origin: "KIRO_CLI", Images: images} + if len(toolResults) > 0 { + userMessage.UserInputMessageContext = &kiroUserMessageContext{ToolResults: toolResults} + } + history = append(history, kiroHistoryMessage{UserInputMessage: userMessage}) + case "assistant": + text, tools := extractAssistantContent(message.Content) + history = append(history, kiroHistoryMessage{AssistantResponseMessage: &kiroAssistantResponseMessage{Content: text, ToolUses: tools}}) + } + } + + tools, nameMap := convertTools(request.Tools) + payload.ToolNameMap = nameMap + if len(currentToolResults) > 0 { + currentText = joinNonEmpty(currentText, readableToolResults(currentToolResults)) + } + current := kiroUserInputMessage{ + Content: fallbackContent(currentText, len(currentImages) > 0), + ModelID: model, + Origin: "KIRO_CLI", + Images: currentImages, + } + if len(tools) > 0 || len(currentToolResults) > 0 { + current.UserInputMessageContext = &kiroUserMessageContext{Tools: tools, ToolResults: currentToolResults} + } + payload.ConversationState.CurrentMessage.UserInputMessage = current + payload.ConversationState.History = trimLeadingAssistant(history) + payload.EstimatedInputTokens = estimateClaudeRequestInputTokens(&request) + if request.MaxTokens > 0 || request.Temperature > 0 || request.TopP > 0 { + payload.InferenceConfig = &kiroInferenceConfig{MaxTokens: request.MaxTokens, Temperature: request.Temperature, TopP: request.TopP} + } + truncatePayload(payload) + return payload, &request, nil +} + +func normalizeKiroModel(model string) string { + model = strings.TrimSpace(model) + model = strings.TrimSuffix(model, "-thinking") + if model == "claude-sonnet-4-20250514" { + return "claude-sonnet-4" + } + if match := modelDashVersionPattern.FindStringSubmatch(model); len(match) == 5 { + return fmt.Sprintf("claude-%s-%s.%s%s", match[1], match[2], match[3], match[4]) + } + return model +} + +func extractSystemPrompt(value any) string { + switch system := value.(type) { + case string: + return strings.TrimSpace(system) + case []any: + parts := make([]string, 0, len(system)) + for _, item := range system { + block, ok := item.(map[string]any) + if !ok { + continue + } + if text, okText := block["text"].(string); okText && text != "" { + parts = append(parts, text) + } + } + return strings.TrimSpace(strings.Join(parts, "\n")) + default: + return "" + } +} + +func extractUserContent(content any) (string, []kiroImage, []kiroToolResult) { + if text, ok := content.(string); ok { + return text, nil, nil + } + var texts []string + var images []kiroImage + var results []kiroToolResult + for _, block := range contentBlocks(content) { + typeName, _ := block["type"].(string) + switch typeName { + case "text", "input_text": + if text, ok := block["text"].(string); ok { + texts = append(texts, text) + } + case "image", "input_image": + if image := extractImage(block); image != nil { + images = append(images, *image) + } + case "tool_result": + toolUseID, _ := block["tool_use_id"].(string) + text, resultImages := extractToolResultContent(block["content"]) + images = append(images, resultImages...) + results = append(results, kiroToolResult{ToolUseID: toolUseID, Content: []kiroResultContent{{Text: fallbackContent(text, len(resultImages) > 0)}}, Status: "success"}) + } + } + return strings.Join(texts, ""), images, results +} + +func extractAssistantContent(content any) (string, []kiroToolUse) { + if text, ok := content.(string); ok { + return text, nil + } + var texts []string + var tools []kiroToolUse + for _, block := range contentBlocks(content) { + typeName, _ := block["type"].(string) + switch typeName { + case "text": + if text, ok := block["text"].(string); ok { + texts = append(texts, text) + } + case "tool_use": + id, _ := block["id"].(string) + name, _ := block["name"].(string) + input, _ := block["input"].(map[string]any) + if input == nil { + input = map[string]any{} + } + tools = append(tools, kiroToolUse{ToolUseID: id, Name: sanitizeToolName(name), Input: input}) + } + } + return strings.Join(texts, ""), tools +} + +func contentBlocks(content any) []map[string]any { + items, ok := content.([]any) + if !ok { + return nil + } + blocks := make([]map[string]any, 0, len(items)) + for _, item := range items { + if block, okBlock := item.(map[string]any); okBlock { + blocks = append(blocks, block) + } + } + return blocks +} + +func extractImage(block map[string]any) *kiroImage { + source, ok := block["source"].(map[string]any) + if !ok { + return nil + } + data, _ := source["data"].(string) + mediaType, _ := source["media_type"].(string) + if data == "" { + return nil + } + if _, errDecode := base64.StdEncoding.DecodeString(data); errDecode != nil { + return nil + } + format := strings.TrimPrefix(strings.ToLower(mediaType), "image/") + if format == "jpg" { + format = "jpeg" + } + if format != "png" && format != "jpeg" && format != "gif" && format != "webp" { + return nil + } + return &kiroImage{Format: format, Source: kiroImageSource{Bytes: data}} +} + +func extractToolResultContent(content any) (string, []kiroImage) { + if text, ok := content.(string); ok { + return text, nil + } + var texts []string + var images []kiroImage + for _, block := range contentBlocks(content) { + if text, ok := block["text"].(string); ok { + texts = append(texts, text) + } + if image := extractImage(block); image != nil { + images = append(images, *image) + } + } + return strings.Join(texts, ""), images +} + +func convertTools(tools []claudeTool) ([]kiroToolWrapper, map[string]string) { + out := make([]kiroToolWrapper, 0, len(tools)) + nameMap := make(map[string]string) + used := make(map[string]int) + for _, tool := range tools { + if strings.HasPrefix(strings.ToLower(tool.Type), "web_search") { + continue + } + name := uniqueToolName(sanitizeToolName(tool.Name), used) + if name != tool.Name { + nameMap[name] = tool.Name + } + description := strings.TrimSpace(tool.Description) + if description == "" { + description = "Call " + name + "." + } + if len(description) > maxToolDescriptionLength { + description = description[:maxToolDescriptionLength] + } + schema, ok := tool.InputSchema.(map[string]any) + if !ok || schema == nil { + schema = map[string]any{"type": "object"} + } + if _, exists := schema["type"]; !exists { + schema["type"] = "object" + } + out = append(out, kiroToolWrapper{ToolSpecification: kiroToolSpecification{ + Name: name, Description: description, InputSchema: kiroInputSchema{JSON: schema}, + }}) + } + if len(nameMap) == 0 { + nameMap = nil + } + return out, nameMap +} + +func sanitizeToolName(name string) string { + name = invalidToolNamePattern.ReplaceAllString(strings.TrimSpace(name), "_") + name = strings.Trim(name, "_") + if name == "" { + name = "tool" + } + if len(name) > 64 { + name = name[:64] + } + return name +} + +func uniqueToolName(name string, used map[string]int) string { + used[name]++ + if used[name] == 1 { + return name + } + suffix := fmt.Sprintf("_%d", used[name]) + if len(name)+len(suffix) > 64 { + name = name[:64-len(suffix)] + } + return name + suffix +} + +func readableToolResults(results []kiroToolResult) string { + parts := make([]string, 0, len(results)) + for _, result := range results { + for _, content := range result.Content { + if text := strings.TrimSpace(content.Text); text != "" { + parts = append(parts, text) + } + } + } + if len(parts) == 0 { + return "" + } + return "Tool results:\n\n" + strings.Join(parts, "\n\n") +} + +func fallbackContent(content string, hasImage bool) string { + if strings.TrimSpace(content) != "" { + return content + } + if hasImage { + return "Please analyze the attached image." + } + return "." +} + +func joinNonEmpty(left, right string) string { + left, right = strings.TrimSpace(left), strings.TrimSpace(right) + if left == "" { + return right + } + if right == "" { + return left + } + return left + "\n\n" + right +} + +func trimLeadingAssistant(history []kiroHistoryMessage) []kiroHistoryMessage { + for len(history) > 0 && history[0].AssistantResponseMessage != nil { + history = history[1:] + } + return history +} + +func firstUserAnchor(messages []claudeMessage) string { + for _, message := range messages { + if strings.EqualFold(message.Role, "user") { + text, _, _ := extractUserContent(message.Content) + if text = strings.TrimSpace(text); text != "" && text != "." { + return text + } + } + } + return "" +} + +func conversationID(model, systemPrompt, anchor string) string { + if strings.TrimSpace(anchor) == "" { + return randomUUID() + } + seed := strings.Join([]string{model, strings.TrimSpace(systemPrompt), strings.TrimSpace(anchor)}, "\n") + sum := sha1.Sum(append([]byte("6ba7b8119dad11d180b400c04fd430c8"), []byte(seed)...)) + b := sum[:16] + b[6] = (b[6] & 0x0f) | 0x50 + b[8] = (b[8] & 0x3f) | 0x80 + return formatUUID(b) +} + +func randomUUID() string { + b := make([]byte, 16) + if _, errRead := rand.Read(b); errRead != nil { + sum := sha1.Sum([]byte(fmt.Sprintf("fallback-%p", &b))) + b = sum[:16] + } + b[6] = (b[6] & 0x0f) | 0x40 + b[8] = (b[8] & 0x3f) | 0x80 + return formatUUID(b) +} + +func formatUUID(b []byte) string { + hexValue := hex.EncodeToString(b) + return hexValue[:8] + "-" + hexValue[8:12] + "-" + hexValue[12:16] + "-" + hexValue[16:20] + "-" + hexValue[20:32] +} + +func truncatePayload(payload *kiroPayload) { + for { + raw, errMarshal := json.Marshal(payload) + if errMarshal != nil || len(raw) <= maxKiroPayloadBytes || len(payload.ConversationState.History) <= 4 { + return + } + removeAt := 0 + if len(payload.ConversationState.History) >= 2 && payload.ConversationState.History[0].UserInputMessage != nil && payload.ConversationState.History[1].AssistantResponseMessage != nil { + removeAt = 2 + } + if removeAt >= len(payload.ConversationState.History) { + return + } + payload.ConversationState.History = append(payload.ConversationState.History[:removeAt], payload.ConversationState.History[removeAt+1:]...) + } +} From 12b289ed712b7fc6e6e8363e16e6a38be3e7c4a3 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 15:32:41 +0800 Subject: [PATCH 02/20] fix(pluginhost): handle executor streaming and usage --- internal/pluginhost/adapters_executors.go | 36 +++++++- internal/pluginhost/adapters_test.go | 67 +++++++++++++++ internal/pluginhost/executor_usage.go | 100 ++++++++++++++++++++++ 3 files changed, 199 insertions(+), 4 deletions(-) create mode 100644 internal/pluginhost/executor_usage.go diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index 38e9e383f85..f4078810d7d 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -581,13 +581,35 @@ func (a *executorAdapter) translateExecutorStreamPayload(ctx context.Context, pr if len(originalRequest) == 0 { originalRequest = prepared.req.Payload } - frames := sdktranslator.TranslateStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) - if executorStreamTranslationFellBack(prepared, payload, frames) { - return nil + translationPayloads := executorStreamTranslationPayloads(payload) + frames := make([][]byte, 0, len(translationPayloads)) + for _, translationPayload := range translationPayloads { + translated := sdktranslator.TranslateStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, translationPayload, param) + if executorStreamTranslationFellBack(prepared, translationPayload, translated) { + continue + } + frames = append(frames, translated...) } return frames } +func executorStreamTranslationPayloads(payload []byte) [][]byte { + if !bytes.Contains(payload, []byte("\n")) { + return [][]byte{payload} + } + var dataLines [][]byte + for _, line := range bytes.Split(payload, []byte("\n")) { + trimmed := bytes.TrimSpace(line) + if bytes.HasPrefix(trimmed, []byte("data:")) { + dataLines = append(dataLines, bytes.Clone(trimmed)) + } + } + if len(dataLines) == 0 { + return [][]byte{payload} + } + return dataLines +} + func executorStreamTranslationFellBack(prepared preparedExecutorCall, payload []byte, frames [][]byte) bool { if prepared.requestedFormat == "" || prepared.outputFormat == "" || prepared.outputFormat == prepared.requestedFormat { return false @@ -636,6 +658,8 @@ func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) } + usageReporter := newPluginExecutorUsage(ctx, a, req.Model, auth) + defer usageReporter.trackFailure(ctx, &err) defer func() { if recovered := recover(); recovered != nil { a.host.fusePlugin(a.pluginID, "Executor.Execute", recovered) @@ -652,6 +676,7 @@ func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req if errExecute != nil { return coreexecutor.Response{}, errExecute } + usageReporter.publishNonStream(ctx, prepared.outputFormat, pluginResp.Payload) return coreexecutor.Response{ Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), Metadata: cloneAnyMap(pluginResp.Metadata), @@ -663,6 +688,8 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) } + usageReporter := newPluginExecutorUsage(ctx, a, req.Model, auth) + defer usageReporter.trackFailure(ctx, &err) defer func() { if recovered := recover(); recovered != nil { a.host.fusePlugin(a.pluginID, "Executor.ExecuteStream", recovered) @@ -679,9 +706,10 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth if errExecuteStream != nil { return nil, errExecuteStream } + nativeChunks := usageReporter.observeStream(ctx, prepared.outputFormat, pluginResp.Chunks) return &coreexecutor.StreamResult{ Headers: cloneHeader(pluginResp.Headers), - Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Chunks)), + Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, nativeChunks)), }, nil } diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 9540040f2d5..386142de549 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -15,6 +15,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -2848,6 +2849,72 @@ func TestExecutorAdapterUsesResponseFormatForOutputTranslation(t *testing.T) { } } +func TestExecutorStreamTranslationPayloadsExtractsDataFromSSEFrame(t *testing.T) { + payload := []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n") + frames := executorStreamTranslationPayloads(payload) + if len(frames) != 1 { + t.Fatalf("translation payload count = %d, want 1", len(frames)) + } + want := []byte(`data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hello"}}`) + if !bytes.Equal(frames[0], want) { + t.Fatalf("translation payload = %q, want %q", frames[0], want) + } +} + +func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { + var usageBuffer helps.StreamUsageBuffer + payload := []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":12,\"output_tokens\":4}}\n\n") + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, payload) + detail, ok := usageBuffer.Detail() + if !ok { + t.Fatal("stream usage was not observed") + } + if detail.InputTokens != 12 || detail.OutputTokens != 4 || detail.TotalTokens != 16 { + t.Fatalf("stream usage = %#v, want input=12 output=4 total=16", detail) + } +} + +func TestExecutorAdapterPublishesNativeUsageForBilling(t *testing.T) { + const model = "plugin-usage-billing-model" + records := make(chan coreusage.Record, 1) + coreusage.RegisterNamedPlugin("test:plugin-executor-usage", coreUsagePluginFunc(func(ctx context.Context, record coreusage.Record) { + if record.Model == model { + select { + case records <- record: + default: + } + } + })) + + host := New() + adapter := newCurrentExecutorAdapterForTest(host, "executor-usage", &fakeExecutor{ + execute: func(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { + return pluginapi.ExecutorResponse{Payload: []byte(`{"id":"msg_usage","type":"message","role":"assistant","model":"claude-test","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":12,"output_tokens":4}}`)}, nil + }, + }, []sdktranslator.Format{sdktranslator.FormatClaude}, []sdktranslator.Format{sdktranslator.FormatClaude}) + + _, errExecute := adapter.Execute(context.Background(), &coreauth.Auth{ID: "auth-usage", Provider: "plugin-provider"}, coreexecutor.Request{ + Model: model, + Format: sdktranslator.FormatClaude, + Payload: []byte(`{"model":"plugin-usage-billing-model","messages":[{"role":"user","content":"hi"}]}`), + }, coreexecutor.Options{SourceFormat: sdktranslator.FormatClaude, ResponseFormat: sdktranslator.FormatClaude}) + if errExecute != nil { + t.Fatal(errExecute) + } + + select { + case record := <-records: + if record.Provider != "plugin-provider" || record.AuthID != "auth-usage" { + t.Fatalf("billing identity = provider:%q auth:%q", record.Provider, record.AuthID) + } + if record.Detail.InputTokens != 12 || record.Detail.OutputTokens != 4 || record.Detail.TotalTokens != 16 { + t.Fatalf("billing usage = %#v, want input=12 output=4 total=16", record.Detail) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for plugin executor billing record") + } +} + func TestExecutorAdapterSelectsCustomOutputWithHostResponseTranslator(t *testing.T) { customOutputFormat := sdktranslator.Format("plugin-custom-output") requestedFormat := sdktranslator.FormatOpenAI diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go new file mode 100644 index 00000000000..06e1674bb37 --- /dev/null +++ b/internal/pluginhost/executor_usage.go @@ -0,0 +1,100 @@ +package pluginhost + +import ( + "bytes" + "context" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +type pluginExecutorUsage struct { + reporter *helps.UsageReporter +} + +func newPluginExecutorUsage(ctx context.Context, adapter *executorAdapter, model string, auth *coreauth.Auth) *pluginExecutorUsage { + return &pluginExecutorUsage{reporter: helps.NewExecutorUsageReporter(ctx, adapter, model, auth)} +} + +func (u *pluginExecutorUsage) trackFailure(ctx context.Context, errPtr *error) { + if u == nil || u.reporter == nil { + return + } + u.reporter.TrackFailure(ctx, errPtr) +} + +func (u *pluginExecutorUsage) publishNonStream(ctx context.Context, format sdktranslator.Format, payload []byte) { + if u == nil || u.reporter == nil { + return + } + detail, ok := pluginExecutorNonStreamUsage(format, payload) + if ok { + u.reporter.Publish(ctx, detail) + return + } + u.reporter.EnsurePublished(ctx) +} + +func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktranslator.Format, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { + if u == nil || u.reporter == nil || in == nil { + return in + } + out := make(chan pluginapi.ExecutorStreamChunk) + go func() { + defer close(out) + var usageBuffer helps.StreamUsageBuffer + var terminalErr error + for chunk := range in { + if chunk.Err != nil { + terminalErr = chunk.Err + } else { + observePluginExecutorStreamUsage(&usageBuffer, format, chunk.Payload) + } + if !sendExecutorPluginStreamChunk(ctx, out, chunk) { + if errContext := ctx.Err(); errContext != nil { + u.reporter.PublishFailure(ctx, errContext) + } + return + } + } + if terminalErr != nil { + u.reporter.PublishFailure(ctx, terminalErr) + return + } + if !usageBuffer.Publish(ctx, u.reporter) { + u.reporter.EnsurePublished(ctx) + } + }() + return out +} + +func pluginExecutorNonStreamUsage(format sdktranslator.Format, payload []byte) (coreusage.Detail, bool) { + switch format { + case sdktranslator.FormatClaude: + return helps.ParseClaudeUsage(payload), bytes.Contains(payload, []byte(`"usage"`)) + case sdktranslator.FormatOpenAI, sdktranslator.FormatOpenAIResponse: + return helps.ParseOpenAIUsage(payload), bytes.Contains(payload, []byte(`"usage"`)) + default: + return coreusage.Detail{}, false + } +} + +func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sdktranslator.Format, payload []byte) { + for _, line := range bytes.Split(payload, []byte("\n")) { + switch format { + case sdktranslator.FormatClaude: + detail, ok := helps.ParseClaudeStreamUsage(line) + buffer.Observe(detail, ok) + case sdktranslator.FormatOpenAI: + detail, ok := helps.ParseOpenAIStreamUsage(line) + buffer.Observe(detail, ok) + case sdktranslator.FormatOpenAIResponse: + jsonPayload := helps.JSONPayload(line) + detail, ok := helps.ParseCodexUsage(jsonPayload) + buffer.Observe(detail, ok) + } + } +} From 09f5355b76c01b8422cc9878a53e8c51cfc711b2 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 15:50:16 +0800 Subject: [PATCH 03/20] docs(plugin): note translator compatibility dependency --- examples/plugin/kiro/README.md | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/examples/plugin/kiro/README.md b/examples/plugin/kiro/README.md index 2a0026296cf..2dc4cc00540 100644 --- a/examples/plugin/kiro/README.md +++ b/examples/plugin/kiro/README.md @@ -1,9 +1,11 @@ # Kiro Provider Plugin This Go plugin adds an API-key-only Kiro provider to CLIProxyAPI. It uses the -Kiro CLI runtime protocol and exposes Claude as its native protocol. The host -translates OpenAI Chat Completions and Responses API requests to and from Claude -when those entry points are used. +Kiro CLI runtime protocol and exposes Claude as its native protocol. Claude +`/v1/messages` requests are supported directly. Complete non-streaming OpenAI +Chat Completions and Responses API compatibility additionally requires the +restricted translator work tracked in +[issue #4815](https://github.com/router-for-me/CLIProxyAPI/issues/4815). The implementation deliberately does not include Kiro OAuth, device login, refresh tokens, account selection, or its own HTTP client. CLIProxyAPI owns @@ -120,8 +122,10 @@ Run these after providing a real key: explicit allow-list is configured, confirm only its intersection appears. 2. Send a non-streaming `/v1/messages` text request. 3. Send the same request with `stream: true` and verify Anthropic SSE ordering. -4. Call `/v1/chat/completions` to verify host-side OpenAI translation. -5. Call `/v1/responses` to verify Responses API translation. +4. After issue #4815 is implemented, call `/v1/chat/completions` to verify + host-side OpenAI translation. +5. After issue #4815 is implemented, call `/v1/responses` to verify Responses + API translation. 6. Exercise one client tool call and return its `tool_result`. 7. Send a small base64 image if the selected account model supports images. 8. Use an invalid key and verify a 401/403 does not leak the key. @@ -133,6 +137,8 @@ Run these after providing a real key: - The Kiro subscription API key is documented, but the raw runtime protocol is not a public model API and may change. +- Complete non-streaming OpenAI-compatible endpoint support depends on the + maintainer-owned translator changes tracked in issue #4815. - Model discovery uses Kiro's internal `ListAvailableModels` service rather than a documented public model API and may change. - Discovery is refreshed when the host registers an auth. The five-minute From c826e4100adf3eba0804dac5b7c1e14e53e32bd7 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 16:41:53 +0800 Subject: [PATCH 04/20] fix(plugin): address Kiro review feedback --- examples/plugin/kiro/go/kiro_test.go | 37 +++++++++++++++++++++++++++ examples/plugin/kiro/go/translate.go | 9 ++++++- internal/pluginhost/adapters_test.go | 17 ++++++++++++ internal/pluginhost/executor_usage.go | 16 +++++++++++- 4 files changed, 77 insertions(+), 2 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 0bcced94a13..428b693739a 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -239,6 +239,9 @@ func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { if context == nil || len(context.Tools) != 1 || len(context.ToolResults) != 1 { t.Fatalf("current context = %#v", context) } + if context.ToolResults[0].Status != "success" { + t.Fatalf("tool result status = %q, want success", context.ToolResults[0].Status) + } toolName := context.Tools[0].ToolSpecification.Name if toolName != "math_add_unsafe" || payload.ToolNameMap[toolName] != "math.add/unsafe" { t.Fatalf("tool name=%q map=%#v", toolName, payload.ToolNameMap) @@ -248,6 +251,40 @@ func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { } } +func TestClaudeToKiroPreservesToolResultErrorStatus(t *testing.T) { + raw := []byte(`{ + "model":"claude-sonnet-4-5", + "messages":[ + {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"lookup failed","is_error":true}]} + ] + }`) + payload, _, errTranslate := claudeToKiro(raw, "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + context := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext + if context == nil || len(context.ToolResults) != 1 { + t.Fatalf("current context = %#v", context) + } + if context.ToolResults[0].Status != "error" { + t.Fatalf("tool result status = %q, want error", context.ToolResults[0].Status) + } +} + +func TestClaudeToKiroRejectsAssistantPrefill(t *testing.T) { + raw := []byte(`{ + "model":"claude-sonnet-4-5", + "messages":[ + {"role":"user","content":"Complete this sentence"}, + {"role":"assistant","content":"The answer starts with"} + ] + }`) + if _, _, errTranslate := claudeToKiro(raw, ""); errTranslate == nil || !strings.Contains(errTranslate.Error(), "assistant prefills") { + t.Fatalf("claudeToKiro() error = %v, want unsupported assistant prefill error", errTranslate) + } +} + func TestNormalizeKiroModelDoesNotRewriteDatedSnapshotAsDecimal(t *testing.T) { if got := normalizeKiroModel("claude-sonnet-4-20250514"); got != "claude-sonnet-4" { t.Fatalf("normalizeKiroModel() = %q", got) diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index f5931ffe5ff..7f39dae4ee4 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -33,6 +33,9 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if len(request.Messages) == 0 { return nil, nil, fmt.Errorf("messages must not be empty") } + if lastRole := strings.ToLower(strings.TrimSpace(request.Messages[len(request.Messages)-1].Role)); lastRole != "user" { + return nil, nil, fmt.Errorf("last message must have role user; Kiro does not support assistant prefills") + } systemPrompt := extractSystemPrompt(request.System) if request.Thinking != nil { @@ -158,7 +161,11 @@ func extractUserContent(content any) (string, []kiroImage, []kiroToolResult) { toolUseID, _ := block["tool_use_id"].(string) text, resultImages := extractToolResultContent(block["content"]) images = append(images, resultImages...) - results = append(results, kiroToolResult{ToolUseID: toolUseID, Content: []kiroResultContent{{Text: fallbackContent(text, len(resultImages) > 0)}}, Status: "success"}) + status := "success" + if isError, _ := block["is_error"].(bool); isError { + status = "error" + } + results = append(results, kiroToolResult{ToolUseID: toolUseID, Content: []kiroResultContent{{Text: fallbackContent(text, len(resultImages) > 0)}}, Status: status}) } } return strings.Join(texts, ""), images, results diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 386142de549..74e1c5d2800 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2874,6 +2874,23 @@ func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { } } +func TestPluginExecutorUsageBuffersSplitClaudeStreamCounts(t *testing.T) { + var usageBuffer helps.StreamUsageBuffer + var pending []byte + pending = observePluginExecutorStreamChunk(&usageBuffer, sdktranslator.FormatClaude, pending, []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":12,")) + pending = observePluginExecutorStreamChunk(&usageBuffer, sdktranslator.FormatClaude, pending, []byte("\"output_tokens\":4}}\n\n")) + if len(pending) != 0 { + t.Fatalf("pending stream payload = %q, want empty", pending) + } + detail, ok := usageBuffer.Detail() + if !ok { + t.Fatal("stream usage was not observed") + } + if detail.InputTokens != 12 || detail.OutputTokens != 4 || detail.TotalTokens != 16 { + t.Fatalf("stream usage = %#v, want input=12 output=4 total=16", detail) + } +} + func TestExecutorAdapterPublishesNativeUsageForBilling(t *testing.T) { const model = "plugin-usage-billing-model" records := make(chan coreusage.Record, 1) diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go index 06e1674bb37..98840c3bf2f 100644 --- a/internal/pluginhost/executor_usage.go +++ b/internal/pluginhost/executor_usage.go @@ -46,12 +46,13 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans go func() { defer close(out) var usageBuffer helps.StreamUsageBuffer + var pendingUsagePayload []byte var terminalErr error for chunk := range in { if chunk.Err != nil { terminalErr = chunk.Err } else { - observePluginExecutorStreamUsage(&usageBuffer, format, chunk.Payload) + pendingUsagePayload = observePluginExecutorStreamChunk(&usageBuffer, format, pendingUsagePayload, chunk.Payload) } if !sendExecutorPluginStreamChunk(ctx, out, chunk) { if errContext := ctx.Err(); errContext != nil { @@ -64,6 +65,9 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans u.reporter.PublishFailure(ctx, terminalErr) return } + if len(pendingUsagePayload) > 0 { + observePluginExecutorStreamUsage(&usageBuffer, format, pendingUsagePayload) + } if !usageBuffer.Publish(ctx, u.reporter) { u.reporter.EnsurePublished(ctx) } @@ -71,6 +75,16 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans return out } +func observePluginExecutorStreamChunk(buffer *helps.StreamUsageBuffer, format sdktranslator.Format, pending, payload []byte) []byte { + pending = append(pending, payload...) + lastNewline := bytes.LastIndexByte(pending, '\n') + if lastNewline < 0 { + return pending + } + observePluginExecutorStreamUsage(buffer, format, pending[:lastNewline+1]) + return bytes.Clone(pending[lastNewline+1:]) +} + func pluginExecutorNonStreamUsage(format sdktranslator.Format, payload []byte) (coreusage.Detail, bool) { switch format { case sdktranslator.FormatClaude: From ac0ff1719746dae28bf5ad96e931a1790d60d512 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 17:38:17 +0800 Subject: [PATCH 05/20] docs(plugin): add Kiro client setup guide --- examples/plugin/kiro/README.md | 287 +++++++++++++++++++++++++++------ 1 file changed, 234 insertions(+), 53 deletions(-) diff --git a/examples/plugin/kiro/README.md b/examples/plugin/kiro/README.md index 2dc4cc00540..d9cfa444e04 100644 --- a/examples/plugin/kiro/README.md +++ b/examples/plugin/kiro/README.md @@ -1,30 +1,59 @@ # Kiro Provider Plugin This Go plugin adds an API-key-only Kiro provider to CLIProxyAPI. It uses the -Kiro CLI runtime protocol and exposes Claude as its native protocol. Claude -`/v1/messages` requests are supported directly. Complete non-streaming OpenAI -Chat Completions and Responses API compatibility additionally requires the -restricted translator work tracked in -[issue #4815](https://github.com/router-for-me/CLIProxyAPI/issues/4815). +Kiro CLI runtime protocol and exposes Claude as its native protocol, so Claude +Code can use it through `/v1/messages`. -The implementation deliberately does not include Kiro OAuth, device login, -refresh tokens, account selection, or its own HTTP client. CLIProxyAPI owns -credential selection, session affinity, cooldowns, protocol translation, proxy -configuration, and outbound request logging. +Codex CLI uses the OpenAI Responses API. CLIProxyAPI translates those requests +to the plugin's Claude-native protocol. Complete OpenAI Chat Completions and +Responses compatibility also depends on the translator work tracked in +[issue #4815](https://github.com/router-for-me/CLIProxyAPI/issues/4815). If a +build does not contain that work, Claude Code can still work while Codex CLI +may fail or return incomplete formatting or usage information. -## Build +## Before you start -Build for the current platform: +You need: + +- a Kiro API key and the region for that account; +- a CLIProxyAPI client key, such as `local-dev-key`; +- Go 1.26 or later to build the plugin; and +- Codex CLI or Claude Code for the client workflow you want to test. + +The two keys have different purposes: + +- `KIRO_API_KEY` authenticates CLIProxyAPI to Kiro. Only the server needs it. +- `local-dev-key` authenticates Codex CLI or Claude Code to CLIProxyAPI. It is + configured under `api-keys` in `config.yaml`. + +## 1. Build and install the plugin + +For a CLIProxyAPI process running directly on the current machine: ```bash make -C examples/plugin build-kiro + +plugin_os="$(go env GOOS)" +plugin_arch="$(go env GOARCH)" +case "$plugin_os" in + darwin) plugin_ext="dylib" ;; + windows) plugin_ext="dll" ;; + *) plugin_ext="so" ;; +esac + +mkdir -p "plugins/$plugin_os/$plugin_arch" +cp "examples/plugin/bin/kiro-go.$plugin_ext" "plugins/$plugin_os/$plugin_arch/" ``` -On Linux, the extension is `.so`. The plugin must be built for the operating -system and architecture where CLIProxyAPI runs. For a Podman deployment on an -Apple Silicon Mac, build the Linux ARM64 plugin inside the Podman VM: +The plugin must be built for the operating system and architecture where +CLIProxyAPI runs, not necessarily where the source directory is located. + +For a Podman deployment on an Apple Silicon Mac, CLIProxyAPI runs inside a +Linux ARM64 virtual machine. Build the Linux ARM64 plugin with Podman: ```bash +mkdir -p plugins/linux/arm64 + podman run --rm \ -v "$PWD:/src" \ -w /src/examples/plugin/kiro/go \ @@ -32,14 +61,17 @@ podman run --rm \ sh -c 'CGO_ENABLED=1 go build -buildmode=c-shared -o /src/plugins/linux/arm64/kiro-go.so . && rm -f /src/plugins/linux/arm64/kiro-go.h' ``` -The image name uses an OCI registry reference; the command is executed by -Podman and does not require a Docker daemon. +The image name is an OCI registry reference. This command uses Podman and does +not require a Docker daemon. -## Configure the plugin +## 2. Configure CLIProxyAPI -Dynamic plugins are disabled by default. Add this to `config.yaml`: +Enable the plugin and add a client-facing proxy key in `config.yaml`: ```yaml +api-keys: + - "local-dev-key" + plugins: enabled: true dir: "plugins" @@ -51,9 +83,8 @@ plugins: - "*" ``` -The plugin discovers models separately for every Kiro account. Use `"*"` to -expose every model returned for that account, or list specific model IDs to use -the configured list as an allow-list: +`"*"` exposes every model discovered for each Kiro account. To limit the +models, replace it with exact IDs returned for your account: ```yaml models: @@ -61,14 +92,31 @@ the configured list as an allow-list: - "claude-haiku-4.5" ``` -Successful results are cached for five minutes. If a refresh fails, the plugin -keeps the last successful result. On a cold-start failure it falls back to the -configured models; when `"*"` is configured, the cold-start fallback is Claude +Model discovery results are cached for five minutes. If a refresh fails, the +plugin keeps the last successful result. On a cold-start failure it falls back +to the configured models; when `"*"` is configured, the fallback is Claude Sonnet 4.5 and Claude Haiku 4.5. -## Add an auth record +## 3. Add the Kiro credential -Create `auths/kiro-pro.json`: +Create `.env` in the directory from which CLIProxyAPI starts: + +```dotenv +KIRO_API_KEY=replace-with-your-kiro-key +``` + +CLIProxyAPI automatically loads that file at startup. `.env` is ignored by +this repository and should not be committed. + +If CLIProxyAPI runs in Podman, the host's shell variables and `.env` file are +not automatically available inside the container. Pass the variable with +`--env-file .env`, `--env KIRO_API_KEY`, or a Podman secret in the container's +normal launch configuration. + +Create the auth record in CLIProxyAPI's configured `auth-dir`. For a native +process, the default path is `~/.cli-proxy-api/kiro-pro.json`. For the +repository's Podman bind mount, create `auths/kiro-pro.json` on the host so it +appears under `/root/.cli-proxy-api` in the container: ```json { @@ -79,24 +127,174 @@ Create `auths/kiro-pro.json`: } ``` -Set restrictive permissions: +Protect the auth record and restart CLIProxyAPI. Use the path for your launch +method: ```bash +chmod 600 ~/.cli-proxy-api/kiro-pro.json +# Podman bind-mount path: chmod 600 auths/kiro-pro.json ``` -Pass `KIRO_API_KEY` to the CLIProxyAPI process or Podman container. The plugin -resolves the environment variable only when it sends a request. The resolved -key is not written back to the auth JSON. +`api_key_env` stores only the environment-variable name. The plugin resolves +the value when it sends a request and never writes the resolved key back to the +auth JSON. An inline `api_key` field is also accepted, but an environment +variable or Podman secret is preferred. + +## 4. Verify the server first + +Set the client-facing proxy key in your current shell: + +```bash +export CLIPROXY_API_KEY='local-dev-key' +``` + +Check the server and list the models available through it: + +```bash +curl -fsS http://127.0.0.1:8317/healthz + +curl -fsS http://127.0.0.1:8317/v1/models \ + -H "Authorization: Bearer $CLIPROXY_API_KEY" +``` -For a temporary local test: +Use an exact Kiro model ID from this response in the client examples below. +If `jq` is installed, this prints only Kiro-owned model IDs: ```bash -export KIRO_API_KEY='replace-me' +curl -fsS http://127.0.0.1:8317/v1/models \ + -H "Authorization: Bearer $CLIPROXY_API_KEY" \ + | jq -r '.data[] | select(.owned_by == "kiro") | .id' ``` -An `api_key` field is also accepted for installations that already protect the -auth directory, but `api_key_env` or a Podman secret is preferred. +## Use with Claude Code + +Claude Code speaks the plugin's native Anthropic Messages protocol, so this is +the shortest path. + +### Configure the session + +Run these commands in the shell that will start Claude Code: + +```bash +export ANTHROPIC_BASE_URL='http://127.0.0.1:8317' +export ANTHROPIC_AUTH_TOKEN="$CLIPROXY_API_KEY" +export ANTHROPIC_MODEL='claude-sonnet-4.5' +export ANTHROPIC_DEFAULT_HAIKU_MODEL='claude-haiku-4.5' +export CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1 +``` + +Important details: + +- Do not add `/v1` to `ANTHROPIC_BASE_URL`; Claude Code adds + `/v1/messages` itself. +- `ANTHROPIC_AUTH_TOKEN` is the CLIProxyAPI client key, not the Kiro key. +- Set `ANTHROPIC_MODEL` and `ANTHROPIC_DEFAULT_HAIKU_MODEL` only to models + returned by `/v1/models` for the account. +- Gateway model discovery requires Claude Code 2.1.129 or later. It adds the + proxy's models to the `/model` picker. + +For persistent configuration, place the same variables in the `env` object of +`~/.claude/settings.json` or a gitignored `.claude/settings.local.json`. Do not +put credentials in a committed `.claude/settings.json`. + +See Anthropic's +[gateway connection guide](https://code.claude.com/docs/en/llm-gateway-connect) +for Claude Code's configuration precedence. + +### Run a test + +Start an interactive session: + +```bash +claude --model "$ANTHROPIC_MODEL" +``` + +Or run one non-interactive request: + +```bash +claude -p --model "$ANTHROPIC_MODEL" \ + "Find the main entry point and name its file." +``` + +Inside Claude Code, run `/status`. `Anthropic base URL` should show +`http://127.0.0.1:8317`, and the credential source should be +`ANTHROPIC_AUTH_TOKEN`. CLIProxyAPI logs should show `/v1/messages`, a Kiro auth +record, and an upstream request to `https://runtime..kiro.dev/`. + +## Use with Codex CLI + +Codex CLI sends Responses API requests, so use a custom model provider with +`wire_api = "responses"`. + +### Configure a Codex profile + +Create `~/.codex/kiro.config.toml`: + +```toml +model = "claude-sonnet-4.5" +model_provider = "cliproxy_kiro" + +[model_providers.cliproxy_kiro] +name = "CLIProxyAPI Kiro" +base_url = "http://127.0.0.1:8317/v1" +wire_api = "responses" +env_key = "CLIPROXY_API_KEY" +``` + +Keep `CLIPROXY_API_KEY` exported in the shell that starts Codex: + +```bash +export CLIPROXY_API_KEY='local-dev-key' +``` + +Use a user profile rather than a project `.codex/config.toml`. Current Codex +versions ignore provider and authentication settings in project configuration +files. Also do not set `requires_openai_auth = true`: that would make Codex use +OpenAI authentication and ignore `env_key`. + +See OpenAI's +[custom model provider documentation](https://developers.openai.com/codex/config-advanced/#custom-model-providers) +for the full provider configuration reference. + +### Run a test + +Start an interactive session: + +```bash +codex --profile kiro +``` + +Or run one read-only, non-interactive request: + +```bash +codex exec --profile kiro \ + "Find the main entry point and name its file." +``` + +Override the profile's model for one session with an exact model returned by +`/v1/models`: + +```bash +codex --profile kiro --model claude-haiku-4.5 +``` + +CLIProxyAPI logs should show `/v1/responses`, provider `kiro`, and an upstream +request to `https://runtime..kiro.dev/`. + +## Troubleshooting + +| Symptom | What to check | +| --- | --- | +| `connection refused` or `/healthz` fails | Confirm CLIProxyAPI is running on port 8317. With Podman, publish the port and check `podman ps`. | +| The plugin does not load | Confirm `plugins.enabled` and `kiro-go.enabled` are true, and that the plugin binary matches the server's OS and architecture. Check startup logs for `kiro-go`. | +| `401` from CLIProxyAPI | The client must send a value listed under `api-keys`. Check `CLIPROXY_API_KEY`, `ANTHROPIC_AUTH_TOKEN`, or the Codex provider's `env_key`. | +| Kiro returns `401` or `403` | Check that `KIRO_API_KEY` is visible inside the CLIProxyAPI process or Podman container and that the auth record uses the correct region. | +| `auth_unavailable: no auth available` | Confirm the model appears in `/v1/models`, the auth JSON has `"type":"kiro"`, and the plugin model allow-list includes the exact model ID. Restart after changing credentials. | +| Claude Code opens its normal login or uses a subscription | Start it from the shell containing the gateway variables, then check `/status`. `ANTHROPIC_AUTH_TOKEN` takes precedence over a saved login. | +| Models do not appear in Claude Code's `/model` picker | Check `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1`, update Claude Code if needed, and ensure `CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC` is not set because it disables gateway discovery. | +| Codex uses OpenAI auth or the wrong endpoint | Put the provider in `~/.codex/kiro.config.toml`, select it with `--profile kiro`, keep `requires_openai_auth` unset, and verify the base URL ends in `/v1`. | +| Claude Code works but Codex formatting or usage is incomplete | The build is missing the OpenAI-to-Claude translator compatibility work tracked in issue #4815. | ## Request flow @@ -114,25 +312,6 @@ auth directory, but `api_key_env` or a Podman secret is preferred. 8. It returns Claude JSON or emits Claude SSE. CLIProxyAPI translates that back to the original client protocol when necessary. -## Initial test cases - -Run these after providing a real key: - -1. List models and confirm the account's discovered models appear. If an - explicit allow-list is configured, confirm only its intersection appears. -2. Send a non-streaming `/v1/messages` text request. -3. Send the same request with `stream: true` and verify Anthropic SSE ordering. -4. After issue #4815 is implemented, call `/v1/chat/completions` to verify - host-side OpenAI translation. -5. After issue #4815 is implemented, call `/v1/responses` to verify Responses - API translation. -6. Exercise one client tool call and return its `tool_result`. -7. Send a small base64 image if the selected account model supports images. -8. Use an invalid key and verify a 401/403 does not leak the key. -9. Use a model outside the account entitlement and verify the upstream error is - surfaced without disabling unrelated credentials. -10. Cancel a streaming request and verify the upstream stream closes. - ## Limitations - The Kiro subscription API key is documented, but the raw runtime protocol is @@ -148,6 +327,8 @@ Run these after providing a real key: per-model credit multiplier. - `/v1/messages/count_tokens` returns an explicitly marked local estimate. - Web search server tools are not forwarded in the initial implementation. +- Assistant-prefill requests are rejected because Kiro's current message must + be a user message. - Kiro client compatibility versions are configurable because upstream header expectations may change. From c1af02ed9c18195ec05be407b708ebe4aaa61103 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 19:41:04 +0800 Subject: [PATCH 06/20] fix(plugin): address follow-up review feedback --- examples/plugin/Makefile | 6 ++ examples/plugin/kiro/go/kiro_test.go | 50 ++++++++++++++++ examples/plugin/kiro/go/translate.go | 23 ++++---- internal/pluginhost/adapters_executors.go | 71 +++++++++++++++++++++-- internal/pluginhost/adapters_test.go | 33 +++++++++++ 5 files changed, 168 insertions(+), 15 deletions(-) diff --git a/examples/plugin/Makefile b/examples/plugin/Makefile index 201066a3803..20fc6b3e602 100644 --- a/examples/plugin/Makefile +++ b/examples/plugin/Makefile @@ -38,6 +38,12 @@ $(BIN_DIR): $(BUILD_DIR): mkdir -p $(BUILD_DIR) +KIRO_GO_SOURCES := $(wildcard kiro/go/*.go) + +$(BIN_DIR)/kiro-go.$(PLUGIN_EXT): $(KIRO_GO_SOURCES) kiro/go/go.mod kiro/go/go.sum | $(BIN_DIR) + cd kiro/go && go build -buildmode=c-shared -o $(abspath $@) . + rm -f $(BIN_DIR)/kiro-go.h + $(BIN_DIR)/%-go.$(PLUGIN_EXT): %/go/main.go %/go/go.mod | $(BIN_DIR) cd $*/go && go build -buildmode=c-shared -o $(abspath $@) . rm -f $(BIN_DIR)/$*-go.h diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 428b693739a..52da33eac1a 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -251,6 +251,56 @@ func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { } } +func TestTruncatePayloadEvictsOldestConversationTurn(t *testing.T) { + payload := &kiroPayload{} + payload.ConversationState.History = []kiroHistoryMessage{ + {UserInputMessage: &kiroUserInputMessage{Content: strings.Repeat("x", maxKiroPayloadBytes), ModelID: "model", Origin: "KIRO_CLI"}}, + {AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "old response"}}, + {UserInputMessage: &kiroUserInputMessage{Content: "new request", ModelID: "model", Origin: "KIRO_CLI"}}, + {AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "new response"}}, + } + + truncatePayload(payload, false) + + if len(payload.ConversationState.History) != 2 { + t.Fatalf("history length = %d, want newest turn only", len(payload.ConversationState.History)) + } + if got := payload.ConversationState.History[0].UserInputMessage.Content; got != "new request" { + t.Fatalf("oldest retained request = %q, want newest request", got) + } + raw, errMarshal := json.Marshal(payload) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if len(raw) > maxKiroPayloadBytes { + t.Fatalf("truncated payload size = %d, limit = %d", len(raw), maxKiroPayloadBytes) + } +} + +func TestTruncatePayloadPreservesSyntheticSystemPair(t *testing.T) { + payload := &kiroPayload{} + payload.ConversationState.History = []kiroHistoryMessage{ + {UserInputMessage: &kiroUserInputMessage{Content: "system prompt", ModelID: "model", Origin: "KIRO_CLI"}}, + {AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "I will follow these instructions."}}, + {UserInputMessage: &kiroUserInputMessage{Content: strings.Repeat("x", maxKiroPayloadBytes), ModelID: "model", Origin: "KIRO_CLI"}}, + {AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "old response"}}, + {UserInputMessage: &kiroUserInputMessage{Content: "new request", ModelID: "model", Origin: "KIRO_CLI"}}, + {AssistantResponseMessage: &kiroAssistantResponseMessage{Content: "new response"}}, + } + + truncatePayload(payload, true) + + if len(payload.ConversationState.History) != 4 { + t.Fatalf("history length = %d, want system pair plus newest turn", len(payload.ConversationState.History)) + } + if got := payload.ConversationState.History[0].UserInputMessage.Content; got != "system prompt" { + t.Fatalf("first history request = %q, want preserved system prompt", got) + } + if got := payload.ConversationState.History[2].UserInputMessage.Content; got != "new request" { + t.Fatalf("retained conversation request = %q, want newest request", got) + } +} + func TestClaudeToKiroPreservesToolResultErrorStatus(t *testing.T) { raw := []byte(`{ "model":"claude-sonnet-4-5", diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 7f39dae4ee4..242abdd0bf0 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -102,7 +102,7 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if request.MaxTokens > 0 || request.Temperature > 0 || request.TopP > 0 { payload.InferenceConfig = &kiroInferenceConfig{MaxTokens: request.MaxTokens, Temperature: request.Temperature, TopP: request.TopP} } - truncatePayload(payload) + truncatePayload(payload, systemPrompt != "") return payload, &request, nil } @@ -394,19 +394,22 @@ func formatUUID(b []byte) string { return hexValue[:8] + "-" + hexValue[8:12] + "-" + hexValue[12:16] + "-" + hexValue[16:20] + "-" + hexValue[20:32] } -func truncatePayload(payload *kiroPayload) { +func truncatePayload(payload *kiroPayload, preserveSystemPair bool) { + protected := 0 + if preserveSystemPair { + protected = 2 + } for { raw, errMarshal := json.Marshal(payload) - if errMarshal != nil || len(raw) <= maxKiroPayloadBytes || len(payload.ConversationState.History) <= 4 { + if errMarshal != nil || len(raw) <= maxKiroPayloadBytes || len(payload.ConversationState.History) <= protected { return } - removeAt := 0 - if len(payload.ConversationState.History) >= 2 && payload.ConversationState.History[0].UserInputMessage != nil && payload.ConversationState.History[1].AssistantResponseMessage != nil { - removeAt = 2 - } - if removeAt >= len(payload.ConversationState.History) { - return + removeCount := 1 + if protected+1 < len(payload.ConversationState.History) && + payload.ConversationState.History[protected].UserInputMessage != nil && + payload.ConversationState.History[protected+1].AssistantResponseMessage != nil { + removeCount = 2 } - payload.ConversationState.History = append(payload.ConversationState.History[:removeAt], payload.ConversationState.History[removeAt+1:]...) + payload.ConversationState.History = append(payload.ConversationState.History[:protected], payload.ConversationState.History[protected+removeCount:]...) } } diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index f4078810d7d..0d7fc25cba1 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -537,7 +537,7 @@ func (a *executorAdapter) translateExecutorResponse(ctx context.Context, prepare return sdktranslator.TranslateNonStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) } -func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, prepared preparedExecutorCall, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { +func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, prepared preparedExecutorCall, headers http.Header, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { if prepared.requestedFormat == "" || prepared.outputFormat == prepared.requestedFormat { return in } @@ -551,12 +551,30 @@ func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, pre go func() { defer close(out) var param any + var sseBuffer executorSSERecordBuffer + isSSE := strings.Contains(strings.ToLower(headers.Get("Content-Type")), "text/event-stream") + sendTranslated := func(payload []byte) bool { + frames := a.translateExecutorStreamPayload(ctx, prepared, payload, ¶m) + for _, frame := range frames { + if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { + return false + } + } + return true + } for { select { case <-ctx.Done(): return case chunk, ok := <-in: if !ok { + if isSSE { + for _, payload := range sseBuffer.Flush() { + if !sendTranslated(payload) { + return + } + } + } a.emitTranslatedExecutorStreamTail(ctx, prepared, out, ¶m) return } @@ -564,9 +582,12 @@ func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, pre _ = sendExecutorPluginStreamChunk(ctx, out, chunk) continue } - frames := a.translateExecutorStreamPayload(ctx, prepared, chunk.Payload, ¶m) - for _, frame := range frames { - if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { + payloads := [][]byte{chunk.Payload} + if isSSE { + payloads = sseBuffer.Push(chunk.Payload) + } + for _, payload := range payloads { + if !sendTranslated(payload) { return } } @@ -576,6 +597,46 @@ func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, pre return out } +type executorSSERecordBuffer struct { + pending []byte +} + +func (b *executorSSERecordBuffer) Push(payload []byte) [][]byte { + if b == nil || len(payload) == 0 { + return nil + } + b.pending = append(b.pending, payload...) + var records [][]byte + for { + end := executorSSERecordEnd(b.pending) + if end < 0 { + return records + } + records = append(records, bytes.Clone(b.pending[:end])) + b.pending = append(b.pending[:0], b.pending[end:]...) + } +} + +func (b *executorSSERecordBuffer) Flush() [][]byte { + if b == nil || len(bytes.TrimSpace(b.pending)) == 0 { + return nil + } + payload := bytes.Clone(b.pending) + b.pending = b.pending[:0] + return [][]byte{payload} +} + +func executorSSERecordEnd(payload []byte) int { + end := -1 + if index := bytes.Index(payload, []byte("\n\n")); index >= 0 { + end = index + 2 + } + if index := bytes.Index(payload, []byte("\r\n\r\n")); index >= 0 && (end < 0 || index+4 < end) { + end = index + 4 + } + return end +} + func (a *executorAdapter) translateExecutorStreamPayload(ctx context.Context, prepared preparedExecutorCall, payload []byte, param *any) [][]byte { originalRequest := prepared.opts.OriginalRequest if len(originalRequest) == 0 { @@ -709,7 +770,7 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth nativeChunks := usageReporter.observeStream(ctx, prepared.outputFormat, pluginResp.Chunks) return &coreexecutor.StreamResult{ Headers: cloneHeader(pluginResp.Headers), - Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, nativeChunks)), + Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Headers, nativeChunks)), }, nil } diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 74e1c5d2800..12672ddfdc5 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2861,6 +2861,39 @@ func TestExecutorStreamTranslationPayloadsExtractsDataFromSSEFrame(t *testing.T) } } +func TestExecutorSSERecordBufferRetainsSplitDataLine(t *testing.T) { + var buffer executorSSERecordBuffer + first := []byte("event: message_delta\ndata: {\"type\":") + if records := buffer.Push(first); len(records) != 0 { + t.Fatalf("first chunk produced %d records, want none", len(records)) + } + + second := []byte("\"message_delta\",\"usage\":{\"output_tokens\":4}}\n\n") + records := buffer.Push(second) + if len(records) != 1 { + t.Fatalf("second chunk produced %d records, want one", len(records)) + } + want := append(bytes.Clone(first), second...) + if !bytes.Equal(records[0], want) { + t.Fatalf("buffered record = %q, want %q", records[0], want) + } + if tail := buffer.Flush(); len(tail) != 0 { + t.Fatalf("buffer tail = %q, want empty", tail) + } +} + +func TestExecutorSSERecordBufferFlushesFinalRecordWithoutDelimiter(t *testing.T) { + var buffer executorSSERecordBuffer + payload := []byte("data: [DONE]") + if records := buffer.Push(payload); len(records) != 0 { + t.Fatalf("unterminated chunk produced %d records, want none", len(records)) + } + tail := buffer.Flush() + if len(tail) != 1 || !bytes.Equal(tail[0], payload) { + t.Fatalf("buffer tail = %q, want %q", tail, payload) + } +} + func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { var usageBuffer helps.StreamUsageBuffer payload := []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":12,\"output_tokens\":4}}\n\n") From 5cb7bdeb692251df6592c877439ab69e9afb2d35 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 20:00:09 +0800 Subject: [PATCH 07/20] fix(plugin): address additional review feedback --- examples/plugin/kiro/go/kiro_test.go | 24 +++++++++++++++++++ examples/plugin/kiro/go/protocol.go | 10 ++++---- examples/plugin/kiro/go/translate.go | 2 +- internal/pluginhost/adapters_test.go | 34 +++++++++++++++++++++++++++ internal/pluginhost/executor_usage.go | 17 ++++++++++++++ 5 files changed, 81 insertions(+), 6 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 52da33eac1a..3663c673d4a 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -1,6 +1,7 @@ package main import ( + "bytes" "encoding/binary" "encoding/json" "hash/crc32" @@ -277,6 +278,29 @@ func TestTruncatePayloadEvictsOldestConversationTurn(t *testing.T) { } } +func TestClaudeToKiroPreservesExplicitZeroSamplingValues(t *testing.T) { + payload, _, errTranslate := claudeToKiro([]byte(`{ + "model":"claude-sonnet-4-5", + "messages":[{"role":"user","content":"hello"}], + "temperature":0, + "top_p":0 + }`), "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + if payload.InferenceConfig == nil || payload.InferenceConfig.Temperature == nil || payload.InferenceConfig.TopP == nil { + t.Fatalf("inference config = %#v, want explicit sampling values", payload.InferenceConfig) + } + + raw, errMarshal := json.Marshal(payload) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if !bytes.Contains(raw, []byte(`"temperature":0`)) || !bytes.Contains(raw, []byte(`"topP":0`)) { + t.Fatalf("Kiro payload = %s, want explicit zero sampling values", raw) + } +} + func TestTruncatePayloadPreservesSyntheticSystemPair(t *testing.T) { payload := &kiroPayload{} payload.ConversationState.History = []kiroHistoryMessage{ diff --git a/examples/plugin/kiro/go/protocol.go b/examples/plugin/kiro/go/protocol.go index 479f943919e..f738247ba5d 100644 --- a/examples/plugin/kiro/go/protocol.go +++ b/examples/plugin/kiro/go/protocol.go @@ -83,17 +83,17 @@ type kiroToolUse struct { } type kiroInferenceConfig struct { - MaxTokens int `json:"maxTokens,omitempty"` - Temperature float64 `json:"temperature,omitempty"` - TopP float64 `json:"topP,omitempty"` + MaxTokens int `json:"maxTokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"topP,omitempty"` } type claudeRequest struct { Model string `json:"model"` Messages []claudeMessage `json:"messages"` MaxTokens int `json:"max_tokens"` - Temperature float64 `json:"temperature,omitempty"` - TopP float64 `json:"top_p,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` Stream bool `json:"stream,omitempty"` System any `json:"system,omitempty"` Thinking *struct { diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 242abdd0bf0..70c2e387adc 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -99,7 +99,7 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque payload.ConversationState.CurrentMessage.UserInputMessage = current payload.ConversationState.History = trimLeadingAssistant(history) payload.EstimatedInputTokens = estimateClaudeRequestInputTokens(&request) - if request.MaxTokens > 0 || request.Temperature > 0 || request.TopP > 0 { + if request.MaxTokens > 0 || request.Temperature != nil || request.TopP != nil { payload.InferenceConfig = &kiroInferenceConfig{MaxTokens: request.MaxTokens, Temperature: request.Temperature, TopP: request.TopP} } truncatePayload(payload, systemPrompt != "") diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 12672ddfdc5..c7325be11bb 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2924,6 +2924,40 @@ func TestPluginExecutorUsageBuffersSplitClaudeStreamCounts(t *testing.T) { } } +func TestPluginExecutorUsageParsesGeminiCounts(t *testing.T) { + payload := []byte(`{"usageMetadata":{"promptTokenCount":12,"candidatesTokenCount":4,"totalTokenCount":16}}`) + detail, ok := pluginExecutorNonStreamUsage(sdktranslator.FormatGemini, payload) + if !ok { + t.Fatal("non-stream Gemini usage was not observed") + } + if detail.InputTokens != 12 || detail.OutputTokens != 4 || detail.TotalTokens != 16 { + t.Fatalf("non-stream Gemini usage = %#v, want input=12 output=4 total=16", detail) + } + + var usageBuffer helps.StreamUsageBuffer + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatGemini, append([]byte("data: "), payload...)) + detail, ok = usageBuffer.Detail() + if !ok { + t.Fatal("stream Gemini usage was not observed") + } + if detail.InputTokens != 12 || detail.OutputTokens != 4 || detail.TotalTokens != 16 { + t.Fatalf("stream Gemini usage = %#v, want input=12 output=4 total=16", detail) + } +} + +func TestPluginExecutorUsageDoesNotBufferUnsupportedStreamFormat(t *testing.T) { + var usageBuffer helps.StreamUsageBuffer + pending := observePluginExecutorStreamChunk( + &usageBuffer, + sdktranslator.Format("plugin-custom-output"), + nil, + []byte("custom stream payload without a newline"), + ) + if len(pending) != 0 { + t.Fatalf("pending unsupported stream payload = %q, want empty", pending) + } +} + func TestExecutorAdapterPublishesNativeUsageForBilling(t *testing.T) { const model = "plugin-usage-billing-model" records := make(chan coreusage.Record, 1) diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go index 98840c3bf2f..bb8fa29f5db 100644 --- a/internal/pluginhost/executor_usage.go +++ b/internal/pluginhost/executor_usage.go @@ -76,6 +76,9 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans } func observePluginExecutorStreamChunk(buffer *helps.StreamUsageBuffer, format sdktranslator.Format, pending, payload []byte) []byte { + if !pluginExecutorSupportsStreamUsage(format) { + return nil + } pending = append(pending, payload...) lastNewline := bytes.LastIndexByte(pending, '\n') if lastNewline < 0 { @@ -91,11 +94,22 @@ func pluginExecutorNonStreamUsage(format sdktranslator.Format, payload []byte) ( return helps.ParseClaudeUsage(payload), bytes.Contains(payload, []byte(`"usage"`)) case sdktranslator.FormatOpenAI, sdktranslator.FormatOpenAIResponse: return helps.ParseOpenAIUsage(payload), bytes.Contains(payload, []byte(`"usage"`)) + case sdktranslator.FormatGemini: + return helps.ParseGeminiUsage(payload), bytes.Contains(payload, []byte(`"usageMetadata"`)) || bytes.Contains(payload, []byte(`"usage_metadata"`)) default: return coreusage.Detail{}, false } } +func pluginExecutorSupportsStreamUsage(format sdktranslator.Format) bool { + switch format { + case sdktranslator.FormatClaude, sdktranslator.FormatOpenAI, sdktranslator.FormatOpenAIResponse, sdktranslator.FormatGemini: + return true + default: + return false + } +} + func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sdktranslator.Format, payload []byte) { for _, line := range bytes.Split(payload, []byte("\n")) { switch format { @@ -109,6 +123,9 @@ func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sd jsonPayload := helps.JSONPayload(line) detail, ok := helps.ParseCodexUsage(jsonPayload) buffer.Observe(detail, ok) + case sdktranslator.FormatGemini: + detail, ok := helps.ParseGeminiStreamUsage(line) + buffer.Observe(detail, ok) } } } From 27913efb761c032a6bc6e0de22ba344ac5d66adb Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 20:30:00 +0800 Subject: [PATCH 08/20] fix(plugin): address latest review feedback --- examples/plugin/kiro/go/kiro_test.go | 41 +++++++++++++++++++++++ examples/plugin/kiro/go/protocol.go | 3 +- examples/plugin/kiro/go/translate.go | 22 +++++++++--- internal/pluginhost/adapters_executors.go | 30 ++++++++++++++--- internal/pluginhost/adapters_test.go | 12 +++++++ 5 files changed, 98 insertions(+), 10 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 3663c673d4a..aa55c2451ae 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -301,6 +301,47 @@ func TestClaudeToKiroPreservesExplicitZeroSamplingValues(t *testing.T) { } } +func TestClaudeToKiroPreservesThinkingBudget(t *testing.T) { + payload, _, errTranslate := claudeToKiro([]byte(`{ + "model":"claude-sonnet-4-5", + "messages":[{"role":"user","content":"hello"}], + "thinking":{"type":"enabled","budget_tokens":1024} + }`), "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + if len(payload.ConversationState.History) < 1 || payload.ConversationState.History[0].UserInputMessage == nil { + t.Fatalf("history = %#v, want thinking system prompt", payload.ConversationState.History) + } + content := payload.ConversationState.History[0].UserInputMessage.Content + if !strings.Contains(content, "1024") { + t.Fatalf("thinking system prompt = %q, want requested budget", content) + } +} + +func TestClaudeToKiroRejectsPayloadAboveSizeLimit(t *testing.T) { + tests := []struct { + name string + raw []byte + }{ + { + name: "oversized current message", + raw: []byte(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"` + strings.Repeat("x", maxKiroPayloadBytes) + `"}]}`), + }, + { + name: "oversized protected system prompt", + raw: []byte(`{"model":"claude-sonnet-4-5","system":"` + strings.Repeat("x", maxKiroPayloadBytes) + `","messages":[{"role":"user","content":"hello"}]}`), + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if _, _, errTranslate := claudeToKiro(test.raw, ""); errTranslate == nil || !strings.Contains(errTranslate.Error(), "limit is") { + t.Fatalf("claudeToKiro() error = %v, want local size-limit error", errTranslate) + } + }) + } +} + func TestTruncatePayloadPreservesSyntheticSystemPair(t *testing.T) { payload := &kiroPayload{} payload.ConversationState.History = []kiroHistoryMessage{ diff --git a/examples/plugin/kiro/go/protocol.go b/examples/plugin/kiro/go/protocol.go index f738247ba5d..2d6d6f0d25d 100644 --- a/examples/plugin/kiro/go/protocol.go +++ b/examples/plugin/kiro/go/protocol.go @@ -97,7 +97,8 @@ type claudeRequest struct { Stream bool `json:"stream,omitempty"` System any `json:"system,omitempty"` Thinking *struct { - Type string `json:"type,omitempty"` + Type string `json:"type,omitempty"` + BudgetTokens int `json:"budget_tokens,omitempty"` } `json:"thinking,omitempty"` Tools []claudeTool `json:"tools,omitempty"` } diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 70c2e387adc..81ff1b8ddac 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -41,7 +41,11 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if request.Thinking != nil { kind := strings.ToLower(strings.TrimSpace(request.Thinking.Type)) if kind == "enabled" || kind == "adaptive" { - systemPrompt = strings.TrimSpace("enabled\n200000\n\n" + systemPrompt) + thinkingBudget := 200000 + if request.Thinking.BudgetTokens > 0 { + thinkingBudget = request.Thinking.BudgetTokens + } + systemPrompt = strings.TrimSpace(fmt.Sprintf("enabled\n%d\n\n%s", thinkingBudget, systemPrompt)) } } @@ -102,7 +106,9 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if request.MaxTokens > 0 || request.Temperature != nil || request.TopP != nil { payload.InferenceConfig = &kiroInferenceConfig{MaxTokens: request.MaxTokens, Temperature: request.Temperature, TopP: request.TopP} } - truncatePayload(payload, systemPrompt != "") + if errTruncate := truncatePayload(payload, systemPrompt != ""); errTruncate != nil { + return nil, nil, errTruncate + } return payload, &request, nil } @@ -394,15 +400,21 @@ func formatUUID(b []byte) string { return hexValue[:8] + "-" + hexValue[8:12] + "-" + hexValue[12:16] + "-" + hexValue[16:20] + "-" + hexValue[20:32] } -func truncatePayload(payload *kiroPayload, preserveSystemPair bool) { +func truncatePayload(payload *kiroPayload, preserveSystemPair bool) error { protected := 0 if preserveSystemPair { protected = 2 } for { raw, errMarshal := json.Marshal(payload) - if errMarshal != nil || len(raw) <= maxKiroPayloadBytes || len(payload.ConversationState.History) <= protected { - return + if errMarshal != nil { + return fmt.Errorf("encode Kiro payload: %w", errMarshal) + } + if len(raw) <= maxKiroPayloadBytes { + return nil + } + if len(payload.ConversationState.History) <= protected { + return fmt.Errorf("Kiro request payload is %d bytes after truncating removable history; limit is %d bytes", len(raw), maxKiroPayloadBytes) } removeCount := 1 if protected+1 < len(payload.ConversationState.History) && diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index 0d7fc25cba1..87ab054bd35 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -658,17 +658,39 @@ func executorStreamTranslationPayloads(payload []byte) [][]byte { if !bytes.Contains(payload, []byte("\n")) { return [][]byte{payload} } - var dataLines [][]byte + var translationPayloads [][]byte + var dataValues [][]byte + sawData := false + flushData := func() { + if len(dataValues) == 0 { + return + } + joined := bytes.Join(dataValues, []byte("\n")) + frame := []byte("data:") + if len(joined) > 0 { + frame = append(frame, ' ') + frame = append(frame, joined...) + } + translationPayloads = append(translationPayloads, frame) + dataValues = nil + } for _, line := range bytes.Split(payload, []byte("\n")) { trimmed := bytes.TrimSpace(line) + if len(trimmed) == 0 { + flushData() + continue + } if bytes.HasPrefix(trimmed, []byte("data:")) { - dataLines = append(dataLines, bytes.Clone(trimmed)) + sawData = true + value := bytes.TrimPrefix(trimmed[len("data:"):], []byte(" ")) + dataValues = append(dataValues, bytes.Clone(value)) } } - if len(dataLines) == 0 { + flushData() + if !sawData { return [][]byte{payload} } - return dataLines + return translationPayloads } func executorStreamTranslationFellBack(prepared preparedExecutorCall, payload []byte, frames [][]byte) bool { diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index c7325be11bb..4aab9cf0759 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2861,6 +2861,18 @@ func TestExecutorStreamTranslationPayloadsExtractsDataFromSSEFrame(t *testing.T) } } +func TestExecutorStreamTranslationPayloadsReassemblesMultilineData(t *testing.T) { + payload := []byte("event: content_block_delta\ndata: {\ndata: \"type\": \"content_block_delta\",\ndata: \"delta\": {\"type\": \"text_delta\", \"text\": \"hello\"}\ndata: }\n\n") + frames := executorStreamTranslationPayloads(payload) + if len(frames) != 1 { + t.Fatalf("translation payload count = %d, want 1", len(frames)) + } + want := []byte("data: {\n \"type\": \"content_block_delta\",\n \"delta\": {\"type\": \"text_delta\", \"text\": \"hello\"}\n}") + if !bytes.Equal(frames[0], want) { + t.Fatalf("translation payload = %q, want %q", frames[0], want) + } +} + func TestExecutorSSERecordBufferRetainsSplitDataLine(t *testing.T) { var buffer executorSSERecordBuffer first := []byte("event: message_delta\ndata: {\"type\":") From eaa59db44d99c17fe6ed214caddb1e6e45390492 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 21:21:17 +0800 Subject: [PATCH 09/20] fix(plugin): close remaining review gaps --- examples/plugin/kiro/go/kiro_test.go | 34 ++++++++++++ examples/plugin/kiro/go/translate.go | 23 +++++--- internal/pluginhost/adapters_executors.go | 2 +- internal/pluginhost/adapters_test.go | 26 +++++++++ internal/pluginhost/executor_usage.go | 64 ++++++++++++++++++++--- 5 files changed, 134 insertions(+), 15 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index aa55c2451ae..f2e87aad220 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -252,6 +252,40 @@ func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { } } +func TestClaudeToKiroDisambiguatesHistoricalToolNames(t *testing.T) { + payload, _, errTranslate := claudeToKiro([]byte(`{ + "model":"claude-sonnet-4-5", + "tools":[ + {"name":"foo","input_schema":{"type":"object"}}, + {"name":"foo_","input_schema":{"type":"object"}} + ], + "messages":[ + {"role":"user","content":"Use the second tool"}, + {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"foo_","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"done"}]} + ] + }`), "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + context := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext + if context == nil || len(context.Tools) != 2 { + t.Fatalf("current context = %#v, want two tools", context) + } + if got := context.Tools[1].ToolSpecification.Name; got != "foo_2" { + t.Fatalf("second declared tool name = %q, want foo_2", got) + } + if len(payload.ConversationState.History) != 2 || len(payload.ConversationState.History[1].AssistantResponseMessage.ToolUses) != 1 { + t.Fatalf("history = %#v, want assistant tool call", payload.ConversationState.History) + } + if got := payload.ConversationState.History[1].AssistantResponseMessage.ToolUses[0].Name; got != "foo_2" { + t.Fatalf("historical tool name = %q, want foo_2", got) + } + if got := payload.ToolNameMap["foo_2"]; got != "foo_" { + t.Fatalf("response tool name mapping = %q, want foo_", got) + } +} + func TestTruncatePayloadEvictsOldestConversationTurn(t *testing.T) { payload := &kiroPayload{} payload.ConversationState.History = []kiroHistoryMessage{ diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 81ff1b8ddac..5b192df4d0a 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -63,6 +63,7 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque ) } + tools, nameMap, requestToolNames := convertTools(request.Tools) var currentText string var currentImages []kiroImage var currentToolResults []kiroToolResult @@ -81,12 +82,11 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque } history = append(history, kiroHistoryMessage{UserInputMessage: userMessage}) case "assistant": - text, tools := extractAssistantContent(message.Content) - history = append(history, kiroHistoryMessage{AssistantResponseMessage: &kiroAssistantResponseMessage{Content: text, ToolUses: tools}}) + text, assistantTools := extractAssistantContent(message.Content, requestToolNames) + history = append(history, kiroHistoryMessage{AssistantResponseMessage: &kiroAssistantResponseMessage{Content: text, ToolUses: assistantTools}}) } } - tools, nameMap := convertTools(request.Tools) payload.ToolNameMap = nameMap if len(currentToolResults) > 0 { currentText = joinNonEmpty(currentText, readableToolResults(currentToolResults)) @@ -177,7 +177,7 @@ func extractUserContent(content any) (string, []kiroImage, []kiroToolResult) { return strings.Join(texts, ""), images, results } -func extractAssistantContent(content any) (string, []kiroToolUse) { +func extractAssistantContent(content any, toolNames map[string]string) (string, []kiroToolUse) { if text, ok := content.(string); ok { return text, nil } @@ -193,11 +193,15 @@ func extractAssistantContent(content any) (string, []kiroToolUse) { case "tool_use": id, _ := block["id"].(string) name, _ := block["name"].(string) + kiroName := toolNames[name] + if kiroName == "" { + kiroName = sanitizeToolName(name) + } input, _ := block["input"].(map[string]any) if input == nil { input = map[string]any{} } - tools = append(tools, kiroToolUse{ToolUseID: id, Name: sanitizeToolName(name), Input: input}) + tools = append(tools, kiroToolUse{ToolUseID: id, Name: kiroName, Input: input}) } } return strings.Join(texts, ""), tools @@ -257,15 +261,17 @@ func extractToolResultContent(content any) (string, []kiroImage) { return strings.Join(texts, ""), images } -func convertTools(tools []claudeTool) ([]kiroToolWrapper, map[string]string) { +func convertTools(tools []claudeTool) ([]kiroToolWrapper, map[string]string, map[string]string) { out := make([]kiroToolWrapper, 0, len(tools)) nameMap := make(map[string]string) + requestNameMap := make(map[string]string) used := make(map[string]int) for _, tool := range tools { if strings.HasPrefix(strings.ToLower(tool.Type), "web_search") { continue } name := uniqueToolName(sanitizeToolName(tool.Name), used) + requestNameMap[tool.Name] = name if name != tool.Name { nameMap[name] = tool.Name } @@ -290,7 +296,10 @@ func convertTools(tools []claudeTool) ([]kiroToolWrapper, map[string]string) { if len(nameMap) == 0 { nameMap = nil } - return out, nameMap + if len(requestNameMap) == 0 { + requestNameMap = nil + } + return out, nameMap, requestNameMap } func sanitizeToolName(name string) string { diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index 87ab054bd35..d5556d1b821 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -789,7 +789,7 @@ func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth if errExecuteStream != nil { return nil, errExecuteStream } - nativeChunks := usageReporter.observeStream(ctx, prepared.outputFormat, pluginResp.Chunks) + nativeChunks := usageReporter.observeStream(ctx, prepared.outputFormat, pluginResp.Headers, pluginResp.Chunks) return &coreexecutor.StreamResult{ Headers: cloneHeader(pluginResp.Headers), Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Headers, nativeChunks)), diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 4aab9cf0759..464436e65ba 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2919,6 +2919,32 @@ func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { } } +func TestPluginExecutorUsageMergesClaudeStreamCounts(t *testing.T) { + var usageBuffer helps.StreamUsageBuffer + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, []byte("event: message_start\ndata: {\"type\":\"message_start\",\"usage\":{\"input_tokens\":12,\"output_tokens\":1,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":2}}\n\n")) + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":4}}\n\n")) + detail, ok := usageBuffer.Detail() + if !ok { + t.Fatal("stream usage was not observed") + } + if detail.InputTokens != 12 || detail.OutputTokens != 4 || detail.CacheReadTokens != 3 || detail.CacheCreationTokens != 2 || detail.TotalTokens != 21 { + t.Fatalf("merged Claude stream usage = %#v, want input=12 output=4 cache-read=3 cache-create=2 total=21", detail) + } +} + +func TestPluginExecutorUsageReassemblesMultilineSSEData(t *testing.T) { + var usageBuffer helps.StreamUsageBuffer + payload := []byte("event: message_delta\ndata: {\ndata: \"type\": \"message_delta\",\ndata: \"usage\": {\"output_tokens\": 4}\ndata: }\n\n") + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, payload) + detail, ok := usageBuffer.Detail() + if !ok { + t.Fatal("multi-line SSE usage was not observed") + } + if detail.OutputTokens != 4 || detail.TotalTokens != 4 { + t.Fatalf("multi-line SSE usage = %#v, want output=4 total=4", detail) + } +} + func TestPluginExecutorUsageBuffersSplitClaudeStreamCounts(t *testing.T) { var usageBuffer helps.StreamUsageBuffer var pending []byte diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go index bb8fa29f5db..3453b762835 100644 --- a/internal/pluginhost/executor_usage.go +++ b/internal/pluginhost/executor_usage.go @@ -3,6 +3,8 @@ package pluginhost import ( "bytes" "context" + "net/http" + "strings" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -38,7 +40,7 @@ func (u *pluginExecutorUsage) publishNonStream(ctx context.Context, format sdktr u.reporter.EnsurePublished(ctx) } -func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktranslator.Format, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { +func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktranslator.Format, headers http.Header, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { if u == nil || u.reporter == nil || in == nil { return in } @@ -46,11 +48,17 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans go func() { defer close(out) var usageBuffer helps.StreamUsageBuffer + var sseBuffer executorSSERecordBuffer var pendingUsagePayload []byte var terminalErr error + isSSE := strings.Contains(strings.ToLower(headers.Get("Content-Type")), "text/event-stream") for chunk := range in { if chunk.Err != nil { terminalErr = chunk.Err + } else if isSSE { + for _, record := range sseBuffer.Push(chunk.Payload) { + observePluginExecutorStreamUsage(&usageBuffer, format, record) + } } else { pendingUsagePayload = observePluginExecutorStreamChunk(&usageBuffer, format, pendingUsagePayload, chunk.Payload) } @@ -65,7 +73,11 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans u.reporter.PublishFailure(ctx, terminalErr) return } - if len(pendingUsagePayload) > 0 { + if isSSE { + for _, record := range sseBuffer.Flush() { + observePluginExecutorStreamUsage(&usageBuffer, format, record) + } + } else if len(pendingUsagePayload) > 0 { observePluginExecutorStreamUsage(&usageBuffer, format, pendingUsagePayload) } if !usageBuffer.Publish(ctx, u.reporter) { @@ -111,21 +123,59 @@ func pluginExecutorSupportsStreamUsage(format sdktranslator.Format) bool { } func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sdktranslator.Format, payload []byte) { - for _, line := range bytes.Split(payload, []byte("\n")) { + usagePayloads := executorStreamTranslationPayloads(payload) + if len(usagePayloads) == 1 && bytes.Equal(usagePayloads[0], payload) { + usagePayloads = bytes.Split(payload, []byte("\n")) + } + for _, usagePayload := range usagePayloads { switch format { case sdktranslator.FormatClaude: - detail, ok := helps.ParseClaudeStreamUsage(line) + detail, ok := helps.ParseClaudeStreamUsage(usagePayload) + if previous, exists := buffer.Detail(); ok && exists { + detail = mergeClaudePluginStreamUsage(previous, detail) + } buffer.Observe(detail, ok) case sdktranslator.FormatOpenAI: - detail, ok := helps.ParseOpenAIStreamUsage(line) + detail, ok := helps.ParseOpenAIStreamUsage(usagePayload) buffer.Observe(detail, ok) case sdktranslator.FormatOpenAIResponse: - jsonPayload := helps.JSONPayload(line) + jsonPayload := helps.JSONPayload(usagePayload) detail, ok := helps.ParseCodexUsage(jsonPayload) buffer.Observe(detail, ok) case sdktranslator.FormatGemini: - detail, ok := helps.ParseGeminiStreamUsage(line) + detail, ok := helps.ParseGeminiStreamUsage(usagePayload) buffer.Observe(detail, ok) } } } + +func mergeClaudePluginStreamUsage(previous, current coreusage.Detail) coreusage.Detail { + current.InputTokens = max(previous.InputTokens, current.InputTokens) + current.OutputTokens = max(previous.OutputTokens, current.OutputTokens) + current.ReasoningTokens = max(previous.ReasoningTokens, current.ReasoningTokens) + current.CacheReadTokens = max(previous.CacheReadTokens, current.CacheReadTokens) + current.CacheCreationTokens = max(previous.CacheCreationTokens, current.CacheCreationTokens) + current.CachedTokens = current.CacheReadTokens + if current.CachedTokens == 0 { + current.CachedTokens = current.CacheCreationTokens + } + if current.ResponseServiceTier == "" { + current.ResponseServiceTier = previous.ResponseServiceTier + } + nonReasoningOutput := current.OutputTokens + if current.ReasoningTokens > 0 && current.ReasoningTokens <= current.OutputTokens { + nonReasoningOutput -= current.ReasoningTokens + } else if current.ReasoningTokens > current.OutputTokens { + nonReasoningOutput = 0 + } + current.TotalTokens = current.InputTokens + current.OutputTokens + current.CacheReadTokens + current.CacheCreationTokens + current.TokenBreakdown = coreusage.NewIndependentTokenBreakdown( + current.InputTokens, + current.CacheReadTokens, + current.CacheCreationTokens, + nonReasoningOutput, + current.ReasoningTokens, + current.TotalTokens, + ) + return current +} From 0daef8ad4acf22eb8434a05b73bd75353c1395fa Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 22:43:51 +0800 Subject: [PATCH 10/20] fix(plugin): finalize Kiro review edge cases --- examples/plugin/kiro/go/kiro_test.go | 45 +++++++++++++++++++++------- examples/plugin/kiro/go/translate.go | 26 ++++++++++------ 2 files changed, 52 insertions(+), 19 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index f2e87aad220..a5c123d7918 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -252,16 +252,17 @@ func TestClaudeToKiroTransformsSystemToolsAndResults(t *testing.T) { } } -func TestClaudeToKiroDisambiguatesHistoricalToolNames(t *testing.T) { +func TestClaudeToKiroGloballyDisambiguatesHistoricalToolNames(t *testing.T) { payload, _, errTranslate := claudeToKiro([]byte(`{ "model":"claude-sonnet-4-5", "tools":[ {"name":"foo","input_schema":{"type":"object"}}, - {"name":"foo_","input_schema":{"type":"object"}} + {"name":"foo_","input_schema":{"type":"object"}}, + {"name":"foo_2","input_schema":{"type":"object"}} ], "messages":[ - {"role":"user","content":"Use the second tool"}, - {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"foo_","input":{}}]}, + {"role":"user","content":"Use the third tool"}, + {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"foo_2","input":{}}]}, {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"done"}]} ] }`), "") @@ -269,20 +270,44 @@ func TestClaudeToKiroDisambiguatesHistoricalToolNames(t *testing.T) { t.Fatalf("claudeToKiro() error = %v", errTranslate) } context := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext - if context == nil || len(context.Tools) != 2 { - t.Fatalf("current context = %#v, want two tools", context) + if context == nil || len(context.Tools) != 3 { + t.Fatalf("current context = %#v, want three tools", context) } if got := context.Tools[1].ToolSpecification.Name; got != "foo_2" { t.Fatalf("second declared tool name = %q, want foo_2", got) } + if got := context.Tools[2].ToolSpecification.Name; got != "foo_2_2" { + t.Fatalf("third declared tool name = %q, want foo_2_2", got) + } if len(payload.ConversationState.History) != 2 || len(payload.ConversationState.History[1].AssistantResponseMessage.ToolUses) != 1 { t.Fatalf("history = %#v, want assistant tool call", payload.ConversationState.History) } - if got := payload.ConversationState.History[1].AssistantResponseMessage.ToolUses[0].Name; got != "foo_2" { - t.Fatalf("historical tool name = %q, want foo_2", got) + if got := payload.ConversationState.History[1].AssistantResponseMessage.ToolUses[0].Name; got != "foo_2_2" { + t.Fatalf("historical tool name = %q, want foo_2_2", got) + } + if got := payload.ToolNameMap["foo_2_2"]; got != "foo_2" { + t.Fatalf("response tool name mapping = %q, want foo_2", got) + } +} + +func TestClaudeToKiroRecalculatesInputEstimateAfterTruncation(t *testing.T) { + payload, request, errTranslate := claudeToKiro([]byte(`{ + "model":"claude-sonnet-4-5", + "messages":[ + {"role":"user","content":"`+strings.Repeat("x", maxKiroPayloadBytes)+`"}, + {"role":"assistant","content":"old response"}, + {"role":"user","content":"new request"} + ] + }`), "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + if len(payload.ConversationState.History) != 0 { + t.Fatalf("retained history length = %d, want 0", len(payload.ConversationState.History)) } - if got := payload.ToolNameMap["foo_2"]; got != "foo_" { - t.Fatalf("response tool name mapping = %q, want foo_", got) + originalEstimate := estimateClaudeRequestInputTokens(request) + if originalEstimate < 100_000 || payload.EstimatedInputTokens >= 1_000 { + t.Fatalf("input estimates = original:%d retained:%d, want large original and small retained estimate", originalEstimate, payload.EstimatedInputTokens) } } diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 5b192df4d0a..8bb058c6e2a 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -102,13 +102,13 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque } payload.ConversationState.CurrentMessage.UserInputMessage = current payload.ConversationState.History = trimLeadingAssistant(history) - payload.EstimatedInputTokens = estimateClaudeRequestInputTokens(&request) if request.MaxTokens > 0 || request.Temperature != nil || request.TopP != nil { payload.InferenceConfig = &kiroInferenceConfig{MaxTokens: request.MaxTokens, Temperature: request.Temperature, TopP: request.TopP} } if errTruncate := truncatePayload(payload, systemPrompt != ""); errTruncate != nil { return nil, nil, errTruncate } + payload.EstimatedInputTokens = estimateJSONTokens(payload) return payload, &request, nil } @@ -265,7 +265,7 @@ func convertTools(tools []claudeTool) ([]kiroToolWrapper, map[string]string, map out := make([]kiroToolWrapper, 0, len(tools)) nameMap := make(map[string]string) requestNameMap := make(map[string]string) - used := make(map[string]int) + used := make(map[string]struct{}) for _, tool := range tools { if strings.HasPrefix(strings.ToLower(tool.Type), "web_search") { continue @@ -314,16 +314,24 @@ func sanitizeToolName(name string) string { return name } -func uniqueToolName(name string, used map[string]int) string { - used[name]++ - if used[name] == 1 { +func uniqueToolName(name string, used map[string]struct{}) string { + if _, exists := used[name]; !exists { + used[name] = struct{}{} return name } - suffix := fmt.Sprintf("_%d", used[name]) - if len(name)+len(suffix) > 64 { - name = name[:64-len(suffix)] + for index := 2; ; index++ { + suffix := fmt.Sprintf("_%d", index) + base := name + if len(base)+len(suffix) > 64 { + base = base[:64-len(suffix)] + } + candidate := base + suffix + if _, exists := used[candidate]; exists { + continue + } + used[candidate] = struct{}{} + return candidate } - return name + suffix } func readableToolResults(results []kiroToolResult) string { From e220811ec86e199154b7c40c6b7b731d6733a6bf Mon Sep 17 00:00:00 2001 From: ianwangye Date: Fri, 7 Aug 2026 17:28:07 +0800 Subject: [PATCH 11/20] fix(plugin): address latest review feedback --- examples/plugin/kiro/go/kiro_test.go | 8 ++++---- examples/plugin/kiro/go/response.go | 9 ++++----- internal/pluginhost/adapters_executors.go | 9 +++++++-- internal/pluginhost/adapters_test.go | 21 +++++++++++++++++++++ internal/pluginhost/executor_usage.go | 2 +- 5 files changed, 37 insertions(+), 12 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index a5c123d7918..4271b1dd356 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -580,7 +580,7 @@ func TestExecuteStreamEmitsAnthropicSSE(t *testing.T) { } } -func TestAccumulatorMergesAdjacentFragments(t *testing.T) { +func TestAccumulatorOmitsUnsignedReasoningAndMergesTextFragments(t *testing.T) { payload := &kiroPayload{} payload.ConversationState.CurrentMessage.UserInputMessage.ModelID = "claude-sonnet-4.5" accumulator := newAccumulator(payload) @@ -596,10 +596,10 @@ func TestAccumulatorMergesAdjacentFragments(t *testing.T) { t.Fatal(errAccept) } } - if len(accumulator.Blocks) != 2 { - t.Fatalf("content blocks = %#v, want one thinking and one text block", accumulator.Blocks) + if len(accumulator.Blocks) != 1 { + t.Fatalf("content blocks = %#v, want one text block", accumulator.Blocks) } - if accumulator.Blocks[0].Thinking != "Think first." || accumulator.Blocks[1].Text != "Hello world." { + if accumulator.Blocks[0].Type != "text" || accumulator.Blocks[0].Text != "Hello world." { t.Fatalf("merged content blocks = %#v", accumulator.Blocks) } } diff --git a/examples/plugin/kiro/go/response.go b/examples/plugin/kiro/go/response.go index d268b14842c..6d403b92696 100644 --- a/examples/plugin/kiro/go/response.go +++ b/examples/plugin/kiro/go/response.go @@ -25,11 +25,10 @@ func (a *responseAccumulator) accept(event kiroEvent) ([]claudeContentBlock, err return []claudeContentBlock{block}, nil } case "reasoningContentEvent": - if text := firstStringField(event.Payload, "text", "content"); text != "" { - block := claudeContentBlock{Type: "thinking", Thinking: text} - a.appendFragment(block) - return []claudeContentBlock{block}, nil - } + // Kiro does not return an Anthropic-verifiable signature. Omit its + // reasoning instead of emitting a Claude thinking block that clients + // cannot safely replay. + return nil, nil case "toolUseEvent": tools, errTools := a.pendingTools.accept(event.Payload) if errTools != nil { diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index d5556d1b821..b5ea9824a8c 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -634,6 +634,9 @@ func executorSSERecordEnd(payload []byte) int { if index := bytes.Index(payload, []byte("\r\n\r\n")); index >= 0 && (end < 0 || index+4 < end) { end = index + 4 } + if index := bytes.Index(payload, []byte("\r\r")); index >= 0 && (end < 0 || index+2 < end) { + end = index + 2 + } return end } @@ -655,7 +658,9 @@ func (a *executorAdapter) translateExecutorStreamPayload(ctx context.Context, pr } func executorStreamTranslationPayloads(payload []byte) [][]byte { - if !bytes.Contains(payload, []byte("\n")) { + normalized := bytes.ReplaceAll(payload, []byte("\r\n"), []byte("\n")) + normalized = bytes.ReplaceAll(normalized, []byte("\r"), []byte("\n")) + if !bytes.Contains(normalized, []byte("\n")) { return [][]byte{payload} } var translationPayloads [][]byte @@ -674,7 +679,7 @@ func executorStreamTranslationPayloads(payload []byte) [][]byte { translationPayloads = append(translationPayloads, frame) dataValues = nil } - for _, line := range bytes.Split(payload, []byte("\n")) { + for _, line := range bytes.Split(normalized, []byte("\n")) { trimmed := bytes.TrimSpace(line) if len(trimmed) == 0 { flushData() diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 464436e65ba..394eab36032 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2906,6 +2906,27 @@ func TestExecutorSSERecordBufferFlushesFinalRecordWithoutDelimiter(t *testing.T) } } +func TestExecutorSSERecordBufferRecognizesCROnlyDelimiter(t *testing.T) { + var buffer executorSSERecordBuffer + payload := []byte("event: message_delta\rdata: {\"type\":\"message_delta\"}\r\r") + records := buffer.Push(payload) + if len(records) != 1 || !bytes.Equal(records[0], payload) { + t.Fatalf("CR-only records = %q, want one complete record %q", records, payload) + } + if tail := buffer.Flush(); len(tail) != 0 { + t.Fatalf("buffer tail = %q, want empty", tail) + } +} + +func TestExecutorStreamTranslationPayloadsExtractsDataFromCROnlySSEFrame(t *testing.T) { + payload := []byte("event: content_block_delta\rdata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\r\r") + frames := executorStreamTranslationPayloads(payload) + want := []byte(`data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hello"}}`) + if len(frames) != 1 || !bytes.Equal(frames[0], want) { + t.Fatalf("CR-only translation payloads = %q, want %q", frames, want) + } +} + func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { var usageBuffer helps.StreamUsageBuffer payload := []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":12,\"output_tokens\":4}}\n\n") diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go index 3453b762835..2e246442765 100644 --- a/internal/pluginhost/executor_usage.go +++ b/internal/pluginhost/executor_usage.go @@ -51,7 +51,7 @@ func (u *pluginExecutorUsage) observeStream(ctx context.Context, format sdktrans var sseBuffer executorSSERecordBuffer var pendingUsagePayload []byte var terminalErr error - isSSE := strings.Contains(strings.ToLower(headers.Get("Content-Type")), "text/event-stream") + isSSE := pluginExecutorSupportsStreamUsage(format) && strings.Contains(strings.ToLower(headers.Get("Content-Type")), "text/event-stream") for chunk := range in { if chunk.Err != nil { terminalErr = chunk.Err From 828ffda703d1bc8314a9a2df1d8446da408d74ce Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 22:52:44 +0800 Subject: [PATCH 12/20] fix(plugin): support Claude assistant prefills --- examples/plugin/kiro/go/kiro_test.go | 14 +++++++++++--- examples/plugin/kiro/go/translate.go | 3 --- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 4271b1dd356..24dfc66e4ad 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -446,7 +446,7 @@ func TestClaudeToKiroPreservesToolResultErrorStatus(t *testing.T) { } } -func TestClaudeToKiroRejectsAssistantPrefill(t *testing.T) { +func TestClaudeToKiroSupportsAssistantPrefill(t *testing.T) { raw := []byte(`{ "model":"claude-sonnet-4-5", "messages":[ @@ -454,8 +454,16 @@ func TestClaudeToKiroRejectsAssistantPrefill(t *testing.T) { {"role":"assistant","content":"The answer starts with"} ] }`) - if _, _, errTranslate := claudeToKiro(raw, ""); errTranslate == nil || !strings.Contains(errTranslate.Error(), "assistant prefills") { - t.Fatalf("claudeToKiro() error = %v, want unsupported assistant prefill error", errTranslate) + payload, _, errTranslate := claudeToKiro(raw, "") + if errTranslate != nil { + t.Fatalf("claudeToKiro() error = %v", errTranslate) + } + if got := payload.ConversationState.CurrentMessage.UserInputMessage.Content; got != "." { + t.Fatalf("current content = %q, want fallback continuation", got) + } + history := payload.ConversationState.History + if len(history) != 2 || history[1].AssistantResponseMessage == nil || history[1].AssistantResponseMessage.Content != "The answer starts with" { + t.Fatalf("history = %#v, want assistant prefill preserved", history) } } diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 8bb058c6e2a..06642fc272b 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -33,9 +33,6 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if len(request.Messages) == 0 { return nil, nil, fmt.Errorf("messages must not be empty") } - if lastRole := strings.ToLower(strings.TrimSpace(request.Messages[len(request.Messages)-1].Role)); lastRole != "user" { - return nil, nil, fmt.Errorf("last message must have role user; Kiro does not support assistant prefills") - } systemPrompt := extractSystemPrompt(request.System) if request.Thinking != nil { From 0f4621e686c23b51a36b9682b3498fcfe56fa9a9 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Thu, 6 Aug 2026 22:53:07 +0800 Subject: [PATCH 13/20] deploy(gcp): add Cloud Run container configuration --- .containerignore | 44 ++++++++++++++++++++++++++ Containerfile | 51 ++++++++++++++++++++++++++++++ deploy/gcp/config.yaml.template | 55 +++++++++++++++++++++++++++++++++ 3 files changed, 150 insertions(+) create mode 100644 .containerignore create mode 100644 Containerfile create mode 100644 deploy/gcp/config.yaml.template diff --git a/.containerignore b/.containerignore new file mode 100644 index 00000000000..30a9f24d99a --- /dev/null +++ b/.containerignore @@ -0,0 +1,44 @@ +.git +.github +.agents +.codex +.claude +.gemini +.idea +.vscode +.worktrees + +.env +config.yaml +auths +logs +conv +pgstore +gitstore +objectstore +plugins +static +temp +refs + +bin +examples/plugin/bin +*.exe +cli-proxy-api +test-output + +Dockerfile +docker-compose.yml +docker-compose.cluster.yml +.dockerignore +.containerignore +.gitignore +.goreleaser.yml + +docs +README.md +README_CN.md +README_JA.md +AGENTS.md +CLAUDE.md +GEMINI.md diff --git a/Containerfile b/Containerfile new file mode 100644 index 00000000000..818bb16f635 --- /dev/null +++ b/Containerfile @@ -0,0 +1,51 @@ +FROM golang:1.26-bookworm AS builder + +WORKDIR /src + +RUN apt-get update \ + && apt-get install -y --no-install-recommends build-essential git \ + && rm -rf /var/lib/apt/lists/* + +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . + +ARG VERSION=dev +ARG COMMIT=none +ARG BUILD_DATE=unknown + +RUN CGO_ENABLED=1 GOOS=linux go build \ + -buildvcs=false \ + -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.Commit=${COMMIT}' -X 'main.BuildDate=${BUILD_DATE}'" \ + -o /out/CLIProxyAPI \ + ./cmd/server/ + +RUN plugin_arch="$(go env GOARCH)" \ + && plugin_dir="/out/plugins/linux/${plugin_arch}" \ + && mkdir -p "${plugin_dir}" \ + && cd examples/plugin/kiro/go \ + && CGO_ENABLED=1 GOOS=linux go build \ + -buildmode=c-shared \ + -o "${plugin_dir}/kiro-go.so" \ + . \ + && rm -f "${plugin_dir}/kiro-go.h" + +FROM debian:bookworm-slim + +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates tzdata \ + && rm -rf /var/lib/apt/lists/* \ + && mkdir -p /CLIProxyAPI + +COPY --from=builder /out/CLIProxyAPI /CLIProxyAPI/CLIProxyAPI +COPY --from=builder /out/plugins /CLIProxyAPI/plugins +COPY config.example.yaml /CLIProxyAPI/config.example.yaml + +WORKDIR /CLIProxyAPI + +EXPOSE 8317 + +ENV TZ=UTC + +CMD ["./CLIProxyAPI"] diff --git a/deploy/gcp/config.yaml.template b/deploy/gcp/config.yaml.template new file mode 100644 index 00000000000..acb3af8bd88 --- /dev/null +++ b/deploy/gcp/config.yaml.template @@ -0,0 +1,55 @@ +host: "" +port: 8317 + +tls: + enable: false + cert: "" + key: "" + +remote-management: + allow-remote: false + secret-key: "" + disable-control-panel: true + disable-auto-update-panel: true + +auth-dir: "/root/.cli-proxy-api" + +api-keys: + - "__CLIPROXY_API_KEY__" + +debug: false + +pprof: + enable: false + addr: "127.0.0.1:8316" + +commercial-mode: false +logging-to-file: false +request-log: false +usage-statistics-enabled: false +passthrough-headers: false + +request-retry: 3 +max-retry-credentials: 0 +max-retry-interval: 30 +disable-cooling: false +save-cooldown-status: false + +streaming: + keepalive-seconds: 15 + bootstrap-retries: 1 + +routing: + strategy: "round-robin" + session-affinity: true + session-affinity-ttl: "1h" + +plugins: + enabled: true + dir: "plugins" + configs: + kiro-go: + enabled: true + priority: 1 + models: + - "*" From 020f69db01174140ef37dfcc94fee0c568611596 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Fri, 7 Aug 2026 18:42:44 +0800 Subject: [PATCH 14/20] fix(plugin): address Kiro review constraints --- examples/plugin/kiro/go/auth.go | 5 ++- examples/plugin/kiro/go/eventstream.go | 27 +++++------- examples/plugin/kiro/go/executor.go | 12 +++++- examples/plugin/kiro/go/kiro_test.go | 60 ++++++++++++++++++++++++++ examples/plugin/kiro/go/protocol.go | 24 +++++++---- examples/plugin/kiro/go/translate.go | 6 +++ 6 files changed, 107 insertions(+), 27 deletions(-) diff --git a/examples/plugin/kiro/go/auth.go b/examples/plugin/kiro/go/auth.go index 64c6043d431..ede658e64f5 100644 --- a/examples/plugin/kiro/go/auth.go +++ b/examples/plugin/kiro/go/auth.go @@ -4,6 +4,7 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "fmt" "os" "path/filepath" @@ -15,6 +16,8 @@ import ( var regionPattern = regexp.MustCompile(`^[a-z]{2}(?:-gov)?-[a-z]+-\d+$`) +var errKiroAPIKeyUnavailable = errors.New("Kiro API key unavailable") + type kiroCredential struct { Type string `json:"type"` APIKey string `json:"api_key,omitempty"` @@ -134,7 +137,7 @@ func resolveAPIKey(credential kiroCredential) (string, error) { } key := strings.TrimSpace(os.Getenv(credential.APIKeyEnv)) if key == "" { - return "", fmt.Errorf("environment variable %s is empty", credential.APIKeyEnv) + return "", fmt.Errorf("%w: environment variable %s is empty", errKiroAPIKeyUnavailable, credential.APIKeyEnv) } return key, nil } diff --git a/examples/plugin/kiro/go/eventstream.go b/examples/plugin/kiro/go/eventstream.go index 6783a4ee7ba..e2ea1d7e830 100644 --- a/examples/plugin/kiro/go/eventstream.go +++ b/examples/plugin/kiro/go/eventstream.go @@ -287,7 +287,13 @@ func readFloat(values map[string]any, keys ...string) (float64, bool) { func updateUsage(event map[string]any, inputTokens, outputTokens int) (int, int) { candidates := []map[string]any{event} - collectUsageMaps(event, &candidates) + appendDirectUsageMaps(event, &candidates) + for _, key := range []string{"metrics", "metadata"} { + if container, ok := event[key].(map[string]any); ok { + candidates = append(candidates, container) + appendDirectUsageMaps(container, &candidates) + } + } for _, candidate := range candidates { if value, ok := readNumber(candidate, "outputTokens", "completionTokens", "totalOutputTokens", "output_tokens", "completion_tokens", "total_output_tokens"); ok { outputTokens = value @@ -310,21 +316,10 @@ func updateUsage(event map[string]any, inputTokens, outputTokens int) (int, int) return inputTokens, outputTokens } -func collectUsageMaps(value any, candidates *[]map[string]any) { - switch typed := value.(type) { - case map[string]any: - for key, child := range typed { - normalized := strings.ToLower(key) - if normalized == "usage" || normalized == "tokenusage" || normalized == "token_usage" { - if nested, ok := child.(map[string]any); ok { - *candidates = append(*candidates, nested) - } - } - collectUsageMaps(child, candidates) - } - case []any: - for _, child := range typed { - collectUsageMaps(child, candidates) +func appendDirectUsageMaps(value map[string]any, candidates *[]map[string]any) { + for _, key := range []string{"usage", "tokenUsage", "token_usage"} { + if nested, ok := value[key].(map[string]any); ok { + *candidates = append(*candidates, nested) } } } diff --git a/examples/plugin/kiro/go/executor.go b/examples/plugin/kiro/go/executor.go index 6ac028b1a2c..6b1dcdbe8c0 100644 --- a/examples/plugin/kiro/go/executor.go +++ b/examples/plugin/kiro/go/executor.go @@ -3,6 +3,7 @@ package main import ( "bytes" "encoding/json" + "errors" "fmt" "net/http" "net/url" @@ -70,7 +71,7 @@ func execute(raw []byte) ([]byte, error) { } upstream, payload, errPrepare := prepareUpstreamRequest(request.ExecutorRequest) if errPrepare != nil { - return errorEnvelope("invalid_request", errPrepare.Error(), http.StatusBadRequest, false), nil + return prepareRequestErrorEnvelope(errPrepare), nil } var lastErr error @@ -118,12 +119,19 @@ func executeStream(raw []byte) ([]byte, error) { } upstream, payload, errPrepare := prepareUpstreamRequest(request.ExecutorRequest) if errPrepare != nil { - return errorEnvelope("invalid_request", errPrepare.Error(), http.StatusBadRequest, false), nil + return prepareRequestErrorEnvelope(errPrepare), nil } go runStream(request, upstream, payload) return okEnvelope(map[string]any{"headers": http.Header{"Content-Type": []string{"text/event-stream"}}}) } +func prepareRequestErrorEnvelope(errPrepare error) []byte { + if errors.Is(errPrepare, errKiroAPIKeyUnavailable) { + return errorEnvelope("invalid_auth", errPrepare.Error(), http.StatusUnauthorized, false) + } + return errorEnvelope("invalid_request", errPrepare.Error(), http.StatusBadRequest, false) +} + func runStream(request rpcExecutorRequest, upstream upstreamRequest, payload *kiroPayload) { var terminalErr error defer func() { diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 24dfc66e4ad..bc07c5cc161 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/binary" "encoding/json" + "errors" "hash/crc32" "net/http" "net/url" @@ -56,6 +57,23 @@ func TestDecodeCredentialRejectsHostInjectionRegion(t *testing.T) { } } +func TestMissingEnvironmentAPIKeyIsAuthenticationFailure(t *testing.T) { + t.Setenv("KIRO_MISSING_TEST_KEY", "") + _, errKey := resolveAPIKey(kiroCredential{APIKeyEnv: "KIRO_MISSING_TEST_KEY"}) + if !errors.Is(errKey, errKiroAPIKeyUnavailable) { + t.Fatalf("resolveAPIKey() error = %v, want API-key unavailable", errKey) + } + + raw := prepareRequestErrorEnvelope(errKey) + var env envelope + if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if env.Error == nil || env.Error.Code != "invalid_auth" || env.Error.HTTPStatus != http.StatusUnauthorized { + t.Fatalf("prepare error envelope = %#v, want invalid_auth/401", env.Error) + } +} + func TestModelsForAuthDiscoversPaginatesAndCaches(t *testing.T) { configureKiroTest(t, "models:\n - '*'\n") t.Setenv("KIRO_TEST_KEY", "test-discovery-key") @@ -360,6 +378,38 @@ func TestClaudeToKiroPreservesExplicitZeroSamplingValues(t *testing.T) { } } +func TestClaudeToKiroRejectsUnsupportedGenerationControls(t *testing.T) { + tests := []struct { + name string + body string + want string + }{ + { + name: "stop sequences", + body: `{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":"hello"}],"stop_sequences":["END"]}`, + want: "stop_sequences", + }, + { + name: "forced tool", + body: `{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":"hello"}],"tool_choice":{"type":"tool","name":"lookup"}}`, + want: "automatic tool choice", + }, + { + name: "disabled parallel tools", + body: `{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":"hello"}],"tool_choice":{"type":"auto","disable_parallel_tool_use":true}}`, + want: "parallel-tool restrictions", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, _, errTranslate := claudeToKiro([]byte(test.body), "") + if errTranslate == nil || !strings.Contains(errTranslate.Error(), test.want) { + t.Fatalf("claudeToKiro() error = %v, want %q", errTranslate, test.want) + } + }) + } +} + func TestClaudeToKiroPreservesThinkingBudget(t *testing.T) { payload, _, errTranslate := claudeToKiro([]byte(`{ "model":"claude-sonnet-4-5", @@ -660,6 +710,16 @@ func TestUpdateUsageReadsNestedCacheBuckets(t *testing.T) { } } +func TestUpdateUsageIgnoresNestedToolInputUsage(t *testing.T) { + event := map[string]any{"input": map[string]any{"usage": map[string]any{ + "inputTokens": 100000.0, "outputTokens": 100000.0, + }}} + inputTokens, outputTokens := updateUsage(event, 12, 2) + if inputTokens != 12 || outputTokens != 2 { + t.Fatalf("usage = input:%d output:%d, want existing provider counts", inputTokens, outputTokens) + } +} + func configureKiroTest(t *testing.T, configYAML string) { t.Helper() request, errMarshal := json.Marshal(lifecycleRequest{ConfigYAML: []byte(configYAML)}) diff --git a/examples/plugin/kiro/go/protocol.go b/examples/plugin/kiro/go/protocol.go index 2d6d6f0d25d..c1e1cf17445 100644 --- a/examples/plugin/kiro/go/protocol.go +++ b/examples/plugin/kiro/go/protocol.go @@ -89,20 +89,28 @@ type kiroInferenceConfig struct { } type claudeRequest struct { - Model string `json:"model"` - Messages []claudeMessage `json:"messages"` - MaxTokens int `json:"max_tokens"` - Temperature *float64 `json:"temperature,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - Stream bool `json:"stream,omitempty"` - System any `json:"system,omitempty"` - Thinking *struct { + Model string `json:"model"` + Messages []claudeMessage `json:"messages"` + MaxTokens int `json:"max_tokens"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + Stream bool `json:"stream,omitempty"` + System any `json:"system,omitempty"` + StopSequences []string `json:"stop_sequences,omitempty"` + ToolChoice *claudeToolChoice `json:"tool_choice,omitempty"` + Thinking *struct { Type string `json:"type,omitempty"` BudgetTokens int `json:"budget_tokens,omitempty"` } `json:"thinking,omitempty"` Tools []claudeTool `json:"tools,omitempty"` } +type claudeToolChoice struct { + Type string `json:"type"` + Name string `json:"name,omitempty"` + DisableParallelToolUse bool `json:"disable_parallel_tool_use,omitempty"` +} + type claudeMessage struct { Role string `json:"role"` Content any `json:"content"` diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 06642fc272b..1f1d87919d4 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -33,6 +33,12 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if len(request.Messages) == 0 { return nil, nil, fmt.Errorf("messages must not be empty") } + if len(request.StopSequences) > 0 { + return nil, nil, fmt.Errorf("Kiro provider does not support stop_sequences") + } + if request.ToolChoice != nil && (strings.ToLower(strings.TrimSpace(request.ToolChoice.Type)) != "auto" || request.ToolChoice.DisableParallelToolUse) { + return nil, nil, fmt.Errorf("Kiro provider supports only automatic tool choice without parallel-tool restrictions") + } systemPrompt := extractSystemPrompt(request.System) if request.Thinking != nil { From af3d87321282cb95f37b15c1b571500cbedf9f50 Mon Sep 17 00:00:00 2001 From: ianwangye Date: Sat, 8 Aug 2026 00:45:40 +0800 Subject: [PATCH 15/20] fix(pluginhost): preserve auth HTTP client in callbacks --- internal/pluginhost/callback_contexts.go | 38 +++++++++--- internal/pluginhost/host_callbacks.go | 12 +++- internal/pluginhost/host_callbacks_test.go | 60 +++++++++++++++++++ internal/pluginhost/rpc_client.go | 8 +-- internal/pluginhost/rpc_client_stream.go | 2 +- internal/pluginhost/rpc_client_stream_test.go | 6 +- 6 files changed, 111 insertions(+), 15 deletions(-) diff --git a/internal/pluginhost/callback_contexts.go b/internal/pluginhost/callback_contexts.go index 27c5aaded12..c1cd5bb11ee 100644 --- a/internal/pluginhost/callback_contexts.go +++ b/internal/pluginhost/callback_contexts.go @@ -6,6 +6,8 @@ import ( "strings" "sync" "sync/atomic" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) type callbackContextRegistry struct { @@ -15,16 +17,17 @@ type callbackContextRegistry struct { } type callbackContextEntry struct { - ctx context.Context - pluginID string - cleanup []func() + ctx context.Context + pluginID string + httpClient pluginapi.HostHTTPClient + cleanup []func() } func newCallbackContextRegistry() *callbackContextRegistry { return &callbackContextRegistry{contexts: make(map[string]callbackContextEntry)} } -func (r *callbackContextRegistry) open(ctx context.Context, pluginID string) (string, func()) { +func (r *callbackContextRegistry) open(ctx context.Context, pluginID string, httpClients ...pluginapi.HostHTTPClient) (string, func()) { if r == nil { return "", func() {} } @@ -33,9 +36,13 @@ func (r *callbackContextRegistry) open(ctx context.Context, pluginID string) (st } pluginID = strings.TrimSpace(pluginID) ctx = withHostCallbackPluginID(ctx, pluginID) + var httpClient pluginapi.HostHTTPClient + if len(httpClients) > 0 { + httpClient = httpClients[0] + } id := strconv.FormatUint(r.next.Add(1), 10) r.mu.Lock() - r.contexts[id] = callbackContextEntry{ctx: ctx, pluginID: pluginID} + r.contexts[id] = callbackContextEntry{ctx: ctx, pluginID: pluginID, httpClient: httpClient} r.mu.Unlock() var once sync.Once @@ -56,6 +63,16 @@ func (r *callbackContextRegistry) open(ctx context.Context, pluginID string) (st } } +func (r *callbackContextRegistry) httpClient(id string) pluginapi.HostHTTPClient { + if r == nil || id == "" { + return nil + } + r.mu.RLock() + httpClient := r.contexts[id].httpClient + r.mu.RUnlock() + return httpClient +} + func (r *callbackContextRegistry) pluginID(id string) string { if r == nil || id == "" { return "" @@ -104,11 +121,11 @@ func (h *Host) openCallbackContext(ctx context.Context) (string, func()) { return h.openCallbackContextForPlugin(ctx, "") } -func (h *Host) openCallbackContextForPlugin(ctx context.Context, pluginID string) (string, func()) { +func (h *Host) openCallbackContextForPlugin(ctx context.Context, pluginID string, httpClients ...pluginapi.HostHTTPClient) (string, func()) { if h == nil || h.callbackContexts == nil { return "", func() {} } - return h.callbackContexts.open(ctx, pluginID) + return h.callbackContexts.open(ctx, pluginID, httpClients...) } func (h *Host) addCallbackCleanup(id string, cleanup func()) bool { @@ -137,3 +154,10 @@ func (h *Host) callbackContextPluginID(id string) string { } return h.callbackContexts.pluginID(id) } + +func (h *Host) callbackHTTPClient(id string) pluginapi.HostHTTPClient { + if h == nil || h.callbackContexts == nil { + return nil + } + return h.callbackContexts.httpClient(id) +} diff --git a/internal/pluginhost/host_callbacks.go b/internal/pluginhost/host_callbacks.go index 53c3bf544a1..0acc9309c4e 100644 --- a/internal/pluginhost/host_callbacks.go +++ b/internal/pluginhost/host_callbacks.go @@ -145,7 +145,11 @@ func (h *Host) callHostHTTPDo(ctx context.Context, request []byte) ([]byte, erro return nil, errDecode } ctx = h.resolveCallbackContext(callbackID, ctx) - resp, errDo := h.newHTTPClient(nil).Do(ctx, httpReq) + httpClient := h.callbackHTTPClient(callbackID) + if httpClient == nil { + httpClient = h.newHTTPClient(nil) + } + resp, errDo := httpClient.Do(ctx, httpReq) if errDo != nil { return nil, errDo } @@ -162,7 +166,11 @@ func (h *Host) callHostHTTPDoStream(ctx context.Context, request []byte) ([]byte ctx = context.Background() } streamCtx, cancel := context.WithCancel(ctx) - resp, errDo := h.newHTTPClient(nil).DoStream(streamCtx, httpReq) + httpClient := h.callbackHTTPClient(callbackID) + if httpClient == nil { + httpClient = h.newHTTPClient(nil) + } + resp, errDo := httpClient.DoStream(streamCtx, httpReq) if errDo != nil { cancel() return nil, errDo diff --git a/internal/pluginhost/host_callbacks_test.go b/internal/pluginhost/host_callbacks_test.go index 827b5694f08..bbbb6389223 100644 --- a/internal/pluginhost/host_callbacks_test.go +++ b/internal/pluginhost/host_callbacks_test.go @@ -26,6 +26,23 @@ type fakeHostModelExecutor struct { executeModelStream func(context.Context, handlers.ModelExecutionRequest) (handlers.ModelExecutionStream, *interfaces.ErrorMessage) } +type callbackHTTPClient struct { + doCalls int + doStreamCalls int +} + +func (c *callbackHTTPClient) Do(context.Context, pluginapi.HTTPRequest) (pluginapi.HTTPResponse, error) { + c.doCalls++ + return pluginapi.HTTPResponse{StatusCode: http.StatusCreated, Body: []byte("selected-client")}, nil +} + +func (c *callbackHTTPClient) DoStream(context.Context, pluginapi.HTTPRequest) (pluginapi.HTTPStreamResponse, error) { + c.doStreamCalls++ + chunks := make(chan pluginapi.HTTPStreamChunk) + close(chunks) + return pluginapi.HTTPStreamResponse{StatusCode: http.StatusAccepted, Chunks: chunks}, nil +} + func (e *fakeHostModelExecutor) ExecuteModel(ctx context.Context, req handlers.ModelExecutionRequest) (handlers.ModelExecutionResponse, *interfaces.ErrorMessage) { return e.executeModel(ctx, req) } @@ -71,6 +88,49 @@ func TestHostHTTPDoCallbackUsesHostHTTPClient(t *testing.T) { } } +func TestHostHTTPCallbacksUseRegisteredAuthAwareClient(t *testing.T) { + host := New() + httpClient := &callbackHTTPClient{} + callbackID, closeCallback := host.openCallbackContextForPlugin(context.Background(), "test-plugin", httpClient) + defer closeCallback() + + rawReq, errMarshal := json.Marshal(rpcHostHTTPRequest{ + HostCallbackID: callbackID, + Method: http.MethodPost, + URL: "https://example.test", + }) + if errMarshal != nil { + t.Fatal(errMarshal) + } + rawResp, errCall := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPDo, rawReq) + if errCall != nil { + t.Fatalf("HTTP callback error = %v", errCall) + } + resp, errDecode := decodeRPCEnvelope[pluginapi.HTTPResponse](rawResp) + if errDecode != nil { + t.Fatal(errDecode) + } + if resp.StatusCode != http.StatusCreated || string(resp.Body) != "selected-client" { + t.Fatalf("HTTP response = %#v, want registered client response", resp) + } + + rawStreamResp, errCallStream := host.callFromPlugin(context.Background(), pluginabi.MethodHostHTTPDoStream, rawReq) + if errCallStream != nil { + t.Fatalf("stream callback error = %v", errCallStream) + } + streamResp, errDecodeStream := decodeRPCEnvelope[rpcHostHTTPStreamResponse](rawStreamResp) + if errDecodeStream != nil { + t.Fatal(errDecodeStream) + } + if streamResp.StatusCode != http.StatusAccepted || streamResp.StreamID == "" { + t.Fatalf("stream response = %#v, want registered client response", streamResp) + } + host.httpStreams.close(streamResp.StreamID) + if httpClient.doCalls != 1 || httpClient.doStreamCalls != 1 { + t.Fatalf("registered client calls = do:%d stream:%d, want one each", httpClient.doCalls, httpClient.doStreamCalls) + } +} + func TestHostHTTPDoCallbackRestoresRegisteredRequestContext(t *testing.T) { gin.SetMode(gin.TestMode) ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 01319fd78ba..42935fc9361 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -336,11 +336,11 @@ func marshalRPCError(code, message string) []byte { return raw } -func (a *rpcPluginAdapter) openHostCallbackContext(ctx context.Context) (string, func()) { +func (a *rpcPluginAdapter) openHostCallbackContext(ctx context.Context, httpClients ...pluginapi.HostHTTPClient) (string, func()) { if a == nil || a.host == nil { return "", func() {} } - return a.host.openCallbackContextForPlugin(ctx, a.id) + return a.host.openCallbackContextForPlugin(ctx, a.id, httpClients...) } func (a *rpcPluginAdapter) RegisterModels(ctx context.Context, req pluginapi.ModelRegistrationRequest) (pluginapi.ModelRegistrationResponse, error) { @@ -352,7 +352,7 @@ func (a *rpcPluginAdapter) StaticModels(ctx context.Context, req pluginapi.Stati } func (a *rpcPluginAdapter) ModelsForAuth(ctx context.Context, req pluginapi.AuthModelRequest) (pluginapi.ModelResponse, error) { - callbackID, closeCallback := a.openHostCallbackContext(ctx) + callbackID, closeCallback := a.openHostCallbackContext(ctx, req.HTTPClient) defer closeCallback() return callPlugin[pluginapi.ModelResponse](ctx, a.client, pluginabi.MethodModelForAuth, rpcAuthModelRequest{ AuthModelRequest: req, @@ -433,7 +433,7 @@ func (a *rpcPluginAdapter) Authenticate(ctx context.Context, req pluginapi.Front } func (a *rpcPluginAdapter) Execute(ctx context.Context, req pluginapi.ExecutorRequest) (pluginapi.ExecutorResponse, error) { - callbackID, closeCallback := a.openHostCallbackContext(ctx) + callbackID, closeCallback := a.openHostCallbackContext(ctx, req.HTTPClient) defer closeCallback() return callPlugin[pluginapi.ExecutorResponse](ctx, a.client, pluginabi.MethodExecutorExecute, rpcExecutorRequest{ ExecutorRequest: req, diff --git a/internal/pluginhost/rpc_client_stream.go b/internal/pluginhost/rpc_client_stream.go index 87939146a01..c0a11af4618 100644 --- a/internal/pluginhost/rpc_client_stream.go +++ b/internal/pluginhost/rpc_client_stream.go @@ -14,7 +14,7 @@ func (a *rpcPluginAdapter) ExecuteStream(ctx context.Context, req pluginapi.Exec return pluginapi.ExecutorStreamResponse{}, fmt.Errorf("plugin stream bridge is unavailable") } streamID, chunks, cleanupStream := a.host.streams.open(ctx) - callbackID, closeCallback := a.openHostCallbackContext(ctx) + callbackID, closeCallback := a.openHostCallbackContext(ctx, req.HTTPClient) cleanup := combinedCleanup(cleanupStream, closeCallback) rpcReq := rpcExecutorRequest{ ExecutorRequest: req, diff --git a/internal/pluginhost/rpc_client_stream_test.go b/internal/pluginhost/rpc_client_stream_test.go index 6e293a248a2..8b95ab8924f 100644 --- a/internal/pluginhost/rpc_client_stream_test.go +++ b/internal/pluginhost/rpc_client_stream_test.go @@ -14,6 +14,7 @@ import ( func TestRPCExecuteStreamKeepsHostCallbackScopeUntilStreamCloses(t *testing.T) { host := New() + httpClient := host.newHTTPClient(nil) client := newStreamCallbackPluginClient() adapter := &rpcPluginAdapter{ id: "stream-plugin", @@ -21,7 +22,7 @@ func TestRPCExecuteStreamKeepsHostCallbackScopeUntilStreamCloses(t *testing.T) { client: client, } - stream, errStream := adapter.ExecuteStream(context.Background(), pluginapi.ExecutorRequest{Stream: true}) + stream, errStream := adapter.ExecuteStream(context.Background(), pluginapi.ExecutorRequest{Stream: true, HTTPClient: httpClient}) if errStream != nil { t.Fatalf("ExecuteStream() error = %v", errStream) } @@ -32,6 +33,9 @@ func TestRPCExecuteStreamKeepsHostCallbackScopeUntilStreamCloses(t *testing.T) { if !callbackContextExists(host, client.callbackID) { t.Fatal("host callback scope closed before plugin stream closed") } + if got := host.callbackHTTPClient(client.callbackID); got != httpClient { + t.Fatalf("callback HTTP client = %T, want selected auth-aware client", got) + } closeReq, errMarshal := json.Marshal(rpcStreamCloseRequest{StreamID: client.streamID}) if errMarshal != nil { From f11a111bc34e653229dc4070082754ac393ad551 Mon Sep 17 00:00:00 2001 From: wangyedev Date: Sun, 9 Aug 2026 19:49:31 +0800 Subject: [PATCH 16/20] fix(plugin): reject unsupported URL images --- examples/plugin/kiro/go/kiro_test.go | 24 ++++++++++++++++++++++++ examples/plugin/kiro/go/translate.go | 25 +++++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index bc07c5cc161..c25cf2daa1c 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -410,6 +410,30 @@ func TestClaudeToKiroRejectsUnsupportedGenerationControls(t *testing.T) { } } +func TestClaudeToKiroRejectsURLBackedImages(t *testing.T) { + tests := []struct { + name string + body string + }{ + { + name: "user image", + body: `{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":[{"type":"image","source":{"type":"url","url":"https://example.com/image.png"}}]}]}`, + }, + { + name: "tool result image", + body: `{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":[{"type":"image","source":{"type":"url","url":"https://example.com/tool.png"}}]}]}]}`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, _, errTranslate := claudeToKiro([]byte(test.body), "") + if errTranslate == nil || !strings.Contains(errTranslate.Error(), "URL-backed image") { + t.Fatalf("claudeToKiro() error = %v, want URL-backed image error", errTranslate) + } + }) + } +} + func TestClaudeToKiroPreservesThinkingBudget(t *testing.T) { payload, _, errTranslate := claudeToKiro([]byte(`{ "model":"claude-sonnet-4-5", diff --git a/examples/plugin/kiro/go/translate.go b/examples/plugin/kiro/go/translate.go index 1f1d87919d4..12d0600e464 100644 --- a/examples/plugin/kiro/go/translate.go +++ b/examples/plugin/kiro/go/translate.go @@ -33,6 +33,11 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque if len(request.Messages) == 0 { return nil, nil, fmt.Errorf("messages must not be empty") } + for _, message := range request.Messages { + if errImages := validateClaudeImageSources(message.Content); errImages != nil { + return nil, nil, errImages + } + } if len(request.StopSequences) > 0 { return nil, nil, fmt.Errorf("Kiro provider does not support stop_sequences") } @@ -115,6 +120,26 @@ func claudeToKiro(raw []byte, requestedModel string) (*kiroPayload, *claudeReque return payload, &request, nil } +func validateClaudeImageSources(content any) error { + for _, block := range contentBlocks(content) { + typeName, _ := block["type"].(string) + switch typeName { + case "image", "input_image": + source, _ := block["source"].(map[string]any) + sourceType, _ := source["type"].(string) + urlValue, _ := source["url"].(string) + if strings.EqualFold(strings.TrimSpace(sourceType), "url") || strings.TrimSpace(urlValue) != "" { + return fmt.Errorf("Kiro provider does not support URL-backed image sources; use base64 image data") + } + case "tool_result": + if errNested := validateClaudeImageSources(block["content"]); errNested != nil { + return errNested + } + } + } + return nil +} + func normalizeKiroModel(model string) string { model = strings.TrimSpace(model) model = strings.TrimSuffix(model, "-thinking") From ba9e808625ca9fe8d0d86a1c50a91f53a318626d Mon Sep 17 00:00:00 2001 From: wangyedev Date: Sun, 9 Aug 2026 20:58:03 +0800 Subject: [PATCH 17/20] fix(plugin): distinguish dated Claude snapshots --- examples/plugin/kiro/go/kiro_test.go | 17 +++++++++++++++++ examples/plugin/kiro/go/token_estimator.go | 12 ++++++++++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index c25cf2daa1c..03c13dd1145 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -550,6 +550,23 @@ func TestNormalizeKiroModelDoesNotRewriteDatedSnapshotAsDecimal(t *testing.T) { } } +func TestContextWindowTokensDistinguishesDatedSnapshots(t *testing.T) { + tests := []struct { + model string + want int + }{ + {model: "claude-opus-4-20250514", want: 200_000}, + {model: "claude-sonnet-4.6", want: 1_000_000}, + {model: "claude-opus-5", want: 1_000_000}, + } + + for _, test := range tests { + if got := contextWindowTokens(test.model); got != test.want { + t.Errorf("contextWindowTokens(%q) = %d, want %d", test.model, got, test.want) + } + } +} + func TestEventStreamDecoderHandlesSplitFramesAndCRC(t *testing.T) { frame := testEventFrame(t, "assistantResponseEvent", map[string]any{"content": "hello"}) decoder := &eventStreamDecoder{} diff --git a/examples/plugin/kiro/go/token_estimator.go b/examples/plugin/kiro/go/token_estimator.go index 837d89e4408..b31a248a213 100644 --- a/examples/plugin/kiro/go/token_estimator.go +++ b/examples/plugin/kiro/go/token_estimator.go @@ -8,7 +8,10 @@ import ( "strings" ) -var claudeVersionPattern = regexp.MustCompile(`claude-(?:opus|sonnet|haiku)-(\d+)(?:[.-](\d+))?`) +var ( + claudeVersionPattern = regexp.MustCompile(`claude-(?:opus|sonnet|haiku)-(\d+)(?:[.-](\d+))?`) + claudeDatedSnapshotPattern = regexp.MustCompile(`claude-(?:opus|sonnet|haiku)-\d+-\d{8}(?:[^0-9]|$)`) +) func estimateApproxTokens(text string) int { if text == "" { @@ -122,7 +125,12 @@ func stringValue(value any) string { } func contextWindowTokens(model string) int { - match := claudeVersionPattern.FindStringSubmatch(strings.ToLower(model)) + model = strings.ToLower(model) + if claudeDatedSnapshotPattern.MatchString(model) { + return 200_000 + } + + match := claudeVersionPattern.FindStringSubmatch(model) if len(match) == 3 { major, errMajor := strconv.Atoi(match[1]) minor := 0 From fb390c56e96756ba0917eab529854825d7807c84 Mon Sep 17 00:00:00 2001 From: wangyedev Date: Sun, 9 Aug 2026 22:02:51 +0800 Subject: [PATCH 18/20] fix(plugin): correct Kiro usage and proxy handling --- examples/plugin/kiro/go/kiro_test.go | 7 +++ examples/plugin/kiro/go/protocol.go | 1 + examples/plugin/kiro/go/response.go | 3 +- internal/pluginhost/adapters_test.go | 2 +- internal/pluginhost/executor_usage.go | 49 ++++++++++++++++--- internal/pluginhost/rpc_client.go | 2 +- internal/pluginhost/rpc_client_stream_test.go | 38 ++++++++++++++ 7 files changed, 92 insertions(+), 10 deletions(-) diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 03c13dd1145..5593d5177ae 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -701,6 +701,13 @@ func TestAccumulatorOmitsUnsignedReasoningAndMergesTextFragments(t *testing.T) { if accumulator.Blocks[0].Type != "text" || accumulator.Blocks[0].Text != "Hello world." { t.Fatalf("merged content blocks = %#v", accumulator.Blocks) } + if errFinish := accumulator.finish(); errFinish != nil { + t.Fatal(errFinish) + } + visibleTokens := estimateClaudeOutputTokens(accumulator.Blocks) + if accumulator.OutputTokens <= visibleTokens { + t.Fatalf("output tokens = %d, want hidden reasoning included beyond %d visible tokens", accumulator.OutputTokens, visibleTokens) + } } func TestAccumulatorEstimatesUsageWhenUpstreamOmitsTokens(t *testing.T) { diff --git a/examples/plugin/kiro/go/protocol.go b/examples/plugin/kiro/go/protocol.go index c1e1cf17445..4de4c6794a9 100644 --- a/examples/plugin/kiro/go/protocol.go +++ b/examples/plugin/kiro/go/protocol.go @@ -162,6 +162,7 @@ type responseAccumulator struct { StopReason string Credits float64 EstimatedInputTokens int + HiddenReasoningTokens int ContextUsagePercentage float64 ToolNames map[string]string pendingTools pendingToolUses diff --git a/examples/plugin/kiro/go/response.go b/examples/plugin/kiro/go/response.go index 6d403b92696..07819add67a 100644 --- a/examples/plugin/kiro/go/response.go +++ b/examples/plugin/kiro/go/response.go @@ -28,6 +28,7 @@ func (a *responseAccumulator) accept(event kiroEvent) ([]claudeContentBlock, err // Kiro does not return an Anthropic-verifiable signature. Omit its // reasoning instead of emitting a Claude thinking block that clients // cannot safely replay. + a.HiddenReasoningTokens += estimateApproxTokens(firstStringField(event.Payload, "content", "text")) return nil, nil case "toolUseEvent": tools, errTools := a.pendingTools.accept(event.Payload) @@ -108,7 +109,7 @@ func (a *responseAccumulator) finish() error { a.InputTokens = a.currentInputTokens() } if a.OutputTokens <= 0 { - a.OutputTokens = estimateClaudeOutputTokens(a.Blocks) + a.OutputTokens = estimateClaudeOutputTokens(a.Blocks) + a.HiddenReasoningTokens } return nil } diff --git a/internal/pluginhost/adapters_test.go b/internal/pluginhost/adapters_test.go index 394eab36032..5710e25cd63 100644 --- a/internal/pluginhost/adapters_test.go +++ b/internal/pluginhost/adapters_test.go @@ -2942,7 +2942,7 @@ func TestPluginExecutorUsageParsesFinalClaudeStreamCounts(t *testing.T) { func TestPluginExecutorUsageMergesClaudeStreamCounts(t *testing.T) { var usageBuffer helps.StreamUsageBuffer - observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, []byte("event: message_start\ndata: {\"type\":\"message_start\",\"usage\":{\"input_tokens\":12,\"output_tokens\":1,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":2}}\n\n")) + observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, []byte("event: message_start\ndata: {\"type\":\"message_start\",\"usage\":{\"input_tokens\":12,\"output_tokens\":10,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":2}}\n\n")) observePluginExecutorStreamUsage(&usageBuffer, sdktranslator.FormatClaude, []byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":4}}\n\n")) detail, ok := usageBuffer.Detail() if !ok { diff --git a/internal/pluginhost/executor_usage.go b/internal/pluginhost/executor_usage.go index 2e246442765..467389cd913 100644 --- a/internal/pluginhost/executor_usage.go +++ b/internal/pluginhost/executor_usage.go @@ -3,6 +3,7 @@ package pluginhost import ( "bytes" "context" + "encoding/json" "net/http" "strings" @@ -132,7 +133,7 @@ func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sd case sdktranslator.FormatClaude: detail, ok := helps.ParseClaudeStreamUsage(usagePayload) if previous, exists := buffer.Detail(); ok && exists { - detail = mergeClaudePluginStreamUsage(previous, detail) + detail = mergeClaudePluginStreamUsage(previous, detail, usagePayload) } buffer.Observe(detail, ok) case sdktranslator.FormatOpenAI: @@ -149,12 +150,23 @@ func observePluginExecutorStreamUsage(buffer *helps.StreamUsageBuffer, format sd } } -func mergeClaudePluginStreamUsage(previous, current coreusage.Detail) coreusage.Detail { - current.InputTokens = max(previous.InputTokens, current.InputTokens) - current.OutputTokens = max(previous.OutputTokens, current.OutputTokens) - current.ReasoningTokens = max(previous.ReasoningTokens, current.ReasoningTokens) - current.CacheReadTokens = max(previous.CacheReadTokens, current.CacheReadTokens) - current.CacheCreationTokens = max(previous.CacheCreationTokens, current.CacheCreationTokens) +func mergeClaudePluginStreamUsage(previous, current coreusage.Detail, payload []byte) coreusage.Detail { + fields := claudePluginStreamUsageFields(payload) + if !fields["input_tokens"] { + current.InputTokens = previous.InputTokens + } + if !fields["output_tokens"] { + current.OutputTokens = previous.OutputTokens + } + if !fields["reasoning_tokens"] { + current.ReasoningTokens = previous.ReasoningTokens + } + if !fields["cache_read_input_tokens"] { + current.CacheReadTokens = previous.CacheReadTokens + } + if !fields["cache_creation_input_tokens"] { + current.CacheCreationTokens = previous.CacheCreationTokens + } current.CachedTokens = current.CacheReadTokens if current.CachedTokens == 0 { current.CachedTokens = current.CacheCreationTokens @@ -179,3 +191,26 @@ func mergeClaudePluginStreamUsage(previous, current coreusage.Detail) coreusage. ) return current } + +func claudePluginStreamUsageFields(payload []byte) map[string]bool { + var envelope struct { + Usage map[string]json.RawMessage `json:"usage"` + } + if errUnmarshal := json.Unmarshal(helps.JSONPayload(payload), &envelope); errUnmarshal != nil { + return nil + } + fields := make(map[string]bool, len(envelope.Usage)+1) + for name := range envelope.Usage { + fields[name] = true + } + if fields["thinking_tokens"] { + fields["reasoning_tokens"] = true + } + if rawDetails, ok := envelope.Usage["output_tokens_details"]; ok { + var details map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(rawDetails, &details); errUnmarshal == nil { + fields["reasoning_tokens"] = fields["reasoning_tokens"] || details["thinking_tokens"] != nil || details["reasoning_tokens"] != nil + } + } + return fields +} diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 42935fc9361..f891d64f6cd 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -451,7 +451,7 @@ func (a *rpcPluginAdapter) CountTokens(ctx context.Context, req pluginapi.Execut } func (a *rpcPluginAdapter) HttpRequest(ctx context.Context, req pluginapi.ExecutorHTTPRequest) (pluginapi.ExecutorHTTPResponse, error) { - callbackID, closeCallback := a.openHostCallbackContext(ctx) + callbackID, closeCallback := a.openHostCallbackContext(ctx, req.HTTPClient) defer closeCallback() return callPlugin[pluginapi.ExecutorHTTPResponse](ctx, a.client, pluginabi.MethodExecutorHTTPRequest, rpcExecutorHTTPRequest{ ExecutorHTTPRequest: req, diff --git a/internal/pluginhost/rpc_client_stream_test.go b/internal/pluginhost/rpc_client_stream_test.go index 8b95ab8924f..a96ac73aae5 100644 --- a/internal/pluginhost/rpc_client_stream_test.go +++ b/internal/pluginhost/rpc_client_stream_test.go @@ -83,6 +83,23 @@ func TestRPCExecuteStreamClosesHostCallbackScopeOnContextCancelWhileChunkPending } } +func TestRPCExecutorHTTPRequestPassesHostHTTPClientToCallbackScope(t *testing.T) { + host := New() + httpClient := &callbackHTTPClient{} + client := &executorHTTPCallbackPluginClient{host: host} + adapter := &rpcPluginAdapter{id: "http-plugin", host: host, client: client} + + if _, errRequest := adapter.HttpRequest(context.Background(), pluginapi.ExecutorHTTPRequest{HTTPClient: httpClient}); errRequest != nil { + t.Fatalf("HttpRequest() error = %v", errRequest) + } + if client.httpClient != httpClient { + t.Fatalf("callback HTTP client = %T, want selected auth-aware client", client.httpClient) + } + if callbackContextExists(host, client.callbackID) { + t.Fatal("HTTP callback scope remained open after request completed") + } +} + func callbackContextExists(host *Host, callbackID string) bool { if host == nil || host.callbackContexts == nil { return false @@ -99,6 +116,27 @@ type streamCallbackPluginClient struct { callbackID string } +type executorHTTPCallbackPluginClient struct { + host *Host + callbackID string + httpClient pluginapi.HostHTTPClient +} + +func (c *executorHTTPCallbackPluginClient) Call(_ context.Context, method string, request []byte) ([]byte, error) { + if method != pluginabi.MethodExecutorHTTPRequest { + return nil, fmt.Errorf("method = %s, want %s", method, pluginabi.MethodExecutorHTTPRequest) + } + var req rpcExecutorHTTPRequest + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode executor HTTP request: %w", errUnmarshal) + } + c.callbackID = req.HostCallbackID + c.httpClient = c.host.callbackHTTPClient(req.HostCallbackID) + return marshalRPCResult(pluginapi.ExecutorHTTPResponse{StatusCode: http.StatusOK}) +} + +func (c *executorHTTPCallbackPluginClient) Shutdown() {} + func newStreamCallbackPluginClient() *streamCallbackPluginClient { return &streamCallbackPluginClient{called: make(chan struct{})} } From e99dc770a9e85bcdaa4794212baa21ea96e16f3a Mon Sep 17 00:00:00 2001 From: ianwangye Date: Mon, 10 Aug 2026 15:56:11 +0800 Subject: [PATCH 19/20] fix(plugin): align Kiro model discovery with Kiro-Go --- .env.example | 3 +- docker-compose.yml | 4 ++- examples/plugin/kiro/go/executor.go | 2 +- examples/plugin/kiro/go/kiro_test.go | 30 ++++++++++++++++++ examples/plugin/kiro/go/model_discovery.go | 37 ++++++++++++++++++---- 5 files changed, 66 insertions(+), 10 deletions(-) diff --git a/.env.example b/.env.example index a8b03899d6b..9d84939bb99 100644 --- a/.env.example +++ b/.env.example @@ -6,7 +6,8 @@ # ------------------------------------------------------------------------------ # Kiro provider plugin (optional) # ------------------------------------------------------------------------------ -# KIRO_API_KEY=replace-with-your-kiro-api-key +# KIRO_API_KEY_1=REPLACE_WITH_KIRO_ACCOUNT_1_API_KEY +# KIRO_API_KEY_2=REPLACE_WITH_KIRO_ACCOUNT_2_API_KEY # ------------------------------------------------------------------------------ # Management Web UI diff --git a/docker-compose.yml b/docker-compose.yml index d4c6cc27055..b92446955dc 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,6 +15,8 @@ services: environment: DEPLOY: ${DEPLOY:-} KIRO_API_KEY: ${KIRO_API_KEY:-} + KIRO_API_KEY_1: ${KIRO_API_KEY_1:-} + KIRO_API_KEY_2: ${KIRO_API_KEY_2:-} ports: - "8317:8317" - "8085:8085" @@ -24,7 +26,7 @@ services: - "11451:11451" volumes: - ${CLI_PROXY_CONFIG_PATH:-./config.yaml}:/CLIProxyAPI/config.yaml - - ${CLI_PROXY_AUTH_PATH:-./auths}:/root/.cli-proxy-api + - ${CLI_PROXY_AUTH_PATH:-./auths}:/CLIProxyAPI/auths - ${CLI_PROXY_LOG_PATH:-./logs}:/CLIProxyAPI/logs - ${CLI_PROXY_PLUGIN_PATH:-./plugins}:/CLIProxyAPI/plugins restart: unless-stopped diff --git a/examples/plugin/kiro/go/executor.go b/examples/plugin/kiro/go/executor.go index 6b1dcdbe8c0..bfc90d52ad5 100644 --- a/examples/plugin/kiro/go/executor.go +++ b/examples/plugin/kiro/go/executor.go @@ -255,7 +255,7 @@ func prepareUpstreamRequest(request pluginapi.ExecutorRequest) (upstreamRequest, "User-Agent": []string{userAgent}, "X-Amz-Target": []string{kiroAmzTarget}, "X-Amz-User-Agent": []string{amzUserAgent}, - "X-Amzn-Codewhisperer-Optout": []string{"false"}, + "X-Amzn-Codewhisperer-Optout": []string{"true"}, } return upstreamRequest{URL: endpoint, Headers: headers, Body: body}, payload, nil } diff --git a/examples/plugin/kiro/go/kiro_test.go b/examples/plugin/kiro/go/kiro_test.go index 5593d5177ae..6bf6ea4bac9 100644 --- a/examples/plugin/kiro/go/kiro_test.go +++ b/examples/plugin/kiro/go/kiro_test.go @@ -89,6 +89,13 @@ func TestModelsForAuthDiscoversPaginatesAndCaches(t *testing.T) { if request.Headers.Get("Authorization") != "Bearer test-discovery-key" || request.Headers.Get("TokenType") != "API_KEY" { t.Fatalf("discovery headers = %#v", request.Headers) } + machine := machineID("test-discovery-key") + if !strings.Contains(request.Headers.Get("User-Agent"), "api/codewhispererruntime#1.0.0") || + !strings.HasSuffix(request.Headers.Get("User-Agent"), "-"+machine) || + !strings.HasSuffix(request.Headers.Get("X-Amz-User-Agent"), "-"+machine) || + request.Headers.Get("X-Amzn-Codewhisperer-Optout") != "true" { + t.Fatalf("Kiro identity headers = %#v", request.Headers) + } parsed, errParse := url.Parse(request.URL) if errParse != nil { t.Fatal(errParse) @@ -147,6 +154,29 @@ func TestModelsForAuthDiscoversPaginatesAndCaches(t *testing.T) { } } +func TestModelDiscoveryURLUsesKiroGoRegionalHosts(t *testing.T) { + tests := []struct { + region string + host string + }{ + {region: "us-east-1", host: "codewhisperer.us-east-1.amazonaws.com"}, + {region: "eu-central-1", host: "q.eu-central-1.amazonaws.com"}, + } + for _, test := range tests { + rawURL, errURL := modelDiscoveryURL(test.region, "") + if errURL != nil { + t.Fatalf("modelDiscoveryURL(%q) error = %v", test.region, errURL) + } + parsed, errParse := url.Parse(rawURL) + if errParse != nil { + t.Fatal(errParse) + } + if parsed.Host != test.host { + t.Fatalf("modelDiscoveryURL(%q) host = %q, want %q", test.region, parsed.Host, test.host) + } + } +} + func TestModelsForAuthAppliesConfiguredAllowList(t *testing.T) { configureKiroTest(t, "models:\n - claude-sonnet-4.5\n") t.Setenv("KIRO_TEST_KEY", "test-allow-list-key") diff --git a/examples/plugin/kiro/go/model_discovery.go b/examples/plugin/kiro/go/model_discovery.go index 8ef0e927d3b..aaee4e49003 100644 --- a/examples/plugin/kiro/go/model_discovery.go +++ b/examples/plugin/kiro/go/model_discovery.go @@ -19,6 +19,7 @@ const ( modelCacheTTL = 5 * time.Minute modelDiscoveryPageSize = 50 modelDiscoveryMaxPages = 10 + kiroRuntimeSDKVersion = "1.0.0" ) type kiroModelListResponse struct { @@ -111,12 +112,7 @@ func fetchAvailableModels(hostCallbackID string, credential kiroCredential, key HostCallbackID: hostCallbackID, Method: http.MethodGet, URL: endpoint, - Headers: http.Header{ - "Accept": []string{"application/json"}, - "Authorization": []string{"Bearer " + key}, - "Tokentype": []string{"API_KEY"}, - "X-Amzn-Codewhisperer-Optout": []string{"false"}, - }, + Headers: modelDiscoveryHeaders(key), }) if errCall != nil { return nil, fmt.Errorf("call model-list endpoint: %w", errCall) @@ -170,7 +166,34 @@ func modelDiscoveryURL(region, nextToken string) (string, error) { if nextToken != "" { query.Set("nextToken", nextToken) } - return "https://codewhisperer." + region + ".amazonaws.com/ListAvailableModels?" + query.Encode(), nil + host := "codewhisperer.us-east-1.amazonaws.com" + if region != "us-east-1" { + host = "q." + region + ".amazonaws.com" + } + return "https://" + host + "/ListAvailableModels?" + query.Encode(), nil +} + +func modelDiscoveryHeaders(key string) http.Header { + cfg := loadedConfig() + machine := machineID(key) + userAgent := fmt.Sprintf( + "aws-sdk-js/%s ua/2.1 os/%s lang/js md/nodejs#%s api/codewhispererruntime#%s m/N,E KiroIDE-%s-%s", + kiroRuntimeSDKVersion, + cfg.SystemVersion, + cfg.NodeVersion, + kiroRuntimeSDKVersion, + cfg.KiroVersion, + machine, + ) + amzUserAgent := fmt.Sprintf("aws-sdk-js/%s KiroIDE-%s-%s", kiroRuntimeSDKVersion, cfg.KiroVersion, machine) + return http.Header{ + "Accept": []string{"application/json"}, + "Authorization": []string{"Bearer " + key}, + "Tokentype": []string{"API_KEY"}, + "User-Agent": []string{userAgent}, + "X-Amz-User-Agent": []string{amzUserAgent}, + "X-Amzn-Codewhisperer-Optout": []string{"true"}, + } } func discoveredModelInfo(model kiroAvailableModel) pluginapi.ModelInfo { From 3414685c36fd58f5b0524a3e7a3d2a1395f237df Mon Sep 17 00:00:00 2001 From: ianwangye Date: Mon, 10 Aug 2026 16:14:51 +0800 Subject: [PATCH 20/20] fix(compose): restore default auth directory mount --- docker-compose.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index b92446955dc..e500dd157ea 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -26,7 +26,7 @@ services: - "11451:11451" volumes: - ${CLI_PROXY_CONFIG_PATH:-./config.yaml}:/CLIProxyAPI/config.yaml - - ${CLI_PROXY_AUTH_PATH:-./auths}:/CLIProxyAPI/auths + - ${CLI_PROXY_AUTH_PATH:-./auths}:/root/.cli-proxy-api - ${CLI_PROXY_LOG_PATH:-./logs}:/CLIProxyAPI/logs - ${CLI_PROXY_PLUGIN_PATH:-./plugins}:/CLIProxyAPI/plugins restart: unless-stopped