agents_and_robots/shell/avatar/fetch.go

// Package avatar fetches avatar images from external providers.
// Impure shell: performs HTTP I/O.
package avatar

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/enmanuel/agents/pkg/avatar"
)

const (
	fetchTimeout = 30 * time.Second
	maxBytes     = 10 * 1024 * 1024 // 10 MB
)

// Fetch downloads an avatar image from the given provider using seed as the
// generation key (typically the agent ID or display name). Returns the path
// to a temporary file containing the image.
// The caller is responsible for removing the temp file when done.
func Fetch(ctx context.Context, provider avatar.Provider, seed string, opts avatar.Options) (string, error) {
	imageURL := avatar.URL(provider, seed, opts)
	return Download(ctx, imageURL)
}

// Download fetches the image at imageURL and writes it to a temp file.
// Returns the temp file path. The caller must remove the file when done.
func Download(ctx context.Context, imageURL string) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, fetchTimeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
	if err != nil {
		return "", fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("User-Agent", "agents-avatar-fetcher/1.0")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("fetch %s: %w", imageURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("fetch %s: HTTP %d", imageURL, resp.StatusCode)
	}

	ext := extensionFromContentType(resp.Header.Get("Content-Type"))

	tmp, err := os.CreateTemp("", "avatar-*"+ext)
	if err != nil {
		return "", fmt.Errorf("create temp file: %w", err)
	}

	if _, err := io.Copy(tmp, io.LimitReader(resp.Body, maxBytes)); err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return "", fmt.Errorf("write temp file: %w", err)
	}

	if err := tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return "", fmt.Errorf("close temp file: %w", err)
	}

	return tmp.Name(), nil
}

// FetchToDir downloads the avatar and saves it to dir/<seed>.<ext>.
// Creates dir if it doesn't exist. Returns the final file path.
func FetchToDir(ctx context.Context, provider avatar.Provider, seed string, opts avatar.Options, dir string) (string, error) {
	tmpPath, err := Fetch(ctx, provider, seed, opts)
	if err != nil {
		return "", err
	}
	defer os.Remove(tmpPath)

	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", fmt.Errorf("create dir %s: %w", dir, err)
	}

	ext := filepath.Ext(tmpPath)
	finalPath := filepath.Join(dir, seed+ext)

	src, err := os.Open(tmpPath)
	if err != nil {
		return "", fmt.Errorf("open temp: %w", err)
	}
	defer src.Close()

	dst, err := os.Create(finalPath)
	if err != nil {
		return "", fmt.Errorf("create %s: %w", finalPath, err)
	}

	if _, err := io.Copy(dst, src); err != nil {
		dst.Close()
		os.Remove(finalPath)
		return "", fmt.Errorf("copy to %s: %w", finalPath, err)
	}

	if err := dst.Close(); err != nil {
		os.Remove(finalPath)
		return "", err
	}

	return finalPath, nil
}

func extensionFromContentType(ct string) string {
	switch ct {
	case "image/png":
		return ".png"
	case "image/jpeg":
		return ".jpg"
	case "image/svg+xml":
		return ".svg"
	case "image/webp":
		return ".webp"
	case "image/gif":
		return ".gif"
	default:
		return ".png"
	}
}