warmbly/internal/app/aitools/tools_web.go

package aitools

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/microcosm-cc/bluemonday"

	"github.com/warmbly/warmbly/internal/app/webhook"
	"github.com/warmbly/warmbly/internal/pkg/generation"
	"github.com/warmbly/warmbly/internal/pkg/safehttp"
)

const (
	fetchMaxBytes   = 2 << 20 // 2 MB cap on a fetched page
	fetchTimeout    = 10 * time.Second
	fetchCacheTTL   = 15 * time.Minute
	fetchCachePrefx = "aitools:fetch:"
)

// fetchTextPolicy strips all tags, leaving text content only.
var fetchTextPolicy = bluemonday.StrictPolicy()

var wsCollapse = regexp.MustCompile(`[ \t]*\n\s*\n\s*`)

func (d Deps) registerWebTools(r *Registry) {
	r.Register(Tool{
		Name:        "search_web",
		Description: "Search the public web and return the top results (title, url, snippet). Use to find current information about a company or person.",
		InputSchema: objectSchema(map[string]any{
			"query": strProp("The search query."),
			"limit": intProp("Max results (1-10, default 5)."),
		}, "query"),
		Risk:    generation.RiskRead,
		Handler: d.searchWeb,
	})

	r.Register(Tool{
		Name:        "fetch_url",
		Description: "Fetch a public web page and return its text content (HTML stripped). Only https public URLs are allowed; private and internal addresses are blocked.",
		InputSchema: objectSchema(map[string]any{
			"url": strProp("The https URL to fetch."),
		}, "url"),
		Risk:    generation.RiskRead,
		Handler: d.fetchURL,
	})
}

func (d Deps) searchWeb(ctx context.Context, _ Invocation, args json.RawMessage) (string, error) {
	in, err := decodeArgs[struct {
		Query string `json:"query"`
		Limit int    `json:"limit"`
	}](args)
	if err != nil {
		return "", err
	}
	if strings.TrimSpace(in.Query) == "" {
		return "", ErrInvalidArgs
	}
	if d.Search == nil {
		return "", generation.ErrSearchNotConfigured
	}
	results, err := d.Search.Search(ctx, in.Query, in.Limit)
	if err != nil {
		return "", err
	}
	return generation.FormatSearchResults(results), nil
}

func (d Deps) fetchURL(ctx context.Context, _ Invocation, args json.RawMessage) (string, error) {
	in, err := decodeArgs[struct {
		URL string `json:"url"`
	}](args)
	if err != nil {
		return "", err
	}
	raw := strings.TrimSpace(in.URL)
	// SSRF guard: reuse the webhook validator (https only, publicly routable,
	// no credentials, blocks localhost/metadata/private IPs).
	if err := webhook.ValidateOutboundURL(raw); err != nil {
		return "", err
	}

	cacheKey := fetchCachePrefx + hashURL(raw)
	if d.Cache != nil {
		var cached string
		if cerr := d.Cache.GetJSON(ctx, cacheKey, &cached); cerr == nil && cached != "" {
			return cached, nil
		}
	}

	// safehttp.Client re-blocks private IPs at dial time (defends against DNS
	// rebinding after the URL-string check).
	client := safehttp.Client(fetchTimeout)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, raw, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", "WarmblyBot/1.0 (+https://warmbly.com)")
	req.Header.Set("Accept", "text/html,text/plain;q=0.9,*/*;q=0.5")

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, fetchMaxBytes))
	if err != nil {
		return "", err
	}

	text := string(body)
	if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "html") {
		text = htmlToText(text)
	}
	text = strings.TrimSpace(text)
	text = truncateRunes(text, 12000)

	result, err := jsonResult(map[string]any{"url": raw, "status": resp.StatusCode, "text": text})
	if err != nil {
		return "", err
	}
	if d.Cache != nil {
		_ = d.Cache.SetJSON(ctx, cacheKey, result, fetchCacheTTL)
	}
	return result, nil
}

// htmlToText strips tags and collapses whitespace so a page fits in the model
// context.
func htmlToText(html string) string {
	stripped := fetchTextPolicy.Sanitize(html)
	stripped = wsCollapse.ReplaceAllString(stripped, "\n\n")
	return strings.TrimSpace(stripped)
}

func hashURL(u string) string {
	sum := sha256.Sum256([]byte(u))
	return hex.EncodeToString(sum[:])
}

// truncateRunes caps s to maxRunes runes (not bytes), appending an ellipsis
// when it trims, so a multi-byte UTF-8 sequence is never split mid-rune.
func truncateRunes(s string, maxRunes int) string {
	if len(s) <= maxRunes {
		return s // fast path: byte length <= rune cap means rune count is too
	}
	r := []rune(s)
	if len(r) <= maxRunes {
		return s
	}
	return string(r[:maxRunes]) + "…"
}