From 7f562127799fc8814fc3f2dc743f0b8e0cd12039 Mon Sep 17 00:00:00 2001 From: Asher Date: Sun, 21 Sep 2025 15:12:57 -0800 Subject: [PATCH] feat: add coder_workspace_port_forward MCP tool (#19863) Closes https://github.com/coder/internal/issues/784 --- codersdk/toolsdk/bash_test.go | 10 ++--- codersdk/toolsdk/toolsdk.go | 69 +++++++++++++++++++++++++--- codersdk/toolsdk/toolsdk_test.go | 77 ++++++++++++++++++++++++++++---- 3 files changed, 138 insertions(+), 18 deletions(-) diff --git a/codersdk/toolsdk/bash_test.go b/codersdk/toolsdk/bash_test.go index da05a71ce3..003dd7fcbc 100644 --- a/codersdk/toolsdk/bash_test.go +++ b/codersdk/toolsdk/bash_test.go @@ -217,7 +217,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) { // Scenario: echo "123"; sleep 60; echo "456" with 5s timeout // In this scenario, we'd expect to see "123" in the output and a cancellation message - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start the agent and wait for it to be fully ready _ = agenttest.New(t, client.URL, agentToken) @@ -259,7 +259,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) { // Test that normal commands still work with timeout functionality present - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start the agent and wait for it to be fully ready _ = agenttest.New(t, client.URL, agentToken) @@ -304,7 +304,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) { t.Run("BackgroundCommandCapturesOutput", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start the agent and wait for it to be fully ready _ = agenttest.New(t, client.URL, agentToken) @@ -345,7 +345,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) { t.Run("BackgroundVsNormalExecution", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start the agent and wait for it to be fully ready _ = agenttest.New(t, client.URL, agentToken) @@ -391,7 +391,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) { t.Run("BackgroundCommandContinuesAfterTimeout", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start the agent and wait for it to be fully ready _ = agenttest.New(t, client.URL, agentToken) diff --git a/codersdk/toolsdk/toolsdk.go b/codersdk/toolsdk/toolsdk.go index 5cd437eb36..acc7a9b102 100644 --- a/codersdk/toolsdk/toolsdk.go +++ b/codersdk/toolsdk/toolsdk.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "runtime/debug" + "strconv" "strings" "github.com/google/uuid" @@ -17,6 +18,7 @@ import ( "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/coderd/workspaceapps/appurl" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" ) @@ -47,6 +49,7 @@ const ( ToolNameWorkspaceWriteFile = "coder_workspace_write_file" ToolNameWorkspaceEditFile = "coder_workspace_edit_file" ToolNameWorkspaceEditFiles = "coder_workspace_edit_files" + ToolNameWorkspacePortForward = "coder_workspace_port_forward" ) func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) { @@ -219,6 +222,7 @@ var All = []GenericTool{ WorkspaceWriteFile.Generic(), WorkspaceEditFile.Generic(), WorkspaceEditFiles.Generic(), + WorkspacePortForward.Generic(), } type ReportTaskArgs struct { @@ -1389,6 +1393,8 @@ type WorkspaceLSResponse struct { Contents []WorkspaceLSFile `json:"contents"` } +const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used." + var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{ Tool: aisdk.Tool{ Name: ToolNameWorkspaceLS, @@ -1397,7 +1403,7 @@ var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{ Properties: map[string]any{ "workspace": map[string]any{ "type": "string", - "description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.", + "description": workspaceDescription, }, "path": map[string]any{ "type": "string", @@ -1454,7 +1460,7 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{ Properties: map[string]any{ "workspace": map[string]any{ "type": "string", - "description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.", + "description": workspaceDescription, }, "path": map[string]any{ "type": "string", @@ -1519,7 +1525,7 @@ var WorkspaceWriteFile = Tool[WorkspaceWriteFileArgs, codersdk.Response]{ Properties: map[string]any{ "workspace": map[string]any{ "type": "string", - "description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.", + "description": workspaceDescription, }, "path": map[string]any{ "type": "string", @@ -1567,7 +1573,7 @@ var WorkspaceEditFile = Tool[WorkspaceEditFileArgs, codersdk.Response]{ Properties: map[string]any{ "workspace": map[string]any{ "type": "string", - "description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.", + "description": workspaceDescription, }, "path": map[string]any{ "type": "string", @@ -1634,7 +1640,7 @@ var WorkspaceEditFiles = Tool[WorkspaceEditFilesArgs, codersdk.Response]{ Properties: map[string]any{ "workspace": map[string]any{ "type": "string", - "description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.", + "description": workspaceDescription, }, "files": map[string]any{ "type": "array", @@ -1691,6 +1697,59 @@ var WorkspaceEditFiles = Tool[WorkspaceEditFilesArgs, codersdk.Response]{ }, } +type WorkspacePortForwardArgs struct { + Workspace string `json:"workspace"` + Port int `json:"port"` +} + +type WorkspacePortForwardResponse struct { + URL string `json:"url"` +} + +var WorkspacePortForward = Tool[WorkspacePortForwardArgs, WorkspacePortForwardResponse]{ + Tool: aisdk.Tool{ + Name: ToolNameWorkspacePortForward, + Description: `Fetch URLs that forward to the specified port.`, + Schema: aisdk.Schema{ + Properties: map[string]any{ + "workspace": map[string]any{ + "type": "string", + "description": workspaceDescription, + }, + "port": map[string]any{ + "type": "number", + "description": "The port to forward.", + }, + }, + Required: []string{"workspace", "port"}, + }, + }, + UserClientOptional: true, + Handler: func(ctx context.Context, deps Deps, args WorkspacePortForwardArgs) (WorkspacePortForwardResponse, error) { + workspaceName := NormalizeWorkspaceInput(args.Workspace) + workspace, workspaceAgent, err := findWorkspaceAndAgent(ctx, deps.coderClient, workspaceName) + if err != nil { + return WorkspacePortForwardResponse{}, xerrors.Errorf("failed to find workspace: %w", err) + } + res, err := deps.coderClient.AppHost(ctx) + if err != nil { + return WorkspacePortForwardResponse{}, xerrors.Errorf("failed to get app host: %w", err) + } + if res.Host == "" { + return WorkspacePortForwardResponse{}, xerrors.New("no app host for forwarding has been configured") + } + url := appurl.ApplicationURL{ + AppSlugOrPort: strconv.Itoa(args.Port), + AgentName: workspaceAgent.Name, + WorkspaceName: workspace.Name, + Username: workspace.OwnerName, + } + return WorkspacePortForwardResponse{ + URL: deps.coderClient.URL.Scheme + "://" + strings.Replace(res.Host, "*", url.String(), 1), + }, nil + }, +} + // NormalizeWorkspaceInput converts workspace name input to standard format. // Handles the following input formats: // - workspace → workspace diff --git a/codersdk/toolsdk/toolsdk_test.go b/codersdk/toolsdk/toolsdk_test.go index 39e2a08e4a..f89f22e008 100644 --- a/codersdk/toolsdk/toolsdk_test.go +++ b/codersdk/toolsdk/toolsdk_test.go @@ -3,6 +3,7 @@ package toolsdk_test import ( "context" "encoding/json" + "fmt" "os" "path/filepath" "runtime" @@ -35,10 +36,10 @@ import ( // setupWorkspaceForAgent creates a workspace setup exactly like main SSH tests // nolint:gocritic // This is in a test package and does not end up in the build -func setupWorkspaceForAgent(t *testing.T) (*codersdk.Client, database.WorkspaceTable, string) { +func setupWorkspaceForAgent(t *testing.T, opts *coderdtest.Options) (*codersdk.Client, database.WorkspaceTable, string) { t.Helper() - client, store := coderdtest.NewWithDatabase(t, nil) + client, store := coderdtest.NewWithDatabase(t, opts) client.SetLogger(testutil.Logger(t).Named("client")) first := coderdtest.CreateFirstUser(t, client) userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) { @@ -405,7 +406,7 @@ func TestTools(t *testing.T) { t.Skip("WorkspaceSSHExec is not supported on Windows") } // Setup workspace exactly like main SSH tests - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) // Start agent and wait for it to be ready (following main SSH test pattern) _ = agenttest.New(t, client.URL, agentToken) @@ -457,7 +458,7 @@ func TestTools(t *testing.T) { t.Run("WorkspaceLS", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) fs := afero.NewMemMapFs() _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { opts.Filesystem = fs @@ -503,7 +504,7 @@ func TestTools(t *testing.T) { t.Run("WorkspaceReadFile", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) fs := afero.NewMemMapFs() _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { opts.Filesystem = fs @@ -606,7 +607,7 @@ func TestTools(t *testing.T) { t.Run("WorkspaceWriteFile", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) fs := afero.NewMemMapFs() _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { opts.Filesystem = fs @@ -633,7 +634,7 @@ func TestTools(t *testing.T) { t.Run("WorkspaceEditFile", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) fs := afero.NewMemMapFs() _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { opts.Filesystem = fs @@ -673,7 +674,7 @@ func TestTools(t *testing.T) { t.Run("WorkspaceEditFiles", func(t *testing.T) { t.Parallel() - client, workspace, agentToken := setupWorkspaceForAgent(t) + client, workspace, agentToken := setupWorkspaceForAgent(t, nil) fs := afero.NewMemMapFs() _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { opts.Filesystem = fs @@ -730,6 +731,66 @@ func TestTools(t *testing.T) { require.NoError(t, err) require.Equal(t, "bar2 bar2", string(b)) }) + + t.Run("WorkspacePortForward", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + workspace string + host string + port int + expect string + error string + }{ + { + name: "OK", + workspace: "myuser/myworkspace", + port: 1234, + host: "*.test.coder.com", + expect: "%s://1234--dev--myworkspace--myuser.test.coder.com:%s", + }, + { + name: "NonExistentWorkspace", + workspace: "doesnotexist", + port: 1234, + host: "*.test.coder.com", + error: "failed to find workspace", + }, + { + name: "NoAppHost", + host: "", + workspace: "myuser/myworkspace", + port: 1234, + error: "no app host", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + client, workspace, agentToken := setupWorkspaceForAgent(t, &coderdtest.Options{ + AppHostname: tt.host, + }) + _ = agenttest.New(t, client.URL, agentToken) + coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait() + tb, err := toolsdk.NewDeps(client) + require.NoError(t, err) + + res, err := testTool(t, toolsdk.WorkspacePortForward, tb, toolsdk.WorkspacePortForwardArgs{ + Workspace: tt.workspace, + Port: tt.port, + }) + if tt.error != "" { + require.Error(t, err) + require.ErrorContains(t, err, tt.error) + } else { + require.NoError(t, err) + require.Equal(t, fmt.Sprintf(tt.expect, client.URL.Scheme, client.URL.Port()), res.URL) + } + }) + } + }) } // TestedTools keeps track of which tools have been tested.