diff --git a/internal/app/azldev/cmds/component/changed.go b/internal/app/azldev/cmds/component/changed.go index 63a1b896..c74cd49e 100644 --- a/internal/app/azldev/cmds/component/changed.go +++ b/internal/app/azldev/cmds/component/changed.go @@ -94,7 +94,7 @@ detected via lock file presence in the compared refs when using -a.`, // it inspects historical locks at arbitrary refs. _ = cmd.Flags().MarkHidden("skip-lock-validation") - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/cmds/component/history.go b/internal/app/azldev/cmds/component/history.go index 2ea21ef5..e5389683 100644 --- a/internal/app/azldev/cmds/component/history.go +++ b/internal/app/azldev/cmds/component/history.go @@ -131,7 +131,7 @@ hand-picking entries to document.`, // History is read-only; the lock validation flag is meaningless here. _ = cmd.Flags().MarkHidden("skip-lock-validation") - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/cmds/component/list.go b/internal/app/azldev/cmds/component/list.go index ac348c81..04cc0e7f 100644 --- a/internal/app/azldev/cmds/component/list.go +++ b/internal/app/azldev/cmds/component/list.go @@ -53,7 +53,7 @@ Component name patterns support glob syntax (*, ?, []).`, ValidArgsFunction: components.GenerateComponentNameCompletions, } - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) components.AddComponentFilterOptionsToCommand(cmd, &options.ComponentFilter) diff --git a/internal/app/azldev/cmds/config/dump.go b/internal/app/azldev/cmds/config/dump.go index 4550f43d..0fa58c43 100644 --- a/internal/app/azldev/cmds/config/dump.go +++ b/internal/app/azldev/cmds/config/dump.go @@ -91,7 +91,7 @@ issues or inspecting effective values.`, cmd.Flags().VarP(&configDumpFormat, "format", "f", "Output format {json, toml}") - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/cmds/config/schema.go b/internal/app/azldev/cmds/config/schema.go index fc09a70a..857224e6 100644 --- a/internal/app/azldev/cmds/config/schema.go +++ b/internal/app/azldev/cmds/config/schema.go @@ -43,7 +43,7 @@ editors or linters that support JSON Schema validation of TOML files.`, }, } - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/cmds/image/list.go b/internal/app/azldev/cmds/image/list.go index d0d6bc78..0f7e5166 100644 --- a/internal/app/azldev/cmds/image/list.go +++ b/internal/app/azldev/cmds/image/list.go @@ -93,7 +93,7 @@ If no patterns are provided, all images are listed.`, ValidArgsFunction: generateImageNameCompletions, } - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/cmds/pkg/list.go b/internal/app/azldev/cmds/pkg/list.go index 07258423..3bf36a6d 100644 --- a/internal/app/azldev/cmds/pkg/list.go +++ b/internal/app/azldev/cmds/pkg/list.go @@ -124,7 +124,7 @@ and are emitted as empty arrays (never null) when there are no memberships.`, // Help shells complete '--rpm-file' with .json paths. _ = cmd.MarkFlagFilename("rpm-file", "json") - azldev.ExportAsMCPTool(cmd) + azldev.ExportAsReadOnlyMCPTool(cmd) return cmd } diff --git a/internal/app/azldev/command.go b/internal/app/azldev/command.go index 1880b332..977dee40 100644 --- a/internal/app/azldev/command.go +++ b/internal/app/azldev/command.go @@ -36,6 +36,15 @@ var ErrInvalidUsage = errors.New("invalid usage") // should be enabled in MCP server mode. The value associated with the key is ignored. const CmdAnnotationMCPEnabled = "azldev.mcp.enabled" +// CmdAnnotationMCPReadOnly is the [cobra.Command] annotation key used to indicate that a command, +// when exposed as an MCP tool, is read-only (does not mutate any state). MCP clients may use this +// hint to auto-approve the tool. The value associated with the key is ignored. +const CmdAnnotationMCPReadOnly = "azldev.mcp.readonly" + +// cmdMCPAnnotationValue is the placeholder value stored for MCP command annotations; only the +// presence of the key matters. +const cmdMCPAnnotationValue = "true" + type ( cobraRunFuncType = func(command *cobra.Command, args []string) error @@ -147,14 +156,17 @@ func GetEnvFromCommand(cmd *cobra.Command) (*Env, error) { } // ExportAsMCPTool updates the provided command (and all descendant commands), -// opting it into being advertised as a tool in MCP server mode. +// opting it into being advertised as a tool in MCP server mode. Use it for commands +// that can write, but only ever to a caller-provided output path (e.g. 'docs markdown', +// 'component diff-sources --output-file'); use [ExportAsReadOnlyMCPTool] for commands +// that never modify their environment so clients can auto-approve them. func ExportAsMCPTool(cmd *cobra.Command) { if cmd.Annotations == nil { cmd.Annotations = make(map[string]string) } // The value doesn't matter. - cmd.Annotations[CmdAnnotationMCPEnabled] = "true" + cmd.Annotations[CmdAnnotationMCPEnabled] = cmdMCPAnnotationValue // If the command has subcommands, then recursively opt them in as well. for _, subCmd := range cmd.Commands() { @@ -162,6 +174,24 @@ func ExportAsMCPTool(cmd *cobra.Command) { } } +// ExportAsReadOnlyMCPTool is like [ExportAsMCPTool], but additionally marks the command (and all +// descendant commands) as read-only -- it does not modify its environment -- so that MCP clients +// can advertise and auto-approve them as non-mutating tools. +func ExportAsReadOnlyMCPTool(cmd *cobra.Command) { + if cmd.Annotations == nil { + cmd.Annotations = make(map[string]string) + } + + // The values don't matter. + cmd.Annotations[CmdAnnotationMCPEnabled] = cmdMCPAnnotationValue + cmd.Annotations[CmdAnnotationMCPReadOnly] = cmdMCPAnnotationValue + + // If the command has subcommands, then recursively opt them in as well. + for _, subCmd := range cmd.Commands() { + ExportAsReadOnlyMCPTool(subCmd) + } +} + // Displays the results of a command in the appropriate format to stdout. func reportResults(env *Env, results interface{}) error { switch env.defaultReportFormat { diff --git a/internal/app/azldev/core/mcp/mcpserver.go b/internal/app/azldev/core/mcp/mcpserver.go index 51948319..d0a0cd02 100644 --- a/internal/app/azldev/core/mcp/mcpserver.go +++ b/internal/app/azldev/core/mcp/mcpserver.go @@ -19,11 +19,13 @@ import ( "github.com/samber/lo" "github.com/spf13/cobra" "github.com/spf13/pflag" + "go.szostok.io/version" ) -// Performs the file download requested by options. +// RunMCPServer starts the azldev MCP server, advertising the exported commands as +// tools over stdio until the context is canceled. func RunMCPServer(env *azldev.Env, cmd *cobra.Command) error { - srv := server.NewMCPServer(cmd.Short, "1.0.0", + srv := server.NewMCPServer(cmd.Root().Name(), version.Get().Version, server.WithLogging(), server.WithRecovery(), server.WithResourceCapabilities(true, true), @@ -47,6 +49,13 @@ func RunMCPServer(env *azldev.Env, cmd *cobra.Command) error { stdioServer := server.NewStdioServer(srv) + // handleToolCall runs the shared cobra command tree and swaps the process-global + // os.Stdout to capture output, so two handlers must not run at once. mcp-go + // dispatches tool calls to a worker pool (default 5) -- an agent firing parallel + // skill/tool calls would hit it -- so pin the pool to a single worker and let calls + // serialize instead of corrupting each other's output. + server.WithWorkerPoolSize(1)(stdioServer) + slog.Info("Starting MCP server") // Run the server until canceled. @@ -79,6 +88,12 @@ func addToolForCmd(env *azldev.Env, srv *server.MCPServer, leaf *cobra.Command) toolOptions = append(toolOptions, mcp.WithDescription(toolDesc)) + // Mirror our read-only annotation (set by [azldev.ExportAsReadOnlyMCPTool]) into the MCP tool + // schema so that clients can treat and auto-approve the tool as non-mutating. + if _, readOnly := leaf.Annotations[azldev.CmdAnnotationMCPReadOnly]; readOnly { + toolOptions = append(toolOptions, mcp.WithReadOnlyHintAnnotation(true)) + } + flags := getAllFlagDefs(leaf) for _, flag := range flags { var propOptions []mcp.PropertyOption @@ -132,6 +147,8 @@ func addToolForCmd(env *azldev.Env, srv *server.MCPServer, leaf *cobra.Command) func handleToolCall( env *azldev.Env, cmd *cobra.Command, ) func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + flagDefaults := captureFlagDefaults(getAllFlagDefs(cmd)) + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { slog.Info("Invoking tool", "tool", cmd.Name(), "params", request.Params.Arguments) @@ -148,6 +165,10 @@ func handleToolCall( }, nil } + if err := restoreFlagDefaults(flagDefaults, args); err != nil { + return nil, fmt.Errorf("failed to reset command flags:\n%w", err) + } + for key, val := range args { if key == "args" { continue @@ -165,32 +186,144 @@ func handleToolCall( slog.Info("Executing command", "args", fullArgs) cmd.Root().SetArgs(fullArgs) - reader, writer, err := os.Pipe() - if err != nil { - return nil, fmt.Errorf("failed to create pipe for command output:\n%w", err) + // LIMITATION: the [azldev.Env] here is created once when the MCP server starts and is reused + // for every tool call. Per-call global flags parsed by this Execute (e.g. '--dry-run', + // '--project') update the App's flag fields but are NOT re-threaded into this reused Env, so + // env.DryRun() and the loaded config still reflect the server's startup values. As a result a + // mutating tool invoked with '--dry-run=true' would still write. Until the MCP rework gives + // each call its own Env, only expose commands that do not modify managed project state + // (specs, locks, config) on their own: read-only commands (marked with + // [azldev.ExportAsReadOnlyMCPTool]), or ones that write solely to a caller-provided output + // path, such as 'docs markdown' or 'component diff-sources --output-file'. + capturedText, execErr := captureStdout(func() error { + env.SetReportFile(os.Stdout) // os.Stdout is the capture pipe for the duration of this call + + return cmd.Root().Execute() + }) + if execErr != nil { + slog.Error("Error executing command", "error", execErr) + + errorText := capturedText + if errorText != "" && !strings.HasSuffix(errorText, "\n") { + errorText += "\n" + } + + errorText += execErr.Error() + + result := mcp.NewToolResultText(errorText) + result.IsError = true + + return result, nil + } + + return mcp.NewToolResultText(capturedText), nil + } +} + +type flagDefault struct { + flag *pflag.Flag + value string + sliceValues []string + changed bool +} + +func (value flagDefault) restore(supplied bool) error { + if sliceValue, ok := value.flag.Value.(pflag.SliceValue); ok { + sliceValues := value.sliceValues + if supplied { + sliceValues = nil } - origStdout := os.Stdout - os.Stdout = writer + if err := sliceValue.Replace(sliceValues); err != nil { + return fmt.Errorf("failed to restore slice value:\n%w", err) + } - env.SetReportFile(writer) + return nil + } - err = cmd.Root().Execute() + if err := value.flag.Value.Set(value.value); err != nil { + return fmt.Errorf("failed to restore value:\n%w", err) + } - os.Stdout = origStdout + return nil +} - if err != nil { - slog.Error("Error executing command", "error", err) +func captureFlagDefaults(flags []*pflag.Flag) []flagDefault { + defaults := make([]flagDefault, 0, len(flags)) - return mcp.NewToolResultText(err.Error()), nil + for _, flag := range flags { + value := flagDefault{ + flag: flag, + value: flag.Value.String(), + changed: flag.Changed, + } + if sliceValue, ok := flag.Value.(pflag.SliceValue); ok { + value.sliceValues = append([]string(nil), sliceValue.GetSlice()...) + } + + defaults = append(defaults, value) + } + + return defaults +} + +func restoreFlagDefaults(defaults []flagDefault, args map[string]any) error { + for _, value := range defaults { + _, supplied := args[value.flag.Name] + if err := value.restore(supplied); err != nil { + return fmt.Errorf("failed to reset flag %#q:\n%w", value.flag.Name, err) } - writer.Close() + value.flag.Changed = value.changed + } - capturedText, _ := io.ReadAll(reader) + return nil +} - return mcp.NewToolResultText(string(capturedText)), nil +// captureStdout runs action with os.Stdout redirected to a pipe and returns everything +// written to it. The pipe is drained by a concurrent goroutine so that output larger +// than the OS pipe buffer (~64KB, e.g. 'config dump' on a large distro) does not block +// the write and hang the caller. Not safe for concurrent use: it mutates the global +// os.Stdout, so callers must serialize (the MCP server pins its worker pool to one). +func captureStdout(action func() error) (string, error) { + reader, writer, err := os.Pipe() + if err != nil { + return "", fmt.Errorf("failed to create pipe for command output:\n%w", err) } + + origStdout := os.Stdout + os.Stdout = writer + + captured := make(chan []byte, 1) + + go func() { + data, _ := io.ReadAll(reader) + captured <- data + }() + + cleanedUp := false + cleanup := func() string { + if cleanedUp { + return "" + } + + cleanedUp = true + os.Stdout = origStdout + _ = writer.Close() // signal EOF so the drain goroutine finishes + output := <-captured + _ = reader.Close() + + return string(output) + } + + defer func() { + _ = cleanup() + }() + + actionErr := action() + output := cleanup() + + return output, actionErr } func getLeafCommands(cmd *cobra.Command) []*cobra.Command { diff --git a/internal/app/azldev/core/mcp/mcpserver_internal_test.go b/internal/app/azldev/core/mcp/mcpserver_internal_test.go new file mode 100644 index 00000000..c0fe4dde --- /dev/null +++ b/internal/app/azldev/core/mcp/mcpserver_internal_test.go @@ -0,0 +1,161 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +package mcp + +import ( + "errors" + "fmt" + "os" + "strings" + "testing" + "time" + + mcpapi "github.com/mark3labs/mcp-go/mcp" + "github.com/microsoft/azure-linux-dev-tools/internal/app/azldev/core/testutils" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestHandleToolCallMarksCommandError(t *testing.T) { + testEnv := testutils.NewTestEnv(t) + root := &cobra.Command{Use: "azldev"} + cmd := &cobra.Command{ + Use: "fail", + RunE: func(_ *cobra.Command, _ []string) error { + fmt.Fprint(os.Stdout, "partial") + + return errors.New("boom") + }, + } + root.AddCommand(cmd) + + result, err := handleToolCall(testEnv.Env, cmd)(t.Context(), mcpapi.CallToolRequest{ + Params: mcpapi.CallToolParams{Arguments: map[string]any{}}, + }) + + require.NoError(t, err) + require.NotNil(t, result) + assert.True(t, result.IsError) + require.Len(t, result.Content, 1) + text, ok := result.Content[0].(mcpapi.TextContent) + require.True(t, ok) + assert.Equal(t, "partial\nboom", text.Text) +} + +func TestHandleToolCallResetsCommandFlags(t *testing.T) { + testEnv := testutils.NewTestEnv(t) + root := &cobra.Command{Use: "azldev"} + + var ( + value string + items []string + ) + + cmd := &cobra.Command{ + Use: "show", + RunE: func(command *cobra.Command, _ []string) error { + fmt.Fprintf(os.Stdout, "%s:%t|%s:%t", + value, command.Flags().Changed("value"), + strings.Join(items, ","), command.Flags().Changed("item")) + + return nil + }, + } + cmd.Flags().StringVar(&value, "value", "default", "value to print") + cmd.Flags().StringArrayVar(&items, "item", []string{"base"}, "item to print") + root.AddCommand(cmd) + handler := handleToolCall(testEnv.Env, cmd) + + first, err := handler(t.Context(), mcpapi.CallToolRequest{ + Params: mcpapi.CallToolParams{Arguments: map[string]any{ + "value": "first", + "item": "one", + }}, + }) + require.NoError(t, err) + require.Len(t, first.Content, 1) + firstText, isText := first.Content[0].(mcpapi.TextContent) + require.True(t, isText) + assert.Equal(t, "first:true|one:true", firstText.Text) + + second, err := handler(t.Context(), mcpapi.CallToolRequest{ + Params: mcpapi.CallToolParams{Arguments: map[string]any{}}, + }) + require.NoError(t, err) + require.Len(t, second.Content, 1) + secondText, isText := second.Content[0].(mcpapi.TextContent) + require.True(t, isText) + assert.Equal(t, "default:false|base:false", secondText.Text) + + third, err := handler(t.Context(), mcpapi.CallToolRequest{ + Params: mcpapi.CallToolParams{Arguments: map[string]any{"item": "two"}}, + }) + require.NoError(t, err) + require.Len(t, third.Content, 1) + thirdText, isText := third.Content[0].(mcpapi.TextContent) + require.True(t, isText) + assert.Equal(t, "default:false|two:true", thirdText.Text) +} + +// TestCaptureStdoutLargeOutput is a regression guard for a pipe deadlock: capturing +// output larger than the OS pipe buffer (~64KB) must not block. A command such as +// 'config dump -f json' on a large distro emits >1MB; without a concurrent drain +// the write blocks and hangs the server. The timeout turns a regression into a +// clean failure instead of a hang. +func TestCaptureStdoutLargeOutput(t *testing.T) { + want := strings.Repeat("x", 1<<20) // 1 MiB, well beyond the pipe buffer + + type result struct { + out string + err error + } + + done := make(chan result, 1) + + go func() { + out, err := captureStdout(func() error { + _, writeErr := fmt.Fprint(os.Stdout, want) + + return writeErr + }) + done <- result{out: out, err: err} + }() + + select { + case got := <-done: + require.NoError(t, got.err) + assert.Equal(t, want, got.out) + case <-time.After(10 * time.Second): + t.Fatal("captureStdout deadlocked on output larger than the pipe buffer") + } +} + +// TestCaptureStdoutReturnsFnError confirms captureStdout returns fn's error alongside +// whatever was written before it failed. +func TestCaptureStdoutReturnsFnError(t *testing.T) { + out, err := captureStdout(func() error { + fmt.Fprint(os.Stdout, "partial") + + return errors.New("boom") + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), "boom") + assert.Equal(t, "partial", out) +} + +func TestCaptureStdoutRestoresStdoutAfterPanic(t *testing.T) { + origStdout := os.Stdout + + assert.PanicsWithValue(t, "boom", func() { + _, _ = captureStdout(func() error { + fmt.Fprint(os.Stdout, "partial") + + panic("boom") + }) + }) + + assert.Same(t, origStdout, os.Stdout) +} diff --git a/scenario/__snapshots__/TestMCPServerMode_1.snap.json b/scenario/__snapshots__/TestMCPServerMode_1.snap.json index 0972acfe..876ca6e7 100755 --- a/scenario/__snapshots__/TestMCPServerMode_1.snap.json +++ b/scenario/__snapshots__/TestMCPServerMode_1.snap.json @@ -8,7 +8,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "Detect which components changed between two git refs", "inputSchema": { @@ -205,7 +205,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "Report per-component change activity and customization detail", "inputSchema": { @@ -300,7 +300,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "List components in this project", "inputSchema": { @@ -385,7 +385,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "Dump the current configuration", "inputSchema": { @@ -457,7 +457,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "Generates JSON schema for validating .toml config files", "inputSchema": { @@ -610,7 +610,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "List images in this project", "inputSchema": { @@ -678,7 +678,7 @@ "destructiveHint": true, "idempotentHint": false, "openWorldHint": true, - "readOnlyHint": false + "readOnlyHint": true }, "description": "List resolved configuration for packages (RPMs and SRPMs)", "inputSchema": {