feat(agent): thread the new turn through Loop as llm.Message

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.
This commit is contained in:
Victor Hugo Vargas Servin 2026-07-16 22:24:03 -07:00
parent 2f6f5fab1c
commit 49353485e5
6 changed files with 32 additions and 29 deletions

View file

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

View file

@ -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
}

View file

@ -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")
}

View file

@ -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.

View file

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

View file

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