test: add coder agent

kujtimiihoxha created

Change summary

internal/agent/agent_test.go                            | 145 ++++++++--
internal/agent/coder.go                                 |  18 +
internal/agent/testdata/TestCoderAgent/simple_test.yaml |  69 +++++
internal/agent/testdata/TestSessionSimpleAgent.yaml     |  65 +++-
internal/config/load.go                                 |   7 
5 files changed, 257 insertions(+), 47 deletions(-)

Detailed changes

internal/agent/agent_test.go 🔗

@@ -1,14 +1,19 @@
 package agent
 
 import (
-	"database/sql"
 	"net/http"
 	"os"
 	"testing"
 
 	"github.com/charmbracelet/catwalk/pkg/catwalk"
+	"github.com/charmbracelet/crush/internal/agent/tools"
+	"github.com/charmbracelet/crush/internal/config"
+	"github.com/charmbracelet/crush/internal/csync"
 	"github.com/charmbracelet/crush/internal/db"
+	"github.com/charmbracelet/crush/internal/history"
+	"github.com/charmbracelet/crush/internal/lsp"
 	"github.com/charmbracelet/crush/internal/message"
+	"github.com/charmbracelet/crush/internal/permission"
 	"github.com/charmbracelet/crush/internal/session"
 	"github.com/charmbracelet/fantasy/ai"
 	"github.com/charmbracelet/fantasy/anthropic"
@@ -19,29 +24,72 @@ import (
 	_ "github.com/joho/godotenv/autoload"
 )
 
+type env struct {
+	workingDir  string
+	sessions    session.Service
+	messages    message.Service
+	permissions permission.Service
+	history     history.Service
+	lspClients  *csync.Map[string, *lsp.Client]
+}
+
 type builderFunc func(r *recorder.Recorder) (ai.LanguageModel, error)
 
-func TestSessionSimpleAgent(t *testing.T) {
-	r := newRecorder(t)
-	sonnet, err := anthropicBuilder("claude-sonnet-4-5-20250929")(r)
-	require.NoError(t, err)
-	haiku, err := anthropicBuilder("claude-3-5-haiku-20241022")(r)
-	require.NoError(t, err)
-	agent, sessions, messages := testSessionAgent(t, sonnet, haiku, "You are a helpful assistant")
-	session, err := sessions.Create(t.Context(), "New Session")
-	require.NoError(t, err)
+func TestSessionAgent(t *testing.T) {
+	t.Run("simple test", func(t *testing.T) {
+		r := newRecorder(t)
+		sonnet, err := anthropicBuilder("claude-sonnet-4-5-20250929")(r)
+		require.NoError(t, err)
+		haiku, err := anthropicBuilder("claude-3-5-haiku-20241022")(r)
+		require.NoError(t, err)
+
+		env := testEnv(t)
+		agent := testSessionAgent(env, sonnet, haiku, "You are a helpful assistant")
+		session, err := env.sessions.Create(t.Context(), "New Session")
+		require.NoError(t, err)
 
-	res, err := agent.Run(t.Context(), SessionAgentCall{
-		Prompt:          "Hello",
-		SessionID:       session.ID,
-		MaxOutputTokens: 10000,
+		res, err := agent.Run(t.Context(), SessionAgentCall{
+			Prompt:          "Hello",
+			SessionID:       session.ID,
+			MaxOutputTokens: 10000,
+		})
+
+		require.NoError(t, err)
+		assert.NotNil(t, res)
+
+		t.Run("should create session messages", func(t *testing.T) {
+			msgs, err := env.messages.List(t.Context(), session.ID)
+			require.NoError(t, err)
+			// Should have the agent and user message
+			assert.Equal(t, len(msgs), 2)
+		})
 	})
+}
 
-	require.NoError(t, err)
-	assert.NotNil(t, res)
+func TestCoderAgent(t *testing.T) {
+	t.Run("simple test", func(t *testing.T) {
+		r := newRecorder(t)
+		sonnet, err := anthropicBuilder("claude-sonnet-4-5-20250929")(r)
+		require.NoError(t, err)
+		haiku, err := anthropicBuilder("claude-3-5-haiku-20241022")(r)
+		require.NoError(t, err)
+
+		env := testEnv(t)
+		agent, err := coderAgent(env, sonnet, haiku)
+		require.NoError(t, err)
+		session, err := env.sessions.Create(t.Context(), "New Session")
+		require.NoError(t, err)
+
+		res, err := agent.Run(t.Context(), SessionAgentCall{
+			Prompt:          "Hello",
+			SessionID:       session.ID,
+			MaxOutputTokens: 10000,
+		})
 
-	t.Run("should create session messages", func(t *testing.T) {
-		msgs, err := messages.List(t.Context(), session.ID)
+		require.NoError(t, err)
+		assert.NotNil(t, res)
+
+		msgs, err := env.messages.List(t.Context(), session.ID)
 		require.NoError(t, err)
 		// Should have the agent and user message
 		assert.Equal(t, len(msgs), 2)
@@ -58,17 +106,27 @@ func anthropicBuilder(model string) builderFunc {
 	}
 }
 
-func testDBConn(t *testing.T) (*sql.DB, error) {
-	return db.Connect(t.Context(), t.TempDir())
-}
-
-func testSessionAgent(t *testing.T, large, small ai.LanguageModel, systemPrompt string, tools ...ai.AgentTool) (SessionAgent, session.Service, message.Service) {
-	conn, err := testDBConn(t)
-	require.Nil(t, err)
+func testEnv(t *testing.T) env {
+	workingDir := t.TempDir()
+	conn, err := db.Connect(t.Context(), t.TempDir())
+	require.NoError(t, err)
 	q := db.New(conn)
 	sessions := session.NewService(q)
 	messages := message.NewService(q)
+	permissions := permission.NewPermissionService(workingDir, true, []string{})
+	history := history.NewService(q, conn)
+	lspClients := csync.NewMap[string, *lsp.Client]()
+	return env{
+		workingDir,
+		sessions,
+		messages,
+		permissions,
+		history,
+		lspClients,
+	}
+}
 
+func testSessionAgent(env env, large, small ai.LanguageModel, systemPrompt string, tools ...ai.AgentTool) SessionAgent {
 	largeModel := Model{
 		model:  large,
 		config: catwalk.Model{
@@ -76,11 +134,42 @@ func testSessionAgent(t *testing.T, large, small ai.LanguageModel, systemPrompt
 		},
 	}
 	smallModel := Model{
-		model:  large,
+		model:  small,
 		config: catwalk.Model{
 			// todo: add values
 		},
 	}
-	agent := NewSessionAgent(largeModel, smallModel, systemPrompt, sessions, messages, tools...)
-	return agent, sessions, messages
+	agent := NewSessionAgent(largeModel, smallModel, systemPrompt, env.sessions, env.messages, tools...)
+	return agent
+}
+
+func coderAgent(env env, large, small ai.LanguageModel) (SessionAgent, error) {
+	prompt, err := coderPrompt()
+	if err != nil {
+		return nil, err
+	}
+	cfg, err := config.Init(env.workingDir, "", false)
+	if err != nil {
+		return nil, err
+	}
+
+	systemPrompt, err := prompt.Build(large.Provider(), large.Model(), *cfg)
+	if err != nil {
+		return nil, err
+	}
+	allTools := []ai.AgentTool{
+		tools.NewBashTool(env.permissions, env.workingDir, cfg.Options.Attribution),
+		tools.NewDownloadTool(env.permissions, env.workingDir),
+		tools.NewEditTool(env.lspClients, env.permissions, env.history, env.workingDir),
+		tools.NewMultiEditTool(env.lspClients, env.permissions, env.history, env.workingDir),
+		tools.NewFetchTool(env.permissions, env.workingDir),
+		tools.NewGlobTool(env.workingDir),
+		tools.NewGrepTool(env.workingDir),
+		tools.NewLsTool(env.permissions, env.workingDir),
+		tools.NewSourcegraphTool(),
+		tools.NewViewTool(env.lspClients, env.permissions, env.workingDir),
+		tools.NewWriteTool(env.lspClients, env.permissions, env.history, env.workingDir),
+	}
+
+	return testSessionAgent(env, large, small, systemPrompt, allTools...), nil
 }

internal/agent/coder.go 🔗

@@ -0,0 +1,18 @@
+package agent
+
+import (
+	_ "embed"
+
+	"github.com/charmbracelet/crush/internal/agent/prompt"
+)
+
+//go:embed templates/coder.gotmpl
+var coderPromptTmpl []byte
+
+func coderPrompt() (*prompt.Prompt, error) {
+	systemPrompt, err := prompt.NewPrompt("coder", string(coderPromptTmpl))
+	if err != nil {
+		return nil, err
+	}
+	return systemPrompt, nil
+}

internal/agent/testdata/TestCoderAgent/simple_test.yaml 🔗

@@ -0,0 +1,121 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 620
+    host: ""
+    body: '{"max_tokens":40,"messages":[{"content":[{"text":"Generate a concise title for the following content:\n\nHello","type":"text"}],"role":"user"}],"model":"claude-3-5-haiku-20241022","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}'
+    headers:
+      Accept:
+      - application/json
+      Content-Type:
+      - application/json
+      User-Agent:
+      - Anthropic/Go 1.12.0
+    url: https://api.anthropic.com/v1/messages
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1
+    body: |+
+      event: message_start
+      data: {"type":"message_start","message":{"id":"msg_01VNJNnqH4TvCQ2YLgKUQrRT","type":"message","role":"assistant","model":"claude-3-5-haiku-20241022","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":108,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard"}}}
+
+      event: content_block_start
+      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}              }
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"First"}      }
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" Contact"}}
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: content_block_stop
+      data: {"type":"content_block_stop","index":0       }
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: message_delta
+      data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":108,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5}            }
+
+      event: message_stop
+      data: {"type":"message_stop"               }
+
+    headers:
+      Content-Type:
+      - text/event-stream; charset=utf-8
+    status: 200 OK
+    code: 200
+    duration: 619.065458ms
+- id: 1
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 31989
+    host: ""

internal/agent/testdata/TestSessionSimpleAgent.yaml 🔗

@@ -6,9 +6,9 @@ interactions:
     proto: HTTP/1.1
     proto_major: 1
     proto_minor: 1
-    content_length: 281
+    content_length: 620
     host: ""
-    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}'
+    body: '{"max_tokens":40,"messages":[{"content":[{"text":"Generate a concise title for the following content:\n\nHello","type":"text"}],"role":"user"}],"model":"claude-3-5-haiku-20241022","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}'
     headers:
       Accept:
       - application/json
@@ -25,40 +25,55 @@ interactions:
     content_length: -1
     body: |+
       event: message_start
-      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"}}      }
+      data: {"type":"message_start","message":{"id":"msg_011Y36fsS8dYbr8fbjbwBHLA","type":"message","role":"assistant","model":"claude-3-5-haiku-20241022","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":108,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard"}}        }
 
       event: content_block_start
-      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}         }
+      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}      }
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"New"} }
 
       event: content_block_delta
-      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello! How can I help you today"}}
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" Conversation"}           }
+
+      event: ping
+      data: {"type": "ping"}
 
       event: content_block_delta
