1package core
2
3import (
4 "context"
5 "encoding/json"
6 "testing"
7
8 "github.com/patriceckhart/zot/packages/provider"
9)
10
11type preludeClient struct {
12 messages []provider.Message
13 calls int
14}
15
16func (c *preludeClient) Name() string { return "prelude" }
17func (c *preludeClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
18 c.calls++
19 c.messages = append([]provider.Message(nil), req.Messages...)
20 ch := make(chan provider.Event, 1)
21 ch <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "done"}}}}
22 close(ch)
23 return ch, nil
24}
25
26func TestPromptWithToolPairsResultBeforeModelAndPersists(t *testing.T) {
27 client := &preludeClient{}
28 tool := &recordingTool{}
29 a := NewAgent(client, "test", "", Registry{"echo": tool})
30 session, err := NewSession(t.TempDir(), t.TempDir(), "test", "test", "test")
31 if err != nil {
32 t.Fatal(err)
33 }
34 a.OnMessageAppended = func(msg provider.Message) {
35 if err := session.AppendMessage(msg); err != nil {
36 t.Error(err)
37 }
38 }
39 call := provider.ToolCallBlock{ID: "ext-prompt-1", Name: "echo", Arguments: json.RawMessage(`{"x":1}`)}
40 var events []AgentEvent
41 if err := a.PromptWithTool(context.Background(), "and push", call, "my-ext", func(ev AgentEvent) { events = append(events, ev) }); err != nil {
42 t.Fatal(err)
43 }
44 if client.calls != 1 || len(client.messages) != 3 || client.messages[0].Role != provider.RoleUser || client.messages[1].Role != provider.RoleAssistant || client.messages[2].Role != provider.RoleTool {
45 t.Fatalf("model input: calls=%d messages=%+v", client.calls, client.messages)
46 }
47 gotCall := client.messages[1].Content[0].(provider.ToolCallBlock)
48 gotResult := client.messages[2].Content[0].(provider.ToolResultBlock)
49 if gotCall.ID != gotResult.CallID || gotResult.IsError || client.messages[1].Meta["origin_extension"] != "my-ext" || string(tool.lastArgs) != string(call.Arguments) {
50 t.Fatalf("call/result mismatch: %v %v", gotCall, gotResult)
51 }
52 if _, ok := events[1].(EvToolCall); !ok {
53 t.Fatalf("missing call event: %T", events[1])
54 }
55 if _, ok := events[2].(EvToolResult); !ok {
56 t.Fatalf("missing result event: %T", events[2])
57 }
58 path := session.Path
59 if err := session.Close(); err != nil {
60 t.Fatal(err)
61 }
62 reopened, msgs, err := OpenSession(path)
63 if err != nil {
64 t.Fatal(err)
65 }
66 t.Cleanup(func() { _ = reopened.Close() })
67 if len(msgs) < 3 || msgs[1].Meta["origin_extension"] != "my-ext" || msgs[1].Meta["synthetic_tool_call"] != "true" || msgs[2].Content[0].(provider.ToolResultBlock).CallID != gotCall.ID {
68 t.Fatalf("restored transcript: %+v", msgs)
69 }
70}
71
72func TestPromptWithToolCancellationKeepsPair(t *testing.T) {
73 client := &preludeClient{}
74 tool := &blockingTool{started: make(chan struct{}), release: make(chan struct{})}
75 a := NewAgent(client, "test", "", Registry{"echo": tool})
76 ctx, cancel := context.WithCancel(context.Background())
77 done := make(chan error, 1)
78 go func() {
79 done <- a.PromptWithTool(ctx, "request", provider.ToolCallBlock{ID: "id", Name: "echo", Arguments: json.RawMessage(`{}`)}, "ext", nil)
80 }()
81 <-tool.started
82 cancel()
83 if err := <-done; err != context.Canceled {
84 t.Fatalf("error = %v", err)
85 }
86 msgs := a.Messages()
87 if client.calls != 0 || len(msgs) != 3 || msgs[1].Content[0].(provider.ToolCallBlock).ID != msgs[2].Content[0].(provider.ToolResultBlock).CallID || !msgs[2].Content[0].(provider.ToolResultBlock).IsError {
88 t.Fatalf("cancelled transcript: calls=%d messages=%+v", client.calls, msgs)
89 }
90}
91
92func TestPromptWithToolDeniedAndInvalid(t *testing.T) {
93 client := &preludeClient{}
94 tool := &recordingTool{}
95 a := NewAgent(client, "test", "", Registry{"echo": tool})
96 a.BeforeToolExecute = func(provider.ToolCallBlock) (bool, string, json.RawMessage) { return false, "denied", nil }
97 call := provider.ToolCallBlock{ID: "id", Name: "echo", Arguments: json.RawMessage(`{}`)}
98 if err := a.PromptWithTool(context.Background(), "request", call, "ext", nil); err != nil {
99 t.Fatal(err)
100 }
101 if tool.lastArgs != nil || !client.messages[2].Content[0].(provider.ToolResultBlock).IsError {
102 t.Fatalf("guard bypassed: %+v", client.messages)
103 }
104 before := len(a.Messages())
105 call.Arguments = json.RawMessage(`[]`)
106 if err := a.PromptWithTool(context.Background(), "request", call, "ext", nil); err == nil || len(a.Messages()) != before {
107 t.Fatalf("invalid input mutated transcript: %v", err)
108 }
109}