1package server
2
3import (
4 "context"
5 "encoding/json"
6 "log/slog"
7 "net/http"
8 "net/http/httptest"
9 "strings"
10 "testing"
11 "time"
12
13 "shelley.exe.dev/claudetool"
14 "shelley.exe.dev/db"
15 "shelley.exe.dev/db/generated"
16 "shelley.exe.dev/llm"
17 "shelley.exe.dev/loop"
18)
19
20// TestMessageQueuedDuringThinking tests that messages sent while the LLM is
21// processing (thinking/tool execution) are properly queued and eventually processed.
22func TestMessageQueuedDuringThinking(t *testing.T) {
23 database, cleanup := setupTestDB(t)
24 defer cleanup()
25
26 predictableService := loop.NewPredictableService()
27 llmManager := &testLLMManager{service: predictableService}
28 logger := slog.Default()
29 server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
30
31 // Create conversation
32 conversation, err := database.CreateConversation(context.Background(), nil, true, nil, nil)
33 if err != nil {
34 t.Fatalf("failed to create conversation: %v", err)
35 }
36 conversationID := conversation.ConversationID
37
38 // Send first message that triggers a slow response via "delay:" prefix
39 chatReq := ChatRequest{
40 Message: "delay: 2",
41 Model: "predictable",
42 }
43 chatBody, _ := json.Marshal(chatReq)
44
45 req := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/chat", strings.NewReader(string(chatBody)))
46 req.Header.Set("Content-Type", "application/json")
47 w := httptest.NewRecorder()
48
49 server.handleChatConversation(w, req, conversationID)
50 if w.Code != http.StatusAccepted {
51 t.Fatalf("expected status 202 for first message, got %d: %s", w.Code, w.Body.String())
52 }
53
54 // Wait for the LLM to start processing (but still be in the delay)
55 time.Sleep(200 * time.Millisecond)
56
57 // Now send a SECOND message while the first is still processing
58 // This is the bug: this message should be immediately recorded and visible,
59 // not lost until the first message finishes processing
60 secondReq := ChatRequest{
61 Message: "echo: second message while thinking",
62 Model: "predictable",
63 }
64 secondBody, _ := json.Marshal(secondReq)
65
66 req2 := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/chat", strings.NewReader(string(secondBody)))
67 req2.Header.Set("Content-Type", "application/json")
68 w2 := httptest.NewRecorder()
69
70 server.handleChatConversation(w2, req2, conversationID)
71 if w2.Code != http.StatusAccepted {
72 t.Fatalf("expected status 202 for second message, got %d: %s", w2.Code, w2.Body.String())
73 }
74
75 // The second message should be recorded in the database IMMEDIATELY
76 // (or at least very soon), not waiting for the first message to finish
77 // Wait a short time for the message to be recorded
78 time.Sleep(100 * time.Millisecond)
79
80 var messages []generated.Message
81 err = database.Queries(context.Background(), func(q *generated.Queries) error {
82 var qerr error
83 messages, qerr = q.ListMessages(context.Background(), conversationID)
84 return qerr
85 })
86 if err != nil {
87 t.Fatalf("failed to get messages: %v", err)
88 }
89
90 // Look for the second user message in the database
91 foundSecondUserMessage := false
92 for _, msg := range messages {
93 if msg.Type != string(db.MessageTypeUser) {
94 continue
95 }
96 if msg.LlmData == nil {
97 continue
98 }
99 var llmMsg llm.Message
100 if err := json.Unmarshal([]byte(*msg.LlmData), &llmMsg); err != nil {
101 continue
102 }
103 for _, content := range llmMsg.Content {
104 if content.Type == llm.ContentTypeText && strings.Contains(content.Text, "second message while thinking") {
105 foundSecondUserMessage = true
106 break
107 }
108 }
109 }
110
111 if !foundSecondUserMessage {
112 t.Error("BUG: second user message sent during LLM processing was not immediately recorded to database")
113 t.Logf("Found %d messages total:", len(messages))
114 for i, msg := range messages {
115 t.Logf(" Message %d: type=%s", i, msg.Type)
116 }
117 }
118
119 // Wait for everything to complete
120 deadline := time.Now().Add(5 * time.Second)
121 for time.Now().Before(deadline) {
122 err = database.Queries(context.Background(), func(q *generated.Queries) error {
123 var qerr error
124 messages, qerr = q.ListMessages(context.Background(), conversationID)
125 return qerr
126 })
127 if err != nil {
128 t.Fatalf("failed to get messages: %v", err)
129 }
130 // Look for response to second message
131 for _, msg := range messages {
132 if msg.Type != string(db.MessageTypeAgent) || msg.LlmData == nil {
133 continue
134 }
135 var llmMsg llm.Message
136 if err := json.Unmarshal([]byte(*msg.LlmData), &llmMsg); err != nil {
137 continue
138 }
139 for _, content := range llmMsg.Content {
140 if content.Type == llm.ContentTypeText && strings.Contains(content.Text, "second message while thinking") {
141 // Found the response
142 return
143 }
144 }
145 }
146 time.Sleep(100 * time.Millisecond)
147 }
148 t.Error("timed out waiting for response to second message")
149}
150
151// TestContextPreservedAfterCancel tests that conversation context is properly
152// preserved after cancellation and the conversation can be resumed correctly.
153func TestContextPreservedAfterCancel(t *testing.T) {
154 database, cleanup := setupTestDB(t)
155 defer cleanup()
156
157 predictableService := loop.NewPredictableService()
158 llmManager := &testLLMManager{service: predictableService}
159 logger := slog.Default()
160 server := NewServer(database, llmManager, claudetool.ToolSetConfig{}, logger, true, "", "predictable", "", nil)
161
162 // Create conversation
163 conversation, err := database.CreateConversation(context.Background(), nil, true, nil, nil)
164 if err != nil {
165 t.Fatalf("failed to create conversation: %v", err)
166 }
167 conversationID := conversation.ConversationID
168
169 // Send first message and let it complete
170 chatReq := ChatRequest{
171 Message: "echo: initial context message",
172 Model: "predictable",
173 }
174 chatBody, _ := json.Marshal(chatReq)
175
176 req := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/chat", strings.NewReader(string(chatBody)))
177 req.Header.Set("Content-Type", "application/json")
178 w := httptest.NewRecorder()
179
180 server.handleChatConversation(w, req, conversationID)
181 if w.Code != http.StatusAccepted {
182 t.Fatalf("expected status 202, got %d: %s", w.Code, w.Body.String())
183 }
184
185 // Wait for first message to complete
186 time.Sleep(300 * time.Millisecond)
187
188 // Now start a slow operation and cancel it
189 slowReq := ChatRequest{
190 Message: "bash: sleep 5",
191 Model: "predictable",
192 }
193 slowBody, _ := json.Marshal(slowReq)
194
195 req2 := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/chat", strings.NewReader(string(slowBody)))
196 req2.Header.Set("Content-Type", "application/json")
197 w2 := httptest.NewRecorder()
198
199 server.handleChatConversation(w2, req2, conversationID)
200 if w2.Code != http.StatusAccepted {
201 t.Fatalf("expected status 202, got %d: %s", w2.Code, w2.Body.String())
202 }
203
204 // Wait for tool to start
205 time.Sleep(200 * time.Millisecond)
206
207 // Cancel the conversation
208 cancelReq := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/cancel", nil)
209 cancelW := httptest.NewRecorder()
210 server.handleCancelConversation(cancelW, cancelReq, conversationID)
211
212 if cancelW.Code != http.StatusOK {
213 t.Fatalf("expected cancel status 200, got %d: %s", cancelW.Code, cancelW.Body.String())
214 }
215
216 // Wait for cancellation to complete
217 time.Sleep(200 * time.Millisecond)
218
219 // Clear the predictable service request history so we can inspect the next request
220 predictableService.ClearRequests()
221
222 // Resume the conversation
223 resumeReq := ChatRequest{
224 Message: "echo: after cancel",
225 Model: "predictable",
226 }
227 resumeBody, _ := json.Marshal(resumeReq)
228
229 req3 := httptest.NewRequest("POST", "/api/conversation/"+conversationID+"/chat", strings.NewReader(string(resumeBody)))
230 req3.Header.Set("Content-Type", "application/json")
231 w3 := httptest.NewRecorder()
232
233 server.handleChatConversation(w3, req3, conversationID)
234 if w3.Code != http.StatusAccepted {
235 t.Fatalf("expected status 202 for resume, got %d: %s", w3.Code, w3.Body.String())
236 }
237
238 // Wait for the request to be processed
239 time.Sleep(300 * time.Millisecond)
240
241 // Check that the LLM request included the conversation history
242 lastReq := predictableService.GetLastRequest()
243 if lastReq == nil {
244 t.Fatal("BUG: no LLM request was made after resume")
245 }
246
247 // The request should include ALL previous messages:
248 // 1. Initial context message (user)
249 // 2. Response to initial context (assistant)
250 // 3. bash: sleep 5 (user)
251 // 4. Assistant response with tool use
252 // 5. Cancelled tool result (user)
253 // 6. [Operation cancelled] (assistant)
254 // 7. echo: after cancel (user)
255 //
256 // If context is lost, we'll only have the last message (#7)
257
258 if len(lastReq.Messages) < 3 {
259 t.Errorf("BUG: context lost after cancellation! Expected at least 3 messages in LLM request, got %d", len(lastReq.Messages))
260 t.Log("Messages in request:")
261 for i, msg := range lastReq.Messages {
262 t.Logf(" Message %d: role=%s, content_count=%d", i, msg.Role, len(msg.Content))
263 for j, content := range msg.Content {
264 if content.Type == llm.ContentTypeText {
265 // Truncate long text
266 text := content.Text
267 if len(text) > 100 {
268 text = text[:100] + "..."
269 }
270 t.Logf(" Content %d: type=%s, text=%q", j, content.Type, text)
271 } else {
272 t.Logf(" Content %d: type=%s", j, content.Type)
273 }
274 }
275 }
276 }
277
278 // Check that "initial context message" appears somewhere in the history
279 foundInitialContext := false
280 for _, msg := range lastReq.Messages {
281 for _, content := range msg.Content {
282 if content.Type == llm.ContentTypeText && strings.Contains(content.Text, "initial context message") {
283 foundInitialContext = true
284 break
285 }
286 }
287 }
288
289 if !foundInitialContext {
290 t.Error("BUG: initial context message was not preserved after cancellation")
291 }
292}