-      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"?"}    }
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" Starter"}             }
+
+      event: ping
+      data: {"type": "ping"}
 
       event: content_block_stop
-      data: {"type":"content_block_stop","index":0  }
+      data: {"type":"content_block_stop","index":0               }
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: ping
+      data: {"type": "ping"}
 
       event: message_delta
-      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}       }
+      data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":108,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":8}       }
 
       event: message_stop
-      data: {"type":"message_stop"        }
+      data: {"type":"message_stop" }
 
     headers:
       Content-Type:
       - text/event-stream; charset=utf-8
     status: 200 OK
     code: 200
-    duration: 1.891647s
+    duration: 583.077125ms
 - id: 1
   request:
     proto: HTTP/1.1
     proto_major: 1
     proto_minor: 1
-    content_length: 621
+    content_length: 281
     host: ""
-    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}'
+    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}'
     headers:
       Accept:
       - application/json
@@ -75,26 +90,38 @@ interactions:
     content_length: -1
     body: |+
       event: message_start
-      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"}}    }
+      data: {"type":"message_start","message":{"id":"msg_01VPFYRaiH21mFLcGUeUk3Gv","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":3,"service_tier":"standard"}}  }
 
       event: content_block_start
-      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}       }
+      data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}           }
+
+      event: content_block_delta
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello! How"}            }
 
       event: content_block_delta
