git.iko.soy/crush
on main help sign in
test: initial test setup
commitcf42f88da036a51c8194e7ca05b6383523f9d92e authorkujtimiihoxha <kujtimii.h@gmail.com> dateTue 30 Sep 2025, 13:25 (1 year ago) parent477f846
7 files changed+326−11
compare with main
unifiedsplit
go.mod+2 go.sum+4 internal/agent/.env.sample+1 internal/agent/agent.go+30−11 internal/agent/agent_test.go+85 internal/agent/recorder_test.go+104 internal/agent/testdata/TestSessionSimpleAgent.yaml+100
go.mod+2
@@ -41,7 +41,9 @@ require (
4141 github.com/stretchr/testify v1.11.1
4242 github.com/tidwall/sjson v1.2.5
4343 github.com/zeebo/xxh3 v1.0.2
44+ go.yaml.in/yaml/v4 v4.0.0-rc.2
4445 google.golang.org/genai v1.26.0
46+ gopkg.in/dnaeon/go-vcr.v4 v4.0.6-0.20250923044825-7b4892dd3117
4547 gopkg.in/natefinch/lumberjack.v2 v2.2.1
4648 mvdan.cc/sh/v3 v3.12.1-0.20250902163504-3cf4fd5717a5
4749)
go.sum+4
@@ -367,6 +367,8 @@ go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mx
367367go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
368368go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
369369go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
370+go.yaml.in/yaml/v4 v4.0.0-rc.2 h1:/FrI8D64VSr4HtGIlUtlFMGsm7H7pWTbj6vOLVZcA6s=
371+go.yaml.in/yaml/v4 v4.0.0-rc.2/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=
370372golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
371373golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
372374golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
@@ -514,6 +516,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
514516gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
515517gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
516518gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
519+gopkg.in/dnaeon/go-vcr.v4 v4.0.6-0.20250923044825-7b4892dd3117 h1:fbE/sTnBb9UNfE8cJsOzrYYPqVWVHb7jWH4SI1W//cM=
520+gopkg.in/dnaeon/go-vcr.v4 v4.0.6-0.20250923044825-7b4892dd3117/go.mod h1:YuVT9NPq7t3oT2WpUemB0DbNL7djIjgajZycxoDLnqs=
517521gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
518522gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
519523gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
internal/agent/.env.samplenew file+1
@@ -0,0 +1 @@
1+CRUSH_ANTHROPIC_API_KEY=
internal/agent/agent.go+30−11
@@ -7,6 +7,7 @@ import (
77 "fmt"
88 "log/slog"
99 "strings"
10+ "sync"
1011 "time"
1112
1213 "github.com/charmbracelet/catwalk/pkg/catwalk"
@@ -57,13 +58,12 @@ type Model struct {
5758}
5859
5960type sessionAgent struct {
60− largeModel Model
61− smallModel Model
62− systemPrompt string
63− tools []ai.AgentTool
64− maxOutputTokens int64
65− sessions session.Service
66− messages message.Service
61+ largeModel Model
62+ smallModel Model
63+ systemPrompt string
64+ tools []ai.AgentTool
65+ sessions session.Service
66+ messages message.Service
6767
6868 messageQueue *csync.Map[string, []SessionAgentCall]
6969 activeRequests *csync.Map[string, context.CancelFunc]
@@ -71,8 +71,24 @@ type sessionAgent struct {
7171
7272type SessionAgentOption func(*sessionAgent)
7373
74−func NewSessionAgent() SessionAgent {
75− return &sessionAgent{}
74+func NewSessionAgent(
75+ largeModel Model,
76+ smallModel Model,
77+ systemPrompt string,
78+ sessions session.Service,
79+ messages message.Service,
80+ tools ...ai.AgentTool,
81+) SessionAgent {
82+ return &sessionAgent{
83+ largeModel: largeModel,
84+ smallModel: smallModel,
85+ systemPrompt: systemPrompt,
86+ sessions: sessions,
87+ messages: messages,
88+ tools: tools,
89+ messageQueue: csync.NewMap[string, []SessionAgentCall](),
90+ activeRequests: csync.NewMap[string, context.CancelFunc](),
91+ }
7692}
7793
7894func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.AgentResult, error) {
@@ -103,7 +119,6 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
103119 a.largeModel.model,
104120 ai.WithSystemPrompt(a.systemPrompt),
105121 ai.WithTools(a.tools...),
106− ai.WithMaxOutputTokens(a.maxOutputTokens),
107122 )
108123
109124 currentSession, err := a.sessions.Get(ctx, call.SessionID)
@@ -116,9 +131,12 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
116131 return nil, fmt.Errorf("failed to get session messages: %w", err)
117132 }
118133
134+ var wg sync.WaitGroup
119135 // Generate title if first message
120136 if len(msgs) == 0 {
121− go a.generateTitle(ctx, currentSession, call.Prompt)
137+ wg.Go(func() {
138+ a.generateTitle(ctx, currentSession, call.Prompt)
139+ })
122140 }
123141
124142 // Add the user message to the session
@@ -325,6 +343,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
325343 if err != nil {
326344 return nil, err
327345 }
346+ wg.Wait()
328347
329348 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
330349 if !ok || len(queuedMessages) == 0 {
internal/agent/agent_test.gonew file+85
@@ -0,0 +1,85 @@
1+package agent
2+
3+import (
4+ "database/sql"
5+ "net/http"
6+ "os"
7+ "testing"
8+
9+ "github.com/charmbracelet/catwalk/pkg/catwalk"
10+ "github.com/charmbracelet/crush/internal/db"
11+ "github.com/charmbracelet/crush/internal/message"
12+ "github.com/charmbracelet/crush/internal/session"
13+ "github.com/charmbracelet/fantasy/ai"
14+ "github.com/charmbracelet/fantasy/anthropic"
15+ "github.com/stretchr/testify/require"
16+ "gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
17+
18+ _ "github.com/joho/godotenv/autoload"
19+)
20+
21+type builderFunc func(r *recorder.Recorder) (ai.LanguageModel, error)
22+
23+func TestSessionSimpleAgent(t *testing.T) {
24+ r := newRecorder(t)
25+ sonnet, err := anthropicBuilder("claude-sonnet-4-5-20250929")(r)
26+ require.Nil(t, err)
27+ haiku, err := anthropicBuilder("claude-3-5-haiku-20241022")(r)
28+ require.Nil(t, err)
29+ agent, sessions, messages := testSessionAgent(t, sonnet, haiku, "You are a helpful assistant")
30+ session, err := sessions.Create(t.Context(), "New Session")
31+ require.Nil(t, err)
32+
33+ res, err := agent.Run(t.Context(), SessionAgentCall{
34+ Prompt: "Hello",
35+ SessionID: session.ID,
36+ MaxOutputTokens: 10000,
37+ })
38+
39+ require.Nil(t, err)
40+ require.NotNil(t, res)
41+
42+ t.Run("should create session messages", func(t *testing.T) {
43+ msgs, err := messages.List(t.Context(), session.ID)
44+ require.Nil(t, err)
45+ // Should have the agent and user message
46+ require.Equal(t, len(msgs), 2)
47+ })
48+}
49+
50+func anthropicBuilder(model string) builderFunc {
51+ return func(r *recorder.Recorder) (ai.LanguageModel, error) {
52+ provider := anthropic.New(
53+ anthropic.WithAPIKey(os.Getenv("CRUSH_ANTHROPIC_API_KEY")),
54+ anthropic.WithHTTPClient(&http.Client{Transport: r}),
55+ )
56+ return provider.LanguageModel(model)
57+ }
58+}
59+
60+func testDBConn(t *testing.T) (*sql.DB, error) {
61+ return db.Connect(t.Context(), t.TempDir())
62+}
63+
64+func testSessionAgent(t *testing.T, large, small ai.LanguageModel, systemPrompt string, tools ...ai.AgentTool) (SessionAgent, session.Service, message.Service) {
65+ conn, err := testDBConn(t)
66+ require.Nil(t, err)
67+ q := db.New(conn)
68+ sessions := session.NewService(q)
69+ messages := message.NewService(q)
70+
71+ largeModel := Model{
72+ model: large,
73+ config: catwalk.Model{
74+ // todo: add values
75+ },
76+ }
77+ smallModel := Model{
78+ model: large,
79+ config: catwalk.Model{
80+ // todo: add values
81+ },
82+ }
83+ agent := NewSessionAgent(largeModel, smallModel, systemPrompt, sessions, messages, tools...)
84+ return agent, sessions, messages
85+}
internal/agent/recorder_test.gonew file+104
@@ -0,0 +1,104 @@
1+package agent
2+
3+import (
4+ "bytes"
5+ "encoding/json"
6+ "io"
7+ "net/http"
8+ "path/filepath"
9+ "reflect"
10+ "strings"
11+ "testing"
12+
13+ "go.yaml.in/yaml/v4"
14+ "gopkg.in/dnaeon/go-vcr.v4/pkg/cassette"
15+ "gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
16+)
17+
18+func newRecorder(t *testing.T) *recorder.Recorder {
19+ cassetteName := filepath.Join("testdata", t.Name())
20+
21+ r, err := recorder.New(
22+ cassetteName,
23+ recorder.WithMode(recorder.ModeRecordOnce),
24+ recorder.WithMatcher(customMatcher(t)),
25+ recorder.WithMarshalFunc(marshalFunc),
26+ recorder.WithSkipRequestLatency(true), // disable sleep to simulate response time, makes tests faster
27+ recorder.WithHook(hookRemoveHeaders, recorder.AfterCaptureHook),
28+ )
29+ if err != nil {
30+ t.Fatalf("recorder: failed to create recorder: %v", err)
31+ }
32+
33+ t.Cleanup(func() {
34+ if err := r.Stop(); err != nil {
35+ t.Errorf("recorder: failed to stop recorder: %v", err)
36+ }
37+ })
38+
39+ return r
40+}
41+
42+func customMatcher(t *testing.T) recorder.MatcherFunc {
43+ return func(r *http.Request, i cassette.Request) bool {
44+ if r.Body == nil || r.Body == http.NoBody {
45+ return cassette.DefaultMatcher(r, i)
46+ }
47+ if r.Method != i.Method || r.URL.String() != i.URL {
48+ return false
49+ }
50+
51+ reqBody, err := io.ReadAll(r.Body)
52+ if err != nil {
53+ t.Fatalf("recorder: failed to read request body")
54+ }
55+ r.Body.Close()
56+ r.Body = io.NopCloser(bytes.NewBuffer(reqBody))
57+
58+ // Some providers can sometimes generate JSON requests with keys in
59+ // a different order, which means a direct string comparison will fail.
60+ // Falling back to deserializing the content if we don't have a match.
61+ if string(reqBody) == i.Body { // hot path
62+ return true
63+ }
64+ var content1, content2 any
65+ if err := json.Unmarshal(reqBody, &content1); err != nil {
66+ return false
67+ }
68+ if err := json.Unmarshal([]byte(i.Body), &content2); err != nil {
69+ return false
70+ }
71+ return reflect.DeepEqual(content1, content2)
72+ }
73+}
74+
75+func marshalFunc(in any) ([]byte, error) {
76+ var buff bytes.Buffer
77+ enc := yaml.NewEncoder(&buff)
78+ enc.SetIndent(2)
79+ enc.CompactSeqIndent()
80+ if err := enc.Encode(in); err != nil {
81+ return nil, err
82+ }
83+ return buff.Bytes(), nil
84+}
85+
86+var headersToKeep = map[string]struct{}{
87+ "accept": {},
88+ "content-type": {},
89+ "user-agent": {},
90+}
91+
92+func hookRemoveHeaders(i *cassette.Interaction) error {
93+ for k := range i.Request.Headers {
94+ if _, ok := headersToKeep[strings.ToLower(k)]; !ok {
95+ delete(i.Request.Headers, k)
96+ }
97+ }
98+ for k := range i.Response.Headers {
99+ if _, ok := headersToKeep[strings.ToLower(k)]; !ok {
100+ delete(i.Response.Headers, k)
101+ }
102+ }
103+ return nil
104+}
internal/agent/testdata/TestSessionSimpleAgent.yamlnew file+100
@@ -0,0 +1,100 @@
1+---
2+version: 2
3+interactions:
4+- id: 0
5+ request:
6+ proto: HTTP/1.1
7+ proto_major: 1
8+ proto_minor: 1
9+ content_length: 281
10+ host: ""
11+ body: '{"max_tokens":10000,"messages":[{"content":[{"text":"Hello","cache_control":{"type":"ephemeral"},"type":"text"}],"role":"user"}],"model":"claude-sonnet-4-5-20250929","system":[{"text":"You are a helpful assistant","cache_control":{"type":"ephemeral"},"type":"text"}],"stream":true}'
12+ headers:
13+ Accept:
14+ - application/json
15+ Content-Type:
16+ - application/json
17+ User-Agent:
18+ - Anthropic/Go 1.12.0
19+ url: https://api.anthropic.com/v1/messages
20+ method: POST
21+ response:
22+ proto: HTTP/2.0
23+ proto_major: 2
24+ proto_minor: 0
25+ content_length: -1
26+ body: |+
27+ event: message_start
28+ data: {"type":"message_start","message":{"id":"msg_01WGBTmd2Q5E2ajXUoHZYg6K","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":13,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":8,"service_tier":"standard"}} }
29+
30+ event: content_block_start
31+ data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }
32+
33+ event: content_block_delta
34+ data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello! How can I help you today"}}
35+
36+ event: content_block_delta
37+ data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"?"} }
38+
39+ event: content_block_stop
40+ data: {"type":"content_block_stop","index":0 }
41+
42+ event: message_delta
43+ data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":13,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":12} }
44+
45+ event: message_stop
46+ data: {"type":"message_stop" }
47+
48+ headers:
49+ Content-Type:
50+ - text/event-stream; charset=utf-8
51+ status: 200 OK
52+ code: 200
53+ duration: 1.891647s
54+- id: 1
55+ request:
56+ proto: HTTP/1.1
57+ proto_major: 1
58+ proto_minor: 1
59+ content_length: 621
60+ host: ""
61+ body: '{"max_tokens":40,"messages":[{"content":[{"text":"Generate a concise title for the following content:\n\nHello","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-5-20250929","system":[{"text":"you will generate a short title based on the first message a user begins a conversation with\n\n- ensure it is not more than 50 characters long\n- the title should be a summary of the user''s message\n- it should be one line long\n- do not use quotes or colons\n- the entire text you return will be used as the title\n- never return anything that is more than one sentence (one line) long\n","type":"text"}],"stream":true}'
62+ headers:
63+ Accept:
64+ - application/json
65+ Content-Type:
66+ - application/json
67+ User-Agent:
68+ - Anthropic/Go 1.12.0
69+ url: https://api.anthropic.com/v1/messages
70+ method: POST
71+ response:
72+ proto: HTTP/2.0
73+ proto_major: 2
74+ proto_minor: 0
75+ content_length: -1
76+ body: |+
77+ event: message_start
78+ data: {"type":"message_start","message":{"id":"msg_01CZb9drep7yMKkc2wzNDt5G","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":109,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":10,"service_tier":"standard"}} }
79+
80+ event: content_block_start
81+ data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }
82+
83+ event: content_block_delta
84+ data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Greeting or Starting a Conversation"} }
85+
86+ event: content_block_stop
87+ data: {"type":"content_block_stop","index":0}
88+
89+ event: message_delta
90+ data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":109,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":10} }
91+
92+ event: message_stop
93+ data: {"type":"message_stop" }
94+
95+ headers:
96+ Content-Type:
97+ - text/event-stream; charset=utf-8
98+ status: 200 OK
99+ code: 200
100+ duration: 2.547489s