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
|
// 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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
Reference in New Issue
Block a user