Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit ea1f503

Browse files
ibuildthecloudibuildthecloud
committedApr 19, 2025·
chore: add MCP support
1 parent 7e668d5 commit ea1f503

File tree

13 files changed

+784
-20
lines changed

13 files changed

+784
-20
lines changed
 

‎go.mod

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,11 @@ require (
1818
github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d
1919
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb
2020
github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61
21-
github.com/gptscript-ai/tui v0.0.0-20250204145344-33cd15de4cee
21+
github.com/gptscript-ai/tui v0.0.0-20250419050840-5e79e16786c9
2222
github.com/hexops/autogold/v2 v2.2.1
2323
github.com/hexops/valast v1.4.4
2424
github.com/jaytaylor/html2text v0.0.0-20230321000545-74c2419ad056
25+
github.com/mark3labs/mcp-go v0.21.1
2526
github.com/mholt/archives v0.1.0
2627
github.com/pkoukk/tiktoken-go v0.1.7
2728
github.com/pkoukk/tiktoken-go-loader v0.0.2-0.20240522064338-c17e8bc0f699
@@ -122,6 +123,7 @@ require (
122123
github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect
123124
github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
124125
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
126+
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
125127
github.com/yuin/goldmark v1.5.4 // indirect
126128
github.com/yuin/goldmark-emoji v1.0.2 // indirect
127129
go4.org v0.0.0-20230225012048-214862532bf5 // indirect

‎go.sum

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -203,8 +203,8 @@ github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb h1:ky2J2CzBOskC7J
203203
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
204204
github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61 h1:QxLjsLOYlsVLPwuRkP0Q8EcAoZT1s8vU2ZBSX0+R6CI=
205205
github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61/go.mod h1:/FVuLwhz+sIfsWUgUHWKi32qT0i6+IXlUlzs70KKt/Q=
206-
github.com/gptscript-ai/tui v0.0.0-20250204145344-33cd15de4cee h1:70PHW6Xw70yNNZ5aX936XqcMLwNmfMZpCV3FCOGKpxE=
207-
github.com/gptscript-ai/tui v0.0.0-20250204145344-33cd15de4cee/go.mod h1:iwHxuueg2paOak7zIg0ESBWx7A0wIHGopAratbgaPNY=
206+
github.com/gptscript-ai/tui v0.0.0-20250419050840-5e79e16786c9 h1:wQC8sKyeGA50WnCEG+Jo5FNRIkuX3HX8d3ubyWCCoI8=
207+
github.com/gptscript-ai/tui v0.0.0-20250419050840-5e79e16786c9/go.mod h1:iwHxuueg2paOak7zIg0ESBWx7A0wIHGopAratbgaPNY=
208208
github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
209209
github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I=
210210
github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
@@ -270,6 +270,8 @@ github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69
270270
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
271271
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
272272
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
273+
github.com/mark3labs/mcp-go v0.21.1 h1:7Ek6KPIIbMhEYHRiRIg6K6UAgNZCJaHKQp926MNr6V0=
274+
github.com/mark3labs/mcp-go v0.21.1/go.mod h1:KmJndYv7GIgcPVwEKJjNcbhVQ+hJGJhrCCB/9xITzpE=
273275
github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
274276
github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
275277
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
@@ -406,6 +408,8 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
406408
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
407409
github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
408410
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
411+
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
412+
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
409413
github.com/yuin/goldmark v1.3.7/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
410414
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
411415
github.com/yuin/goldmark v1.5.4 h1:2uY/xC0roWy8IBEGLgB1ywIoEJFGmRrX21YQcvGZzjU=

‎pkg/cli/gptscript.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ func (r *GPTScript) listTools(ctx context.Context, gptScript *gptscript.GPTScrip
215215
// Don't print instructions
216216
tool.Instructions = ""
217217

218-
lines = append(lines, tool.String())
218+
lines = append(lines, tool.Print())
219219
}
220220
fmt.Println(strings.Join(lines, "\n---\n"))
221221
return nil

‎pkg/engine/engine.go

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"sync"
1212

1313
"github.com/gptscript-ai/gptscript/pkg/counter"
14+
"github.com/gptscript-ai/gptscript/pkg/mcp"
1415
"github.com/gptscript-ai/gptscript/pkg/types"
1516
"github.com/gptscript-ai/gptscript/pkg/version"
1617
)
@@ -41,6 +42,11 @@ type Engine struct {
4142
RuntimeManager RuntimeManager
4243
Env []string
4344
Progress chan<- types.CompletionStatus
45+
MCPRunner MCPRunner
46+
}
47+
48+
type MCPRunner interface {
49+
Run(ctx context.Context, progress chan<- types.CompletionStatus, tool types.Tool, input string) (string, error)
4450
}
4551

4652
type State struct {
@@ -307,6 +313,21 @@ func populateMessageParams(ctx Context, completion *types.CompletionRequest, too
307313
return nil
308314
}
309315

316+
func (e *Engine) runMCPInvoke(ctx Context, tool types.Tool, input string) (*Return, error) {
317+
runner := e.MCPRunner
318+
if runner == nil {
319+
runner = mcp.DefaultRunner
320+
}
321+
output, err := runner.Run(ctx.Ctx, e.Progress, tool, input)
322+
if err != nil {
323+
return nil, fmt.Errorf("failed to run MCP invoke: %w", err)
324+
}
325+
326+
return &Return{
327+
Result: &output,
328+
}, nil
329+
}
330+
310331
func (e *Engine) runCommandTools(ctx Context, tool types.Tool, input string) (*Return, error) {
311332
if tool.IsHTTP() {
312333
return e.runHTTP(ctx, tool, input)
@@ -342,6 +363,10 @@ func (e *Engine) Start(ctx Context, input string) (ret *Return, err error) {
342363
}
343364
}()
344365

366+
if tool.IsMCPInvoke() {
367+
return e.runMCPInvoke(ctx, tool, input)
368+
}
369+
345370
if tool.IsCommand() {
346371
return e.runCommandTools(ctx, tool, input)
347372
}
@@ -378,6 +403,7 @@ func addUpdateSystem(ctx Context, tool types.Tool, msgs []types.CompletionMessag
378403
instructions = append(instructions, context.Content)
379404
}
380405

406+
tool.Instructions = strings.TrimPrefix(tool.Instructions, types.PromptPrefix)
381407
if tool.Instructions != "" {
382408
instructions = append(instructions, tool.Instructions)
383409
}

‎pkg/loader/loader.go

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"github.com/gptscript-ai/gptscript/pkg/builtin"
2121
"github.com/gptscript-ai/gptscript/pkg/cache"
2222
"github.com/gptscript-ai/gptscript/pkg/hash"
23+
"github.com/gptscript-ai/gptscript/pkg/mcp"
2324
"github.com/gptscript-ai/gptscript/pkg/openapi"
2425
"github.com/gptscript-ai/gptscript/pkg/parser"
2526
"github.com/gptscript-ai/gptscript/pkg/system"
@@ -155,7 +156,23 @@ func loadOpenAPI(prg *types.Program, data []byte) *openapi3.T {
155156
return openAPIDocument
156157
}
157158

158-
func readTool(ctx context.Context, cache *cache.Client, prg *types.Program, base *source, targetToolName, defaultModel string) ([]types.Tool, error) {
159+
func processMCP(ctx context.Context, tool []types.Tool, mcpLoader MCPLoader) (result []types.Tool, _ error) {
160+
for _, t := range tool {
161+
if t.IsMCP() {
162+
mcpTools, err := mcpLoader.Load(ctx, t)
163+
if err != nil {
164+
return nil, fmt.Errorf("error loading MCP tools: %w", err)
165+
}
166+
result = append(result, mcpTools...)
167+
} else {
168+
result = append(result, t)
169+
}
170+
}
171+
172+
return result, nil
173+
}
174+
175+
func readTool(ctx context.Context, cache *cache.Client, mcp MCPLoader, prg *types.Program, base *source, targetToolName, defaultModel string) ([]types.Tool, error) {
159176
data := base.Content
160177

161178
var (
@@ -212,6 +229,11 @@ func readTool(ctx context.Context, cache *cache.Client, prg *types.Program, base
212229
return nil, fmt.Errorf("no tools found in %s", base)
213230
}
214231

232+
tools, err := processMCP(ctx, tools, mcp)
233+
if err != nil {
234+
return nil, err
235+
}
236+
215237
var (
216238
localTools = types.ToolSet{}
217239
targetTools []types.Tool
@@ -279,17 +301,17 @@ func readTool(ctx context.Context, cache *cache.Client, prg *types.Program, base
279301
localTools[strings.ToLower(tool.Name)] = tool
280302
}
281303

282-
return linkAll(ctx, cache, prg, base, targetTools, localTools, defaultModel)
304+
return linkAll(ctx, cache, mcp, prg, base, targetTools, localTools, defaultModel)
283305
}
284306

285-
func linkAll(ctx context.Context, cache *cache.Client, prg *types.Program, base *source, tools []types.Tool, localTools types.ToolSet, defaultModel string) (result []types.Tool, _ error) {
307+
func linkAll(ctx context.Context, cache *cache.Client, mcp MCPLoader, prg *types.Program, base *source, tools []types.Tool, localTools types.ToolSet, defaultModel string) (result []types.Tool, _ error) {
286308
localToolsMapping := make(map[string]string, len(tools))
287309
for _, localTool := range localTools {
288310
localToolsMapping[strings.ToLower(localTool.Name)] = localTool.ID
289311
}
290312

291313
for _, tool := range tools {
292-
tool, err := link(ctx, cache, prg, base, tool, localTools, localToolsMapping, defaultModel)
314+
tool, err := link(ctx, cache, mcp, prg, base, tool, localTools, localToolsMapping, defaultModel)
293315
if err != nil {
294316
return nil, err
295317
}
@@ -298,7 +320,7 @@ func linkAll(ctx context.Context, cache *cache.Client, prg *types.Program, base
298320
return
299321
}
300322

301-
func link(ctx context.Context, cache *cache.Client, prg *types.Program, base *source, tool types.Tool, localTools types.ToolSet, localToolsMapping map[string]string, defaultModel string) (types.Tool, error) {
323+
func link(ctx context.Context, cache *cache.Client, mcp MCPLoader, prg *types.Program, base *source, tool types.Tool, localTools types.ToolSet, localToolsMapping map[string]string, defaultModel string) (types.Tool, error) {
302324
if existing, ok := prg.ToolSet[tool.ID]; ok {
303325
return existing, nil
304326
}
@@ -323,7 +345,7 @@ func link(ctx context.Context, cache *cache.Client, prg *types.Program, base *so
323345
linkedTool = existing
324346
} else {
325347
var err error
326-
linkedTool, err = link(ctx, cache, prg, base, localTool, localTools, localToolsMapping, defaultModel)
348+
linkedTool, err = link(ctx, cache, mcp, prg, base, localTool, localTools, localToolsMapping, defaultModel)
327349
if err != nil {
328350
return types.Tool{}, fmt.Errorf("failed linking %s at %s: %w", targetToolName, base, err)
329351
}
@@ -333,7 +355,7 @@ func link(ctx context.Context, cache *cache.Client, prg *types.Program, base *so
333355
toolNames[targetToolName] = struct{}{}
334356
} else {
335357
toolName, subTool := types.SplitToolRef(targetToolName)
336-
resolvedTools, err := resolve(ctx, cache, prg, base, toolName, subTool, defaultModel)
358+
resolvedTools, err := resolve(ctx, cache, mcp, prg, base, toolName, subTool, defaultModel)
337359
if err != nil {
338360
return types.Tool{}, fmt.Errorf("failed resolving %s from %s: %w", targetToolName, base, err)
339361
}
@@ -373,7 +395,7 @@ func ProgramFromSource(ctx context.Context, content, subToolName string, opts ..
373395
prg := types.Program{
374396
ToolSet: types.ToolSet{},
375397
}
376-
tools, err := readTool(ctx, opt.Cache, &prg, &source{
398+
tools, err := readTool(ctx, opt.Cache, opt.MCPLoader, &prg, &source{
377399
Content: []byte(content),
378400
Path: locationPath,
379401
Name: locationName,
@@ -390,13 +412,19 @@ type Options struct {
390412
Cache *cache.Client
391413
Location string
392414
DefaultModel string
415+
MCPLoader MCPLoader
416+
}
417+
418+
type MCPLoader interface {
419+
Load(ctx context.Context, tool types.Tool) ([]types.Tool, error)
393420
}
394421

395422
func complete(opts ...Options) (result Options) {
396423
for _, opt := range opts {
397424
result.Cache = types.FirstSet(opt.Cache, result.Cache)
398425
result.Location = types.FirstSet(opt.Location, result.Location)
399426
result.DefaultModel = types.FirstSet(opt.DefaultModel, result.DefaultModel)
427+
result.MCPLoader = types.FirstSet(opt.MCPLoader, result.MCPLoader)
400428
}
401429

402430
if result.Location == "" {
@@ -407,6 +435,10 @@ func complete(opts ...Options) (result Options) {
407435
result.DefaultModel = builtin.GetDefaultModel()
408436
}
409437

438+
if result.MCPLoader == nil {
439+
result.MCPLoader = mcp.DefaultLoader
440+
}
441+
410442
return
411443
}
412444

@@ -430,15 +462,15 @@ func Program(ctx context.Context, name, subToolName string, opts ...Options) (ty
430462
Name: name,
431463
ToolSet: types.ToolSet{},
432464
}
433-
tools, err := resolve(ctx, opt.Cache, &prg, &source{}, name, subToolName, opt.DefaultModel)
465+
tools, err := resolve(ctx, opt.Cache, opt.MCPLoader, &prg, &source{}, name, subToolName, opt.DefaultModel)
434466
if err != nil {
435467
return types.Program{}, err
436468
}
437469
prg.EntryToolID = tools[0].ID
438470
return prg, nil
439471
}
440472

441-
func resolve(ctx context.Context, cache *cache.Client, prg *types.Program, base *source, name, subTool, defaultModel string) ([]types.Tool, error) {
473+
func resolve(ctx context.Context, cache *cache.Client, mcp MCPLoader, prg *types.Program, base *source, name, subTool, defaultModel string) ([]types.Tool, error) {
442474
if subTool == "" {
443475
t, ok := builtin.DefaultModel(name, defaultModel)
444476
if ok {
@@ -452,7 +484,7 @@ func resolve(ctx context.Context, cache *cache.Client, prg *types.Program, base
452484
return nil, err
453485
}
454486

455-
result, err := readTool(ctx, cache, prg, s, subTool, defaultModel)
487+
result, err := readTool(ctx, cache, mcp, prg, s, subTool, defaultModel)
456488
if err != nil {
457489
return nil, err
458490
}

‎pkg/mcp/loader.go

Lines changed: 264 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,264 @@
1+
package mcp
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"maps"
8+
"slices"
9+
"strings"
10+
"sync"
11+
12+
"github.com/getkin/kin-openapi/openapi3"
13+
"github.com/gptscript-ai/gptscript/pkg/hash"
14+
"github.com/gptscript-ai/gptscript/pkg/types"
15+
"github.com/gptscript-ai/gptscript/pkg/version"
16+
"github.com/mark3labs/mcp-go/client"
17+
"github.com/mark3labs/mcp-go/mcp"
18+
)
19+
20+
var (
21+
DefaultLoader = &Local{}
22+
DefaultRunner = DefaultLoader
23+
)
24+
25+
type Local struct {
26+
nextID int64
27+
lock sync.Mutex
28+
sessions map[string]*Session
29+
}
30+
31+
type Session struct {
32+
ID string
33+
InitResult *mcp.InitializeResult
34+
Client client.MCPClient
35+
Config ServerConfig
36+
}
37+
38+
type Config struct {
39+
MCPServers map[string]ServerConfig `json:"mcpServers"`
40+
}
41+
42+
type ServerConfig struct {
43+
DisableInstruction bool `json:"disableInstruction"`
44+
Command string `json:"command"`
45+
Args []string `json:"args"`
46+
Env map[string]string `json:"env"`
47+
Server string `json:"server"`
48+
URL string `json:"url"`
49+
BaseURL string `json:"baseURL,omitempty"`
50+
Headers map[string]string `json:"headers"`
51+
}
52+
53+
func (s *ServerConfig) GetBaseURL() string {
54+
if s.BaseURL != "" {
55+
return s.BaseURL
56+
}
57+
if s.Server != "" {
58+
return s.Server
59+
}
60+
return s.URL
61+
}
62+
63+
func (l *Local) Load(ctx context.Context, tool types.Tool) (result []types.Tool, _ error) {
64+
if !tool.IsMCP() {
65+
return []types.Tool{tool}, nil
66+
}
67+
68+
_, configData, _ := strings.Cut(tool.Instructions, "\n")
69+
var servers Config
70+
71+
if err := json.Unmarshal([]byte(strings.TrimSpace(configData)), &servers); err != nil {
72+
return nil, fmt.Errorf("failed to parse MCP configuration: %w\n%s", err, configData)
73+
}
74+
75+
if len(servers.MCPServers) == 0 {
76+
// Try to load just one server
77+
var server ServerConfig
78+
if err := json.Unmarshal([]byte(strings.TrimSpace(configData)), &server); err != nil {
79+
return nil, fmt.Errorf("failed to parse single MCP server configuration: %w\n%s", err, configData)
80+
}
81+
if server.Command == "" && server.URL == "" && server.Server == "" {
82+
return nil, fmt.Errorf("no MCP server configuration found in tool instructions: %s", configData)
83+
}
84+
servers.MCPServers = map[string]ServerConfig{
85+
"default": server,
86+
}
87+
}
88+
89+
if len(servers.MCPServers) > 1 {
90+
return nil, fmt.Errorf("only a single MCP server definition is support")
91+
}
92+
93+
for _, server := range slices.Sorted(maps.Keys(servers.MCPServers)) {
94+
session, err := l.loadSession(ctx, servers.MCPServers[server])
95+
if err != nil {
96+
return nil, fmt.Errorf("failed to load MCP session for server %s: %w", server, err)
97+
}
98+
99+
return l.sessionToTools(ctx, session, tool.Name)
100+
}
101+
102+
// This should never happen, but just in case
103+
return nil, fmt.Errorf("no MCP server configuration found in tool instructions: %s", configData)
104+
}
105+
106+
func (l *Local) sessionToTools(ctx context.Context, session *Session, toolName string) ([]types.Tool, error) {
107+
tools, err := session.Client.ListTools(ctx, mcp.ListToolsRequest{})
108+
if err != nil {
109+
return nil, fmt.Errorf("failed to list tools: %w", err)
110+
}
111+
112+
toolDefs := []types.Tool{{ /* this is a placeholder for main tool */ }}
113+
var toolNames []string
114+
115+
for _, tool := range tools.Tools {
116+
var schema openapi3.Schema
117+
118+
schemaData, err := json.Marshal(tool.InputSchema)
119+
if err != nil {
120+
panic(err)
121+
}
122+
123+
if tool.Name == "" {
124+
// I dunno, bad tool?
125+
continue
126+
}
127+
128+
if err := json.Unmarshal(schemaData, &schema); err != nil {
129+
return nil, fmt.Errorf("failed to unmarshal tool input schema: %w", err)
130+
}
131+
132+
annotations, err := json.Marshal(tool.Annotations)
133+
if err != nil {
134+
return nil, fmt.Errorf("failed to marshal tool annotations: %w", err)
135+
}
136+
137+
toolDef := types.Tool{
138+
ToolDef: types.ToolDef{
139+
Parameters: types.Parameters{
140+
Name: tool.Name,
141+
Description: tool.Description,
142+
Arguments: &schema,
143+
},
144+
Instructions: types.MCPInvokePrefix + "." + tool.Name + " " + session.ID + " " + tool.Name,
145+
},
146+
}
147+
148+
if string(annotations) != "{}" {
149+
toolDef.MetaData = map[string]string{
150+
"mcp-tool-annotations": string(annotations),
151+
}
152+
}
153+
154+
if tool.Annotations.Title != "" && !slices.Contains(strings.Fields(tool.Annotations.Title), "as") {
155+
toolDef.Name = tool.Annotations.Title + " as " + tool.Name
156+
}
157+
158+
toolDefs = append(toolDefs, toolDef)
159+
toolNames = append(toolNames, tool.Name)
160+
}
161+
162+
main := types.Tool{
163+
ToolDef: types.ToolDef{
164+
Parameters: types.Parameters{
165+
Name: toolName,
166+
Description: session.InitResult.ServerInfo.Name,
167+
Export: toolNames,
168+
},
169+
MetaData: map[string]string{
170+
"bundle": "true",
171+
},
172+
},
173+
}
174+
175+
if session.InitResult.Instructions != "" {
176+
data, _ := json.Marshal(map[string]any{
177+
"tools": toolNames,
178+
"instructions": session.InitResult.Instructions,
179+
})
180+
toolDefs = append(toolDefs, types.Tool{
181+
ToolDef: types.ToolDef{
182+
Parameters: types.Parameters{
183+
Name: session.ID,
184+
Type: "context",
185+
},
186+
Instructions: types.EchoPrefix + "\n" + `# START MCP SERVER INFO: ` + session.InitResult.ServerInfo.Name + "\n" +
187+
`You have available the following tools from an MCP Server that has provided the following additional instructions` + "\n" +
188+
string(data) + "\n" +
189+
`# END MCP SERVER INFO` + "\n",
190+
},
191+
})
192+
193+
main.ExportContext = append(main.ExportContext, session.ID)
194+
}
195+
196+
toolDefs[0] = main
197+
return toolDefs, nil
198+
}
199+
200+
func (l *Local) loadSession(ctx context.Context, server ServerConfig) (*Session, error) {
201+
id := hash.Digest(server)
202+
l.lock.Lock()
203+
existing, ok := l.sessions[id]
204+
l.lock.Unlock()
205+
if ok {
206+
return existing, nil
207+
}
208+
209+
var (
210+
c client.MCPClient
211+
err error
212+
)
213+
214+
if server.Command != "" {
215+
env := make([]string, 0, len(server.Env))
216+
for k, v := range server.Env {
217+
env = append(env, fmt.Sprintf("%s=%s", k, v))
218+
}
219+
c, err = client.NewStdioMCPClient(server.Command, env, server.Args...)
220+
if err != nil {
221+
return nil, fmt.Errorf("failed to create MCP stdio client: %w", err)
222+
}
223+
} else {
224+
url := server.URL
225+
if url == "" {
226+
url = server.Server
227+
}
228+
c, err = client.NewSSEMCPClient(url, client.WithHeaders(server.Headers))
229+
if err != nil {
230+
return nil, fmt.Errorf("failed to create MCP HTTP client: %w", err)
231+
}
232+
}
233+
234+
var initRequest mcp.InitializeRequest
235+
initRequest.Params.ClientInfo = mcp.Implementation{
236+
Name: version.ProgramName,
237+
Version: version.Get().String(),
238+
}
239+
240+
initResult, err := c.Initialize(ctx, initRequest)
241+
if err != nil {
242+
return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
243+
}
244+
245+
result := &Session{
246+
ID: id,
247+
InitResult: initResult,
248+
Client: c,
249+
Config: server,
250+
}
251+
252+
l.lock.Lock()
253+
defer l.lock.Unlock()
254+
255+
if existing, ok := l.sessions[id]; ok {
256+
return existing, c.Close()
257+
}
258+
259+
if l.sessions == nil {
260+
l.sessions = make(map[string]*Session)
261+
}
262+
l.sessions[id] = result
263+
return result, nil
264+
}

‎pkg/mcp/runner.go

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package mcp
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"strings"
8+
9+
"github.com/gptscript-ai/gptscript/pkg/types"
10+
"github.com/mark3labs/mcp-go/mcp"
11+
)
12+
13+
func (l *Local) Run(ctx context.Context, _ chan<- types.CompletionStatus, tool types.Tool, input string) (string, error) {
14+
fields := strings.Fields(tool.Instructions)
15+
if len(fields) < 3 {
16+
return "", fmt.Errorf("invalid mcp call, invalid number of fields in %s", tool.Instructions)
17+
}
18+
19+
id := fields[1]
20+
toolName := fields[2]
21+
arguments := map[string]any{}
22+
23+
if input != "" {
24+
if err := json.Unmarshal([]byte(input), &arguments); err != nil {
25+
return "", fmt.Errorf("failed to unmarshal input: %w", err)
26+
}
27+
}
28+
29+
l.lock.Lock()
30+
session, ok := l.sessions[id]
31+
l.lock.Unlock()
32+
if !ok {
33+
return "", fmt.Errorf("session not found for MCP server %s", id)
34+
}
35+
36+
request := mcp.CallToolRequest{}
37+
request.Params.Name = toolName
38+
request.Params.Arguments = arguments
39+
40+
result, err := session.Client.CallTool(ctx, request)
41+
if err != nil {
42+
return "", fmt.Errorf("failed to call tool %s: %w", toolName, err)
43+
}
44+
45+
str, err := json.Marshal(result)
46+
if err != nil {
47+
return "", fmt.Errorf("failed to marshal result: %w", err)
48+
}
49+
50+
return string(str), nil
51+
}

‎pkg/tests/runner2_test.go

Lines changed: 352 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"github.com/gptscript-ai/gptscript/pkg/loader"
99
"github.com/gptscript-ai/gptscript/pkg/runner"
1010
"github.com/gptscript-ai/gptscript/pkg/tests/tester"
11+
"github.com/gptscript-ai/gptscript/pkg/types"
1112
"github.com/hexops/autogold/v2"
1213
"github.com/stretchr/testify/require"
1314
)
@@ -203,3 +204,354 @@ echo "${GPTSCRIPT_INPUT}"
203204
require.NoError(t, err)
204205
autogold.Expect(map[string]interface{}{"foo": "baz", "start": true}).Equal(t, data)
205206
}
207+
208+
func TestMCPLoad(t *testing.T) {
209+
r := tester.NewRunner(t)
210+
prg, err := loader.ProgramFromSource(context.Background(), `
211+
name: mcp
212+
213+
#!mcp
214+
215+
{
216+
"mcpServers": {
217+
"sqlite": {
218+
"command": "docker",
219+
"args": [
220+
"run",
221+
"--rm",
222+
"-i",
223+
"-v",
224+
"mcp-test:/mcp",
225+
"mcp/sqlite@sha256:007ccae941a6f6db15b26ee41d92edda50ce157176d9273449e8b3f51d979c70",
226+
"--db-path",
227+
"/mcp/test.db"
228+
]
229+
}
230+
}
231+
}
232+
`, "")
233+
require.NoError(t, err)
234+
235+
autogold.Expect(types.Tool{
236+
ToolDef: types.ToolDef{
237+
Parameters: types.Parameters{
238+
Name: "mcp",
239+
Description: "sqlite",
240+
ModelName: "gpt-4o",
241+
Export: []string{
242+
"read_query",
243+
"write_query",
244+
"create_table",
245+
"list_tables",
246+
"describe_table",
247+
"append_insight",
248+
},
249+
},
250+
MetaData: map[string]string{"bundle": "true"},
251+
},
252+
ID: "inline:mcp",
253+
ToolMapping: map[string][]types.ToolReference{
254+
"append_insight": {{
255+
Reference: "append_insight",
256+
ToolID: "inline:append_insight",
257+
}},
258+
"create_table": {{
259+
Reference: "create_table",
260+
ToolID: "inline:create_table",
261+
}},
262+
"describe_table": {{
263+
Reference: "describe_table",
264+
ToolID: "inline:describe_table",
265+
}},
266+
"list_tables": {{
267+
Reference: "list_tables",
268+
ToolID: "inline:list_tables",
269+
}},
270+
"read_query": {{
271+
Reference: "read_query",
272+
ToolID: "inline:read_query",
273+
}},
274+
"write_query": {{
275+
Reference: "write_query",
276+
ToolID: "inline:write_query",
277+
}},
278+
},
279+
LocalTools: map[string]string{
280+
"append_insight": "inline:append_insight",
281+
"create_table": "inline:create_table",
282+
"describe_table": "inline:describe_table",
283+
"list_tables": "inline:list_tables",
284+
"mcp": "inline:mcp",
285+
"read_query": "inline:read_query",
286+
"write_query": "inline:write_query",
287+
},
288+
Source: types.ToolSource{Location: "inline"},
289+
WorkingDir: ".",
290+
}).Equal(t, prg.ToolSet[prg.EntryToolID])
291+
autogold.Expect(7).Equal(t, len(prg.ToolSet[prg.EntryToolID].LocalTools))
292+
data, _ := json.MarshalIndent(prg.ToolSet, "", " ")
293+
autogold.Expect(`{
294+
"inline:append_insight": {
295+
"name": "append_insight",
296+
"description": "Add a business insight to the memo",
297+
"modelName": "gpt-4o",
298+
"internalPrompt": null,
299+
"arguments": {
300+
"properties": {
301+
"insight": {
302+
"description": "Business insight discovered from data analysis",
303+
"type": "string"
304+
}
305+
},
306+
"required": [
307+
"insight"
308+
],
309+
"type": "object"
310+
},
311+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f append_insight",
312+
"id": "inline:append_insight",
313+
"localTools": {
314+
"append_insight": "inline:append_insight",
315+
"create_table": "inline:create_table",
316+
"describe_table": "inline:describe_table",
317+
"list_tables": "inline:list_tables",
318+
"mcp": "inline:mcp",
319+
"read_query": "inline:read_query",
320+
"write_query": "inline:write_query"
321+
},
322+
"source": {
323+
"location": "inline"
324+
},
325+
"workingDir": "."
326+
},
327+
"inline:create_table": {
328+
"name": "create_table",
329+
"description": "Create a new table in the SQLite database",
330+
"modelName": "gpt-4o",
331+
"internalPrompt": null,
332+
"arguments": {
333+
"properties": {
334+
"query": {
335+
"description": "CREATE TABLE SQL statement",
336+
"type": "string"
337+
}
338+
},
339+
"required": [
340+
"query"
341+
],
342+
"type": "object"
343+
},
344+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f create_table",
345+
"id": "inline:create_table",
346+
"localTools": {
347+
"append_insight": "inline:append_insight",
348+
"create_table": "inline:create_table",
349+
"describe_table": "inline:describe_table",
350+
"list_tables": "inline:list_tables",
351+
"mcp": "inline:mcp",
352+
"read_query": "inline:read_query",
353+
"write_query": "inline:write_query"
354+
},
355+
"source": {
356+
"location": "inline"
357+
},
358+
"workingDir": "."
359+
},
360+
"inline:describe_table": {
361+
"name": "describe_table",
362+
"description": "Get the schema information for a specific table",
363+
"modelName": "gpt-4o",
364+
"internalPrompt": null,
365+
"arguments": {
366+
"properties": {
367+
"table_name": {
368+
"description": "Name of the table to describe",
369+
"type": "string"
370+
}
371+
},
372+
"required": [
373+
"table_name"
374+
],
375+
"type": "object"
376+
},
377+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f describe_table",
378+
"id": "inline:describe_table",
379+
"localTools": {
380+
"append_insight": "inline:append_insight",
381+
"create_table": "inline:create_table",
382+
"describe_table": "inline:describe_table",
383+
"list_tables": "inline:list_tables",
384+
"mcp": "inline:mcp",
385+
"read_query": "inline:read_query",
386+
"write_query": "inline:write_query"
387+
},
388+
"source": {
389+
"location": "inline"
390+
},
391+
"workingDir": "."
392+
},
393+
"inline:list_tables": {
394+
"name": "list_tables",
395+
"description": "List all tables in the SQLite database",
396+
"modelName": "gpt-4o",
397+
"internalPrompt": null,
398+
"arguments": {
399+
"type": "object"
400+
},
401+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f list_tables",
402+
"id": "inline:list_tables",
403+
"localTools": {
404+
"append_insight": "inline:append_insight",
405+
"create_table": "inline:create_table",
406+
"describe_table": "inline:describe_table",
407+
"list_tables": "inline:list_tables",
408+
"mcp": "inline:mcp",
409+
"read_query": "inline:read_query",
410+
"write_query": "inline:write_query"
411+
},
412+
"source": {
413+
"location": "inline"
414+
},
415+
"workingDir": "."
416+
},
417+
"inline:mcp": {
418+
"name": "mcp",
419+
"description": "sqlite",
420+
"modelName": "gpt-4o",
421+
"internalPrompt": null,
422+
"export": [
423+
"read_query",
424+
"write_query",
425+
"create_table",
426+
"list_tables",
427+
"describe_table",
428+
"append_insight"
429+
],
430+
"metaData": {
431+
"bundle": "true"
432+
},
433+
"id": "inline:mcp",
434+
"toolMapping": {
435+
"append_insight": [
436+
{
437+
"reference": "append_insight",
438+
"toolID": "inline:append_insight"
439+
}
440+
],
441+
"create_table": [
442+
{
443+
"reference": "create_table",
444+
"toolID": "inline:create_table"
445+
}
446+
],
447+
"describe_table": [
448+
{
449+
"reference": "describe_table",
450+
"toolID": "inline:describe_table"
451+
}
452+
],
453+
"list_tables": [
454+
{
455+
"reference": "list_tables",
456+
"toolID": "inline:list_tables"
457+
}
458+
],
459+
"read_query": [
460+
{
461+
"reference": "read_query",
462+
"toolID": "inline:read_query"
463+
}
464+
],
465+
"write_query": [
466+
{
467+
"reference": "write_query",
468+
"toolID": "inline:write_query"
469+
}
470+
]
471+
},
472+
"localTools": {
473+
"append_insight": "inline:append_insight",
474+
"create_table": "inline:create_table",
475+
"describe_table": "inline:describe_table",
476+
"list_tables": "inline:list_tables",
477+
"mcp": "inline:mcp",
478+
"read_query": "inline:read_query",
479+
"write_query": "inline:write_query"
480+
},
481+
"source": {
482+
"location": "inline"
483+
},
484+
"workingDir": "."
485+
},
486+
"inline:read_query": {
487+
"name": "read_query",
488+
"description": "Execute a SELECT query on the SQLite database",
489+
"modelName": "gpt-4o",
490+
"internalPrompt": null,
491+
"arguments": {
492+
"properties": {
493+
"query": {
494+
"description": "SELECT SQL query to execute",
495+
"type": "string"
496+
}
497+
},
498+
"required": [
499+
"query"
500+
],
501+
"type": "object"
502+
},
503+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f read_query",
504+
"id": "inline:read_query",
505+
"localTools": {
506+
"append_insight": "inline:append_insight",
507+
"create_table": "inline:create_table",
508+
"describe_table": "inline:describe_table",
509+
"list_tables": "inline:list_tables",
510+
"mcp": "inline:mcp",
511+
"read_query": "inline:read_query",
512+
"write_query": "inline:write_query"
513+
},
514+
"source": {
515+
"location": "inline"
516+
},
517+
"workingDir": "."
518+
},
519+
"inline:write_query": {
520+
"name": "write_query",
521+
"description": "Execute an INSERT, UPDATE, or DELETE query on the SQLite database",
522+
"modelName": "gpt-4o",
523+
"internalPrompt": null,
524+
"arguments": {
525+
"properties": {
526+
"query": {
527+
"description": "SQL query to execute",
528+
"type": "string"
529+
}
530+
},
531+
"required": [
532+
"query"
533+
],
534+
"type": "object"
535+
},
536+
"instructions": "#!sys.mcp.invoke 441826308787ad271e84a381e90d8eccc3fce0fe94503636e679bd0984c79f2f write_query",
537+
"id": "inline:write_query",
538+
"localTools": {
539+
"append_insight": "inline:append_insight",
540+
"create_table": "inline:create_table",
541+
"describe_table": "inline:describe_table",
542+
"list_tables": "inline:list_tables",
543+
"mcp": "inline:mcp",
544+
"read_query": "inline:read_query",
545+
"write_query": "inline:write_query"
546+
},
547+
"source": {
548+
"location": "inline"
549+
},
550+
"workingDir": "."
551+
}
552+
}`).Equal(t, string(data))
553+
554+
prg.EntryToolID = prg.ToolSet[prg.EntryToolID].LocalTools["read_query"]
555+
resp, err := r.Chat(context.Background(), nil, prg, nil, `{"query": "SELECT 1"}`, runner.RunOptions{})
556+
r.AssertStep(t, resp, err)
557+
}

‎pkg/tests/testdata/TestMCPLoad/call1-resp.golden

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
`{
2+
"role": "assistant",
3+
"content": [
4+
{
5+
"text": "TEST RESULT CALL: 1"
6+
}
7+
],
8+
"usage": {}
9+
}`

‎pkg/tests/testdata/TestMCPLoad/call1.golden

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
`{
2+
"model": "gpt-4o"
3+
}`

‎pkg/tests/testdata/TestMCPLoad/step1.golden

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
`{
2+
"done": true,
3+
"content": "{\"content\":[{\"type\":\"text\",\"text\":\"[{'1': 1}]\"}]}",
4+
"toolID": "",
5+
"state": null
6+
}`

‎pkg/types/tool.go

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,14 @@ import (
1616
)
1717

1818
const (
19-
DaemonPrefix = "#!sys.daemon"
20-
OpenAPIPrefix = "#!sys.openapi"
21-
EchoPrefix = "#!sys.echo"
22-
CallPrefix = "#!sys.call"
23-
CommandPrefix = "#!"
19+
DaemonPrefix = "#!sys.daemon"
20+
OpenAPIPrefix = "#!sys.openapi"
21+
EchoPrefix = "#!sys.echo"
22+
CallPrefix = "#!sys.call"
23+
MCPPrefix = "#!mcp"
24+
MCPInvokePrefix = "#!sys.mcp.invoke"
25+
CommandPrefix = "#!"
26+
PromptPrefix = "!!"
2427
)
2528

2629
var (
@@ -862,6 +865,14 @@ func (t Tool) IsDaemon() bool {
862865
return strings.HasPrefix(t.Instructions, DaemonPrefix)
863866
}
864867

868+
func (t Tool) IsMCP() bool {
869+
return strings.HasPrefix(t.Instructions, MCPPrefix)
870+
}
871+
872+
func (t Tool) IsMCPInvoke() bool {
873+
return strings.HasPrefix(t.Instructions, MCPInvokePrefix)
874+
}
875+
865876
func (t Tool) IsOpenAPI() bool {
866877
return strings.HasPrefix(t.Instructions, OpenAPIPrefix)
867878
}

‎pkg/types/toolstring.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@ func ToDisplayText(tool Tool, input string) string {
4444
}
4545

4646
func ToSysDisplayString(id string, args map[string]string) (string, error) {
47+
if suffix, ok := strings.CutPrefix(id, "sys.mcp.invoke."); ok {
48+
return fmt.Sprintf("Invoking MCP `%s`", suffix), nil
49+
}
50+
4751
switch id {
4852
case "sys.append":
4953
return fmt.Sprintf("Appending to file `%s`", args["filename"]), nil

0 commit comments

Comments
 (0)
Please sign in to comment.