Skip to content

Commit 30ff13b

Browse files
author
Marek Safarik
committed
Streaming communication using socket
Signed-off-by: Marek Safarik <[email protected]>
1 parent 90f33db commit 30ff13b

File tree

9 files changed

+690
-146
lines changed

9 files changed

+690
-146
lines changed

mcp-client/agent.go

Lines changed: 60 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -13,20 +13,22 @@ import (
1313
"github.com/modelcontextprotocol/go-sdk/mcp"
1414
)
1515

16-
var (
17-
ErrMissingAPIKey = errors.New("ANTHROPIC_API_KEY environment variable not set")
18-
ErrServerNotFound = errors.New("MCP server binary not found")
19-
)
16+
var ErrMissingAPIKey = errors.New("ANTHROPIC_API_KEY environment variable not set")
2017

2118
type OutputHandler func(message string)
19+
type ToolApprovalHandler func(toolName string, args any) (approved bool, err error)
20+
type ToolResultHandler func(toolID string, result string)
2221

2322
type Agent struct {
24-
config *Config
25-
session *mcp.ClientSession
26-
cmd *exec.Cmd
27-
anthropicClient anthropic.Client
28-
claudeTools []anthropic.ToolUnionParam
29-
outputHandler OutputHandler
23+
config *Config
24+
session *mcp.ClientSession
25+
cmd *exec.Cmd
26+
anthropicClient anthropic.Client
27+
claudeTools []anthropic.ToolUnionParam
28+
outputHandler OutputHandler
29+
toolApprovalHandler ToolApprovalHandler
30+
toolResultHandler ToolResultHandler
31+
conversationHistory []anthropic.MessageParam
3032
}
3133

