mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat: add coder_workspace_port_forward MCP tool (#19863)
Closes https://github.com/coder/internal/issues/784
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user