158 lines
3.4 KiB
Go
158 lines
3.4 KiB
Go
package mcp
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// Client 调用 TaskPool OpenAPI (/open2api/v1)
|
||
type Client struct {
|
||
BaseURL string
|
||
Token string
|
||
HTTPClient *http.Client
|
||
}
|
||
|
||
func NewClient(baseURL, token string) *Client {
|
||
return &Client{
|
||
BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"),
|
||
Token: strings.TrimSpace(token),
|
||
HTTPClient: &http.Client{
|
||
Timeout: 60 * time.Second,
|
||
},
|
||
}
|
||
}
|
||
|
||
func (c *Client) openBase() string {
|
||
return c.BaseURL + "/open2api/v1"
|
||
}
|
||
|
||
func (c *Client) do(method, path string, query map[string]string, body any) (json.RawMessage, error) {
|
||
if c.BaseURL == "" {
|
||
return nil, fmt.Errorf("TASKPOOL_URL 未配置")
|
||
}
|
||
if c.Token == "" {
|
||
return nil, fmt.Errorf("TASKPOOL_TOKEN 未配置")
|
||
}
|
||
|
||
u, err := url.Parse(c.openBase() + path)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if len(query) > 0 {
|
||
q := u.Query()
|
||
for k, v := range query {
|
||
if strings.TrimSpace(v) != "" {
|
||
q.Set(k, v)
|
||
}
|
||
}
|
||
u.RawQuery = q.Encode()
|
||
}
|
||
|
||
var reader io.Reader
|
||
if body != nil {
|
||
b, err := json.Marshal(body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
reader = bytes.NewReader(b)
|
||
}
|
||
|
||
req, err := http.NewRequest(method, u.String(), reader)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+c.Token)
|
||
req.Header.Set("Accept", "application/json")
|
||
if body != nil {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
}
|
||
|
||
resp, err := c.HTTPClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("请求失败: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
raw, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var envelope struct {
|
||
Code int `json:"code"`
|
||
Msg string `json:"msg"`
|
||
Data json.RawMessage `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(raw, &envelope); err != nil {
|
||
if resp.StatusCode >= 400 {
|
||
return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, truncate(string(raw), 500))
|
||
}
|
||
return raw, nil
|
||
}
|
||
|
||
// TaskPool 统一返回 HTTP 200，用 JSON body 中 code 区分成功/失败
|
||
// code=200 → 成功；其他 → 错误
|
||
if envelope.Code != 200 {
|
||
msg := envelope.Msg
|
||
if msg == "" {
|
||
msg = truncate(string(raw), 500)
|
||
}
|
||
return nil, fmt.Errorf("[%d] %s", envelope.Code, msg)
|
||
}
|
||
if len(envelope.Data) > 0 && string(envelope.Data) != "null" {
|
||
return envelope.Data, nil
|
||
}
|
||
// data 为空但请求成功（如 delete/stop 只返回 msg）
|
||
result := map[string]any{"code": envelope.Code, "msg": envelope.Msg}
|
||
b, _ := json.Marshal(result)
|
||
return b, nil
|
||
}
|
||
|
||
func (c *Client) Get(path string, query map[string]string) (json.RawMessage, error) {
|
||
return c.do(http.MethodGet, path, query, nil)
|
||
}
|
||
|
||
func (c *Client) Post(path string, body any) (json.RawMessage, error) {
|
||
return c.do(http.MethodPost, path, nil, body)
|
||
}
|
||
|
||
func (c *Client) Put(path string, body any) (json.RawMessage, error) {
|
||
return c.do(http.MethodPut, path, nil, body)
|
||
}
|
||
|
||
func (c *Client) Delete(path string) (json.RawMessage, error) {
|
||
return c.do(http.MethodDelete, path, nil, nil)
|
||
}
|
||
|
||
func prettyJSON(v any) string {
|
||
b, err := json.MarshalIndent(v, "", " ")
|
||
if err != nil {
|
||
return fmt.Sprintf("%v", v)
|
||
}
|
||
return string(b)
|
||
}
|
||
|
||
func prettyRaw(raw json.RawMessage) string {
|
||
if len(raw) == 0 {
|
||
return "null"
|
||
}
|
||
var v any
|
||
if err := json.Unmarshal(raw, &v); err != nil {
|
||
return string(raw)
|
||
}
|
||
return prettyJSON(v)
|
||
}
|
||
|
||
func truncate(s string, n int) string {
|
||
if len(s) <= n {
|
||
return s
|
||
}
|
||
return s[:n] + "..."
|
||
}
|