Misc improvements

cmd/root/flags.go

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

cmd/root/mcp.go

@@ -1,21 +1,11 @@
 package root
 
 import (
-	"context"
-	"fmt"
-	"log/slog"
-
-	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/spf13/cobra"
 
-	"github.com/docker/cagent/pkg/agentfile"
 	"github.com/docker/cagent/pkg/config"
-	"github.com/docker/cagent/pkg/runtime"
-	"github.com/docker/cagent/pkg/session"
-	"github.com/docker/cagent/pkg/team"
-	"github.com/docker/cagent/pkg/teamloader"
+	"github.com/docker/cagent/pkg/mcp"
 	"github.com/docker/cagent/pkg/telemetry"
-	"github.com/docker/cagent/pkg/version"
 )
 
 type mcpFlags struct {
@@ -43,127 +33,7 @@ func newMCPCmd() *cobra.Command {
 
 func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) error {
 	telemetry.TrackCommand("mcp", args)
-	return f.runMCP(cmd, args)
-}
-
-func (f *mcpFlags) runMCP(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
-	slog.Debug("Starting MCP server", "agent_ref", args[0])
-
-	agentFilename, err := agentfile.Resolve(ctx, args[0])
-	if err != nil {
-		return err
-	}
-
-	if f.runConfig.RedirectURI == "" {
-		f.runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
-	}
-
-	t, err := teamloader.Load(ctx, agentFilename, f.runConfig)
-	if err != nil {
-		return fmt.Errorf("failed to load agents: %w", err)
-	}
-
-	defer func() {
-		if err := t.StopToolSets(ctx); err != nil {
-			slog.Error("Failed to stop tool sets", "error", err)
-		}
-	}()
-
-	server := mcp.NewServer(&mcp.Implementation{
-		Name:    "cagent",
-		Version: version.Version,
-	}, nil)
-
-	agentNames := t.AgentNames()
-	slog.Debug("Adding MCP tools for agents", "count", len(agentNames))
-
-	for _, agentName := range agentNames {
-		agent, err := t.Agent(agentName)
-		if err != nil {
-			return fmt.Errorf("failed to get agent %s: %w", agentName, err)
-		}
-
-		description := agent.Description()
-		if description == "" {
-			description = fmt.Sprintf("Run the %s agent", agentName)
-		}
-
-		slog.Debug("Adding MCP tool", "agent", agentName, "description", description)
-
-		toolDef := &mcp.Tool{
-			Name:        agentName,
-			Description: description,
-			InputSchema: map[string]any{
-				"type": "object",
-				"properties": map[string]any{
-					"message": map[string]any{
-						"type":        "string",
-						"description": "The message to send to the agent",
-					},
-				},
-				"required": []string{"message"},
-			},
-		}
-
-		mcp.AddTool(server, toolDef, CreateToolHandler(t, agentName, agentFilename))
-	}
-
-	slog.Debug("MCP server starting with stdio transport")
-
-	if err := server.Run(ctx, &mcp.StdioTransport{}); err != nil {
-		return fmt.Errorf("MCP server error: %w", err)
-	}
-
-	return nil
-}
-
-type ToolInput struct {
-	Message string `json:"message" jsonschema:"the message to send to the agent"`
-}
-
-type ToolOutput struct {
-	Response string `json:"response" jsonschema:"the response from the agent"`
-}
-
-func CreateToolHandler(t *team.Team, agentName, agentFilename string) func(context.Context, *mcp.CallToolRequest, ToolInput) (*mcp.CallToolResult, ToolOutput, error) {
-	return func(ctx context.Context, req *mcp.CallToolRequest, input ToolInput) (*mcp.CallToolResult, ToolOutput, error) {
-		slog.Debug("MCP tool called", "agent", agentName, "message", input.Message)
-
-		agent, err := t.Agent(agentName)
-		if err != nil {
-			return nil, ToolOutput{}, fmt.Errorf("failed to get agent: %w", err)
-		}
-
-		sess := session.New(
-			session.WithTitle("MCP tool call"),
-			session.WithMaxIterations(agent.MaxIterations()),
-			session.WithUserMessage(agentFilename, input.Message),
-		)
-		sess.ToolsApproved = true
-
-		rt, err := runtime.New(t,
-			runtime.WithCurrentAgent(agentName),
-			runtime.WithRootSessionID(sess.ID),
-		)
-		if err != nil {
-			return nil, ToolOutput{}, fmt.Errorf("failed to create runtime: %w", err)
-		}
-
-		_, err = rt.Run(ctx, sess)
-		if err != nil {
-			slog.Error("Agent execution failed", "agent", agentName, "error", err)
-			return nil, ToolOutput{}, fmt.Errorf("agent execution failed: %w", err)
-		}
-
-		result := sess.GetLastAssistantMessageContent()
-		if result == "" {
-			result = "No response from agent"
-		}
-
-		slog.Debug("Agent execution completed", "agent", agentName, "response_length", len(result))
-
-		return nil, ToolOutput{Response: result}, nil
-	}
+	return mcp.StartMCPServer(ctx, args[0], f.runConfig)
 }

cmd/root/run.go

@@ -153,10 +153,6 @@ func (f *runExecFlags) resolveAgentFile(ctx context.Context, agentFilename strin
 }
 
 func (f *runExecFlags) loadAgents(ctx context.Context, agentFilename string) (*team.Team, error) {
-	if f.runConfig.RedirectURI == "" {
-		f.runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
-	}
-
 	t, err := teamloader.Load(ctx, agentFilename, f.runConfig, teamloader.WithModelOverrides(f.modelOverrides))
 	if err != nil {
 		return nil, err

e2e/mcp_test.go

@@ -6,11 +6,11 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/docker/cagent/cmd/root"
+	"github.com/docker/cagent/pkg/mcp"
 	"github.com/docker/cagent/pkg/teamloader"
 )
 
-func TestMCPSingleAgent(t *testing.T) {
+func TestMCP_SingleAgent(t *testing.T) {
 	t.Parallel()
 
 	ctx := t.Context()
@@ -22,16 +22,16 @@ func TestMCPSingleAgent(t *testing.T) {
 		require.NoError(t, team.StopToolSets(ctx))
 	})
 
-	handler := root.CreateToolHandler(team, "root", "testdata/basic.yaml")
-	_, output, err := handler(ctx, nil, root.ToolInput{
+	handler := mcp.CreateToolHandler(team, "root", "testdata/basic.yaml")
+	_, output, err := handler(ctx, nil, mcp.ToolInput{
 		Message: "What is 2+2? Answer in one sentence.",
 	})
 
 	require.NoError(t, err)
 	assert.Equal(t, "2+2 equals 4.", output.Response)
 }
 
-func TestMCPMultiAgent(t *testing.T) {
+func TestMCP_MultiAgent(t *testing.T) {
 	t.Parallel()
 
 	ctx := t.Context()
@@ -43,11 +43,11 @@ func TestMCPMultiAgent(t *testing.T) {
 		require.NoError(t, team.StopToolSets(ctx))
 	})
 
-	handler := root.CreateToolHandler(team, "web", "testdata/multi.yaml")
-	_, output, err := handler(ctx, nil, root.ToolInput{
+	handler := mcp.CreateToolHandler(team, "web", "testdata/multi.yaml")
+	_, output, err := handler(ctx, nil, mcp.ToolInput{
 		Message: "Say hello in one sentence.",
 	})
 
 	require.NoError(t, err)
-	assert.Equal(t, "Hello!", output.Response)
+	assert.Equal(t, "Hello, nice to meet you.", output.Response)
 }

e2e/proxy_test.go

@@ -2,6 +2,7 @@ package e2e_test
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"log/slog"
 	"maps"
@@ -190,3 +191,9 @@ func isStreamResponse(resp *http.Response) bool {
 		strings.Contains(ct, "application/x-ndjson") ||
 		strings.Contains(ct, "application/stream+json")
 }
+
+type testEnvProvider map[string]string
+
+func (p *testEnvProvider) Get(_ context.Context, name string) string {
+	return (*p)[name]
+}

e2e/runtime_mistral_test.go → e2e/runtime_test.go

@@ -1,7 +1,6 @@
 package e2e_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -12,13 +11,13 @@ import (
 	"github.com/docker/cagent/pkg/teamloader"
 )
 
-func TestRuntime_BasicMistral(t *testing.T) {
+func TestRuntime_OpenAI_Basic(t *testing.T) {
 	t.Parallel()
 
 	ctx := t.Context()
 	_, runtimeConfig := startRecordingAIProxy(t)
 
-	team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig, teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))
+	team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig)
 	require.NoError(t, err)
 
 	rt, err := runtime.New(team)
@@ -29,12 +28,27 @@ func TestRuntime_BasicMistral(t *testing.T) {
 	require.NoError(t, err)
 
 	response := sess.GetLastAssistantMessageContent()
-	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)
-	assert.Equal(t, `"Inquiry About the Identity of 'Djordje'"`, sess.Title)
+	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)
+	assert.Equal(t, "Understanding identity: Who is Djordje?", sess.Title)
 }
 
-type testEnvProvider map[string]string
+func TestRuntime_Mistral_Basic(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+	_, runtimeConfig := startRecordingAIProxy(t)
+
+	team, err := teamloader.Load(ctx, "testdata/basic.yaml", runtimeConfig, teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))
+	require.NoError(t, err)
+
+	rt, err := runtime.New(team)
+	require.NoError(t, err)
+
+	sess := session.New(session.WithUserMessage("", "Who's djordje?"))
+	_, err = rt.Run(ctx, sess)
+	require.NoError(t, err)
 
-func (p *testEnvProvider) Get(_ context.Context, name string) string {
-	return (*p)[name]
+	response := sess.GetLastAssistantMessageContent()
+	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)
+	assert.Equal(t, `"Inquiry About the Identity of 'Djordje'"`, sess.Title)
 }

e2e/testdata/cassettes/TestMCP_MultiAgent.yaml

@@ -0,0 +1,23 @@
+---
+version: 2
+interactions:
+- id: 0
+  request:
+    proto: HTTP/1.1
+    proto_major: 1
+    proto_minor: 1
+    content_length: 0
+    host: api.openai.com
+    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}}"
+    url: https://api.openai.com/v1/chat/completions
+    method: POST
+  response:
+    proto: HTTP/2.0
+    proto_major: 2
+    proto_minor: 0
+    content_length: -1
+    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"
+    headers: {}
+    status: 200 OK
+    code: 200
+    duration: 3.244890916s

e2e/testdata/multi.yaml

@@ -1,4 +1,3 @@
-#!/usr/bin/env cagent run
 version: "2"
 
 agents:
@@ -13,5 +12,4 @@ agents:
   web:
     model: openai/gpt-5-mini
     instruction: |
-      You are a knowledgeable assistant that helps users with various tasks.
-      Be helpful, accurate, and concise in your responses.
+      You are a knowledgeable assistant that helps users with web tasks.