From 49353485e5611d449bcc001e3c4b7843ea2d509e Mon Sep 17 00:00:00 2001
From: Victor Vargas <victorhvargas@outlook.com>
Date: Thu, 16 Jul 2026 22:24:03 -0700
Subject: [PATCH] feat(agent): thread the new turn through Loop as llm.Message
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Run/RunStream took the new turn as a bare string, which had nowhere
to carry ContentPart attachments. Both now take an llm.Message
(Role is forced to RoleUser regardless of what the caller sets), so a
caller building a multimodal turn just fills in Content/Parts on it
instead of the loop needing a second, parallel parameter.
subagent.go and every test call site are updated to wrap their string
prompt as llm.Message{Role: llm.RoleUser, Content: ...} — SubAgent.Run
itself is untouched, it still takes a plain task string.
---
 pkg/agent/integration_test.go       |  4 ++--
 pkg/agent/loop.go                   | 15 +++++++++------
 pkg/agent/loop_test.go              | 30 ++++++++++++++---------------
 pkg/agent/subagent.go               |  2 +-
 pkg/agent/thinking_budget_test.go   |  4 ++--
 pkg/agent/unparsed_toolcall_test.go |  6 +++---
 6 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/pkg/agent/integration_test.go b/pkg/agent/integration_test.go
index 2a2f052..68532bf 100644
--- a/pkg/agent/integration_test.go
+++ b/pkg/agent/integration_test.go
@@ -103,7 +103,7 @@ func TestIntegration_AgentLoop_Generate(t *testing.T) {
 		MaxIters: 3,
 	})
 
-	resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "What is 20+22? Use the add_numbers tool."})
 	if err != nil {
 		t.Fatalf("run failed: %v", err)
 	}
@@ -130,7 +130,7 @@ func TestIntegration_AgentLoop_Stream(t *testing.T) {
 		MaxIters: 3,
 	})
 
-	stream := loop.RunStream(context.Background(), "Say something interesting.")
+	stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Say something interesting."})
 
 	var chunks []llm.StreamChunk
 	for chunk, err := range stream {
diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 6932298..de38251 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -77,8 +77,9 @@ func New(cfg Config) *Loop {
 
 // Run executes the agent loop and returns the final response.
 // Optional history messages are appended after the system prompt and before
-// the new user input.
-func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (Response, error) {
+// the new user input. input's Role is overwritten to RoleUser regardless of
+// what the caller sets, so callers only need to fill in Content/Parts.
+func (l *Loop) Run(ctx context.Context, input llm.Message, history ...llm.Message) (Response, error) {
 	start := time.Now()
 
 	messages := l.buildInitialMessages(input, history)
@@ -174,8 +175,9 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
 
 // RunStream executes the agent loop with streaming output.
 // Optional history messages are appended after the system prompt and before
-// the new user input.
-func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
+// the new user input. input's Role is overwritten to RoleUser regardless of
+// what the caller sets, so callers only need to fill in Content/Parts.
+func (l *Loop) RunStream(ctx context.Context, input llm.Message, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
 	return func(yield func(llm.StreamChunk, error) bool) {
 		messages := l.buildInitialMessages(input, history)
 		// Same as Run: the schemas are identical on every iteration.
@@ -362,12 +364,13 @@ func containsUnparsedToolCall(s string) bool {
 	return strings.Contains(s, "<tool_call") || strings.Contains(s, "<function=")
 }
 
-func (l *Loop) buildInitialMessages(input string, history []llm.Message) []llm.Message {
+func (l *Loop) buildInitialMessages(input llm.Message, history []llm.Message) []llm.Message {
 	systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, l.cfg.AgentsMD)
 	messages := make([]llm.Message, 0, len(history)+2)
 	messages = append(messages, llm.Message{Role: llm.RoleSystem, Content: systemPrompt})
 	messages = append(messages, history...)
-	messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
+	input.Role = llm.RoleUser
+	messages = append(messages, input)
 	return messages
 }
 
diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go
index 2135b93..d26af7e 100644
--- a/pkg/agent/loop_test.go
+++ b/pkg/agent/loop_test.go
@@ -102,7 +102,7 @@ func TestRun_NoToolCalls(t *testing.T) {
 		Tools:   tools.NewRegistry(),
 	})
 
-	resp, err := loop.Run(context.Background(), "Hello")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Hello"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -133,7 +133,7 @@ func TestRun_IncludesAgentsMD(t *testing.T) {
 		AgentsMD: "Never edit go.mod directly.",
 	})
 
-	if _, err := loop.Run(context.Background(), "Hello"); err != nil {
+	if _, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Hello"}); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 	if !strings.Contains(capturedSystemPrompt, "Never edit go.mod directly.") {
@@ -176,7 +176,7 @@ func TestRun_ToolCalls(t *testing.T) {
 		Tools:   registry,
 	})
 
-	resp, err := loop.Run(context.Background(), "Say hi")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Say hi"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -233,7 +233,7 @@ func TestRun_ToolCalls_RecordsAssistantTurnAndToolCallID(t *testing.T) {
 		Tools:   registry,
 	})
 
-	if _, err := loop.Run(context.Background(), "Say hi"); err != nil {
+	if _, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Say hi"}); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -315,7 +315,7 @@ func TestRun_Stream_ToolCalls_RecordsAssistantTurnAndToolCallID(t *testing.T) {
 		Tools:   registry,
 	})
 
