test(agentcompat): add protocol clients

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
2026-09-21 02:30:14 +00:00 · 2026-07-20 04:48:15 +00:00
parent 8a0e84b214
commit 1c3e926819
23 changed files with 3329 additions and 0 deletions
@@ -0,0 +1,184 @@
+package client
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+)
+
+type TransferURL struct {
+	URL       string    `json:"url"`
+	Method    string    `json:"method"`
+	ExpiresAt time.Time `json:"expires_at"`
+}
+
+type DownloadURLRequest struct {
+	ServerID   uint64 `json:"server_id"`
+	Path       string `json:"path"`
+	TTLSeconds int    `json:"ttl_seconds,omitempty"`
+}
+
+type UploadURLRequest struct {
+	ServerID      uint64 `json:"server_id"`
+	Path          string `json:"path"`
+	TTLSeconds    int    `json:"ttl_seconds,omitempty"`
+	Mode          string `json:"mode,omitempty"`
+	CreateDirs    bool   `json:"create_dirs,omitempty"`
+	IfMatchSHA256 string `json:"if_match_sha256,omitempty"`
+}
+
+type UploadTransfer struct {
+	Body          io.Reader
+	ContentLength int64
+	SHA256        string
+}
+
+type UploadResult struct {
+	Size   int64  `json:"size"`
+	SHA256 string `json:"sha256"`
+}
+
+type uploadResultPayload struct {
+	Size   *int64  `json:"size"`
+	SHA256 *string `json:"sha256"`
+}
+
+func RequestDownloadURL(ctx context.Context, client *Client, request DownloadURLRequest) (TransferURL, error) {
+	result, err := CallTool[DownloadURLRequest, TransferURL](ctx, client, ToolCall[DownloadURLRequest]{Name: "fs.download_url", Arguments: request})
+	if err != nil {
+		return TransferURL{}, err
+	}
+	return client.validateTransferURL(result.StructuredContent, http.MethodGet)
+}
+
+func RequestUploadURL(ctx context.Context, client *Client, request UploadURLRequest) (TransferURL, error) {
+	result, err := CallTool[UploadURLRequest, TransferURL](ctx, client, ToolCall[UploadURLRequest]{Name: "fs.upload_url", Arguments: request})
+	if err != nil {
+		return TransferURL{}, err
+	}
+	return client.validateTransferURL(result.StructuredContent, http.MethodPost)
+}
+
+func (client *Client) DownloadTransfer(ctx context.Context, transfer TransferURL, destination io.Writer) (int64, error) {
+	validated, err := client.validateTransferURL(transfer, http.MethodGet)
+	if err != nil {
+		return 0, err
+	}
+	requestContext, cancel := client.transferContext(ctx)
+	defer cancel()
+	request, err := http.NewRequestWithContext(requestContext, http.MethodGet, validated.URL, nil)
+	if err != nil {
+		return 0, errorsNewRedacted("create transfer request", err)
+	}
+	response, err := client.transferHTTPClient().Do(request)
+	if err != nil {
+		if requestContext.Err() != nil {
+			return 0, fmt.Errorf("download transfer: %w", requestContext.Err())
+		}
+		return 0, errorsNewRedacted("download transfer", err)
+	}
+	defer response.Body.Close()
+	if response.StatusCode < 200 || response.StatusCode >= 300 {
+		return 0, &HTTPError{StatusCode: response.StatusCode}
+	}
+	written, err := io.Copy(destination, io.LimitReader(response.Body, client.maxTransferBytes))
+	if err != nil {
+		return written, fmt.Errorf("copy download transfer: %w", err)
+	}
+	var overflow [1]byte
+	read, err := response.Body.Read(overflow[:])
+	if err != nil && err != io.EOF {
+		return written, fmt.Errorf("probe download transfer size: %w", err)
+	}
+	if read > 0 {
+		return written, ErrTransferTooLarge
+	}
+	return written, nil
+}
+
+func (client *Client) UploadTransfer(ctx context.Context, transfer TransferURL, upload UploadTransfer) (UploadResult, error) {
+	validated, err := client.validateTransferURL(transfer, http.MethodPost)
+	if err != nil {
+		return UploadResult{}, err
+	}
+	if upload.Body == nil || upload.ContentLength <= 0 {
+		return UploadResult{}, fmt.Errorf("upload content length: %w", ErrInvalidConfig)
+	}
+	if upload.ContentLength > client.maxTransferBytes {
+		return UploadResult{}, ErrTransferTooLarge
+	}
+	transferURL, err := url.Parse(validated.URL)
+	if err != nil {
+		return UploadResult{}, errorsNewRedacted("parse transfer URL", err)
+	}
+	if upload.SHA256 != "" {
+		query := transferURL.Query()
+		query.Set("sha256", upload.SHA256)
+		transferURL.RawQuery = query.Encode()
+	}
+	requestContext, cancel := client.transferContext(ctx)
+	defer cancel()
+	request, err := http.NewRequestWithContext(requestContext, http.MethodPost, transferURL.String(), io.LimitReader(upload.Body, upload.ContentLength))
+	if err != nil {
+		return UploadResult{}, errorsNewRedacted("create transfer request", err)
+	}
+	request.ContentLength = upload.ContentLength
+	response, err := client.transferHTTPClient().Do(request)
+	if err != nil {
+		if requestContext.Err() != nil {
+			return UploadResult{}, fmt.Errorf("upload transfer: %w", requestContext.Err())
+		}
+		return UploadResult{}, errorsNewRedacted("upload transfer", err)
+	}
+	defer response.Body.Close()
+	body, err := readBounded(response.Body, client.maxResponseBytes)
+	if err != nil {
+		return UploadResult{}, err
+	}
+	if response.StatusCode < 200 || response.StatusCode >= 300 {
+		return UploadResult{}, &HTTPError{StatusCode: response.StatusCode, Message: Redact(string(body))}
+	}
+	var payload uploadResultPayload
+	if err := json.NewDecoder(bytes.NewReader(body)).Decode(&payload); err != nil {
+		return UploadResult{}, fmt.Errorf("decode upload result: %w", err)
+	}
+	if payload.Size == nil || payload.SHA256 == nil || *payload.Size <= 0 || *payload.SHA256 == "" {
+		return UploadResult{}, fmt.Errorf("decode upload result: %w", ErrSemanticFailure)
+	}
+	return UploadResult{Size: *payload.Size, SHA256: *payload.SHA256}, nil
+}
+
+func (client *Client) transferHTTPClient() *http.Client {
+	clone := *client.httpClient
+	clone.Jar = nil
+	clone.CheckRedirect = rejectRedirect
+	return &clone
+}
+
+func (client *Client) validateTransferURL(transfer TransferURL, expectedMethod string) (TransferURL, error) {
+	parsed, err := url.Parse(transfer.URL)
+	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
+		return TransferURL{}, fmt.Errorf("transfer URL: %w", ErrInvalidConfig)
+	}
+	if parsed.User != nil || !client.sameOrigin(parsed) {
+		return TransferURL{}, fmt.Errorf("transfer origin: %w", ErrInvalidConfig)
+	}
+	if transfer.ExpiresAt.IsZero() || !time.Now().Before(transfer.ExpiresAt) {
+		return TransferURL{}, ErrTransferExpired
+	}
+	if !strings.EqualFold(transfer.Method, expectedMethod) {
+		return TransferURL{}, fmt.Errorf("transfer method: %w", ErrInvalidConfig)
+	}
+	transfer.Method = expectedMethod
+	return transfer, nil
+}
+
+func (client *Client) sameOrigin(candidate *url.URL) bool {
+	return strings.EqualFold(candidate.Scheme, client.baseURL.Scheme) && strings.EqualFold(candidate.Host, client.baseURL.Host)
+}