Streaming communication using socket

Signed-off-by: Marek Safarik <[email protected]>

Author: Marek Safarik

mcp-client/agent.go

@@ -13,20 +13,22 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-var (
-	ErrMissingAPIKey = errors.New("ANTHROPIC_API_KEY environment variable not set")
-	ErrServerNotFound = errors.New("MCP server binary not found")
-)
+var ErrMissingAPIKey = errors.New("ANTHROPIC_API_KEY environment variable not set")
 
 type OutputHandler func(message string)
+type ToolApprovalHandler func(toolName string, args any) (approved bool, err error)
+type ToolResultHandler func(toolID string, result string)
 
 type Agent struct {
-	config          *Config
-	session         *mcp.ClientSession
-	cmd             *exec.Cmd
-	anthropicClient anthropic.Client
-	claudeTools     []anthropic.ToolUnionParam
-	outputHandler   OutputHandler
+	config              *Config
+	session             *mcp.ClientSession
+	cmd                 *exec.Cmd
+	anthropicClient     anthropic.Client
+	claudeTools         []anthropic.ToolUnionParam
+	outputHandler       OutputHandler
+	toolApprovalHandler ToolApprovalHandler
+	toolResultHandler   ToolResultHandler
+	conversationHistory []anthropic.MessageParam
 }
 
 func NewAgent(cfg *Config) *Agent {
@@ -40,6 +42,14 @@ func (a *Agent) SetOutputHandler(handler OutputHandler) {
 	a.outputHandler = handler
 }
 
+func (a *Agent) SetToolApprovalHandler(handler ToolApprovalHandler) {
+	a.toolApprovalHandler = handler
+}
+
+func (a *Agent) SetToolResultHandler(handler ToolResultHandler) {
+	a.toolResultHandler = handler
+}
+
 func (a *Agent) output(message string) {
 	if a.outputHandler != nil {
 		a.outputHandler(message)
@@ -164,16 +174,16 @@ func (a *Agent) convertMCPToolToClaudeTool(tool *mcp.Tool) anthropic.ToolUnionPa
 }
 
 func (a *Agent) ProcessQuery(ctx context.Context, query string) error {
-	messages := []anthropic.MessageParam{
+	a.conversationHistory = append(a.conversationHistory,
 		anthropic.NewUserMessage(anthropic.NewTextBlock(query)),
-	}
+	)
 
-	for _ = range maxAgentTurns {
+	for range maxAgentTurns {
 		message, err := a.anthropicClient.Messages.New(ctx, anthropic.MessageNewParams{
 			Model:     claudeModel,
 			MaxTokens: maxTokens,
 			System:    []anthropic.TextBlockParam{{Type: "text", Text: systemPrompt}},
-			Messages:  messages,
+			Messages:  a.conversationHistory,
 			Tools:     a.claudeTools,
 		})
 		if err != nil {
@@ -182,20 +192,27 @@ func (a *Agent) ProcessQuery(ctx context.Context, query string) error {
 
 		assistantContent, toolResults, hasToolUse := a.processClaudeResponse(ctx, message)
 
+		a.conversationHistory = append(a.conversationHistory,
+			anthropic.NewAssistantMessage(assistantContent...))
+
 		if !hasToolUse {
 			return nil
 		}
 
-		messages = append(messages, anthropic.NewAssistantMessage(assistantContent...))
-		messages = append(messages, anthropic.NewUserMessage(toolResults...))
+		a.conversationHistory = append(a.conversationHistory,
+			anthropic.NewUserMessage(toolResults...))
 	}
 
 	a.output("\n=== Maximum turns reached, requesting summary ===")
-	a.generateFinalSummary(ctx, messages)
+	a.generateFinalSummary(ctx, a.conversationHistory)
 
 	return nil
 }
 
+func (a *Agent) ClearHistory() {
+	a.conversationHistory = nil
+}
+
 func (a *Agent) processClaudeResponse(
 	ctx context.Context,
 	message *anthropic.Message,
@@ -232,6 +249,26 @@ func (a *Agent) executeToolCall(
 	ctx context.Context,
 	toolUse anthropic.ToolUseBlock,
 ) anthropic.ContentBlockParamUnion {
+	if a.toolApprovalHandler != nil {
+		approved, err := a.toolApprovalHandler(toolUse.Name, toolUse.Input)
+		if err != nil {
+			log.Printf("[Error] Tool approval failed: %v", err)
+			return anthropic.NewToolResultBlock(
+				toolUse.ID,
+				fmt.Sprintf("Approval error: %v", err),
+				true,
+			)
+		}
+		if !approved {
+			log.Printf("[Info] Tool %s was denied by user", toolUse.Name)
+			return anthropic.NewToolResultBlock(
+				toolUse.ID,
+				"Tool execution was denied by user",
+				true,
+			)
+		}
+	}
+
 	result, err := a.session.CallTool(ctx, &mcp.CallToolParams{
 		Name:      toolUse.Name,
 		Arguments: toolUse.Input,
@@ -258,11 +295,13 @@ func (a *Agent) executeToolCall(
 
 	resultText := extractTextContent(result.Content)
 	if resultText == "" {
-		log.Printf("[Warning] Tool returned empty content - this might indicate an unexpected response from MCP server")
+		log.Printf("[Warning] %s tool returned empty content - this might indicate an unexpected response from MCP server", toolUse.Name)
+	}
+	log.Printf("[%s] Tool Result: %s", toolUse.Name, resultText)
+
+	if a.toolResultHandler != nil {
+		a.toolResultHandler(toolUse.ID, resultText)
 	}
-	log.Printf("================================================")
-	log.Printf("[Tool Result]\n%s", resultText)
-	log.Printf("================================================")
 
 	return anthropic.NewToolResultBlock(toolUse.ID, resultText, false)
 }

mcp-client/config.go

@@ -4,27 +4,23 @@ import "github.com/anthropics/anthropic-sdk-go"
 
 const (
 	defaultServerPath = "../backend/server"
+	defaultPort       = "8081"
 	mcpClientName     = "mcp-client"
 	mcpClientVersion  = "v1.0.0"
 
 	claudeModel   = anthropic.ModelClaude3_5HaikuLatest
 	maxTokens     = 2048
 	maxAgentTurns = 5
 
-	systemPrompt = `You are an AI assistant with access to Keylime system management tools. Your goal is to help users manage and monitor their Keylime infrastructure.
+	systemPrompt = `You are a helpful assistant for managing Keylime infrastructure. You have access to tools for querying and managing Keylime agents, verifiers, and registrars.
 
-You have a maximum of 5 conversation turns to complete the task. When given a task:
-1. Break it down into steps if needed
-2. Use available tools to gather information and take actions
-3. Chain multiple tool calls together to accomplish complex tasks
-4. Provide clear explanations of what you're doing and what you found
-5. If you encounter failures, investigate and suggest solutions
-6. Work efficiently to complete tasks within the turn limit`
+Answer the user's question directly. If they ask for a list, give them the list. If they ask for details, give them details. Match your response to what they're asking for.`
 )
 
 type Config struct {
 	AnthropicAPIKey string
 	MCPServerPath   string
+	Port            string
 }
 
 func NewConfig() (*Config, error) {
@@ -34,9 +30,11 @@ func NewConfig() (*Config, error) {
 	}
 
 	serverPath := getEnvOrDefault("MCP_SERVER_PATH", defaultServerPath)
+	port := getEnvOrDefault("MCP_CLIENT_PORT", defaultPort)
 
 	return &Config{
 		AnthropicAPIKey: apiKey,
 		MCPServerPath:   serverPath,
+		Port:            ":" + port,
 	}, nil
 }

mcp-client/go.mod

@@ -11,6 +11,7 @@ require (
 require (
 	github.com/go-chi/chi/v5 v5.2.3 // indirect
 	github.com/google/jsonschema-go v0.3.0 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/tidwall/gjson v1.18.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect

mcp-client/go.sum

@@ -8,6 +8,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
 github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
 github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
 github.com/modelcontextprotocol/go-sdk v1.1.0 h1:Qjayg53dnKC4UZ+792W21e4BpwEZBzwgRW6LrjLWSwA=

mcp-client/main.go

@@ -37,29 +37,29 @@ func main() {
 	defer agent.Close()
 
 	if *webMode {
-		runWebMode(ctx, agent)
+		runWebMode(ctx, agent, config)
 	} else {
 		runCLIMode(ctx, agent)
 	}
 }
 
 func setupSignalHandler(cancel context.CancelFunc) {
 	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
 	go func() {
 		<-sigChan
 		log.Println("\nShutting down gracefully...")
 		cancel()
 	}()
 }
 
-func runWebMode(ctx context.Context, agent *Agent) {
+func runWebMode(ctx context.Context, agent *Agent, config *Config) {
 	webServer, err := NewWebServer(agent)
 	if err != nil {
 		log.Fatalf("Failed to create web server: %v", err)
 	}
 
-	if err := webServer.Start(ctx, ":8080"); err != nil {
+	if err := webServer.Start(ctx, config.Port); err != nil {
 		log.Fatalf("Web server failed: %v", err)
 	}
 }