refactor: move logic out of main into command handlers to simplify

cmd/buildkite-mcp-server/main.go

@@ -7,10 +7,8 @@ import (
 	"time"
 
 	"github.com/alecthomas/kong"
-	buildkitelogs "github.com/buildkite/buildkite-logs"
 	"github.com/buildkite/buildkite-mcp-server/internal/commands"
 	"github.com/buildkite/buildkite-mcp-server/pkg/trace"
-	gobuildkite "github.com/buildkite/go-buildkite/v4"
 	"github.com/mattn/go-isatty"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -20,17 +18,12 @@ var (
 	version = "dev"
 
 	cli struct {
-		Stdio                 commands.StdioCmd `cmd:"" help:"stdio mcp server."`
-		HTTP                  commands.HTTPCmd  `cmd:"" help:"http mcp server. (pass --use-sse to use SSE transport"`
-		Tools                 commands.ToolsCmd `cmd:"" help:"list available tools." hidden:""`
-		APIToken              string            `help:"The Buildkite API token to use." env:"BUILDKITE_API_TOKEN"`
-		APITokenFrom1Password string            `help:"The 1Password item to read the Buildkite API token from. Format: 'op://vault/item/field'" env:"BUILDKITE_API_TOKEN_FROM_1PASSWORD"`
-		BaseURL               string            `help:"The base URL of the Buildkite API to use." env:"BUILDKITE_BASE_URL" default:"https://api.buildkite.com/"`
-		CacheURL              string            `help:"The blob storage URL for job logs cache." env:"BKLOG_CACHE_URL"`
-		Debug                 bool              `help:"Enable debug mode." env:"DEBUG"`
-		OTELExporter          string            `help:"OpenTelemetry exporter to enable. Options are 'http/protobuf', 'grpc', or 'noop'." enum:"http/protobuf, grpc, noop" env:"OTEL_EXPORTER_OTLP_PROTOCOL" default:"noop"`
-		HTTPHeaders           []string          `help:"Additional HTTP headers to send with every request. Format: 'Key: Value'" name:"http-header" env:"BUILDKITE_HTTP_HEADERS"`
-		Version               kong.VersionFlag
+		Stdio        commands.StdioCmd `cmd:"" help:"stdio mcp server."`
+		HTTP         commands.HTTPCmd  `cmd:"" help:"http mcp server. (pass --use-sse to use SSE transport"`
+		Tools        commands.ToolsCmd `cmd:"" help:"list available tools." hidden:""`
+		Debug        bool              `help:"Enable debug mode." env:"DEBUG"`
+		OTELExporter string            `help:"OpenTelemetry exporter to enable. Options are 'http/protobuf', 'grpc', or 'noop'." enum:"http/protobuf, grpc, noop" env:"OTEL_EXPORTER_OTLP_PROTOCOL" default:"noop"`
+		Version      kong.VersionFlag
 	}
 )
 
@@ -62,48 +55,7 @@ func run(ctx context.Context, cmd *kong.Context) error {
 		_ = tp.Shutdown(ctx)
 	}()
 
-	// Parse additional headers into a map
-	headers := commands.ParseHeaders(cli.HTTPHeaders)
-
-	// resolve the api token from either the token or 1password flag
-	apiToken, err := commands.ResolveAPIToken(cli.APIToken, cli.APITokenFrom1Password)
-	if err != nil {
-		return fmt.Errorf("failed to resolve Buildkite API token: %w", err)
-	}
-
-	client, err := gobuildkite.NewOpts(
-		gobuildkite.WithTokenAuth(apiToken),
-		gobuildkite.WithUserAgent(commands.UserAgent(version)),
-		gobuildkite.WithHTTPClient(trace.NewHTTPClientWithHeaders(headers)),
-		gobuildkite.WithBaseURL(cli.BaseURL),
-	)
-	if err != nil {
-		return fmt.Errorf("failed to create buildkite client: %w", err)
-	}
-
-	// Create ParquetClient with cache URL from flag/env (uses upstream library's high-level client)
-	buildkiteLogsClient, err := buildkitelogs.NewClient(ctx, client, cli.CacheURL)
-	if err != nil {
-		return fmt.Errorf("failed to create buildkite logs client: %w", err)
-	}
-
-	buildkiteLogsClient.Hooks().AddAfterCacheCheck(func(ctx context.Context, result *buildkitelogs.CacheCheckResult) {
-		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Checked job logs cache")
-	})
-
-	buildkiteLogsClient.Hooks().AddAfterLogDownload(func(ctx context.Context, result *buildkitelogs.LogDownloadResult) {
-		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Downloaded and cached job logs")
-	})
-
-	buildkiteLogsClient.Hooks().AddAfterLogParsing(func(ctx context.Context, result *buildkitelogs.LogParsingResult) {
-		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Parsed logs to Parquet")
-	})
-
-	buildkiteLogsClient.Hooks().AddAfterBlobStorage(func(ctx context.Context, result *buildkitelogs.BlobStorageResult) {
-		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Stored logs to blob storage")
-	})
-
-	return cmd.Run(&commands.Globals{Version: version, Client: client, BuildkiteLogsClient: buildkiteLogsClient})
+	return cmd.Run(&commands.Globals{Version: version})
 }
 
 func setupLogger(debug bool) zerolog.Logger {

go.sum

@@ -78,10 +78,6 @@ github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMU
 github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
 github.com/buildkite/buildkite-logs v0.6.3 h1:wkH+3IBJg36BC44d7fIDlsn6Zpq99mle+SQPdcMoYeU=
 github.com/buildkite/buildkite-logs v0.6.3/go.mod h1:5kxG9E/GTERyOBpMitXNFi9Q0YyFL+niBUO25TkPQa8=
-github.com/buildkite/go-buildkite/v4 v4.11.0 h1:rEvvUwITrqv433W9JWf6mj+NkkcM45s+ObhNs6C17i4=
-github.com/buildkite/go-buildkite/v4 v4.11.0/go.mod h1:DlebrRJqpZttXDjCW+MJ1QyW9AN++ZWt/UbPtKdbSSk=
-github.com/buildkite/go-buildkite/v4 v4.13.0 h1:ZOngFPSNGWAvFjNSqslAx0WKbxRpOhIFxpKCB41ko7I=
-github.com/buildkite/go-buildkite/v4 v4.13.0/go.mod h1:DlebrRJqpZttXDjCW+MJ1QyW9AN++ZWt/UbPtKdbSSk=
 github.com/buildkite/go-buildkite/v4 v4.13.1 h1:3PhrShdQwlZ2+OIWy0QbxbbjP6M97NCvarlIB6Oye+k=
 github.com/buildkite/go-buildkite/v4 v4.13.1/go.mod h1:DlebrRJqpZttXDjCW+MJ1QyW9AN++ZWt/UbPtKdbSSk=
 github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=

internal/commands/command.go

@@ -1,20 +1,28 @@
 package commands
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os/exec"
 	"runtime"
 
 	buildkitelogs "github.com/buildkite/buildkite-logs"
+	"github.com/buildkite/buildkite-mcp-server/pkg/trace"
 	gobuildkite "github.com/buildkite/go-buildkite/v4"
 	"github.com/rs/zerolog/log"
 )
 
