subagent_test.go

  1package server
  2
  3import (
  4	"encoding/json"
  5	"strings"
  6	"testing"
  7
  8	"shelley.exe.dev/db/generated"
  9	"shelley.exe.dev/llm"
 10)
 11
 12func TestBuildConversationSummary(t *testing.T) {
 13	// Create a server to get a SubagentRunner
 14	runner := &SubagentRunner{server: nil} // server not needed for buildConversationSummary
 15
 16	// Create some mock messages
 17	userMsg := llm.Message{
 18		Role:    llm.MessageRoleUser,
 19		Content: []llm.Content{{Type: llm.ContentTypeText, Text: "Hello, please do task X"}},
 20	}
 21	userMsgJSON, _ := json.Marshal(userMsg)
 22	userMsgStr := string(userMsgJSON)
 23
 24	assistantMsg := llm.Message{
 25		Role:    llm.MessageRoleAssistant,
 26		Content: []llm.Content{{Type: llm.ContentTypeText, Text: "I'll start working on task X"}},
 27	}
 28	assistantMsgJSON, _ := json.Marshal(assistantMsg)
 29	assistantMsgStr := string(assistantMsgJSON)
 30
 31	messages := []generated.Message{
 32		{
 33			MessageID: "msg1",
 34			Type:      "user",
 35			LlmData:   &userMsgStr,
 36		},
 37		{
 38			MessageID: "msg2",
 39			Type:      "agent",
 40			LlmData:   &assistantMsgStr,
 41		},
 42	}
 43
 44	summary := runner.buildConversationSummary(messages)
 45
 46	// Check that the summary contains expected content
 47	if summary == "" {
 48		t.Error("Expected non-empty summary")
 49	}
 50	if !strings.Contains(summary, "Hello") {
 51		t.Error("Summary should contain user message content")
 52	}
 53	if !strings.Contains(summary, "task X") {
 54		t.Error("Summary should contain assistant message content")
 55	}
 56}
 57
 58func TestBuildConversationSummary_ToolUse(t *testing.T) {
 59	runner := &SubagentRunner{server: nil}
 60
 61	// Create a message with tool use
 62	toolUseMsg := llm.Message{
 63		Role: llm.MessageRoleAssistant,
 64		Content: []llm.Content{{
 65			Type:      llm.ContentTypeToolUse,
 66			ID:        "tool1",
 67			ToolName:  "bash",
 68			ToolInput: json.RawMessage(`{"command": "ls -la"}`),
 69		}},
 70	}
 71	toolUseMsgJSON, _ := json.Marshal(toolUseMsg)
 72	toolUseMsgStr := string(toolUseMsgJSON)
 73
 74	messages := []generated.Message{
 75		{
 76			MessageID: "msg1",
 77			Type:      "agent",
 78			LlmData:   &toolUseMsgStr,
 79		},
 80	}
 81
 82	summary := runner.buildConversationSummary(messages)
 83
 84	// Check that tool use is included
 85	if !strings.Contains(summary, "bash") {
 86		t.Error("Summary should contain tool name")
 87	}
 88	if !strings.Contains(summary, "ls -la") {
 89		t.Error("Summary should contain tool input")
 90	}
 91}
 92
 93func TestBuildConversationSummary_Truncation(t *testing.T) {
 94	runner := &SubagentRunner{server: nil}
 95
 96	// Create a message with very long content
 97	longText := make([]byte, 10000)
 98	for i := range longText {
 99		longText[i] = 'a'
100	}
101
102	userMsg := llm.Message{
103		Role:    llm.MessageRoleUser,
104		Content: []llm.Content{{Type: llm.ContentTypeText, Text: string(longText)}},
105	}
106	userMsgJSON, _ := json.Marshal(userMsg)
107	userMsgStr := string(userMsgJSON)
108
109	messages := []generated.Message{
110		{
111			MessageID: "msg1",
112			Type:      "user",
113			LlmData:   &userMsgStr,
114		},
115	}
116
117	summary := runner.buildConversationSummary(messages)
118
119	// Check that the summary is truncated
120	if !strings.Contains(summary, "[truncated]") {
121		t.Error("Expected truncation marker in long message")
122	}
123}
124
125func TestBuildConversationSummary_ToolResult(t *testing.T) {
126	runner := &SubagentRunner{server: nil}
127
128	// Create a message with tool result
129	toolResultMsg := llm.Message{
130		Role: llm.MessageRoleUser,
131		Content: []llm.Content{{
132			Type:      llm.ContentTypeToolResult,
133			ToolUseID: "tool1",
134			ToolError: false,
135			ToolResult: []llm.Content{{
136				Type: llm.ContentTypeText,
137				Text: "Command output: file1.txt file2.txt",
138			}},
139		}},
140	}
141	toolResultMsgJSON, _ := json.Marshal(toolResultMsg)
142	toolResultMsgStr := string(toolResultMsgJSON)
143
144	messages := []generated.Message{
145		{
146			MessageID: "msg1",
147			Type:      "user",
148			LlmData:   &toolResultMsgStr,
149		},
150	}
151
152	summary := runner.buildConversationSummary(messages)
153
154	// Check that tool result is included
155	if !strings.Contains(summary, "file1.txt") {
156		t.Error("Summary should contain tool result content")
157	}
158}
159
160func TestBuildConversationSummary_ToolError(t *testing.T) {
161	runner := &SubagentRunner{server: nil}
162
163	// Create a message with tool error
164	toolErrorMsg := llm.Message{
165		Role: llm.MessageRoleUser,
166		Content: []llm.Content{{
167			Type:      llm.ContentTypeToolResult,
168			ToolUseID: "tool1",
169			ToolError: true,
170			ToolResult: []llm.Content{{
171				Type: llm.ContentTypeText,
172				Text: "command not found: xyz",
173			}},
174		}},
175	}
176	toolErrorMsgJSON, _ := json.Marshal(toolErrorMsg)
177	toolErrorMsgStr := string(toolErrorMsgJSON)
178
179	messages := []generated.Message{
180		{
181			MessageID: "msg1",
182			Type:      "user",
183			LlmData:   &toolErrorMsgStr,
184		},
185	}
186
187	summary := runner.buildConversationSummary(messages)
188
189	// Check that error is marked
190	if !strings.Contains(summary, "(error)") {
191		t.Error("Summary should mark tool errors")
192	}
193}
194
195func TestBuildConversationSummary_SkipsSystemMessages(t *testing.T) {
196	runner := &SubagentRunner{server: nil}
197
198	systemMsg := llm.Message{
199		Role:    llm.MessageRoleUser,
200		Content: []llm.Content{{Type: llm.ContentTypeText, Text: "SECRET SYSTEM PROMPT CONTENT"}},
201	}
202	systemMsgJSON, _ := json.Marshal(systemMsg)
203	systemMsgStr := string(systemMsgJSON)
204
205	userMsg := llm.Message{
206		Role:    llm.MessageRoleUser,
207		Content: []llm.Content{{Type: llm.ContentTypeText, Text: "Regular user message"}},
208	}
209	userMsgJSON, _ := json.Marshal(userMsg)
210	userMsgStr := string(userMsgJSON)
211
212	messages := []generated.Message{
213		{
214			MessageID: "sys1",
215			Type:      "system",
216			LlmData:   &systemMsgStr,
217		},
218		{
219			MessageID: "msg1",
220			Type:      "user",
221			LlmData:   &userMsgStr,
222		},
223	}
224
225	summary := runner.buildConversationSummary(messages)
226
227	// System message should be excluded
228	if strings.Contains(summary, "SECRET") {
229		t.Error("Summary should not include system messages")
230	}
231	// User message should be included
232	if !strings.Contains(summary, "Regular user message") {
233		t.Error("Summary should include user messages")
234	}
235}