-      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Greeting or Starting a Conversation"}           }
+      data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" can I help you today?"}         }
+
+      event: ping
+      data: {"type": "ping"}
 
       event: content_block_stop
-      data: {"type":"content_block_stop","index":0}
+      data: {"type":"content_block_stop","index":0     }
+
+      event: ping
+      data: {"type": "ping"}
+
+      event: ping
+      data: {"type": "ping"}
 
       event: message_delta
-      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}          }
+      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}      }
 
       event: message_stop
-      data: {"type":"message_stop"               }
+      data: {"type":"message_stop"     }
 
     headers:
       Content-Type:
       - text/event-stream; charset=utf-8
     status: 200 OK
     code: 200
-    duration: 2.547489s
+    duration: 1.658454625s

internal/config/load.go 🔗

@@ -348,6 +348,13 @@ func (c *Config) setDefaults(workingDir, dataDir string) {
 	if str, ok := os.LookupEnv("CRUSH_DISABLE_PROVIDER_AUTO_UPDATE"); ok {
 		c.Options.DisableProviderAutoUpdate, _ = strconv.ParseBool(str)
 	}
+
+	if c.Options.Attribution == nil {
+		c.Options.Attribution = &Attribution{
+			CoAuthoredBy:  true,
+			GeneratedWith: true,
+		}
+	}
 }
 
 // applyLSPDefaults applies default values from powernap to LSP configurations