Refactor cache, dedupe code

Author: jboolean

contrib/mark3labs/mcp-go/hooks.go

@@ -69,25 +69,27 @@ func (h *Hooks) onBeforeInitialize(ctx context.Context, id any, request *mcp.Ini
 	h.spanCache.Set(id, taskSpan, ttlcache.DefaultTTL)
 }
 
-// onAfterInitialize is called after initialization is successfully executed.
-func (h *Hooks) onAfterInitialize(ctx context.Context, id any, request *mcp.InitializeRequest, result *mcp.InitializeResult) {
+// finishSpanWithIO retrieves a span from cache, annotates it with input/output, and finishes it.
+func finishSpanWithIO[Req any, Res any](h *Hooks, id any, request Req, result Res) {
 	if item := h.spanCache.Get(id); item != nil {
 		span := item.Value()
-		if taskSpan, ok := span.(*llmobs.TaskSpan); ok {
+		if annotator, ok := span.(textIOAnnotator); ok {
 			inputJSON, _ := json.Marshal(request)
-			var outputText string
-			if result != nil {
-				resultJSON, _ := json.Marshal(result)
-				outputText = string(resultJSON)
-			}
+			resultJSON, _ := json.Marshal(result)
+			outputText := string(resultJSON)
 
-			taskSpan.AnnotateTextIO(string(inputJSON), outputText)
-			taskSpan.Finish()
+			annotator.AnnotateTextIO(string(inputJSON), outputText)
+			span.Finish()
 		}
 		h.spanCache.Delete(id)
 	}
 }
 
+// onAfterInitialize is called after initialization is successfully executed.
+func (h *Hooks) onAfterInitialize(ctx context.Context, id any, request *mcp.InitializeRequest, result *mcp.InitializeResult) {
+	finishSpanWithIO(h, id, request, result)
+}
+
 // onBeforeCallTool is called before a tool is executed.
 func (h *Hooks) onBeforeCallTool(ctx context.Context, id any, request *mcp.CallToolRequest) {
 	toolSpan, _ := llmobs.StartToolSpan(ctx, request.Params.Name)
@@ -96,21 +98,7 @@ func (h *Hooks) onBeforeCallTool(ctx context.Context, id any, request *mcp.CallT
 
 // onAfterCallTool is called after a tool is successfully executed.
 func (h *Hooks) onAfterCallTool(ctx context.Context, id any, request *mcp.CallToolRequest, result *mcp.CallToolResult) {
-	if item := h.spanCache.Get(id); item != nil {
-		span := item.Value()
-		if toolSpan, ok := span.(*llmobs.ToolSpan); ok {
-			inputJSON, _ := json.Marshal(request)
-			var outputText string
-			if result != nil {
-				resultJSON, _ := json.Marshal(result)
-				outputText = string(resultJSON)
-			}
-
-			toolSpan.AnnotateTextIO(string(inputJSON), outputText)
-			toolSpan.Finish()
-		}
-		h.spanCache.Delete(id)
-	}
+	finishSpanWithIO(h, id, request, result)
 }
 
 // textIOAnnotator mirrors the internal textIOSpan interface

contrib/mark3labs/mcp-go/hooks_test.go

@@ -27,7 +27,7 @@ func TestNewHooks(t *testing.T) {
 
 	hooks := NewHooks()
 	assert.NotNil(t, hooks)
-	assert.NotNil(t, hooks.toolCache)
+	assert.NotNil(t, hooks.spanCache)
 
 	hooks.Stop()
 }
@@ -43,6 +43,8 @@ func TestAddHooks(t *testing.T) {
 	ddHooks.AddHooks(serverHooks)
 
 	// Verify hooks were added
+	assert.Len(t, serverHooks.OnBeforeInitialize, 1)
+	assert.Len(t, serverHooks.OnAfterInitialize, 1)
 	assert.Len(t, serverHooks.OnBeforeCallTool, 1)
 	assert.Len(t, serverHooks.OnAfterCallTool, 1)
 	assert.Len(t, serverHooks.OnError, 1)