From 415a5fb2497ac4ba6381540c5cc4948b55223308 Mon Sep 17 00:00:00 2001 From: "ark-hand[bot]" <315378070+ark-hand[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:26:29 +0000 Subject: [PATCH] feat(selfhosted): add client-side MCP tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 简述 为 Go self-hosted worker 增加 client-side MCP tools,并适配当前 Managed Agents Custom Tool 契约。核心模块保持 Go 1.20,可选官方 MCP adapter 使用独立 Go 1.23 module;同时修复通用 HTTP retry 的等待时间、服务端提示与重试标记。 ## 改动 - core 增加协议无关 MCP Client、ToolDefinition、结果模型和 Custom Tool 转换。 - 可选 `github.com/volcengine/ark-runtime-go/mcp` 接入官方 MCP Go SDK。 - 支持 text、image、embedded resource、PDF 与 structured content。 - `type/properties/required` 结构化传递,property 内本地引用内联,其余有效顶层约束以紧凑 JSON 写入 tool description。 - 增加根版本发布时自动创建匹配 `mcp/vMAJOR.MINOR.0` tag 的 workflow;adapter 从 `v0.6.0` 开始与 core 同版本安装。 - MCP 凭证和连接仅保留在 worker 侧。 ## SDK retry 修复 - 默认最大重试次数保持 2 次;首次等待改为 0.5 秒基数,之后指数增长并封顶 8 秒,每次减去 0–25% jitter。 - 直接按 `time.Duration` 的纳秒精度计算退避,避免先转整数秒导致首次重试接近 0 秒。 - 优先读取 `Retry-After-Ms`,其次读取 `Retry-After`;支持小数数值以及 HTTP-date。 - 首次请求和后续重试分别携带 `X-Stainless-Retry-Count: 0/1/2`,调用方显式设置该 header 时保留调用方值。 - 支持 `X-Should-Retry` 显式控制,并将 408、409、429、5xx 统一视为可重试状态。 - 错误响应保留 response header 供退避决策使用;流式握手失败时关闭响应体,避免连接占用。 ## 其他行为修复 - 文本 content block 即使内容为空也保留 MA 契约必需的 `text` 字段。 - MCP 返回 `isError=true` 且没有内容时,生成 `tool returned an error`,避免空错误结果。 - embedded resource 转换失败时不再把 URI 写入 tool result,避免 signed URL 或查询凭证进入 MA、模型和日志。 - property 内已成功内联的 `$defs/definitions` 不再重复写入 description;仍被其他约束引用时继续保留。 - 不修改 MA 或 Agent Loop,仍适配现网 `type/properties/required` 三字段 Custom Tool 契约。 ## 测试 - `go test ./...`、`go vet ./...` - `golangci-lint v1.64.8` - MCP module: `go test ./...`、`go vet ./...` - retry 新增亚秒退避、8 秒上限、服务端等待提示、状态码/显式控制及 retry-count 集成测试。 - GitHub Actions `actionlint` - 本地模拟 `v0.6.0` + `mcp/v0.6.0` 后,从干净 consumer 执行 `go get .../mcp@v0.6.0` 成功。 - STG 真实链路:MCP `tools/list -> Agent custom tool -> worker CallTool -> user.custom_tool_result -> 模型继续推理` 通过,并完成 Python/Java 版本工具调用。 ## 示例 - 新增 `examples/self_hosted_mcp_worker`,以最小粒度展示 `tools/list -> Custom Tool 定义 -> MCP Tool 执行 -> EnvironmentWorker.Run`。 - 示例通过独立 module 引入可选 MCP adapter,不提高 core module 的 Go 版本或基础依赖。 - 内置最小 stdio `mcp_echo` server,便于本地直接验证;示例不负责创建 Agent、Session 或消费 SSE。 See merge request: !89 Sync-Source-Commit: d8736eda4005b445901dca9c6e85b4e99f1d675d Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.6.0 --- .github/workflows/ci.yml | 41 ++ .github/workflows/release-mcp.yml | 61 +++ THIRD_PARTY_NOTICES.md | 4 +- arkruntime/client.go | 146 +++-- arkruntime/environment_work.go | 11 +- arkruntime/model/common.go | 4 + arkruntime/model/error.go | 31 +- arkruntime/retry_test.go | 169 ++++++ arkruntime/selfhosted/mcp/mcp.go | 515 ++++++++++++++++++ arkruntime/selfhosted/mcp/mcp_test.go | 229 ++++++++ arkruntime/selfhosted/session_tool_runner.go | 3 + .../selfhosted/session_tool_runner_test.go | 30 + arkruntime/selfhosted/types.go | 18 + arkruntime/toolset/types.go | 3 + arkruntime/utils/retry.go | 50 +- arkruntime/utils/retry_test.go | 86 +++ examples/README.md | 2 + examples/self_hosted_mcp_worker/README.md | 113 ++++ examples/self_hosted_mcp_worker/go.mod | 29 + examples/self_hosted_mcp_worker/go.sum | 131 +++++ examples/self_hosted_mcp_worker/main.go | 126 +++++ .../self_hosted_mcp_worker/server/main.go | 40 ++ mcp/README.md | 114 ++++ mcp/go.mod | 26 + mcp/go.sum | 131 +++++ mcp/mcp.go | 153 ++++++ mcp/mcp_test.go | 160 ++++++ 27 files changed, 2369 insertions(+), 57 deletions(-) create mode 100644 .github/workflows/release-mcp.yml create mode 100644 arkruntime/retry_test.go create mode 100644 arkruntime/selfhosted/mcp/mcp.go create mode 100644 arkruntime/selfhosted/mcp/mcp_test.go create mode 100644 arkruntime/utils/retry_test.go create mode 100644 examples/self_hosted_mcp_worker/README.md create mode 100644 examples/self_hosted_mcp_worker/go.mod create mode 100644 examples/self_hosted_mcp_worker/go.sum create mode 100644 examples/self_hosted_mcp_worker/main.go create mode 100644 examples/self_hosted_mcp_worker/server/main.go create mode 100644 mcp/README.md create mode 100644 mcp/go.mod create mode 100644 mcp/go.sum create mode 100644 mcp/mcp.go create mode 100644 mcp/mcp_test.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a999053..a036255 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,6 +21,24 @@ concurrency: cancel-in-progress: true jobs: + minimum-go-build: + name: Build on Go 1.20 + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: "1.20.x" + cache: true + + - name: Build core SDK + run: go build ./... + go-ci: name: Go CI runs-on: ubuntu-latest @@ -53,3 +71,26 @@ jobs: with: version: v1.64.8 args: --timeout=5m --concurrency=8 --issues-exit-code=1 + + mcp-ci: + name: MCP adapter CI + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: "1.23.x" + cache-dependency-path: mcp/go.sum + + - name: Test optional MCP adapter + working-directory: mcp + run: go test ./... + + - name: Build self-hosted MCP example + working-directory: examples/self_hosted_mcp_worker + run: go build ./... diff --git a/.github/workflows/release-mcp.yml b/.github/workflows/release-mcp.yml new file mode 100644 index 0000000..1341739 --- /dev/null +++ b/.github/workflows/release-mcp.yml @@ -0,0 +1,61 @@ +name: Release MCP module + +on: + push: + tags: + - "v*.*.0" + +permissions: + contents: write + +jobs: + tag-mcp-module: + name: Tag matching MCP module + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - name: Check out release + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: "1.23.x" + cache-dependency-path: mcp/go.sum + + - name: Validate and test MCP module + env: + RELEASE_TAG: ${{ github.ref_name }} + run: | + set -euo pipefail + version="${RELEASE_TAG#v}" + if [[ ! "${version}" =~ ^[0-9]+\.[0-9]+\.0$ ]]; then + echo "unsupported release version: ${version}" >&2 + exit 1 + fi + required_version="$(cd mcp && go list -m -f '{{if eq .Path "github.com/volcengine/ark-runtime-go"}}{{.Version}}{{end}}' all)" + if [[ "${required_version}" != "v${version}" ]]; then + echo "mcp requires core ${required_version}, want v${version}" >&2 + exit 1 + fi + (cd mcp && go test ./...) + + - name: Publish MCP module tag + env: + RELEASE_TAG: ${{ github.ref_name }} + run: | + set -euo pipefail + nested_tag="mcp/${RELEASE_TAG}" + existing="$(git ls-remote origin "refs/tags/${nested_tag}" | cut -f1)" + if [[ -n "${existing}" ]]; then + if [[ "${existing}" != "${GITHUB_SHA}" ]]; then + echo "${nested_tag} already points to ${existing}, want ${GITHUB_SHA}" >&2 + exit 1 + fi + exit 0 + fi + git tag "${nested_tag}" "${GITHUB_SHA}" + git push origin "refs/tags/${nested_tag}" diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 8cfd333..631680e 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -8,8 +8,8 @@ from third-party open-source projects. Portions of the self-hosted worker lifecycle and local agent tool implementations under `arkruntime/selfhosted`, `arkruntime/lib/environments`, `arkruntime/toolset`, -and `arkruntime/tools/agenttoolset` are structurally adapted from Anthropic's -self-hosted worker SDK implementation: +`arkruntime/tools/agenttoolset`, and `mcp` are structurally adapted from +Anthropic's self-hosted worker SDK and client-side MCP helper implementation: https://github.com/anthropics/anthropic-sdk-go diff --git a/arkruntime/client.go b/arkruntime/client.go index dba29e9..6312b68 100644 --- a/arkruntime/client.go +++ b/arkruntime/client.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "io" + "math" "math/rand" "net/http" "net/url" @@ -362,19 +363,16 @@ func (c *Client) sendRequest(client *http.Client, req *http.Request, v model.Res } func (c *Client) Do(ctx context.Context, method, url, resourceType, resourceId string, v model.Response, setters ...requestOption) (err error) { - err = utils.Retry( + err = utils.RetryWithAttempt( ctx, - utils.RetryPolicy{ - MaxAttempts: c.config.RetryTimes, - InitialBackoff: model.ErrorRetryBaseDelay, - MaxBackoff: model.ErrorRetryMaxDelay, - }, + defaultRetryPolicy(c.config.RetryTimes), func() bool { return true }, - func() error { + func(retryCount int) error { req, innerErr := c.newRequest(ctx, method, url, resourceType, resourceId, setters...) if innerErr != nil { return innerErr } + setRetryCountHeader(req, retryCount) return c.sendRequest(c.config.HTTPClient, req, v) }, @@ -508,19 +506,16 @@ func sendCreateResponsesRequestStream(client *Client, httpClient *http.Client, r // ImageGenerationStreamRequestDo executes an /images/generations request // with stream=true and returns a gen-typed reader. func (c *Client) ImageGenerationStreamRequestDo(ctx context.Context, method, url, resourceId string, setters ...requestOption) (streamReader *utils.ImageGenerationStreamReader, err error) { - err = utils.Retry( + err = utils.RetryWithAttempt( ctx, - utils.RetryPolicy{ - MaxAttempts: c.config.RetryTimes, - InitialBackoff: model.ErrorRetryBaseDelay, - MaxBackoff: model.ErrorRetryMaxDelay, - }, + defaultRetryPolicy(c.config.RetryTimes), func() bool { return true }, - func() error { + func(retryCount int) error { req, innerErr := c.newRequest(ctx, method, url, resourceTypeEndpoint, resourceId, setters...) if innerErr != nil { return innerErr } + setRetryCountHeader(req, retryCount) streamReader, err = sendImageGenerationStream(c, c.config.HTTPClient, req) return err @@ -534,19 +529,16 @@ func (c *Client) ImageGenerationStreamRequestDo(ctx context.Context, method, url // ChatGenStreamRequestDo executes a chat-completions stream request and // returns a gen-typed reader. func (c *Client) ChatGenStreamRequestDo(ctx context.Context, method, url, resourceId string, setters ...requestOption) (streamReader *utils.ChatGenStreamReader, err error) { - err = utils.Retry( + err = utils.RetryWithAttempt( ctx, - utils.RetryPolicy{ - MaxAttempts: c.config.RetryTimes, - InitialBackoff: model.ErrorRetryBaseDelay, - MaxBackoff: model.ErrorRetryMaxDelay, - }, + defaultRetryPolicy(c.config.RetryTimes), func() bool { return true }, - func() error { + func(retryCount int) error { req, innerErr := c.newRequest(ctx, method, url, resourceTypeEndpoint, resourceId, setters...) if innerErr != nil { return innerErr } + setRetryCountHeader(req, retryCount) streamReader, err = sendChatGenStream(c, c.config.HTTPClient, req) return err @@ -560,19 +552,16 @@ func (c *Client) ChatGenStreamRequestDo(ctx context.Context, method, url, resour // ResponsesRequestStreamDo executes a request. func (c *Client) ResponsesRequestStreamDo(ctx context.Context, method, url, resourceType, resourceId string, setters ...requestOption) (resp *utils.ResponsesStreamReader, err error) { - err = utils.Retry( + err = utils.RetryWithAttempt( ctx, - utils.RetryPolicy{ - MaxAttempts: c.config.RetryTimes, - InitialBackoff: model.ErrorRetryBaseDelay, - MaxBackoff: model.ErrorRetryMaxDelay, - }, + defaultRetryPolicy(c.config.RetryTimes), func() bool { return true }, - func() error { + func(retryCount int) error { req, innerErr := c.newRequest(ctx, method, url, resourceType, resourceId, setters...) if innerErr != nil { return innerErr } + setRetryCountHeader(req, retryCount) resp, err = sendCreateResponsesRequestStream(c, c.config.HTTPClient, req) return err }, @@ -624,18 +613,33 @@ func isFailureStatusCode(resp *http.Response) bool { } func needRetryError(err error) bool { + if header, ok := responseHeader(err); ok { + switch strings.ToLower(header.Get(model.ShouldRetryHeader)) { + case "true": + return true + case "false": + return false + } + } apiErr := &model.APIError{} reqErr := &model.RequestError{} if errors.As(err, &apiErr) { - return apiErr.HTTPStatusCode >= http.StatusInternalServerError || apiErr.HTTPStatusCode == http.StatusTooManyRequests + return isRetryableStatus(apiErr.HTTPStatusCode) } else if errors.Is(err, io.EOF) { return true } else if errors.As(err, &reqErr) { - return reqErr.HTTPStatusCode >= http.StatusInternalServerError + return isRetryableStatus(reqErr.HTTPStatusCode) } return false } +func isRetryableStatus(statusCode int) bool { + return statusCode == http.StatusRequestTimeout || + statusCode == http.StatusConflict || + statusCode == http.StatusTooManyRequests || + statusCode >= http.StatusInternalServerError +} + func decodeResponse(body io.Reader, v interface{}) error { if v == nil { return nil @@ -663,19 +667,23 @@ func (c *Client) fullURL(suffix string) string { } func (c *Client) handleErrorResp(resp *http.Response) error { + // Streaming callers rely on this close. Non-streaming callers also defer + // Close upstream and rely on response bodies supporting idempotent Close. + defer resp.Body.Close() //nolint:errcheck // response body close errors are non-actionable requestID := responseRequestID(resp) body, readErr := io.ReadAll(resp.Body) if readErr != nil { - return model.NewRequestError( + return newResponseRequestError( resp.StatusCode, fmt.Errorf("read error response body: %w", readErr), requestID, + resp.Header, ) } var errRes model.ErrorResponse if err := json.Unmarshal(body, &errRes); err == nil && errRes.Error != nil { - return setAPIErrorResponseMetadata(errRes.Error, resp.StatusCode, requestID) + return setAPIErrorResponseMetadata(errRes.Error, resp.StatusCode, requestID, resp.Header) } // Some services return the error object directly instead of wrapping it in @@ -683,21 +691,23 @@ func (c *Client) handleErrorResp(resp *http.Response) error { var apiErr model.APIError if err := json.Unmarshal(body, &apiErr); err == nil && (apiErr.Message != "" || apiErr.Code != "" || apiErr.Type != "") { - return setAPIErrorResponseMetadata(&apiErr, resp.StatusCode, requestID) + return setAPIErrorResponseMetadata(&apiErr, resp.StatusCode, requestID, resp.Header) } bodyText := strings.TrimSpace(string(body)) if bodyText == "" { - return model.NewRequestError( + return newResponseRequestError( resp.StatusCode, errors.New("unexpected error response: empty body"), requestID, + resp.Header, ) } - return model.NewRequestError( + return newResponseRequestError( resp.StatusCode, fmt.Errorf("unexpected error response body: %s", bodyText), requestID, + resp.Header, ) } @@ -714,14 +724,78 @@ func responseRequestID(resp *http.Response) string { return "" } -func setAPIErrorResponseMetadata(apiErr *model.APIError, statusCode int, requestID string) error { +func setAPIErrorResponseMetadata(apiErr *model.APIError, statusCode int, requestID string, header http.Header) error { apiErr.HTTPStatusCode = statusCode + apiErr.ResponseHeader = header.Clone() if requestID != "" { apiErr.RequestId = requestID } return apiErr } +func newResponseRequestError(statusCode int, err error, requestID string, header http.Header) error { + requestErr := model.NewRequestError(statusCode, err, requestID) + requestErr.ResponseHeader = header.Clone() + return requestErr +} + +func defaultRetryPolicy(maxAttempts int) utils.RetryPolicy { + return utils.RetryPolicy{ + MaxAttempts: maxAttempts, + InitialBackoff: model.ErrorRetryBaseDelay, + MaxBackoff: model.ErrorRetryMaxDelay, + MaxRetryAfter: model.MaxServerRetryDelay, + RetryAfter: retryAfter, + } +} + +func retryAfter(err error) (time.Duration, bool) { + header, ok := responseHeader(err) + if !ok { + return 0, false + } + for _, retry := range []struct { + name string + unit time.Duration + }{ + {name: model.RetryAfterMSHeader, unit: time.Millisecond}, + {name: model.RetryAfterHeader, unit: time.Second}, + } { + value := header.Get(retry.name) + if value == "" { + continue + } + if parsed, parseErr := strconv.ParseFloat(value, 64); parseErr == nil { + if math.IsNaN(parsed) || math.IsInf(parsed, 0) || + parsed > float64(math.MaxInt64)/float64(retry.unit) || + parsed < float64(math.MinInt64)/float64(retry.unit) { + continue + } + return time.Duration(parsed * float64(retry.unit)), true + } + if retry.name == model.RetryAfterHeader { + if retryAt, parseErr := http.ParseTime(value); parseErr == nil { + return time.Until(retryAt), true + } + } + } + return 0, false +} + +func responseHeader(err error) (http.Header, bool) { + var responseErr interface{ GetHeader() http.Header } + if !errors.As(err, &responseErr) || responseErr == nil || responseErr.GetHeader() == nil { + return nil, false + } + return responseErr.GetHeader(), true +} + +func setRetryCountHeader(req *http.Request, retryCount int) { + if req.Header.Get(model.RetryCountHeader) == "" { + req.Header.Set(model.RetryCountHeader, strconv.Itoa(retryCount)) + } +} + func (c *Client) getRetryAfter(v model.Response) int64 { header := v.GetHeader() retryAfter := header[model.RetryAfterHeader] diff --git a/arkruntime/environment_work.go b/arkruntime/environment_work.go index 962d341..9c055ac 100644 --- a/arkruntime/environment_work.go +++ b/arkruntime/environment_work.go @@ -166,19 +166,16 @@ func (c *Client) doControlPlaneRequest( v model.Response, setters ...requestOption, ) error { - return utils.Retry( + return utils.RetryWithAttempt( ctx, - utils.RetryPolicy{ - MaxAttempts: c.config.RetryTimes, - InitialBackoff: model.ErrorRetryBaseDelay, - MaxBackoff: model.ErrorRetryMaxDelay, - }, + defaultRetryPolicy(c.config.RetryTimes), func() bool { return true }, - func() error { + func(retryCount int) error { req, reqErr := c.newRequest(ctx, method, u, "", "", setters...) if reqErr != nil { return reqErr } + setRetryCountHeader(req, retryCount) return c.sendControlPlaneRequest(req, v) }, nil, diff --git a/arkruntime/model/common.go b/arkruntime/model/common.go index f5d781c..7f6b54e 100644 --- a/arkruntime/model/common.go +++ b/arkruntime/model/common.go @@ -13,6 +13,9 @@ const ( ClientRequestHeader = "X-Client-Request-Id" ServerRequestHeader = "X-Request-Id" RetryAfterHeader = "Retry-After" + RetryAfterMSHeader = "Retry-After-Ms" + RetryCountHeader = "X-Stainless-Retry-Count" + ShouldRetryHeader = "X-Should-Retry" DefaultMandatoryRefreshTimeout = 10 * 60 // 10 min DefaultAdvisoryRefreshTimeout = 30 * 60 // 30 min @@ -23,6 +26,7 @@ const ( ErrorRetryBaseDelay = 500 * time.Millisecond ErrorRetryMaxDelay = 8 * time.Second + MaxServerRetryDelay = 60 * time.Second ) type PromptTokensDetail struct { diff --git a/arkruntime/model/error.go b/arkruntime/model/error.go index d575199..2ed20af 100644 --- a/arkruntime/model/error.go +++ b/arkruntime/model/error.go @@ -7,22 +7,25 @@ import ( "encoding/json" "errors" "fmt" + "net/http" ) type APIError struct { - Code string `json:"code,omitempty"` - Message string `json:"message"` - Param *string `json:"param,omitempty"` - Type string `json:"type"` - HTTPStatusCode int `json:"-"` - RequestId string `json:"request_id"` + Code string `json:"code,omitempty"` + Message string `json:"message"` + Param *string `json:"param,omitempty"` + Type string `json:"type"` + HTTPStatusCode int `json:"-"` + RequestId string `json:"request_id"` + ResponseHeader http.Header `json:"-"` } // RequestError provides information about generic request errors. type RequestError struct { HTTPStatusCode int Err error - RequestId string `json:"request_id"` + RequestId string `json:"request_id"` + ResponseHeader http.Header `json:"-"` } func NewRequestError(httpStatusCode int, rawErr error, requestID string) *RequestError { @@ -50,6 +53,20 @@ func (e *RequestError) Unwrap() error { return e.Err } +func (e *APIError) GetHeader() http.Header { + if e == nil { + return nil + } + return e.ResponseHeader +} + +func (e *RequestError) GetHeader() http.Header { + if e == nil { + return nil + } + return e.ResponseHeader +} + var ( ErrTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages") ErrChatCompletionInvalidModel = errors.New("this model is not supported with this method, please use CreateCompletion client method instead") //nolint:lll diff --git a/arkruntime/retry_test.go b/arkruntime/retry_test.go new file mode 100644 index 0000000..649f32d --- /dev/null +++ b/arkruntime/retry_test.go @@ -0,0 +1,169 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +package arkruntime + +import ( + "context" + "errors" + "fmt" + "math" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/volcengine/ark-runtime-go/arkruntime/model" +) + +func TestRetryAfterPrefersMilliseconds(t *testing.T) { + err := &model.APIError{ + ResponseHeader: http.Header{ + model.RetryAfterMSHeader: []string{"125.5"}, + model.RetryAfterHeader: []string{"9"}, + }, + } + delay, ok := retryAfter(err) + if !ok || delay != 125500*time.Microsecond { + t.Fatalf("retryAfter() = %s, %v; want 125.5ms, true", delay, ok) + } +} + +func TestRetryAfterParsing(t *testing.T) { + future := time.Now().Add(2 * time.Second).UTC().Format(http.TimeFormat) + tests := []struct { + name string + header http.Header + wantOK bool + check func(time.Duration) bool + }{ + { + name: "invalid milliseconds falls back to seconds", + header: http.Header{model.RetryAfterMSHeader: []string{"bad"}, model.RetryAfterHeader: []string{"0.25"}}, + wantOK: true, + check: func(got time.Duration) bool { return got == 250*time.Millisecond }, + }, + { + name: "http date", + header: http.Header{model.RetryAfterHeader: []string{future}}, + wantOK: true, + check: func(got time.Duration) bool { return got > 0 && got <= 2*time.Second }, + }, + { + name: "negative", + header: http.Header{model.RetryAfterHeader: []string{"-5"}}, + wantOK: true, + check: func(got time.Duration) bool { return got == -5*time.Second }, + }, + { + name: "non finite", + header: http.Header{model.RetryAfterHeader: []string{fmt.Sprint(math.Inf(1))}}, + wantOK: false, + check: func(time.Duration) bool { return true }, + }, + { + name: "missing", + header: http.Header{}, + wantOK: false, + check: func(time.Duration) bool { return true }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + delay, ok := retryAfter(&model.APIError{ResponseHeader: test.header}) + if ok != test.wantOK || !test.check(delay) { + t.Fatalf("retryAfter() = %s, %v", delay, ok) + } + }) + } +} + +func TestNeedRetryErrorHonorsServerOverride(t *testing.T) { + err := &model.RequestError{ + HTTPStatusCode: http.StatusBadRequest, + Err: errors.New("bad request"), + ResponseHeader: http.Header{model.ShouldRetryHeader: []string{"true"}}, + } + if !needRetryError(err) { + t.Fatal("expected X-Should-Retry=true to force a retry") + } + err.ResponseHeader.Set(model.ShouldRetryHeader, "false") + err.HTTPStatusCode = http.StatusInternalServerError + if needRetryError(err) { + t.Fatal("expected X-Should-Retry=false to prevent a retry") + } +} + +func TestDoPreservesCustomRetryCount(t *testing.T) { + var retryCounts []string + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + retryCounts = append(retryCounts, r.Header.Get(model.RetryCountHeader)) + calls++ + if calls == 1 { + w.Header().Set(model.RetryAfterMSHeader, "1") + http.Error(w, "retry", http.StatusTooManyRequests) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{}`)) + })) + defer server.Close() + + client := NewClientWithApiKey("placeholder", WithHTTPClient(server.Client()), WithRetryTimes(1)) + err := client.Do( + context.Background(), http.MethodGet, server.URL, "", "", nil, + WithCustomHeader(model.RetryCountHeader, "custom"), + ) + if err != nil { + t.Fatalf("Do() error = %v", err) + } + if want := "[custom custom]"; fmt.Sprint(retryCounts) != want { + t.Fatalf("retry counts = %v, want %s", retryCounts, want) + } +} + +func TestRetryableStatuses(t *testing.T) { + for _, status := range []int{ + http.StatusRequestTimeout, + http.StatusConflict, + http.StatusTooManyRequests, + http.StatusInternalServerError, + } { + if !isRetryableStatus(status) { + t.Fatalf("status %d should be retryable", status) + } + } + if isRetryableStatus(http.StatusBadRequest) { + t.Fatal("400 bad request should not be retried without an explicit server override") + } +} + +func TestDoUsesServerDelayAndIncrementsRetryCount(t *testing.T) { + var retryCounts []string + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + retryCounts = append(retryCounts, r.Header.Get(model.RetryCountHeader)) + calls++ + if calls == 1 { + w.Header().Set(model.RetryAfterMSHeader, "1") + http.Error(w, "rate limited", http.StatusTooManyRequests) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{}`)) + })) + defer server.Close() + + client := NewClientWithApiKey( + "placeholder", + WithHTTPClient(server.Client()), + WithRetryTimes(2), + ) + if err := client.Do(context.Background(), http.MethodGet, server.URL, "", "", nil); err != nil { + t.Fatalf("Do() error = %v", err) + } + if len(retryCounts) != 2 || retryCounts[0] != "0" || retryCounts[1] != "1" { + t.Fatalf("retry counts = %v, want [0 1]", retryCounts) + } +} diff --git a/arkruntime/selfhosted/mcp/mcp.go b/arkruntime/selfhosted/mcp/mcp.go new file mode 100644 index 0000000..840694b --- /dev/null +++ b/arkruntime/selfhosted/mcp/mcp.go @@ -0,0 +1,515 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +// Package mcp 提供与具体 MCP SDK 无关的 self-hosted MCP 工具适配能力。 +package mcp + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "strings" + "unicode/utf8" + + "github.com/volcengine/ark-runtime-go/arkruntime/model/agent" + "github.com/volcengine/ark-runtime-go/arkruntime/tools/agenttoolset" +) + +var supportedImageMIMETypes = map[string]bool{ + "image/gif": true, + "image/jpeg": true, + "image/png": true, + "image/webp": true, +} + +const ( + compatibilityDescriptionPrefix = "\n\nMCP input constraints (JSON Schema): " + maxCustomToolDescriptionRunes = 10_000 + objectSchemaType = "object" +) + +var ignoredTopLevelSchemaKeywords = map[string]bool{ + "$anchor": true, + "$comment": true, + "$dynamicAnchor": true, + "$id": true, + "$schema": true, + "title": true, +} + +// Client 是 self-hosted worker 调用 MCP Server 所需的最小接口。 +type Client interface { + CallTool(ctx context.Context, name string, arguments map[string]any) (*CallToolResult, error) +} + +// ToolDefinition 描述一个 MCP Tool 及其输入 JSON Schema。 +type ToolDefinition struct { + Name string + Description string + InputSchema any +} + +// CallToolResult 是 MCP Tool 调用结果的协议无关表示。 +type CallToolResult struct { + Content []Content + StructuredContent any + IsError bool +} + +// Content 是 MCP Tool 返回的单个内容块。 +type Content struct { + Type string + Text string + MIMEType string + Data []byte + Resource *Resource +} + +// Resource 是 MCP embedded resource 的协议无关表示。 +type Resource struct { + URI string + MIMEType string + Text string + Blob []byte +} + +// NewTool 把 MCP Tool 定义与 Client 包装成 worker Custom Tool。 +func NewTool(tool ToolDefinition, client Client) (agenttoolset.Tool, error) { + if client == nil { + return nil, errors.New("mcp client is required") + } + if tool.Name == "" { + return nil, errors.New("mcp tool name is required") + } + return &runnableTool{tool: tool, client: client}, nil +} + +// NewTools 批量包装 MCP Tools。 +func NewTools(tools []ToolDefinition, client Client) (map[string]agenttoolset.Tool, error) { + out := make(map[string]agenttoolset.Tool, len(tools)) + for _, tool := range tools { + wrapped, err := NewTool(tool, client) + if err != nil { + return nil, err + } + if _, exists := out[wrapped.Name()]; exists { + return nil, fmt.Errorf("duplicate mcp tool name %q", wrapped.Name()) + } + out[wrapped.Name()] = wrapped + } + return out, nil +} + +// CustomToolItem 把 MCP Tool 定义转换成创建或更新 Agent 使用的 custom ToolItem。 +func CustomToolItem(tool ToolDefinition) (agent.ToolItem, error) { + if tool.Name == "" { + return agent.ToolItem{}, errors.New("mcp tool name is required") + } + schema, constraints, err := customToolInputSchema(tool.InputSchema) + if err != nil { + return agent.ToolItem{}, fmt.Errorf("mcp tool %s input schema: %w", tool.Name, err) + } + description := tool.Description + if description == "" { + description = tool.Name + } + if constraints != "" { + description += compatibilityDescriptionPrefix + constraints + } + if utf8.RuneCountInString(description) > maxCustomToolDescriptionRunes { + return agent.ToolItem{}, fmt.Errorf( + "mcp tool %s description exceeds %d characters after adding input constraints", + tool.Name, + maxCustomToolDescriptionRunes, + ) + } + return agent.ToolItem{ + Type: "custom", + Name: agent.NewOptString(tool.Name), + Description: agent.NewOptString(description), + InputSchema: agent.NewOptCustomToolInputSchema(schema), + }, nil +} + +// CustomToolItems 批量生成 Agent custom tool 声明。 +func CustomToolItems(tools []ToolDefinition) ([]agent.ToolItem, error) { + out := make([]agent.ToolItem, 0, len(tools)) + seen := make(map[string]struct{}, len(tools)) + for _, tool := range tools { + item, err := CustomToolItem(tool) + if err != nil { + return nil, err + } + if _, exists := seen[tool.Name]; exists { + return nil, fmt.Errorf("duplicate mcp tool name %q", tool.Name) + } + seen[tool.Name] = struct{}{} + out = append(out, item) + } + return out, nil +} + +type runnableTool struct { + tool ToolDefinition + client Client +} + +func (t *runnableTool) Name() string { return t.tool.Name } + +func (t *runnableTool) Execute(ctx context.Context, input json.RawMessage) agenttoolset.Result { + if len(input) == 0 { + input = json.RawMessage("{}") + } + var arguments map[string]any + if err := json.Unmarshal(input, &arguments); err != nil { + return errorResult(fmt.Sprintf("mcp tool %s: invalid input: %v", t.tool.Name, err)) + } + result, err := t.client.CallTool(ctx, t.tool.Name, arguments) + if err != nil { + return errorResult(fmt.Sprintf("mcp tool %s: %v", t.tool.Name, err)) + } + return convertCallToolResult(result, t.tool.Name) +} + +func customToolInputSchema(value any) (agent.CustomToolInputSchema, string, error) { + if value == nil { + value = map[string]any{"type": objectSchemaType} + } + raw, err := json.Marshal(value) + if err != nil { + return agent.CustomToolInputSchema{}, "", err + } + var object map[string]any + if err := json.Unmarshal(raw, &object); err != nil { + return agent.CustomToolInputSchema{}, "", errors.New("schema must be an object") + } + if object == nil { + object = make(map[string]any) + } + + typeName := objectSchemaType + if rawType, ok := object["type"]; ok { + if rawType == nil { + typeName = objectSchemaType + } else if parsedType, valid := rawType.(string); valid { + typeName = parsedType + } else { + return agent.CustomToolInputSchema{}, "", errors.New("type must be a string") + } + } + if typeName != objectSchemaType { + return agent.CustomToolInputSchema{}, "", errors.New("top-level type must be \"object\"") + } + + var properties map[string]any + if rawProperties, ok := object["properties"]; ok { + var valid bool + properties, valid = rawProperties.(map[string]any) + if !valid { + return agent.CustomToolInputSchema{}, "", errors.New("properties must be an object") + } + } + required, err := requiredNames(object["required"]) + if err != nil { + return agent.CustomToolInputSchema{}, "", err + } + + constraints := unsupportedTopLevelConstraints(object) + resolvedProperties, unresolved := resolveLocalReferences(properties, object) + if unresolved { + constraints["properties"] = properties + } + removeUnreferencedDefinitions(constraints) + constraintJSON, err := json.Marshal(constraints) + if err != nil { + return agent.CustomToolInputSchema{}, "", err + } + if len(constraints) == 0 { + constraintJSON = nil + } + + schema := agent.CustomToolInputSchema{ + Type: agent.NewOptString(objectSchemaType), + Required: required, + } + if resolvedProperties != nil { + converted := make(agent.CustomToolInputSchemaProperties, len(resolvedProperties)) + for name, property := range resolvedProperties { + rawProperty, marshalErr := json.Marshal(property) + if marshalErr != nil { + return agent.CustomToolInputSchema{}, "", marshalErr + } + converted[name] = rawProperty + } + schema.Properties = agent.NewOptCustomToolInputSchemaProperties(converted) + } + return schema, string(constraintJSON), nil +} + +func requiredNames(value any) ([]string, error) { + if value == nil { + return nil, nil + } + values, ok := value.([]any) + if !ok { + return nil, errors.New("required must be an array of strings") + } + result := make([]string, 0, len(values)) + for _, value := range values { + name, ok := value.(string) + if !ok { + return nil, errors.New("required must be an array of strings") + } + result = append(result, name) + } + return result, nil +} + +func unsupportedTopLevelConstraints(schema map[string]any) map[string]any { + constraints := make(map[string]any) + for key, value := range schema { + if key == "type" || key == "properties" || key == "required" || ignoredTopLevelSchemaKeywords[key] { + continue + } + constraints[key] = value + } + return constraints +} + +func removeUnreferencedDefinitions(constraints map[string]any) { + definitionReferences := map[string]string{ + "$defs": "#/$defs", + "definitions": "#/definitions", + } + for definitionKey, referencePrefix := range definitionReferences { + if _, exists := constraints[definitionKey]; !exists { + continue + } + referenced := false + for key, value := range constraints { + if key != definitionKey && containsReference(value, referencePrefix) { + referenced = true + break + } + } + if !referenced { + delete(constraints, definitionKey) + } + } +} + +func containsReference(value any, prefix string) bool { + switch typed := value.(type) { + case map[string]any: + if reference, ok := typed["$ref"].(string); ok && + (reference == prefix || strings.HasPrefix(reference, prefix+"/")) { + return true + } + for _, item := range typed { + if containsReference(item, prefix) { + return true + } + } + case []any: + for _, item := range typed { + if containsReference(item, prefix) { + return true + } + } + } + return false +} + +func resolveLocalReferences(properties map[string]any, root map[string]any) (map[string]any, bool) { + if properties == nil { + return nil, false + } + resolved, unresolved := resolveSchemaValue(properties, root, make(map[string]bool)) + return resolved.(map[string]any), unresolved +} + +func resolveSchemaValue(value any, root map[string]any, resolving map[string]bool) (any, bool) { + switch typed := value.(type) { + case map[string]any: + if reference, ok := typed["$ref"].(string); ok { + target, found := resolveJSONPointer(root, reference) + if found && !resolving[reference] { + resolving[reference] = true + resolvedTarget, unresolved := resolveSchemaValue(target, root, resolving) + delete(resolving, reference) + if targetMap, valid := resolvedTarget.(map[string]any); valid { + merged := make(map[string]any, len(targetMap)+len(typed)-1) + for key, item := range targetMap { + merged[key] = item + } + for key, item := range typed { + if key != "$ref" { + merged[key] = item + } + } + resolved, mergedUnresolved := resolveSchemaValue(merged, root, resolving) + return resolved, unresolved || mergedUnresolved + } + return mapWithoutReference(typed, root, resolving, true) + } + return mapWithoutReference(typed, root, resolving, true) + } + out := make(map[string]any, len(typed)) + unresolved := false + for key, item := range typed { + resolved, itemUnresolved := resolveSchemaValue(item, root, resolving) + out[key] = resolved + unresolved = unresolved || itemUnresolved + } + return out, unresolved + case []any: + out := make([]any, len(typed)) + unresolved := false + for i, item := range typed { + resolved, itemUnresolved := resolveSchemaValue(item, root, resolving) + out[i] = resolved + unresolved = unresolved || itemUnresolved + } + return out, unresolved + default: + return value, false + } +} + +func mapWithoutReference(value map[string]any, root map[string]any, resolving map[string]bool, unresolved bool) (any, bool) { + out := make(map[string]any, len(value)-1) + for key, item := range value { + if key == "$ref" { + continue + } + resolved, itemUnresolved := resolveSchemaValue(item, root, resolving) + out[key] = resolved + unresolved = unresolved || itemUnresolved + } + return out, unresolved +} + +func resolveJSONPointer(root map[string]any, reference string) (any, bool) { + if !strings.HasPrefix(reference, "#/") { + return nil, false + } + var current any = root + for _, pointerPart := range strings.Split(strings.TrimPrefix(reference, "#/"), "/") { + pointerPart = strings.ReplaceAll(strings.ReplaceAll(pointerPart, "~1", "/"), "~0", "~") + object, ok := current.(map[string]any) + if !ok { + return nil, false + } + current, ok = object[pointerPart] + if !ok { + return nil, false + } + } + return current, true +} + +// ConvertCallToolResult 把协议无关的 MCP 结果转换成 worker Tool Result。 +func ConvertCallToolResult(result *CallToolResult) agenttoolset.Result { + return convertCallToolResult(result, "") +} + +func convertCallToolResult(result *CallToolResult, toolName string) agenttoolset.Result { + if result == nil { + return errorResult("mcp tool returned no result") + } + blocks := make([]agenttoolset.ContentBlock, 0, len(result.Content)) + for _, content := range result.Content { + block, err := contentBlock(content) + if err != nil { + return errorResult(err.Error()) + } + blocks = append(blocks, block) + } + if len(blocks) == 0 && result.StructuredContent != nil { + raw, err := json.Marshal(result.StructuredContent) + if err != nil { + return errorResult(fmt.Sprintf("serialize mcp structured content: %v", err)) + } + blocks = append(blocks, agenttoolset.ContentBlock{Type: "text", Text: string(raw)}) + } + if len(blocks) == 0 && result.IsError { + message := "mcp tool returned an error but returned no content" + if toolName != "" { + message = fmt.Sprintf("mcp tool %q reported an error but returned no content", toolName) + } + blocks = append(blocks, agenttoolset.ContentBlock{Type: "text", Text: message}) + } + return agenttoolset.Result{Content: blocks, IsError: result.IsError} +} + +func contentBlock(content Content) (agenttoolset.ContentBlock, error) { + switch content.Type { + case "text": + return agenttoolset.ContentBlock{Type: "text", Text: content.Text}, nil + case "image": + if !supportedImageMIMETypes[content.MIMEType] { + return agenttoolset.ContentBlock{}, fmt.Errorf("unsupported image MIME type %q", content.MIMEType) + } + return base64Block("image", content.MIMEType, content.Data), nil + case "resource": + return resourceBlock(content.Resource) + case "audio", "resource_link": + return agenttoolset.ContentBlock{}, fmt.Errorf("unsupported MCP content type %s", content.Type) + default: + return agenttoolset.ContentBlock{}, fmt.Errorf("unsupported MCP content type %s", content.Type) + } +} + +func resourceBlock(resource *Resource) (agenttoolset.ContentBlock, error) { + if resource == nil { + return agenttoolset.ContentBlock{}, errors.New("embedded MCP resource has no content") + } + if supportedImageMIMETypes[resource.MIMEType] { + if resource.Blob == nil { + return agenttoolset.ContentBlock{}, errors.New("image resource must contain blob data") + } + return base64Block("image", resource.MIMEType, resource.Blob), nil + } + if resource.MIMEType == "application/pdf" { + if resource.Blob == nil { + return agenttoolset.ContentBlock{}, errors.New("PDF resource must contain blob data") + } + return base64Block("document", resource.MIMEType, resource.Blob), nil + } + if resource.MIMEType == "" || strings.HasPrefix(resource.MIMEType, "text/") { + text := resource.Text + if resource.Blob != nil { + // Text blobs are interpreted as UTF-8, matching the other SDK adapters. + text = string(resource.Blob) + } + return agenttoolset.ContentBlock{ + Type: "document", + Source: map[string]any{ + "type": "text", + "media_type": "text/plain", + "data": text, + }, + }, nil + } + return agenttoolset.ContentBlock{}, fmt.Errorf("unsupported resource MIME type %q", resource.MIMEType) +} + +func base64Block(blockType, mimeType string, data []byte) agenttoolset.ContentBlock { + return agenttoolset.ContentBlock{ + Type: blockType, + Source: map[string]any{ + "type": "base64", + "media_type": mimeType, + "data": base64.StdEncoding.EncodeToString(data), + }, + } +} + +func errorResult(message string) agenttoolset.Result { + return agenttoolset.Result{ + Content: []agenttoolset.ContentBlock{{Type: "text", Text: message}}, + IsError: true, + } +} diff --git a/arkruntime/selfhosted/mcp/mcp_test.go b/arkruntime/selfhosted/mcp/mcp_test.go new file mode 100644 index 0000000..e1687a8 --- /dev/null +++ b/arkruntime/selfhosted/mcp/mcp_test.go @@ -0,0 +1,229 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +package mcp + +import ( + "context" + "encoding/json" + "strings" + "testing" +) + +type fakeClient struct { + name string + arguments map[string]any + result *CallToolResult +} + +func (f *fakeClient) CallTool(_ context.Context, name string, arguments map[string]any) (*CallToolResult, error) { + f.name = name + f.arguments = arguments + return f.result, nil +} + +func TestCustomToolItemAdaptsSchemaToCurrentAgentContract(t *testing.T) { + item, err := CustomToolItem(ToolDefinition{ + Name: "lookup_order", + Description: "Lookup an order", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "order_id": map[string]any{"$ref": "#/$defs/order_id"}, + }, + "required": []string{"order_id"}, + "additionalProperties": false, + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$defs": map[string]any{ + "order_id": map[string]any{"type": "string", "minLength": 1}, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + raw, err := item.MarshalJSON() + if err != nil { + t.Fatal(err) + } + var itemJSON map[string]any + if err := json.Unmarshal(raw, &itemJSON); err != nil { + t.Fatal(err) + } + inputSchema, ok := itemJSON["input_schema"].(map[string]any) + if !ok || len(inputSchema) != 3 { + t.Fatalf("input schema does not match current agent contract: %s", raw) + } + if _, exists := inputSchema["additionalProperties"]; exists { + t.Fatalf("unsupported top-level constraint leaked into input schema: %s", raw) + } + properties := inputSchema["properties"].(map[string]any) + orderID := properties["order_id"].(map[string]any) + if orderID["type"] != "string" || orderID["minLength"] != float64(1) || orderID["$ref"] != nil { + t.Fatalf("local schema reference was not inlined: %s", raw) + } + description := itemJSON["description"].(string) + if !strings.Contains(description, `MCP input constraints (JSON Schema): {"additionalProperties":false}`) || + strings.Contains(description, `"$defs"`) || + strings.Contains(description, `"$schema"`) { + t.Fatalf("unexpected compatibility description: %s", description) + } +} + +func TestCustomToolItemRetainsReferencedDefinitions(t *testing.T) { + item, err := CustomToolItem(ToolDefinition{ + Name: "lookup", + InputSchema: map[string]any{ + "type": "object", + "allOf": []any{map[string]any{"$ref": "#/$defs/constraint"}}, + "$defs": map[string]any{ + "constraint": map[string]any{"additionalProperties": false}, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(item.Description.Value, `"$defs"`) || + !strings.Contains(item.Description.Value, `"$ref":"#/$defs/constraint"`) { + t.Fatalf("referenced definitions are absent from description: %s", item.Description.Value) + } +} + +func TestCustomToolItemDescribesUnresolvedReferences(t *testing.T) { + item, err := CustomToolItem(ToolDefinition{ + Name: "lookup", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"order": map[string]any{"$ref": "https://example.com/order.json"}}, + }, + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(item.Description.Value, `"$ref":"https://example.com/order.json"`) { + t.Fatalf("unresolved reference is absent from description: %s", item.Description.Value) + } + if got := string(item.InputSchema.Value.Properties.Value["order"]); got != `{}` { + t.Fatalf("unresolved reference leaked into input schema: %s", got) + } +} + +func TestCustomToolItemRejectsOversizedCompatibilityDescription(t *testing.T) { + _, err := CustomToolItem(ToolDefinition{ + Name: "large", + Description: strings.Repeat("a", maxCustomToolDescriptionRunes), + InputSchema: map[string]any{"type": "object", "additionalProperties": false}, + }) + if err == nil || !strings.Contains(err.Error(), "description exceeds") { + t.Fatalf("expected description length error, got %v", err) + } +} + +func TestNewToolCallsClient(t *testing.T) { + client := &fakeClient{result: &CallToolResult{ + Content: []Content{{Type: "text", Text: "echo: hello"}}, + }} + tool, err := NewTool(ToolDefinition{Name: "echo"}, client) + if err != nil { + t.Fatal(err) + } + result := tool.Execute(context.Background(), json.RawMessage(`{"text":"hello"}`)) + if result.IsError || len(result.Content) != 1 || result.Content[0].Text != "echo: hello" { + t.Fatalf("unexpected tool result: %+v", result) + } + if client.name != "echo" || client.arguments["text"] != "hello" { + t.Fatalf("unexpected client call: name=%q arguments=%v", client.name, client.arguments) + } +} + +func TestNewToolAddsNameToEmptyError(t *testing.T) { + tool, err := NewTool(ToolDefinition{Name: "lookup"}, &fakeClient{ + result: &CallToolResult{IsError: true}, + }) + if err != nil { + t.Fatal(err) + } + result := tool.Execute(context.Background(), json.RawMessage(`{}`)) + if !result.IsError || len(result.Content) != 1 || + result.Content[0].Text != `mcp tool "lookup" reported an error but returned no content` { + t.Fatalf("unexpected result: %+v", result) + } +} + +func TestConvertCallToolResultPreservesRichContent(t *testing.T) { + result := ConvertCallToolResult(&CallToolResult{Content: []Content{ + {Type: "image", MIMEType: "image/png", Data: []byte("image")}, + {Type: "resource", Resource: &Resource{ + URI: "file:///result.txt", + MIMEType: "text/plain", + Text: "resource text", + }}, + }}) + if result.IsError || len(result.Content) != 2 { + t.Fatalf("unexpected result: %+v", result) + } + imageSource, ok := result.Content[0].Source.(map[string]any) + if !ok || imageSource["data"] != "aW1hZ2U=" { + t.Fatalf("unexpected image source: %#v", result.Content[0].Source) + } + documentSource, ok := result.Content[1].Source.(map[string]any) + if !ok || documentSource["data"] != "resource text" { + t.Fatalf("unexpected document source: %#v", result.Content[1].Source) + } +} + +func TestConvertTextResourceNormalizesMIMEType(t *testing.T) { + result := ConvertCallToolResult(&CallToolResult{Content: []Content{{ + Type: "resource", + Resource: &Resource{ + MIMEType: "text/html", + Text: "

hello

", + }, + }}}) + if result.IsError || len(result.Content) != 1 { + t.Fatalf("unexpected result: %+v", result) + } + source, ok := result.Content[0].Source.(map[string]any) + if !ok || source["media_type"] != "text/plain" { + t.Fatalf("unexpected result: %+v", result) + } +} + +func TestConvertStructuredAndErrorResults(t *testing.T) { + structured := ConvertCallToolResult(&CallToolResult{StructuredContent: map[string]any{"status": "ok"}}) + if structured.IsError || len(structured.Content) != 1 || structured.Content[0].Text != `{"status":"ok"}` { + t.Fatalf("unexpected structured result: %+v", structured) + } + + failed := ConvertCallToolResult(&CallToolResult{ + Content: []Content{{Type: "text", Text: "not found"}}, + IsError: true, + }) + if !failed.IsError || len(failed.Content) != 1 || failed.Content[0].Text != "not found" { + t.Fatalf("unexpected error result: %+v", failed) + } + + emptyFailure := ConvertCallToolResult(&CallToolResult{IsError: true}) + if !emptyFailure.IsError || len(emptyFailure.Content) != 1 || + emptyFailure.Content[0].Text != "mcp tool returned an error but returned no content" { + t.Fatalf("unexpected empty error result: %+v", emptyFailure) + } +} + +func TestConvertCallToolResultDoesNotExposeResourceURI(t *testing.T) { + const secretURI = "https://example.com/file?signature=secret" + result := ConvertCallToolResult(&CallToolResult{Content: []Content{{ + Type: "resource", + Resource: &Resource{ + URI: secretURI, + MIMEType: "image/png", + }, + }}}) + if !result.IsError || len(result.Content) != 1 { + t.Fatalf("unexpected result: %+v", result) + } + if strings.Contains(result.Content[0].Text, secretURI) || strings.Contains(result.Content[0].Text, "secret") { + t.Fatalf("resource URI leaked into error: %q", result.Content[0].Text) + } +} diff --git a/arkruntime/selfhosted/session_tool_runner.go b/arkruntime/selfhosted/session_tool_runner.go index 2ad6501..2eb9e96 100644 --- a/arkruntime/selfhosted/session_tool_runner.go +++ b/arkruntime/selfhosted/session_tool_runner.go @@ -1080,6 +1080,9 @@ func runnerContentBlocks(blocks []toolset.ContentBlock) []ContentBlock { Text: block.Text, MediaType: block.MediaType, Data: block.Data, + Source: block.Source, + Title: block.Title, + Context: block.Context, }) } return out diff --git a/arkruntime/selfhosted/session_tool_runner_test.go b/arkruntime/selfhosted/session_tool_runner_test.go index 84d6d7c..3892daa 100644 --- a/arkruntime/selfhosted/session_tool_runner_test.go +++ b/arkruntime/selfhosted/session_tool_runner_test.go @@ -140,6 +140,36 @@ func TestSessionToolRunnerConvertsToolPanicToErrorResult(t *testing.T) { } } +func TestRunnerContentBlocksPreservesStructuredSource(t *testing.T) { + source := map[string]any{ + "type": "base64", + "media_type": "image/png", + "data": "aW1hZ2U=", + } + blocks := runnerContentBlocks([]toolset.ContentBlock{{ + Type: "image", + Source: source, + Title: "result", + Context: "generated by MCP", + }}) + if len(blocks) != 1 || blocks[0].Type != "image" || blocks[0].Title != "result" || blocks[0].Context != "generated by MCP" { + t.Fatalf("unexpected blocks: %+v", blocks) + } + if got, ok := blocks[0].Source.(map[string]any); !ok || got["data"] != "aW1hZ2U=" { + t.Fatalf("unexpected source: %#v", blocks[0].Source) + } +} + +func TestContentBlockMarshalPreservesEmptyText(t *testing.T) { + raw, err := json.Marshal(ContentBlock{Type: "text"}) + if err != nil { + t.Fatal(err) + } + if string(raw) != `{"type":"text","text":""}` { + t.Fatalf("unexpected text block JSON: %s", raw) + } +} + func TestSessionToolRunnerReplayDoesNotResetIdleDeadline(t *testing.T) { maxIdle := time.Second runner := NewSessionToolRunner(context.Background(), &runnerTestAPI{}, "session-id", SessionToolRunnerOptions{ diff --git a/arkruntime/selfhosted/types.go b/arkruntime/selfhosted/types.go index 95f60d5..63fdb05 100644 --- a/arkruntime/selfhosted/types.go +++ b/arkruntime/selfhosted/types.go @@ -349,6 +349,24 @@ type ContentBlock struct { Text string `json:"text,omitempty"` MediaType string `json:"media_type,omitempty"` Data []byte `json:"data,omitempty"` + Source any `json:"source,omitempty"` + Title string `json:"title,omitempty"` + Context string `json:"context,omitempty"` +} + +// MarshalJSON 在文本块中保留必需的空 text 字段。 +func (b ContentBlock) MarshalJSON() ([]byte, error) { + type contentBlockAlias ContentBlock + if b.Type != "text" || b.Text != "" { + return json.Marshal(contentBlockAlias(b)) + } + return json.Marshal(struct { + contentBlockAlias + Text string `json:"text"` + }{ + contentBlockAlias: contentBlockAlias(b), + Text: b.Text, + }) } // RawJSON 保存未解释的 JSON 对象,兼容 wire 上以字符串承载 raw JSON。 diff --git a/arkruntime/toolset/types.go b/arkruntime/toolset/types.go index 8d6326f..d9fb03c 100644 --- a/arkruntime/toolset/types.go +++ b/arkruntime/toolset/types.go @@ -17,6 +17,9 @@ type ContentBlock struct { Text string `json:"text,omitempty"` MediaType string `json:"media_type,omitempty"` Data []byte `json:"data,omitempty"` + Source any `json:"source,omitempty"` + Title string `json:"title,omitempty"` + Context string `json:"context,omitempty"` } // Result 是一次工具执行结果。 diff --git a/arkruntime/utils/retry.go b/arkruntime/utils/retry.go index 0c9baa8..6753ae8 100644 --- a/arkruntime/utils/retry.go +++ b/arkruntime/utils/retry.go @@ -14,6 +14,8 @@ type RetryPolicy struct { MaxAttempts int InitialBackoff time.Duration MaxBackoff time.Duration + MaxRetryAfter time.Duration + RetryAfter func(error) (time.Duration, bool) } func Retry(ctx context.Context, @@ -21,10 +23,30 @@ func Retry(ctx context.Context, isNeedRetry func() bool, doFunc func() error, overRetryLimitError error, isNeedRetryError func(error) bool, +) error { + return retry(ctx, rp, isNeedRetry, func(_ int) error { + return doFunc() + }, overRetryLimitError, isNeedRetryError) +} + +func RetryWithAttempt(ctx context.Context, + rp RetryPolicy, + isNeedRetry func() bool, + doFunc func(int) error, overRetryLimitError error, + isNeedRetryError func(error) bool, +) error { + return retry(ctx, rp, isNeedRetry, doFunc, overRetryLimitError, isNeedRetryError) +} + +func retry(ctx context.Context, + rp RetryPolicy, + isNeedRetry func() bool, + doFunc func(int) error, overRetryLimitError error, + isNeedRetryError func(error) bool, ) error { var err error for numRetriesSincePushback := 0; numRetriesSincePushback <= rp.MaxAttempts; numRetriesSincePushback++ { - err = doFunc() + err = doFunc(numRetriesSincePushback) // no error: just return on this try if err == nil { @@ -45,10 +67,7 @@ func Retry(ctx context.Context, if numRetriesSincePushback == rp.MaxAttempts { break } - nbRetries := numRetriesSincePushback + 1 - sleepSeconds := math.Min(rp.InitialBackoff.Seconds()*math.Pow(2.0, float64(nbRetries)), rp.MaxBackoff.Seconds()) - jitter := 1.0 - 0.25*rand.Float64() - dur := time.Duration(sleepSeconds*jitter) * time.Second + dur := retryDelay(rp, numRetriesSincePushback, err) t := time.NewTimer(dur) select { @@ -65,3 +84,24 @@ func Retry(ctx context.Context, } return overRetryLimitError } + +func retryDelay(rp RetryPolicy, retryCount int, err error) time.Duration { + if rp.RetryAfter != nil { + if delay, ok := rp.RetryAfter(err); ok { + maxRetryAfter := rp.MaxRetryAfter + if maxRetryAfter <= 0 { + maxRetryAfter = 60 * time.Second + } + if delay > 0 && delay <= maxRetryAfter { + return delay + } + } + } + + delay := time.Duration(float64(rp.InitialBackoff) * math.Pow(2.0, float64(retryCount))) + if delay > rp.MaxBackoff { + delay = rp.MaxBackoff + } + jitter := 1.0 - 0.25*rand.Float64() + return time.Duration(float64(delay) * jitter) +} diff --git a/arkruntime/utils/retry_test.go b/arkruntime/utils/retry_test.go new file mode 100644 index 0000000..e593754 --- /dev/null +++ b/arkruntime/utils/retry_test.go @@ -0,0 +1,86 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +package utils + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestRetryDelayPreservesSubsecondBackoff(t *testing.T) { + policy := RetryPolicy{InitialBackoff: 500 * time.Millisecond, MaxBackoff: 8 * time.Second} + for i, bounds := range [][2]time.Duration{ + {375 * time.Millisecond, 500 * time.Millisecond}, + {750 * time.Millisecond, time.Second}, + {6 * time.Second, 8 * time.Second}, + } { + retryCount := i + if i == 2 { + retryCount = 8 + } + delay := retryDelay(policy, retryCount, errors.New("retry")) + if delay < bounds[0] || delay > bounds[1] { + t.Fatalf("retry %d delay %s outside [%s, %s]", retryCount, delay, bounds[0], bounds[1]) + } + } +} + +func TestRetryDelayPrefersServerValue(t *testing.T) { + want := 125 * time.Millisecond + policy := RetryPolicy{ + InitialBackoff: 500 * time.Millisecond, + MaxBackoff: 8 * time.Second, + RetryAfter: func(error) (time.Duration, bool) { + return want, true + }, + } + if got := retryDelay(policy, 0, errors.New("retry")); got != want { + t.Fatalf("retry delay = %s, want %s", got, want) + } +} + +func TestRetryDelayRejectsInvalidServerValues(t *testing.T) { + for _, serverDelay := range []time.Duration{0, -time.Second, 61 * time.Second} { + policy := RetryPolicy{ + InitialBackoff: 500 * time.Millisecond, + MaxBackoff: 8 * time.Second, + MaxRetryAfter: 60 * time.Second, + RetryAfter: func(error) (time.Duration, bool) { + return serverDelay, true + }, + } + got := retryDelay(policy, 0, errors.New("retry")) + if got < 375*time.Millisecond || got > 500*time.Millisecond { + t.Fatalf("server delay %s produced retry delay %s", serverDelay, got) + } + } +} + +func TestRetryWithAttemptReportsRetryCount(t *testing.T) { + var attempts []int + err := RetryWithAttempt( + context.Background(), + RetryPolicy{ + MaxAttempts: 2, + RetryAfter: func(error) (time.Duration, bool) { + return 0, true + }, + }, + func() bool { return true }, + func(retryCount int) error { + attempts = append(attempts, retryCount) + return errors.New("retry") + }, + nil, + func(error) bool { return true }, + ) + if err == nil { + t.Fatal("expected final retry error") + } + if len(attempts) != 3 || attempts[0] != 0 || attempts[1] != 1 || attempts[2] != 2 { + t.Fatalf("attempts = %v, want [0 1 2]", attempts) + } +} diff --git a/examples/README.md b/examples/README.md index b3a94ed..5206c39 100644 --- a/examples/README.md +++ b/examples/README.md @@ -20,3 +20,5 @@ The paired multimodal and sparse embedding examples default to `doubao-embedding MCP is available in both clouds and its examples explicitly send `ark-beta-mcp: true`. Other built-in tools are CN-only: Web Search sends `ark-beta-web-search: true`, and Doubao App sends `ark-beta-doubao-app: true`. The [`self_hosted_worker/`](./self_hosted_worker) example runs a local Managed Agents worker for an existing self-hosted environment. It requires `MA_ENVIRONMENT_ID`; the client defaults to `https://ark.cn-beijing.volces.com/api/v3`. + +The [`self_hosted_mcp_worker/`](./self_hosted_mcp_worker) example discovers tools from a local MCP server, prints their schemas as Agent custom tool declarations for manual configuration, and executes calls through a self-hosted worker. It is an isolated Go 1.23 module so MCP dependencies do not change the core SDK or the other examples. diff --git a/examples/self_hosted_mcp_worker/README.md b/examples/self_hosted_mcp_worker/README.md new file mode 100644 index 0000000..7ffaf97 --- /dev/null +++ b/examples/self_hosted_mcp_worker/README.md @@ -0,0 +1,113 @@ +# Self-hosted MCP worker + +This example follows Anthropic's client-side MCP helper example at the same +level of abstraction: connect to an MCP server, discover its tools, convert +them, and run an existing self-hosted Environment Worker. + +The self-hosted Environment must exist before starting the worker. Create or +update an Agent with the printed `Agent custom tool` declarations before +creating a Session. Printing declarations does not update the Agent +automatically. The same MCP tool list is registered with the worker for +execution, and the example reads every `tools/list` page. + +## Manual end-to-end verification + +The example registers the MCP tool implementation with the self-hosted worker, +but it does not create or update Managed Agents resources. Complete the +following control-plane fields manually: + +1. Create a self-hosted Environment and copy its ID into + `MA_ENVIRONMENT_ID`. +2. Set `ARK_API_KEY`. Set `ARK_BASE_URL` only when using a non-production + endpoint. +3. Start the worker with the MCP server command after `--`: + + ```bash + export ARK_API_KEY=... + export MA_ENVIRONMENT_ID=env_xxx + # Optional, for example when testing against staging: + # export ARK_BASE_URL=https://example.com/api/v3 + + cd examples/self_hosted_mcp_worker + go run . -- go run ./server + ``` + +4. Copy every printed `Agent custom tool: {...}` declaration into the Agent's + tool configuration. For the bundled server, use the declaration below. + Configure it before creating the Session; printing the declaration does not + update the Agent automatically. +5. Create a Session that uses both that Agent and the same self-hosted + Environment from `MA_ENVIRONMENT_ID`. +6. Send a message such as: + + ```text + Call mcp_echo exactly once with text "Hello from MCP echo!" and report the result. + ``` + +The verification passes when the Session shows an `mcp_echo` call with that +input, a `user.custom_tool_result` containing +`MCP echo: Hello from MCP echo!`, a final Agent response, and a final +`session.status_idle` whose stop reason is `end_turn`. A temporary +`session.status_idle` with stop reason `requires_action` means that the Session +is waiting for the external custom-tool result; it is expected and is not an +approval prompt or a failure. At the event level, observe these milestones: + +```text +agent.custom_tool_use +session.status_idle stop_reason=requires_action +user.custom_tool_result posted by the worker +agent.message +session.status_idle stop_reason=end_turn +``` + +Do not depend on the first idle event and the tool-result POST being displayed +in an exact relative order: the worker starts executing as soon as it observes +`agent.custom_tool_use`. + +Keep the worker process running for the whole verification. The command after +`--` is a stdio MCP server command, not a URL; the worker starts the process and +communicates with it through stdin/stdout. + +The bundled server exposes this declaration: + +```json +{ + "type": "custom", + "name": "mcp_echo", + "description": "Echo text through the local MCP server.", + "input_schema": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"] + } +} +``` + +To use another stdio MCP server, replace the command after `--`. Set +`ARK_BASE_URL` only when overriding the SDK's production endpoint. The example +removes `ARK_API_KEY` from the MCP subprocess environment, but inherits other +environment variables. Review or allowlist them before production and use +separate MCP-specific credentials. + +The example opens one MCP process and client session for the lifetime of the +Environment Worker and reuses it for every Managed Agents Session handled by +that worker. MCP calls do not automatically contain the Managed Agents +`session_id` or `work_id`, and Session idle/deletion is not an MCP lifecycle +notification. Use a stateless MCP server or implement explicit tenant/session +isolation, and expect the MCP process to stop only when the worker exits. The +command-line example accepts a stdio child command only; other transports can +be used by constructing an MCP client session programmatically. + +Managed Agents currently accepts at most eight custom tools per Agent. If the +server exposes more, select the same stable subset for both the Agent and the +worker. Custom tools do not use Managed Agents permission policies: the worker +executes matching calls directly, so put approval, authorization, and operation +allowlists in the MCP server or wrapper. Only connect trusted servers, avoid +tool names that collide with built-in Agent tools, and configure an MCP client +timeout. MCP servers run with the worker's OS, filesystem, and network +permissions rather than in a Managed Agents sandbox, so run them with least +privilege and do not pass `ARK_API_KEY` to them. Update the Agent while it is +idle and restart the worker whenever the server's tool list changes. + +This directory is a separate Go 1.23 module so the optional MCP dependency does +not change the core SDK's Go 1.20 baseline. diff --git a/examples/self_hosted_mcp_worker/go.mod b/examples/self_hosted_mcp_worker/go.mod new file mode 100644 index 0000000..1af5bee --- /dev/null +++ b/examples/self_hosted_mcp_worker/go.mod @@ -0,0 +1,29 @@ +module github.com/volcengine/ark-runtime-go/examples/self_hosted_mcp_worker + +go 1.23.0 + +require ( + github.com/modelcontextprotocol/go-sdk v1.3.1 + github.com/volcengine/ark-runtime-go v0.6.0 + github.com/volcengine/ark-runtime-go/mcp v0.6.0 +) + +require ( + github.com/go-faster/errors v0.7.1 // indirect + github.com/go-faster/jx v1.2.0 // indirect + github.com/google/jsonschema-go v0.4.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/jmespath/go-jmespath v0.4.0 // indirect + github.com/segmentio/asm v1.2.1 // indirect + github.com/segmentio/encoding v0.5.3 // indirect + github.com/volcengine/volc-sdk-golang v1.0.23 // indirect + github.com/volcengine/volcengine-go-sdk v1.2.15 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/oauth2 v0.30.0 // indirect + golang.org/x/sys v0.35.0 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect +) + +replace github.com/volcengine/ark-runtime-go => ../.. + +replace github.com/volcengine/ark-runtime-go/mcp => ../../mcp diff --git a/examples/self_hosted_mcp_worker/go.sum b/examples/self_hosted_mcp_worker/go.sum new file mode 100644 index 0000000..1a85a26 --- /dev/null +++ b/examples/self_hosted_mcp_worker/go.sum @@ -0,0 +1,131 @@ +cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= +github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= +github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY= +github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= +github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/go-faster/errors v0.7.1 h1:MkJTnDoEdi9pDabt1dpWf7AA8/BaSYZqibYyhZ20AYg= +github.com/go-faster/errors v0.7.1/go.mod h1:5ySTjWFiphBs07IKuiL69nxdfd5+fzh1u7FPGZP2quo= +github.com/go-faster/jx v1.2.0 h1:T2YHJPrFaYu21fJtUxC9GzmluKu8rVIFDwwGBKTDseI= +github.com/go-faster/jx v1.2.0/go.mod h1:UWLOVDmMG597a5tBFPLIWJdUxz5/2emOpfsj9Neg0PE= +github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= +github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= +github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= +github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= +github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8= +github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= +github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= +github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= +github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U= +github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/modelcontextprotocol/go-sdk v1.3.1 h1:TfqtNKOIWN4Z1oqmPAiWDC2Jq7K9OdJaooe0teoXASI= +github.com/modelcontextprotocol/go-sdk v1.3.1/go.mod h1:DgVX498dMD8UJlseK1S5i1T4tFz2fkBk4xogC3D15nw= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= +github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/segmentio/encoding v0.5.3 h1:OjMgICtcSFuNvQCdwqMCv9Tg7lEOXGwm1J5RPQccx6w= +github.com/segmentio/encoding v0.5.3/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/volcengine/volc-sdk-golang v1.0.23 h1:anOslb2Qp6ywnsbyq9jqR0ljuO63kg9PY+4OehIk5R8= +github.com/volcengine/volc-sdk-golang v1.0.23/go.mod h1:AfG/PZRUkHJ9inETvbjNifTDgut25Wbkm2QoYBTbvyU= +github.com/volcengine/volcengine-go-sdk v1.2.15 h1:duhofGY6gVqcMUfvfa2JTo4uvfixH9rASDlJs4TwQJk= +github.com/volcengine/volcengine-go-sdk v1.2.15/go.mod h1:oxoVo+A17kvkwPkIeIHPVLjSw7EQAm+l/Vau1YGHN+A= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= +golang.org/x/exp v0.0.0-20230116083435-1de6713980de h1:DBWn//IJw30uYCgERoxCg84hWtA97F4wMiKOIh00Uf0= +golang.org/x/exp v0.0.0-20230116083435-1de6713980de/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= +golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= +golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= +golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= +golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= +golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= +golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= +golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= +golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= +google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= +google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= +google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= +google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= +google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= +google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= +google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c= +google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= +google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= +honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/examples/self_hosted_mcp_worker/main.go b/examples/self_hosted_mcp_worker/main.go new file mode 100644 index 0000000..111274c --- /dev/null +++ b/examples/self_hosted_mcp_worker/main.go @@ -0,0 +1,126 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +// Self-hosted worker with client-side MCP tools. +// +// Required: +// +// export ARK_API_KEY=... +// export MA_ENVIRONMENT_ID=env_xxx +// +// Run the bundled MCP server from this directory: +// +// go run . -- go run ./server +package main + +import ( + "context" + "fmt" + "log" + "os" + "os/exec" + "os/signal" + "strings" + "syscall" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/volcengine/ark-runtime-go/arkruntime" + "github.com/volcengine/ark-runtime-go/arkruntime/lib/environments" + arkmcp "github.com/volcengine/ark-runtime-go/mcp" +) + +func main() { + apiKey := mustEnv("ARK_API_KEY") + environmentID := mustEnv("MA_ENVIRONMENT_ID") + commandArgs := mcpCommandArgs(os.Args[1:]) + if len(commandArgs) == 0 { + log.Fatal("MCP server command is required; example: go run . -- go run ./server") + } + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + command := exec.CommandContext(ctx, commandArgs[0], commandArgs[1:]...) + command.Env = environmentWithout(os.Environ(), "ARK_API_KEY") + mcpClient := mcpsdk.NewClient(&mcpsdk.Implementation{ + Name: "ark-self-hosted-worker-example", + Version: "1.0.0", + }, nil) + mcpSession, err := mcpClient.Connect(ctx, &mcpsdk.CommandTransport{Command: command}, nil) + if err != nil { + log.Fatalf("connect MCP server: %v", err) + } + defer func() { + if closeErr := mcpSession.Close(); closeErr != nil { + log.Printf("close MCP session: %v", closeErr) + } + }() + + tools := make([]*mcpsdk.Tool, 0) + for tool, listErr := range mcpSession.Tools(ctx, nil) { + if listErr != nil { + log.Fatalf("list MCP tools: %v", listErr) + } + tools = append(tools, tool) + } + declarations, err := arkmcp.CustomToolItems(tools) + if err != nil { + log.Fatalf("convert MCP tool declarations: %v", err) + } + for _, declaration := range declarations { + raw, marshalErr := declaration.MarshalJSON() + if marshalErr != nil { + log.Fatalf("marshal MCP tool declaration: %v", marshalErr) + } + fmt.Printf("Agent custom tool: %s\n", raw) + } + + customTools, err := arkmcp.NewTools(tools, mcpSession) + if err != nil { + log.Fatalf("create MCP worker tools: %v", err) + } + clientOptions := make([]arkruntime.ConfigOption, 0, 1) + if baseURL := os.Getenv("ARK_BASE_URL"); baseURL != "" { + clientOptions = append(clientOptions, arkruntime.WithBaseUrl(baseURL)) + } + client := arkruntime.NewClientWithApiKey(apiKey, clientOptions...) + worker := environments.NewEnvironmentWorkerForClient(client, environments.EnvironmentWorkerOptions{ + EnvironmentID: environmentID, + Workdir: ".", + CustomTools: customTools, + }) + if err := worker.Run(ctx); err != nil { + log.Fatal(err) + } +} + +func mcpCommandArgs(args []string) []string { + if len(args) > 0 && args[0] == "--" { + return args[1:] + } + return args +} + +func environmentWithout(values []string, names ...string) []string { + removed := make(map[string]struct{}, len(names)) + for _, name := range names { + removed[name] = struct{}{} + } + out := make([]string, 0, len(values)) + for _, value := range values { + name, _, _ := strings.Cut(value, "=") + if _, ok := removed[name]; !ok { + out = append(out, value) + } + } + return out +} + +func mustEnv(name string) string { + value := os.Getenv(name) + if value == "" { + log.Fatalf("%s is required", name) + } + return value +} diff --git a/examples/self_hosted_mcp_worker/server/main.go b/examples/self_hosted_mcp_worker/server/main.go new file mode 100644 index 0000000..e9dc9fd --- /dev/null +++ b/examples/self_hosted_mcp_worker/server/main.go @@ -0,0 +1,40 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func main() { + server := mcp.NewServer(&mcp.Implementation{ + Name: "ark-self-hosted-mcp-example", + Version: "1.0.0", + }, nil) + server.AddTool(&mcp.Tool{ + Name: "mcp_echo", + Description: "Echo text through the local MCP server.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}`), + }, func(_ context.Context, request *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + var input struct { + Text string `json:"text"` + } + if err := json.Unmarshal(request.Params.Arguments, &input); err != nil { + return nil, err + } + return &mcp.CallToolResult{Content: []mcp.Content{ + &mcp.TextContent{Text: "MCP echo: " + input.Text}, + }}, nil + }) + + if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } +} diff --git a/mcp/README.md b/mcp/README.md new file mode 100644 index 0000000..c11a330 --- /dev/null +++ b/mcp/README.md @@ -0,0 +1,114 @@ +# Client-side MCP tools + +This optional module converts tools from the official MCP Go SDK into: + +- `agent.ToolItem` declarations for creating or updating a Managed Agent. +- `agenttoolset.Tool` implementations executed by a self-hosted worker. + +The MCP server only needs to be reachable from the worker. Its credentials stay +in the worker-side MCP transport and are not sent to Managed Agents. + +The adapter wraps an already connected MCP client session, so applications may +use stdio or another transport supported by their MCP client. Keep one client +session open for the worker lifetime. That session is reused across all Managed +Agents Sessions handled by the worker; calls do not automatically include a +Managed Agents `session_id` or `work_id`, and Session idle/deletion is not an +MCP lifecycle notification. Use stateless tools or implement explicit +tenant/session isolation in the MCP server. + +The Agent declaration and worker registry must be built from the same MCP tool +list. Tools are discovered at startup; restart the worker and update the Agent +when the MCP server changes its tool list. + +```go +tools := make([]*mcpsdk.Tool, 0) +for tool, listErr := range session.Tools(ctx, nil) { + if listErr != nil { + return listErr + } + tools = append(tools, tool) +} + +declarations, err := arkmcp.CustomToolItems(tools) +if err != nil { + return err +} +// Use declarations in CreateAgentRequest.Tools or UpdateAgentRequest.Tools. + +customTools, err := arkmcp.NewTools(tools, session) +if err != nil { + return err +} +worker := environments.NewEnvironmentWorkerForClient(client, environments.EnvironmentWorkerOptions{ + EnvironmentID: environmentID, + CustomTools: customTools, +}) +return worker.Run(ctx) +``` + +This module uses the official MCP Go SDK, which requires Go 1.23 or newer. It is +a separate Go module so the main Ark Runtime SDK keeps its Go 1.20 baseline. +Install the adapter with the same release version as the core SDK: + +```bash +go get github.com/volcengine/ark-runtime-go/mcp@v0.6.0 +``` + +Each core `vMAJOR.MINOR.0` release also publishes the matching +`mcp/vMAJOR.MINOR.0` module tag. + +See [`examples/self_hosted_mcp_worker`](../examples/self_hosted_mcp_worker) for +a complete local MCP Server -> Agent custom tool -> self-hosted worker example. + +The main module also exposes the protocol-independent +`arkruntime/selfhosted/mcp.Client` interface. Applications may implement that +interface directly when they use another MCP transport or cannot use the +official Go SDK. + +Managed Agents currently accepts the top-level JSON Schema fields `type`, +`properties`, and `required`. The helper keeps those fields structured, inlines +local `$defs` and `definitions` references used by properties, and appends other +top-level constraints as compact JSON to the tool description. The MCP server +remains the authoritative validator when the worker executes the call. Agent +tool descriptions, including appended constraints, must fit within 10,000 +characters. + +## Tool result support + +The worker preserves MCP `isError` and supports these result blocks: + +- text; +- `image/jpeg`, `image/png`, `image/gif`, and `image/webp` image blocks; +- embedded resources with the same image MIME types; +- embedded `application/pdf` resources; and +- embedded text resources whose MIME type is absent, empty, or starts with + `text/`. + +When a result has no content blocks but has `structuredContent`, the helper +serializes it as compact JSON text. Audio, resource links, unknown content +types, and other resource MIME types become an error result. If a result mixes +supported and unsupported blocks, the whole converted result is an error; the +supported blocks are not returned separately. + +## Operational and security notes + +- Fetch every `tools/list` page. Use the exact same selected tool definitions + for the Agent declaration and worker registry. Managed Agents currently + accepts at most eight custom tools per Agent, so explicitly select a stable + subset when the MCP server exposes more. +- Tool discovery happens at worker startup. When the MCP server changes its + tools, update the Agent while it is idle and restart the worker. +- Tool names must match `[a-zA-Z0-9_-]{1,128}`. Avoid names that collide with + built-in Agent tools, and add your own prefixes when multiple MCP servers + expose the same name. +- Managed Agents permission policies do not apply to custom tools. The worker + executes each matching custom tool call, so implement approval, authorization, + and operation allowlists in the MCP server or a wrapper tool. +- Client-side MCP servers run with the worker's OS, filesystem, and network + permissions; Managed Agents does not put them in a separate sandbox. Run them + with least privilege and a minimal environment. Do not pass `ARK_API_KEY` to + an MCP subprocess; use separate MCP-specific credentials. +- Only wrap MCP servers you trust. Tool names, descriptions, inputs, and results + enter the model context and must be treated as untrusted content. +- Configure MCP transport or client timeouts. The worker `ToolTimeout` remains + the final upper bound, but a shorter MCP timeout gives clearer failures. diff --git a/mcp/go.mod b/mcp/go.mod new file mode 100644 index 0000000..38951a0 --- /dev/null +++ b/mcp/go.mod @@ -0,0 +1,26 @@ +module github.com/volcengine/ark-runtime-go/mcp + +go 1.23.0 + +require ( + github.com/modelcontextprotocol/go-sdk v1.3.1 + github.com/volcengine/ark-runtime-go v0.6.0 +) + +require ( + github.com/go-faster/errors v0.7.1 // indirect + github.com/go-faster/jx v1.2.0 // indirect + github.com/google/jsonschema-go v0.4.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/jmespath/go-jmespath v0.4.0 // indirect + github.com/segmentio/asm v1.2.1 // indirect + github.com/segmentio/encoding v0.5.3 // indirect + github.com/volcengine/volc-sdk-golang v1.0.23 // indirect + github.com/volcengine/volcengine-go-sdk v1.2.15 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/oauth2 v0.30.0 // indirect + golang.org/x/sys v0.35.0 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect +) + +replace github.com/volcengine/ark-runtime-go => .. diff --git a/mcp/go.sum b/mcp/go.sum new file mode 100644 index 0000000..1a85a26 --- /dev/null +++ b/mcp/go.sum @@ -0,0 +1,131 @@ +cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= +github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= +github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY= +github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= +github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/go-faster/errors v0.7.1 h1:MkJTnDoEdi9pDabt1dpWf7AA8/BaSYZqibYyhZ20AYg= +github.com/go-faster/errors v0.7.1/go.mod h1:5ySTjWFiphBs07IKuiL69nxdfd5+fzh1u7FPGZP2quo= +github.com/go-faster/jx v1.2.0 h1:T2YHJPrFaYu21fJtUxC9GzmluKu8rVIFDwwGBKTDseI= +github.com/go-faster/jx v1.2.0/go.mod h1:UWLOVDmMG597a5tBFPLIWJdUxz5/2emOpfsj9Neg0PE= +github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8= +github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk= +github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= +github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= +github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8= +github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= +github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= +github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= +github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U= +github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/modelcontextprotocol/go-sdk v1.3.1 h1:TfqtNKOIWN4Z1oqmPAiWDC2Jq7K9OdJaooe0teoXASI= +github.com/modelcontextprotocol/go-sdk v1.3.1/go.mod h1:DgVX498dMD8UJlseK1S5i1T4tFz2fkBk4xogC3D15nw= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= +github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/segmentio/encoding v0.5.3 h1:OjMgICtcSFuNvQCdwqMCv9Tg7lEOXGwm1J5RPQccx6w= +github.com/segmentio/encoding v0.5.3/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/volcengine/volc-sdk-golang v1.0.23 h1:anOslb2Qp6ywnsbyq9jqR0ljuO63kg9PY+4OehIk5R8= +github.com/volcengine/volc-sdk-golang v1.0.23/go.mod h1:AfG/PZRUkHJ9inETvbjNifTDgut25Wbkm2QoYBTbvyU= +github.com/volcengine/volcengine-go-sdk v1.2.15 h1:duhofGY6gVqcMUfvfa2JTo4uvfixH9rASDlJs4TwQJk= +github.com/volcengine/volcengine-go-sdk v1.2.15/go.mod h1:oxoVo+A17kvkwPkIeIHPVLjSw7EQAm+l/Vau1YGHN+A= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= +golang.org/x/exp v0.0.0-20230116083435-1de6713980de h1:DBWn//IJw30uYCgERoxCg84hWtA97F4wMiKOIh00Uf0= +golang.org/x/exp v0.0.0-20230116083435-1de6713980de/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= +golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= +golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= +golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= +golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= +golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= +golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= +golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= +golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= +google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= +google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= +google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= +google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= +google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= +google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= +google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c= +google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= +google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= +honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/mcp/mcp.go b/mcp/mcp.go new file mode 100644 index 0000000..b8ed578 --- /dev/null +++ b/mcp/mcp.go @@ -0,0 +1,153 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 + +// Package mcp 将官方 MCP Go SDK 适配到方舟 self-hosted MCP 接口。 +package mcp + +import ( + "context" + "errors" + "fmt" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/volcengine/ark-runtime-go/arkruntime/model/agent" + coremcp "github.com/volcengine/ark-runtime-go/arkruntime/selfhosted/mcp" + "github.com/volcengine/ark-runtime-go/arkruntime/tools/agenttoolset" +) + +// Client 把官方 MCP ClientSession 适配成核心 MCP Client 接口。 +type Client struct { + session *mcpsdk.ClientSession +} + +// NewClient 创建官方 MCP ClientSession adapter。 +func NewClient(session *mcpsdk.ClientSession) (*Client, error) { + if session == nil { + return nil, errors.New("mcp client session is required") + } + return &Client{session: session}, nil +} + +// CallTool 调用官方 MCP ClientSession。 +func (c *Client) CallTool(ctx context.Context, name string, arguments map[string]any) (*coremcp.CallToolResult, error) { + result, err := c.session.CallTool(ctx, &mcpsdk.CallToolParams{Name: name, Arguments: arguments}) + if err != nil { + return nil, err + } + return callToolResult(result) +} + +// NewTool 把 MCP Tool 与保持打开的 MCP ClientSession 包装成 worker Custom Tool。 +func NewTool(tool *mcpsdk.Tool, session *mcpsdk.ClientSession) (agenttoolset.Tool, error) { + client, err := NewClient(session) + if err != nil { + return nil, err + } + definition, err := toolDefinition(tool) + if err != nil { + return nil, err + } + return coremcp.NewTool(definition, client) +} + +// NewTools 批量包装 MCP Tools,并返回可传给 EnvironmentWorkerOptions.CustomTools 的映射。 +func NewTools(tools []*mcpsdk.Tool, session *mcpsdk.ClientSession) (map[string]agenttoolset.Tool, error) { + client, err := NewClient(session) + if err != nil { + return nil, err + } + definitions, err := toolDefinitions(tools) + if err != nil { + return nil, err + } + return coremcp.NewTools(definitions, client) +} + +// CustomToolItem 把 MCP Tool 定义转换成创建或更新 Agent 使用的 custom ToolItem。 +func CustomToolItem(tool *mcpsdk.Tool) (agent.ToolItem, error) { + definition, err := toolDefinition(tool) + if err != nil { + return agent.ToolItem{}, err + } + return coremcp.CustomToolItem(definition) +} + +// CustomToolItems 批量生成 Agent custom tool 声明。 +func CustomToolItems(tools []*mcpsdk.Tool) ([]agent.ToolItem, error) { + definitions, err := toolDefinitions(tools) + if err != nil { + return nil, err + } + return coremcp.CustomToolItems(definitions) +} + +func toolDefinitions(tools []*mcpsdk.Tool) ([]coremcp.ToolDefinition, error) { + definitions := make([]coremcp.ToolDefinition, 0, len(tools)) + for _, tool := range tools { + definition, err := toolDefinition(tool) + if err != nil { + return nil, err + } + definitions = append(definitions, definition) + } + return definitions, nil +} + +func toolDefinition(tool *mcpsdk.Tool) (coremcp.ToolDefinition, error) { + if tool == nil { + return coremcp.ToolDefinition{}, errors.New("mcp tool is required") + } + return coremcp.ToolDefinition{ + Name: tool.Name, + Description: tool.Description, + InputSchema: tool.InputSchema, + }, nil +} + +func callToolResult(result *mcpsdk.CallToolResult) (*coremcp.CallToolResult, error) { + if result == nil { + return nil, nil + } + converted := &coremcp.CallToolResult{ + Content: make([]coremcp.Content, 0, len(result.Content)), + StructuredContent: result.StructuredContent, + IsError: result.IsError, + } + for _, content := range result.Content { + item, err := contentValue(content) + if err != nil { + return nil, err + } + converted.Content = append(converted.Content, item) + } + return converted, nil +} + +func contentValue(content mcpsdk.Content) (coremcp.Content, error) { + switch value := content.(type) { + case *mcpsdk.TextContent: + return coremcp.Content{Type: "text", Text: value.Text}, nil + case *mcpsdk.ImageContent: + return coremcp.Content{Type: "image", MIMEType: value.MIMEType, Data: value.Data}, nil + case *mcpsdk.EmbeddedResource: + return coremcp.Content{Type: "resource", Resource: resourceValue(value.Resource)}, nil + case *mcpsdk.AudioContent: + return coremcp.Content{Type: "audio"}, nil + case *mcpsdk.ResourceLink: + return coremcp.Content{Type: "resource_link"}, nil + default: + return coremcp.Content{}, fmt.Errorf("unsupported MCP content type %T", content) + } +} + +func resourceValue(resource *mcpsdk.ResourceContents) *coremcp.Resource { + if resource == nil { + return nil + } + return &coremcp.Resource{ + URI: resource.URI, + MIMEType: resource.MIMEType, + Text: resource.Text, + Blob: resource.Blob, + } +} diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go new file mode 100644 index 0000000..4fab615 --- /dev/null +++ b/mcp/mcp_test.go @@ -0,0 +1,160 @@ +// Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. +// SPDX-License-Identifier: Apache-2.0 +package mcp + +import ( + "context" + "encoding/json" + "strings" + "testing" + + mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" + coremcp "github.com/volcengine/ark-runtime-go/arkruntime/selfhosted/mcp" +) + +func TestCustomToolItem(t *testing.T) { + item, err := CustomToolItem(&mcpsdk.Tool{ + Name: "lookup_order", + Description: "Lookup an order", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "order_id": map[string]any{"type": "string"}, + }, + "required": []string{"order_id"}, + "additionalProperties": false, + }, + }) + if err != nil { + t.Fatal(err) + } + if item.Type != "custom" || item.Name.Value != "lookup_order" || + !strings.HasPrefix(item.Description.Value, "Lookup an order") { + t.Fatalf("unexpected custom tool item: %+v", item) + } + schema := item.InputSchema.Value + if schema.Type.Value != "object" || len(schema.Required) != 1 || schema.Required[0] != "order_id" { + t.Fatalf("unexpected custom tool schema: %+v", schema) + } + if got := string(schema.Properties.Value["order_id"]); got != `{"type":"string"}` { + t.Fatalf("unexpected order_id schema: %s", got) + } + if !strings.Contains(item.Description.Value, `"additionalProperties":false`) { + t.Fatalf("compatibility constraints missing from description: %s", item.Description.Value) + } +} + +func TestCustomToolItemDefaultsDescriptionAndSchema(t *testing.T) { + item, err := CustomToolItem(&mcpsdk.Tool{Name: "ping"}) + if err != nil { + t.Fatal(err) + } + if item.Description.Value != "ping" { + t.Fatalf("description = %q, want ping", item.Description.Value) + } + if item.InputSchema.Value.Type.Value != "object" { + t.Fatalf("schema type = %q, want object", item.InputSchema.Value.Type.Value) + } +} + +func TestCustomToolItemsRejectsDuplicateNames(t *testing.T) { + _, err := CustomToolItems([]*mcpsdk.Tool{{Name: "same"}, {Name: "same"}}) + if err == nil { + t.Fatal("expected duplicate name error") + } +} + +func TestRunnableToolCallsMCPSession(t *testing.T) { + ctx := context.Background() + server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "server", Version: "1.0.0"}, nil) + server.AddTool(&mcpsdk.Tool{ + Name: "echo", + Description: "Echo text", + InputSchema: json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`), + }, func(_ context.Context, request *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { + var input struct { + Text string `json:"text"` + } + if err := json.Unmarshal(request.Params.Arguments, &input); err != nil { + return nil, err + } + return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{ + &mcpsdk.TextContent{Text: "echo: " + input.Text}, + }}, nil + }) + + serverTransport, clientTransport := mcpsdk.NewInMemoryTransports() + serverSession, err := server.Connect(ctx, serverTransport, nil) + if err != nil { + t.Fatal(err) + } + defer serverSession.Close() + client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: "client", Version: "1.0.0"}, nil) + clientSession, err := client.Connect(ctx, clientTransport, nil) + if err != nil { + t.Fatal(err) + } + defer clientSession.Close() + + tool, err := NewTool(&mcpsdk.Tool{Name: "echo"}, clientSession) + if err != nil { + t.Fatal(err) + } + result := tool.Execute(ctx, json.RawMessage(`{"text":"hello"}`)) + if result.IsError || len(result.Content) != 1 || result.Content[0].Text != "echo: hello" { + t.Fatalf("unexpected tool result: %+v", result) + } +} + +func TestConvertCallToolResult(t *testing.T) { + converted, err := callToolResult(&mcpsdk.CallToolResult{ + Content: []mcpsdk.Content{ + &mcpsdk.ImageContent{MIMEType: "image/png", Data: []byte("image")}, + &mcpsdk.EmbeddedResource{Resource: &mcpsdk.ResourceContents{ + URI: "file:///result.txt", + MIMEType: "text/plain", + Text: "resource text", + }}, + }, + }) + if err != nil { + t.Fatal(err) + } + result := coremcp.ConvertCallToolResult(converted) + if result.IsError || len(result.Content) != 2 { + t.Fatalf("unexpected result: %+v", result) + } + imageSource, ok := result.Content[0].Source.(map[string]any) + if !ok || imageSource["type"] != "base64" || imageSource["data"] != "aW1hZ2U=" { + t.Fatalf("unexpected image source: %#v", result.Content[0].Source) + } + documentSource, ok := result.Content[1].Source.(map[string]any) + if !ok || documentSource["type"] != "text" || documentSource["data"] != "resource text" { + t.Fatalf("unexpected document source: %#v", result.Content[1].Source) + } +} + +func TestConvertStructuredAndErrorResults(t *testing.T) { + converted, err := callToolResult(&mcpsdk.CallToolResult{ + StructuredContent: map[string]any{"status": "ok"}, + }) + if err != nil { + t.Fatal(err) + } + structured := coremcp.ConvertCallToolResult(converted) + if structured.IsError || len(structured.Content) != 1 || structured.Content[0].Text != `{"status":"ok"}` { + t.Fatalf("unexpected structured result: %+v", structured) + } + + converted, err = callToolResult(&mcpsdk.CallToolResult{ + Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "not found"}}, + IsError: true, + }) + if err != nil { + t.Fatal(err) + } + failed := coremcp.ConvertCallToolResult(converted) + if !failed.IsError || len(failed.Content) != 1 || failed.Content[0].Text != "not found" { + t.Fatalf("unexpected error result: %+v", failed) + } +}