feat: add coder_workspace_port_forward MCP tool (#19863)

Closes https://github.com/coder/internal/issues/784
This commit is contained in:
Asher
2025-09-21 15:12:57 -08:00
committed by GitHub
parent d464360103
commit 7f56212779
3 changed files with 138 additions and 18 deletions
+5 -5
View File
@@ -217,7 +217,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) {
// Scenario: echo "123"; sleep 60; echo "456" with 5s timeout // 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 // 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 // Start the agent and wait for it to be fully ready
_ = agenttest.New(t, client.URL, agentToken) _ = 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 // 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 // Start the agent and wait for it to be fully ready
_ = agenttest.New(t, client.URL, agentToken) _ = agenttest.New(t, client.URL, agentToken)
@@ -304,7 +304,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
t.Run("BackgroundCommandCapturesOutput", func(t *testing.T) { t.Run("BackgroundCommandCapturesOutput", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
// Start the agent and wait for it to be fully ready // Start the agent and wait for it to be fully ready
_ = agenttest.New(t, client.URL, agentToken) _ = agenttest.New(t, client.URL, agentToken)
@@ -345,7 +345,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
t.Run("BackgroundVsNormalExecution", func(t *testing.T) { t.Run("BackgroundVsNormalExecution", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
// Start the agent and wait for it to be fully ready // Start the agent and wait for it to be fully ready
_ = agenttest.New(t, client.URL, agentToken) _ = agenttest.New(t, client.URL, agentToken)
@@ -391,7 +391,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
t.Run("BackgroundCommandContinuesAfterTimeout", func(t *testing.T) { t.Run("BackgroundCommandContinuesAfterTimeout", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
// Start the agent and wait for it to be fully ready // Start the agent and wait for it to be fully ready
_ = agenttest.New(t, client.URL, agentToken) _ = agenttest.New(t, client.URL, agentToken)
+64 -5
View File
@@ -8,6 +8,7 @@ import (
"fmt" "fmt"
"io" "io"
"runtime/debug" "runtime/debug"
"strconv"
"strings" "strings"
"github.com/google/uuid" "github.com/google/uuid"
@@ -17,6 +18,7 @@ import (
"github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/buildinfo"
"github.com/coder/coder/v2/cli/cliui" "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"
"github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/codersdk/workspacesdk"
) )
@@ -47,6 +49,7 @@ const (
ToolNameWorkspaceWriteFile = "coder_workspace_write_file" ToolNameWorkspaceWriteFile = "coder_workspace_write_file"
ToolNameWorkspaceEditFile = "coder_workspace_edit_file" ToolNameWorkspaceEditFile = "coder_workspace_edit_file"
ToolNameWorkspaceEditFiles = "coder_workspace_edit_files" ToolNameWorkspaceEditFiles = "coder_workspace_edit_files"
ToolNameWorkspacePortForward = "coder_workspace_port_forward"
) )
func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) { func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) {
@@ -219,6 +222,7 @@ var All = []GenericTool{
WorkspaceWriteFile.Generic(), WorkspaceWriteFile.Generic(),
WorkspaceEditFile.Generic(), WorkspaceEditFile.Generic(),
WorkspaceEditFiles.Generic(), WorkspaceEditFiles.Generic(),
WorkspacePortForward.Generic(),
} }
type ReportTaskArgs struct { type ReportTaskArgs struct {
@@ -1389,6 +1393,8 @@ type WorkspaceLSResponse struct {
Contents []WorkspaceLSFile `json:"contents"` 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]{ var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{
Tool: aisdk.Tool{ Tool: aisdk.Tool{
Name: ToolNameWorkspaceLS, Name: ToolNameWorkspaceLS,
@@ -1397,7 +1403,7 @@ var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{
Properties: map[string]any{ Properties: map[string]any{
"workspace": map[string]any{ "workspace": map[string]any{
"type": "string", "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{ "path": map[string]any{
"type": "string", "type": "string",
@@ -1454,7 +1460,7 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
Properties: map[string]any{ Properties: map[string]any{
"workspace": map[string]any{ "workspace": map[string]any{
"type": "string", "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{ "path": map[string]any{
"type": "string", "type": "string",
@@ -1519,7 +1525,7 @@ var WorkspaceWriteFile = Tool[WorkspaceWriteFileArgs, codersdk.Response]{
Properties: map[string]any{ Properties: map[string]any{
"workspace": map[string]any{ "workspace": map[string]any{
"type": "string", "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{ "path": map[string]any{
"type": "string", "type": "string",
@@ -1567,7 +1573,7 @@ var WorkspaceEditFile = Tool[WorkspaceEditFileArgs, codersdk.Response]{
Properties: map[string]any{ Properties: map[string]any{
"workspace": map[string]any{ "workspace": map[string]any{
"type": "string", "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{ "path": map[string]any{
"type": "string", "type": "string",
@@ -1634,7 +1640,7 @@ var WorkspaceEditFiles = Tool[WorkspaceEditFilesArgs, codersdk.Response]{
Properties: map[string]any{ Properties: map[string]any{
"workspace": map[string]any{ "workspace": map[string]any{
"type": "string", "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{ "files": map[string]any{
"type": "array", "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. // NormalizeWorkspaceInput converts workspace name input to standard format.
// Handles the following input formats: // Handles the following input formats:
// - workspace → workspace // - workspace → workspace
+69 -8
View File
@@ -3,6 +3,7 @@ package toolsdk_test
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"fmt"
"os" "os"
"path/filepath" "path/filepath"
"runtime" "runtime"
@@ -35,10 +36,10 @@ import (
// setupWorkspaceForAgent creates a workspace setup exactly like main SSH tests // 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 // 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() t.Helper()
client, store := coderdtest.NewWithDatabase(t, nil) client, store := coderdtest.NewWithDatabase(t, opts)
client.SetLogger(testutil.Logger(t).Named("client")) client.SetLogger(testutil.Logger(t).Named("client"))
first := coderdtest.CreateFirstUser(t, client) first := coderdtest.CreateFirstUser(t, client)
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) { 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") t.Skip("WorkspaceSSHExec is not supported on Windows")
} }
// Setup workspace exactly like main SSH tests // 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) // Start agent and wait for it to be ready (following main SSH test pattern)
_ = agenttest.New(t, client.URL, agentToken) _ = agenttest.New(t, client.URL, agentToken)
@@ -457,7 +458,7 @@ func TestTools(t *testing.T) {
t.Run("WorkspaceLS", func(t *testing.T) { t.Run("WorkspaceLS", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
fs := afero.NewMemMapFs() fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs opts.Filesystem = fs
@@ -503,7 +504,7 @@ func TestTools(t *testing.T) {
t.Run("WorkspaceReadFile", func(t *testing.T) { t.Run("WorkspaceReadFile", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
fs := afero.NewMemMapFs() fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs opts.Filesystem = fs
@@ -606,7 +607,7 @@ func TestTools(t *testing.T) {
t.Run("WorkspaceWriteFile", func(t *testing.T) { t.Run("WorkspaceWriteFile", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
fs := afero.NewMemMapFs() fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs opts.Filesystem = fs
@@ -633,7 +634,7 @@ func TestTools(t *testing.T) {
t.Run("WorkspaceEditFile", func(t *testing.T) { t.Run("WorkspaceEditFile", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
fs := afero.NewMemMapFs() fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs opts.Filesystem = fs
@@ -673,7 +674,7 @@ func TestTools(t *testing.T) {
t.Run("WorkspaceEditFiles", func(t *testing.T) { t.Run("WorkspaceEditFiles", func(t *testing.T) {
t.Parallel() t.Parallel()
client, workspace, agentToken := setupWorkspaceForAgent(t) client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
fs := afero.NewMemMapFs() fs := afero.NewMemMapFs()
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) { _ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
opts.Filesystem = fs opts.Filesystem = fs
@@ -730,6 +731,66 @@ func TestTools(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "bar2 bar2", string(b)) 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. // TestedTools keeps track of which tools have been tested.