Skip to content

Commit 53baee4

Browse files
dgageotdgageot
authored
Merge pull request #714 from dgageot/misc
Misc improvements
2 parents efd5b90 + 485ab02 commit 53baee4

16 files changed

Lines changed: 200 additions & 215 deletions

‎cmd/root/flags.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,6 @@ import (
99
func addRuntimeConfigFlags(cmd *cobra.Command, runConfig *config.RuntimeConfig) {
1010
addGatewayFlags(cmd, runConfig)
1111
cmd.PersistentFlags().StringSliceVar(&runConfig.EnvFiles, "env-from-file", nil, "Set environment variables from file")
12-
cmd.PersistentFlags().StringVar(&runConfig.RedirectURI, "redirect-uri", "", "Set the redirect URI for OAuth2 flows")
12+
cmd.PersistentFlags().StringVar(&runConfig.RedirectURI, "redirect-uri", "http://localhost:8083/oauth-callback", "Set the redirect URI for OAuth2 flows")
1313
cmd.PersistentFlags().BoolVar(&runConfig.GlobalCodeMode, "code-mode-tools", false, "Provide a single tool to call other tools via Javascript")
1414
}

‎cmd/root/mcp.go‎

Lines changed: 2 additions & 132 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,11 @@
11
package root
22

33
import (
4-
"context"
5-
"fmt"
6-
"log/slog"
7-
8-
"github.com/modelcontextprotocol/go-sdk/mcp"
94
"github.com/spf13/cobra"
105

11-
"github.com/docker/cagent/pkg/agentfile"
126
"github.com/docker/cagent/pkg/config"
13-
"github.com/docker/cagent/pkg/runtime"
14-
"github.com/docker/cagent/pkg/session"
15-
"github.com/docker/cagent/pkg/team"
16-
"github.com/docker/cagent/pkg/teamloader"
7+
"github.com/docker/cagent/pkg/mcp"
178
"github.com/docker/cagent/pkg/telemetry"
18-
"github.com/docker/cagent/pkg/version"
199
)
2010

2111
type mcpFlags struct {
@@ -43,127 +33,7 @@ func newMCPCmd() *cobra.Command {
4333

4434
func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) error {
4535
telemetry.TrackCommand("mcp", args)
46-
return f.runMCP(cmd, args)
47-
}
48-
49-
func (f *mcpFlags) runMCP(cmd *cobra.Command, args []string) error {
5036
ctx := cmd.Context()
5137

52-
slog.Debug("Starting MCP server", "agent_ref", args[0])
53-
54-
agentFilename, err := agentfile.Resolve(ctx, args[0])
55-
if err != nil {
56-
return err
57-
}
58-
59-
if f.runConfig.RedirectURI == "" {
60-
f.runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
61-
}
62-
63-
t, err := teamloader.Load(ctx, agentFilename, f.runConfig)
64-
if err != nil {
65-
return fmt.Errorf("failed to load agents: %w", err)
66-
}
67-
68-
defer func() {
69-
if err := t.StopToolSets(ctx); err != nil {
70-
slog.Error("Failed to stop tool sets", "error", err)
71-
}
72-
}()
73-
74-
server := mcp.NewServer(&mcp.Implementation{
75-
Name: "cagent",
76-
Version: version.Version,
77-
}, nil)
78-
79-
agentNames := t.AgentNames()
80-
slog.Debug("Adding MCP tools for agents", "count", len(agentNames))
81-
82-
for _, agentName := range agentNames {
83-
agent, err := t.Agent(agentName)
84-
if err != nil {
85-
return fmt.Errorf("failed to get agent %s: %w", agentName, err)
86-
}
87-
88-
description := agent.Description()
89-
if description == "" {
90-
description = fmt.Sprintf("Run the %s agent", agentName)
91-
}
92-
93-
slog.Debug("Adding MCP tool", "agent", agentName, "description", description)
94-
95-
toolDef := &mcp.Tool{
96-
Name: agentName,
97-
Description: description,
98-
InputSchema: map[string]any{
99-
"type": "object",
100-
"properties": map[string]any{
101-
"message": map[string]any{
102-
"type": "string",
103-
"description": "The message to send to the agent",
104-
},
105-
},
106-
"required": []string{"message"},
107-
},
108-
}
109-
110-
mcp.AddTool(server, toolDef, CreateToolHandler(t, agentName, agentFilename))
111-
}
112-
113-
slog.Debug("MCP server starting with stdio transport")
114-
115-
if err := server.Run(ctx, &mcp.StdioTransport{}); err != nil {
116-
return fmt.Errorf("MCP server error: %w", err)
117-
}
118-
119-
return nil
120-
}
121-
122-
type ToolInput struct {
123-
Message string `json:"message" jsonschema:"the message to send to the agent"`
124-
}
125-
126-
type ToolOutput struct {
127-
Response string `json:"response" jsonschema:"the response from the agent"`
128-
}
129-
130-
func CreateToolHandler(t *team.Team, agentName, agentFilename string) func(context.Context, *mcp.CallToolRequest, ToolInput) (*mcp.CallToolResult, ToolOutput, error) {
131-
return func(ctx context.Context, req *mcp.CallToolRequest, input ToolInput) (*mcp.CallToolResult, ToolOutput, error) {
132-
slog.Debug("MCP tool called", "agent", agentName, "message", input.Message)
133-
134-
agent, err := t.Agent(agentName)
135-
if err != nil {
136-
return nil, ToolOutput{}, fmt.Errorf("failed to get agent: %w", err)
137-
}
138-
139-
sess := session.New(
140-
session.WithTitle("MCP tool call"),
141-
session.WithMaxIterations(agent.MaxIterations()),
142-
session.WithUserMessage(agentFilename, input.Message),
143-
)
144-
sess.ToolsApproved = true
145-
146-
rt, err := runtime.New(t,
147-
runtime.WithCurrentAgent(agentName),
148-
runtime.WithRootSessionID(sess.ID),
149-
)
150-
if err != nil {
151-
return nil, ToolOutput{}, fmt.Errorf("failed to create runtime: %w", err)
152-
}
153-
154-
_, err = rt.Run(ctx, sess)
155-
if err != nil {
156-
slog.Error("Agent execution failed", "agent", agentName, "error", err)
157-
return nil, ToolOutput{}, fmt.Errorf("agent execution failed: %w", err)
158-
}
159-
160-
result := sess.GetLastAssistantMessageContent()
161-
if result == "" {
162-
result = "No response from agent"
163-
}
164-
165-
slog.Debug("Agent execution completed", "agent", agentName, "response_length", len(result))
166-
167-
return nil, ToolOutput{Response: result}, nil
168-
}
38+
return mcp.StartMCPServer(ctx, args[0], f.runConfig)
16939
}

‎cmd/root/run.go‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -153,10 +153,6 @@ func (f *runExecFlags) resolveAgentFile(ctx context.Context, agentFilename strin
153153
}
154154

155155
func (f *runExecFlags) loadAgents(ctx context.Context, agentFilename string) (*team.Team, error) {
156-
if f.runConfig.RedirectURI == "" {
157-
f.runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
158-
}
159-
160156
t, err := teamloader.Load(ctx, agentFilename, f.runConfig, teamloader.WithModelOverrides(f.modelOverrides))
161157
if err != nil {
162158
return nil, err

e2e/cagent_exec_openai_test.go renamed to e2e/cagent_exec_test.go

File renamed without changes.

‎e2e/mcp_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@ import (
66
"github.com/stretchr/testify/assert"
77
"github.com/stretchr/testify/require"
88

9-
"github.com/docker/cagent/cmd/root"
9+
"github.com/docker/cagent/pkg/mcp"
1010
"github.com/docker/cagent/pkg/teamloader"
1111
)
1212

13-
func TestMCPSingleAgent(t *testing.T) {
13+
func TestMCP_SingleAgent(t *testing.T) {
1414
t.Parallel()
1515

1616
ctx := t.Context()
@@ -22,16 +22,16 @@ func TestMCPSingleAgent(t *testing.T) {
2222
require.NoError(t, team.StopToolSets(ctx))
2323
})
2424

25-
handler := root.CreateToolHandler(team, "root", "testdata/basic.yaml")
26-
_, output, err := handler(ctx, nil, root.ToolInput{
25+
handler := mcp.CreateToolHandler(team, "root", "testdata/basic.yaml")
26+
_, output, err := handler(ctx, nil, mcp.ToolInput{
2727
Message: "What is 2+2? Answer in one sentence.",
2828
})
2929

3030
require.NoError(t, err)
3131
assert.Equal(t, "2+2 equals 4.", output.Response)
3232
}
3333

34-
func TestMCPMultiAgent(t *testing.T) {
34+
func TestMCP_MultiAgent(t *testing.T) {
3535
t.Parallel()
3636

3737
ctx := t.Context()
@@ -43,11 +43,11 @@ func TestMCPMultiAgent(t *testing.T) {
4343
require.NoError(t, team.StopToolSets(ctx))
4444
})
4545

46-
handler := root.CreateToolHandler(team, "web", "testdata/multi.yaml")
47-
_, output, err := handler(ctx, nil, root.ToolInput{
46+
handler := mcp.CreateToolHandler(team, "web", "testdata/multi.yaml")
47+
_, output, err := handler(ctx, nil, mcp.ToolInput{
4848
Message: "Say hello in one sentence.",
4949
})
5050

5151
require.NoError(t, err)
52-
assert.Equal(t, "Hello!", output.Response)
52+
assert.Equal(t, "Hello, nice to meet you.", output.Response)
5353
}

‎e2e/proxy_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package e2e_test
22

33
import (
44
"bytes"
5+
"context"
56
"io"
67
"log/slog"
78
"maps"
@@ -190,3 +191,9 @@ func isStreamResponse(resp *http.Response) bool {
190191
strings.Contains(ct, "application/x-ndjson") ||
191192
strings.Contains(ct, "application/stream+json")
192193
}
194+
195+
type testEnvProvider map[string]string
196+
197+
func (p *testEnvProvider) Get(_ context.Context, name string) string {
198+
return (*p)[name]
199+
}

‎e2e/runtime_openai_test.go‎

Lines changed: 0 additions & 33 deletions
This file was deleted.

e2e/runtime_mistral_test.go renamed to e2e/runtime_test.go

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

33
import (
4-
"context"
54
"testing"
65

76
"github.com/stretchr/testify/assert"
@@ -12,13 +11,13 @@ import (
1211
"github.com/docker/cagent/pkg/teamloader"
1312
)
1413

15-
func TestRuntime_BasicMistral(t *testing.T) {
14+
func TestRuntime_OpenAI_Basic(t *testing.T) {
1615
t.Parallel()
1716

1817
ctx := t.Context()
1918
_, runtimeConfig := startRecordingAIProxy(t)
2019

21-
team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig, teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))
20+
team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig)
2221
require.NoError(t, err)
2322

2423
rt, err := runtime.New(team)
@@ -29,12 +28,27 @@ func TestRuntime_BasicMistral(t *testing.T) {
2928
require.NoError(t, err)
3029

3130
response := sess.GetLastAssistantMessageContent()
32-
assert.Equal(t, `It seems like "djordje" is a name, most likely of Slavic origin. It is commonly spelled as "Đorđe" in Serbian language, and it means "farmer" or "earthworker". It is a masculine given name, and it is quite popular in Serbia, Montenegro, and other countries in the region. Without more context, it's hard to say exactly who "djordje" is, as it could refer to any person by that name.`, response)
33-
assert.Equal(t, `"Inquiry About the Identity of 'Djordje'"`, sess.Title)
31+
assert.Equal(t, "Djordje is a popular given name in some Eastern European countries, such as Serbia. If you have more specific information or context, I'd be happy to help further.", response)
32+
assert.Equal(t, "Understanding identity: Who is Djordje?", sess.Title)
3433
}
3534

36-
type testEnvProvider map[string]string
35+
func TestRuntime_Mistral_Basic(t *testing.T) {
36+
t.Parallel()
37+
38+
ctx := t.Context()
39+
_, runtimeConfig := startRecordingAIProxy(t)
40+
41+
team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig, teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))
42+
require.NoError(t, err)
43+
44+
rt, err := runtime.New(team)
45+
require.NoError(t, err)
46+
47+
sess := session.New(session.WithUserMessage("", "Who's djordje?"))
48+
_, err = rt.Run(ctx, sess)
49+
require.NoError(t, err)
3750

38-
func (p *testEnvProvider) Get(_ context.Context, name string) string {
39-
return (*p)[name]
51+
response := sess.GetLastAssistantMessageContent()
52+
assert.Equal(t, `It seems like "djordje" is a name, most likely of Slavic origin. It is commonly spelled as "Đorđe" in Serbian language, and it means "farmer" or "earthworker". It is a masculine given name, and it is quite popular in Serbia, Montenegro, and other countries in the region. Without more context, it's hard to say exactly who "djordje" is, as it could refer to any person by that name.`, response)
53+
assert.Equal(t, `"Inquiry About the Identity of 'Djordje'"`, sess.Title)
4054
}

‎e2e/testdata/cassettes/TestMCPMultiAgent.yaml‎

Lines changed: 0 additions & 23 deletions
This file was deleted.

‎e2e/testdata/cassettes/TestMCP_MultiAgent.yaml‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
---
2+
version: 2
3+
interactions:
4+
- id: 0
5+
request:
6+
proto: HTTP/1.1
7+
proto_major: 1
8+
proto_minor: 1
9+
content_length: 0
10+
host: api.openai.com
11+
body: "{\"model\":\"gpt-5-mini\",\"messages\":[{\"role\":\"system\",\"content\":\"You are a knowledgeable assistant that helps users with web tasks.\\n\"},{\"role\":\"user\",\"content\":\"Say hello in one sentence.\"}],\"stream\":true,\"stream_options\":{\"include_usage\":true}}"
12+
url: https://api.openai.com/v1/chat/completions
13+
method: POST
14+
response:
15+
proto: HTTP/2.0
16+
proto_major: 2
17+
proto_minor: 0
18+
content_length: -1
19+
body: "data: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\",\"refusal\":null},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"9o3Tj\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"ZQ\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\",\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"S9ksF0\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\" nice\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"ny\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\" to\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"mgzx\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\" meet\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"7n\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\" you\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"YAh\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{\"content\":\".\"},\"finish_reason\":null}],\"usage\":null,\"obfuscation\":\"0Kc5Id\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":null,\"obfuscation\":\"r\"}\n\ndata: {\"id\":\"chatcmpl-CY9HhLy35p8nQ0Ul63kaF9VFmdtEb\",\"object\":\"chat.completion.chunk\",\"created\":1762254877,\"model\":\"gpt-5-mini-2025-08-07\",\"service_tier\":\"default\",\"system_fingerprint\":null,\"choices\":[],\"usage\":{\"prompt_tokens\":28,\"completion_tokens\":80,\"total_tokens\":108,\"prompt_tokens_details\":{\"cached_tokens\":0,\"audio_tokens\":0},\"completion_tokens_details\":{\"reasoning_tokens\":64,\"audio_tokens\":0,\"accepted_prediction_tokens\":0,\"rejected_prediction_tokens\":0}},\"obfuscation\":\"fd8D\"}\n\ndata: [DONE]\n\n"
20+
headers: {}
21+
status: 200 OK
22+
code: 200
23+
duration: 3.244890916s

0 commit comments

Comments
 (0)