interactive_tool_test.go RawHistoryFolder

1package ext
2
3import (
4	"context"
5	"encoding/json"
6	"io"
7	"slices"
8	"testing"
9	"time"
10
11	"github.com/patriceckhart/zot/packages/agent/extproto"
12)
13
14func TestInteractiveToolLifecycle(t *testing.T) {
15	for _, end := range []string{"cancel", "shutdown", "eof", "complete"} {
16		t.Run(end, func(t *testing.T) {
17			h := newHarness("interactive")
18			started := make(chan context.Context, 2)
19			finished := make(chan struct{}, 2)
20			answer := make(chan struct{})
21			h.ext.InteractiveTool("ask", "Ask the user", json.RawMessage(`{}`), func(ctx context.Context, _ json.RawMessage) ToolResult {
22				started <- ctx
23				defer func() { finished <- struct{}{} }()
24				select {
25				case <-ctx.Done():
26					return TextErrorResult("cancelled")
27				case <-answer:
28					return TextResult("answer")
29				}
30			})
31			done := make(chan error, 1)
32			go func() { done <- h.ext.Run() }()
33			t.Cleanup(func() {
34				h.hostW.Close()
35				select {
36				case err := <-done:
37					if err != nil {
38						t.Error(err)
39					}
40				case <-time.After(time.Second):
41					t.Error("SDK did not exit")
42				}
43				h.ext.out.(io.Closer).Close()
44				h.ext.in.(io.Closer).Close()
45			})
46			var hello extproto.HelloFromExt
47			if err := json.Unmarshal(h.next(t).raw, &hello); err != nil {
48				t.Fatal(err)
49			}
50			if hello.Type != "hello" || !slices.Contains(hello.Capabilities, "tool_cancel") {
51				t.Fatalf("SDK did not advertise cancellation support: %+v", hello)
52			}
53			h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack"})
54			var registration extproto.RegisterToolFromExt
55			json.Unmarshal(h.drainUntil(t, "register_tool").raw, &registration)
56			if !registration.Interactive || registration.Name != "ask" {
57				t.Fatalf("registration: %+v", registration)
58			}
59			h.drainUntil(t, "ready")
60			h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "one", Name: "ask", Args: json.RawMessage(`{}`)})
61			// Send cancellation without waiting for the goroutine to start.
62			if end == "cancel" {
63				h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
64			}
65			var ctx context.Context
66			select {
67			case ctx = <-started:
68			case <-time.After(time.Second):
69				t.Fatal("handler did not start")
70			}
71			switch end {
72			case "shutdown":
73				h.sendToExt(t, extproto.ShutdownFromHost{Type: "shutdown"})
74				h.drainUntil(t, "shutdown_ack")
75			case "eof":
76				h.hostW.Close()
77			case "complete":
78				close(answer)
79				var result extproto.ToolResultFromExt
80				json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
81				if result.ID != "one" || len(result.Content) != 1 || result.Content[0].Text != "answer" {
82					t.Fatalf("result: %+v", result)
83				}
84			}
85			select {
86			case <-finished:
87			case <-time.After(time.Second):
88				t.Fatal("handler did not finish")
89			}
90			if end != "complete" && ctx.Err() != context.Canceled {
91				t.Fatal("handler context was not cancelled")
92			}
93			if end == "cancel" {
94				// Unknown and duplicate cancellation must not affect another call.
95				h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "one"})
96				h.sendToExt(t, extproto.ToolCancelFromHost{Type: "tool_cancel", ID: "unknown"})
97				h.sendToExt(t, extproto.ToolCallFromHost{Type: "tool_call", ID: "two", Name: "ask", Args: json.RawMessage(`{}`)})
98				select {
99				case ctx = <-started:
100				case <-time.After(time.Second):
101					t.Fatal("second handler did not start")
102				}
103				if ctx.Err() != nil {
104					t.Fatal("second call was cancelled")
105				}
106				close(answer)
107				var result extproto.ToolResultFromExt
108				json.Unmarshal(h.drainUntil(t, "tool_result").raw, &result)
109				if result.ID != "two" {
110					t.Fatalf("unexpected result: %+v", result)
111				}
112			}
113		})
114	}
115}