before_start_test.go RawHistoryFolder

1package ext
2
3import (
4	"encoding/json"
5	"io"
6	"testing"
7
8	"github.com/patriceckhart/zot/packages/agent/extproto"
9)
10
11func TestBeforeAgentStartSDK(t *testing.T) {
12	for _, mode := range []string{"replace", "empty", "omit"} {
13		t.Run(mode, func(t *testing.T) {
14			h := newHarness("start-ext")
15			h.ext.InterceptBeforeAgentStart(func(ev BeforeAgentStartEvent) *string {
16				if ev.SystemPrompt != "base" || ev.CWD != "/work" || ev.Provider != "p" || ev.Model != "m" || ev.SessionID != "s" || ev.AgentRunID != "r" {
17					t.Error("wrong event metadata")
18				}
19				if mode == "omit" {
20					return nil
21				}
22				value := ""
23				if mode == "replace" {
24					value = "  new\n"
25				}
26				return &value
27			})
28			done := make(chan struct{})
29			go func() { defer close(done); h.ext.Run() }()
30			t.Cleanup(func() { h.hostW.Close(); <-done; h.ext.out.(io.Closer).Close() })
31			h.drainUntil(t, "hello")
32			h.sendToExt(t, extproto.HelloAckFromHost{Type: "hello_ack", ProtocolVersion: extproto.ProtocolVersion})
33			sub := h.drainUntil(t, "subscribe")
34			var subscription extproto.SubscribeFromExt
35			if err := json.Unmarshal(sub.raw, &subscription); err != nil {
36				t.Fatal(err)
37			}
38			if len(subscription.Intercept) != 1 || subscription.Intercept[0] != "before_agent_start" {
39				t.Fatalf("subscription = %+v", subscription)
40			}
41			h.drainUntil(t, "ready")
42			base := "base"
43			h.sendToExt(t, extproto.EventInterceptFromHost{Type: "event_intercept", Event: "before_agent_start", ID: "1", SystemPrompt: &base, CWD: "/work", Provider: "p", Model: "m", SessionID: "s", AgentRunID: "r"})
44			frame := h.drainUntil(t, "event_intercept_response")
45			var result extproto.EventInterceptResponseFromExt
46			if err := json.Unmarshal(frame.raw, &result); err != nil {
47				t.Fatal(err)
48			}
49			if result.ID != "1" {
50				t.Fatal("missing correlation ID")
51			}
52			want := `"  new\n"`
53			if mode == "empty" {
54				want = `""`
55			}
56			if mode == "omit" {
57				want = ""
58			}
59			if string(result.SystemPrompt) != want {
60				t.Fatalf("got %s want %s", result.SystemPrompt, want)
61			}
62		})
63	}
64}