nit: Refactor for better structure

Author: 4ndrelim

internal/services/toolkit/client/client.go

@@ -54,12 +54,22 @@ func NewAIClient(
 
 	// Load tools dynamically from backend (TODO: Make URL configurable / Xtramcp url)
 	xtraMCPLoader := xtramcp.NewXtraMCPLoader(db, projectService, "http://localhost:8080/mcp")
-	err := xtraMCPLoader.LoadToolsFromBackend(toolRegistry)
+	
+	// initialize MCP session first and log session ID
+	sessionID, err := xtraMCPLoader.InitializeMCP()
 	if err != nil {
-		logger.Errorf("[AI Client] Failed to load XtraMCP tools: %v", err)
-		// Fallback to static tools or return error based on your preference
+		logger.Errorf("[AI Client] Failed to initialize XtraMCP session: %v", err)
+		// TODO: Fallback to static tools or exit?
 	} else {
-		logger.Info("[AI Client] Successfully loaded XtraMCP tools")
+		logger.Info("[AI Client] XtraMCP session initialized", "sessionID", sessionID)
+		
+		// dynamically load all tools from XtraMCP backend
+		err = xtraMCPLoader.LoadToolsFromBackend(toolRegistry)
+		if err != nil {
+			logger.Errorf("[AI Client] Failed to load XtraMCP tools: %v", err)
+		} else {
+			logger.Info("[AI Client] Successfully loaded XtraMCP tools")
+		}
 	}
 
 	toolCallHandler := handler.NewToolCallHandler(toolRegistry)

internal/services/toolkit/tools/xtramcp/loader.go

@@ -6,13 +6,14 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"paperdebugger/internal/libs/db"
 	"paperdebugger/internal/services"
 	"paperdebugger/internal/services/toolkit/registry"
 )
 