-	for _, err := range loop.RunStream(context.Background(), "test") {
+	for _, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"}) {
 		if err != nil {
 			t.Fatalf("unexpected stream error: %v", err)
 		}
@@ -379,7 +379,7 @@ func TestRun_MaxIterations(t *testing.T) {
 		MaxIters: 3,
 	})
 
-	_, err := loop.Run(context.Background(), "test")
+	_, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err == nil {
 		t.Fatal("expected error, got nil")
 	}
@@ -408,7 +408,7 @@ func TestRun_ToolNotFound(t *testing.T) {
 		Tools:   tools.NewRegistry(),
 	})
 
-	resp, err := loop.Run(context.Background(), "test")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -451,7 +451,7 @@ func TestRun_ApprovalDenied(t *testing.T) {
 		},
 	})
 
-	resp, err := loop.Run(context.Background(), "test")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -496,7 +496,7 @@ func TestRun_SandboxViolation(t *testing.T) {
 		},
 	})
 
-	resp, err := loop.Run(context.Background(), "test")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -540,7 +540,7 @@ func TestRun_OnIterationHook(t *testing.T) {
 		},
 	})
 
-	_, err := loop.Run(context.Background(), "test")
+	_, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -570,7 +570,7 @@ func TestRun_Stream_NoToolCalls(t *testing.T) {
 	})
 
 	var chunks []llm.StreamChunk
-	stream := loop.RunStream(context.Background(), "test")
+	stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	for chunk, err := range stream {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
@@ -608,7 +608,7 @@ func TestRun_Stream_ForwardsTrailingUsageOnlyChunk(t *testing.T) {
 	})
 
 	var chunks []llm.StreamChunk
-	stream := loop.RunStream(context.Background(), "test")
+	stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	for chunk, err := range stream {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
@@ -661,7 +661,7 @@ func TestRun_Stream_WithToolCalls(t *testing.T) {
 	})
 
 	var chunks []llm.StreamChunk
-	stream := loop.RunStream(context.Background(), "test")
+	stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	for chunk, err := range stream {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
@@ -713,7 +713,7 @@ func TestRun_Stream_MaxIterations(t *testing.T) {
 	})
 
 	var chunks []llm.StreamChunk
-	stream := loop.RunStream(context.Background(), "test")
+	stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "test"})
 	for chunk, err := range stream {
 		if err != nil {
 			// expect max iterations error
@@ -748,7 +748,7 @@ func TestRun_Timeout(t *testing.T) {
 		Tools:   tools.NewRegistry(),
 	})
 
-	_, err := loop.Run(ctx, "test")
+	_, err := loop.Run(ctx, llm.Message{Role: llm.RoleUser, Content: "test"})
 	if err == nil {
 		t.Fatal("expected timeout error, got nil")
 	}
diff --git a/pkg/agent/subagent.go b/pkg/agent/subagent.go
index e30c57f..cafcc7a 100644
--- a/pkg/agent/subagent.go
+++ b/pkg/agent/subagent.go
@@ -33,7 +33,7 @@ func (s SubAgent) Run(ctx context.Context, llmClient llm.LLMClient, agentsMD str
 	if cfg.MaxIters == 0 {
 		cfg.MaxIters = DefaultMaxIterations
 	}
-	return New(cfg).Run(ctx, task)
+	return New(cfg).Run(ctx, llm.Message{Role: llm.RoleUser, Content: task})
 }
 
 // SubAgentRegistry looks up SubAgents by name for the delegate tool.
diff --git a/pkg/agent/thinking_budget_test.go b/pkg/agent/thinking_budget_test.go
index 054f30b..ca77bc0 100644
--- a/pkg/agent/thinking_budget_test.go
+++ b/pkg/agent/thinking_budget_test.go
@@ -28,7 +28,7 @@ func TestRunStream_RecoversFromThinkingBudgetCut(t *testing.T) {
 	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
 
 	var final strings.Builder
-	for chunk, err := range loop.RunStream(context.Background(), "arregla x.py") {
+	for chunk, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "arregla x.py"}) {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
@@ -62,7 +62,7 @@ func TestRunStream_ThinkingBudgetNudgeGivesUpAfterLimit(t *testing.T) {
 
 	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
 
-	for _, err := range loop.RunStream(context.Background(), "haz algo") {
+	for _, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "haz algo"}) {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
diff --git a/pkg/agent/unparsed_toolcall_test.go b/pkg/agent/unparsed_toolcall_test.go
index 1ea00b9..83b5e3c 100644
--- a/pkg/agent/unparsed_toolcall_test.go
+++ b/pkg/agent/unparsed_toolcall_test.go
@@ -97,7 +97,7 @@ func TestRunStream_RecoversFromUnparsedToolCall(t *testing.T) {
 	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
 
 	var final strings.Builder
-	for chunk, err := range loop.RunStream(context.Background(), "arregla x.py") {
+	for chunk, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "arregla x.py"}) {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
@@ -133,7 +133,7 @@ func TestRunStream_NudgeGivesUpAfterLimit(t *testing.T) {
 	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
 
 	rounds := 0
-	for _, err := range loop.RunStream(context.Background(), "haz algo") {
+	for _, err := range loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "haz algo"}) {
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
@@ -158,7 +158,7 @@ func TestRun_RecoversFromUnparsedToolCall(t *testing.T) {
 	}}
 
 	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})
-	resp, err := loop.Run(context.Background(), "arregla x.py")
+	resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "arregla x.py"})
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}