Merge pull request #4 from VictorVargas/feat/tools-for-model
Fix tool call tracking and streaming assembly for all providers
This commit is contained in:
commit
eeefe07809
9 changed files with 820 additions and 25 deletions
|
|
@ -14,8 +14,9 @@ import (
|
||||||
)
|
)
|
||||||
|
|
||||||
// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080.
|
// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080.
|
||||||
// Run with: go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate -tags=integration
|
// Run with: INTEGRATION_TESTS=1 go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate
|
||||||
func TestIntegration_LlamaCPP_Generate(t *testing.T) {
|
func TestIntegration_LlamaCPP_Generate(t *testing.T) {
|
||||||
|
skipUnlessIntegration(t)
|
||||||
client, err := llamacpp.New(llamacpp.Config{
|
client, err := llamacpp.New(llamacpp.Config{
|
||||||
BaseURL: "http://localhost:8080/v1",
|
BaseURL: "http://localhost:8080/v1",
|
||||||
})
|
})
|
||||||
|
|
@ -39,6 +40,7 @@ func TestIntegration_LlamaCPP_Generate(t *testing.T) {
|
||||||
|
|
||||||
// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080.
|
// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080.
|
||||||
func TestIntegration_LlamaCPP_Stream(t *testing.T) {
|
func TestIntegration_LlamaCPP_Stream(t *testing.T) {
|
||||||
|
skipUnlessIntegration(t)
|
||||||
client, err := llamacpp.New(llamacpp.Config{
|
client, err := llamacpp.New(llamacpp.Config{
|
||||||
BaseURL: "http://localhost:8080/v1",
|
BaseURL: "http://localhost:8080/v1",
|
||||||
})
|
})
|
||||||
|
|
@ -67,6 +69,7 @@ func TestIntegration_LlamaCPP_Stream(t *testing.T) {
|
||||||
|
|
||||||
// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp.
|
// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp.
|
||||||
func TestIntegration_AgentLoop_Generate(t *testing.T) {
|
func TestIntegration_AgentLoop_Generate(t *testing.T) {
|
||||||
|
skipUnlessIntegration(t)
|
||||||
client, err := llamacpp.New(llamacpp.Config{
|
client, err := llamacpp.New(llamacpp.Config{
|
||||||
BaseURL: "http://localhost:8080/v1",
|
BaseURL: "http://localhost:8080/v1",
|
||||||
})
|
})
|
||||||
|
|
@ -112,6 +115,7 @@ func TestIntegration_AgentLoop_Generate(t *testing.T) {
|
||||||
|
|
||||||
// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming.
|
// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming.
|
||||||
func TestIntegration_AgentLoop_Stream(t *testing.T) {
|
func TestIntegration_AgentLoop_Stream(t *testing.T) {
|
||||||
|
skipUnlessIntegration(t)
|
||||||
client, err := llamacpp.New(llamacpp.Config{
|
client, err := llamacpp.New(llamacpp.Config{
|
||||||
BaseURL: "http://localhost:8080/v1",
|
BaseURL: "http://localhost:8080/v1",
|
||||||
})
|
})
|
||||||
|
|
@ -148,10 +152,21 @@ func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestMain(m *testing.M) {
|
// skipUnlessIntegration skips a single integration test (one that needs a
|
||||||
// Skip integration tests unless explicitly enabled
|
// live llama.cpp server on localhost:8080) unless explicitly enabled.
|
||||||
|
//
|
||||||
|
// This used to be done in TestMain by returning early without calling
|
||||||
|
// m.Run() when INTEGRATION_TESTS wasn't set - but a package's TestMain
|
||||||
|
// covers its *entire* test binary (both package agent_test, here, and
|
||||||
|
// package agent, e.g. loop_test.go, get linked together), so that skipped
|
||||||
|
// every test in the package, not just the four integration ones. In
|
||||||
|
// practice that meant `go test ./...` reported this package as passing
|
||||||
|
// while silently running zero of its tests, including all the mock-based
|
||||||
|
// coverage in loop_test.go for approvals, sandboxing, and tool-call
|
||||||
|
// handling.
|
||||||
|
func skipUnlessIntegration(t *testing.T) {
|
||||||
|
t.Helper()
|
||||||
if os.Getenv("INTEGRATION_TESTS") != "1" {
|
if os.Getenv("INTEGRATION_TESTS") != "1" {
|
||||||
return
|
t.Skip("set INTEGRATION_TESTS=1 to run (requires a live llama.cpp server on localhost:8080)")
|
||||||
}
|
}
|
||||||
os.Exit(m.Run())
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -107,18 +107,30 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Record the assistant's own turn (including which tools it asked
|
||||||
|
// for) before the results, so the next request has a coherent
|
||||||
|
// assistant-tool_calls / tool-result pair instead of a dangling
|
||||||
|
// tool message the model can't attribute to anything.
|
||||||
|
messages = append(messages, llm.Message{
|
||||||
|
Role: llm.RoleAssistant,
|
||||||
|
Content: resp.Content,
|
||||||
|
ToolCalls: resp.ToolCalls,
|
||||||
|
})
|
||||||
|
|
||||||
for _, call := range resp.ToolCalls {
|
for _, call := range resp.ToolCalls {
|
||||||
result, err := l.executeTool(ctx, call)
|
result, err := l.executeTool(ctx, call)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
messages = append(messages, llm.Message{
|
messages = append(messages, llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
Content: fmt.Sprintf("Error: %v", err),
|
ToolCallID: call.ID,
|
||||||
|
Content: fmt.Sprintf("Error: %v", err),
|
||||||
})
|
})
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
messages = append(messages, llm.Message{
|
messages = append(messages, llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
Content: result.Content,
|
ToolCallID: call.ID,
|
||||||
|
Content: result.Content,
|
||||||
})
|
})
|
||||||
allToolCalls = append(allToolCalls, call)
|
allToolCalls = append(allToolCalls, call)
|
||||||
}
|
}
|
||||||
|
|
@ -165,18 +177,31 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
|
||||||
|
|
||||||
if len(chunk.ToolCalls) > 0 {
|
if len(chunk.ToolCalls) > 0 {
|
||||||
hasToolCalls = true
|
hasToolCalls = true
|
||||||
|
|
||||||
|
// Same reasoning as in Run: without recording the
|
||||||
|
// assistant's own tool_calls turn first, the tool
|
||||||
|
// results that follow have nothing for the model to
|
||||||
|
// attribute them to on the next request.
|
||||||
|
messages = append(messages, llm.Message{
|
||||||
|
Role: llm.RoleAssistant,
|
||||||
|
Content: responseBuilder.String(),
|
||||||
|
ToolCalls: chunk.ToolCalls,
|
||||||
|
})
|
||||||
|
|
||||||
for _, tc := range chunk.ToolCalls {
|
for _, tc := range chunk.ToolCalls {
|
||||||
result, err := l.executeTool(ctx, tc)
|
result, err := l.executeTool(ctx, tc)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
messages = append(messages, llm.Message{
|
messages = append(messages, llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
Content: fmt.Sprintf("Error: %v", err),
|
ToolCallID: tc.ID,
|
||||||
|
Content: fmt.Sprintf("Error: %v", err),
|
||||||
})
|
})
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
messages = append(messages, llm.Message{
|
messages = append(messages, llm.Message{
|
||||||
Role: llm.RoleTool,
|
Role: llm.RoleTool,
|
||||||
Content: result.Content,
|
ToolCallID: tc.ID,
|
||||||
|
Content: result.Content,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -166,6 +166,167 @@ func TestRun_ToolCalls(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestRun_ToolCalls_RecordsAssistantTurnAndToolCallID is the regression test
|
||||||
|
// for a bug where the follow-up request sent to the model, after executing
|
||||||
|
// a tool, never included the assistant message that requested the call
|
||||||
|
// (with its ToolCalls) nor set ToolCallID on the tool-result message. Some
|
||||||
|
// chat templates get confused by a "tool" message with nothing to attribute
|
||||||
|
// it to and the model loses track of what it already tried, which produced
|
||||||
|
// exactly the symptom reported in production: the model re-greeting and
|
||||||
|
// re-attempting the same search over and over instead of ever converging.
|
||||||
|
func TestRun_ToolCalls_RecordsAssistantTurnAndToolCallID(t *testing.T) {
|
||||||
|
registry := tools.NewRegistry()
|
||||||
|
registry.Register(tools.Tool{
|
||||||
|
Name: "greet",
|
||||||
|
Description: "Greet someone",
|
||||||
|
InputSchema: json.RawMessage(`{}`),
|
||||||
|
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||||
|
return tools.ToolResult{Content: "Hello!"}, nil
|
||||||
|
},
|
||||||
|
Permission: tools.Allow,
|
||||||
|
})
|
||||||
|
|
||||||
|
var requests []llm.CompletionRequest
|
||||||
|
callCount := 0
|
||||||
|
mockClient := &mockLLM{
|
||||||
|
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
requests = append(requests, req)
|
||||||
|
callCount++
|
||||||
|
if callCount == 1 {
|
||||||
|
return llm.CompletionResponse{
|
||||||
|
Content: "Voy a saludar.",
|
||||||
|
ToolCalls: []llm.ToolCall{{ID: "call-1", Name: "greet", Arguments: json.RawMessage(`{"name":"World"}`)}},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
return llm.CompletionResponse{Content: "Done!"}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
loop := New(Config{
|
||||||
|
LLM: mockClient,
|
||||||
|
Persona: persona.DefaultPersona(),
|
||||||
|
Tools: registry,
|
||||||
|
})
|
||||||
|
|
||||||
|
if _, err := loop.Run(context.Background(), "Say hi"); err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(requests) != 2 {
|
||||||
|
t.Fatalf("expected 2 requests to the model, got %d", len(requests))
|
||||||
|
}
|
||||||
|
|
||||||
|
// The second request (the follow-up after the tool ran) must contain
|
||||||
|
// the assistant's own tool_calls turn, immediately followed by a tool
|
||||||
|
// message whose ToolCallID matches it.
|
||||||
|
second := requests[1].Messages
|
||||||
|
var assistantIdx, toolIdx = -1, -1
|
||||||
|
for i, m := range second {
|
||||||
|
if m.Role == llm.RoleAssistant && len(m.ToolCalls) > 0 {
|
||||||
|
assistantIdx = i
|
||||||
|
}
|
||||||
|
if m.Role == llm.RoleTool {
|
||||||
|
toolIdx = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if assistantIdx == -1 {
|
||||||
|
t.Fatalf("expected an assistant message carrying ToolCalls in the follow-up request, got %+v", second)
|
||||||
|
}
|
||||||
|
if second[assistantIdx].Content != "Voy a saludar." {
|
||||||
|
t.Errorf("expected the assistant message to keep its original content, got %q", second[assistantIdx].Content)
|
||||||
|
}
|
||||||
|
if second[assistantIdx].ToolCalls[0].ID != "call-1" || second[assistantIdx].ToolCalls[0].Name != "greet" {
|
||||||
|
t.Errorf("expected the recorded tool call to match what was requested, got %+v", second[assistantIdx].ToolCalls[0])
|
||||||
|
}
|
||||||
|
if toolIdx == -1 {
|
||||||
|
t.Fatalf("expected a tool-result message in the follow-up request, got %+v", second)
|
||||||
|
}
|
||||||
|
if second[toolIdx].ToolCallID != "call-1" {
|
||||||
|
t.Errorf("expected the tool message's ToolCallID to be %q, got %q", "call-1", second[toolIdx].ToolCallID)
|
||||||
|
}
|
||||||
|
if toolIdx <= assistantIdx {
|
||||||
|
t.Errorf("expected the tool-result message to come after the assistant's tool_calls message")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestRun_Stream_ToolCalls_RecordsAssistantTurnAndToolCallID is the
|
||||||
|
// streaming counterpart of the test above.
|
||||||
|
func TestRun_Stream_ToolCalls_RecordsAssistantTurnAndToolCallID(t *testing.T) {
|
||||||
|
registry := tools.NewRegistry()
|
||||||
|
registry.Register(tools.Tool{
|
||||||
|
Name: "greet",
|
||||||
|
Description: "Greet",
|
||||||
|
InputSchema: json.RawMessage(`{}`),
|
||||||
|
Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) {
|
||||||
|
return tools.ToolResult{Content: "greeted"}, nil
|
||||||
|
},
|
||||||
|
Permission: tools.Allow,
|
||||||
|
})
|
||||||
|
|
||||||
|
var requests []llm.CompletionRequest
|
||||||
|
callCount := 0
|
||||||
|
mockClient := &mockLLM{
|
||||||
|
streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
|
requests = append(requests, req)
|
||||||
|
callCount++
|
||||||
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
|
if callCount == 1 {
|
||||||
|
yield(llm.StreamChunk{Delta: "Voy a saludar."}, nil)
|
||||||
|
yield(llm.StreamChunk{
|
||||||
|
ToolCalls: []llm.ToolCall{{ID: "call-9", Name: "greet", Arguments: json.RawMessage("{}")}},
|
||||||
|
FinishReason: "tool_calls",
|
||||||
|
}, nil)
|
||||||
|
} else {
|
||||||
|
yield(llm.StreamChunk{Delta: "done"}, nil)
|
||||||
|
yield(llm.StreamChunk{FinishReason: "stop"}, nil)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
loop := New(Config{
|
||||||
|
LLM: mockClient,
|
||||||
|
Persona: persona.DefaultPersona(),
|
||||||
|
Tools: registry,
|
||||||
|
})
|
||||||
|
|
||||||
|
for _, err := range loop.RunStream(context.Background(), "test") {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected stream error: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(requests) != 2 {
|
||||||
|
t.Fatalf("expected 2 requests to the model, got %d", len(requests))
|
||||||
|
}
|
||||||
|
|
||||||
|
second := requests[1].Messages
|
||||||
|
var assistantIdx, toolIdx = -1, -1
|
||||||
|
for i, m := range second {
|
||||||
|
if m.Role == llm.RoleAssistant && len(m.ToolCalls) > 0 {
|
||||||
|
assistantIdx = i
|
||||||
|
}
|
||||||
|
if m.Role == llm.RoleTool {
|
||||||
|
toolIdx = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if assistantIdx == -1 {
|
||||||
|
t.Fatalf("expected an assistant message carrying ToolCalls in the follow-up request, got %+v", second)
|
||||||
|
}
|
||||||
|
if second[assistantIdx].Content != "Voy a saludar." {
|
||||||
|
t.Errorf("expected the assistant message to carry the content streamed before the tool call, got %q", second[assistantIdx].Content)
|
||||||
|
}
|
||||||
|
if second[assistantIdx].ToolCalls[0].ID != "call-9" {
|
||||||
|
t.Errorf("expected the recorded tool call ID to be %q, got %q", "call-9", second[assistantIdx].ToolCalls[0].ID)
|
||||||
|
}
|
||||||
|
if toolIdx == -1 || second[toolIdx].ToolCallID != "call-9" {
|
||||||
|
t.Fatalf("expected a tool-result message with ToolCallID %q, got %+v", "call-9", second)
|
||||||
|
}
|
||||||
|
if toolIdx <= assistantIdx {
|
||||||
|
t.Errorf("expected the tool-result message to come after the assistant's tool_calls message")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestRun_MaxIterations(t *testing.T) {
|
func TestRun_MaxIterations(t *testing.T) {
|
||||||
mockClient := &mockLLM{
|
mockClient := &mockLLM{
|
||||||
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
|
|
||||||
|
|
@ -105,6 +105,40 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// toolCallAccum buffers one tool call's fragments as they stream in:
|
||||||
|
// the OpenAI-compatible SSE format sends the id/name in the first
|
||||||
|
// delta for a given tool-call index and the (potentially large)
|
||||||
|
// arguments JSON in pieces across many subsequent deltas, so it
|
||||||
|
// can't be handed to a tool handler until it's fully assembled.
|
||||||
|
type toolCallAccum struct {
|
||||||
|
id string
|
||||||
|
name string
|
||||||
|
args strings.Builder
|
||||||
|
}
|
||||||
|
toolCallFrags := map[int]*toolCallAccum{}
|
||||||
|
var toolCallOrder []int
|
||||||
|
|
||||||
|
// flushToolCalls assembles the buffered fragments into complete
|
||||||
|
// tool calls (called once finish_reason arrives) and resets the
|
||||||
|
// accumulator for any further choices/events.
|
||||||
|
flushToolCalls := func() []llm.ToolCall {
|
||||||
|
if len(toolCallOrder) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
calls := make([]llm.ToolCall, 0, len(toolCallOrder))
|
||||||
|
for _, idx := range toolCallOrder {
|
||||||
|
frag := toolCallFrags[idx]
|
||||||
|
calls = append(calls, llm.ToolCall{
|
||||||
|
ID: frag.id,
|
||||||
|
Name: frag.name,
|
||||||
|
Arguments: json.RawMessage(frag.args.String()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
toolCallFrags = map[int]*toolCallAccum{}
|
||||||
|
toolCallOrder = nil
|
||||||
|
return calls
|
||||||
|
}
|
||||||
|
|
||||||
scanner := bufio.NewScanner(resp.Body)
|
scanner := bufio.NewScanner(resp.Body)
|
||||||
for scanner.Scan() {
|
for scanner.Scan() {
|
||||||
line := scanner.Text()
|
line := scanner.Text()
|
||||||
|
|
@ -143,6 +177,23 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, choice := range event.Choices {
|
for _, choice := range event.Choices {
|
||||||
|
hasFragment := len(choice.Delta.ToolCalls) > 0
|
||||||
|
for _, tc := range choice.Delta.ToolCalls {
|
||||||
|
frag, ok := toolCallFrags[tc.Index]
|
||||||
|
if !ok {
|
||||||
|
frag = &toolCallAccum{}
|
||||||
|
toolCallFrags[tc.Index] = frag
|
||||||
|
toolCallOrder = append(toolCallOrder, tc.Index)
|
||||||
|
}
|
||||||
|
if tc.ID != "" {
|
||||||
|
frag.id = tc.ID
|
||||||
|
}
|
||||||
|
if tc.Function.Name != "" {
|
||||||
|
frag.name = tc.Function.Name
|
||||||
|
}
|
||||||
|
frag.args.WriteString(tc.Function.Arguments)
|
||||||
|
}
|
||||||
|
|
||||||
chunk := llm.StreamChunk{
|
chunk := llm.StreamChunk{
|
||||||
Delta: choice.Delta.Content,
|
Delta: choice.Delta.Content,
|
||||||
ReasoningDelta: choice.Delta.ReasoningContent,
|
ReasoningDelta: choice.Delta.ReasoningContent,
|
||||||
|
|
@ -150,7 +201,17 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
||||||
}
|
}
|
||||||
if choice.FinishReason != "" {
|
if choice.FinishReason != "" {
|
||||||
chunk.FinishReason = choice.FinishReason
|
chunk.FinishReason = choice.FinishReason
|
||||||
|
chunk.ToolCalls = flushToolCalls()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// A fragment-only event (a piece of a tool call's streamed
|
||||||
|
// arguments, with nothing else in this delta) has nothing
|
||||||
|
// yet for the agent loop to act on: it was buffered above,
|
||||||
|
// so skip yielding an empty chunk for it.
|
||||||
|
if hasFragment && chunk.Delta == "" && chunk.ReasoningDelta == "" && chunk.FinishReason == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
if !yield(chunk, nil) {
|
if !yield(chunk, nil) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
@ -180,8 +241,24 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader
|
||||||
messages := make([]llamaMessage, len(req.Messages))
|
messages := make([]llamaMessage, len(req.Messages))
|
||||||
for i, m := range req.Messages {
|
for i, m := range req.Messages {
|
||||||
messages[i] = llamaMessage{
|
messages[i] = llamaMessage{
|
||||||
Role: string(m.Role),
|
Role: string(m.Role),
|
||||||
Content: m.Content,
|
Content: m.Content,
|
||||||
|
ToolCallID: m.ToolCallID,
|
||||||
|
Name: m.Name,
|
||||||
|
}
|
||||||
|
if len(m.ToolCalls) > 0 {
|
||||||
|
calls := make([]llamaToolCall, len(m.ToolCalls))
|
||||||
|
for j, tc := range m.ToolCalls {
|
||||||
|
calls[j] = llamaToolCall{
|
||||||
|
ID: tc.ID,
|
||||||
|
Type: "function",
|
||||||
|
Function: llamaFunction{
|
||||||
|
Name: tc.Name,
|
||||||
|
Arguments: string(tc.Arguments),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
messages[i].ToolCalls = calls
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -281,8 +358,11 @@ type llamaStreamOptions struct {
|
||||||
}
|
}
|
||||||
|
|
||||||
type llamaMessage struct {
|
type llamaMessage struct {
|
||||||
Role string `json:"role"`
|
Role string `json:"role"`
|
||||||
Content string `json:"content"`
|
Content string `json:"content"`
|
||||||
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
|
Name string `json:"name,omitempty"`
|
||||||
|
ToolCalls []llamaToolCall `json:"tool_calls,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type llamaTool struct {
|
type llamaTool struct {
|
||||||
|
|
|
||||||
|
|
@ -82,6 +82,89 @@ func TestClient_Generate(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestClient_BuildRequest_SendsToolCallHistory is the regression test for a
|
||||||
|
// bug where an assistant message's ToolCalls and a tool message's
|
||||||
|
// ToolCallID were silently dropped when building the wire request: the
|
||||||
|
// model would see a "tool" message with nothing tying it to a prior
|
||||||
|
// assistant turn, lose track of what it had already tried, and re-attempt
|
||||||
|
// the same thing over and over (reported in production as Rony repeatedly
|
||||||
|
// re-greeting and re-searching for a file instead of ever finishing).
|
||||||
|
func TestClient_BuildRequest_SendsToolCallHistory(t *testing.T) {
|
||||||
|
var gotBody string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
gotBody = string(body)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(llamaChatResponse{
|
||||||
|
Choices: []llamaChoice{{Message: llamaMessageResult{Content: "ok"}}},
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{BaseURL: server.URL + "/v1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = client.Generate(context.Background(), llm.CompletionRequest{
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "busca el archivo"},
|
||||||
|
{
|
||||||
|
Role: llm.RoleAssistant,
|
||||||
|
Content: "voy a buscar",
|
||||||
|
ToolCalls: []llm.ToolCall{
|
||||||
|
{ID: "call-1", Name: "glob", Arguments: json.RawMessage(`{"pattern":"*.md"}`)},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{Role: llm.RoleTool, ToolCallID: "call-1", Content: "No matches found."},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var sent struct {
|
||||||
|
Messages []struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ToolCallID string `json:"tool_call_id"`
|
||||||
|
ToolCalls []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Function struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
} `json:"function"`
|
||||||
|
} `json:"tool_calls"`
|
||||||
|
} `json:"messages"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(gotBody), &sent); err != nil {
|
||||||
|
t.Fatalf("failed to parse sent body: %v\nbody: %s", err, gotBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(sent.Messages) != 3 {
|
||||||
|
t.Fatalf("expected 3 messages sent, got %d: %s", len(sent.Messages), gotBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
assistantMsg := sent.Messages[1]
|
||||||
|
if assistantMsg.Role != "assistant" {
|
||||||
|
t.Fatalf("expected message 1 to be the assistant turn, got role %q", assistantMsg.Role)
|
||||||
|
}
|
||||||
|
if len(assistantMsg.ToolCalls) != 1 || assistantMsg.ToolCalls[0].ID != "call-1" {
|
||||||
|
t.Fatalf("expected the assistant message to carry its tool_calls with id 'call-1', got %+v", assistantMsg.ToolCalls)
|
||||||
|
}
|
||||||
|
if assistantMsg.ToolCalls[0].Function.Name != "glob" {
|
||||||
|
t.Errorf("expected function name 'glob', got %q", assistantMsg.ToolCalls[0].Function.Name)
|
||||||
|
}
|
||||||
|
|
||||||
|
toolMsg := sent.Messages[2]
|
||||||
|
if toolMsg.Role != "tool" {
|
||||||
|
t.Fatalf("expected message 2 to be the tool result, got role %q", toolMsg.Role)
|
||||||
|
}
|
||||||
|
if toolMsg.ToolCallID != "call-1" {
|
||||||
|
t.Errorf("expected tool_call_id 'call-1' on the tool message, got %q (body: %s)", toolMsg.ToolCallID, gotBody)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestClient_Generate_Error(t *testing.T) {
|
func TestClient_Generate_Error(t *testing.T) {
|
||||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
w.WriteHeader(http.StatusInternalServerError)
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
|
@ -233,6 +316,183 @@ func TestClient_Stream_FinishReason(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestClient_Stream_ToolCallSingleEvent covers the simplest case: a server
|
||||||
|
// that sends the whole tool call (id, name, complete arguments) in one delta
|
||||||
|
// followed immediately by finish_reason "tool_calls".
|
||||||
|
func TestClient_Stream_ToolCallSingleEvent(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"write","arguments":"{\"path\":\"a.txt\",\"content\":\"hi\"}"}}]},"finish_reason":"tool_calls"}]}` + "\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{BaseURL: server.URL + "/v1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 1 {
|
||||||
|
t.Fatalf("expected 1 chunk, got %d: %+v", len(chunks), chunks)
|
||||||
|
}
|
||||||
|
if len(chunks[0].ToolCalls) != 1 {
|
||||||
|
t.Fatalf("expected 1 tool call in the chunk, got %d", len(chunks[0].ToolCalls))
|
||||||
|
}
|
||||||
|
call := chunks[0].ToolCalls[0]
|
||||||
|
if call.ID != "call-1" || call.Name != "write" {
|
||||||
|
t.Errorf("expected call-1/write, got %+v", call)
|
||||||
|
}
|
||||||
|
if string(call.Arguments) != `{"path":"a.txt","content":"hi"}` {
|
||||||
|
t.Errorf("unexpected arguments: %s", call.Arguments)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClient_Stream_ToolCallFragmentsAssembled is the regression test for
|
||||||
|
// the actual bug reported in production: llama.cpp (like any OpenAI-style
|
||||||
|
// server) streams a tool call's arguments in many small deltas keyed by
|
||||||
|
// index, with the name/id only present in the first fragment. The old
|
||||||
|
// Stream() implementation never even read choice.Delta.ToolCalls, so every
|
||||||
|
// fragment was silently dropped and the agent loop never saw a tool call at
|
||||||
|
// all - the model would narrate "I'll write the file" and nothing would
|
||||||
|
// happen. This verifies the fragments are buffered and only surfaced, fully
|
||||||
|
// assembled, once finish_reason arrives.
|
||||||
|
func TestClient_Stream_ToolCallFragmentsAssembled(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"write","arguments":""}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.txt\",\"content\""}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"hi\"}"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{BaseURL: server.URL + "/v1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
// The four fragment-only events must not surface as separate empty
|
||||||
|
// chunks; only the finish_reason event, carrying the fully assembled
|
||||||
|
// call, should be yielded.
|
||||||
|
if len(chunks) != 1 {
|
||||||
|
t.Fatalf("expected 1 chunk (fragments buffered, only the assembled call yielded), got %d: %+v", len(chunks), chunks)
|
||||||
|
}
|
||||||
|
if chunks[0].FinishReason != "tool_calls" {
|
||||||
|
t.Errorf("expected finish_reason 'tool_calls', got %q", chunks[0].FinishReason)
|
||||||
|
}
|
||||||
|
if len(chunks[0].ToolCalls) != 1 {
|
||||||
|
t.Fatalf("expected 1 assembled tool call, got %d", len(chunks[0].ToolCalls))
|
||||||
|
}
|
||||||
|
call := chunks[0].ToolCalls[0]
|
||||||
|
if call.ID != "call-1" || call.Name != "write" {
|
||||||
|
t.Errorf("expected call-1/write, got %+v", call)
|
||||||
|
}
|
||||||
|
if string(call.Arguments) != `{"path":"a.txt","content":"hi"}` {
|
||||||
|
t.Errorf("expected assembled arguments %q, got %q", `{"path":"a.txt","content":"hi"}`, call.Arguments)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClient_Stream_ToolCallWithPrecedingContent verifies that reasoning or
|
||||||
|
// content deltas that arrive before a tool call (e.g. a model "thinking"
|
||||||
|
// before deciding to call a tool) are still streamed normally, and don't get
|
||||||
|
// mixed up with the buffered tool-call fragments.
|
||||||
|
func TestClient_Stream_ToolCallWithPrecedingContent(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"content":"Voy a escribir el archivo."}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-9","type":"function","function":{"name":"write","arguments":"{}"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{BaseURL: server.URL + "/v1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 2 {
|
||||||
|
t.Fatalf("expected 2 chunks (content, then the assembled tool call), got %d: %+v", len(chunks), chunks)
|
||||||
|
}
|
||||||
|
if chunks[0].Delta != "Voy a escribir el archivo." {
|
||||||
|
t.Errorf("expected the content delta first, got %q", chunks[0].Delta)
|
||||||
|
}
|
||||||
|
if len(chunks[0].ToolCalls) != 0 {
|
||||||
|
t.Errorf("expected the content chunk to carry no tool calls, got %+v", chunks[0].ToolCalls)
|
||||||
|
}
|
||||||
|
if len(chunks[1].ToolCalls) != 1 || chunks[1].ToolCalls[0].Name != "write" {
|
||||||
|
t.Errorf("expected the second chunk to carry the assembled write call, got %+v", chunks[1].ToolCalls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestClient_Stream_ParallelToolCalls verifies two tool calls streamed in
|
||||||
|
// parallel (interleaved by index) are assembled independently and returned
|
||||||
|
// in call order.
|
||||||
|
func TestClient_Stream_ParallelToolCalls(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-a","type":"function","function":{"name":"read","arguments":"{\"path\":"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"call-b","type":"function","function":{"name":"glob","arguments":"{\"pattern\":"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.txt\"}"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"*.go\"}"}}]}}]}` + "\n"))
|
||||||
|
w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{BaseURL: server.URL + "/v1"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 1 {
|
||||||
|
t.Fatalf("expected 1 chunk, got %d: %+v", len(chunks), chunks)
|
||||||
|
}
|
||||||
|
if len(chunks[0].ToolCalls) != 2 {
|
||||||
|
t.Fatalf("expected 2 assembled tool calls, got %d", len(chunks[0].ToolCalls))
|
||||||
|
}
|
||||||
|
if chunks[0].ToolCalls[0].Name != "read" || string(chunks[0].ToolCalls[0].Arguments) != `{"path":"a.txt"}` {
|
||||||
|
t.Errorf("unexpected first call: %+v", chunks[0].ToolCalls[0])
|
||||||
|
}
|
||||||
|
if chunks[0].ToolCalls[1].Name != "glob" || string(chunks[0].ToolCalls[1].Arguments) != `{"pattern":"*.go"}` {
|
||||||
|
t.Errorf("unexpected second call: %+v", chunks[0].ToolCalls[1])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestClient_Stream_RequestsAndParsesUsage(t *testing.T) {
|
func TestClient_Stream_RequestsAndParsesUsage(t *testing.T) {
|
||||||
var gotBody string
|
var gotBody string
|
||||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
|
|
||||||
|
|
@ -165,8 +165,24 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
||||||
messages := make([]openaiMessage, len(req.Messages))
|
messages := make([]openaiMessage, len(req.Messages))
|
||||||
for i, m := range req.Messages {
|
for i, m := range req.Messages {
|
||||||
messages[i] = openaiMessage{
|
messages[i] = openaiMessage{
|
||||||
Role: string(m.Role),
|
Role: string(m.Role),
|
||||||
Content: m.Content,
|
Content: m.Content,
|
||||||
|
ToolCallID: m.ToolCallID,
|
||||||
|
Name: m.Name,
|
||||||
|
}
|
||||||
|
if len(m.ToolCalls) > 0 {
|
||||||
|
calls := make([]openaiToolCall, len(m.ToolCalls))
|
||||||
|
for j, tc := range m.ToolCalls {
|
||||||
|
calls[j] = openaiToolCall{
|
||||||
|
ID: tc.ID,
|
||||||
|
Type: "function",
|
||||||
|
Function: openaiFunction{
|
||||||
|
Name: tc.Name,
|
||||||
|
Arguments: string(tc.Arguments),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
messages[i].ToolCalls = calls
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -250,8 +266,11 @@ type openaiChatRequest struct {
|
||||||
}
|
}
|
||||||
|
|
||||||
type openaiMessage struct {
|
type openaiMessage struct {
|
||||||
Role string `json:"role"`
|
Role string `json:"role"`
|
||||||
Content string `json:"content"`
|
Content string `json:"content"`
|
||||||
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
|
Name string `json:"name,omitempty"`
|
||||||
|
ToolCalls []openaiToolCall `json:"tool_calls,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type openaiTool struct {
|
type openaiTool struct {
|
||||||
|
|
|
||||||
|
|
@ -18,10 +18,17 @@ const (
|
||||||
|
|
||||||
// Message is a single message in a conversation.
|
// Message is a single message in a conversation.
|
||||||
type Message struct {
|
type Message struct {
|
||||||
Role Role `json:"role"`
|
Role Role `json:"role"`
|
||||||
Content string `json:"content"`
|
Content string `json:"content"`
|
||||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
Name string `json:"name,omitempty"`
|
Name string `json:"name,omitempty"`
|
||||||
|
// ToolCalls records the calls an assistant message requested, so the
|
||||||
|
// agent loop can replay them on the next request: without this, the
|
||||||
|
// conversation sent back to the model has tool-result messages with no
|
||||||
|
// assistant turn that requested them, which confuses (or is outright
|
||||||
|
// rejected by) the chat template - the model loses track of what it
|
||||||
|
// already asked for and re-attempts it, or restarts from scratch.
|
||||||
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ProviderCapabilities describes what a model supports.
|
// ProviderCapabilities describes what a model supports.
|
||||||
|
|
|
||||||
|
|
@ -129,6 +129,160 @@ func TestValidatePath_InvalidJSON(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_RelativeTraversalEscapes is the important case the
|
||||||
|
// absolute-path check in TestValidatePath_EscapesSandbox doesn't cover: a
|
||||||
|
// *relative* path that climbs out of the sandbox root with "..". This must
|
||||||
|
// be caught by the join+clean+prefix check in validatePath, not by the
|
||||||
|
// early filepath.IsAbs rejection.
|
||||||
|
func TestValidatePath_RelativeTraversalEscapes(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cases := []string{
|
||||||
|
"../outside.txt",
|
||||||
|
"../../etc/passwd",
|
||||||
|
"sub/../../outside.txt",
|
||||||
|
}
|
||||||
|
for _, path := range cases {
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "read_file",
|
||||||
|
Arguments: json.RawMessage(`{"path": "` + path + `"}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err == nil {
|
||||||
|
t.Errorf("expected %q to be rejected as a sandbox escape, got no error", path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_RelativeTraversalStayingInsideIsAllowed makes sure the
|
||||||
|
// traversal check isn't so strict it rejects "../" segments that still
|
||||||
|
// resolve back inside the sandbox root once cleaned.
|
||||||
|
func TestValidatePath_RelativeTraversalStayingInsideIsAllowed(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
if err := os.MkdirAll(filepath.Join(dir, "sub"), 0o755); err != nil {
|
||||||
|
t.Fatalf("setup mkdir: %v", err)
|
||||||
|
}
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "read_file",
|
||||||
|
Arguments: json.RawMessage(`{"path": "sub/../file.txt"}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err != nil {
|
||||||
|
t.Fatalf("expected path resolving back inside the sandbox to be allowed, got: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_PathInsideStringArray covers extractPaths' handling of
|
||||||
|
// []interface{} arguments (e.g. a tool that takes a list of file paths),
|
||||||
|
// which none of the single-"path"-key tests above exercise.
|
||||||
|
func TestValidatePath_PathInsideStringArray(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "read_many",
|
||||||
|
Arguments: json.RawMessage(`{"paths": ["ok.txt", "../../etc/passwd"]}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err == nil {
|
||||||
|
t.Fatal("expected the escaping path inside the array to be rejected")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_AllPathsInStringArrayAllowed is the allowed counterpart:
|
||||||
|
// every element of the array stays inside the sandbox.
|
||||||
|
func TestValidatePath_AllPathsInStringArrayAllowed(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "read_many",
|
||||||
|
Arguments: json.RawMessage(`{"paths": ["a.txt", "b/c.txt"]}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err != nil {
|
||||||
|
t.Fatalf("expected all-inside array to be allowed, got: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_OneOfSeveralArgsEscapes verifies that a call with several
|
||||||
|
// argument keys is rejected if *any* of them is a path-like value that
|
||||||
|
// escapes, not just when the single "path" key does.
|
||||||
|
func TestValidatePath_OneOfSeveralArgsEscapes(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "copy_file",
|
||||||
|
Arguments: json.RawMessage(`{"from": "safe.txt", "to": "../../etc/passwd", "note": "hello world"}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err == nil {
|
||||||
|
t.Fatal("expected the escaping 'to' argument to reject the whole call")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestValidatePath_NonPathStringsIgnored ensures ordinary string arguments
|
||||||
|
// that don't look like paths (no leading ./, ../, /, and no "word.ext"
|
||||||
|
// shape) are never treated as paths and can't accidentally trip the
|
||||||
|
// sandbox check.
|
||||||
|
func TestValidatePath_NonPathStringsIgnored(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
call := llm.ToolCall{
|
||||||
|
Name: "search",
|
||||||
|
Arguments: json.RawMessage(`{"query": "hello world", "count": 5, "enabled": true}`),
|
||||||
|
}
|
||||||
|
if err := sb.ValidateToolCall(tools.Tool{}, call); err != nil {
|
||||||
|
t.Fatalf("expected non-path-like arguments to be ignored, got: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestNewSandbox_RootIsAFileNotADirectory exercises the MkdirAll error
|
||||||
|
// branch: passing a path that already exists as a regular file can't be
|
||||||
|
// turned into a sandbox root.
|
||||||
|
func TestNewSandbox_RootIsAFileNotADirectory(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
filePath := filepath.Join(dir, "not-a-dir")
|
||||||
|
if err := os.WriteFile(filePath, []byte("x"), 0o644); err != nil {
|
||||||
|
t.Fatalf("setup: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := sandbox.NewSandbox(filePath); err == nil {
|
||||||
|
t.Fatal("expected an error when the sandbox root is an existing file")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSandboxOpen checks the escape hatch used by tests: it must return a
|
||||||
|
// usable, non-nil os.Root for the sandbox that was created.
|
||||||
|
func TestSandboxOpen(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if sb.Open() == nil {
|
||||||
|
t.Fatal("expected Open() to return a non-nil os.Root")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestValidatePath_NoPathsInArgs(t *testing.T) {
|
func TestValidatePath_NoPathsInArgs(t *testing.T) {
|
||||||
dir := t.TempDir()
|
dir := t.TempDir()
|
||||||
sb, err := sandbox.NewSandbox(dir)
|
sb, err := sandbox.NewSandbox(dir)
|
||||||
|
|
|
||||||
74
pkg/tools/types_test.go
Normal file
74
pkg/tools/types_test.go
Normal file
|
|
@ -0,0 +1,74 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestPermissionString(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
perm Permission
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{Allow, "allow"},
|
||||||
|
{Ask, "ask"},
|
||||||
|
{Deny, "deny"},
|
||||||
|
{Permission(99), "unknown"},
|
||||||
|
}
|
||||||
|
for _, c := range cases {
|
||||||
|
if got := c.perm.String(); got != c.want {
|
||||||
|
t.Errorf("Permission(%d).String() = %q, want %q", c.perm, got, c.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestPermissionZeroValueIsAllow guards an easy-to-miss footgun: a Tool
|
||||||
|
// literal that forgets to set Permission defaults to Allow (iota 0), not to
|
||||||
|
// the safer Ask/Deny, so any code relying on the zero value must be
|
||||||
|
// deliberate about it.
|
||||||
|
func TestPermissionZeroValueIsAllow(t *testing.T) {
|
||||||
|
var p Permission
|
||||||
|
if p != Allow {
|
||||||
|
t.Fatalf("expected zero-value Permission to be Allow, got %v", p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSentinelErrorsAreDistinctAndNonNil(t *testing.T) {
|
||||||
|
if ErrToolNotFound == nil {
|
||||||
|
t.Fatal("ErrToolNotFound must not be nil")
|
||||||
|
}
|
||||||
|
if ErrDuplicateTool == nil {
|
||||||
|
t.Fatal("ErrDuplicateTool must not be nil")
|
||||||
|
}
|
||||||
|
if ErrToolNotFound.Error() == ErrDuplicateTool.Error() {
|
||||||
|
t.Fatal("expected distinct error messages")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestToolHandlerSignatureIsCallable is a compile-time-flavored smoke test:
|
||||||
|
// it just confirms a plain function value satisfies ToolHandler and can be
|
||||||
|
// invoked through the type, so a signature change here would be caught.
|
||||||
|
func TestToolHandlerSignatureIsCallable(t *testing.T) {
|
||||||
|
var h ToolHandler = func(ctx context.Context, args json.RawMessage) (ToolResult, error) {
|
||||||
|
return ToolResult{Content: string(args)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
res, err := h(context.Background(), json.RawMessage(`{"ok":true}`))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if res.Content != `{"ok":true}` {
|
||||||
|
t.Errorf("unexpected content: %q", res.Content)
|
||||||
|
}
|
||||||
|
if res.IsError {
|
||||||
|
t.Error("expected IsError to be false")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestToolResultZeroValue(t *testing.T) {
|
||||||
|
var res ToolResult
|
||||||
|
if res.Content != "" || res.IsError || res.Metadata != nil || res.Artifacts != nil {
|
||||||
|
t.Fatalf("expected zero-value ToolResult to be empty, got %+v", res)
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue