Skip to content

Commit fdf534e

Browse files
MegaGrindStoneMegaGrindStone
committed
Test prompts functionality
1 parent 9aad1e8 commit fdf534e

File tree

5 files changed

+137
-152
lines changed

5 files changed

+137
-152
lines changed

‎cmd/server/main.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,6 @@ func main() {
9393
mux.HandleFunc("/", m.HandleHome)
9494
mux.HandleFunc("/chats", m.HandleChats)
9595
mux.HandleFunc("/refresh-title", m.HandleRefreshTitle)
96-
mux.HandleFunc("/use-prompt", m.HandleUsePrompt)
9796
mux.HandleFunc("/sse/messages", m.HandleSSE)
9897
mux.HandleFunc("/sse/chats", m.HandleSSE)
9998

‎internal/handlers/chat.go‎

Lines changed: 100 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,106 @@ func (m Main) HandleChats(w http.ResponseWriter, r *http.Request) {
173173
m.renderExistingChatResponse(w, messages, addedMessageIDs, am, aiMsgID)
174174
}
175175

176+
// HandleRefreshTitle handles requests to regenerate a chat title. It accepts POST requests with a chat_id
177+
// parameter, retrieves the first user message from the chat history, and uses the title generator to create
178+
// a new title. The handler updates the chat title in the database and returns the new title to be displayed.
179+
//
180+
// The function expects a "chat_id" form field identifying which chat's title should be refreshed.
181+
// After updating the database, it asynchronously notifies all connected clients through Server-Sent Events (SSE)
182+
// to maintain UI consistency across sessions while immediately returning the new title text to the requesting client.
183+
//
184+
// The function returns appropriate HTTP error responses for invalid methods, missing required fields,
185+
// or when no messages are found for title generation. On success, it returns just the title text to be
186+
// inserted into the targeted span element via HTMX.
187+
func (m Main) HandleRefreshTitle(w http.ResponseWriter, r *http.Request) {
188+
if r.Method != http.MethodPost {
189+
m.logger.Error("Method not allowed", slog.String("method", r.Method))
190+
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
191+
return
192+
}
193+
194+
chatID := r.FormValue("chat_id")
195+
if chatID == "" {
196+
m.logger.Error("Chat ID is required")
197+
http.Error(w, "Chat ID is required", http.StatusBadRequest)
198+
return
199+
}
200+
201+
// Get messages to find first user message
202+
messages, err := m.store.Messages(r.Context(), chatID)
203+
if err != nil {
204+
m.logger.Error("Failed to get messages",
205+
slog.String("chatID", chatID),
206+
slog.String(errLoggerKey, err.Error()))
207+
http.Error(w, err.Error(), http.StatusInternalServerError)
208+
return
209+
}
210+
211+
if len(messages) == 0 {
212+
m.logger.Error("No messages found for chat", slog.String("chatID", chatID))
213+
http.Error(w, "No messages found for chat", http.StatusNotFound)
214+
return
215+
}
216+
217+
// Find first user message for title generation
218+
var firstUserMessage string
219+
for _, msg := range messages {
220+
if msg.Role == models.RoleUser && len(msg.Contents) > 0 && msg.Contents[0].Type == models.ContentTypeText {
221+
firstUserMessage = msg.Contents[0].Text
222+
break
223+
}
224+
}
225+
226+
if firstUserMessage == "" {
227+
m.logger.Error("No user message found for title generation", slog.String("chatID", chatID))
228+
http.Error(w, "No user message found for title generation", http.StatusInternalServerError)
229+
return
230+
}
231+
232+
// Generate and update title
233+
title, err := m.titleGenerator.GenerateTitle(r.Context(), firstUserMessage)
234+
if err != nil {
235+
m.logger.Error("Error generating chat title",
236+
slog.String("message", firstUserMessage),
237+
slog.String(errLoggerKey, err.Error()))
238+
http.Error(w, "Failed to generate title", http.StatusInternalServerError)
239+
return
240+
}
241+
242+
updatedChat := models.Chat{
243+
ID: chatID,
244+
Title: title,
245+
}
246+
if err := m.store.UpdateChat(r.Context(), updatedChat); err != nil {
247+
m.logger.Error("Failed to update chat title",
248+
slog.String(errLoggerKey, err.Error()))
249+
http.Error(w, "Failed to update chat title", http.StatusInternalServerError)
250+
return
251+
}
252+
253+
// Update all clients via SSE asynchronously
254+
go func() {
255+
divs, err := m.chatDivs(chatID)
256+
if err != nil {
257+
m.logger.Error("Failed to generate chat divs",
258+
slog.String(errLoggerKey, err.Error()))
259+
return
260+
}
261+
262+
msg := sse.Message{
263+
Type: chatsSSEType,
264+
}
265+
msg.AppendData(divs)
266+
if err := m.sseSrv.Publish(&msg, chatsSSETopic); err != nil {
267+
m.logger.Error("Failed to publish chats",
268+
slog.String(errLoggerKey, err.Error()))
269+
}
270+
}()
271+
272+
// Return just the title text for HTMX to insert into the span
273+
fmt.Fprintf(w, "%s", title)
274+
}
275+
176276
// processPromptInput handles prompt-based inputs, extracting arguments and retrieving
177277
// prompt messages from the MCP client.
178278
func (m Main) processPromptInput(ctx context.Context, promptName, promptArgs string) ([]models.Message, string, error) {
@@ -334,106 +434,6 @@ func (m Main) renderExistingChatResponse(w http.ResponseWriter, messages []model
334434
}
335435
}
336436

337-
// HandleRefreshTitle handles requests to regenerate a chat title. It accepts POST requests with a chat_id
338-
// parameter, retrieves the first user message from the chat history, and uses the title generator to create
339-
// a new title. The handler updates the chat title in the database and returns the new title to be displayed.
340-
//
341-
// The function expects a "chat_id" form field identifying which chat's title should be refreshed.
342-
// After updating the database, it asynchronously notifies all connected clients through Server-Sent Events (SSE)
343-
// to maintain UI consistency across sessions while immediately returning the new title text to the requesting client.
344-
//
345-
// The function returns appropriate HTTP error responses for invalid methods, missing required fields,
346-
// or when no messages are found for title generation. On success, it returns just the title text to be
347-
// inserted into the targeted span element via HTMX.
348-
func (m Main) HandleRefreshTitle(w http.ResponseWriter, r *http.Request) {
349-
if r.Method != http.MethodPost {
350-
m.logger.Error("Method not allowed", slog.String("method", r.Method))
351-
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
352-
return
353-
}
354-
355-
chatID := r.FormValue("chat_id")
356-
if chatID == "" {
357-
m.logger.Error("Chat ID is required")
358-
http.Error(w, "Chat ID is required", http.StatusBadRequest)
359-
return
360-
}
361-
362-
// Get messages to find first user message
363-
messages, err := m.store.Messages(r.Context(), chatID)
364-
if err != nil {
365-
m.logger.Error("Failed to get messages",
366-
slog.String("chatID", chatID),
367-
slog.String(errLoggerKey, err.Error()))
368-
http.Error(w, err.Error(), http.StatusInternalServerError)
369-
return
370-
}
371-
372-
if len(messages) == 0 {
373-
m.logger.Error("No messages found for chat", slog.String("chatID", chatID))
374-
http.Error(w, "No messages found for chat", http.StatusNotFound)
375-
return
376-
}
377-
378-
// Find first user message for title generation
379-
var firstUserMessage string
380-
for _, msg := range messages {
381-
if msg.Role == models.RoleUser && len(msg.Contents) > 0 && msg.Contents[0].Type == models.ContentTypeText {
382-
firstUserMessage = msg.Contents[0].Text
383-
break
384-
}
385-
}
386-
387-
if firstUserMessage == "" {
388-
m.logger.Error("No user message found for title generation", slog.String("chatID", chatID))
389-
http.Error(w, "No user message found for title generation", http.StatusInternalServerError)
390-
return
391-
}
392-
393-
// Generate and update title
394-
title, err := m.titleGenerator.GenerateTitle(r.Context(), firstUserMessage)
395-
if err != nil {
396-
m.logger.Error("Error generating chat title",
397-
slog.String("message", firstUserMessage),
398-
slog.String(errLoggerKey, err.Error()))
399-
http.Error(w, "Failed to generate title", http.StatusInternalServerError)
400-
return
401-
}
402-
403-
updatedChat := models.Chat{
404-
ID: chatID,
405-
Title: title,
406-
}
407-
if err := m.store.UpdateChat(r.Context(), updatedChat); err != nil {
408-
m.logger.Error("Failed to update chat title",
409-
slog.String(errLoggerKey, err.Error()))
410-
http.Error(w, "Failed to update chat title", http.StatusInternalServerError)
411-
return
412-
}
413-
414-
// Update all clients via SSE asynchronously
415-
go func() {
416-
divs, err := m.chatDivs(chatID)
417-
if err != nil {
418-
m.logger.Error("Failed to generate chat divs",
419-
slog.String(errLoggerKey, err.Error()))
420-
return
421-
}
422-
423-
msg := sse.Message{
424-
Type: chatsSSEType,
425-
}
426-
msg.AppendData(divs)
427-
if err := m.sseSrv.Publish(&msg, chatsSSETopic); err != nil {
428-
m.logger.Error("Failed to publish chats",
429-
slog.String(errLoggerKey, err.Error()))
430-
}
431-
}()
432-
433-
// Return just the title text for HTMX to insert into the span
434-
fmt.Fprintf(w, "%s", title)
435-
}
436-
437437
func (m Main) newChat() (string, error) {
438438
newChat := models.Chat{
439439
ID: uuid.New().String(),

‎internal/handlers/home.go‎

Lines changed: 0 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
package handlers
22

33
import (
4-
"context"
5-
"encoding/json"
64
"fmt"
75
"log/slog"
86
"net/http"
@@ -111,38 +109,3 @@ func (m Main) HandleHome(w http.ResponseWriter, r *http.Request) {
111109
func (m Main) HandleSSE(w http.ResponseWriter, r *http.Request) {
112110
m.sseSrv.ServeHTTP(w, r)
113111
}
114-
115-
// HandleUsePrompt processes a prompt request and returns formatted text to be inserted in the message textarea.
116-
func (m Main) HandleUsePrompt(w http.ResponseWriter, r *http.Request) {
117-
if r.Method != http.MethodPost {
118-
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
119-
return
120-
}
121-
122-
var params mcp.GetPromptParams
123-
if err := json.NewDecoder(r.Body).Decode(&params); err != nil {
124-
m.logger.Error("Failed to decode prompt request", slog.String(errLoggerKey, err.Error()))
125-
http.Error(w, "Invalid request", http.StatusBadRequest)
126-
return
127-
}
128-
129-
clientIdx, ok := m.promptsMap[params.Name]
130-
if !ok {
131-
m.logger.Error("Prompt not found", slog.String("promptName", params.Name))
132-
http.Error(w, "Prompt not found", http.StatusNotFound)
133-
}
134-
135-
res, err := m.mcpClients[clientIdx].GetPrompt(context.Background(), params)
136-
if err != nil {
137-
m.logger.Error("Failed to get prompt", slog.String("promptName", params.Name), slog.String(errLoggerKey, err.Error()))
138-
http.Error(w, "Failed to get prompt", http.StatusInternalServerError)
139-
return
140-
}
141-
142-
w.Header().Set("Content-Type", "application/json")
143-
if err := json.NewEncoder(w).Encode(res); err != nil {
144-
m.logger.Error("Failed to encode prompt response", slog.String(errLoggerKey, err.Error()))
145-
http.Error(w, "Failed to encode prompt response", http.StatusInternalServerError)
146-
return
147-
}
148-
}

‎internal/handlers/main.go‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,10 @@ type Main struct {
5555
resources []mcp.Resource
5656
prompts []mcp.Prompt
5757

58-
promptsMap map[string]int // Map of prompt names to mcpClients index.
59-
toolsMap map[string]int // Map of tool names to mcpClients index.
60-
logger *slog.Logger
58+
promptsMap map[string]int // Map of prompt names to mcpClients index.
59+
resourcesMap map[string]int // Map of resource uri to mcpClients index.
60+
toolsMap map[string]int // Map of tool names to mcpClients index.
61+
logger *slog.Logger
6162
}
6263

6364
const (
@@ -91,6 +92,7 @@ func NewMain(
9192
resources := make([]mcp.Resource, 0, len(mcpClients))
9293
prompts := make([]mcp.Prompt, 0, len(mcpClients))
9394
pm := make(map[string]int)
95+
rm := make(map[string]int)
9496
tm := make(map[string]int)
9597
for i := range mcpClients {
9698
servers[i] = mcpClients[i].ServerInfo()
@@ -115,6 +117,9 @@ func NewMain(
115117
return Main{}, fmt.Errorf("failed to list resources from server %s: %w", serverName, err)
116118
}
117119
rs = listResources.Resources
120+
for _, resource := range rs {
121+
rm[resource.URI] = i
122+
}
118123
}
119124

120125
var ps []mcp.Prompt
@@ -159,6 +164,7 @@ func NewMain(
159164
store: store,
160165
mcpClients: mcpClients,
161166
promptsMap: pm,
167+
resourcesMap: rm,
162168
toolsMap: tm,
163169
logger: logger.With(slog.String("module", "main")),
164170
servers: servers,

‎internal/handlers/main_test.go‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -114,40 +114,57 @@ func TestHandleChats(t *testing.T) {
114114
tests := []struct {
115115
name string
116116
method string
117-
message string
118-
chatID string
117+
formData string
119118
wantStatus int
120119
}{
121120
{
122121
name: "Invalid method",
123122
method: http.MethodGet,
123+
formData: "",
124124
wantStatus: http.StatusMethodNotAllowed,
125125
},
126126
{
127-
name: "Empty message",
127+
name: "Empty message and no prompt",
128128
method: http.MethodPost,
129+
formData: "chat_id=",
129130
wantStatus: http.StatusBadRequest,
130131
},
131132
{
132-
name: "New chat",
133+
name: "New chat with message",
133134
method: http.MethodPost,
134-
message: "Hello",
135+
formData: "message=Hello",
135136
wantStatus: http.StatusOK,
136137
},
137138
{
138-
name: "Existing chat",
139+
name: "Existing chat with message",
139140
method: http.MethodPost,
140-
message: "Hello",
141-
chatID: "1",
141+
formData: "message=Hello&chat_id=1",
142142
wantStatus: http.StatusOK,
143143
},
144+
// Testing prompt functionality with invalid inputs
145+
{
146+
name: "Invalid prompt arguments",
147+
method: http.MethodPost,
148+
formData: `prompt_name=test_prompt&prompt_args=invalid_json`,
149+
wantStatus: http.StatusInternalServerError,
150+
},
151+
{
152+
name: "Prompt without args",
153+
method: http.MethodPost,
154+
formData: `prompt_name=test_prompt`,
155+
wantStatus: http.StatusInternalServerError,
156+
},
157+
{
158+
name: "Prompt not found",
159+
method: http.MethodPost,
160+
formData: `prompt_name=unknown_prompt&prompt_args={"key":"value"}`,
161+
wantStatus: http.StatusInternalServerError,
162+
},
144163
}
145164

146165
for _, tt := range tests {
147166
t.Run(tt.name, func(t *testing.T) {
148-
form := strings.NewReader(
149-
"message=" + tt.message + "&chat_id=" + tt.chatID,
150-
)
167+
form := strings.NewReader(tt.formData)
151168
req := httptest.NewRequest(tt.method, "/chat", form)
152169
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
153170
w := httptest.NewRecorder()

0 commit comments

Comments
 (0)