Skip to content

Commit 7f56212

Browse files
code-ashercode-asher
authored
feat: add coder_workspace_port_forward MCP tool (coder#19863)
Closes coder/internal#784
1 parent d464360 commit 7f56212

File tree

3 files changed

+138
-18
lines changed

3 files changed

+138
-18
lines changed

‎codersdk/toolsdk/bash_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) {
217217
// Scenario: echo "123"; sleep 60; echo "456" with 5s timeout
218218
// In this scenario, we'd expect to see "123" in the output and a cancellation message
219219

220-
client, workspace, agentToken := setupWorkspaceForAgent(t)
220+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
221221

222222
// Start the agent and wait for it to be fully ready
223223
_ = agenttest.New(t, client.URL, agentToken)
@@ -259,7 +259,7 @@ func TestWorkspaceBashTimeoutIntegration(t *testing.T) {
259259

260260
// Test that normal commands still work with timeout functionality present
261261

262-
client, workspace, agentToken := setupWorkspaceForAgent(t)
262+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
263263

264264
// Start the agent and wait for it to be fully ready
265265
_ = agenttest.New(t, client.URL, agentToken)
@@ -304,7 +304,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
304304
t.Run("BackgroundCommandCapturesOutput", func(t *testing.T) {
305305
t.Parallel()
306306

307-
client, workspace, agentToken := setupWorkspaceForAgent(t)
307+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
308308

309309
// Start the agent and wait for it to be fully ready
310310
_ = agenttest.New(t, client.URL, agentToken)
@@ -345,7 +345,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
345345
t.Run("BackgroundVsNormalExecution", func(t *testing.T) {
346346
t.Parallel()
347347

348-
client, workspace, agentToken := setupWorkspaceForAgent(t)
348+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
349349

350350
// Start the agent and wait for it to be fully ready
351351
_ = agenttest.New(t, client.URL, agentToken)
@@ -391,7 +391,7 @@ func TestWorkspaceBashBackgroundIntegration(t *testing.T) {
391391
t.Run("BackgroundCommandContinuesAfterTimeout", func(t *testing.T) {
392392
t.Parallel()
393393

394-
client, workspace, agentToken := setupWorkspaceForAgent(t)
394+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
395395

396396
// Start the agent and wait for it to be fully ready
397397
_ = agenttest.New(t, client.URL, agentToken)

‎codersdk/toolsdk/toolsdk.go‎

Lines changed: 64 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"fmt"
99
"io"
1010
"runtime/debug"
11+
"strconv"
1112
"strings"
1213

1314
"github.com/google/uuid"
@@ -17,6 +18,7 @@ import (
1718

1819
"github.com/coder/coder/v2/buildinfo"
1920
"github.com/coder/coder/v2/cli/cliui"
21+
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
2022
"github.com/coder/coder/v2/codersdk"
2123
"github.com/coder/coder/v2/codersdk/workspacesdk"
2224
)
@@ -47,6 +49,7 @@ const (
4749
ToolNameWorkspaceWriteFile = "coder_workspace_write_file"
4850
ToolNameWorkspaceEditFile = "coder_workspace_edit_file"
4951
ToolNameWorkspaceEditFiles = "coder_workspace_edit_files"
52+
ToolNameWorkspacePortForward = "coder_workspace_port_forward"
5053
)
5154

5255
func NewDeps(client *codersdk.Client, opts ...func(*Deps)) (Deps, error) {
@@ -219,6 +222,7 @@ var All = []GenericTool{
219222
WorkspaceWriteFile.Generic(),
220223
WorkspaceEditFile.Generic(),
221224
WorkspaceEditFiles.Generic(),
225+
WorkspacePortForward.Generic(),
222226
}
223227

224228
type ReportTaskArgs struct {
@@ -1389,6 +1393,8 @@ type WorkspaceLSResponse struct {
13891393
Contents []WorkspaceLSFile `json:"contents"`
13901394
}
13911395

1396+
const workspaceDescription = "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used."
1397+
13921398
var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{
13931399
Tool: aisdk.Tool{
13941400
Name: ToolNameWorkspaceLS,
@@ -1397,7 +1403,7 @@ var WorkspaceLS = Tool[WorkspaceLSArgs, WorkspaceLSResponse]{
13971403
Properties: map[string]any{
13981404
"workspace": map[string]any{
13991405
"type": "string",
1400-
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
1406+
"description": workspaceDescription,
14011407
},
14021408
"path": map[string]any{
14031409
"type": "string",
@@ -1454,7 +1460,7 @@ var WorkspaceReadFile = Tool[WorkspaceReadFileArgs, WorkspaceReadFileResponse]{
14541460
Properties: map[string]any{
14551461
"workspace": map[string]any{
14561462
"type": "string",
1457-
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
1463+
"description": workspaceDescription,
14581464
},
14591465
"path": map[string]any{
14601466
"type": "string",
@@ -1519,7 +1525,7 @@ var WorkspaceWriteFile = Tool[WorkspaceWriteFileArgs, codersdk.Response]{
15191525
Properties: map[string]any{
15201526
"workspace": map[string]any{
15211527
"type": "string",
1522-
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
1528+
"description": workspaceDescription,
15231529
},
15241530
"path": map[string]any{
15251531
"type": "string",
@@ -1567,7 +1573,7 @@ var WorkspaceEditFile = Tool[WorkspaceEditFileArgs, codersdk.Response]{
15671573
Properties: map[string]any{
15681574
"workspace": map[string]any{
15691575
"type": "string",
1570-
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
1576+
"description": workspaceDescription,
15711577
},
15721578
"path": map[string]any{
15731579
"type": "string",
@@ -1634,7 +1640,7 @@ var WorkspaceEditFiles = Tool[WorkspaceEditFilesArgs, codersdk.Response]{
16341640
Properties: map[string]any{
16351641
"workspace": map[string]any{
16361642
"type": "string",
1637-
"description": "The workspace name in the format [owner/]workspace[.agent]. If an owner is not specified, the authenticated user is used.",
1643+
"description": workspaceDescription,
16381644
},
16391645
"files": map[string]any{
16401646
"type": "array",
@@ -1691,6 +1697,59 @@ var WorkspaceEditFiles = Tool[WorkspaceEditFilesArgs, codersdk.Response]{
16911697
},
16921698
}
16931699

1700+
type WorkspacePortForwardArgs struct {
1701+
Workspace string `json:"workspace"`
1702+
Port int `json:"port"`
1703+
}
1704+
1705+
type WorkspacePortForwardResponse struct {
1706+
URL string `json:"url"`
1707+
}
1708+
1709+
var WorkspacePortForward = Tool[WorkspacePortForwardArgs, WorkspacePortForwardResponse]{
1710+
Tool: aisdk.Tool{
1711+
Name: ToolNameWorkspacePortForward,
1712+
Description: `Fetch URLs that forward to the specified port.`,
1713+
Schema: aisdk.Schema{
1714+
Properties: map[string]any{
1715+
"workspace": map[string]any{
1716+
"type": "string",
1717+
"description": workspaceDescription,
1718+
},
1719+
"port": map[string]any{
1720+
"type": "number",
1721+
"description": "The port to forward.",
1722+
},
1723+
},
1724+
Required: []string{"workspace", "port"},
1725+
},
1726+
},
1727+
UserClientOptional: true,
1728+
Handler: func(ctx context.Context, deps Deps, args WorkspacePortForwardArgs) (WorkspacePortForwardResponse, error) {
1729+
workspaceName := NormalizeWorkspaceInput(args.Workspace)
1730+
workspace, workspaceAgent, err := findWorkspaceAndAgent(ctx, deps.coderClient, workspaceName)
1731+
if err != nil {
1732+
return WorkspacePortForwardResponse{}, xerrors.Errorf("failed to find workspace: %w", err)
1733+
}
1734+
res, err := deps.coderClient.AppHost(ctx)
1735+
if err != nil {
1736+
return WorkspacePortForwardResponse{}, xerrors.Errorf("failed to get app host: %w", err)
1737+
}
1738+
if res.Host == "" {
1739+
return WorkspacePortForwardResponse{}, xerrors.New("no app host for forwarding has been configured")
1740+
}
1741+
url := appurl.ApplicationURL{
1742+
AppSlugOrPort: strconv.Itoa(args.Port),
1743+
AgentName: workspaceAgent.Name,
1744+
WorkspaceName: workspace.Name,
1745+
Username: workspace.OwnerName,
1746+
}
1747+
return WorkspacePortForwardResponse{
1748+
URL: deps.coderClient.URL.Scheme + "://" + strings.Replace(res.Host, "*", url.String(), 1),
1749+
}, nil
1750+
},
1751+
}
1752+
16941753
// NormalizeWorkspaceInput converts workspace name input to standard format.
16951754
// Handles the following input formats:
16961755
// - workspace → workspace

‎codersdk/toolsdk/toolsdk_test.go‎

Lines changed: 69 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package toolsdk_test
33
import (
44
"context"
55
"encoding/json"
6+
"fmt"
67
"os"
78
"path/filepath"
89
"runtime"
@@ -35,10 +36,10 @@ import (
3536

3637
// setupWorkspaceForAgent creates a workspace setup exactly like main SSH tests
3738
// nolint:gocritic // This is in a test package and does not end up in the build
38-
func setupWorkspaceForAgent(t *testing.T) (*codersdk.Client, database.WorkspaceTable, string) {
39+
func setupWorkspaceForAgent(t *testing.T, opts *coderdtest.Options) (*codersdk.Client, database.WorkspaceTable, string) {
3940
t.Helper()
4041

41-
client, store := coderdtest.NewWithDatabase(t, nil)
42+
client, store := coderdtest.NewWithDatabase(t, opts)
4243
client.SetLogger(testutil.Logger(t).Named("client"))
4344
first := coderdtest.CreateFirstUser(t, client)
4445
userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
@@ -405,7 +406,7 @@ func TestTools(t *testing.T) {
405406
t.Skip("WorkspaceSSHExec is not supported on Windows")
406407
}
407408
// Setup workspace exactly like main SSH tests
408-
client, workspace, agentToken := setupWorkspaceForAgent(t)
409+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
409410

410411
// Start agent and wait for it to be ready (following main SSH test pattern)
411412
_ = agenttest.New(t, client.URL, agentToken)
@@ -457,7 +458,7 @@ func TestTools(t *testing.T) {
457458
t.Run("WorkspaceLS", func(t *testing.T) {
458459
t.Parallel()
459460

460-
client, workspace, agentToken := setupWorkspaceForAgent(t)
461+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
461462
fs := afero.NewMemMapFs()
462463
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
463464
opts.Filesystem = fs
@@ -503,7 +504,7 @@ func TestTools(t *testing.T) {
503504
t.Run("WorkspaceReadFile", func(t *testing.T) {
504505
t.Parallel()
505506

506-
client, workspace, agentToken := setupWorkspaceForAgent(t)
507+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
507508
fs := afero.NewMemMapFs()
508509
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
509510
opts.Filesystem = fs
@@ -606,7 +607,7 @@ func TestTools(t *testing.T) {
606607
t.Run("WorkspaceWriteFile", func(t *testing.T) {
607608
t.Parallel()
608609

609-
client, workspace, agentToken := setupWorkspaceForAgent(t)
610+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
610611
fs := afero.NewMemMapFs()
611612
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
612613
opts.Filesystem = fs
@@ -633,7 +634,7 @@ func TestTools(t *testing.T) {
633634
t.Run("WorkspaceEditFile", func(t *testing.T) {
634635
t.Parallel()
635636

636-
client, workspace, agentToken := setupWorkspaceForAgent(t)
637+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
637638
fs := afero.NewMemMapFs()
638639
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
639640
opts.Filesystem = fs
@@ -673,7 +674,7 @@ func TestTools(t *testing.T) {
673674
t.Run("WorkspaceEditFiles", func(t *testing.T) {
674675
t.Parallel()
675676

676-
client, workspace, agentToken := setupWorkspaceForAgent(t)
677+
client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
677678
fs := afero.NewMemMapFs()
678679
_ = agenttest.New(t, client.URL, agentToken, func(opts *agent.Options) {
679680
opts.Filesystem = fs
@@ -730,6 +731,66 @@ func TestTools(t *testing.T) {
730731
require.NoError(t, err)
731732
require.Equal(t, "bar2 bar2", string(b))
732733
})
734+
735+
t.Run("WorkspacePortForward", func(t *testing.T) {
736+
t.Parallel()
737+
738+
tests := []struct {
739+
name string
740+
workspace string
741+
host string
742+
port int
743+
expect string
744+
error string
745+
}{
746+
{
747+
name: "OK",
748+
workspace: "myuser/myworkspace",
749+
port: 1234,
750+
host: "*.test.coder.com",
751+
expect: "%s://1234--dev--myworkspace--myuser.test.coder.com:%s",
752+
},
753+
{
754+
name: "NonExistentWorkspace",
755+
workspace: "doesnotexist",
756+
port: 1234,
757+
host: "*.test.coder.com",
758+
error: "failed to find workspace",
759+
},
760+
{
761+
name: "NoAppHost",
762+
host: "",
763+
workspace: "myuser/myworkspace",
764+
port: 1234,
765+
error: "no app host",
766+
},
767+
}
768+
769+
for _, tt := range tests {
770+
t.Run(tt.name, func(t *testing.T) {
771+
t.Parallel()
772+
client, workspace, agentToken := setupWorkspaceForAgent(t, &coderdtest.Options{
773+
AppHostname: tt.host,
774+
})
775+
_ = agenttest.New(t, client.URL, agentToken)
776+
coderdtest.NewWorkspaceAgentWaiter(t, client, workspace.ID).Wait()
777+
tb, err := toolsdk.NewDeps(client)
778+
require.NoError(t, err)
779+
780+
res, err := testTool(t, toolsdk.WorkspacePortForward, tb, toolsdk.WorkspacePortForwardArgs{
781+
Workspace: tt.workspace,
782+
Port: tt.port,
783+
})
784+
if tt.error != "" {
785+
require.Error(t, err)
786+
require.ErrorContains(t, err, tt.error)
787+
} else {
788+
require.NoError(t, err)
789+
require.Equal(t, fmt.Sprintf(tt.expect, client.URL.Scheme, client.URL.Port()), res.URL)
790+
}
791+
})
792+
}
793+
})
733794
}
734795

735796
// TestedTools keeps track of which tools have been tested.

0 commit comments

Comments
 (0)