1package extensions
2
3import (
4 "context"
5 "encoding/json"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/agent/extproto"
11 "github.com/patriceckhart/zot/packages/provider"
12)
13
14func TestCallToolDispatchAndRefusal(t *testing.T) {
15 m, ext, frames, replies := toolPeer(t)
16 ext.callTool = true
17 m.SetToolCaller(func(ctx context.Context, origin, name string, args json.RawMessage) extproto.ToolResultFromHost {
18 if origin != ext.Manifest.Name || name != "read" {
19 t.Errorf("unexpected request: %s %s %s", origin, name, args)
20 }
21 return extproto.ToolResultFromHost{Content: []extproto.ContentBlock{{Type: "text", Text: "ok"}}}
22 })
23 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "one", Name: "read", Args: json.RawMessage(`{"path":"a"}`)})
24 var result extproto.ToolResultFromHost
25 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || result.ID != "one" || result.IsError || result.Content[0].Text != "ok" {
26 t.Fatalf("result: %+v, %v", result, err)
27 }
28 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "two", Name: "read", Args: json.RawMessage(`{}`), ParentID: "unknown"})
29 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "parent") {
30 t.Fatalf("parent refusal: %+v, %v", result, err)
31 }
32 ext.mu.Lock()
33 ext.outboundCalls = map[string]outboundCall{"parent": {ctx: context.Background(), chain: []string{"ask"}}}
34 ext.mu.Unlock()
35 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "three", Name: "read", Args: json.RawMessage(`{}`)})
36 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "parent_id") {
37 t.Fatalf("missing parent refusal: %+v, %v", result, err)
38 }
39 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "four", ParentID: "parent", Name: "ask", Args: json.RawMessage(`{}`)})
40 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "cycle") {
41 t.Fatalf("cycle refusal: %+v, %v", result, err)
42 }
43 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "five", ParentID: "parent", Name: "read", Args: json.RawMessage(`{}`)})
44 result = extproto.ToolResultFromHost{}
45 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || result.IsError || result.Content[0].Text != "ok" {
46 t.Fatalf("nested call: %+v, %v", result, err)
47 }
48 ext.mu.Lock()
49 ext.outboundCalls["parent"] = outboundCall{ctx: context.Background(), chain: []string{"one", "two", "three", "four"}}
50 ext.mu.Unlock()
51 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "six", ParentID: "parent", Name: "read", Args: json.RawMessage(`{}`)})
52 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "depth") {
53 t.Fatalf("depth refusal: %+v, %v", result, err)
54 }
55}
56
57func TestCallToolRequiresOptIn(t *testing.T) {
58 m, _, frames, replies := toolPeer(t)
59 m.SetToolCaller(func(context.Context, string, string, json.RawMessage) extproto.ToolResultFromHost {
60 t.Error("unadvertised call reached tool caller")
61 return extproto.ToolResultFromHost{}
62 })
63 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "one", Name: "read", Args: json.RawMessage(`{}`)})
64 var result extproto.ToolResultFromHost
65 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "opt in") {
66 t.Fatalf("opt-in refusal: %+v, %v", result, err)
67 }
68}
69
70func TestCallToolToSameExtension(t *testing.T) {
71 m, ext, frames, replies := toolPeer(t)
72 ext.callTool = true
73 m.SetToolCaller(func(ctx context.Context, _, name string, args json.RawMessage) extproto.ToolResultFromHost {
74 tool, err := NewTool(m, ToolInfo{Name: name, Extension: ext.Manifest.Name}).Execute(ctx, args, nil)
75 if err != nil {
76 return extproto.ToolResultFromHost{IsError: true}
77 }
78 return extproto.ToolResultFromHost{IsError: tool.IsError, Content: []extproto.ContentBlock{{Type: "text", Text: tool.Content[0].(provider.TextBlock).Text}}}
79 })
80 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "one", Name: "ask", Args: json.RawMessage(`{}`)})
81 call := toolCallFrame(t, frames)
82 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: call.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "answer"}}})
83 var result extproto.ToolResultFromHost
84 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || result.IsError || result.Content[0].Text != "answer" {
85 t.Fatalf("round trip: %+v, %v", result, err)
86 }
87}
88
89func TestNestedCallAcrossHostAndExtension(t *testing.T) {
90 m, ext, frames, replies := toolPeer(t)
91 ext.callTool = true
92 m.SetToolCaller(func(ctx context.Context, _, name string, args json.RawMessage) extproto.ToolResultFromHost {
93 if name == "read" {
94 chain, _ := ctx.Value(toolChainKey{}).([]string)
95 if len(chain) != 2 || chain[0] != "ask" || chain[1] != "read" {
96 t.Errorf("nested chain: %v", chain)
97 }
98 return extproto.ToolResultFromHost{Content: []extproto.ContentBlock{{Type: "text", Text: "loaded"}}}
99 }
100 result, err := NewTool(m, ToolInfo{Name: name, Extension: ext.Manifest.Name}).Execute(ctx, args, nil)
101 if err != nil || result.IsError {
102 return extproto.ToolResultFromHost{IsError: true}
103 }
104 return extproto.ToolResultFromHost{Content: []extproto.ContentBlock{{Type: "text", Text: "done"}}}
105 })
106 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "root", Name: "ask", Args: json.RawMessage(`{}`)})
107 parent := toolCallFrame(t, frames)
108 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "cycle", ParentID: parent.ID, Name: "ask", Args: json.RawMessage(`{}`)})
109 var refused extproto.ToolResultFromHost
110 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &refused); err != nil || refused.ID != "cycle" || !refused.IsError || !strings.Contains(refused.Content[0].Text, "cycle") {
111 t.Fatalf("cycle: %+v, %v", refused, err)
112 }
113 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "child", ParentID: parent.ID, Name: "read", Args: json.RawMessage(`{}`)})
114 var child extproto.ToolResultFromHost
115 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &child); err != nil || child.ID != "child" || child.IsError || child.Content[0].Text != "loaded" {
116 t.Fatalf("child: %+v, %v", child, err)
117 }
118 sendToolFrame(t, replies, extproto.ToolResultFromExt{Type: "tool_result", ID: parent.ID, Content: []extproto.ContentBlock{{Type: "text", Text: "ok"}}})
119 var root extproto.ToolResultFromHost
120 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &root); err != nil || root.ID != "root" || root.IsError || root.Content[0].Text != "done" {
121 t.Fatalf("root: %+v, %v", root, err)
122 }
123}
124
125func TestNestedCallUsesParentContext(t *testing.T) {
126 m, ext, frames, replies := toolPeer(t)
127 ext.callTool = true
128 parentCtx, cancel := context.WithCancel(context.Background())
129 defer cancel()
130 ext.mu.Lock()
131 ext.outboundCalls = map[string]outboundCall{"parent": {ctx: parentCtx, chain: []string{"ask"}}}
132 ext.mu.Unlock()
133 started := make(chan context.Context, 1)
134 m.SetToolCaller(func(ctx context.Context, _, _ string, _ json.RawMessage) extproto.ToolResultFromHost {
135 started <- ctx
136 <-ctx.Done()
137 return extproto.ToolResultFromHost{IsError: true, Content: []extproto.ContentBlock{{Type: "text", Text: ctx.Err().Error()}}}
138 })
139 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "child", ParentID: "parent", Name: "read", Args: json.RawMessage(`{}`)})
140 select {
141 case ctx := <-started:
142 if chain, _ := ctx.Value(toolChainKey{}).([]string); len(chain) != 2 || chain[0] != "ask" || chain[1] != "read" {
143 t.Fatalf("chain: %v", chain)
144 }
145 case <-time.After(time.Second):
146 t.Fatal("nested call did not start")
147 }
148 cancel()
149 var result extproto.ToolResultFromHost
150 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "canceled") {
151 t.Fatalf("parent cancellation: %+v, %v", result, err)
152 }
153}
154
155func TestCallToolCancel(t *testing.T) {
156 m, ext, frames, replies := toolPeer(t)
157 ext.callTool = true
158 started := make(chan struct{})
159 m.SetToolCaller(func(ctx context.Context, _, _ string, _ json.RawMessage) extproto.ToolResultFromHost {
160 close(started)
161 <-ctx.Done()
162 return extproto.ToolResultFromHost{IsError: true, Content: []extproto.ContentBlock{{Type: "text", Text: ctx.Err().Error()}}}
163 })
164 sendToolFrame(t, replies, extproto.CallToolFromExt{Type: "call_tool", ID: "one", Name: "read", Args: json.RawMessage(`{}`)})
165 select {
166 case <-started:
167 case <-time.After(time.Second):
168 t.Fatal("call did not start")
169 }
170 sendToolFrame(t, replies, extproto.Frame{Type: "call_tool_cancel", ID: "one"})
171 var result extproto.ToolResultFromHost
172 if err := json.Unmarshal(nextLifecycleFrame(t, frames), &result); err != nil || !result.IsError || !strings.Contains(result.Content[0].Text, "canceled") {
173 t.Fatalf("cancel: %+v, %v", result, err)
174 }
175}