3234
func NewAgent(cfg *Config) *Agent {
@@ -40,6 +42,14 @@ func (a *Agent) SetOutputHandler(handler OutputHandler) {
4042
a.outputHandler = handler
4143
}
4244

45+
func (a *Agent) SetToolApprovalHandler(handler ToolApprovalHandler) {
46+
a.toolApprovalHandler = handler
47+
}
48+
49+
func (a *Agent) SetToolResultHandler(handler ToolResultHandler) {
50+
a.toolResultHandler = handler
51+
}
52+
4353
func (a *Agent) output(message string) {
4454
if a.outputHandler != nil {
4555
a.outputHandler(message)
@@ -164,16 +174,16 @@ func (a *Agent) convertMCPToolToClaudeTool(tool *mcp.Tool) anthropic.ToolUnionPa
164174
}
165175

166176
func (a *Agent) ProcessQuery(ctx context.Context, query string) error {
167-
messages := []anthropic.MessageParam{
177+
a.conversationHistory = append(a.conversationHistory,
168178
anthropic.NewUserMessage(anthropic.NewTextBlock(query)),
169-
}
179+
)
170180

171-
for _ = range maxAgentTurns {
181+
for range maxAgentTurns {
172182
message, err := a.anthropicClient.Messages.New(ctx, anthropic.MessageNewParams{
173183
Model: claudeModel,
174184
MaxTokens: maxTokens,
175185
System: []anthropic.TextBlockParam{{Type: "text", Text: systemPrompt}},
176-
Messages: messages,
186+
Messages: a.conversationHistory,
177187
Tools: a.claudeTools,
178188
})
179189
if err != nil {
@@ -182,20 +192,27 @@ func (a *Agent) ProcessQuery(ctx context.Context, query string) error {
182192

183193
assistantContent, toolResults, hasToolUse := a.processClaudeResponse(ctx, message)
184194

195+
a.conversationHistory = append(a.conversationHistory,
196+
anthropic.NewAssistantMessage(assistantContent...))
197+
185198
if !hasToolUse {
186199
return nil
187200
}
188201

189-
messages = append(messages, anthropic.NewAssistantMessage(assistantContent...))
190-
messages = append(messages, anthropic.NewUserMessage(toolResults...))
202+
a.conversationHistory = append(a.conversationHistory,
203+
anthropic.NewUserMessage(toolResults...))
191204
}
192205

193206
a.output("\n=== Maximum turns reached, requesting summary ===")
194-
a.generateFinalSummary(ctx, messages)
207+
a.generateFinalSummary(ctx, a.conversationHistory)
195208

196209
return nil
197210
}
198211

212+
func (a *Agent) ClearHistory() {
213+
a.conversationHistory = nil
214+
}
215+
199216
func (a *Agent) processClaudeResponse(
200217
ctx context.Context,
201218
message *anthropic.Message,
@@ -232,6 +249,26 @@ func (a *Agent) executeToolCall(
232249
ctx context.Context,
233250
toolUse anthropic.ToolUseBlock,
234251
) anthropic.ContentBlockParamUnion {
252+
if a.toolApprovalHandler != nil {
253+
approved, err := a.toolApprovalHandler(toolUse.Name, toolUse.Input)
254+
if err != nil {
255+
log.Printf("[Error] Tool approval failed: %v", err)
256+
return anthropic.NewToolResultBlock(
257+
toolUse.ID,
258+
fmt.Sprintf("Approval error: %v", err),
259+
true,
260+
)
261+
}
262+
if !approved {
263+
log.Printf("[Info] Tool %s was denied by user", toolUse.Name)
264+
return anthropic.NewToolResultBlock(
265+
toolUse.ID,
266+
"Tool execution was denied by user",
267+
true,
268+
)
269+
}
270+
}
271+
235272
result, err := a.session.CallTool(ctx, &mcp.CallToolParams{
236273
Name: toolUse.Name,
237274
Arguments: toolUse.Input,
@@ -258,11 +295,13 @@ func (a *Agent) executeToolCall(
258295

259296
resultText := extractTextContent(result.Content)
260297
if resultText == "" {
261-
log.Printf("[Warning] Tool returned empty content - this might indicate an unexpected response from MCP server")
298+
log.Printf("[Warning] %s tool returned empty content - this might indicate an unexpected response from MCP server", toolUse.Name)
299+
}
300+
log.Printf("[%s] Tool Result: %s", toolUse.Name, resultText)
301+
302+
if a.toolResultHandler != nil {
303+
a.toolResultHandler(toolUse.ID, resultText)
262304
}
263-
log.Printf("================================================")
264-
log.Printf("[Tool Result]\n%s", resultText)
265-
log.Printf("================================================")
266305

267306
return anthropic.NewToolResultBlock(toolUse.ID, resultText, false)
268307
}

mcp-client/config.go

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,27 +4,23 @@ import "github.com/anthropics/anthropic-sdk-go"
44

55
const (
66
defaultServerPath = "../backend/server"
7+
defaultPort = "8081"
78
mcpClientName = "mcp-client"
89
mcpClientVersion = "v1.0.0"
910

1011
claudeModel = anthropic.ModelClaude3_5HaikuLatest
1112
maxTokens = 2048
1213
maxAgentTurns = 5
1314

14-
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.
15+
systemPrompt = `You are a helpful assistant for managing Keylime infrastructure. You have access to tools for querying and managing Keylime agents, verifiers, and registrars.
1516
16-
You have a maximum of 5 conversation turns to complete the task. When given a task:
17-
1. Break it down into steps if needed
18-
2. Use available tools to gather information and take actions
19-
3. Chain multiple tool calls together to accomplish complex tasks
20-
4. Provide clear explanations of what you're doing and what you found
21-
5. If you encounter failures, investigate and suggest solutions
22-
6. Work efficiently to complete tasks within the turn limit`
17+
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.`
2318
)
2419

2520
type Config struct {
2621
AnthropicAPIKey string
2722
MCPServerPath string
23+
Port string
2824
}
2925

3026
func NewConfig() (*Config, error) {
@@ -34,9 +30,11 @@ func NewConfig() (*Config, error) {
3430
}
3531

3632
serverPath := getEnvOrDefault("MCP_SERVER_PATH", defaultServerPath)
33+
port := getEnvOrDefault("MCP_CLIENT_PORT", defaultPort)
3734

3835
return &Config{
3936
AnthropicAPIKey: apiKey,
4037
MCPServerPath: serverPath,
38+
Port: ":" + port,
4139
}, nil
4240
}

mcp-client/go.mod

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ require (
1111
require (
1212
github.com/go-chi/chi/v5 v5.2.3 // indirect
1313
github.com/google/jsonschema-go v0.3.0 // indirect
14+
github.com/gorilla/websocket v1.5.3 // indirect
1415
github.com/tidwall/gjson v1.18.0 // indirect
1516
github.com/tidwall/match v1.1.1 // indirect
1617
github.com/tidwall/pretty v1.2.1 // indirect

mcp-client/go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
88
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
99
github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
1010
github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
11+
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
12+
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
1113
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
1214
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
1315
github.com/modelcontextprotocol/go-sdk v1.1.0 h1:Qjayg53dnKC4UZ+792W21e4BpwEZBzwgRW6LrjLWSwA=

mcp-client/main.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,29 +37,29 @@ func main() {
3737
defer agent.Close()
3838

3939
if *webMode {
40-
runWebMode(ctx, agent)
40+
runWebMode(ctx, agent, config)
4141
} else {
4242
runCLIMode(ctx, agent)
4343
}
4444
}
4545

4646
func setupSignalHandler(cancel context.CancelFunc) {
4747
sigChan := make(chan os.Signal, 1)
48-
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
48+
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
4949
go func() {
5050
<-sigChan
5151
log.Println("\nShutting down gracefully...")
5252
cancel()
5353
}()
5454
}
5555

56-
func runWebMode(ctx context.Context, agent *Agent) {
56+
func runWebMode(ctx context.Context, agent *Agent, config *Config) {
5757
webServer, err := NewWebServer(agent)
5858
if err != nil {
5959
log.Fatalf("Failed to create web server: %v", err)
6060
}
6161

62-
if err := webServer.Start(ctx, ":8080"); err != nil {
62+
if err := webServer.Start(ctx, config.Port); err != nil {
6363
log.Fatalf("Web server failed: %v", err)
6464
}
6565
}

mcp-client/mcp-client

160 KB
Binary file not shown.

0 commit comments

Comments
 (0)