1package ext
2
3import (
4 "context"
5 "encoding/json"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/agent/extproto"
11)
12
13func TestCallToolRoundTrip(t *testing.T) {
14 h := newHarness("caller")
15 done := make(chan error, 1)
16 go func() { done <- h.ext.Run() }()
17 t.Cleanup(func() {
18 h.hostW.Close()
19 select {
20 case err := <-done:
21 if err != nil {
22 t.Error(err)
23 }
24 case <-time.After(time.Second):
25 t.Error("SDK did not exit")
26 }
27 })
28 h.drainUntil(t, "hello")
29 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", Capabilities: []string{"call_tool"}})
30 h.drainUntil(t, "ready")
31 result := make(chan ToolResult, 1)
32 errs := make(chan error, 1)
33 go func() {
34 r, err := h.ext.CallTool(context.Background(), "skill", map[string]string{"name": "test"})
35 result <- r
36 errs <- err
37 }()
38 frame := h.drainUntil(t, "call_tool")
39 var call extproto.CallToolFromExt
40 if err := json.Unmarshal(frame.raw, &call); err != nil || call.Name != "skill" || string(call.Args) != `{"name":"test"}` {
41 t.Fatalf("call: %+v, %v", call, err)
42 }
43 h.sendToExt(t, extproto.ToolResultFromHost{Type: "tool_result", ID: call.ID, Name: "skill", Content: []extproto.ContentBlock{{Type: "text", Text: "loaded"}}})
44 if err := <-errs; err != nil {
45 t.Fatal(err)
46 }
47 if r := <-result; len(r.Content) != 1 || r.Content[0].Text != "loaded" {
48 t.Fatalf("result: %+v", r)
49 }
50}
51
52func TestCallToolFromToolHandlerCarriesParent(t *testing.T) {
53 h := newHarness("caller")
54 h.ext.InteractiveTool("ask", "ask", json.RawMessage(`{}`), func(ctx context.Context, _ json.RawMessage) ToolResult {
55 res, err := h.ext.CallTool(ctx, "read", map[string]string{"path": "a"})
56 if err != nil {
57 return TextErrorResult(err.Error())
58 }
59 return res
60 })
61 done := make(chan error, 1)
62 go func() { done <- h.ext.Run() }()
63 t.Cleanup(func() {
64 h.hostW.Close()
65 select {
66 case err := <-done:
67 if err != nil {
68 t.Error(err)
69 }
70 case <-time.After(time.Second):
71 t.Error("SDK did not exit")
72 }
73 })
74 h.drainUntil(t, "hello")
75 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", Capabilities: []string{"call_tool"}})
76 h.drainUntil(t, "ready")
77 h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "parent", Name: "ask", Args: json.RawMessage(`{}`)})
78 var call extproto.CallToolFromExt
79 if err := json.Unmarshal(h.drainUntil(t, "call_tool").raw, &call); err != nil || call.ParentID != "parent" || call.Name != "read" {
80 t.Fatalf("nested request: %+v, %v", call, err)
81 }
82 h.sendToExt(t, extproto.ToolResultFromHost{Type: "tool_result", ID: call.ID, Name: "read", Content: []extproto.ContentBlock{{Type: "text", Text: "loaded"}}})
83 var reply extproto.ToolResultFromExt
84 if err := json.Unmarshal(h.drainUntil(t, "tool_result").raw, &reply); err != nil || reply.ID != "parent" || reply.Content[0].Text != "loaded" {
85 t.Fatalf("parent result: %+v, %v", reply, err)
86 }
87}
88
89func TestCallToolCancellation(t *testing.T) {
90 h := newHarness("caller")
91 done := make(chan error, 1)
92 go func() { done <- h.ext.Run() }()
93 t.Cleanup(func() {
94 h.hostW.Close()
95 select {
96 case err := <-done:
97 if err != nil {
98 t.Error(err)
99 }
100 case <-time.After(time.Second):
101 t.Error("SDK did not exit")
102 }
103 })
104 h.drainUntil(t, "hello")
105 h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", Capabilities: []string{"call_tool"}})
106 h.drainUntil(t, "ready")
107 ctx, cancel := context.WithCancel(context.Background())
108 result := make(chan error, 1)
109 go func() { _, err := h.ext.CallTool(ctx, "read", map[string]string{}); result <- err }()
110 var call extproto.CallToolFromExt
111 if err := json.Unmarshal(h.drainUntil(t, "call_tool").raw, &call); err != nil {
112 t.Fatal(err)
113 }
114 cancel()
115 frame := h.drainUntil(t, "call_tool_cancel")
116 if frame.hdr.ID != call.ID {
117 t.Fatalf("canceled ID = %q, want %q", frame.hdr.ID, call.ID)
118 }
119 select {
120 case err := <-result:
121 if err != context.Canceled {
122 t.Fatalf("cancel error: %v", err)
123 }
124 case <-time.After(time.Second):
125 t.Fatal("call did not cancel")
126 }
127}
128
129func TestCallToolUnsupportedHost(t *testing.T) {
130 e := New("caller", "test")
131 _, err := e.CallTool(context.Background(), "read", map[string]string{"path": "a"})
132 if err == nil || !strings.Contains(err.Error(), "does not support") {
133 t.Fatalf("unsupported host: %v", err)
134 }
135}