feat: complete v0.3.1 agent workflow
This commit is contained in:
loveuer
@@ -3,8 +3,13 @@ package agent
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"agentu/internal/tools"
|
||||
@@ -177,6 +182,173 @@ func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, e
|
||||
return emit(llm.StreamEvent{Content: p.text})
|
||||
}
|
||||
|
||||
type btwProvider struct {
|
||||
req llm.ChatRequest
|
||||
err error
|
||||
}
|
||||
|
||||
func (p *btwProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.req = req
|
||||
if p.err != nil {
|
||||
return p.err
|
||||
}
|
||||
return emit(llm.StreamEvent{Content: "side answer", Usage: &llm.Usage{TotalTokens: 999}})
|
||||
}
|
||||
|
||||
func TestBTWUsesContextWithoutMutatingConversation(t *testing.T) {
|
||||
provider := &btwProvider{}
|
||||
a := New(Options{Provider: provider, Model: "test", Thinking: "high", ThinkingEnabled: true})
|
||||
a.SetMessages([]llm.Message{
|
||||
{Role: llm.RoleUser, Content: "main question"},
|
||||
{Role: llm.RoleAssistant, Content: "main answer"},
|
||||
})
|
||||
a.SetLastUsage(&llm.Usage{TotalTokens: 7})
|
||||
before := a.Messages()
|
||||
|
||||
var out strings.Builder
|
||||
if err := a.BTW(context.Background(), "what does that mean?", &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.String() != "side answer" {
|
||||
t.Fatalf("output = %q", out.String())
|
||||
}
|
||||
if !reflect.DeepEqual(a.Messages(), before) {
|
||||
t.Fatalf("conversation changed:\nbefore: %+v\nafter: %+v", before, a.Messages())
|
||||
}
|
||||
if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 7 {
|
||||
t.Fatalf("usage changed: %+v", usage)
|
||||
}
|
||||
if len(provider.req.Tools) != 0 || provider.req.ToolChoice != "" {
|
||||
t.Fatalf("btw exposed tools: %+v / %q", provider.req.Tools, provider.req.ToolChoice)
|
||||
}
|
||||
last := provider.req.Messages[len(provider.req.Messages)-1]
|
||||
if last.Role != llm.RoleUser || last.Content != "what does that mean?" {
|
||||
t.Fatalf("last request message = %+v", last)
|
||||
}
|
||||
if provider.req.Extra["thinking"] != "high" {
|
||||
t.Fatalf("thinking extra = %+v", provider.req.Extra)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBTWRejectsEmptyQuestionAndReturnsProviderError(t *testing.T) {
|
||||
provider := &btwProvider{err: errors.New("boom")}
|
||||
a := New(Options{Provider: provider, Model: "test"})
|
||||
if err := a.BTW(context.Background(), " ", nil); err == nil {
|
||||
t.Fatal("empty question did not fail")
|
||||
}
|
||||
if err := a.BTW(context.Background(), "question", nil); err == nil || !strings.Contains(err.Error(), "boom") {
|
||||
t.Fatalf("provider error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDotContinuesCurrentTaskWithoutRewritingHistory(t *testing.T) {
|
||||
provider := &btwProvider{}
|
||||
a := New(Options{Provider: provider, Model: "test"})
|
||||
if err := a.RunTurn(context.Background(), " . ", nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
requestMessage := provider.req.Messages[len(provider.req.Messages)-1]
|
||||
if requestMessage.Content != continuationPrompt {
|
||||
t.Fatalf("request content = %q, want %q", requestMessage.Content, continuationPrompt)
|
||||
}
|
||||
history := a.Messages()
|
||||
if history[0].Content != "." {
|
||||
t.Fatalf("stored input = %q, want literal dot", history[0].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpandContinuationMessagesOnlyExpandsSingleDot(t *testing.T) {
|
||||
messages := []llm.Message{
|
||||
{Role: llm.RoleUser, Content: "."},
|
||||
{Role: llm.RoleUser, Content: " . "},
|
||||
{Role: llm.RoleUser, Content: ".."},
|
||||
{Role: llm.RoleUser, Content: ". continue"},
|
||||
{Role: llm.RoleAssistant, Content: "."},
|
||||
}
|
||||
expanded := expandContinuationMessages(messages)
|
||||
if expanded[0].Content != continuationPrompt || expanded[1].Content != continuationPrompt {
|
||||
t.Fatalf("single dots were not expanded: %+v", expanded)
|
||||
}
|
||||
for _, index := range []int{2, 3, 4} {
|
||||
if expanded[index].Content != messages[index].Content {
|
||||
t.Fatalf("message %d changed from %q to %q", index, messages[index].Content, expanded[index].Content)
|
||||
}
|
||||
}
|
||||
if messages[0].Content != "." || messages[1].Content != " . " {
|
||||
t.Fatalf("source messages were mutated: %+v", messages)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompleteConversationPrefixDropsPendingToolRound(t *testing.T) {
|
||||
messages := []llm.Message{
|
||||
{Role: llm.RoleUser, Content: "first"},
|
||||
{Role: llm.RoleAssistant, Content: "answer"},
|
||||
{Role: llm.RoleUser, Content: "main"},
|
||||
{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "call_1"}, {ID: "call_2"}}},
|
||||
{Role: llm.RoleTool, ToolCallID: "call_1", Content: "one"},
|
||||
}
|
||||
got := completeConversationPrefix(messages)
|
||||
if len(got) != 3 || got[len(got)-1].Content != "main" {
|
||||
t.Fatalf("complete prefix = %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
type concurrentBTWProvider struct {
|
||||
mainStarted chan struct{}
|
||||
releaseMain chan struct{}
|
||||
once sync.Once
|
||||
mu sync.Mutex
|
||||
btwRequest llm.ChatRequest
|
||||
}
|
||||
|
||||
func (p *concurrentBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
last := req.Messages[len(req.Messages)-1].Content
|
||||
if last == "main task" {
|
||||
p.once.Do(func() { close(p.mainStarted) })
|
||||
select {
|
||||
case <-p.releaseMain:
|
||||
return emit(llm.StreamEvent{Content: "main answer"})
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
p.mu.Lock()
|
||||
p.btwRequest = req
|
||||
p.mu.Unlock()
|
||||
return emit(llm.StreamEvent{Content: "side answer"})
|
||||
}
|
||||
|
||||
func TestBTWCanRunDuringMainTurn(t *testing.T) {
|
||||
provider := &concurrentBTWProvider{mainStarted: make(chan struct{}), releaseMain: make(chan struct{})}
|
||||
a := New(Options{Provider: provider, Model: "test"})
|
||||
mainDone := make(chan error, 1)
|
||||
go func() { mainDone <- a.RunTurn(context.Background(), "main task", nil, nil) }()
|
||||
<-provider.mainStarted
|
||||
|
||||
var side strings.Builder
|
||||
if err := a.BTW(context.Background(), "side question", &side); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if side.String() != "side answer" {
|
||||
t.Fatalf("side output = %q", side.String())
|
||||
}
|
||||
provider.mu.Lock()
|
||||
req := provider.btwRequest
|
||||
provider.mu.Unlock()
|
||||
if len(req.Messages) < 2 || req.Messages[len(req.Messages)-2].Content != "main task" {
|
||||
t.Fatalf("btw request did not include active main turn: %+v", req.Messages)
|
||||
}
|
||||
|
||||
close(provider.releaseMain)
|
||||
if err := <-mainDone; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(a.Messages()) != 2 {
|
||||
t.Fatalf("main conversation contains btw messages: %+v", a.Messages())
|
||||
}
|
||||
}
|
||||
|
||||
type compactProvider struct {
|
||||
summaryCalls int
|
||||
answerCalls int
|
||||
@@ -629,6 +801,120 @@ type parallelProvider struct {
|
||||
calls int
|
||||
}
|
||||
|
||||
type authorizationProvider struct {
|
||||
call llm.ToolCallDelta
|
||||
calls int
|
||||
toolOutput string
|
||||
}
|
||||
|
||||
func (p *authorizationProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.calls++
|
||||
if p.calls == 1 {
|
||||
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{p.call}})
|
||||
}
|
||||
p.toolOutput = req.Messages[len(req.Messages)-1].Content
|
||||
return emit(llm.StreamEvent{Content: "done"})
|
||||
}
|
||||
|
||||
type recordingAuthorizer struct {
|
||||
decision ToolAuthorizationDecision
|
||||
requests []ToolAuthorizationRequest
|
||||
}
|
||||
|
||||
func (a *recordingAuthorizer) AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error) {
|
||||
a.requests = append(a.requests, req)
|
||||
return a.decision, ctx.Err()
|
||||
}
|
||||
|
||||
func TestAgentAuthorizesMutationAndCreatesRollbackSnapshot(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
|
||||
Arguments: `{"path":"new.txt","content":"hello"}`,
|
||||
}}
|
||||
authorizer := &recordingAuthorizer{decision: ToolAuthorizationAllow}
|
||||
rollback := tools.NewRollbackStore()
|
||||
a := New(Options{
|
||||
Provider: provider,
|
||||
Model: "test",
|
||||
ToolRegistry: tools.Builtins(dir),
|
||||
Authorizer: authorizer,
|
||||
RollbackStore: rollback,
|
||||
})
|
||||
|
||||
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(authorizer.requests) != 1 || !authorizer.requests[0].RollbackAvailable {
|
||||
t.Fatalf("authorization requests = %+v", authorizer.requests)
|
||||
}
|
||||
if !strings.Contains(authorizer.requests[0].Preview, "write 5 bytes to new.txt") {
|
||||
t.Fatalf("preview = %q", authorizer.requests[0].Preview)
|
||||
}
|
||||
rollbackID := rollback.LastID()
|
||||
if rollbackID == "" || !strings.Contains(provider.toolOutput, rollbackID) {
|
||||
t.Fatalf("tool output %q missing rollback ID %q", provider.toolOutput, rollbackID)
|
||||
}
|
||||
if _, err := rollback.Restore(rollbackID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "new.txt")); !os.IsNotExist(err) {
|
||||
t.Fatalf("new file still exists after rollback: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentDeniedMutationDoesNotExecute(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||
Index: 0, ID: "write_1", Type: "function", Name: "file_write",
|
||||
Arguments: `{"path":"denied.txt","content":"no"}`,
|
||||
}}
|
||||
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
|
||||
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
|
||||
|
||||
if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "denied.txt")); !os.IsNotExist(err) {
|
||||
t.Fatalf("denied file exists: %v", err)
|
||||
}
|
||||
if !strings.Contains(provider.toolOutput, "denied by user") {
|
||||
t.Fatalf("tool output = %q", provider.toolOutput)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentDoesNotAuthorizeApplyPatchDryRun(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
provider := &authorizationProvider{call: llm.ToolCallDelta{
|
||||
Index: 0, ID: "patch_1", Type: "function", Name: "apply_patch",
|
||||
Arguments: `{"patch":"*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch","dry_run":true}`,
|
||||
}}
|
||||
authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
|
||||
a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})
|
||||
|
||||
if err := a.RunTurn(context.Background(), "preview", nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(authorizer.requests) != 0 {
|
||||
t.Fatalf("dry run requested authorization: %+v", authorizer.requests)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
|
||||
t.Fatalf("dry run created a file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetMessagesRefreshesCurrentSystemInstructions(t *testing.T) {
|
||||
a := New(Options{SystemPrompt: "current instructions"})
|
||||
a.SetMessages([]llm.Message{
|
||||
{Role: llm.RoleSystem, Content: "old instructions"},
|
||||
{Role: llm.RoleUser, Content: "hello"},
|
||||
})
|
||||
messages := a.Messages()
|
||||
if messages[0].Content != "current instructions" || messages[1].Content != "hello" {
|
||||
t.Fatalf("messages = %+v", messages)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.calls++
|
||||
if p.calls <= len(p.responses) {
|
||||
|
||||
Reference in New Issue
Block a user