1package agent
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "encoding/json"
8 "fmt"
9 "os"
10 "path/filepath"
11 "testing"
12 "time"
13
14 "github.com/patriceckhart/zot/packages/agent/extensions"
15 "github.com/patriceckhart/zot/packages/agent/extproto"
16 "github.com/patriceckhart/zot/packages/core"
17 "github.com/patriceckhart/zot/packages/provider"
18)
19
20func TestStartExtensionProcessHelper(t *testing.T) {
21 if os.Getenv("ZOT_TEST_START_EXTENSION") != "1" {
22 return
23 }
24 enc := json.NewEncoder(os.Stdout)
25 enc.Encode(map[string]any{"type": "hello", "name": "start-test", "version": "1"})
26 enc.Encode(extproto.SubscribeFromExt{Type: "subscribe", Intercept: []string{"before_agent_start"}})
27 enc.Encode(extproto.ReadyFromExt{Type: "ready"})
28 scanner := bufio.NewScanner(os.Stdin)
29 count := 0
30 for scanner.Scan() {
31 var ev extproto.EventInterceptFromHost
32 if json.Unmarshal(scanner.Bytes(), &ev) != nil {
33 os.Exit(2)
34 }
35 if ev.Type == "shutdown" {
36 os.Exit(0)
37 }
38 if ev.Event != "before_agent_start" {
39 continue
40 }
41 if ev.SystemPrompt == nil || ev.CWD == "" || ev.Provider == "" || ev.Model == "" || ev.SessionID == "" || ev.AgentRunID == "" {
42 os.Exit(3)
43 }
44 count++
45 value, _ := json.Marshal(fmt.Sprintf("%s/%d/%s", *ev.SystemPrompt, count, ev.Model))
46 enc.Encode(extproto.EventInterceptResponseFromExt{Type: "event_intercept_response", ID: ev.ID, SystemPrompt: value})
47 }
48 os.Exit(0)
49}
50
51type startCaptureClient struct{ requests []provider.Request }
52
53func (*startCaptureClient) Name() string { return "test" }
54func (c *startCaptureClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
55 c.requests = append(c.requests, req)
56 ch := make(chan provider.Event, 1)
57 ch <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "ok"}}}}
58 close(ch)
59 return ch, nil
60}
61
62func TestRPCBeforeAgentStart(t *testing.T) {
63 t.Setenv("ZOT_TEST_START_EXTENSION", "1")
64 exe, err := os.Executable()
65 if err != nil {
66 t.Fatal(err)
67 }
68 dir := t.TempDir()
69 manifest, _ := json.Marshal(map[string]any{"name": "start-test", "version": "1", "exec": exe, "args": []string{"-test.run=^TestStartExtensionProcessHelper$"}, "enabled": true})
70 if err := os.WriteFile(filepath.Join(dir, "extension.json"), manifest, 0600); err != nil {
71 t.Fatal(err)
72 }
73 models := provider.ModelsForProvider("anthropic")
74 if len(models) < 2 {
75 t.Fatal("need two catalog models")
76 }
77 m := extensions.New(t.TempDir(), dir, "test", "anthropic", models[0].ID, nil)
78 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
79 defer cancel()
80 defer m.Stop(time.Second)
81 if errs := m.LoadExplicit(ctx, []string{dir}); len(errs) > 0 {
82 t.Fatal(errs)
83 }
84 m.WaitForReady(time.Second)
85 client := &startCaptureClient{}
86 ag := core.NewAgent(client, models[0].ID, "complete base", nil)
87 wireNonInteractiveAgentExtHooks(ctx, ag, m)
88 var out bytes.Buffer
89 server := &rpcServer{agent: ag, provider: "anthropic", model: ag.Model, out: &out}
90 prompt := func(want string) {
91 t.Helper()
92 if err := ag.Prompt(ctx, "hello", nil, nil); err != nil {
93 t.Fatal(err)
94 }
95 if got := client.requests[len(client.requests)-1].System; got != want {
96 t.Fatalf("request system=%q want %q", got, want)
97 }
98 }
99 prompt("complete base/1/" + models[0].ID)
100 prompt("complete base/1/" + models[0].ID)
101 server.dispatch("clear", "1", nil)
102 prompt("complete base/2/" + models[0].ID)
103 raw, _ := json.Marshal(map[string]string{"model": models[1].ID})
104 server.dispatch("set_model", "2", raw)
105 prompt("complete base/3/" + models[1].ID)
106 prompt("complete base/3/" + models[1].ID)
107}