Why: Steering hooks existed on the loop but nothing owned the queues, abort, or a model change during a run. Evidence: go test -count=1 ./runtime/ Impact: Run still emits the existing event names. message_start, tool_execution_update, and agent_settled are added beside them. Refs: feature/pi-headless N3 Co-authored-by: Cursor <cursoragent@cursor.com>
89 lines
2.6 KiB
Go
89 lines
2.6 KiB
Go
package runtime
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"testing"
|
|
|
|
"agent-native/runtime/llm"
|
|
)
|
|
|
|
func TestAgentSessionSteerFollowAndAbort(t *testing.T) {
|
|
var mu sync.Mutex
|
|
step := 0
|
|
stream := llm.StreamFunc(func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
|
ch := make(chan llm.Event, 2)
|
|
go func() {
|
|
defer close(ch)
|
|
mu.Lock()
|
|
n := step
|
|
step++
|
|
mu.Unlock()
|
|
if opts.ThinkingLevel != "low" && n == 0 {
|
|
t.Errorf("thinking level not forwarded: %s", opts.ThinkingLevel)
|
|
}
|
|
if n == 0 {
|
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
|
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
|
ToolCalls: []llm.ToolCall{{ID: "1", Name: "missing", Arguments: "{}"}},
|
|
}}
|
|
return
|
|
}
|
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
|
Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd,
|
|
}}
|
|
}()
|
|
return ch
|
|
})
|
|
var types []string
|
|
sess := NewAgentSession(Config{
|
|
Model: "m",
|
|
ThinkingLevel: "low",
|
|
Stream: stream,
|
|
MaxTurns: 6,
|
|
OnEvent: func(ev Event) {
|
|
types = append(types, ev.Type)
|
|
},
|
|
})
|
|
sess.SetQueueMode(DeliverOne, DeliverAll)
|
|
sess.Steer(llm.Message{Role: llm.RoleUser, Content: "steer-me"})
|
|
sess.FollowUp(llm.Message{Role: llm.RoleUser, Content: "later"})
|
|
res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "go"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if res.StopReason != llm.StopEnd {
|
|
t.Fatalf("stop %s err %v", res.StopReason, res.Err)
|
|
}
|
|
found := map[string]bool{}
|
|
for _, typ := range types {
|
|
found[typ] = true
|
|
}
|
|
for _, want := range []string{"message_start", "queue_update", "tool_execution_update", "agent_settled"} {
|
|
if !found[want] {
|
|
t.Fatalf("missing %s in %v", want, types)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestClearQueueAndSetModel(t *testing.T) {
|
|
sess := NewAgentSession(Config{Model: "old", Stream: func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
|
ch := make(chan llm.Event, 1)
|
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: model, StopReason: llm.StopEnd}}
|
|
close(ch)
|
|
return ch
|
|
}})
|
|
sess.Steer(llm.Message{Role: llm.RoleUser, Content: "drop"})
|
|
sess.ClearQueue()
|
|
if err := sess.SetModel("new-model", "high"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
res, err := sess.Prompt(context.Background(), llm.Message{Role: llm.RoleUser, Content: "hi"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if res.Text != "new-model" {
|
|
t.Fatalf("model switch did not apply: %q", res.Text)
|
|
}
|
|
}
|