3588AdminBackend/internal/service/agent_client.go
2026-01-10 21:30:28 +08:00

65 lines
1.3 KiB
Go

package service
import (
"bytes"
"fmt"
"io"
"net/http"
"strings"
"time"
"3588AdminBackend/internal/config"
)
type AgentClient struct {
cfg *config.Config
client *http.Client
}
func NewAgentClient(cfg *config.Config) *AgentClient {
return &AgentClient{
cfg: cfg,
client: &http.Client{
Timeout: 3 * time.Second, // Overall timeout
Transport: &http.Transport{
DialContext: (&http.Transport{}).DialContext, // Default dialer
// Connect timeout is usually handled at dialer level but 1s is tight.
// For now using the simple client timeout.
},
},
}
}
func (c *AgentClient) Do(method, ip string, port int, path string, body []byte) ([]byte, int, error) {
url := fmt.Sprintf("http://%s:%d%s", ip, port, path)
req, err := http.NewRequest(method, url, bytes.NewBuffer(body))
if err != nil {
return nil, 0, err
}
if c.cfg.AgentToken != "" {
req.Header.Set("X-RK-Token", c.cfg.AgentToken)
}
if len(body) > 0 {
contentType := "application/json"
if strings.HasPrefix(path, "/v1/models/") {
contentType = "application/octet-stream"
}
req.Header.Set("Content-Type", contentType)
}
resp, err := c.client.Do(req)
if err != nil {
return nil, 0, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, resp.StatusCode, err
}
return respBody, resp.StatusCode, nil
}