Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 426f81a

Browse files
johnstcnjohnstcn
committed
simplify toolbox impl, rename to deps
1 parent 720f981 commit 426f81a

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)