diff --git a/internal/gateway/kiro.go b/internal/gateway/kiro.go index 49f1d8ff..1ccbe3cb 100644 --- a/internal/gateway/kiro.go +++ b/internal/gateway/kiro.go @@ -88,7 +88,7 @@ func (s *Server) askKiro(p provider.Provider, model string) round { events := make(chan Event, 16) go func() { defer res.Body.Close() - decodeKiro(ctx, res.Body, events, model, window, thinking) + decodeKiro(ctx, res.Body, events, model, window, thinking, kiroNames(req)) }() return events, 0, "" } @@ -245,6 +245,47 @@ func kiroToolID(id string) string { return "t_" + base64.RawURLEncoding.EncodeToString(sum[:])[:32] } +var ( + kiroToolNameRe = regexp.MustCompile(`^[a-zA-Z0-9_-]{1,64}$`) + kiroToolNameBad = regexp.MustCompile(`[^a-zA-Z0-9_-]`) +) + +// kiroToolName is a tool's name as Kiro takes one: a longer one (an MCP +// server's, prefixed) is rejected once the history holds a call, so it is +// cut and told apart by a hash, the same each time; kiroNames maps it back. +func kiroToolName(name string) string { + if kiroToolNameRe.MatchString(name) { + return name + } + sum := sha256.Sum256([]byte(name)) + short := kiroToolNameBad.ReplaceAllString(name, "_") + if len(short) > 55 { + short = short[:55] + } + return short + "_" + fmt.Sprintf("%x", sum[:4]) +} + +// kiroNames is the caller's name for each name kiroToolName changed. +func kiroNames(r *Request) map[string]string { + names := map[string]string{} + add := func(n string) { + if k := kiroToolName(n); k != n { + names[k] = n + } + } + for _, t := range r.Tools { + add(t.Name) + } + for _, m := range r.Messages { + for _, p := range m.Parts { + if p.Kind == ToolCall { + add(p.Name) + } + } + } + return names +} + // buildKiro is the body of a generateAssistantResponse call. func buildKiro(r *Request, model, profile string, thinking bool) []byte { var entries []kiroEntry @@ -264,7 +305,7 @@ func buildKiro(r *Request, model, profile string, thinking bool) []byte { if !bytes.HasPrefix(bytes.TrimSpace(args), []byte("{")) { args = json.RawMessage("{}") } - a.ToolUses = append(a.ToolUses, kiroToolUse{Name: p.Name, ToolUseID: kiroToolID(p.ID), Input: args}) + a.ToolUses = append(a.ToolUses, kiroToolUse{Name: kiroToolName(p.Name), ToolUseID: kiroToolID(p.ID), Input: args}) } } a.Content = strings.Join(texts, "\n\n") @@ -309,6 +350,15 @@ func buildKiro(r *Request, model, profile string, thinking bool) []byte { status = "error" } results = append(results, kiroToolResult{ToolUseID: kiroToolID(p.CallID), Status: status, Content: []kiroText{{Text: out}}}) + // a tool's images (a screenshot, a file read) go with the + // message, as Kiro's tool results carry text alone + for _, im := range p.Images { + if im.Data != "" { + img := kiroImage{Format: kiroImageFormat(im.MediaType)} + img.Source.Bytes = im.Data + u.Images = append(u.Images, img) + } + } } } u.Content = strings.Join(texts, "\n\n") @@ -414,7 +464,7 @@ func buildKiro(r *Request, model, profile string, thinking bool) []byte { offered := map[string]bool{} for _, t := range r.Tools { kt := kiroTool{} - kt.Spec.Name, kt.Spec.Description = t.Name, t.Description + kt.Spec.Name, kt.Spec.Description = kiroToolName(t.Name), t.Description if kt.Spec.Description == "" { kt.Spec.Description = t.Name } @@ -422,7 +472,7 @@ func buildKiro(r *Request, model, profile string, thinking bool) []byte { if len(bytes.TrimSpace(t.Schema)) == 0 || string(bytes.TrimSpace(t.Schema)) == "null" { kt.Spec.InputSchema.JSON = json.RawMessage(`{"type":"object","properties":{}}`) } - offered[t.Name] = true + offered[kt.Spec.Name] = true tools = append(tools, kt) } for _, e := range entries { @@ -576,7 +626,7 @@ func readKiroFrame(r *bufio.Reader) (kiroFrame, error) { } // decodeKiro turns Kiro's reply into events. -func decodeKiro(ctx context.Context, body io.Reader, out chan<- Event, model string, window int, thinking bool) { +func decodeKiro(ctx context.Context, body io.Reader, out chan<- Event, model string, window int, thinking bool, names map[string]string) { defer close(out) send := func(ev Event) bool { select { @@ -589,7 +639,7 @@ func decodeKiro(ctx context.Context, body io.Reader, out chan<- Event, model str if !send(Event{Kind: KStart, MsgID: "msg_" + randomToken()[:24], Model: model}) { return } - d := kiroDecoder{thinking: thinking} + d := kiroDecoder{thinking: thinking, names: names} br := bufio.NewReaderSize(body, 64<<10) var failed string for { @@ -647,6 +697,7 @@ func decodeKiro(ctx context.Context, body io.Reader, out chan<- Event, model str type kiroDecoder struct { thinking bool think kiroThinkParser + names map[string]string // the caller's name for a tool kiroToolName cut tool string // the id of the call being streamed tools int @@ -694,6 +745,9 @@ func (d *kiroDecoder) frame(f kiroFrame) (evs []Event, failed string) { } case "toolUseEvent": id, name := str("toolUseId"), str("name") + if n, ok := d.names[name]; ok { + name = n + } if id != "" && id != d.tool { // the text before a call is all said evs = append(evs, d.think.end()...) diff --git a/internal/gateway/kiro_test.go b/internal/gateway/kiro_test.go index c7234397..785d5a9f 100644 --- a/internal/gateway/kiro_test.go +++ b/internal/gateway/kiro_test.go @@ -43,7 +43,7 @@ func kiroEvent(kind, payload string) []byte { func kiroEvents(t *testing.T, thinking bool, frames ...[]byte) []Event { t.Helper() out := make(chan Event, 256) - decodeKiro(context.Background(), bytes.NewReader(bytes.Join(frames, nil)), out, "claude-sonnet-4.5", 200000, thinking) + decodeKiro(context.Background(), bytes.NewReader(bytes.Join(frames, nil)), out, "claude-sonnet-4.5", 200000, thinking, nil) var evs []Event for ev := range out { evs = append(evs, ev) @@ -262,6 +262,45 @@ func TestBuildKiro(t *testing.T) { } } +func TestBuildKiroToolImagesAndLongNames(t *testing.T) { + long := "mcp__plugin_vercel_vercel__create_sandboxes_sessions_by_session_id_snapshot_v2" + r := &Request{ + Tools: []Tool{{Name: long, Description: "Snapshot", Schema: json.RawMessage(`{"type":"object"}`)}, {Name: "Read", Schema: json.RawMessage(`{"type":"object"}`)}}, + Messages: []Message{ + {Role: "user", Parts: []Part{{Kind: Text, Text: "look"}}}, + {Role: "assistant", Parts: []Part{{Kind: ToolCall, ID: "c1", Name: "Read", Args: json.RawMessage(`{}`)}, {Kind: ToolCall, ID: "c2", Name: long, Args: json.RawMessage(`{}`)}}}, + {Role: "user", Parts: []Part{ + {Kind: ToolResult, CallID: "c1", Images: []Part{{Kind: Image, MediaType: "image/png", Data: "PNG"}}}, + {Kind: ToolResult, CallID: "c2", Text: "ok"}, + }}, + }, + } + b := buildKiroBody(t, r, false) + cur := b.ConversationState.CurrentMessage.UserInputMessage + // the image a tool read goes with the message answering it + if len(cur.Images) != 1 || cur.Images[0].Format != "png" || cur.Images[0].Source.Bytes != "PNG" { + t.Fatalf("images = %+v", cur.Images) + } + short := kiroToolName(long) + if len(short) > 64 || !kiroToolNameRe.MatchString(short) || short != kiroToolName(long) { + t.Fatalf("short = %q", short) + } + if calls := b.ConversationState.History[1].Asst.ToolUses; calls[0].Name != "Read" || calls[1].Name != short { + t.Fatalf("calls = %+v", calls) + } + if tools := cur.Context.Tools; len(tools) != 2 || tools[0].Spec.Name != short || tools[1].Spec.Name != "Read" { + t.Fatalf("tools = %+v", tools) + } + // and a call Kiro makes by the short name comes back by the caller's + out := make(chan Event, 16) + decodeKiro(context.Background(), bytes.NewReader(kiroEvent("toolUseEvent", `{"toolUseId":"t1","name":"`+short+`","input":"{}"}`)), out, "claude-sonnet-4.5", 200000, false, kiroNames(r)) + for ev := range out { + if ev.Kind == KToolStart && ev.Name != long { + t.Fatalf("name = %q", ev.Name) + } + } +} + func TestBuildKiroThinkingAndOneMessage(t *testing.T) { r := &Request{Effort: "high", Messages: []Message{{Role: "user", Parts: []Part{{Kind: Text, Text: "why?"}}}}} b := buildKiroBody(t, r, true)