codewhisperer: image content block support in user messages and tool results

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) {

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, 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},
 			})
 		}
-	return results
+	}
+	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 {

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"`

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}},
 	})