diff --git a/pkg/agent/core_test.go b/pkg/agent/core_test.go index 86c27c6..2640e48 100644 --- a/pkg/agent/core_test.go +++ b/pkg/agent/core_test.go @@ -677,7 +677,11 @@ func TestCommand_I_SetsInject(t *testing.T) { return textStream("ok"), nil } c := newCore(t, be, nil) - go c.Submit(context.Background(), "hello") + submitDone := make(chan struct{}) + go func() { + c.Submit(context.Background(), "hello") + close(submitDone) + }() waitState(t, c, "thinking") c.handleCommand(context.Background(), "/i my inject") p := c.pendingInject.Load() @@ -685,7 +689,7 @@ func TestCommand_I_SetsInject(t *testing.T) { t.Errorf("pendingInject = %v; want %q", p, "my inject") } close(done) - waitState(t, c, "idle") + <-submitDone } func TestCommand_I_SubmitsWhenIdle(t *testing.T) { diff --git a/pkg/backend/codewhisperer.go b/pkg/backend/codewhisperer.go index 6f43a9b..355518c 100644 --- a/pkg/backend/codewhisperer.go +++ b/pkg/backend/codewhisperer.go @@ -282,8 +282,8 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN if err != nil { return nil, err } - toolResults := encodeKiroToolResults(messages[cutoff+1:]) - current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults)} + toolResults, images := encodeKiroToolResults(messages[cutoff+1:]) + current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults, images)} default: return nil, fmt.Errorf("last message role %q is not supported (expected user, system, or tool)", lastRole) @@ -326,8 +326,8 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) { for i+1 < len(messages) && strings.EqualFold(messages[i+1].Role, "tool") { i++ } - results := encodeKiroToolResults(messages[start : i+1]) - out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results)}) + results, imgs := encodeKiroToolResults(messages[start : i+1]) + out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results, imgs)}) } } return out, nil @@ -335,7 +335,7 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) { func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResults []kiroToolResult) *kiroUserInputMessage { msg := &kiroUserInputMessage{ - Content: m.Content, + Content: kiroEncodeContent(m), Origin: kiroDefaultOrigin, ModelID: modelID, } @@ -352,11 +352,12 @@ func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResu return msg } -func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult) *kiroUserInputMessage { +func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult, images []kiroImageBlock) *kiroUserInputMessage { msg := &kiroUserInputMessage{ - Content: "", + Content: kiroEncodeStringContent(""), Origin: kiroDefaultOrigin, ModelID: modelID, + Images: images, } ctx := &kiroUserInputContext{} if len(tools) > 0 { @@ -412,19 +413,49 @@ func encodeKiroToolDefinitions(tools []Tool) []kiroTool { return out } -func encodeKiroToolResults(messages []Message) []kiroToolResult { +func encodeKiroToolResults(messages []Message) ([]kiroToolResult, []kiroImageBlock) { results := make([]kiroToolResult, 0, len(messages)) + var images []kiroImageBlock for _, m := range messages { if !strings.EqualFold(m.Role, "tool") { continue } toolUseID := strings.TrimSpace(m.ToolCallID) + content, imgs := encodeKiroToolResultBlocks(m) results = append(results, kiroToolResult{ ToolUseID: toolUseID, - Content: []kiroToolResultContent{encodeKiroToolResultContent(m.Content)}, + Content: content, }) + images = append(images, imgs...) } - return results + return results, images +} + +func encodeKiroToolResultBlocks(m Message) ([]kiroToolResultContent, []kiroImageBlock) { + if len(m.ContentBlocks) == 0 { + return []kiroToolResultContent{encodeKiroToolResultContent(m.Content)}, nil + } + out := make([]kiroToolResultContent, 0, len(m.ContentBlocks)) + var images []kiroImageBlock + for _, cb := range m.ContentBlocks { + switch cb.Type { + case "text": + out = append(out, encodeKiroToolResultContent(cb.Text)) + case "image": + if cb.ImageSource == nil { + continue + } + format := strings.TrimPrefix(cb.ImageSource.MediaType, "image/") + images = append(images, kiroImageBlock{ + Format: format, + Source: kiroImageSource{Bytes: cb.ImageSource.Data}, + }) + } + } + if len(out) == 0 { + out = []kiroToolResultContent{{Text: ""}} + } + return out, images } func encodeKiroToolResultContent(content string) kiroToolResultContent { @@ -435,6 +466,58 @@ func encodeKiroToolResultContent(content string) kiroToolResultContent { return kiroToolResultContent{Text: content} } +// kiroEncodeContent returns a json.RawMessage that is either a JSON string +// (plain text) or a JSON array of Bedrock Converse content blocks (when images +// are present). +func kiroEncodeContent(m Message) json.RawMessage { + if len(m.ContentBlocks) == 0 { + return kiroEncodeStringContent(m.Content) + } + hasImage := false + for _, cb := range m.ContentBlocks { + if cb.Type == "image" && cb.ImageSource != nil { + hasImage = true + break + } + } + if !hasImage { + // Text-only blocks: concatenate and send as string. + var sb strings.Builder + for _, cb := range m.ContentBlocks { + if cb.Type == "text" { + sb.WriteString(cb.Text) + } + } + return kiroEncodeStringContent(sb.String()) + } + // Mixed content: encode as Bedrock Converse content block array. + blocks := make([]kiroContentBlock, 0, len(m.ContentBlocks)) + for _, cb := range m.ContentBlocks { + switch cb.Type { + case "text": + blocks = append(blocks, kiroContentBlock{Text: cb.Text}) + case "image": + if cb.ImageSource == nil { + continue + } + format := strings.TrimPrefix(cb.ImageSource.MediaType, "image/") + blocks = append(blocks, kiroContentBlock{ + Image: &kiroImageBlock{ + Format: format, + Source: kiroImageSource{Bytes: cb.ImageSource.Data}, + }, + }) + } + } + data, _ := json.Marshal(blocks) + return json.RawMessage(data) +} + +func kiroEncodeStringContent(s string) json.RawMessage { + data, _ := json.Marshal(s) + return json.RawMessage(data) +} + // ── Auth sources ────────────────────────────────────────────────────────────── type kiroAuthSource interface { diff --git a/pkg/backend/codewhisperer_internal.go b/pkg/backend/codewhisperer_internal.go index 841d4a0..7afbaa3 100644 --- a/pkg/backend/codewhisperer_internal.go +++ b/pkg/backend/codewhisperer_internal.go @@ -73,12 +73,27 @@ type kiroChatMessage struct { } type kiroUserInputMessage struct { - Content string `json:"content"` + Content json.RawMessage `json:"content"` UserInputMessageContext *kiroUserInputContext `json:"userInputMessageContext,omitempty"` Origin string `json:"origin,omitempty"` + Images []kiroImageBlock `json:"images,omitempty"` ModelID string `json:"modelId,omitempty"` } +type kiroContentBlock struct { + Text string `json:"text,omitempty"` + Image *kiroImageBlock `json:"image,omitempty"` +} + +type kiroImageBlock struct { + Format string `json:"format"` + Source kiroImageSource `json:"source"` +} + +type kiroImageSource struct { + Bytes string `json:"bytes"` +} + type kiroUserInputContext struct { ToolResults []kiroToolResult `json:"toolResults,omitempty"` Tools []kiroTool `json:"tools,omitempty"` @@ -117,8 +132,8 @@ type kiroToolResult struct { } type kiroToolResultContent struct { - Text string `json:"text,omitempty"` - JSON json.RawMessage `json:"json,omitempty"` + Text string `json:"text,omitempty"` + JSON json.RawMessage `json:"json,omitempty"` } // stream event types diff --git a/pkg/tools/execute/server.go b/pkg/tools/execute/server.go index 19c809b..92b5197 100644 --- a/pkg/tools/execute/server.go +++ b/pkg/tools/execute/server.go @@ -315,6 +315,13 @@ func (e *Server) CallTool(ctx context.Context, tool string, args json.RawMessage "content": []map[string]string{{"type": "text", "text": err.Error()}}, }) } + // If the tool output is already structured content-block JSON, pass it through. + var structured map[string]json.RawMessage + if json.Unmarshal([]byte(result), &structured) == nil { + if _, ok := structured["content"]; ok { + return []byte(result), nil + } + } return json.Marshal(map[string]any{ "content": []map[string]string{{"type": "text", "text": result}}, })