Skip to content

Commit 5647b8b

Browse files
johnstcnjohnstcn
committed
simplify toolbox impl, rename to deps
1 parent def3fcb commit 5647b8b

File tree

3 files changed

+80
-113
lines changed

3 files changed

+80
-113
lines changed

‎cli/exp_mcp.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -400,21 +400,23 @@ func mcpServerHandler(inv *serpent.Invocation, client *codersdk.Client, instruct
400400
)
401401

402402
// Create a new context for the tools with all relevant information.
403-
tb := toolsdk.NewToolbox(client)
403+
tb := toolsdk.Deps{
404+
CoderClient: client,
405+
}
404406
// Get the workspace agent token from the environment.
405407
var hasAgentClient bool
406408
if agentToken, err := getAgentToken(fs); err == nil && agentToken != "" {
407409
hasAgentClient = true
408410
agentClient := agentsdk.New(client.URL)
409411
agentClient.SetSessionToken(agentToken)
410-
tb = tb.WithAgentClient(agentClient)
412+
tb.AgentClient = agentClient
411413
} else {
412414
cliui.Warnf(inv.Stderr, "CODER_AGENT_TOKEN is not set, task reporting will not be available")
413415
}
414416
if appStatusSlug == "" {
415417
cliui.Warnf(inv.Stderr, "CODER_MCP_APP_STATUS_SLUG is not set, task reporting will not be available.")
416418
} else {
417-
tb = tb.WithAppStatusSlug(appStatusSlug)
419+
tb.AppStatusSlug = appStatusSlug
418420
}
419421