+type APIFlags struct {
+	APIToken              string   `help:"The Buildkite API token to use." env:"BUILDKITE_API_TOKEN"`
+	APITokenFrom1Password string   `help:"The 1Password item to read the Buildkite API token from. Format: 'op://vault/item/field'" env:"BUILDKITE_API_TOKEN_FROM_1PASSWORD"`
+	BaseURL               string   `help:"The base URL of the Buildkite API to use." env:"BUILDKITE_BASE_URL" default:"https://api.buildkite.com/"`
+	CacheURL              string   `help:"The blob storage URL for job logs cache." env:"BKLOG_CACHE_URL"`
+	HTTPHeaders           []string `help:"Additional HTTP headers to send with every request. Format: 'Key: Value'" name:"http-header" env:"BUILDKITE_HTTP_HEADERS"`
+}
+
 type Globals struct {
-	Client              *gobuildkite.Client
-	BuildkiteLogsClient *buildkitelogs.Client
-	Version             string
+	Version string
 }
 
 func UserAgent(version string) string {
@@ -24,7 +32,7 @@ func UserAgent(version string) string {
 	return fmt.Sprintf("buildkite-mcp-server/%s (%s; %s)", version, os, arch)
 }
 
-func ResolveAPIToken(token, tokenFrom1Password string) (string, error) {
+func ResolveAPIToken(ctx context.Context, token, tokenFrom1Password string) (string, error) {
 	if token != "" && tokenFrom1Password != "" {
 		return "", fmt.Errorf("cannot specify both --api-token and --api-token-from-1password")
 	}
@@ -36,16 +44,16 @@ func ResolveAPIToken(token, tokenFrom1Password string) (string, error) {
 	}
 
 	// Fetch the token from 1Password
-	opToken, err := fetchTokenFrom1Password(tokenFrom1Password)
+	opToken, err := fetchTokenFrom1Password(ctx, tokenFrom1Password)
 	if err != nil {
 		return "", fmt.Errorf("failed to fetch API token from 1Password: %w", err)
 	}
 	return opToken, nil
 }
 
-func fetchTokenFrom1Password(opID string) (string, error) {
+func fetchTokenFrom1Password(ctx context.Context, opID string) (string, error) {
 	// read the token using the 1Password CLI with `-n` to avoid a trailing newline
-	out, err := exec.Command("op", "read", "-n", opID).Output()
+	out, err := exec.CommandContext(ctx, "op", "read", "-n", opID).Output()
 	if err != nil {
 		return "", expandExecErr(err)
 	}
@@ -62,3 +70,56 @@ func expandExecErr(err error) error {
 	}
 	return err
 }
+
+func setupBuildkiteAPIClient(ctx context.Context, cli APIFlags, version string) (*gobuildkite.Client, error) {
+	// Parse additional headers into a map
+	headers := ParseHeaders(cli.HTTPHeaders)
+
+	// resolve the api token from either the token or 1password flag
+	apiToken, err := ResolveAPIToken(ctx, cli.APIToken, cli.APITokenFrom1Password)
+	if err != nil {
+		return nil, fmt.Errorf("failed to resolve Buildkite API token: %w", err)
+	}
+
+	client, err := gobuildkite.NewOpts(
+		gobuildkite.WithTokenAuth(apiToken),
+		gobuildkite.WithUserAgent(UserAgent(version)),
+		gobuildkite.WithHTTPClient(trace.NewHTTPClientWithHeaders(headers)),
+		gobuildkite.WithBaseURL(cli.BaseURL),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create buildkite client: %w", err)
+	}
+	return client, nil
+}
+
+func setupBuildkiteLogsClient(ctx context.Context, cli APIFlags, buildkiteClient *gobuildkite.Client) (*buildkitelogs.Client, error) {
+	// Create ParquetClient with cache URL from flag/env (uses upstream library's high-level client)
+	buildkiteLogsClient, err := buildkitelogs.NewClient(
+		ctx,
+		buildkiteClient,
+		cli.CacheURL,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create buildkite logs client: %w", err)
+	}
+
+	// Register debug logging hooks for observability
+	buildkiteLogsClient.Hooks().AddAfterCacheCheck(func(ctx context.Context, result *buildkitelogs.CacheCheckResult) {
+		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Checked job logs cache")
+	})
+
+	buildkiteLogsClient.Hooks().AddAfterLogDownload(func(ctx context.Context, result *buildkitelogs.LogDownloadResult) {
+		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Downloaded and cached job logs")
+	})
+
+	buildkiteLogsClient.Hooks().AddAfterLogParsing(func(ctx context.Context, result *buildkitelogs.LogParsingResult) {
+		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Parsed logs to Parquet")
+	})
+
+	buildkiteLogsClient.Hooks().AddAfterBlobStorage(func(ctx context.Context, result *buildkitelogs.BlobStorageResult) {
+		log.Ctx(ctx).Debug().Str("org", result.Org).Str("pipeline", result.Pipeline).Str("build", result.Build).Str("job", result.Job).Dur("time_taken", result.Duration).Msg("Stored logs to blob storage")
+	})
+
+	return buildkiteLogsClient, nil
+}

internal/commands/http.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/buildkite/buildkite-mcp-server/pkg/middleware"
 	"github.com/buildkite/buildkite-mcp-server/pkg/server"
 	"github.com/buildkite/buildkite-mcp-server/pkg/toolsets"
 	mcpserver "github.com/mark3labs/mcp-go/server"
@@ -15,19 +16,32 @@ import (
 )
 
 type HTTPCmd struct {
+	APIFlags
 	Listen          string   `help:"The address to listen on." default:"localhost:3000" env:"HTTP_LISTEN_ADDR"`
 	UseSSE          bool     `help:"Use deprecated SSS transport instead of Streamable HTTP." default:"false"`
 	EnabledToolsets []string `help:"Comma-separated list of toolsets to enable (e.g., 'pipelines,builds,clusters'). Use 'all' to enable all toolsets." default:"all" env:"BUILDKITE_TOOLSETS"`
 	ReadOnly        bool     `help:"Enable read-only mode, which filters out write operations from all toolsets." default:"false" env:"BUILDKITE_READ_ONLY"`
+	AuthToken       string   `help:"Optional token used to authenticate requests to this HTTP server." env:"BUILDKITE_MCP_AUTH_TOKEN"`
+	TrustProxy      bool     `help:"Trust X-Forwarded-For and other proxy headers for client IP logging. Only enable when behind a trusted reverse proxy." default:"false" env:"BUILDKITE_TRUST_PROXY"`
 }
 
 func (c *HTTPCmd) Run(ctx context.Context, globals *Globals) error {
+	buildkiteClient, err := setupBuildkiteAPIClient(ctx, c.APIFlags, globals.Version)
+	if err != nil {
+		return fmt.Errorf("http server setup: %w", err)
+	}
+
+	buildkiteLogsClient, err := setupBuildkiteLogsClient(ctx, c.APIFlags, buildkiteClient)
+	if err != nil {
+		return fmt.Errorf("http server setup: %w", err)
+	}
+
 	// Validate the enabled toolsets
 	if err := toolsets.ValidateToolsets(c.EnabledToolsets); err != nil {
 		return err
 	}
 
-	mcpServer := server.NewMCPServer(globals.Version, globals.Client, globals.BuildkiteLogsClient,
+	mcpServer := server.NewMCPServer(globals.Version, buildkiteClient, buildkiteLogsClient,
 		server.WithReadOnly(c.ReadOnly), server.WithToolsets(c.EnabledToolsets...))
 
 	listener, err := net.Listen("tcp", c.Listen)
@@ -41,16 +55,25 @@ func (c *HTTPCmd) Run(ctx context.Context, globals *Globals) error {
 
 	mux.HandleFunc("/health", healthHandler)
 
+	// Build middleware chain
+	chain := middleware.NewChain().
+		Use(middleware.ClientIP(c.TrustProxy)).
+		Use(middleware.RequestLog()).
+		UseIf(c.AuthToken != "", middleware.Auth(c.AuthToken))
+
+	var handler http.Handler
 	if c.UseSSE {
-		handler := mcpserver.NewSSEServer(mcpServer)
+		handler = chain.Then(mcpserver.NewSSEServer(mcpServer))
 		mux.Handle("/sse", handler)
 		logEvent.Str("transport", "sse").Str("endpoint", fmt.Sprintf("http://%s/sse", listener.Addr())).Msg("Starting SSE HTTP server")
 	} else {
-		handler := mcpserver.NewStreamableHTTPServer(mcpServer)
+		handler = chain.Then(mcpserver.NewStreamableHTTPServer(mcpServer))
 		mux.Handle("/mcp", handler)
 		logEvent.Str("transport", "streamable-http").Str("endpoint", fmt.Sprintf("http://%s/mcp", listener.Addr())).Msg("Starting Streamable HTTP server")
 	}
 
+	log.Info().Str("address", c.Listen).Msg("starting HTTP server")
+
 	return srv.Serve(listener)
 }
 

internal/commands/stdio.go

@@ -2,6 +2,7 @@ package commands
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/buildkite/buildkite-mcp-server/pkg/server"
 	"github.com/buildkite/buildkite-mcp-server/pkg/toolsets"
@@ -10,17 +11,28 @@ import (
 )
 
 type StdioCmd struct {
+	APIFlags
 	EnabledToolsets []string `help:"Comma-separated list of toolsets to enable (e.g., 'pipelines,builds,clusters'). Use 'all' to enable all toolsets." default:"all" env:"BUILDKITE_TOOLSETS"`
 	ReadOnly        bool     `help:"Enable read-only mode, which filters out write operations from all toolsets." default:"false" env:"BUILDKITE_READ_ONLY"`
 }
 
 func (c *StdioCmd) Run(ctx context.Context, globals *Globals) error {
+	buildkiteClient, err := setupBuildkiteAPIClient(ctx, c.APIFlags, globals.Version)
+	if err != nil {
+		return fmt.Errorf("stdio server setup: %w", err)
+	}
+
+	buildkiteLogsClient, err := setupBuildkiteLogsClient(ctx, c.APIFlags, buildkiteClient)
+	if err != nil {
+		return fmt.Errorf("stdio server setup: %w", err)
+	}
+
 	// Validate the enabled toolsets
 	if err := toolsets.ValidateToolsets(c.EnabledToolsets); err != nil {
 		return err
 	}
 
-	s := server.NewMCPServer(globals.Version, globals.Client, globals.BuildkiteLogsClient,
+	s := server.NewMCPServer(globals.Version, buildkiteClient, buildkiteLogsClient,
 		server.WithReadOnly(c.ReadOnly), server.WithToolsets(c.EnabledToolsets...))
 
 	return mcpserver.ServeStdio(s,

pkg/middleware/auth.go

@@ -0,0 +1,37 @@
+package middleware
+
+import (
+	"crypto/hmac"
+	"net/http"
+	"strings"
+
+	"github.com/rs/zerolog/log"
+)
+
+// Auth creates an HTTP middleware that validates Bearer token authentication.
+// It uses constant-time comparison to prevent timing attacks.
+// The client IP for logging is read from the request context (set by ClientIP middleware).
+func Auth(token string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// extract the token from the Authorization header
+			authHeader := r.Header.Get("Authorization")
+			if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
+				http.Error(w, "Unauthorized", http.StatusUnauthorized)
+				log.Warn().Str("client_ip", GetClientIPFromContext(r.Context())).Msg("Unauthorized access attempt to MCP HTTP server")
+				return
+			}
+
+			// extract the bearer token
+			bearerToken := strings.TrimPrefix(authHeader, "Bearer ")
+
+			// constant time comparison to prevent timing attacks
+			if !hmac.Equal([]byte(bearerToken), []byte(token)) {
+				http.Error(w, "Unauthorized", http.StatusUnauthorized)
+				log.Warn().Str("client_ip", GetClientIPFromContext(r.Context())).Msg("Unauthorized access attempt to MCP HTTP server")
+				return
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}