diff --git a/pkg/agent/integration_test.go b/pkg/agent/integration_test.go index 98cfd49..6de55ad 100644 --- a/pkg/agent/integration_test.go +++ b/pkg/agent/integration_test.go @@ -14,8 +14,9 @@ import ( ) // 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) { + skipUnlessIntegration(t) client, err := llamacpp.New(llamacpp.Config{ 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. func TestIntegration_LlamaCPP_Stream(t *testing.T) { + skipUnlessIntegration(t) client, err := llamacpp.New(llamacpp.Config{ 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. func TestIntegration_AgentLoop_Generate(t *testing.T) { + skipUnlessIntegration(t) client, err := llamacpp.New(llamacpp.Config{ 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. func TestIntegration_AgentLoop_Stream(t *testing.T) { + skipUnlessIntegration(t) client, err := llamacpp.New(llamacpp.Config{ BaseURL: "http://localhost:8080/v1", }) @@ -148,10 +152,21 @@ func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error return nil } -func TestMain(m *testing.M) { - // Skip integration tests unless explicitly enabled +// skipUnlessIntegration skips a single integration test (one that needs a +// 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" { - return + t.Skip("set INTEGRATION_TESTS=1 to run (requires a live llama.cpp server on localhost:8080)") } - os.Exit(m.Run()) } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 64981c3..0d98e22 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -107,18 +107,30 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R 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 { result, err := l.executeTool(ctx, call) if err != nil { messages = append(messages, llm.Message{ - Role: llm.RoleTool, - Content: fmt.Sprintf("Error: %v", err), + Role: llm.RoleTool, + ToolCallID: call.ID, + Content: fmt.Sprintf("Error: %v", err), }) continue } messages = append(messages, llm.Message{ - Role: llm.RoleTool, - Content: result.Content, + Role: llm.RoleTool, + ToolCallID: call.ID, + Content: result.Content, }) 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 { 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 { result, err := l.executeTool(ctx, tc) if err != nil { messages = append(messages, llm.Message{ - Role: llm.RoleTool, - Content: fmt.Sprintf("Error: %v", err), + Role: llm.RoleTool, + ToolCallID: tc.ID, + Content: fmt.Sprintf("Error: %v", err), }) continue } messages = append(messages, llm.Message{ - Role: llm.RoleTool, - Content: result.Content, + Role: llm.RoleTool, + ToolCallID: tc.ID, + Content: result.Content, }) } } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 29e0de7..e81f128 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -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) { mockClient := &mockLLM{ generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { diff --git a/pkg/llm/providers/llamacpp/client.go b/pkg/llm/providers/llamacpp/client.go index 9554ff3..716b078 100644 --- a/pkg/llm/providers/llamacpp/client.go +++ b/pkg/llm/providers/llamacpp/client.go @@ -105,6 +105,40 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq 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) for scanner.Scan() { line := scanner.Text() @@ -143,6 +177,23 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq } 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{ Delta: choice.Delta.Content, ReasoningDelta: choice.Delta.ReasoningContent, @@ -150,7 +201,17 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq } if 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) { return } @@ -180,8 +241,24 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader messages := make([]llamaMessage, len(req.Messages)) for i, m := range req.Messages { messages[i] = llamaMessage{ - Role: string(m.Role), - Content: m.Content, + Role: string(m.Role), + 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 { - Role string `json:"role"` - Content string `json:"content"` + Role string `json:"role"` + 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 { diff --git a/pkg/llm/providers/llamacpp/client_test.go b/pkg/llm/providers/llamacpp/client_test.go index 0de722a..7e66c3b 100644 --- a/pkg/llm/providers/llamacpp/client_test.go +++ b/pkg/llm/providers/llamacpp/client_test.go @@ -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) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 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) { var gotBody string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/pkg/llm/providers/openai/client.go b/pkg/llm/providers/openai/client.go index d3967cf..291e179 100644 --- a/pkg/llm/providers/openai/client.go +++ b/pkg/llm/providers/openai/client.go @@ -165,8 +165,24 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { messages := make([]openaiMessage, len(req.Messages)) for i, m := range req.Messages { messages[i] = openaiMessage{ - Role: string(m.Role), - Content: m.Content, + Role: string(m.Role), + 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 { - Role string `json:"role"` - Content string `json:"content"` + Role string `json:"role"` + 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 { diff --git a/pkg/llm/types.go b/pkg/llm/types.go index f9f222d..3e571de 100644 --- a/pkg/llm/types.go +++ b/pkg/llm/types.go @@ -18,10 +18,17 @@ const ( // Message is a single message in a conversation. type Message struct { - Role Role `json:"role"` - Content string `json:"content"` - ToolCallID string `json:"tool_call_id,omitempty"` - Name string `json:"name,omitempty"` + Role Role `json:"role"` + Content string `json:"content"` + ToolCallID string `json:"tool_call_id,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. diff --git a/pkg/tools/sandbox/sandbox_test.go b/pkg/tools/sandbox/sandbox_test.go index 905e9ca..91cf42b 100644 --- a/pkg/tools/sandbox/sandbox_test.go +++ b/pkg/tools/sandbox/sandbox_test.go @@ -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) { dir := t.TempDir() sb, err := sandbox.NewSandbox(dir) diff --git a/pkg/tools/types_test.go b/pkg/tools/types_test.go new file mode 100644 index 0000000..09f4262 --- /dev/null +++ b/pkg/tools/types_test.go @@ -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) + } +}