420422
// Register tools based on the allowlist (if specified)
@@ -695,7 +697,7 @@ func getAgentToken(fs afero.Fs) (string, error) {
695697

696698
// mcpFromSDK adapts a toolsdk.Tool to go-mcp's server.ServerTool.
697699
// It assumes that the tool responds with a valid JSON object.
698-
func mcpFromSDK(sdkTool toolsdk.Tool[any, any], tb toolsdk.Toolbox) server.ServerTool {
700+
func mcpFromSDK(sdkTool toolsdk.Tool[any, any], tb toolsdk.Deps) server.ServerTool {
699701
// NOTE: some clients will silently refuse to use tools if there is an issue
700702
// with the tool's schema or configuration.
701703
if sdkTool.Schema.Properties == nil {

‎codersdk/toolsdk/toolsdk.go

Lines changed: 45 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -13,52 +13,15 @@ import (
1313
"github.com/coder/coder/v2/codersdk/agentsdk"
1414
)
1515

16-
// Toolbox provides access to tool dependencies.
17-
type Toolbox interface {
18-
CoderClient() *codersdk.Client
19-
AgentClient() (*agentsdk.Client, bool)
20-
AppStatusSlug() (string, bool)
21-
22-
WithAgentClient(*agentsdk.Client) Toolbox
23-
WithAppStatusSlug(string) Toolbox
24-
}
25-
26-
// toolbox is the concrete implementation of Toolbox.
27-
type toolbox struct {
28-
coderClient *codersdk.Client
29-
agentClient *agentsdk.Client
30-
appStatusSlug string
31-
}
32-
33-
// NewToolbox constructs a Toolbox with a required CoderClient.
34-
func NewToolbox(coder *codersdk.Client) Toolbox {
35-
return &toolbox{coderClient: coder}
36-
}
37-
38-
func (tb *toolbox) CoderClient() *codersdk.Client {
39-
return tb.coderClient
40-
}
41-
42-
func (tb *toolbox) AgentClient() (*agentsdk.Client, bool) {
43-
return tb.agentClient, tb.agentClient != nil
44-
}
45-
46-
func (tb *toolbox) AppStatusSlug() (string, bool) {
47-
return tb.appStatusSlug, tb.appStatusSlug != ""
48-
}
49-
50-
func (tb *toolbox) WithAgentClient(agent *agentsdk.Client) Toolbox {
51-
tb.agentClient = agent
52-
return tb
53-
}
54-
55-
func (tb *toolbox) WithAppStatusSlug(slug string) Toolbox {
56-
tb.appStatusSlug = slug
57-
return tb
16+
// Deps provides access to tool dependencies.
17+
type Deps struct {
18+
CoderClient *codersdk.Client
19+
AgentClient *agentsdk.Client
20+
AppStatusSlug string
5821
}
5922

6023
// HandlerFunc is a function that handles a tool call.
61-
type HandlerFunc[Arg, Ret any] func(tb Toolbox, args Arg) (Ret, error)
24+
type HandlerFunc[Arg, Ret any] func(tb Deps, args Arg) (Ret, error)
6225

6326
type Tool[Arg, Ret any] struct {
6427
aisdk.Tool
@@ -69,7 +32,7 @@ type Tool[Arg, Ret any] struct {
6932
func (t Tool[Arg, Ret]) Generic() Tool[any, any] {
7033
return Tool[any, any]{
7134
Tool: t.Tool,
72-
Handler: func(tb Toolbox, args any) (any, error) {
35+
Handler: func(tb Deps, args any) (any, error) {
7336
typedArg, ok := args.(Arg)
7437
if !ok {
7538
return nil, xerrors.Errorf("developer error: invalid argument type for tool %s", t.Tool.Name)
@@ -152,7 +115,7 @@ type UploadTarFileArgs struct {
152115

153116
// WithRecover wraps a HandlerFunc to recover from panics and return an error.
154117
func WithRecover[Arg, Ret any](h HandlerFunc[Arg, Ret]) HandlerFunc[Arg, Ret] {
155-
return func(tb Toolbox, args Arg) (ret Ret, err error) {
118+
return func(tb Deps, args Arg) (ret Ret, err error) {
156119
defer func() {
157120
if r := recover(); r != nil {
158121
err = xerrors.Errorf("tool handler panic: %v", r)
@@ -220,17 +183,15 @@ var (
220183
Required: []string{"summary", "link", "state"},
221184
},
222185
},
223-
Handler: func(tb Toolbox, args ReportTaskArgs) (string, error) {
224-
agentClient, ok := tb.AgentClient()
225-
if !ok {
186+
Handler: func(tb Deps, args ReportTaskArgs) (string, error) {
187+
if tb.AgentClient == nil {
226188
return "", xerrors.New("tool unavailable as CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE not set")
227189
}
228-
appStatusSlug, ok := tb.AppStatusSlug()
229-
if !ok {
190+
if tb.AppStatusSlug == "" {
230191
return "", xerrors.New("workspace app status slug not found in toolbox")
231192
}
232-
if err := agentClient.PatchAppStatus(context.TODO(), agentsdk.PatchAppStatus{
233-
AppSlug: appStatusSlug,
193+
if err := tb.AgentClient.PatchAppStatus(context.TODO(), agentsdk.PatchAppStatus{
194+
AppSlug: tb.AppStatusSlug,
234195
Message: args.Summary,
235196
URI: args.Link,
236197
State: codersdk.WorkspaceAppStatusState(args.State),
@@ -256,12 +217,12 @@ This returns more data than list_workspaces to reduce token usage.`,
256217
Required: []string{"workspace_id"},
257218
},
258219
},
259-
Handler: func(tb Toolbox, args GetWorkspaceArgs) (codersdk.Workspace, error) {
220+
Handler: func(tb Deps, args GetWorkspaceArgs) (codersdk.Workspace, error) {
260221
wsID, err := uuid.Parse(args.WorkspaceID)
261222
if err != nil {
262223
return codersdk.Workspace{}, xerrors.New("workspace_id must be a valid UUID")
263224
}
264-
return tb.CoderClient().Workspace(context.TODO(), wsID)
225+
return tb.CoderClient.Workspace(context.TODO(), wsID)
265226
},
266227
}
267228

@@ -296,7 +257,7 @@ is provisioned correctly and the agent can connect to the control plane.
296257
Required: []string{"user", "template_version_id", "name", "rich_parameters"},
297258
},
298259
},
299-
Handler: func(tb Toolbox, args CreateWorkspaceArgs) (codersdk.Workspace, error) {
260+
Handler: func(tb Deps, args CreateWorkspaceArgs) (codersdk.Workspace, error) {
300261
tvID, err := uuid.Parse(args.TemplateVersionID)
301262
if err != nil {
302263
return codersdk.Workspace{}, xerrors.New("template_version_id must be a valid UUID")
@@ -311,7 +272,7 @@ is provisioned correctly and the agent can connect to the control plane.
311272
Value: v,
312273
})
313274
}
314-
workspace, err := tb.CoderClient().CreateUserWorkspace(context.TODO(), args.User, codersdk.CreateWorkspaceRequest{
275+
workspace, err := tb.CoderClient.CreateUserWorkspace(context.TODO(), args.User, codersdk.CreateWorkspaceRequest{
315276
TemplateVersionID: tvID,
316277
Name: args.Name,
317278
RichParameterValues: buildParams,
@@ -336,12 +297,12 @@ is provisioned correctly and the agent can connect to the control plane.
336297
},
337298
},
338299
},
339-
Handler: func(tb Toolbox, args ListWorkspacesArgs) ([]MinimalWorkspace, error) {
300+
Handler: func(tb Deps, args ListWorkspacesArgs) ([]MinimalWorkspace, error) {
340301
owner := args.Owner
341302
if owner == "" {
342303
owner = codersdk.Me
343304
}
344-
workspaces, err := tb.CoderClient().Workspaces(context.TODO(), codersdk.WorkspaceFilter{
305+
workspaces, err := tb.CoderClient.Workspaces(context.TODO(), codersdk.WorkspaceFilter{
345306
Owner: owner,
346307
})
347308
if err != nil {
@@ -373,8 +334,8 @@ is provisioned correctly and the agent can connect to the control plane.
373334
Required: []string{},
374335
},
375336
},
376-
Handler: func(tb Toolbox, _ NoArgs) ([]MinimalTemplate, error) {
377-
templates, err := tb.CoderClient().Templates(context.TODO(), codersdk.TemplateFilter{})
337+
Handler: func(tb Deps, _ NoArgs) ([]MinimalTemplate, error) {
338+
templates, err := tb.CoderClient.Templates(context.TODO(), codersdk.TemplateFilter{})
378339
if err != nil {
379340
return nil, err
380341
}
@@ -406,12 +367,12 @@ is provisioned correctly and the agent can connect to the control plane.
406367
Required: []string{"template_version_id"},
407368
},
408369
},
409-
Handler: func(tb Toolbox, args ListTemplateVersionParametersArgs) ([]codersdk.TemplateVersionParameter, error) {
370+
Handler: func(tb Deps, args ListTemplateVersionParametersArgs) ([]codersdk.TemplateVersionParameter, error) {
410371
templateVersionID, err := uuid.Parse(args.TemplateVersionID)
411372
if err != nil {
412373
return nil, xerrors.Errorf("template_version_id must be a valid UUID: %w", err)
413374
}
414-
parameters, err := tb.CoderClient().TemplateVersionRichParameters(context.TODO(), templateVersionID)
375+
parameters, err := tb.CoderClient.TemplateVersionRichParameters(context.TODO(), templateVersionID)
415376
if err != nil {
416377
return nil, err
417378
}
@@ -428,8 +389,8 @@ is provisioned correctly and the agent can connect to the control plane.
428389
Required: []string{},
429390
},
430391
},
431-
Handler: func(tb Toolbox, _ NoArgs) (codersdk.User, error) {
432-
return tb.CoderClient().User(context.TODO(), "me")
392+
Handler: func(tb Deps, _ NoArgs) (codersdk.User, error) {
393+
return tb.CoderClient.User(context.TODO(), "me")
433394
},
434395
}
435396

@@ -455,7 +416,7 @@ is provisioned correctly and the agent can connect to the control plane.
455416
Required: []string{"workspace_id", "transition"},
456417
},
457418
},
458-
Handler: func(tb Toolbox, args CreateWorkspaceBuildArgs) (codersdk.WorkspaceBuild, error) {
419+
Handler: func(tb Deps, args CreateWorkspaceBuildArgs) (codersdk.WorkspaceBuild, error) {
459420
workspaceID, err := uuid.Parse(args.WorkspaceID)
460421
if err != nil {
461422
return codersdk.WorkspaceBuild{}, xerrors.Errorf("workspace_id must be a valid UUID: %w", err)
@@ -474,7 +435,7 @@ is provisioned correctly and the agent can connect to the control plane.
474435
if templateVersionID != uuid.Nil {
475436
cbr.TemplateVersionID = templateVersionID
476437
}
477-
return tb.CoderClient().CreateWorkspaceBuild(context.TODO(), workspaceID, cbr)
438+
return tb.CoderClient.CreateWorkspaceBuild(context.TODO(), workspaceID, cbr)
478439
},
479440
}
480441

@@ -936,8 +897,8 @@ The file_id provided is a reference to a tar file you have uploaded containing t
936897
Required: []string{"file_id"},
937898
},
938899
},
939-
Handler: func(tb Toolbox, args CreateTemplateVersionArgs) (codersdk.TemplateVersion, error) {
940-
me, err := tb.CoderClient().User(context.TODO(), "me")
900+
Handler: func(tb Deps, args CreateTemplateVersionArgs) (codersdk.TemplateVersion, error) {
901+
me, err := tb.CoderClient.User(context.TODO(), "me")
941902
if err != nil {
942903
return codersdk.TemplateVersion{}, err
943904
}
@@ -949,7 +910,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
949910
if err != nil {
950911
return codersdk.TemplateVersion{}, xerrors.Errorf("template_id must be a valid UUID: %w", err)
951912
}
952-
templateVersion, err := tb.CoderClient().CreateTemplateVersion(context.TODO(), me.OrganizationIDs[0], codersdk.CreateTemplateVersionRequest{
913+
templateVersion, err := tb.CoderClient.CreateTemplateVersion(context.TODO(), me.OrganizationIDs[0], codersdk.CreateTemplateVersionRequest{
953914
Message: "Created by AI",
954915
StorageMethod: codersdk.ProvisionerStorageMethodFile,
955916
FileID: fileID,
@@ -978,12 +939,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
978939
Required: []string{"workspace_agent_id"},
979940
},
980941
},
981-
Handler: func(tb Toolbox, args GetWorkspaceAgentLogsArgs) ([]string, error) {
942+
Handler: func(tb Deps, args GetWorkspaceAgentLogsArgs) ([]string, error) {
982943
workspaceAgentID, err := uuid.Parse(args.WorkspaceAgentID)
983944
if err != nil {
984945
return nil, xerrors.Errorf("workspace_agent_id must be a valid UUID: %w", err)
985946
}
986-
logs, closer, err := tb.CoderClient().WorkspaceAgentLogsAfter(context.TODO(), workspaceAgentID, 0, false)
947+
logs, closer, err := tb.CoderClient.WorkspaceAgentLogsAfter(context.TODO(), workspaceAgentID, 0, false)
987948
if err != nil {
988949
return nil, err
989950
}
@@ -1013,12 +974,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
1013974
Required: []string{"workspace_build_id"},
1014975
},
1015976
},
1016-
Handler: func(tb Toolbox, args GetWorkspaceBuildLogsArgs) ([]string, error) {
977+
Handler: func(tb Deps, args GetWorkspaceBuildLogsArgs) ([]string, error) {
1017978
workspaceBuildID, err := uuid.Parse(args.WorkspaceBuildID)
1018979
if err != nil {
1019980
return nil, xerrors.Errorf("workspace_build_id must be a valid UUID: %w", err)
1020981
}
1021-
logs, closer, err := tb.CoderClient().WorkspaceBuildLogsAfter(context.TODO(), workspaceBuildID, 0)
982+
logs, closer, err := tb.CoderClient.WorkspaceBuildLogsAfter(context.TODO(), workspaceBuildID, 0)
1022983
if err != nil {
1023984
return nil, err
1024985
}
@@ -1044,13 +1005,13 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10441005
Required: []string{"template_version_id"},
10451006
},
10461007
},
1047-
Handler: func(tb Toolbox, args GetTemplateVersionLogsArgs) ([]string, error) {
1008+
Handler: func(tb Deps, args GetTemplateVersionLogsArgs) ([]string, error) {
10481009
templateVersionID, err := uuid.Parse(args.TemplateVersionID)
10491010
if err != nil {
10501011
return nil, xerrors.Errorf("template_version_id must be a valid UUID: %w", err)
10511012
}
10521013

1053-
logs, closer, err := tb.CoderClient().TemplateVersionLogsAfter(context.TODO(), templateVersionID, 0)
1014+
logs, closer, err := tb.CoderClient.TemplateVersionLogsAfter(context.TODO(), templateVersionID, 0)
10541015
if err != nil {
10551016
return nil, err
10561017
}
@@ -1079,7 +1040,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10791040
Required: []string{"template_id", "template_version_id"},
10801041
},
10811042
},
1082-
Handler: func(tb Toolbox, args UpdateTemplateActiveVersionArgs) (string, error) {
1043+
Handler: func(tb Deps, args UpdateTemplateActiveVersionArgs) (string, error) {
10831044
templateID, err := uuid.Parse(args.TemplateID)
10841045
if err != nil {
10851046
return "", xerrors.Errorf("template_id must be a valid UUID: %w", err)
@@ -1088,7 +1049,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10881049
if err != nil {
10891050
return "", xerrors.Errorf("template_version_id must be a valid UUID: %w", err)
10901051
}
1091-
err = tb.CoderClient().UpdateActiveTemplateVersion(context.TODO(), templateID, codersdk.UpdateActiveTemplateVersion{
1052+
err = tb.CoderClient.UpdateActiveTemplateVersion(context.TODO(), templateID, codersdk.UpdateActiveTemplateVersion{
10921053
ID: templateVersionID,
10931054
})
10941055
if err != nil {
@@ -1112,7 +1073,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
11121073
Required: []string{"mime_type", "files"},
11131074
},
11141075
},
1115-
Handler: func(tb Toolbox, args UploadTarFileArgs) (codersdk.UploadResponse, error) {
1076+
Handler: func(tb Deps, args UploadTarFileArgs) (codersdk.UploadResponse, error) {
11161077
pipeReader, pipeWriter := io.Pipe()
11171078
go func() {
11181079
defer pipeWriter.Close()
@@ -1137,7 +1098,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
11371098
}
11381099
}()
11391100

1140-
resp, err := tb.CoderClient().Upload(context.TODO(), codersdk.ContentTypeTar, pipeReader)
1101+
resp, err := tb.CoderClient.Upload(context.TODO(), codersdk.ContentTypeTar, pipeReader)
11411102
if err != nil {
11421103
return codersdk.UploadResponse{}, err
11431104
}
@@ -1172,16 +1133,16 @@ The file_id provided is a reference to a tar file you have uploaded containing t
11721133
Required: []string{"name", "display_name", "description", "version_id"},
11731134
},
11741135
},
1175-
Handler: func(tb Toolbox, args CreateTemplateArgs) (codersdk.Template, error) {
1176-
me, err := tb.CoderClient().User(context.TODO(), "me")
1136+
Handler: func(tb Deps, args CreateTemplateArgs) (codersdk.Template, error) {
1137+
me, err := tb.CoderClient.User(context.TODO(), "me")
11771138
if err != nil {
11781139
return codersdk.Template{}, err
11791140
}
11801141
versionID, err := uuid.Parse(args.VersionID)
11811142
if err != nil {
11821143
return codersdk.Template{}, xerrors.Errorf("version_id must be a valid UUID: %w", err)
11831144
}
1184-
template, err := tb.CoderClient().CreateTemplate(context.TODO(), me.OrganizationIDs[0], codersdk.CreateTemplateRequest{
1145+
template, err := tb.CoderClient.CreateTemplate(context.TODO(), me.OrganizationIDs[0], codersdk.CreateTemplateRequest{
11851146
Name: args.Name,
11861147
DisplayName: args.DisplayName,
11871148
Description: args.Description,
@@ -1206,12 +1167,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
12061167
},
12071168
},
12081169
},
1209-
Handler: func(tb Toolbox, args DeleteTemplateArgs) (string, error) {
1170+
Handler: func(tb Deps, args DeleteTemplateArgs) (string, error) {
12101171
templateID, err := uuid.Parse(args.TemplateID)
12111172
if err != nil {
12121173
return "", xerrors.Errorf("template_id must be a valid UUID: %w", err)
12131174
}
1214-
err = tb.CoderClient().DeleteTemplate(context.TODO(), templateID)
1175+
err = tb.CoderClient.DeleteTemplate(context.TODO(), templateID)
12151176
if err != nil {
12161177
return "", err
12171178
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy