diff --git a/cmd/buildkite-mcp-server/main.go b/cmd/buildkite-mcp-server/main.go index a63c99b..7f666a4 100644 --- a/cmd/buildkite-mcp-server/main.go +++ b/cmd/buildkite-mcp-server/main.go @@ -2,13 +2,16 @@ package main import ( "context" + "fmt" "os" + "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" ) @@ -17,15 +20,16 @@ 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"` - 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."` - 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:""` + APIToken string `help:"The Buildkite API token to use." env:"BUILDKITE_API_TOKEN"` + 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 } ) @@ -42,24 +46,23 @@ func main() { kong.BindTo(ctx, (*context.Context)(nil)), ) - logger := zerolog.New(os.Stderr).With().Timestamp().Logger() + log.Logger = setupLogger(cli.Debug) - if cli.Debug { - logger = logger.Level(zerolog.DebugLevel).With().Caller().Logger() - } - log.Logger = logger - zerolog.DefaultContextLogger = &logger + err := run(ctx, cmd) + cmd.FatalIfErrorf(err) +} - tp, err := trace.NewProvider(ctx, "buildkite-mcp-server", version) +func run(ctx context.Context, cmd *kong.Context) error { + tp, err := trace.NewProvider(ctx, cli.OTELExporter, "buildkite-mcp-server", version) if err != nil { - logger.Fatal().Err(err).Msg("failed to create trace provider") + return fmt.Errorf("failed to create trace provider: %w", err) } defer func() { _ = tp.Shutdown(ctx) }() // Parse additional headers into a map - headers := commands.ParseHeaders(cli.HTTPHeaders, logger) + headers := commands.ParseHeaders(cli.HTTPHeaders) client, err := gobuildkite.NewOpts( gobuildkite.WithTokenAuth(cli.APIToken), @@ -68,31 +71,49 @@ func main() { gobuildkite.WithBaseURL(cli.BaseURL), ) if err != nil { - logger.Fatal().Err(err).Msg("failed to create buildkite client") + 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 { - logger.Fatal().Err(err).Msg("failed to create buildkite logs client") + return fmt.Errorf("failed to create buildkite logs client: %w", err) } buildkiteLogsClient.Hooks().AddAfterCacheCheck(func(ctx context.Context, result *buildkitelogs.CacheCheckResult) { - log.Ctx(ctx).Info().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") + 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).Info().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") + 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).Info().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") + 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).Info().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") + 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") }) - err = cmd.Run(&commands.Globals{Version: version, Client: client, BuildkiteLogsClient: buildkiteLogsClient, Logger: logger}) - cmd.FatalIfErrorf(err) + return cmd.Run(&commands.Globals{Version: version, Client: client, BuildkiteLogsClient: buildkiteLogsClient}) +} + +func setupLogger(debug bool) zerolog.Logger { + var logger zerolog.Logger + level := zerolog.InfoLevel + if debug { + level = zerolog.DebugLevel + } + + logger = zerolog.New(os.Stderr).Level(level).With().Timestamp().Stack().Logger() + + // are we in an interactive terminal use a console writer + if isatty.IsTerminal(os.Stdout.Fd()) { + logger = logger.Output(zerolog.ConsoleWriter{Out: os.Stderr, FormatTimestamp: func(i any) string { + return time.Now().Format(time.Stamp) + }}).Level(level).With().Stack().Logger() + } + + return logger } diff --git a/go.mod b/go.mod index 001bba0..88fb264 100644 --- a/go.mod +++ b/go.mod @@ -8,13 +8,15 @@ require ( github.com/buildkite/go-buildkite/v4 v4.5.1 github.com/cenkalti/backoff/v5 v5.0.3 github.com/mark3labs/mcp-go v0.39.1 + github.com/mattn/go-isatty v0.0.20 github.com/rs/zerolog v1.34.0 github.com/stretchr/testify v1.11.1 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0 - go.opentelemetry.io/otel v1.37.0 - go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.37.0 - go.opentelemetry.io/otel/sdk v1.37.0 - go.opentelemetry.io/otel/trace v1.37.0 + go.opentelemetry.io/otel v1.38.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.38.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 + go.opentelemetry.io/otel/sdk v1.38.0 + go.opentelemetry.io/otel/trace v1.38.0 ) require ( @@ -56,7 +58,7 @@ require ( github.com/google/uuid v1.6.0 // indirect github.com/google/wire v0.7.0 // indirect github.com/googleapis/gax-go/v2 v2.15.0 // indirect - github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/klauspost/asmfmt v1.3.2 // indirect @@ -64,7 +66,6 @@ require ( github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/mattn/go-colorable v0.1.14 // indirect - github.com/mattn/go-isatty v0.0.20 // indirect github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect github.com/pierrec/lz4/v4 v4.1.22 // indirect @@ -74,10 +75,10 @@ require ( github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/zeebo/xxh3 v1.0.2 // indirect go.opentelemetry.io/auto/sdk v1.1.0 // indirect - go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0 // indirect - go.opentelemetry.io/otel/metric v1.37.0 // indirect - go.opentelemetry.io/otel/sdk/metric v1.37.0 // indirect - go.opentelemetry.io/proto/otlp v1.7.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 // indirect + go.opentelemetry.io/otel/metric v1.38.0 // indirect + go.opentelemetry.io/otel/sdk/metric v1.38.0 // indirect + go.opentelemetry.io/proto/otlp v1.7.1 // indirect gocloud.dev v0.43.0 // indirect golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b // indirect golang.org/x/mod v0.27.0 // indirect @@ -88,7 +89,7 @@ require ( golang.org/x/tools v0.36.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect google.golang.org/api v0.248.0 // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20250715232539-7130f93afb79 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5 // indirect google.golang.org/grpc v1.75.0 // indirect google.golang.org/protobuf v1.36.8 // indirect diff --git a/go.sum b/go.sum index 044b611..3874db3 100644 --- a/go.sum +++ b/go.sum @@ -138,8 +138,8 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81vgd/bo= github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1 h1:X5VWvz21y3gzm9Nw/kaUeku/1+uBhcekkmy4IkffJww= -github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1/go.mod h1:Zanoh4+gvIgluNqcfMVTJueD4wSS5hT7zTt4Mrutd90= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnVTyacbefKhmbLhIhU= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= @@ -216,22 +216,24 @@ go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.6 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.62.0/go.mod h1:ru6KHrNtNHxM4nD/vd6QrLVWgKhxPYgblq4VAtNawTQ= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0 h1:Hf9xI/XLML9ElpiHVDNwvqI0hIFlzV8dgIr35kV1kRU= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0/go.mod h1:NfchwuyNoMcZ5MLHwPrODwUF1HWCXWrL31s8gSAdIKY= -go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ= -go.opentelemetry.io/otel v1.37.0/go.mod h1:ehE/umFRLnuLa/vSccNq9oS1ErUlkkK71gMcN34UG8I= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0 h1:Ahq7pZmv87yiyn3jeFz/LekZmPLLdKejuO3NcK9MssM= -go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0/go.mod h1:MJTqhM0im3mRLw1i8uGHnCvUEeS7VwRyxlLC78PA18M= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.37.0 h1:EtFWSnwW9hGObjkIdmlnWSydO+Qs8OwzfzXLUPg4xOc= -go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.37.0/go.mod h1:QjUEoiGCPkvFZ/MjK6ZZfNOS6mfVEVKYE99dFhuN2LI= -go.opentelemetry.io/otel/metric v1.37.0 h1:mvwbQS5m0tbmqML4NqK+e3aDiO02vsf/WgbsdpcPoZE= -go.opentelemetry.io/otel/metric v1.37.0/go.mod h1:04wGrZurHYKOc+RKeye86GwKiTb9FKm1WHtO+4EVr2E= -go.opentelemetry.io/otel/sdk v1.37.0 h1:ItB0QUqnjesGRvNcmAcU0LyvkVyGJ2xftD29bWdDvKI= -go.opentelemetry.io/otel/sdk v1.37.0/go.mod h1:VredYzxUvuo2q3WRcDnKDjbdvmO0sCzOvVAiY+yUkAg= -go.opentelemetry.io/otel/sdk/metric v1.37.0 h1:90lI228XrB9jCMuSdA0673aubgRobVZFhbjxHHspCPc= -go.opentelemetry.io/otel/sdk/metric v1.37.0/go.mod h1:cNen4ZWfiD37l5NhS+Keb5RXVWZWpRE+9WyVCpbo5ps= -go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4= -go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0= -go.opentelemetry.io/proto/otlp v1.7.0 h1:jX1VolD6nHuFzOYso2E73H85i92Mv8JQYk0K9vz09os= -go.opentelemetry.io/proto/otlp v1.7.0/go.mod h1:fSKjH6YJ7HDlwzltzyMj036AJ3ejJLCgCSHGj4efDDo= +go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8= +go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 h1:GqRJVj7UmLjCVyVJ3ZFLdPRmhDUp2zFmQe3RHIOsw24= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0/go.mod h1:ri3aaHSmCTVYu2AWv44YMauwAQc0aqI9gHKIcSbI1pU= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.38.0 h1:lwI4Dc5leUqENgGuQImwLo4WnuXFPetmPpkLi2IrX54= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.38.0/go.mod h1:Kz/oCE7z5wuyhPxsXDuaPteSWqjSBD5YaSdbxZYGbGk= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 h1:aTL7F04bJHUlztTsNGJ2l+6he8c+y/b//eR0jjjemT4= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0/go.mod h1:kldtb7jDTeol0l3ewcmd8SDvx3EmIE7lyvqbasU3QC4= +go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA= +go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI= +go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E= +go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg= +go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM= +go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA= +go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE= +go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs= +go.opentelemetry.io/proto/otlp v1.7.1 h1:gTOMpGDb0WTBOP8JaO72iL3auEZhVmAQg4ipjOVAtj4= +go.opentelemetry.io/proto/otlp v1.7.1/go.mod h1:b2rVh6rfI/s2pHWNlB7ILJcRALpcNDzKhACevjI+ZnE= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= gocloud.dev v0.43.0 h1:aW3eq4RMyehbJ54PMsh4hsp7iX8cO/98ZRzJJOzN/5M= @@ -268,8 +270,8 @@ google.golang.org/api v0.248.0 h1:hUotakSkcwGdYUqzCRc5yGYsg4wXxpkKlW5ryVqvC1Y= google.golang.org/api v0.248.0/go.mod h1:yAFUAF56Li7IuIQbTFoLwXTCI6XCFKueOlS7S9e4F9k= google.golang.org/genproto v0.0.0-20250715232539-7130f93afb79 h1:Nt6z9UHqSlIdIGJdz6KhTIs2VRx/iOsA5iE8bmQNcxs= google.golang.org/genproto v0.0.0-20250715232539-7130f93afb79/go.mod h1:kTmlBHMPqR5uCZPBvwa2B18mvubkjyY3CRLI0c6fj0s= -google.golang.org/genproto/googleapis/api v0.0.0-20250715232539-7130f93afb79 h1:iOye66xuaAK0WnkPuhQPUFy8eJcmwUXqGGP3om6IxX8= -google.golang.org/genproto/googleapis/api v0.0.0-20250715232539-7130f93afb79/go.mod h1:HKJDgKsFUnv5VAGeQjz8kxcgDP0HoE0iZNp0OdZNlhE= +google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5 h1:BIRfGDEjiHRrk0QKZe3Xv2ieMhtgRGeLcZQ0mIVn4EY= +google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5/go.mod h1:j3QtIyytwqGr1JUDtYXwtMXWPKsEa5LtzIFN1Wn5WvE= google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5 h1:eaY8u2EuxbRv7c3NiGK0/NedzVsCcV6hDuU5qPX5EGE= google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5/go.mod h1:M4/wBTSeyLxupu3W3tJtOgB14jILAS/XWPSSa3TAlJc= google.golang.org/grpc v1.75.0 h1:+TW+dqTd2Biwe6KKfhE5JpiYIBWq865PhKGSXiivqt4= diff --git a/internal/commands/command.go b/internal/commands/command.go index ed57283..c8600d2 100644 --- a/internal/commands/command.go +++ b/internal/commands/command.go @@ -6,14 +6,12 @@ import ( buildkitelogs "github.com/buildkite/buildkite-logs" gobuildkite "github.com/buildkite/go-buildkite/v4" - "github.com/rs/zerolog" ) type Globals struct { Client *gobuildkite.Client BuildkiteLogsClient *buildkitelogs.Client Version string - Logger zerolog.Logger } func UserAgent(version string) string { diff --git a/internal/commands/headers.go b/internal/commands/headers.go index 7e02930..28bf359 100644 --- a/internal/commands/headers.go +++ b/internal/commands/headers.go @@ -3,13 +3,13 @@ package commands import ( "strings" - "github.com/rs/zerolog" + "github.com/rs/zerolog/log" ) // ParseHeaders takes a slice of header strings in the format "Key: Value" // and returns a map of headers. This is used to parse additional HTTP headers // that can be sent with every request to the Buildkite API. -func ParseHeaders(headerStrings []string, logger zerolog.Logger) map[string]string { +func ParseHeaders(headerStrings []string) map[string]string { headers := make(map[string]string) for _, h := range headerStrings { parts := strings.SplitN(h, ":", 2) @@ -17,9 +17,9 @@ func ParseHeaders(headerStrings []string, logger zerolog.Logger) map[string]stri key := strings.TrimSpace(parts[0]) value := strings.TrimSpace(parts[1]) headers[key] = value - logger.Debug().Str("key", key).Str("value", value).Msg("parsed header") + log.Debug().Str("key", key).Str("value", value).Msg("parsed header") } else { - logger.Warn().Str("header", h).Msg("invalid header format, expected 'Key: Value'") + log.Warn().Str("header", h).Msg("invalid header format, expected 'Key: Value'") } } return headers diff --git a/internal/commands/headers_test.go b/internal/commands/headers_test.go index 2c49a58..e466500 100644 --- a/internal/commands/headers_test.go +++ b/internal/commands/headers_test.go @@ -2,8 +2,6 @@ package commands import ( "testing" - - "github.com/rs/zerolog" ) func TestParseHeaders(t *testing.T) { @@ -22,10 +20,8 @@ func TestParseHeaders(t *testing.T) { {[]string{"A:1", "NoColon", "B:2"}, map[string]string{"A": "1", "B": "2"}}, } - logger := zerolog.Nop() - for _, tt := range tests { - got := ParseHeaders(tt.input, logger) + got := ParseHeaders(tt.input) if len(got) != len(tt.want) { t.Errorf("parseHeaders(%v) = %v, want %v", tt.input, got, tt.want) continue diff --git a/internal/commands/stdio.go b/internal/commands/stdio.go index 67760f1..e84dced 100644 --- a/internal/commands/stdio.go +++ b/internal/commands/stdio.go @@ -6,6 +6,7 @@ import ( "github.com/buildkite/buildkite-mcp-server/pkg/server" "github.com/buildkite/buildkite-mcp-server/pkg/toolsets" mcpserver "github.com/mark3labs/mcp-go/server" + "github.com/rs/zerolog/log" ) type StdioCmd struct { @@ -31,8 +32,9 @@ func (c *StdioCmd) Run(ctx context.Context, globals *Globals) error { func setupContext(globals *Globals) mcpserver.StdioContextFunc { return func(ctx context.Context) context.Context { + log.Info().Msg("Starting MCP server over stdio") // add the logger to the context - return globals.Logger.WithContext(ctx) + return log.Logger.WithContext(ctx) } } diff --git a/pkg/server/mcp.go b/pkg/server/mcp.go index abd7ed2..0b097db 100644 --- a/pkg/server/mcp.go +++ b/pkg/server/mcp.go @@ -53,6 +53,8 @@ func NewMCPServer(version string, client *gobuildkite.Client, buildkiteLogsClien server.WithToolCapabilities(true), server.WithPromptCapabilities(true), server.WithResourceCapabilities(true, true), + server.WithToolHandlerMiddleware(trace.ToolHandlerFunc), + server.WithResourceHandlerMiddleware(trace.WithResourceHandlerFunc), server.WithHooks(trace.NewHooks()), server.WithLogging()) diff --git a/pkg/trace/trace.go b/pkg/trace/trace.go index fd8c183..dc22ab5 100644 --- a/pkg/trace/trace.go +++ b/pkg/trace/trace.go @@ -7,23 +7,26 @@ import ( "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" + "github.com/rs/zerolog/log" "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp" "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/sdk/resource" sdktrace "go.opentelemetry.io/otel/sdk/trace" - semconv "go.opentelemetry.io/otel/semconv/v1.34.0" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + semconv "go.opentelemetry.io/otel/semconv/v1.37.0" "go.opentelemetry.io/otel/trace" ) // set a default tracer name var tracerName = "buildkite-mcp-server" -func NewProvider(ctx context.Context, name, version string) (*sdktrace.TracerProvider, error) { - exp, err := otlptracegrpc.New(ctx) +func NewProvider(ctx context.Context, exporter, name, version string) (*sdktrace.TracerProvider, error) { + exp, err := newExporter(ctx, exporter) if err != nil { return nil, fmt.Errorf("failed to create exporter: %w", err) } @@ -55,24 +58,6 @@ func Start(ctx context.Context, name string) (context.Context, trace.Span) { return otel.GetTracerProvider().Tracer(tracerName).Start(ctx, name) } -func newResource(cxt context.Context, name, version string) (*resource.Resource, error) { - options := []resource.Option{ - resource.WithSchemaURL(semconv.SchemaURL), - } - options = append(options, resource.WithHost()) - options = append(options, resource.WithFromEnv()) - options = append(options, resource.WithAttributes( - semconv.TelemetrySDKNameKey.String("otelconfig"), - semconv.TelemetrySDKLanguageGo, - semconv.TelemetrySDKVersionKey.String(version), - )) - - return resource.New( - cxt, - options..., - ) -} - func NewError(span trace.Span, msg string, args ...any) error { if span == nil { return fmt.Errorf("span is nil: %w", fmt.Errorf(msg, args...)) @@ -112,17 +97,96 @@ func (h *headerInjector) RoundTrip(req *http.Request) (*http.Response, error) { return h.wrapped.RoundTrip(req) } +func newResource(cxt context.Context, name, version string) (*resource.Resource, error) { + options := []resource.Option{ + resource.WithSchemaURL(semconv.SchemaURL), + } + options = append(options, resource.WithHost()) + options = append(options, resource.WithFromEnv()) + options = append(options, resource.WithAttributes( + semconv.TelemetrySDKNameKey.String("otelconfig"), + semconv.TelemetrySDKLanguageGo, + semconv.TelemetrySDKVersionKey.String(version), + )) + + return resource.New( + cxt, + options..., + ) +} + +func newExporter(ctx context.Context, exporter string) (sdktrace.SpanExporter, error) { + switch exporter { + case "http/protobuf": + return otlptracehttp.New(ctx) + case "grpc": + return otlptracegrpc.New(ctx) + default: + return tracetest.NewNoopExporter(), nil + } +} + func NewHooks() *server.Hooks { hooks := &server.Hooks{} - hooks.AddOnError(func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) { + hooks.AddOnRegisterSession(func(ctx context.Context, session server.ClientSession) { span := trace.SpanFromContext(ctx) if span != nil { - span.SetAttributes(attribute.String("mcp.method", string(method))) - span.RecordError(err) - span.SetStatus(codes.Error, err.Error()) + span.SetAttributes(attribute.String("mcp.session.id", session.SessionID())) } }) return hooks } + +func ToolHandlerFunc(thf server.ToolHandlerFunc) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + ctx, span := Start(ctx, "mcp.ToolHandler") + defer span.End() + + span.SetAttributes( + attribute.String("mcp.method.name", request.Method), + attribute.String("mcp.tool.name", request.Params.Name), + ) + + log.Debug().Str("mcp.tool.name", request.Params.Name).Msg("Handling MCP tool call") + + res, err := thf(ctx, request) + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + log.Error().Err(err).Str("mcp.tool.name", request.Params.Name).Msg("Error in MCP tool call") + } else { + span.SetStatus(codes.Ok, "OK") + log.Debug().Str("mcp.tool.name", request.Params.Name).Msg("Completed MCP tool call successfully") + } + + return res, err + } +} + +func WithResourceHandlerFunc(rhf server.ResourceHandlerFunc) server.ResourceHandlerFunc { + return func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) { + ctx, span := Start(ctx, "mcp.ResourceHandler") + defer span.End() + + span.SetAttributes( + attribute.String("mcp.method.name", request.Method), + attribute.String("mcp.resource.uri", request.Params.URI), + ) + + log.Debug().Str("mcp.resource.uri", request.Params.URI).Str("mcp.method.name", request.Method).Msg("Handling MCP resource call") + + res, err := rhf(ctx, request) + if err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + log.Error().Err(err).Str("mcp.resource.uri", request.Params.URI).Msg("Error in MCP resource call") + } else { + span.SetStatus(codes.Ok, "OK") + log.Debug().Str("mcp.resource.uri", request.Params.URI).Msg("Completed MCP resource call successfully") + } + + return res, err + } +} diff --git a/pkg/trace/trace_test.go b/pkg/trace/trace_test.go index ac2a277..d073acd 100644 --- a/pkg/trace/trace_test.go +++ b/pkg/trace/trace_test.go @@ -3,15 +3,24 @@ package trace import ( "context" "testing" + + "github.com/stretchr/testify/require" ) func TestNewProvider(t *testing.T) { - provider, err := NewProvider(context.Background(), "test", "1.2.3") - if err != nil { - t.Fatalf("expected no error, got %v", err) - } - - if provider == nil { - t.Fatal("expected non-nil TracerProvider") - } + assert := require.New(t) + + provider, err := NewProvider(context.Background(), "http/protobuf", "test", "1.2.3") + assert.NoError(err) + + assert.NotNil(provider) + + provider, err = NewProvider(context.Background(), "grpc", "test", "1.2.3") + assert.NoError(err) + + assert.NotNil(provider) + + _, err = NewProvider(context.Background(), "", "test", "1.2.3") + assert.NoError(err) + }