Merge pull request #21 from buildkite/feat_add_user_orgs

feat: added tool which returns the orginisation associated with a token

Author: wolfeidau

README.md

@@ -13,6 +13,7 @@ This is an [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introd
 * `list_artifacts` - List all artifacts for a specific job in Buildkite
 * `get_artifact` - Get a specific artifact for a specific job in Buildkite
 * `current_user` - Get details of the current user in Buildkite
+* `user_token_organization` - Get the organization associated with the user token used for this request
 
 Example of the `get_pipeline` tool in action.
 

internal/buildkite/organizations.go

@@ -0,0 +1,62 @@
+package buildkite
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/buildkite/go-buildkite/v4"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+	"github.com/rs/zerolog/log"
+)
+
+type OrganizationsClient interface {
+	List(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error)
+}
+
+func UserTokenOrganization(ctx context.Context, client OrganizationsClient) (tool mcp.Tool, handler server.ToolHandlerFunc) {
+	return mcp.NewTool("user_token_organization",
+			mcp.WithDescription("Get the organization associated with the user token used for this request"),
+		), func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			log.Ctx(ctx).Debug().Msg("Getting current user token organization")
+
+			orgs, resp, err := client.List(ctx, &buildkite.OrganizationListOptions{})
+			if err != nil {
+				return mcp.NewToolResultError(err.Error()), nil
+			}
+
+			if resp.StatusCode != 200 {
+				return mcp.NewToolResultError("failed to get current user organizations"), nil
+			}
+
+			if len(orgs) == 0 {
+				return mcp.NewToolResultError("no organization found for the current user token"), nil
+			}
+
+			r, err := json.Marshal(&orgs[0])
+			if err != nil {
+				return nil, fmt.Errorf("failed to marshal user organizations: %w", err)
+			}
+
+			return mcp.NewToolResultText(string(r)), nil
+		}
+}
+
+func HandleUserTokenOrganizationPrompt(
+	ctx context.Context,
+	request mcp.GetPromptRequest,
+) (*mcp.GetPromptResult, error) {
+	return &mcp.GetPromptResult{
+		Description: "When asked for detail of a users pipelines start by looking up the user's token organization",
+		Messages: []mcp.PromptMessage{
+			{
+				Role: mcp.RoleUser,
+				Content: mcp.TextContent{
+					Type: "text",
+					Text: "When asked for detail of a users pipelines start by looking up the user's token organization",
+				},
+			},
+		},
+	}, nil
+}

internal/buildkite/organizations_test.go

@@ -0,0 +1,107 @@
+package buildkite
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/buildkite/go-buildkite/v4"
+	"github.com/stretchr/testify/require"
+)
+
+type MockOrganizationsClient struct {
+	ListFunc func(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error)
+}
+
+func (m *MockOrganizationsClient) List(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error) {
+	if m.ListFunc != nil {
+		return m.ListFunc(ctx, options)
+	}
+	return nil, nil, nil
+}
+
+func TestUserTokenOrganization(t *testing.T) {
+	assert := require.New(t)
+
+	ctx := context.Background()
+	client := &MockOrganizationsClient{
+		ListFunc: func(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error) {
+			return []buildkite.Organization{
+					{
+						Slug: "test-org",
+						Name: "Test Organization",
+					},
+				}, &buildkite.Response{
+					Response: &http.Response{
+						StatusCode: 200,
+					},
+				}, nil
+		},
+	}
+
+	tool, handler := UserTokenOrganization(ctx, client)
+	assert.NotNil(tool)
+	assert.NotNil(handler)
+
+	request := createMCPRequest(t, map[string]any{})
+	result, err := handler(ctx, request)
+	assert.NoError(err)
+
+	textContent := getTextResult(t, result)
+
+	assert.Equal(`{"name":"Test Organization","slug":"test-org"}`, textContent.Text)
+}
+
+func TestUserTokenOrganizationError(t *testing.T) {
+	assert := require.New(t)
+
+	ctx := context.Background()
+	client := &MockOrganizationsClient{
+		ListFunc: func(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error) {
+			return nil, &buildkite.Response{
+				Response: &http.Response{
+					StatusCode: 500,
+				},
+			}, nil
+		},
+	}
+
+	tool, handler := UserTokenOrganization(ctx, client)
+	assert.NotNil(tool)
+	assert.NotNil(handler)
+
+	request := createMCPRequest(t, map[string]any{})
+	result, err := handler(ctx, request)
+	assert.NoError(err)
+
+	textContent := getTextResult(t, result)
+
+	assert.Equal("failed to get current user organizations", textContent.Text)
+}
+
+func TestUserTokenOrganizationErrorNoOrganization(t *testing.T) {
+	assert := require.New(t)
+
+	ctx := context.Background()
+	client := &MockOrganizationsClient{
+		ListFunc: func(ctx context.Context, options *buildkite.OrganizationListOptions) ([]buildkite.Organization, *buildkite.Response, error) {
+			return nil, &buildkite.Response{
+				Response: &http.Response{
+					StatusCode: 200,
+				},
+			}, nil
+		},
+	}
+
+	tool, handler := UserTokenOrganization(ctx, client)
+	assert.NotNil(tool)
+	assert.NotNil(handler)
+
+	request := createMCPRequest(t, map[string]any{})
+	result, err := handler(ctx, request)
+	assert.NoError(err)
+
+	textContent := getTextResult(t, result)
+
+	assert.Equal("no organization found for the current user token", textContent.Text)
+}

internal/commands/stdio.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/buildkite/buildkite-mcp-server/internal/buildkite"
+	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
 	"github.com/rs/zerolog/log"
 )
@@ -26,14 +27,24 @@ func (c *StdioCmd) Run(ctx context.Context, globals *Globals) error {
 
 	s.AddTool(buildkite.GetPipeline(ctx, globals.Client.Pipelines))
 	s.AddTool(buildkite.ListPipelines(ctx, globals.Client.Pipelines))
+
 	s.AddTool(buildkite.ListBuilds(ctx, globals.Client.Builds))
 	s.AddTool(buildkite.GetBuild(ctx, globals.Client.Builds))
+
 	s.AddTool(buildkite.CurrentUser(ctx, globals.Client.User))
+
 	s.AddTool(buildkite.GetJobLogs(ctx, globals.Client))
+
 	s.AddTool(buildkite.AccessToken(ctx, globals.Client.AccessTokens))
 
 	s.AddTool(buildkite.ListArtifacts(ctx, clientAdapter))
 	s.AddTool(buildkite.GetArtifact(ctx, clientAdapter))
 
+	s.AddTool(buildkite.UserTokenOrganization(ctx, globals.Client.Organizations))
+
+	s.AddPrompt(mcp.NewPrompt("user_token_organization_prompt",
+		mcp.WithPromptDescription("When asked for detail of a users pipelines start by looking up the user's token organization"),
+	), buildkite.HandleUserTokenOrganizationPrompt)
+
 	return server.ServeStdio(s)
 }