-// MCPToolsResponse represents the JSON-RPC response from your backend
-type MCPToolsResponse struct {
+// MCPListToolsResponse represents the JSON-RPC response from tools/list method
+type MCPListToolsResponse struct {
 	JSONRPC string `json:"jsonrpc"`
 	ID      int    `json:"id"`
 	Result  struct {
@@ -41,14 +42,11 @@ func NewXtraMCPLoader(db *db.DB, projectService *services.ProjectService, baseUR
 
 // LoadToolsFromBackend fetches tool schemas from backend and registers them
 func (loader *XtraMCPLoader) LoadToolsFromBackend(toolRegistry *registry.ToolRegistry) error {
-	// Initialize MCP session ONCE
-	sessionID, err := loader.initializeMCP()
-	if err != nil {
-		return fmt.Errorf("failed to initialize MCP: %w", err)
+	if loader.sessionID == "" {
+		return fmt.Errorf("MCP session not initialized - call InitializeMCP first")
 	}
-	loader.sessionID = sessionID
 
-	// Fetch tools from backend using the session (currently returns mock data)
+	// Fetch tools from backend using the established session
 	toolSchemas, err := loader.fetchAvailableTools()
 	if err != nil {
 		return fmt.Errorf("failed to fetch tools from backend: %w", err)
@@ -67,8 +65,8 @@ func (loader *XtraMCPLoader) LoadToolsFromBackend(toolRegistry *registry.ToolReg
 	return nil
 }
 
-// initializeMCP performs the full MCP initialization handshake
-func (loader *XtraMCPLoader) initializeMCP() (string, error) {
+// InitializeMCP performs the full MCP initialization handshake, stores session ID, and returns it
+func (loader *XtraMCPLoader) InitializeMCP() (string, error) {
 	// Step 1: Initialize
 	sessionID, err := loader.performInitialize()
 	if err != nil {
@@ -81,6 +79,9 @@ func (loader *XtraMCPLoader) initializeMCP() (string, error) {
 		return "", fmt.Errorf("step 2 - notifications/initialized failed: %w", err)
 	}
 
+	// Store session ID for future use and return it
+	loader.sessionID = sessionID
+
 	return sessionID, nil
 }
 
@@ -189,15 +190,35 @@ func (loader *XtraMCPLoader) fetchAvailableTools() ([]ToolSchema, error) {
 	}
 	defer resp.Body.Close()
 
-	// Parse response
-	var mcpResponse MCPToolsResponse
-	err = json.NewDecoder(resp.Body).Decode(&mcpResponse)
+	// Read the raw response body (SSE format) for debugging
+	bodyBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	// Parse SSE format - extract JSON from "data: " lines
+	lines := strings.Split(string(bodyBytes), "\n")
+	var extractedJSON string
+	for _, line := range lines {
+		if strings.HasPrefix(line, "data: ") {
+			extractedJSON = strings.TrimPrefix(line, "data: ")
+			break
+		}
+	}
+
+	if extractedJSON == "" {
+		return nil, fmt.Errorf("no data line found in SSE response")
+	}
+
+	// Parse the extracted JSON
+	var mcpResponse MCPListToolsResponse
+	err = json.Unmarshal([]byte(extractedJSON), &mcpResponse)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse response: %w", err)
+		return nil, fmt.Errorf("failed to parse JSON from SSE data: %w. JSON data: %s", err, extractedJSON)
 	}
 
 	return mcpResponse.Result.Tools, nil
-	
+
 	// mock data; return hardcoded tool schemas for testing
 	// mockToolsJSON := `[{"name":"get_user_papers","description":"Fetch all papers published by a specific user identified by email. Supports 'summary' (abstract truncated to 150 words) and 'detailed' (full abstract).","inputSchema":{"properties":{"email":{"description":"Email address of the user whose papers to fetch. Must be a valid email string.","examples":["alice@example.com","bob@university.edu"],"title":"Email","type":"string"},"format":{"default":"detailed","description":"Format of the response. 'summary' shows title, venue, authors, URL, and the first 150 words of the abstract (default). 'detailed' shows the full abstract.","enum":["summary","detailed"],"examples":["summary","detailed"],"title":"Format","type":"string"}},"required":["email"],"title":"get_user_papers_toolArguments","type":"object"},"outputSchema":{"properties":{"result":{"title":"Result","type":"string"}},"required":["result"],"title":"get_user_papers_toolOutput","type":"object"}},{"name":"search_papers_on_openreview","description":"Search for academic papers on OpenReview by keywords within specific conference venues. This tool supports various matching modes and is ideal for discovering recent or broader papers beyond those available in the local database. Use this tool when results from search_relevant_papers are insufficient.","inputSchema":{"properties":{"query":{"description":"Keywords, topics, content, or a chunk of text to search for.","examples":["time series token merging","neural networks"],"title":"Query","type":"string"},"venues":{"description":"List of conference venues and years to search in. Each entry must be a dict with 'venue' and 'year'.","examples":[{"venue":"ICLR.cc","year":"2024"},{"venue":"ICML","year":"2024"},{"venue":"NeurIPS.cc","year":"2023"},{"venue":"NeurIPS.cc","year":"2022"}],"items":{"additionalProperties":{"type":"string"},"type":"object"},"minItems":1,"title":"Venues","type":"array"},"search_fields":{"default":["title","abstract"],"description":"Fields to search within each paper. Options: 'title', 'abstract', 'authors'.","items":{"enum":["title","abstract","authors"],"type":"string"},"title":"Search Fields","type":"array"},"match_mode":{"default":"majority","description":"Match mode:\n- any: At least one keyword must match\n- all: All keywords must match\n- exact: Match the entire phrase exactly\n- majority: Match majority of keywords (>50%)\n- threshold: Match percentage of terms based on 'match_threshold'.","enum":["any","all","exact","majority","threshold"],"title":"Match Mode","type":"string"},"match_threshold":{"default":0.5,"description":"Minimum fraction (0.0-1.0) of search terms that must match when using 'threshold' mode. Example: 0.5 = 50% of terms must match.","maximum":1,"minimum":0,"title":"Match Threshold","type":"number"},"limit":{"default":10,"description":"Maximum number of results to return (1-16).","maximum":16,"minimum":1,"title":"Limit","type":"integer"},"min_score":{"default":0.6,"description":"Minimum match score (0.0-1.0). Lower values allow looser matches; higher values enforce stricter matches.","maximum":1,"minimum":0,"title":"Min Score","type":"number"}},"required":["query","venues"],"title":"search_papers_openreview_toolArguments","type":"object"},"outputSchema":{"properties":{"result":{"title":"Result","type":"string"}},"required":["result"],"title":"search_papers_openreview_toolOutput","type":"object"}},{"name":"search_relevant_papers","description":"Search for similar or relevant papers by keywords against the local database of academic papers. This tool uses semantic search with vector embeddings to find the most relevant results. It is the default and recommended tool for paper searches.","inputSchema":{"properties":{"query":{"description":"Keywords, topics, content, or a chunk of text to search for.","examples":["time series token merging","neural networks","...when trained on first-order Markov chains, transformers with two or more layers consistently develop an induction head mechanism to estimate the in-context bigram conditional distribution"],"title":"Query","type":"string"},"top_k":{"description":"Number of top relevant or similar papers to return.","title":"Top K","type":"integer"},"date_min":{"description":"Minimum publication date (YYYY-MM-DD) to filter papers.","examples":["2023-01-01","2022-06-25"],"title":"Date Min","type":"string"},"date_max":{"description":"Maximum publication date (YYYY-MM-DD) to filter papers.","examples":["2024-12-31","2023-06-25"],"title":"Date Max","type":"string"},"countries":{"anyOf":[{"items":{"type":"string"},"type":"array"},{"type":"null"}],"description":"List of country codes in ISO ALPHA-3 format to filter papers by author affiliations.","examples":[["USA","CHN","SGP","GBR","DEU","KOR","JPN"]],"title":"Countries"},"min_similarity":{"description":"Minimum similarity score (0.0-1.0) for returned papers. Higher values yield more relevant results but fewer papers.","examples":[0.3,0.5,0.7,0.9],"title":"Min Similarity","type":"number"}},"required":["query","top_k","countries","min_similarity"],"title":"search_papers_toolArguments","type":"object"},"outputSchema":{"properties":{"result":{"title":"Result","type":"string"}},"required":["result"],"title":"search_papers_toolOutput","type":"object"}},{"name":"identify_improvements","description":"Analyzes a draft academic paper against the standards of top-tier ML conferences (ICLR, ICML, NeurIPS). Identifies issues in structure, completeness, clarity, and argumentation, then provides prioritized, actionable suggestions.","inputSchema":{"properties":{"paper_content":{"description":"The full text content of the academic paper draft. Paper content should not be truncated.","title":"Paper Content","type":"string"},"target_venue":{"default":"NeurIPS","description":"The target top-tier conference to tailor the feedback for.","enum":["ICLR","ICML","NeurIPS"],"title":"Target Venue","type":"string"},"focus_areas":{"anyOf":[{"items":{"enum":["Structure","Clarity","Evidence","Positioning","Style","Completeness","Soundness","Limitations"],"type":"string"},"type":"array"},{"type":"null"}],"default":null,"description":"List of specific areas to focus the analysis on. If empty, default areas are: {DEFAULT_FOCUS_AREAS}.","title":"Focus Areas"},"severity_threshold":{"default":"major","description":"The minimum severity level to report. 'major' will show blockers and major issues.","enum":["blocker","major","minor","nit"],"title":"Severity Threshold","type":"string"}},"required":["paper_content"],"title":"identify_improvementsArguments","type":"object"},"outputSchema":{"properties":{"result":{"title":"Result","type":"string"}},"required":["result"],"title":"identify_improvementsOutput","type":"object"}},{"name":"enhance_academic_writing","description":"Suggest context-aware academic paper writing enhancements for selected text.","inputSchema":{"properties":{"full_paper_content":{"description":"Surrounding context from the manuscript (e.g., abstract, background, or several sections). This need not be the entire paper; providing a substantial excerpt helps tailor the tone, terminology, and level of detail to academic venues (journals and conferences).","examples":["In reinforcement learning, one could structure these metrics (previously for evaluation) as rewards that could be boosted during training (Sharma et al., 2021; Yadav et al., 2021; Deng et al., 2022; Liu et al., 2023a; Xu et al., 2024; Wang et al., 2024b), to optimize complex objective functions even at testing time (OpenAI, 2024). However, when reward weights remain static, the weakest metric (the 'short-board') becomes a bottleneck that restricts overall LLM effectiveness, which introduces the short-board effect in multi-reward optimization. For example, in Figure 2, when the scaled reward itself (or its growth trend) has not yet reached saturation, its update magnitude should accordingly be increased."],"title":"Full Paper Content","type":"string"},"selected_content":{"description":"The specific text excerpt selected for improvement from the paper.","examples":["...when the scaled reward itself (or its growth trend) has not yet reached saturation, its update magnitude should accordingly be increased..."],"title":"Selected Content","type":"string"}},"required":["full_paper_content","selected_content"],"title":"improve_academic_passage_toolArguments","type":"object"},"outputSchema":{"properties":{"result":{"title":"Result","type":"string"}},"required":["result"],"title":"improve_academic_passage_toolOutput","type":"object"}}]`