diff --git a/.github/workflows/dev-beta-release.yaml b/.github/workflows/dev-beta-release.yaml index 2963f5f0..b8b9665c 100644 --- a/.github/workflows/dev-beta-release.yaml +++ b/.github/workflows/dev-beta-release.yaml @@ -289,12 +289,14 @@ jobs: enable_embedui: "true" enable_python: "true" enable_full_skills: "false" + enable_media_probes: "true" - variant: full suffix: "-full" enable_otel: "false" enable_embedui: "true" enable_python: "true" enable_full_skills: "true" + enable_media_probes: "true" steps: - uses: actions/checkout@v4 with: @@ -348,6 +350,7 @@ jobs: ENABLE_EMBEDUI=${{ matrix.enable_embedui }} ENABLE_PYTHON=${{ matrix.enable_python }} ENABLE_FULL_SKILLS=${{ matrix.enable_full_skills }} + ENABLE_MEDIA_PROBES=${{ matrix.enable_media_probes }} VERSION=${{ needs.beta_version.outputs.tag }} cache-from: type=gha,scope=dev-beta-${{ matrix.variant }} cache-to: type=gha,mode=max,scope=dev-beta-${{ matrix.variant }} diff --git a/.github/workflows/release-beta.yaml b/.github/workflows/release-beta.yaml index 82b6193d..2ace79ed 100644 --- a/.github/workflows/release-beta.yaml +++ b/.github/workflows/release-beta.yaml @@ -98,12 +98,14 @@ jobs: enable_embedui: "true" enable_python: "true" enable_full_skills: "false" + enable_media_probes: "true" - variant: full suffix: "-full" enable_otel: "false" enable_embedui: "true" enable_python: "true" enable_full_skills: "true" + enable_media_probes: "true" steps: - uses: actions/checkout@v4 @@ -149,6 +151,7 @@ jobs: ENABLE_EMBEDUI=${{ matrix.enable_embedui }} ENABLE_PYTHON=${{ matrix.enable_python }} ENABLE_FULL_SKILLS=${{ matrix.enable_full_skills }} + ENABLE_MEDIA_PROBES=${{ matrix.enable_media_probes }} VERSION=${{ github.ref_name }} cache-from: type=gha,scope=beta-${{ matrix.variant }} cache-to: type=gha,mode=max,scope=beta-${{ matrix.variant }} diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 0ace176c..def05209 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -167,16 +167,19 @@ jobs: enable_embedui: "false" enable_python: "false" enable_full_skills: "false" + enable_media_probes: "false" - variant: latest suffix: "" enable_embedui: "true" enable_python: "true" enable_full_skills: "false" + enable_media_probes: "true" - variant: full suffix: "-full" enable_embedui: "true" enable_python: "true" enable_full_skills: "true" + enable_media_probes: "true" runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -223,6 +226,7 @@ jobs: ENABLE_EMBEDUI=${{ matrix.enable_embedui }} ENABLE_PYTHON=${{ matrix.enable_python }} ENABLE_FULL_SKILLS=${{ matrix.enable_full_skills }} + ENABLE_MEDIA_PROBES=${{ matrix.enable_media_probes }} VERSION=v${{ needs.release.outputs.version }} cache-from: type=gha,scope=release-${{ matrix.variant }} cache-to: type=gha,mode=max,scope=release-${{ matrix.variant }} diff --git a/Dockerfile b/Dockerfile index d734b39e..9672f982 100644 --- a/Dockerfile +++ b/Dockerfile @@ -74,6 +74,7 @@ ARG ENABLE_PYTHON=false ARG ENABLE_NODE=false ARG ENABLE_FULL_SKILLS=false ARG ENABLE_CLAUDE_CLI=false +ARG ENABLE_MEDIA_PROBES=true # Copy pinned Python deps (cleaned up after install). # requirements-base.txt: shared deps for ENABLE_PYTHON and ENABLE_FULL_SKILLS. @@ -84,8 +85,18 @@ COPY docker/requirements-base.txt docker/requirements-skills.txt /tmp/ # time.LoadLocation and Python zoneinfo in skill scripts) + optional runtimes. # ENABLE_FULL_SKILLS=true pre-installs all skill deps (larger image, no on-demand install needed). # Otherwise, skill packages are installed on-demand via the admin UI. +# ENABLE_MEDIA_PROBES=true installs ffmpeg (for ffprobe) and poppler-utils (for pdfinfo), the +# two binaries the media budget guard measures with. Without them read_video and read_audio +# refuse the call and read_document falls back to the provider's 1000-page ceiling. +# poppler-utils is skipped here when ENABLE_FULL_SKILLS already installs it below. RUN set -eux; \ apk add --no-cache ca-certificates wget su-exec tzdata; \ + if [ "$ENABLE_MEDIA_PROBES" = "true" ]; then \ + apk add --no-cache ffmpeg; \ + if [ "$ENABLE_FULL_SKILLS" != "true" ]; then \ + apk add --no-cache poppler-utils; \ + fi; \ + fi; \ if [ "$ENABLE_SANDBOX" = "true" ]; then \ apk add --no-cache docker-cli; \ fi; \ diff --git a/README.md b/README.md index 97291a15..c9eabf7e 100644 --- a/README.md +++ b/README.md @@ -223,7 +223,7 @@ When `GOCLAW_*_API_KEY` environment variables are set, the gateway auto-onboards > | Image | Description | > |-------|-------------| > | `latest` | Backend + embedded web UI + Python (**recommended**) | -> | `latest-base` | Backend API-only, no web UI, no runtimes | +> | `latest-base` | Backend API-only, no web UI, no runtimes, no media probes (`read_video` and `read_audio` refuse without `ffprobe`; `read_document` falls back to the 1000-page ceiling without `pdfinfo`) | > | `latest-full` | All runtimes + skill dependencies pre-installed | > | `latest-otel` | Latest + OpenTelemetry tracing | > | `goclaw-web` | Standalone nginx + React SPA (for custom reverse proxy) | diff --git a/docs/14-skills-runtime.md b/docs/14-skills-runtime.md index 8e3580d7..d08cdda8 100644 --- a/docs/14-skills-runtime.md +++ b/docs/14-skills-runtime.md @@ -44,12 +44,21 @@ Pre-installed runtimes depend on the Docker image variant you deploy. The Packag | Variant | Published tag | Build args | Pre-installed runtimes | |---------|---------------|------------|------------------------| -| Latest | `latest` | `ENABLE_PYTHON=true`, `ENABLE_NODE=false`, `ENABLE_FULL_SKILLS=false` | `python3`, `py3-pip`, shared Python deps | -| Base | `base` | `ENABLE_PYTHON=false`, `ENABLE_NODE=false`, `ENABLE_FULL_SKILLS=false` | No Python or Node.js runtimes | -| Full | `full` | `ENABLE_FULL_SKILLS=true` | `python3`, `py3-pip`, `nodejs`, `npm`, `pandoc`, `github-cli`, bundled skill deps, Workspace CLI | +| Latest | `latest` | `ENABLE_PYTHON=true`, `ENABLE_NODE=false`, `ENABLE_FULL_SKILLS=false`, `ENABLE_MEDIA_PROBES=true` | `python3`, `py3-pip`, shared Python deps, `ffprobe`, `pdfinfo` | +| Base | `base` | `ENABLE_PYTHON=false`, `ENABLE_NODE=false`, `ENABLE_FULL_SKILLS=false`, `ENABLE_MEDIA_PROBES=false` | No Python, Node.js, `ffprobe`, or `pdfinfo` | +| Full | `full` | `ENABLE_FULL_SKILLS=true`, `ENABLE_MEDIA_PROBES=true` | `python3`, `py3-pip`, `nodejs`, `npm`, `pandoc`, `github-cli`, `poppler-utils`, bundled skill deps, Workspace CLI, `ffprobe`, `pdfinfo` | | Custom Python | not published | `ENABLE_PYTHON=true` | `python3`, `py3-pip`, shared Python deps | | Custom Node | not published | `ENABLE_NODE=true` | `nodejs`, `npm` | +`ENABLE_MEDIA_PROBES` (default `true`) installs `ffmpeg` for its `ffprobe` binary and +`poppler-utils` for its `pdfinfo` binary. The media budget guard charges a measured duration or +page count, never a number derived from a payload's byte size. Without these binaries, which is +the case for the `base` variant, the desktop build, and bare-binary deployments: + +- `read_video` and `read_audio` refuse the call and name the missing `ffprobe` in the error. +- `read_document` charges every PDF the provider's 1000-page ceiling (258,000 tokens), so it + still works wherever the agent's context window can hold that. + ### Full Variant Extras #### Python Packages diff --git a/internal/mediabudget/mediabudget.go b/internal/mediabudget/mediabudget.go new file mode 100644 index 00000000..b22280c1 --- /dev/null +++ b/internal/mediabudget/mediabudget.go @@ -0,0 +1,574 @@ +// Package mediabudget estimates the token cost of media payloads that are +// uploaded out-of-band to a model provider (read_video, read_audio, +// read_document), so they never appear as bytes in the request the token +// counter otherwise sees. +// +// Only two kinds of number are used here: published per-unit rates, and +// published provider limits. Nothing is derived from a payload's byte size +// except the one case that genuinely is bytes, a non-PDF document, and even +// there the result is capped by a provider limit. A byte size carries no +// information about a duration: 120 seconds of static image fitted in 20,627 +// bytes on the machine this policy was measured on, so any bytes-per-second +// figure is a guess that can be off by orders of magnitude in either +// direction. Where a real measurement is unavailable the payload is either +// charged a proven provider ceiling or refused outright. +package mediabudget + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "math" + "os" + "os/exec" + "path/filepath" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + videoTokensPerSecond = 263 + audioTokensPerSecond = 32 + pdfTokensPerPage = 258 + + // maxPDFPages is Gemini's published PDF limit. It bounds what this package + // CHARGES for a document, not what a provider bills for one. + maxPDFPages = 1000 + maxPDFTokensPerPayload = maxPDFPages * pdfTokensPerPage + + // pdfSniffWindow: readers accept the %PDF- header anywhere in the first 1024 bytes. + pdfSniffWindow = 1024 + + // maxPageScanBytes bounds the page-object scan so one charge cannot become a long read. + maxPageScanBytes = 32 << 20 + + // documentBytesPerToken is the standard rough text ratio, an approximation + // rather than a bound. The proven half of the non-PDF document charge is + // maxPDFTokensPerPayload above it. + documentBytesPerToken = 4 + + probeTimeout = 5 * time.Second + + // maxProbeSeconds: ffprobe parses "inf" as +Inf, so reject any duration past this before it reaches token arithmetic. + maxProbeSeconds = 1e7 // about 115 days + + // maxTokenEstimate: clamp so overflowed or implausible arithmetic never surfaces as a negative charge (read as free by a budget gate). + maxTokenEstimate = 1_000_000_000 +) + +// HeadBytes and TailBytes are the number of leading and trailing bytes a +// caller should read from a remote payload before building a Payload, so +// Tokens has enough of the container to probe. +const ( + HeadBytes = 524288 + TailBytes = 524288 +) + +// ErrDurationUnmeasurable is the parent of every reason a video or audio +// payload could not be measured, so a caller matches the class with errors.Is +// while the operator still reads the specific remedy below. +// +// Refusing rather than charging a ceiling preserves a guarantee that predates +// this package: unverifiable native media was already refused wholesale, for +// every URL. This narrows that refusal to what genuinely cannot be measured +// instead of removing it. A PDF is charged a ceiling instead because its page +// count is cheaply measurable and that ceiling fits an ordinary window. +var ErrDurationUnmeasurable = errors.New("media duration could not be determined") + +// ErrProbeUnavailable means this host has no probe binary to measure with. +var ErrProbeUnavailable = fmt.Errorf("%w: ffprobe (from the ffmpeg package) is not installed, so this tool cannot enforce its token budget", ErrDurationUnmeasurable) + +// ErrNoProbeableBytes means the probe exists but the source yielded nothing to measure. +var ErrNoProbeableBytes = fmt.Errorf("%w: the source provided no bytes that could be probed, so its duration is unknowable here", ErrDurationUnmeasurable) + +// Kind names the media class a payload belongs to. The tool that reads the +// payload knows this; deriving it from a caller-supplied MIME would let a +// mislabelled video be charged on the far cheaper document tier. The zero +// value is the strictest tier, so a Kind left unset cannot under-charge. +type Kind int + +const ( + KindVideo Kind = iota + KindAudio + KindDocument +) + +func (k Kind) String() string { + switch k { + case KindAudio: + return "audio" + case KindDocument: + return "document" + default: + return "video" + } +} + +// Payload describes one out-of-band media payload whose cost must be estimated. +type Payload struct { + Kind Kind // media class; the zero value is video + MIME string // e.g. "video/mp4"; for a document it is the fallback when no bytes can be sniffed + Size int64 // total payload size in bytes, always known + Path string // path to a complete local file, or "" when the payload is remote + Head []byte // leading bytes of a remote payload, the whole file for a document already in memory, or nil + Tail []byte // trailing bytes of a remote payload, or nil +} + +// Tokens charges the tokens a provider will bill for one payload, never less +// than 1 for a non-empty one. +// +// Video and audio are charged from a measured duration or not at all. A PDF is +// charged from a measured page count, or from the 1000-page document maximum +// the provider itself enforces. Any other document is charged as text and +// capped by that same maximum. +func Tokens(ctx context.Context, p Payload) (int, error) { + if p.Size <= 0 { + return 0, nil + } + + if p.Kind == KindDocument { + return documentTokens(ctx, p), nil + } + + rate := int64(videoTokensPerSecond) + if p.Kind == KindAudio { + rate = audioTokensPerSecond + } + + d, err := probeSeconds(ctx, p) + if err != nil { + return 0, fmt.Errorf("%s budget: %w", p.Kind, err) + } + seconds, ok := sanitizeSeconds(d.Seconds()) + if !ok { + return 0, fmt.Errorf("%s budget: %w", p.Kind, ErrNoProbeableBytes) + } + return clampTokens(int64(math.Ceil(seconds)) * rate), nil +} + +// documentTokens never fails: unlike a duration, a document has a ceiling the +// provider publishes and enforces, so an unmeasurable one can be charged that +// ceiling instead of being refused. At 258,000 tokens it still fits the kind of +// window an agent that reads documents is normally given. +func documentTokens(ctx context.Context, p Payload) int { + if !isPDFPayload(p, documentSegments(p, pdfSniffWindow)) { + // docx, pptx, xlsx and friends carry neither a duration nor a page + // count, so they are charged as the text they mostly are. The cap is + // what this package will charge, not a promise about what a provider + // bills for them. + return clampTokens(min(ceilDiv(p.Size, documentBytesPerToken), maxPDFTokensPerPayload)) + } + + pages := 0 + if probed, ok := probePages(ctx, p); ok { + pages = probed + } + // pdfinfo trusts the page tree root's /Count, which a crafted file sets to + // 1 while carrying a thousand real pages. Both halves count something real, + // so taking the larger can only raise the charge toward what the provider + // will bill. + if scanned := countPageObjects(documentSegments(p, maxPageScanBytes)); scanned > pages { + pages = scanned + } + if pages <= 0 { + return clampTokens(maxPDFTokensPerPayload) + } + return clampTokens(int64(min(pages, maxPDFPages)) * pdfTokensPerPage) +} + +// isPDFPayload sniffs the %PDF- magic rather than trusting p.MIME, which is +// caller-supplied and decides which pricing tier applies: the same file +// labelled application/octet-stream would otherwise be charged as loose text. +// The label is consulted only when no bytes are available to sniff. +func isPDFPayload(p Payload, segments [][]byte) bool { + if len(segments) == 0 || len(segments[0]) == 0 { + return isPDF(p.MIME) + } + head := segments[0] + if int64(len(head)) > pdfSniffWindow { + head = head[:pdfSniffWindow] + } + return bytes.Contains(head, []byte("%PDF-")) +} + +func isPDF(mime string) bool { + return mime == "application/pdf" +} + +// documentSegments returns the payload bytes available for inspection here, +// without any network trip: a remote payload carries samples in Head/Tail, a +// local one is read from disk up to limit. +func documentSegments(p Payload, limit int64) [][]byte { + if len(p.Head) > 0 || len(p.Tail) > 0 { + head := capBytes(p.Head, limit) + if int64(len(p.Head)) >= p.Size { + // Head already holds the whole payload; adding Tail would count + // the same page objects a second time. + return [][]byte{head} + } + return [][]byte{head, capBytes(p.Tail, limit)} + } + if p.Path == "" { + return nil + } + f, err := os.Open(p.Path) + if err != nil { + return nil + } + defer f.Close() + data, err := io.ReadAll(io.LimitReader(f, limit)) + if err != nil || len(data) == 0 { + return nil + } + return [][]byte{data} +} + +func capBytes(b []byte, limit int64) []byte { + if int64(len(b)) > limit { + return b[:limit] + } + return b +} + +// pdfPageObjectPattern matches the /Type /Page entry of a page object. PDF +// allows any whitespace, or none, between the two names. +var pdfPageObjectPattern = regexp.MustCompile(`/Type[\x00\t\n\f\r ]*/Page`) + +// countPageObjects counts the page objects visible in the bytes at hand. It is +// a floor and never a bound: pages living in a compressed object stream are +// invisible to it, which is why it only ever raises a probed count. +func countPageObjects(segments [][]byte) int { + total := 0 + for _, seg := range segments { + for _, loc := range pdfPageObjectPattern.FindAllIndex(seg, -1) { + // A name token runs on until a delimiter, so "/Pages" (the page + // tree node, not a page) and any other longer name must not count. + if loc[1] < len(seg) && isPDFNameChar(seg[loc[1]]) { + continue + } + total++ + } + } + return total +} + +func isPDFNameChar(b byte) bool { + switch b { + case 0, '\t', '\n', '\f', '\r', ' ', '(', ')', '<', '>', '[', ']', '{', '}', '/', '%': + return false + } + return true +} + +// ceilDiv avoids the n+d-1 overflow that a Size near math.MaxInt64 triggers. +func ceilDiv(n, d int64) int64 { + q := n / d + if n%d != 0 { + q++ + } + return q +} + +// sanitizeSeconds rejects a duration no real media file would have, whether +// from a crafted ffprobe output or a bogus probe implementation. +func sanitizeSeconds(seconds float64) (float64, bool) { + if math.IsNaN(seconds) || math.IsInf(seconds, 0) || seconds <= 0 || seconds > maxProbeSeconds { + return 0, false + } + return seconds, true +} + +func clampTokens(n int64) int { + if n < 1 { + n = 1 + } + if n > maxTokenEstimate { + n = maxTokenEstimate + } + return int(n) +} + +// DurationProbe measures the duration of a local media file. +type DurationProbe func(ctx context.Context, path string) (time.Duration, bool) + +// PageProbe counts the pages of a local PDF file. +type PageProbe func(ctx context.Context, path string) (int, bool) + +// probeSet is swapped atomically rather than field by field so a test in +// another package cannot race a concurrent Tokens call under -race. +type probeSet struct { + duration DurationProbe + durationReady func() bool + pages PageProbe + pagesReady func() bool +} + +var activeProbes atomic.Pointer[probeSet] + +func init() { + activeProbes.Store(&probeSet{ + duration: ffprobeDuration, + durationReady: ffprobeInstalled, + pages: pdfinfoPageCount, + pagesReady: pdfinfoInstalled, + }) +} + +// SetProbesForTest replaces the external-binary probes, so no test anywhere +// needs ffprobe or pdfinfo on PATH. A nil probe simulates a missing binary. +// It returns a function restoring the previous set. +func SetProbesForTest(duration DurationProbe, pages PageProbe) (restore func()) { + previous := activeProbes.Load() + next := &probeSet{ + duration: duration, + durationReady: func() bool { return duration != nil }, + pages: pages, + pagesReady: func() bool { return pages != nil }, + } + if duration == nil { + next.duration = func(context.Context, string) (time.Duration, bool) { return 0, false } + } + if pages == nil { + next.pages = func(context.Context, string) (int, bool) { return 0, false } + } + activeProbes.Store(next) + return func() { activeProbes.Store(previous) } +} + +// probeSeconds measures the duration of a payload, naming which of the two +// causes stopped it so the operator is pointed at the remedy that applies. +func probeSeconds(ctx context.Context, p Payload) (time.Duration, error) { + probes := activeProbes.Load() + path, cleanup, err := localFile(p, probes.durationReady) + if err != nil { + return 0, err + } + defer cleanup() + d, ok := probes.duration(ctx, path) + if !ok { + return 0, ErrNoProbeableBytes + } + return d, nil +} + +// probePages counts the pages of a PDF payload. +func probePages(ctx context.Context, p Payload) (int, bool) { + probes := activeProbes.Load() + path, cleanup, err := localFile(p, probes.pagesReady) + if err != nil { + return 0, false + } + defer cleanup() + return probes.pages(ctx, path) +} + +// localFile hands a probe a path on disk. A local payload is probed in place; a +// remote one is first materialised into a sparse temp file the probe can seek +// into. Both guards precede materialize: a temp file plus a forked process buys +// nothing with no sampled bytes to read, or with no binary to read them. +func localFile(p Payload, ready func() bool) (path string, cleanup func(), err error) { + noop := func() {} + if p.Path == "" && len(p.Head) == 0 && len(p.Tail) == 0 { + return "", noop, ErrNoProbeableBytes + } + if !ready() { + return "", noop, ErrProbeUnavailable + } + if p.Path != "" { + return p.Path, noop, nil + } + + path, cleanup, materializeErr := materialize(p) + if materializeErr != nil { + return "", noop, ErrNoProbeableBytes + } + return path, cleanup, nil +} + +// materialize writes a remote payload's known head and tail bytes into a +// full-size sparse temp file. A plain file holding only a prefix of the +// payload makes ffprobe under-report duration: it clamps to the bytes it can +// see, and a 64 KB prefix of a 30-second WAV measured as 0.74 seconds. +// Sizing the file to the payload's true size with Truncate, and writing the +// head and tail at their correct offsets, restores the correct answer. +func materialize(p Payload) (path string, cleanup func(), err error) { + dir, err := os.MkdirTemp("", "goclaw-mediabudget-*") + if err != nil { + return "", nil, err + } + cleanup = func() { os.RemoveAll(dir) } + + f, err := os.Create(filepath.Join(dir, "payload")) + if err != nil { + cleanup() + return "", nil, err + } + + if err := writePayload(f, p); err != nil { + f.Close() + cleanup() + return "", nil, err + } + + if err := f.Close(); err != nil { + cleanup() + return "", nil, err + } + + return f.Name(), cleanup, nil +} + +// maxMaterializeBytes bounds the sparse file for a remote payload, whose declared Size is attacker-supplied and not guaranteed sparse-cheap on every filesystem or quota scheme. Tokens still charges from the probe result, not from this clamp. +const maxMaterializeBytes = 4 << 30 // 4 GiB, above any payload this system streams + +func writePayload(f *os.File, p Payload) error { + size := p.Size + if size > maxMaterializeBytes { + size = maxMaterializeBytes + } + if err := f.Truncate(size); err != nil { + return err + } + + head := p.Head + if int64(len(head)) > size { + head = head[:size] + } + if len(head) > 0 { + if _, err := f.WriteAt(head, 0); err != nil { + return err + } + } + + tail := p.Tail + if int64(len(tail)) > size { + tail = tail[int64(len(tail))-size:] + } + if len(tail) > 0 { + if _, err := f.WriteAt(tail, size-int64(len(tail))); err != nil { + return err + } + } + + return nil +} + +// lookupOnce caches one exec.LookPath result, so a per-payload probe does not +// walk PATH on every call. +type lookupOnce struct { + once sync.Once + name string + found bool +} + +func (l *lookupOnce) installed() bool { + l.once.Do(func() { + _, err := exec.LookPath(l.name) + l.found = err == nil + }) + return l.found +} + +var ( + ffprobeLookup = &lookupOnce{name: "ffprobe"} + pdfinfoLookup = &lookupOnce{name: "pdfinfo"} +) + +func ffprobeInstalled() bool { return ffprobeLookup.installed() } + +func pdfinfoInstalled() bool { return pdfinfoLookup.installed() } + +// ffprobeDuration shells out to ffprobe for the duration of a local file. The +// path must never be a URL: -protocol_whitelist file is defence in depth so +// even a malformed argument cannot make ffprobe open a network connection. +func ffprobeDuration(ctx context.Context, path string) (time.Duration, bool) { + if !ffprobeInstalled() { + return 0, false + } + + ctx, cancel := context.WithTimeout(ctx, probeTimeout) + defer cancel() + + cmd := exec.CommandContext(ctx, "ffprobe", + "-v", "error", + "-protocol_whitelist", "file", + "-show_entries", "format=duration", + "-of", "default=nw=1:nk=1", + hardenPath(path), + ) + // Bound Wait() after a context kill even if a child inherits the stdout pipe. + cmd.WaitDelay = time.Second + + out, err := cmd.Output() + if err != nil { + return 0, false + } + + parsed, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64) + if err != nil { + return 0, false + } + seconds, ok := sanitizeSeconds(parsed) + if !ok { + return 0, false + } + + return time.Duration(seconds * float64(time.Second)), true +} + +// pdfinfoPageCount shells out to poppler's pdfinfo for the page count of a +// local PDF. An external parser is used rather than a byte scan because a page +// tree may live inside compressed object streams: pdfinfo reads 19 pages from +// the shared-mime-info-spec.pdf this was measured against, where a hand-rolled +// /Count scan over the same bytes finds nothing at all. Like ffprobeDuration +// the argument is a local path, never a URL, and it is run through argv rather +// than a shell. +func pdfinfoPageCount(ctx context.Context, path string) (int, bool) { + if !pdfinfoInstalled() { + return 0, false + } + + ctx, cancel := context.WithTimeout(ctx, probeTimeout) + defer cancel() + + cmd := exec.CommandContext(ctx, "pdfinfo", hardenPath(path)) + cmd.WaitDelay = time.Second + + out, err := cmd.Output() + if err != nil { + return 0, false + } + return parsePDFInfoPages(string(out)) +} + +// pdfInfoPagesPattern matches pdfinfo's own "Pages:" field. The field is named +// rather than taken as the first integer in the output, which would otherwise +// pick up a page size, a PDF version or a creation date. +var pdfInfoPagesPattern = regexp.MustCompile(`(?m)^Pages:[ \t]+(\d+)[ \t]*$`) + +func parsePDFInfoPages(out string) (int, bool) { + m := pdfInfoPagesPattern.FindStringSubmatch(out) + if m == nil { + return 0, false + } + pages, err := strconv.Atoi(m[1]) + if err != nil || pages <= 0 { + return 0, false + } + return pages, true +} + +// hardenPath stops a path starting with "-" from being parsed as an option by +// the probe binary. +func hardenPath(path string) string { + if filepath.IsAbs(path) || strings.HasPrefix(path, "./") || strings.HasPrefix(path, "../") { + return path + } + return "./" + path +} diff --git a/internal/mediabudget/mediabudget_test.go b/internal/mediabudget/mediabudget_test.go new file mode 100644 index 00000000..c91b581d --- /dev/null +++ b/internal/mediabudget/mediabudget_test.go @@ -0,0 +1,508 @@ +package mediabudget + +import ( + "bytes" + "context" + "errors" + "fmt" + "math" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// stubProbes replaces both external binaries for one test, so no test here +// needs ffprobe or pdfinfo on PATH. A nil probe means "binary missing". +func stubProbes(t *testing.T, duration DurationProbe, pages PageProbe) { + t.Helper() + t.Cleanup(SetProbesForTest(duration, pages)) +} + +func constantDuration(d time.Duration) DurationProbe { + return func(context.Context, string) (time.Duration, bool) { return d, true } +} + +func constantPages(n int) PageProbe { + return func(context.Context, string) (int, bool) { return n, true } +} + +func mustTokens(t *testing.T, p Payload) int { + t.Helper() + got, err := Tokens(context.Background(), p) + if err != nil { + t.Fatalf("Tokens() returned error: %v", err) + } + return got +} + +func TestTokens_VideoChargesMeasuredDuration(t *testing.T) { + stubProbes(t, constantDuration(40*time.Second), nil) + + got := mustTokens(t, Payload{Kind: KindVideo, MIME: "video/mp4", Size: 1, Path: "/does/not/matter"}) + if want := 40 * videoTokensPerSecond; got != want { + t.Fatalf("Tokens() = %d, want %d", got, want) + } +} + +func TestTokens_AudioChargesMeasuredDuration(t *testing.T) { + stubProbes(t, constantDuration(120*time.Second), nil) + + got := mustTokens(t, Payload{Kind: KindAudio, MIME: "audio/mpeg", Size: 1, Path: "/does/not/matter"}) + if want := 120 * audioTokensPerSecond; got != want { + t.Fatalf("Tokens() = %d, want %d", got, want) + } +} + +func TestTokens_FractionalSecondsRoundUp(t *testing.T) { + stubProbes(t, constantDuration(1200*time.Millisecond), nil) + + got := mustTokens(t, Payload{Kind: KindVideo, MIME: "video/mp4", Size: 1, Path: "/does/not/matter"}) + if want := 2 * videoTokensPerSecond; got != want { + t.Fatalf("Tokens() = %d, want %d (1.2s must round up to 2s, not down to 1s)", got, want) + } +} + +// A byte size says nothing about a duration in either direction, so an +// unmeasurable clip is refused with a fix the operator can act on rather than +// charged some number derived from its size. +func TestTokens_UnmeasurableVideoIsRefused(t *testing.T) { + stubProbes(t, nil, nil) + + _, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 200_000_000, Head: []byte("bytes")}) + if !errors.Is(err, ErrDurationUnmeasurable) { + t.Fatalf("Tokens() error = %v, want ErrDurationUnmeasurable", err) + } + if !strings.Contains(err.Error(), "ffprobe") { + t.Fatalf("error %q does not name the missing probe", err) + } + if !strings.Contains(err.Error(), "video") { + t.Fatalf("error %q does not name the media kind", err) + } +} + +func TestTokens_UnmeasurableAudioIsRefused(t *testing.T) { + stubProbes(t, nil, nil) + + _, err := Tokens(context.Background(), Payload{Kind: KindAudio, MIME: "audio/mpeg", Size: 20_000, Head: []byte("bytes")}) + if !errors.Is(err, ErrDurationUnmeasurable) { + t.Fatalf("Tokens() error = %v, want ErrDurationUnmeasurable", err) + } + if !strings.Contains(err.Error(), "ffprobe") { + t.Fatalf("error %q does not name the missing probe", err) + } + if !strings.Contains(err.Error(), "audio") { + t.Fatalf("error %q does not name the media kind", err) + } +} + +// A tiny clip is refused on exactly the same terms as a large one: the policy +// charges measurements, and a small file is no more measurable than a big one. +func TestTokens_UnmeasurableTinyVideoIsRefusedToo(t *testing.T) { + stubProbes(t, nil, nil) + + if _, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 20_000, Head: []byte("bytes")}); !errors.Is(err, ErrDurationUnmeasurable) { + t.Fatalf("Tokens() error = %v, want ErrDurationUnmeasurable", err) + } +} + +// The zero value of Kind must be the strictest tier, so a payload built +// without one cannot slip onto the cheap document charge. +func TestTokens_ZeroValueKindIsChargedAsVideo(t *testing.T) { + stubProbes(t, constantDuration(10*time.Second), nil) + + got := mustTokens(t, Payload{Size: 1, Path: "/does/not/matter"}) + if want := 10 * videoTokensPerSecond; got != want { + t.Fatalf("Tokens() with an unset Kind = %d, want %d (the video rate)", got, want) + } +} + +func TestTokens_NonPositiveSizeReturnsZero(t *testing.T) { + stubProbes(t, nil, nil) + + for _, size := range []int64{0, -1, -100} { + got, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: size}) + if err != nil { + t.Fatalf("Tokens() with Size=%d returned error: %v", size, err) + } + if got != 0 { + t.Fatalf("Tokens() with Size=%d = %d, want 0", size, got) + } + } +} + +func TestTokens_PDFChargesMeasuredPageCount(t *testing.T) { + stubProbes(t, nil, constantPages(19)) + + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: 149_000, Path: "/does/not/matter"}) + if want := 19 * pdfTokensPerPage; got != want { + t.Fatalf("Tokens() = %d, want %d", got, want) + } +} + +// A measured count past the provider limit still stops at the limit: the +// provider refuses the extra pages, so they can never be billed. +func TestTokens_PDFPagesCappedAtProviderLimit(t *testing.T) { + stubProbes(t, nil, constantPages(5000)) + + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: 5 << 20, Path: "/does/not/matter"}) + if want := maxPDFTokensPerPayload; got != want { + t.Fatalf("Tokens() with 5000 measured pages = %d, want %d", got, want) + } +} + +// Unlike a duration, a document has a published ceiling, so an unmeasurable +// one is charged that ceiling instead of being refused. +func TestTokens_UnmeasurablePDFChargesProviderMaximum(t *testing.T) { + stubProbes(t, nil, nil) + + for _, size := range []int64{5 << 20, 500 << 20, math.MaxInt64} { + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: size}) + if got != maxPDFTokensPerPayload { + t.Fatalf("Tokens() for a %d-byte unmeasurable PDF = %d, want %d", size, got, maxPDFTokensPerPayload) + } + } +} + +// docx, pptx and friends have neither a duration nor pages, so they are priced +// as the text they mostly are and bounded by the document maximum. +func TestTokens_NonPDFDocumentChargedAsText(t *testing.T) { + stubProbes(t, nil, nil) + + const size = int64(400_000) + const docxMIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + + got := mustTokens(t, Payload{Kind: KindDocument, MIME: docxMIME, Size: size}) + if want := int(size / documentBytesPerToken); got != want { + t.Fatalf("Tokens() for a 400KB docx = %d, want %d", got, want) + } + if got >= maxPDFTokensPerPayload { + t.Fatalf("Tokens() = %d, want well under the document maximum %d", got, maxPDFTokensPerPayload) + } +} + +func TestTokens_NonPDFDocumentCappedAtDocumentMaximum(t *testing.T) { + stubProbes(t, nil, nil) + + for _, size := range []int64{100 << 20, math.MaxInt64} { + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/octet-stream", Size: size}) + if got != maxPDFTokensPerPayload { + t.Fatalf("Tokens() for a %d-byte document = %d, want the document maximum %d", size, got, maxPDFTokensPerPayload) + } + } +} + +// A remote payload that sampled no bytes has nothing for a probe to read, so +// neither probe may build a full-size sparse file and fork for it. +func TestProbes_RemotePayloadWithoutBytesSkipsMaterialize(t *testing.T) { + durationCalls, pageCalls := 0, 0 + stubProbes(t, + func(context.Context, string) (time.Duration, bool) { durationCalls++; return 30 * time.Second, true }, + func(context.Context, string) (int, bool) { pageCalls++; return 19, true }, + ) + + if _, err := probeSeconds(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 50 << 20}); !errors.Is(err, ErrNoProbeableBytes) { + t.Fatalf("probeSeconds() for a payload carrying no bytes: err = %v, want ErrNoProbeableBytes", err) + } + if _, ok := probePages(context.Background(), Payload{Kind: KindDocument, MIME: "application/pdf", Size: 50 << 20}); ok { + t.Fatal("probePages() reported a page count for a payload carrying no bytes") + } + if durationCalls != 0 || pageCalls != 0 { + t.Fatalf("probes ran %d/%d times for a payload carrying no bytes, want 0/0", durationCalls, pageCalls) + } +} + +// A missing binary must be detected before a temp file is created, so the +// probe is never reached for a remote payload it could not read anyway. +func TestProbes_MissingBinarySkipsMaterialize(t *testing.T) { + stubProbes(t, nil, nil) + + if _, err := probeSeconds(context.Background(), Payload{Kind: KindVideo, Size: 1 << 20, Head: []byte("HEAD")}); !errors.Is(err, ErrProbeUnavailable) { + t.Fatalf("probeSeconds() with no ffprobe available: err = %v, want ErrProbeUnavailable", err) + } + if _, ok := probePages(context.Background(), Payload{Kind: KindDocument, MIME: "application/pdf", Size: 1 << 20, Head: []byte("%PDF")}); ok { + t.Fatal("probePages() succeeded with no pdfinfo available") + } +} + +func TestTokens_RemotePayloadMaterializesSparseFileAndCleansUp(t *testing.T) { + const size = int64(1 << 20) + head := []byte("HEADHEADHEAD") + tail := []byte("TAILTAILTAIL") + + var inspectedDir string + stubProbes(t, func(ctx context.Context, path string) (time.Duration, bool) { + inspectedDir = filepath.Dir(path) + + info, err := os.Stat(path) + if err != nil { + t.Fatalf("stat temp file: %v", err) + } + if info.Size() != size { + t.Fatalf("temp file size = %d, want %d", info.Size(), size) + } + + f, err := os.Open(path) + if err != nil { + t.Fatalf("open temp file: %v", err) + } + defer f.Close() + + gotHead := make([]byte, len(head)) + if _, err := f.ReadAt(gotHead, 0); err != nil { + t.Fatalf("read head: %v", err) + } + if string(gotHead) != string(head) { + t.Fatalf("temp file head = %q, want %q", gotHead, head) + } + + gotTail := make([]byte, len(tail)) + if _, err := f.ReadAt(gotTail, size-int64(len(tail))); err != nil { + t.Fatalf("read tail: %v", err) + } + if string(gotTail) != string(tail) { + t.Fatalf("temp file tail = %q, want %q", gotTail, tail) + } + + return 10 * time.Second, true + }, nil) + + got := mustTokens(t, Payload{Kind: KindVideo, MIME: "video/mp4", Size: size, Head: head, Tail: tail}) + if want := 10 * videoTokensPerSecond; got != want { + t.Fatalf("Tokens() = %d, want %d", got, want) + } + + if inspectedDir == "" { + t.Fatal("probe stub was never called") + } + if _, err := os.Stat(inspectedDir); !os.IsNotExist(err) { + t.Fatalf("temp dir %q still exists after Tokens returned: err=%v", inspectedDir, err) + } +} + +// A stubbed probe returning +Inf must never turn into a negative token count: +// ffprobe itself can report "inf" for a crafted duration field. +func TestTokens_ProbeReturningInfiniteDurationIsRejected(t *testing.T) { + stubProbes(t, constantDuration(time.Duration(math.Inf(1))), nil) + + _, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 1, Path: "/does/not/matter"}) + if !errors.Is(err, ErrDurationUnmeasurable) { + t.Fatalf("Tokens() with an infinite probe duration: err = %v, want ErrDurationUnmeasurable", err) + } +} + +// A stubbed probe claiming 1e15 seconds overflows time.Duration's nanosecond +// range; the result must be a refusal, never a wrapped-around charge. +func TestTokens_ProbeReturningAbsurdDurationIsRejected(t *testing.T) { + var absurdSeconds float64 = 1e15 + stubProbes(t, constantDuration(time.Duration(absurdSeconds*float64(time.Second))), nil) + + _, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 1, Path: "/does/not/matter"}) + if !errors.Is(err, ErrDurationUnmeasurable) { + t.Fatalf("Tokens() with a 1e15s probe duration: err = %v, want ErrDurationUnmeasurable", err) + } +} + +func TestSanitizeSeconds_RejectsInfAndAbsurdSeconds(t *testing.T) { + cases := []struct { + name string + seconds float64 + }{ + {"positive infinity", math.Inf(1)}, + {"negative infinity", math.Inf(-1)}, + {"NaN", math.NaN()}, + {"past maxProbeSeconds", maxProbeSeconds + 1}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, ok := sanitizeSeconds(tc.seconds); ok { + t.Fatalf("sanitizeSeconds(%v) = ok, want rejected", tc.seconds) + } + }) + } +} + +// pdfinfo prints several integers before and after the page count, so the +// field has to be matched by name rather than by position. +func TestParsePDFInfoPages_ReadsTheNamedField(t *testing.T) { + const out = `Title: shared mime info spec +Creator: DocBook XSL Stylesheets V1.79.2 +Producer: Apache FOP Version 2.8 +CreationDate: Mon Jun 3 11:22:33 2024 UTC +Custom Metadata: no +Pages: 19 +Encrypted: no +Page size: 595.276 x 841.89 pts (A4) +PDF version: 1.4 +` + pages, ok := parsePDFInfoPages(out) + if !ok { + t.Fatal("parsePDFInfoPages() found no page count") + } + if pages != 19 { + t.Fatalf("parsePDFInfoPages() = %d, want 19", pages) + } +} + +func TestParsePDFInfoPages_RejectsOutputWithoutTheField(t *testing.T) { + for _, out := range []string{"", "Encrypted: no\n", "Pages:\n", "Pages: 0\n"} { + if pages, ok := parsePDFInfoPages(out); ok { + t.Fatalf("parsePDFInfoPages(%q) = %d, want rejected", out, pages) + } + } +} + +func TestMaterialize_ClampsOversizedHeadAndTail(t *testing.T) { + cases := []struct { + name string + size int64 + head []byte + tail []byte + }{ + {"head longer than size", 4, []byte("HELLOWORLD"), nil}, + {"tail longer than size", 4, nil, []byte("HELLOWORLD")}, + {"head and tail overlap and together exceed size", 6, []byte("ABCDEFGH"), []byte("WXYZ")}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + path, cleanup, err := materialize(Payload{Size: tc.size, Head: tc.head, Tail: tc.tail}) + if err != nil { + t.Fatalf("materialize() returned error: %v", err) + } + defer cleanup() + + info, err := os.Stat(path) + if err != nil { + t.Fatalf("stat temp file: %v", err) + } + if info.Size() != tc.size { + t.Fatalf("temp file size = %d, want %d", info.Size(), tc.size) + } + }) + } +} + +// A URL payload's Size is attacker-supplied and can name a file far larger +// than any that exists, so the sparse temp file must stay bounded. +func TestMaterialize_ClampsHugeDeclaredSize(t *testing.T) { + const declared = int64(100 << 40) // 100 TB, far beyond maxMaterializeBytes + + path, cleanup, err := materialize(Payload{Size: declared}) + if err != nil { + t.Fatalf("materialize() returned error: %v", err) + } + info, statErr := os.Stat(path) + cleanup() + if statErr != nil { + t.Fatalf("stat temp file: %v", statErr) + } + if info.Size() != maxMaterializeBytes { + t.Fatalf("temp file size = %d, want exactly the clamp %d", info.Size(), maxMaterializeBytes) + } +} + +// craftedPDF builds a PDF whose page tree root understates its own size: the +// /Count says 1 while the file carries pageObjects real page objects. pdfinfo +// reports the /Count, ghostscript reports the objects. +func craftedPDF(pageObjects int) []byte { + var b bytes.Buffer + b.WriteString("%PDF-1.7\n1 0 obj\n<< /Type /Pages /Count 1 /Kids [] >>\nendobj\n") + for i := range pageObjects { + fmt.Fprintf(&b, "%d 0 obj\n<< /Type /Page /Parent 1 0 R >>\nendobj\n", i+2) + } + b.WriteString("trailer\n<< /Root 1 0 R >>\n%%EOF\n") + return b.Bytes() +} + +// A crafted /Count must not buy a thousand pages at the price of one. +func TestTokens_CraftedPageCountIsRaisedByTheObjectScan(t *testing.T) { + stubProbes(t, nil, constantPages(1)) + + data := craftedPDF(1000) + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: int64(len(data)), Head: data}) + if got != maxPDFTokensPerPayload { + t.Fatalf("Tokens() for a PDF claiming 1 page and carrying 1000 = %d, want %d", got, maxPDFTokensPerPayload) + } +} + +// The scan reads a local file too, where the whole payload is on disk rather +// than sampled into Head. +func TestTokens_CraftedPageCountIsScannedFromALocalFile(t *testing.T) { + stubProbes(t, nil, constantPages(1)) + + path := filepath.Join(t.TempDir(), "crafted.pdf") + if err := os.WriteFile(path, craftedPDF(400), 0o600); err != nil { + t.Fatal(err) + } + + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: 1, Path: path}) + if want := 400 * pdfTokensPerPage; got != want { + t.Fatalf("Tokens() = %d, want %d (400 page objects on disk)", got, want) + } +} + +// "/Type /Pages" is the page tree node, not a page, and must never be counted. +func TestTokens_PageTreeNodeIsNotCountedAsAPage(t *testing.T) { + stubProbes(t, nil, constantPages(3)) + + data := []byte("%PDF-1.7\n<< /Type /Pages /Count 3 >>\n<< /Type/Pages >>\n") + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: int64(len(data)), Head: data}) + if want := 3 * pdfTokensPerPage; got != want { + t.Fatalf("Tokens() = %d, want %d (the probed count, unchanged by page tree nodes)", got, want) + } +} + +// The MIME label is caller-supplied and would otherwise pick the pricing tier, +// so a PDF labelled application/octet-stream is sniffed and charged as a PDF. +func TestTokens_MislabelledPDFIsStillChargedAsAPDF(t *testing.T) { + stubProbes(t, nil, constantPages(1)) + + data := craftedPDF(1000) + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/octet-stream", Size: int64(len(data)), Head: data}) + if got != maxPDFTokensPerPayload { + t.Fatalf("Tokens() for a mislabelled 1000-page PDF = %d, want %d", got, maxPDFTokensPerPayload) + } +} + +// The opposite label lies the same way: bytes that are plainly not a PDF are +// charged as the text they are, whatever the caller called them. +func TestTokens_NonPDFBytesLabelledPDFAreChargedAsText(t *testing.T) { + stubProbes(t, nil, constantPages(1000)) + + data := []byte(strings.Repeat("plain text, no header here. ", 64)) + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: int64(len(data)), Head: data}) + if want := ceilDiv(int64(len(data)), documentBytesPerToken); got != int(want) { + t.Fatalf("Tokens() = %d, want %d (the text tier)", got, want) + } +} + +// With nothing to sniff, the label is all there is, and a PDF label keeps the +// PDF ceiling rather than falling to the cheaper text tier. +func TestTokens_LabelDecidesTheTierWhenNoBytesAreAvailable(t *testing.T) { + stubProbes(t, nil, nil) + + got := mustTokens(t, Payload{Kind: KindDocument, MIME: "application/pdf", Size: 5 << 20}) + if got != maxPDFTokensPerPayload { + t.Fatalf("Tokens() = %d, want the PDF ceiling %d", got, maxPDFTokensPerPayload) + } +} + +// The two causes of an unmeasurable duration have different remedies, so the +// error must not send an operator to install a binary that is already there. +func TestTokens_UnmeasurableCausesNameTheirOwnRemedy(t *testing.T) { + stubProbes(t, nil, nil) + _, err := Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 20_000, Path: "/does/not/matter"}) + if !errors.Is(err, ErrProbeUnavailable) { + t.Fatalf("with no probe binary: err = %v, want ErrProbeUnavailable", err) + } + + stubProbes(t, constantDuration(30*time.Second), nil) + _, err = Tokens(context.Background(), Payload{Kind: KindVideo, MIME: "video/mp4", Size: 1_900_000_000}) + if !errors.Is(err, ErrNoProbeableBytes) { + t.Fatalf("with a probe but no bytes: err = %v, want ErrNoProbeableBytes", err) + } + if strings.Contains(err.Error(), "not installed") { + t.Fatalf("error %q blames a missing binary that is present", err) + } +} diff --git a/internal/tokencount/budget_counter.go b/internal/tokencount/budget_counter.go index a1f32cee..d311cdae 100644 --- a/internal/tokencount/budget_counter.go +++ b/internal/tokencount/budget_counter.go @@ -20,7 +20,10 @@ import ( const ( budgetEncodingPattern = `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` budgetMessageOverhead = 4 - budgetInlineMediaUnit = 1600 + // InlineMediaUnit is the flat token cost this counter charges for one image, + // video or audio part carried by a structured message, whose bytes it cannot + // size. It is a placeholder unit, not a measurement of the payload. + InlineMediaUnit = 1600 ) //go:embed cl100k_base.tiktoken.gz @@ -80,12 +83,12 @@ func (c *fixedBudgetCounter) CountMessages(messages []providers.Message) (int, e } for _, image := range message.Images { if image.Data != "" || image.URL != "" { - total += budgetInlineMediaUnit + total += InlineMediaUnit } } for _, video := range message.Videos { if video.Data != "" || video.URL != "" { - total += budgetInlineMediaUnit + total += InlineMediaUnit } } } diff --git a/internal/tokencount/budget_counter_test.go b/internal/tokencount/budget_counter_test.go new file mode 100644 index 00000000..002fc457 --- /dev/null +++ b/internal/tokencount/budget_counter_test.go @@ -0,0 +1,51 @@ +package tokencount + +import ( + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +// TestCountMessages_ChargesFlatUnitPerStructuredMedia pins the flat media unit +// that both inline and out-of-band media are charged. Counting media as text +// would scale with payload size and make any real video exceed every window. +func TestCountMessages_ChargesFlatUnitPerStructuredMedia(t *testing.T) { + counter := NewBudgetCounter() + text := []providers.Message{{Role: "user", Content: "describe this"}} + + base, err := counter.CountMessages(text) + if err != nil { + t.Fatalf("CountMessages(text) returned error: %v", err) + } + + withVideo := []providers.Message{{ + Role: "user", + Content: "describe this", + Videos: []providers.VideoContent{ + {MimeType: "video/mp4", URL: "https://example.com/clip.mp4"}, + }, + }} + got, err := counter.CountMessages(withVideo) + if err != nil { + t.Fatalf("CountMessages(video) returned error: %v", err) + } + if got-base != InlineMediaUnit { + t.Fatalf("one video item cost %d tokens, want %d", got-base, InlineMediaUnit) + } + + withTwo := []providers.Message{{ + Role: "user", + Content: "describe this", + Videos: []providers.VideoContent{ + {MimeType: "video/mp4", URL: "https://example.com/a.mp4"}, + {MimeType: "video/mp4", URL: "https://example.com/b.mp4"}, + }, + }} + gotTwo, err := counter.CountMessages(withTwo) + if err != nil { + t.Fatalf("CountMessages(two videos) returned error: %v", err) + } + if gotTwo-base != 2*InlineMediaUnit { + t.Fatalf("two video items cost %d tokens, want %d", gotTwo-base, 2*InlineMediaUnit) + } +} diff --git a/internal/tools/media_budget_chain_test.go b/internal/tools/media_budget_chain_test.go new file mode 100644 index 00000000..b19228aa --- /dev/null +++ b/internal/tools/media_budget_chain_test.go @@ -0,0 +1,256 @@ +package tools + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" + "github.com/nextlevelbuilder/goclaw/internal/providers" + "github.com/nextlevelbuilder/goclaw/internal/store" + usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" +) + +// countingChatProvider answers every Chat as if it were healthy, so a test +// that sees zero calls knows the budget stopped the payload and not the +// provider being unavailable. +type countingChatProvider struct { + name string + calls int +} + +func (p *countingChatProvider) Chat(context.Context, providers.ChatRequest) (*providers.ChatResponse, error) { + p.calls++ + return &providers.ChatResponse{Content: "analysed"}, nil +} + +func (p *countingChatProvider) ChatStream(ctx context.Context, req providers.ChatRequest, _ func(providers.StreamChunk)) (*providers.ChatResponse, error) { + return p.Chat(ctx, req) +} + +func (p *countingChatProvider) DefaultModel() string { return "m" } +func (p *countingChatProvider) Name() string { return p.name } + +// chainOver registers each named provider and returns the chain the tools build +// for them, mirroring videoProviderPriority / documentProviderPriority order. +func chainOver(t *testing.T, names ...string) (*providers.Registry, []MediaProviderEntry, []*countingChatProvider) { + t.Helper() + registry := providers.NewRegistry(func(context.Context) uuid.UUID { return uuid.Nil }) + chain := make([]MediaProviderEntry, 0, len(names)) + fakes := make([]*countingChatProvider, 0, len(names)) + for _, name := range names { + fake := &countingChatProvider{name: name} + registry.Register(fake) + fakes = append(fakes, fake) + entry := MediaProviderEntry{Provider: name, Model: "m", Enabled: true} + entry.applyDefaults() + chain = append(chain, entry) + } + return registry, chain, fakes +} + +func assertNoProviderCalled(t *testing.T, fakes []*countingChatProvider) { + t.Helper() + for _, fake := range fakes { + if fake.calls != 0 { + t.Fatalf("provider %q was called %d time(s) with a refused payload, want 0", fake.name, fake.calls) + } + } +} + +// TestExecuteWithChain_RefusedVideoReachesNoProviderInTheChain is the +// regression test for the bypass: a refusal used to read as "this provider +// failed", so the chain fell through to an entry that charged the payload a +// flat unit and sent 40 MB of video to it. +func TestExecuteWithChain_RefusedVideoReachesNoProviderInTheChain(t *testing.T) { + stubMediaProbes(t, 3600, 1) // an hour of video: 946,800 tokens + + registry, chain, fakes := chainOver(t, videoProviderPriority...) + tool := NewReadVideoTool(registry, nil) + + ctx := withToolAgentBudget(20_000, 8_192) + for i := range chain { + chain[i].Params = map[string]any{ + "prompt": "describe this video", + "data": make([]byte, 40*1024*1024), + "mime": "video/mp4", + videoLocalPathParam: "/does/not/matter", + } + } + + _, err := ExecuteWithChain(ctx, chain, registry, tool.callProvider) + var windowErr *usagecaps.ContextWindowExceededError + if !errors.As(err, &windowErr) { + t.Fatalf("expected the whole chain to refuse on the context window, got %T: %v", err, err) + } + assertNoProviderCalled(t, fakes) +} + +// TestExecuteWithChain_RefusedDocumentReachesNoProviderInTheChain is the same +// bypass for read_document, whose chain is five entries long. +func TestExecuteWithChain_RefusedDocumentReachesNoProviderInTheChain(t *testing.T) { + stubMediaProbes(t, 1, 1000) // the provider's 1000-page maximum: 258,000 tokens + + registry, chain, fakes := chainOver(t, documentProviderPriority...) + tool := NewReadDocumentTool(registry, nil) + + ctx := withToolAgentBudget(20_000, 8_192) + for i := range chain { + chain[i].Params = map[string]any{ + "prompt": "summarise this", + "data": []byte("%PDF-1.7 crafted"), + "mime": "application/pdf", + documentLocalPathParam: "/does/not/matter", + } + } + + _, err := ExecuteWithChain(ctx, chain, registry, tool.callProvider) + var windowErr *usagecaps.ContextWindowExceededError + if !errors.As(err, &windowErr) { + t.Fatalf("expected the whole chain to refuse on the context window, got %T: %v", err, err) + } + assertNoProviderCalled(t, fakes) +} + +// TestReadDocument_RefusedPDFNeverReachesItsProvider covers the non-Gemini +// branch on its own: it builds a base64 document the token counter cannot see, +// so without its own media charge it is the cheap way past the gate. +func TestReadDocument_RefusedPDFNeverReachesItsProvider(t *testing.T) { + stubMediaProbes(t, 1, 1000) + + registry, _, fakes := chainOver(t, "openrouter") + tool := NewReadDocumentTool(registry, nil) + + ctx := withToolAgentBudget(20_000, 8_192) + _, _, err := tool.callProvider(ctx, nil, "openrouter", "m", map[string]any{ + "prompt": "summarise this", + "data": []byte("%PDF-1.7 crafted"), + "mime": "application/pdf", + documentLocalPathParam: "/does/not/matter", + }) + if !isBudgetRefusal(err) { + t.Fatalf("expected a terminal budget refusal, got %T: %v", err, err) + } + assertNoProviderCalled(t, fakes) +} + +// TestReadAudio_RefusedAudioNeverReachesItsProvider proves the audio tool +// stops before transport too. The provider here is an HTTP endpoint, so a hit +// count of zero is proof no request left the process. +func TestReadAudio_RefusedAudioNeverReachesItsProvider(t *testing.T) { + stubMediaProbes(t, 7200, 1) // two hours of audio: 230,400 tokens + + hits := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hits++ + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + tool := NewReadAudioTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key", apiBase: ts.URL} + ctx := withToolAgentBudget(200_000, 8_192) + + _, _, err := tool.callProvider(ctx, cp, "openai", "whisper-1", map[string]any{ + "prompt": "transcribe this", + "data": make([]byte, 4_000_000), + "mime": "audio/mpeg", + audioLocalPathParam: "/does/not/matter", + "_provider_type": "openai", + }) + if !isBudgetRefusal(err) { + t.Fatalf("expected a terminal budget refusal, got %T: %v", err, err) + } + if hits != 0 { + t.Fatalf("the transcription endpoint was contacted %d time(s), want 0", hits) + } +} + +// TestBudgetRefusal_ProviderFailureStillFallsThrough guards the other +// direction: an ordinary provider error must keep the fallback chain working. +func TestBudgetRefusal_ProviderFailureStillFallsThrough(t *testing.T) { + registry, chain, _ := chainOver(t, "first", "second") + + calls := 0 + fn := func(_ context.Context, _ credentialProvider, providerName, _ string, _ map[string]any) ([]byte, *providers.Usage, error) { + calls++ + if providerName == "first" { + return nil, nil, errors.New("upstream 502") + } + return []byte("ok"), nil, nil + } + + result, err := ExecuteWithChain(context.Background(), chain, registry, fn) + if err != nil { + t.Fatalf("a plain provider failure must still fall through, got %v", err) + } + if result.Provider != "second" || calls != 2 { + t.Fatalf("fell through to %q after %d call(s), want second after 2", result.Provider, calls) + } +} + +// TestExecuteWithChain_UnmeasurableMediaIsTerminal covers the refusal class the +// window guard never sees: nothing could be measured, so no number exists to +// compare against a window, and the chain must still stop. +func TestExecuteWithChain_UnmeasurableMediaIsTerminal(t *testing.T) { + registry, chain, _ := chainOver(t, "first", "second") + + calls := 0 + fn := func(context.Context, credentialProvider, string, string, map[string]any) ([]byte, *providers.Usage, error) { + calls++ + return nil, nil, refuseUnpriceableMedia(mediabudget.KindVideo, mediabudget.ErrNoProbeableBytes) + } + + _, err := ExecuteWithChain(context.Background(), chain, registry, fn) + if !errors.Is(err, mediabudget.ErrDurationUnmeasurable) { + t.Fatalf("expected the unmeasurable refusal to surface, got %T: %v", err, err) + } + if calls != 1 { + t.Fatalf("the chain made %d call(s) after a budget refusal, want 1", calls) + } +} + +// TestReserveToolLLMUsageWithMedia_MediaChargeTripsTenantTokenCap joins the two +// halves the other tests cover separately: a charge produced by mediabudget is +// what the tenant's token cap sees, through a real usagecaps.Service. +func TestReserveToolLLMUsageWithMedia_MediaChargeTripsTenantTokenCap(t *testing.T) { + stubMediaProbes(t, 3600, 1) // 946,800 tokens of video + + tenantID := uuid.New() + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: tenantID, MaxTokens: int64Ptr(100_000), Enabled: true} + capStore := &fakeToolUsageCapStore{policies: []store.UsageCapPolicy{policy}, enforceTokenCaps: true} + providerStore := &fakeToolProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: uuid.New()}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + svc := usagecaps.NewService(capStore, providerStore) + + // The window is large enough to admit the charge, so only the tenant cap + // can refuse it. + ctx := store.WithTenantID(withToolAgentBudget(2_000_000, 8), tenantID) + model := "some/model" + + if _, err := reserveToolLLMUsage(ctx, svc, "read_video", "openrouter", model, mediaTestRequest(model)); err != nil { + t.Fatalf("the same request without media must pass the tenant cap, got %v", err) + } + + _, err := reserveToolLLMUsageWithMedia(ctx, svc, "read_video", "openrouter", model, mediaTestRequest(model), bigVideoPayload()) + if !errors.Is(err, usagecaps.ErrCapExceeded) { + t.Fatalf("expected the tenant token cap to block the media charge, got %T: %v", err, err) + } + if !isBudgetRefusal(err) { + t.Fatalf("a tenant cap denial must be terminal for the fallback chain, got %T: %v", err, err) + } + if capStore.reserved.EstimatedTokens < 946_800 { + t.Fatalf("the reservation saw %d tokens, want at least the 946,800-token media charge", capStore.reserved.EstimatedTokens) + } + if len(capStore.events) == 0 || capStore.events[0].Decision != store.UsageCapEventBlock { + t.Fatalf("expected a block event recorded for the tenant, got %+v", capStore.events) + } +} diff --git a/internal/tools/media_provider_chain.go b/internal/tools/media_provider_chain.go index f6a8e992..afa10dca 100644 --- a/internal/tools/media_provider_chain.go +++ b/internal/tools/media_provider_chain.go @@ -163,6 +163,7 @@ type ChainResult struct { // ExecuteWithChain tries each provider in the chain sequentially. // For each provider, it retries up to MaxRetries times (with the configured timeout). // Returns the first successful result or the last error encountered. +// A budget refusal ends the chain immediately: see the check inside the retry loop. func ExecuteWithChain( ctx context.Context, chain []MediaProviderEntry, @@ -221,6 +222,16 @@ func ExecuteWithChain( lastErr = callErr + // A budget refusal is a decision about this payload, not a sick + // provider: falling through would hand the same bytes to the next + // entry and bill them there instead. + if isBudgetRefusal(callErr) { + slog.Warn("media_chain: budget refused the payload, not falling through", + "provider", entry.Provider, "model", entry.Model, + "error", truncateError(callErr)) + return nil, callErr + } + // Don't retry on context cancellation (parent ctx cancelled) if ctx.Err() != nil { return nil, fmt.Errorf("context cancelled: %w", lastErr) diff --git a/internal/tools/read_audio.go b/internal/tools/read_audio.go index 81230316..f2c51e5e 100644 --- a/internal/tools/read_audio.go +++ b/internal/tools/read_audio.go @@ -50,6 +50,9 @@ type ReadAudioTool struct { usageCaps *usagecaps.Service } +// audioLocalPathParam forwards the resolved local file path so the budget estimator can measure the real file. +const audioLocalPathParam = "_local_path" + func NewReadAudioTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadAudioTool { return &ReadAudioTool{registry: registry, mediaLoader: mediaLoader} } @@ -134,6 +137,7 @@ func (t *ReadAudioTool) Execute(ctx context.Context, args map[string]any) *Resul chain[i].Params["prompt"] = prompt chain[i].Params["data"] = data chain[i].Params["mime"] = audioMime + chain[i].Params[audioLocalPathParam] = audioPath } chainResult, err := ExecuteWithChain(ctx, chain, t.registry, t.callProvider) diff --git a/internal/tools/read_audio_resolve.go b/internal/tools/read_audio_resolve.go index 35c9c8b3..a0dbb912 100644 --- a/internal/tools/read_audio_resolve.go +++ b/internal/tools/read_audio_resolve.go @@ -9,6 +9,7 @@ import ( "strings" "time" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" "github.com/nextlevelbuilder/goclaw/internal/providers" ) @@ -74,6 +75,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, prompt := GetParamString(params, "prompt", "Analyze this audio and describe its contents.") data, _ := params["data"].([]byte) mime := GetParamString(params, "mime", "audio/mpeg") + audioPath, _ := params[audioLocalPathParam].(string) // Provider-specific paths require API credentials. Fail-fast (no silent // fallback to chat/completions) for any path we know won't work without @@ -95,7 +97,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, Model: model, Options: map[string]any{"max_tokens": 16384}, } - reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mime, data) + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mediabudget.Payload{Kind: mediabudget.KindAudio, MIME: mime, Size: int64(len(data)), Path: audioPath}) if reserveErr != nil { return nil, nil, reserveErr } @@ -117,7 +119,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, Model: model, Options: map[string]any{"max_tokens": 16384}, } - reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mime, data) + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mediabudget.Payload{Kind: mediabudget.KindAudio, MIME: mime, Size: int64(len(data)), Path: audioPath}) if reserveErr != nil { return nil, nil, reserveErr } @@ -139,7 +141,7 @@ func (t *ReadAudioTool) callProvider(ctx context.Context, cp credentialProvider, Model: model, Options: map[string]any{"max_tokens": 16384}, } - reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mime, data) + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mediabudget.Payload{Kind: mediabudget.KindAudio, MIME: mime, Size: int64(len(data)), Path: audioPath}) if reserveErr != nil { return nil, nil, reserveErr } diff --git a/internal/tools/read_document.go b/internal/tools/read_document.go index ac8d05f4..3c96f012 100644 --- a/internal/tools/read_document.go +++ b/internal/tools/read_document.go @@ -52,6 +52,9 @@ func MediaDocRefsFromCtx(ctx context.Context) []providers.MediaRef { // documentMaxBytes is the max file size for document analysis (20MB). const documentMaxBytes = 20 * 1024 * 1024 +// documentLocalPathParam forwards the resolved local file path so the budget estimator can measure the real file. +const documentLocalPathParam = "_local_path" + // documentProviderPriority is the order in which providers are tried for document analysis. // Gemini has best native PDF support (50MB, 258 tokens/page). claude-cli is // included so installations with only Claude CLI configured can still analyze @@ -204,6 +207,7 @@ func (t *ReadDocumentTool) Execute(ctx context.Context, args map[string]any) *Re chain[i].Params["prompt"] = prompt chain[i].Params["data"] = data chain[i].Params["mime"] = docMime + chain[i].Params[documentLocalPathParam] = docPath } chainResult, err := ExecuteWithChain(ctx, chain, t.registry, t.callProvider) diff --git a/internal/tools/read_document_resolve.go b/internal/tools/read_document_resolve.go index 1b54825e..0380af97 100644 --- a/internal/tools/read_document_resolve.go +++ b/internal/tools/read_document_resolve.go @@ -8,6 +8,7 @@ import ( "path/filepath" "strings" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" "github.com/nextlevelbuilder/goclaw/internal/providers" ) @@ -109,6 +110,7 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid prompt := GetParamString(params, "prompt", "Analyze this document and describe its contents.") data, _ := params["data"].([]byte) mime := GetParamString(params, "mime", "application/octet-stream") + docPath, _ := params[documentLocalPathParam].(string) // Gemini: use native API (requires credentials; OpenAI-compat endpoint doesn't support non-image MIME types). ptype := GetParamString(params, "_provider_type", providerTypeFromName(providerName)) @@ -125,7 +127,7 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid Model: model, Options: map[string]any{"max_tokens": 16384}, } - reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mime, data) + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mediabudget.Payload{Kind: mediabudget.KindDocument, MIME: mime, Size: int64(len(data)), Path: docPath, Head: data}) if reserveErr != nil { return nil, nil, reserveErr } @@ -170,7 +172,11 @@ func (t *ReadDocumentTool) callProvider(ctx context.Context, cp credentialProvid Model: model, Options: opts, } - reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + // Charged on the same terms as the Gemini branch: the base64 document + // rides in a content part the token counter does not price, so an + // unpriced branch here is a way round the gate the Gemini branch applies. + payload := mediabudget.Payload{Kind: mediabudget.KindDocument, MIME: mime, Size: int64(len(data)), Path: docPath, Head: data} + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, payload) if reserveErr != nil { return nil, nil, reserveErr } diff --git a/internal/tools/read_video.go b/internal/tools/read_video.go index 5f660d90..2666e358 100644 --- a/internal/tools/read_video.go +++ b/internal/tools/read_video.go @@ -3,11 +3,13 @@ package tools import ( "context" "fmt" + "io" "log/slog" "net" "os" "path/filepath" "strings" + "time" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/security" @@ -36,6 +38,9 @@ const videoMaxBytes = 100 * 1024 * 1024 const videoURLPinnedIPParam = "_pinned_ip" +// videoLocalPathParam forwards the resolved local file path so the budget probe can read the real file. +const videoLocalPathParam = "_local_path" + // videoProviderPriority is the order in which providers are tried for video analysis. // OpenAI excluded — no native video upload in chat completions. var videoProviderPriority = []string{"gemini", "openrouter"} @@ -52,10 +57,21 @@ type ReadVideoTool struct { registry *providers.Registry mediaLoader MediaPathLoader usageCaps *usagecaps.Service + // streamToGemini is injected per instance rather than through a package + // global, so parallel tests cannot race each other over one seam. + streamToGemini geminiStreamFn } +// geminiStreamFn is the upload a gate-passing URL is handed to. +type geminiStreamFn func(ctx context.Context, apiKey, model, prompt string, reader io.Reader, contentLength int64, mime string, httpTimeout time.Duration) (*providers.ChatResponse, error) + func NewReadVideoTool(registry *providers.Registry, mediaLoader MediaPathLoader) *ReadVideoTool { - return &ReadVideoTool{registry: registry, mediaLoader: mediaLoader} + return &ReadVideoTool{registry: registry, mediaLoader: mediaLoader, streamToGemini: geminiFileAPICallStream} +} + +// setStreamToGeminiForTest observes what reaches the upload without contacting a provider. +func (t *ReadVideoTool) setStreamToGeminiForTest(fn geminiStreamFn) { + t.streamToGemini = fn } func (t *ReadVideoTool) SetUsageCapService(svc *usagecaps.Service) { @@ -119,6 +135,7 @@ func (t *ReadVideoTool) Execute(ctx context.Context, args map[string]any) *Resul var data []byte var videoMime string + var videoPath string var pinnedIP net.IP if videoURL != "" { @@ -132,7 +149,7 @@ func (t *ReadVideoTool) Execute(ctx context.Context, args map[string]any) *Resul ext := filepath.Ext(validatedURL.Path) videoMime = mimeFromVideoExt(ext) } else { - var videoPath, mime string + var mime string var err error if videoArg != "" { videoPath, err = resolveStructuredMediaPath(ctx, videoArg, "video") @@ -171,6 +188,7 @@ func (t *ReadVideoTool) Execute(ctx context.Context, args map[string]any) *Resul chain[i].Params["data"] = data chain[i].Params["url"] = videoURL chain[i].Params["mime"] = videoMime + chain[i].Params[videoLocalPathParam] = videoPath if pinnedIP != nil { chain[i].Params[videoURLPinnedIPParam] = pinnedIP } diff --git a/internal/tools/read_video_resolve.go b/internal/tools/read_video_resolve.go index 548fd799..cd10e1c2 100644 --- a/internal/tools/read_video_resolve.go +++ b/internal/tools/read_video_resolve.go @@ -4,13 +4,16 @@ import ( "context" "encoding/base64" "fmt" + "io" "log/slog" "net" "net/http" "path/filepath" + "strconv" "strings" "time" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/security" usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" @@ -74,6 +77,130 @@ func (t *ReadVideoTool) resolveVideoFile(ctx context.Context, mediaID string) (p return p, mime, nil } +// probeVideoURL learns a remote video's size, and where possible a head and +// tail sample, from two small ranged reads instead of buffering the whole +// file. Head and tail are requested independently so a server that fails one +// range can still answer the other. A prefix alone is not enough: it +// under-reports duration by 98% for mpeg and 79% for ogg, because the index +// several formats need (moov, the ogg/opus final granule, the AVI index) +// sits at the end of the file, not the start. +func probeVideoURL(ctx context.Context, pinnedIP net.IP, rawURL, mime string) (mediabudget.Payload, bool) { + head, headTotal, _ := probeVideoRange(ctx, pinnedIP, rawURL, + fmt.Sprintf("bytes=0-%d", mediabudget.HeadBytes-1), mediabudget.HeadBytes) + tail, tailTotal, tailIsHead := probeVideoRange(ctx, pinnedIP, rawURL, + fmt.Sprintf("bytes=-%d", mediabudget.TailBytes), mediabudget.TailBytes) + if tailIsHead { + // A 200 to a suffix range is the start of the file, not its end; + // written at the tail offset it can only mislead ffprobe. + tail = nil + } + + total := headTotal + if tailTotal > total { + total = tailTotal + } + if total <= 0 { + // Both ranged reads came back empty-handed, which is what an origin + // that rejects Range outright looks like. A plain HEAD still reports + // the size on most of them; that is not a duration, so such a payload + // is refused rather than charged, but reporting it as known keeps the + // caller from taking a second guess at the same unknowable number. + if size, ok := headVideoURLSize(ctx, pinnedIP, rawURL); ok { + return mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: mime, Size: size}, true + } + return mediabudget.Payload{}, false + } + // Accepted risk: a crafted container can make ffprobe under-report duration; Reservation.Reconcile corrects it afterward. + return mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: mime, Size: total, Head: head, Tail: tail}, true +} + +// probeRangeTimeout bounds a probe read (at most one window); the streaming GET below has no deadline since it legitimately transfers a large body. +const probeRangeTimeout = 10 * time.Second + +// probeVideoRange issues one ranged GET and returns whatever bytes and total +// size it could learn. It goes through the same pinned-IP context and +// SSRF-safe client as the streaming GET below: this is a second network trip +// to a caller-supplied URL, and it must not open an unpinned path to it. +// +// Any failure, at any stage, is reported as total == 0 rather than an error: a +// probe is best-effort, and its caller has a HEAD fallback behind it, then +// refuses. isHead marks a body that is the start of the file whatever range +// was asked for, so a caller cannot mistake it for the end. +func probeVideoRange(ctx context.Context, pinnedIP net.IP, rawURL, rangeHeader string, capBytes int64) (data []byte, total int64, isHead bool) { + reqCtx, cancel := context.WithTimeout(security.WithPinnedIP(ctx, pinnedIP), probeRangeTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(reqCtx, "GET", rawURL, nil) + if err != nil { + return nil, 0, false + } + req.Header.Set("Range", rangeHeader) + + resp, err := security.NewSafeClient(0).Do(req) + if err != nil { + return nil, 0, false + } + defer resp.Body.Close() + + switch resp.StatusCode { + case http.StatusPartialContent: + total, _ = parseContentRangeTotal(resp.Header.Get("Content-Range")) + body, _ := io.ReadAll(io.LimitReader(resp.Body, capBytes)) + return body, total, false + case http.StatusOK: + // No range support: cap the read to the head window and close + // immediately, so a server that ignores Range never streams its + // whole file into a probe. + body, _ := io.ReadAll(io.LimitReader(resp.Body, mediabudget.HeadBytes)) + return body, resp.ContentLength, true + default: + return nil, 0, false + } +} + +// headVideoURLSize asks for the size alone, for an origin that refuses ranged +// reads but still answers a HEAD. It uses the same pinned-IP, SSRF-safe client +// as every other trip to this URL. +func headVideoURLSize(ctx context.Context, pinnedIP net.IP, rawURL string) (int64, bool) { + reqCtx, cancel := context.WithTimeout(security.WithPinnedIP(ctx, pinnedIP), probeRangeTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(reqCtx, http.MethodHead, rawURL, nil) + if err != nil { + return 0, false + } + + resp, err := security.NewSafeClient(0).Do(req) + if err != nil { + return 0, false + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 || resp.ContentLength <= 0 { + return 0, false + } + return resp.ContentLength, true +} + +// parseContentRangeTotal extracts total from a "bytes a-b/total" Content-Range +// header. An unparsable header or an unknown total ("bytes a-b/*") both +// report as an error, which the caller treats as size-unknown. +func parseContentRangeTotal(headerValue string) (int64, error) { + const prefix = "bytes " + if !strings.HasPrefix(headerValue, prefix) { + return 0, fmt.Errorf("range probe: missing Content-Range") + } + idx := strings.IndexByte(headerValue, '/') + if idx < 0 { + return 0, fmt.Errorf("range probe: malformed Content-Range %q", headerValue) + } + totalStr := headerValue[idx+1:] + if totalStr == "*" { + return 0, fmt.Errorf("range probe: unknown total in Content-Range %q", headerValue) + } + return strconv.ParseInt(totalStr, 10, 64) +} + // callProvider dispatches video analysis to the appropriate provider API. // Gemini: uses File API (upload → poll → file_data in generateContent). // Others: falls back to base64 or URL in video_url (OpenRouter routes to Gemini which handles video). @@ -101,16 +228,20 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, Model: model, Options: map[string]any{"max_tokens": 16384}, } - // The reservation differs by transport. The base64 path has the real - // bytes now and counts them. The URL-stream path never buffers the - // payload, so it cannot prove completeInput and fails closed under an - // agent budget rather than trusting Content-Length. var reservation *usagecaps.Reservation var reserveErr error if videoURL != "" { - reservation, reserveErr = reserveToolLLMUsageUnverifiableMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + // An origin that answers neither a ranged read nor a HEAD leaves + // nothing to measure and nothing a later gate could measure + // either, so refuse before issuing a GET to a caller-supplied URL. + payload, ok := probeVideoURL(ctx, pinnedIP, videoURL, mime) + if !ok { + return nil, nil, refuseUnpriceableMedia(mediabudget.KindVideo, mediabudget.ErrNoProbeableBytes) + } + reservation, reserveErr = reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, payload) } else { - reservation, reserveErr = reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mime, data) + localPath, _ := params[videoLocalPathParam].(string) + reservation, reserveErr = reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: mime, Size: int64(len(data)), Path: localPath}) } if reserveErr != nil { return nil, nil, reserveErr @@ -174,7 +305,10 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, mime = contentType } - resp, err = geminiFileAPICallStream(ctx, cp.APIKey(), model, prompt, httpResp.Body, contentLength, mime, 300*time.Second) + // No re-gate here: the front gate refused anything it could not + // measure, and Content-Length says nothing about a duration that + // could revise the charge already taken. + resp, err = t.streamToGemini(ctx, cp.APIKey(), model, prompt, httpResp.Body, contentLength, mime, 300*time.Second) } else { slog.Info("read_video: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) resp, err = geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 180*time.Second) @@ -195,13 +329,24 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, return nil, nil, fmt.Errorf("provider %q not available: %w", providerName, err) } + // The payload is charged here too. This branch sends it base64-inlined or + // by URL, either way out of sight of the token counter, so without its own + // media charge it would be the cheap way around every gate above. var vidContent providers.VideoContent + var payload mediabudget.Payload if videoURL != "" { slog.Info("read_video: using chat API with direct video URL", "provider", providerName, "model", model, "url", videoURL) vidContent = providers.VideoContent{MimeType: mime, URL: videoURL} + probed, ok := probeVideoURL(ctx, pinnedIP, videoURL, mime) + if !ok { + return nil, nil, refuseUnpriceableMedia(mediabudget.KindVideo, mediabudget.ErrNoProbeableBytes) + } + payload = probed } else { slog.Info("read_video: using chat API fallback with base64", "provider", providerName, "model", model, "size", len(data)) vidContent = providers.VideoContent{MimeType: mime, Data: base64.StdEncoding.EncodeToString(data)} + localPath, _ := params[videoLocalPathParam].(string) + payload = mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: mime, Size: int64(len(data)), Path: localPath} } chatReq := providers.ChatRequest{ @@ -218,7 +363,7 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, "temperature": 0.2, }, } - reservation, reserveErr := reserveToolLLMUsage(ctx, t.usageCaps, t.Name(), providerName, model, chatReq) + reservation, reserveErr := reserveToolLLMUsageWithMedia(ctx, t.usageCaps, t.Name(), providerName, model, chatReq, payload) if reserveErr != nil { return nil, nil, reserveErr } diff --git a/internal/tools/read_video_test.go b/internal/tools/read_video_test.go index abd55adf..cca4f1c6 100644 --- a/internal/tools/read_video_test.go +++ b/internal/tools/read_video_test.go @@ -2,11 +2,18 @@ package tools import ( "context" + "errors" + "fmt" + "io" "net/http" "net/http/httptest" + "strconv" "strings" "testing" + "time" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" + "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/security" "github.com/nextlevelbuilder/goclaw/internal/store" ) @@ -19,6 +26,39 @@ type mockCredentialProvider struct { func (m *mockCredentialProvider) APIKey() string { return m.apiKey } func (m *mockCredentialProvider) APIBase() string { return m.apiBase } +// parseTestByteRange resolves a "bytes=A-B" or suffix "bytes=-N" spec against +// a known total, mirroring what a real range-capable file server would do. +func parseTestByteRange(header string, total int64) (start, end int64, ok bool) { + const prefix = "bytes=" + if !strings.HasPrefix(header, prefix) { + return 0, 0, false + } + spec := strings.TrimPrefix(header, prefix) + if strings.HasPrefix(spec, "-") { + n, err := strconv.ParseInt(spec[1:], 10, 64) + if err != nil { + return 0, 0, false + } + if n > total { + n = total + } + return total - n, total - 1, true + } + parts := strings.SplitN(spec, "-", 2) + if len(parts) != 2 { + return 0, 0, false + } + s, err1 := strconv.ParseInt(parts[0], 10, 64) + e, err2 := strconv.ParseInt(parts[1], 10, 64) + if err1 != nil || err2 != nil { + return 0, 0, false + } + if e >= total { + e = total - 1 + } + return s, e, true +} + func TestReadVideo_BothMediaIdAndUrl_Error(t *testing.T) { tool := NewReadVideoTool(nil, nil) @@ -53,27 +93,21 @@ func TestReadVideo_PrivateURL_Error(t *testing.T) { } } -// TestReadVideo_GeminiURL_FailsClosedUnderAgentBudget locks the fail-closed -// contract for a streamed video URL. The stream is never buffered, so its bytes -// cannot be counted into completeInput; under an agent budget the call refuses -// before any transport rather than undercounting. Because every tool LLM call -// is agent-scoped in production, this refusal — not the downstream -// Content-Length / 2 GB / HTTP-status checks — is the reachable behavior. Those -// transport validations remain in callProvider for defense in depth but are -// unreachable once an agent budget is present. -func TestReadVideo_GeminiURL_FailsClosedUnderAgentBudget(t *testing.T) { +// TestReadVideo_GeminiURL_UnmeasurableOriginNeverGetsAStreamingGET locks the +// front gate: every request to this server fails, so both range probes and the +// HEAD learn nothing. Nothing later could price the payload either, so the +// streaming GET to a caller-supplied URL must never be issued. +func TestReadVideo_GeminiURL_UnmeasurableOriginNeverGetsAStreamingGET(t *testing.T) { security.SetAllowLoopbackForTest(true) defer security.SetAllowLoopbackForTest(false) - // A server that would otherwise satisfy the static-streaming constraints - // (valid Content-Length, 2xx). The call must still fail closed before it is - // ever contacted, because the payload cannot be counted. - hits := 0 + hits, streamHits := 0, 0 ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits++ - w.Header().Set("Content-Length", "1024") - w.WriteHeader(http.StatusOK) - w.Write(make([]byte, 1024)) + if r.Method == http.MethodGet && r.Header.Get("Range") == "" { + streamHits++ + } + w.WriteHeader(http.StatusInternalServerError) })) defer ts.Close() @@ -90,13 +124,452 @@ func TestReadVideo_GeminiURL_FailsClosedUnderAgentBudget(t *testing.T) { } _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", params) - if err == nil { - t.Fatalf("expected fail-closed for an unverifiable streamed video URL under an agent budget") + if !errors.Is(err, mediabudget.ErrNoProbeableBytes) { + t.Fatalf("expected the refusal to name the unprobeable source, got %T: %v", err, err) } - if !strings.Contains(err.Error(), "cannot verify streamed native media") { - t.Errorf("unexpected error, want fail-closed refusal: %v", err) + if !isBudgetRefusal(err) { + t.Fatalf("the refusal must be terminal for the fallback chain, got %T: %v", err, err) } - if hits != 0 { - t.Errorf("fail-closed must refuse before any transport, but the URL was contacted %d time(s)", hits) + if streamHits != 0 { + t.Fatalf("the streaming GET reached the origin %d time(s), want 0", streamHits) + } + if hits != 3 { + t.Fatalf("the URL was contacted %d time(s), want 3 (head probe, tail probe, HEAD)", hits) + } +} + +// TestReadVideo_GeminiURL_RejectsMissingContentLength covers the Content-Length +// check in read_video_resolve.go, which was dead code while the fail-closed +// reservation refused every streamed URL before transport. +func TestReadVideo_GeminiURL_RejectsMissingContentLength(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + stubMediaProbes(t, 20, 1) + const total = int64(215_000) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Range") == "" { + // Flushing before the body makes Go answer chunked, which is what + // an origin with no declared length looks like. + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + w.Write(make([]byte, 16)) + return + } + start, end, ok := parseTestByteRange(r.Header.Get("Range"), total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(make([]byte, end-start+1)) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if err == nil || !strings.Contains(err.Error(), "does not support static streaming") { + t.Fatalf("expected the Content-Length check to reject, got: %v", err) + } +} + +// TestReadVideo_GeminiURL_RejectsOversizedStream covers the 2 GB stream ceiling. +// The clip measures 20 s, so the budget lets it through and this ceiling is +// what rejects the 3 GB the origin declares. +func TestReadVideo_GeminiURL_RejectsOversizedStream(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + stubMediaProbes(t, 20, 1) + + const total = int64(3_221_225_472) // 3 GB, over the 2 GB ceiling + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Range") == "" { + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + return + } + start, end, ok := parseTestByteRange(r.Header.Get("Range"), total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(make([]byte, 1024)) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if err == nil || !strings.Contains(err.Error(), "exceeds the maximum limit of 2 GB") { + t.Fatalf("expected the 2 GB ceiling to reject, got: %v", err) + } +} + +// TestReadVideo_ProbeVideoURL_RequestsHeadAndTailRanges proves the probe +// issues both windows against a range-capable server and uses a suffix spec +// for the tail, matching what the mediabudget estimator needs. +func TestReadVideo_ProbeVideoURL_RequestsHeadAndTailRanges(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + const total = int64(2_000_000) + + var ranges []string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rangeHeader := r.Header.Get("Range") + ranges = append(ranges, rangeHeader) + + start, end, ok := parseTestByteRange(rangeHeader, total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(make([]byte, end-start+1)) + })) + defer ts.Close() + + _, pinnedIP, err := security.Validate(ts.URL) + if err != nil { + t.Fatalf("validate test server URL: %v", err) + } + + payload, ok := probeVideoURL(context.Background(), pinnedIP, ts.URL, "video/mp4") + if !ok { + t.Fatal("expected probing to succeed against a range-capable server") + } + if payload.Size != total { + t.Fatalf("payload size = %d, want %d", payload.Size, total) + } + if len(ranges) != 2 { + t.Fatalf("got %d range requests, want 2: %v", len(ranges), ranges) + } + if ranges[0] != "bytes=0-524287" { + t.Errorf("head range = %q, want %q", ranges[0], "bytes=0-524287") + } + if ranges[1] != "bytes=-524288" { + t.Errorf("tail range = %q, want %q", ranges[1], "bytes=-524288") + } +} + +// TestReadVideo_ProbeVideoRange_BoundsReadOnUnsupportedRanges covers a server +// that answers 200 to a ranged request instead of 206: the probe must cap its +// read at the head window and close immediately rather than pull the whole +// declared body, even though the requested range (here, a suffix spec far +// bigger than the head window) would otherwise permit reading more. +func TestReadVideo_ProbeVideoRange_BoundsReadOnUnsupportedRanges(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + const declared = int64(8 * mediabudget.HeadBytes) + big := make([]byte, declared) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", strconv.FormatInt(declared, 10)) + w.WriteHeader(http.StatusOK) // ignores Range: no range support + w.Write(big) + })) + defer ts.Close() + + _, pinnedIP, err := security.Validate(ts.URL) + if err != nil { + t.Fatalf("validate test server URL: %v", err) + } + + rangeHeader := fmt.Sprintf("bytes=-%d", mediabudget.TailBytes) + data, total, isHead := probeVideoRange(context.Background(), pinnedIP, ts.URL, rangeHeader, declared) + if total != declared { + t.Fatalf("total = %d, want %d", total, declared) + } + if len(data) != mediabudget.HeadBytes { + t.Fatalf("read %d bytes, want exactly the head window (%d)", len(data), mediabudget.HeadBytes) + } + if !isHead { + t.Fatal("a 200 to a suffix range returned the start of the file and must be reported as head bytes") + } +} + +// TestReadVideo_GeminiURL_RejectsHugeMediaBeforeStreaming covers the front +// gate: a video the probe measures as an hour long, too much for the agent's +// window, is refused by the reservation itself, before the real streaming GET +// (the request with no Range header) is ever issued. +func TestReadVideo_GeminiURL_RejectsHugeMediaBeforeStreaming(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + stubMediaProbes(t, 3600, 1) + + const total = int64(3_221_225_472) // 3 GB + + streamHits := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rangeHeader := r.Header.Get("Range") + if rangeHeader == "" { + streamHits++ + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + return + } + start, end, ok := parseTestByteRange(rangeHeader, total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(make([]byte, end-start+1)) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if err == nil { + t.Fatal("expected the reservation to reject a video this large under a small window") + } + if !strings.Contains(err.Error(), "context window exceeded") { + t.Fatalf("expected a context window rejection, got: %v", err) + } + if streamHits != 0 { + t.Fatalf("the streaming GET reached the server %d time(s), want 0", streamHits) + } +} + +// stubGeminiVideoStream replaces the upload so a gate-passing URL can be +// followed all the way to transport without contacting a provider. +func stubGeminiVideoStream(tool *ReadVideoTool, record func(contentLength int64, mime string)) { + tool.setStreamToGeminiForTest(func(_ context.Context, _, _, _ string, body io.Reader, contentLength int64, mime string, _ time.Duration) (*providers.ChatResponse, error) { + io.Copy(io.Discard, body) + record(contentLength, mime) + return &providers.ChatResponse{Content: "a short clip"}, nil + }) +} + +// TestReadVideo_GeminiURL_MeasuredPayloadReachesTransportInBudget pins the +// rule that a measurement, once taken, is the charge: a 20-second clip costs +// 5,260 tokens and must reach transport under a 200k window, with the re-gate +// after the streaming response forbidden to re-derive anything from its +// Content-Length and overrule that. +func TestReadVideo_GeminiURL_MeasuredPayloadReachesTransportInBudget(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + stubMediaProbes(t, 20, 1) + + const total = int64(215_000) + body := make([]byte, total) + + var ranges []string + streamHits := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rangeHeader := r.Header.Get("Range") + if rangeHeader == "" { + streamHits++ + w.Header().Set("Content-Type", "video/mp4") + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + w.Write(body) + return + } + ranges = append(ranges, rangeHeader) + start, end, ok := parseTestByteRange(rangeHeader, total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(body[start : end+1]) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + var uploadedLength int64 + var uploadedMIME string + stubGeminiVideoStream(tool, func(contentLength int64, mime string) { + uploadedLength, uploadedMIME = contentLength, mime + }) + + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + out, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if err != nil { + t.Fatalf("a measured 20s clip must reach transport under a 200k window, got: %v", err) + } + if string(out) != "a short clip" { + t.Fatalf("provider output = %q, want the stubbed upload's answer", out) + } + if len(ranges) != 2 || ranges[0] != "bytes=0-524287" || ranges[1] != "bytes=-524288" { + t.Fatalf("unexpected range requests, want head then tail: %v", ranges) + } + if streamHits != 1 { + t.Fatalf("streaming GET reached %d time(s), want 1", streamHits) + } + if uploadedLength != total || uploadedMIME != "video/mp4" { + t.Fatalf("upload got %d bytes of %q, want %d of video/mp4", uploadedLength, uploadedMIME, total) + } +} + +// TestReadVideo_GeminiURL_RangeRefusingOriginBlockedBeforeStreaming is the +// regression test for the cap bypass. A WAF or CDN that answers 403 to any +// ranged read leaves both probes empty-handed, so nothing about this 1.9 GB +// payload can be measured and the streaming GET must never be issued at all. +func TestReadVideo_GeminiURL_RangeRefusingOriginBlockedBeforeStreaming(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + const total = int64(1_900_000_000) + + headHits, streamHits := 0, 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Range") != "" { + w.WriteHeader(http.StatusForbidden) + return + } + if r.Method == http.MethodHead { + headHits++ + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + return + } + streamHits++ + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if !errors.Is(err, mediabudget.ErrDurationUnmeasurable) { + t.Fatalf("expected ErrDurationUnmeasurable for an unmeasurable 1.9 GB payload, got %T: %v", err, err) + } + if headHits != 1 { + t.Fatalf("HEAD reached the origin %d time(s), want 1", headHits) + } + if streamHits != 0 { + t.Fatalf("the streaming request reached the origin %d time(s), want 0", streamHits) + } +} + +// TestReadVideo_GeminiURL_RefusesBeforeFetchingAnUnmeasurableOrigin covers an +// origin that reveals nothing before the stream: no ranged read, no HEAD. A +// payload with neither size nor bytes can never be priced, so the refusal +// belongs at the front gate, before a GET is made to an attacker-chosen URL. +func TestReadVideo_GeminiURL_RefusesBeforeFetchingAnUnmeasurableOrigin(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + + const total = int64(1_900_000_000) + + streamHits := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Range") != "" || r.Method == http.MethodHead { + w.WriteHeader(http.StatusForbidden) + return + } + streamHits++ + w.Header().Set("Content-Length", strconv.FormatInt(total, 10)) + w.WriteHeader(http.StatusOK) + w.Write(make([]byte, 4096)) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + uploads := 0 + stubGeminiVideoStream(tool, func(int64, string) { uploads++ }) + + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if !errors.Is(err, mediabudget.ErrNoProbeableBytes) { + t.Fatalf("expected the front gate to refuse an unpriceable payload, got %T: %v", err, err) + } + if streamHits != 0 { + t.Fatalf("the origin was fetched %d time(s) for a payload nothing could price, want 0", streamHits) + } + if uploads != 0 { + t.Fatalf("the body was forwarded to the upload %d time(s), want 0", uploads) + } +} + +// TestReadVideo_GeminiURL_StreamingGetSurfacesOriginStatus covers the +// streaming GET's non-2xx branch (read_video_resolve.go:274-280): an origin +// that answers both ranged probes with real, measurable bytes is priced and +// let through the front gate, and when the streaming GET itself then fails, +// the error must name the origin's status code. +func TestReadVideo_GeminiURL_StreamingGetSurfacesOriginStatus(t *testing.T) { + security.SetAllowLoopbackForTest(true) + defer security.SetAllowLoopbackForTest(false) + stubMediaProbes(t, 20, 1) + + const total = int64(215_000) + body := make([]byte, total) + + streamHits := 0 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rangeHeader := r.Header.Get("Range") + if rangeHeader == "" { + streamHits++ + w.WriteHeader(http.StatusInternalServerError) + return + } + start, end, ok := parseTestByteRange(rangeHeader, total) + if !ok { + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + return + } + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, total)) + w.WriteHeader(http.StatusPartialContent) + w.Write(body[start : end+1]) + })) + defer ts.Close() + + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + ctx := store.WithAgentContextWindow(context.Background(), 200_000) + ctx = store.WithAgentMaxTokens(ctx, 32_000) + + _, _, err := tool.callProvider(ctx, cp, "gemini", "gemini-2.5-flash", map[string]any{ + "prompt": "describe this video", "url": ts.URL, "_provider_type": "gemini", + }) + if err == nil || !strings.Contains(err.Error(), "status code 500") { + t.Fatalf("expected the streaming GET's status check to reject, got: %v", err) + } + if streamHits != 1 { + t.Fatalf("streaming GET reached %d time(s), want 1", streamHits) } } diff --git a/internal/tools/usage_caps.go b/internal/tools/usage_caps.go index 597a6ccb..ac6f3d59 100644 --- a/internal/tools/usage_caps.go +++ b/internal/tools/usage_caps.go @@ -2,16 +2,53 @@ package tools import ( "context" - "encoding/base64" + "errors" "fmt" - "strings" "github.com/google/uuid" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" ) +// budgetRefusalError marks a denial that came from the budget, not from a +// provider being unwell. A media fallback chain must stop on one: trying the +// next provider only moves the same payload to another endpoint, and the next +// entry may not price media at all. +type budgetRefusalError struct{ err error } + +func (e *budgetRefusalError) Error() string { return e.err.Error() } +func (e *budgetRefusalError) Unwrap() error { return e.err } + +// asBudgetRefusal marks err when it is one of the budget denials, and returns +// it untouched otherwise. Every gate on the media path funnels through here: +// the unmeasurable-media refusal, the agent context window, the missing-budget +// wiring failure, and the tenant usage cap. +func asBudgetRefusal(err error) error { + if err == nil { + return nil + } + var windowErr *usagecaps.ContextWindowExceededError + var wiringErr *usagecaps.AgentBudgetWiringError + switch { + case errors.Is(err, mediabudget.ErrDurationUnmeasurable), + errors.Is(err, usagecaps.ErrCapExceeded), + errors.As(err, &windowErr), + errors.As(err, &wiringErr): + return &budgetRefusalError{err: err} + default: + return err + } +} + +// isBudgetRefusal reports whether err is a budget denial, however deeply it has +// been wrapped on its way back up. +func isBudgetRefusal(err error) bool { + var refusal *budgetRefusalError + return errors.As(err, &refusal) +} + func agentBudgetFromContext(ctx context.Context) usagecaps.AgentBudget { return usagecaps.AgentBudget{ ContextWindow: store.AgentContextWindowFromContext(ctx), @@ -20,83 +57,57 @@ func agentBudgetFromContext(ctx context.Context) usagecaps.AgentBudget { } // reserveToolLLMUsage guards and reserves a tool-internal model call whose full -// input already lives in the ChatRequest (text plus the fixed inline-media -// convention). The fixed local BudgetCounter counts the request itself. +// input already lives in the ChatRequest. The fixed local BudgetCounter counts +// the request itself. func reserveToolLLMUsage(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest) (*usagecaps.Reservation, error) { - return reserveToolLLMUsageWithMedia(ctx, svc, toolName, providerName, model, req, "", nil) + return reserveToolLLMUsageWithMediaTokens(ctx, svc, toolName, providerName, model, req, 0) } // reserveToolLLMUsageWithMedia guards a native-media tool call whose payload is // sent OUT-OF-BAND (native provider JSON body or File API upload) and is thus -// invisible to the ChatRequest the counter would otherwise see. To keep the -// budget authority model/provider-independent, the fixed local BudgetCounter -// counts the standard-base64 representation of the raw bytes as one extra -// synthetic input message. That synthetic message is added ONLY to the -// guard/reservation copy — it is never transported, because native paths build -// their own provider payload from the raw bytes separately. -// -// mediaData must be the real bytes that will be sent. A native-media call whose -// payload cannot be buffered here (a streamed remote URL) must NOT route through -// this helper with nil data — use reserveToolLLMUsageUnverifiableMedia, which -// fails closed under an agent budget instead of undercounting. -func reserveToolLLMUsageWithMedia(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest, mediaMIME string, mediaData []byte) (*usagecaps.Reservation, error) { - if mediaMIME != "" || len(mediaData) > 0 { - req = appendNativeMediaBudget(req, mediaMIME, mediaData) +// invisible to the ChatRequest the counter would otherwise see. A payload +// mediabudget cannot price aborts the call rather than being waved through at +// an invented number. +func reserveToolLLMUsageWithMedia(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest, media ...mediabudget.Payload) (*usagecaps.Reservation, error) { + mediaTokens := 0 + for _, p := range media { + tokens, err := mediabudget.Tokens(ctx, p) + if err != nil { + return nil, asBudgetRefusal(err) + } + mediaTokens += tokens } + return reserveToolLLMUsageWithMediaTokens(ctx, svc, toolName, providerName, model, req, mediaTokens) +} + +// refuseUnpriceableMedia is the front gate for a payload whose cost cannot be +// established at all, used where there is not even a Payload to hand to +// mediabudget. It carries the same terminal marker as a priced refusal. +func refuseUnpriceableMedia(kind mediabudget.Kind, cause error) error { + return asBudgetRefusal(fmt.Errorf("%s budget: %w", kind, cause)) +} + +// reserveToolLLMUsageWithMediaTokens is the shared path for a media charge that is already a token count, not a mediabudget.Payload. +func reserveToolLLMUsageWithMediaTokens(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest, mediaTokens int) (*usagecaps.Reservation, error) { budget := agentBudgetFromContext(ctx) req = clampToolRequestMaxTokens(req, budget.MaxTokens) - if guardErr := usagecaps.GuardContextWindow(req, providerName, model, "tool:"+toolName, budget); guardErr != nil { - return nil, guardErr + if guardErr := usagecaps.GuardContextWindowWithMediaTokens(req, providerName, model, "tool:"+toolName, budget, mediaTokens); guardErr != nil { + return nil, asBudgetRefusal(guardErr) } if svc == nil { return nil, nil } - return svc.Preflight(ctx, usagecaps.Request{ - TenantID: store.TenantIDFromContext(ctx), - AgentID: store.AgentIDFromContext(ctx), - ProviderName: providerName, - ModelID: model, - ReservationKey: fmt.Sprintf("tool:%s:%s", toolName, uuid.NewString()), - Messages: req.Messages, - MaxOutputTokens: budget.MaxTokens, + reservation, err := svc.Preflight(ctx, usagecaps.Request{ + TenantID: store.TenantIDFromContext(ctx), + AgentID: store.AgentIDFromContext(ctx), + ProviderName: providerName, + ModelID: model, + ReservationKey: fmt.Sprintf("tool:%s:%s", toolName, uuid.NewString()), + Messages: req.Messages, + MaxOutputTokens: budget.MaxTokens, + ExtraInputTokens: mediaTokens, }) -} - -// reserveToolLLMUsageUnverifiableMedia handles a native-media call whose payload -// cannot be counted before transport (e.g. a remote video streamed straight to -// the provider without buffering). The complete-input invariant cannot be proven -// for such a call, so under an agent budget it fails closed with an explicit -// streaming error rather than trusting a byte-size/Content-Length estimate. -// Without a propagated agent budget it falls through to the text path, which -// itself fails closed with an AgentBudgetWiringError — every tool LLM call is -// agent-scoped, so there is no path here that reaches transport uncounted. -func reserveToolLLMUsageUnverifiableMedia(ctx context.Context, svc *usagecaps.Service, toolName, providerName, model string, req providers.ChatRequest) (*usagecaps.Reservation, error) { - budget := agentBudgetFromContext(ctx) - if budget.ContextWindow > 0 || budget.MaxTokens > 0 { - return nil, fmt.Errorf( - "tool:%s: cannot verify streamed native media against the agent context budget (no in-memory payload to count); refusing to send", - toolName, - ) - } - return reserveToolLLMUsage(ctx, svc, toolName, providerName, model, req) -} - -// appendNativeMediaBudget returns a copy of req with one extra synthetic user -// message carrying the media MIME and the standard-base64 encoding of the raw -// payload, so the fixed BudgetCounter counts the real out-of-band input. The -// original message slice is not mutated; the returned request is guard-only. -func appendNativeMediaBudget(req providers.ChatRequest, mime string, data []byte) providers.ChatRequest { - var b strings.Builder - b.WriteString(mime) - if len(data) > 0 { - b.WriteByte('\n') - b.WriteString(base64.StdEncoding.EncodeToString(data)) - } - msgs := make([]providers.Message, len(req.Messages), len(req.Messages)+1) - copy(msgs, req.Messages) - msgs = append(msgs, providers.Message{Role: "user", Content: b.String()}) - req.Messages = msgs - return req + return reservation, asBudgetRefusal(err) } // clampToolRequestMaxTokens enforces max_tokens <= agentMaxTokens for every diff --git a/internal/tools/usage_caps_service_test.go b/internal/tools/usage_caps_service_test.go new file mode 100644 index 00000000..9aacb5e9 --- /dev/null +++ b/internal/tools/usage_caps_service_test.go @@ -0,0 +1,127 @@ +package tools + +import ( + "context" + "database/sql" + + "github.com/google/uuid" + + "github.com/nextlevelbuilder/goclaw/internal/store" +) + +func int64Ptr(v int64) *int64 { return &v } + +// fakeToolUsageCapStore is the minimum of store.UsageCapStore a Preflight needs: +// the policies in force, and a ReserveUsage that enforces their token cap. +type fakeToolUsageCapStore struct { + policies []store.UsageCapPolicy + enforceTokenCaps bool + reserved store.UsageReserveRequest + events []store.UsageCapEvent +} + +func (s *fakeToolUsageCapStore) UpsertPricingCatalog(context.Context, []store.UsagePricingCatalogEntry) (int, error) { + return 0, nil +} + +func (s *fakeToolUsageCapStore) ListPricingCatalog(context.Context, store.UsagePricingQuery) ([]store.UsagePricingCatalogEntry, error) { + return nil, nil +} + +func (s *fakeToolUsageCapStore) PutPricingOverride(context.Context, *store.UsagePricingOverride) error { + return nil +} + +func (s *fakeToolUsageCapStore) ListPricingOverrides(context.Context, store.UsagePricingQuery) ([]store.UsagePricingOverride, error) { + return nil, nil +} + +func (s *fakeToolUsageCapStore) DeletePricingOverride(context.Context, uuid.UUID, uuid.UUID) error { + return nil +} + +func (s *fakeToolUsageCapStore) ResolvePricing(context.Context, uuid.UUID, uuid.UUID, string, string, string) (*store.ResolvedUsagePricing, error) { + return nil, sql.ErrNoRows +} + +func (s *fakeToolUsageCapStore) CreateUsageCapPolicy(context.Context, *store.UsageCapPolicy) error { + return nil +} + +func (s *fakeToolUsageCapStore) ListUsageCapPolicies(context.Context, store.UsageCapScope, bool) ([]store.UsageCapPolicy, error) { + return s.policies, nil +} + +func (s *fakeToolUsageCapStore) UpdateUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID, store.UsageCapPolicyPatch) (*store.UsageCapPolicy, error) { + return nil, nil +} + +func (s *fakeToolUsageCapStore) DeleteUsageCapPolicy(context.Context, uuid.UUID, uuid.UUID) error { + return nil +} + +func (s *fakeToolUsageCapStore) ReserveUsage(_ context.Context, req store.UsageReserveRequest, policies []store.UsageCapPolicy) (*store.UsageReservationResult, error) { + s.reserved = req + if s.enforceTokenCaps { + for _, p := range policies { + if p.MaxTokens != nil && req.EstimatedTokens > *p.MaxTokens { + return nil, &store.UsageCapExceededError{PolicyID: p.ID, Reason: "max_tokens"} + } + } + } + return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil +} + +func (s *fakeToolUsageCapStore) ReconcileUsage(context.Context, store.UsageReconcileRequest) error { + return nil +} + +func (s *fakeToolUsageCapStore) ListUsageCapUtilization(context.Context, uuid.UUID) ([]store.UsageCapUtilization, error) { + return nil, nil +} + +func (s *fakeToolUsageCapStore) ListUsageCapEvents(context.Context, uuid.UUID, int) ([]store.UsageCapEvent, error) { + return nil, nil +} + +func (s *fakeToolUsageCapStore) InsertUsageCapEvent(_ context.Context, event *store.UsageCapEvent) error { + if event != nil { + s.events = append(s.events, *event) + } + return nil +} + +// fakeToolProviderStore serves the provider metadata Preflight scopes a +// reservation by. Only the lookup by name is ever reached. +type fakeToolProviderStore struct { + provider *store.LLMProviderData +} + +func (s *fakeToolProviderStore) CreateProvider(context.Context, *store.LLMProviderData) error { + return nil +} + +func (s *fakeToolProviderStore) GetProvider(context.Context, uuid.UUID) (*store.LLMProviderData, error) { + return s.provider, nil +} + +func (s *fakeToolProviderStore) GetProviderByName(context.Context, string) (*store.LLMProviderData, error) { + if s.provider == nil { + return nil, sql.ErrNoRows + } + return s.provider, nil +} + +func (s *fakeToolProviderStore) ListProviders(context.Context) ([]store.LLMProviderData, error) { + return nil, nil +} + +func (s *fakeToolProviderStore) ListAllProviders(context.Context) ([]store.LLMProviderData, error) { + return nil, nil +} + +func (s *fakeToolProviderStore) UpdateProvider(context.Context, uuid.UUID, map[string]any) error { + return nil +} + +func (s *fakeToolProviderStore) DeleteProvider(context.Context, uuid.UUID) error { return nil } diff --git a/internal/tools/usage_caps_test.go b/internal/tools/usage_caps_test.go index 5844ffb1..1472e966 100644 --- a/internal/tools/usage_caps_test.go +++ b/internal/tools/usage_caps_test.go @@ -1,12 +1,13 @@ package tools import ( - "bytes" "context" "errors" "strings" "testing" + "time" + "github.com/nextlevelbuilder/goclaw/internal/mediabudget" "github.com/nextlevelbuilder/goclaw/internal/providers" "github.com/nextlevelbuilder/goclaw/internal/store" usagecaps "github.com/nextlevelbuilder/goclaw/internal/usage/caps" @@ -79,36 +80,38 @@ func TestReserveToolLLMUsage_FailsClosedWithoutAgentBudget(t *testing.T) { } } -// TestReserveToolLLMUsageWithMedia_CountsNativeMediaBytes proves the out-of-band -// media payload is part of completeInput. The fixed BudgetCounter counts the -// standard-base64 representation of the raw bytes (a model/provider-independent -// rule), so a small prompt with a large native-media payload is blocked when it -// no longer fits the CALLING agent's window. The bytes are counted, never -// transported — the guard copy carries them, the real request does not. -func TestReserveToolLLMUsageWithMedia_CountsNativeMediaBytes(t *testing.T) { +// stubMediaProbes fixes what the external binaries report, so a charge under test is the policy's arithmetic and not whatever is installed on the runner. +func stubMediaProbes(t *testing.T, seconds int, pages int) { + t.Helper() + t.Cleanup(mediabudget.SetProbesForTest( + func(context.Context, string) (time.Duration, bool) { + return time.Duration(seconds) * time.Second, true + }, + func(context.Context, string) (int, bool) { return pages, true }, + )) +} + +// An hour of video is 946,800 tokens at the published rate: far past a 20k +// window, comfortably inside a 2M one. +func bigVideoPayload() mediabudget.Payload { + return mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: "video/mp4", Size: 200_000_000, Path: "/does/not/matter"} +} + +func mediaTestRequest(model string) providers.ChatRequest { + return providers.ChatRequest{ + Model: model, + Messages: []providers.Message{{Role: "user", Content: "Describe this video."}}, + Options: map[string]any{providers.OptMaxTokens: 8}, + } +} + +func TestReserveToolLLMUsageWithMedia_LargePayloadTripsContextWindowGuard(t *testing.T) { + stubMediaProbes(t, 3600, 1) model := "gpt-4o" - req := func() providers.ChatRequest { - return providers.ChatRequest{ - Model: model, - Messages: []providers.Message{{Role: "user", Content: "Transcribe this."}}, - Options: map[string]any{providers.OptMaxTokens: 4096}, - } - } - - // A small media payload under a large agent window is allowed. - small := bytes.Repeat([]byte{0xAB}, 1024) - ctxBig := withToolAgentBudget(128_000, 8_192) - if _, err := reserveToolLLMUsageWithMedia(ctxBig, nil, "read_document", "openai", model, req(), "application/pdf", small); err != nil { - t.Fatalf("small native media must fit a 128k window, got %v", err) - } - - // A large media payload against a small agent window is blocked BEFORE - // transport: base64(256 KiB) is well over the 20k window's input cap. - big := bytes.Repeat([]byte{0xCD}, 256*1024) - ctxSmall := withToolAgentBudget(20_000, 8_192) - _, err := reserveToolLLMUsageWithMedia(ctxSmall, nil, "read_document", "openai", model, req(), "application/pdf", big) + ctxSmall := withToolAgentBudget(20_000, 8) + _, err := reserveToolLLMUsageWithMedia(ctxSmall, nil, "read_video", "openai", model, mediaTestRequest(model), bigVideoPayload()) if err == nil { - t.Fatal("expected abort: large native media must exceed a 20k agent window") + t.Fatal("expected abort: an hour of measured video must exceed a 20k window") } var ctxErr *usagecaps.ContextWindowExceededError if !errors.As(err, &ctxErr) { @@ -119,41 +122,124 @@ func TestReserveToolLLMUsageWithMedia_CountsNativeMediaBytes(t *testing.T) { } } -// TestReserveToolLLMUsageUnverifiableMedia_FailsClosedUnderAgentBudget proves a -// native-media call whose payload cannot be buffered (a streamed remote URL) -// fails closed under an agent budget rather than undercounting — the -// complete-input invariant cannot be proven, so it refuses to send. -func TestReserveToolLLMUsageUnverifiableMedia_FailsClosedUnderAgentBudget(t *testing.T) { - model := "gemini-2.0-flash" - req := providers.ChatRequest{ - Model: model, - Messages: []providers.Message{{Role: "user", Content: "Analyze this video."}}, - Options: map[string]any{providers.OptMaxTokens: 4096}, - } - - // Under an agent budget: no in-memory payload to count -> refuse with an - // explicit streaming error (not a wiring error). - ctx := withToolAgentBudget(200_000, 8_192) - _, err := reserveToolLLMUsageUnverifiableMedia(ctx, nil, "read_video", "gemini", model, req) - if err == nil { - t.Fatal("expected fail-closed for unverifiable streamed media under an agent budget") - } - var wiringErr *usagecaps.AgentBudgetWiringError - if errors.As(err, &wiringErr) { - t.Fatalf("under a valid agent budget the failure must be the streaming refusal, not a wiring error: %v", err) - } - if !strings.Contains(err.Error(), "refusing to send") { - t.Fatalf("expected an explicit streaming refusal, got %v", err) - } - - // Without a propagated agent budget it falls through to the text path, which - // itself fails closed with a wiring error — every tool LLM call is - // agent-scoped, so no path here reaches transport uncounted. - _, err = reserveToolLLMUsageUnverifiableMedia(context.Background(), nil, "read_video", "gemini", model, req) - if err == nil { - t.Fatal("expected fail-closed without a propagated agent budget") - } - if !errors.As(err, &wiringErr) { - t.Fatalf("expected *AgentBudgetWiringError without a budget, got %T: %v", err, err) +func TestReserveToolLLMUsageWithMedia_LargePayloadAllowedUnderLargeWindow(t *testing.T) { + stubMediaProbes(t, 3600, 1) + model := "gpt-4o" + ctxBig := withToolAgentBudget(2_000_000, 8) + if _, err := reserveToolLLMUsageWithMedia(ctxBig, nil, "read_video", "openai", model, mediaTestRequest(model), bigVideoPayload()); err != nil { + t.Fatalf("expected allow under a 2M window, got %v", err) + } +} + +// Valid because reserveToolLLMUsageWithMediaTokens feeds the same mediaTokens value to both the guard and Request.ExtraInputTokens. +func TestReserveToolLLMUsageWithMedia_ChargesSummedMediaBudgetEstimate(t *testing.T) { + stubMediaProbes(t, 12, 1) + model := "gpt-4o" + video := mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: "video/mp4", Size: 32_000, Path: "/does/not/matter"} + audio := mediabudget.Payload{Kind: mediabudget.KindAudio, MIME: "audio/mpeg", Size: 6_000, Path: "/does/not/matter"} + videoTokens, err := mediabudget.Tokens(context.Background(), video) + if err != nil { + t.Fatalf("pricing the video payload: %v", err) + } + audioTokens, err := mediabudget.Tokens(context.Background(), audio) + if err != nil { + t.Fatalf("pricing the audio payload: %v", err) + } + wantMediaTokens := videoTokens + audioTokens + + ctxTiny := withToolAgentBudget(1, 1) + + _, textOnlyErr := reserveToolLLMUsage(ctxTiny, nil, "read_video", "openai", model, mediaTestRequest(model)) + var textOnlyErrCtx *usagecaps.ContextWindowExceededError + if !errors.As(textOnlyErr, &textOnlyErrCtx) { + t.Fatalf("expected text-only call to abort under a 1-token window, got %v", textOnlyErr) + } + + _, mediaErr := reserveToolLLMUsageWithMedia(ctxTiny, nil, "read_video", "openai", model, mediaTestRequest(model), video, audio) + var mediaErrCtx *usagecaps.ContextWindowExceededError + if !errors.As(mediaErr, &mediaErrCtx) { + t.Fatalf("expected media call to abort under a 1-token window, got %v", mediaErr) + } + + gotMediaTokens := mediaErrCtx.InputTokens - textOnlyErrCtx.InputTokens + if gotMediaTokens != wantMediaTokens { + t.Fatalf("media charge reaching the guard = %d, want %d (sum of mediabudget.Tokens per payload)", gotMediaTokens, wantMediaTokens) + } +} + +// mediaChargeIsTheBlocker runs one payload through the media reservation under a +// window an identical text-only request passes, so a refusal can only come from +// the media charge. +func mediaChargeIsTheBlocker(t *testing.T, window int, payload mediabudget.Payload) { + t.Helper() + model := "gpt-4o" + ctx := withToolAgentBudget(window, 8) + + if _, err := reserveToolLLMUsage(ctx, nil, "read_media", "openai", model, mediaTestRequest(model)); err != nil { + t.Fatalf("the same request without media must pass under a %d-token window, got %v", window, err) + } + + _, err := reserveToolLLMUsageWithMedia(ctx, nil, "read_media", "openai", model, mediaTestRequest(model), payload) + var ctxErr *usagecaps.ContextWindowExceededError + if !errors.As(err, &ctxErr) { + t.Fatalf("expected *ContextWindowExceededError from the media charge, got %T: %v", err, err) + } + if ctxErr.ContextWindow != window { + t.Fatalf("guard used window %d, want %d", ctxErr.ContextWindow, window) + } +} + +// Two hours of measured podcast is 230,400 tokens, which no 200k window holds. +func TestReserveToolLLMUsageWithMedia_LargeAudioCannotProceed(t *testing.T) { + stubMediaProbes(t, 7200, 1) + mediaChargeIsTheBlocker(t, 200_000, mediabudget.Payload{Kind: mediabudget.KindAudio, MIME: "audio/mpeg", Size: 40_000_000, Path: "/does/not/matter"}) +} + +// A PDF whose page count cannot be measured is charged the provider's +// 1000-page maximum, and 258,000 tokens does not fit a 200k window. +func TestReserveToolLLMUsageWithMedia_UnmeasurablePDFCannotProceed(t *testing.T) { + mediaChargeIsTheBlocker(t, 200_000, mediabudget.Payload{Kind: mediabudget.KindDocument, MIME: "application/pdf", Size: 5_000_000}) +} + +// The same 1000-page ceiling must stay reachable: an agent whose window can +// hold 258,000 tokens still gets to read an unmeasurable PDF. +func TestReserveToolLLMUsageWithMedia_UnmeasurablePDFAllowedUnderLargeWindow(t *testing.T) { + model := "gpt-4o" + ctx := withToolAgentBudget(400_000, 8) + pdf := mediabudget.Payload{Kind: mediabudget.KindDocument, MIME: "application/pdf", Size: 5_000_000} + if _, err := reserveToolLLMUsageWithMedia(ctx, nil, "read_document", "openai", model, mediaTestRequest(model), pdf); err != nil { + t.Fatalf("expected the 1000-page ceiling to fit a 400k window, got %v", err) + } +} + +// A few hundred KB of docx is priced as text, well inside an ordinary window. +func TestReserveToolLLMUsageWithMedia_NonPDFDocumentFitsOrdinaryWindow(t *testing.T) { + model := "gpt-4o" + ctx := withToolAgentBudget(200_000, 8) + docx := mediabudget.Payload{ + Kind: mediabudget.KindDocument, + MIME: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + Size: 400_000, + } + if _, err := reserveToolLLMUsageWithMedia(ctx, nil, "read_document", "openai", model, mediaTestRequest(model), docx); err != nil { + t.Fatalf("expected a 400KB docx to pass under a 200k window, got %v", err) + } +} + +// A video nothing can measure is refused outright, with the missing binary +// named, rather than charged a number derived from its bytes. +func TestReserveToolLLMUsageWithMedia_UnmeasurableVideoIsRefused(t *testing.T) { + t.Cleanup(mediabudget.SetProbesForTest(nil, nil)) + + model := "gpt-4o" + ctx := withToolAgentBudget(2_000_000, 8) + video := mediabudget.Payload{Kind: mediabudget.KindVideo, MIME: "video/mp4", Size: 20_000, Path: "/does/not/matter"} + + _, err := reserveToolLLMUsageWithMedia(ctx, nil, "read_video", "openai", model, mediaTestRequest(model), video) + if !errors.Is(err, mediabudget.ErrDurationUnmeasurable) { + t.Fatalf("expected ErrDurationUnmeasurable even under a 2M window, got %T: %v", err, err) + } + if !strings.Contains(err.Error(), "ffprobe") { + t.Fatalf("error %q does not name the missing probe", err) } } diff --git a/internal/usage/caps/context_guard.go b/internal/usage/caps/context_guard.go index bd212658..94a936ce 100644 --- a/internal/usage/caps/context_guard.go +++ b/internal/usage/caps/context_guard.go @@ -53,6 +53,13 @@ func (e *ContextWindowExceededError) ContextBudgetExceeded() bool { return true // GuardContextWindow enforces completeInput + agentMaxTokens <= agentWindow. // Provider/model are retained only for diagnostics. func GuardContextWindow(req providers.ChatRequest, providerName, model, purpose string, budget AgentBudget) error { + return GuardContextWindowWithMediaTokens(req, providerName, model, purpose, budget, 0) +} + +// GuardContextWindowWithMediaTokens is GuardContextWindow for a call that also +// carries media out-of-band, where the payload never appears in the request the +// counter can see. mediaTokens is that payload's budget charge. +func GuardContextWindowWithMediaTokens(req providers.ChatRequest, providerName, model, purpose string, budget AgentBudget, mediaTokens int) error { if !budget.valid() { missing := "agent_context_window,agent_max_tokens" switch { @@ -67,6 +74,9 @@ func GuardContextWindow(req providers.ChatRequest, providerName, model, purpose if err != nil { return fmt.Errorf("count complete request: %w", err) } + if mediaTokens > 0 { + input += mediaTokens + } if input+budget.MaxTokens <= budget.ContextWindow { return nil } @@ -75,6 +85,7 @@ func GuardContextWindow(req providers.ChatRequest, providerName, model, purpose "model", model, "purpose", purpose, "input_tokens", input, + "media_tokens", mediaTokens, "output_reserve_tokens", budget.MaxTokens, "context_window", budget.ContextWindow, "action", "abort", diff --git a/internal/usage/caps/context_guard_test.go b/internal/usage/caps/context_guard_test.go index 268211f7..86c1e9b0 100644 --- a/internal/usage/caps/context_guard_test.go +++ b/internal/usage/caps/context_guard_test.go @@ -150,3 +150,32 @@ func TestCapsChat_AgentBudgetSameRequestDiffersByWindow(t *testing.T) { t.Fatal("200k agent: provider should have been called") } } + +// TestGuardContextWindowWithMediaTokens_AddsMediaToCompleteInput proves an +// out-of-band payload counts toward completeInput even though the ChatRequest +// never shows it. The same request passes without the media charge and fails +// with it. +func TestGuardContextWindowWithMediaTokens_AddsMediaToCompleteInput(t *testing.T) { + req := providers.ChatRequest{ + Model: "gemini-2.0-flash", + Messages: []providers.Message{{Role: "user", Content: "describe this video"}}, + Options: map[string]any{providers.OptMaxTokens: 1000}, + } + budget := AgentBudget{ContextWindow: 5_000, MaxTokens: 1_000} + + if err := GuardContextWindowWithMediaTokens(req, "gemini", req.Model, "tool:read_video", budget, 0); err != nil { + t.Fatalf("text-only request must fit a 5k window: %v", err) + } + + err := GuardContextWindowWithMediaTokens(req, "gemini", req.Model, "tool:read_video", budget, 4_500) + if err == nil { + t.Fatal("expected abort: 4500 media tokens plus a 1000 output reserve exceed a 5k window") + } + var exceeded *ContextWindowExceededError + if !errors.As(err, &exceeded) { + t.Fatalf("expected *ContextWindowExceededError, got %T: %v", err, err) + } + if exceeded.InputTokens < 4_500 { + t.Fatalf("InputTokens = %d, want at least the 4500 media tokens", exceeded.InputTokens) + } +} diff --git a/internal/usage/caps/service.go b/internal/usage/caps/service.go index 6af4c68a..ff8b8d67 100644 --- a/internal/usage/caps/service.go +++ b/internal/usage/caps/service.go @@ -41,6 +41,9 @@ type Request struct { ReservationKey string Messages []providers.Message MaxOutputTokens int + // ExtraInputTokens is input the Messages do not show, such as a media + // payload uploaded out-of-band. + ExtraInputTokens int } type Reservation struct { @@ -90,7 +93,7 @@ func (s *Service) Preflight(ctx context.Context, req Request) (*Reservation, err return skippedScopedReservation(req, scope, "no_policy"), nil } usage := pricing.BillableUsage{ - InputTokens: int64(EstimateInputTokens(req.Messages)), + InputTokens: int64(EstimateInputTokens(req.Messages) + req.ExtraInputTokens), OutputTokens: int64(req.MaxOutputTokens), ImageCount: int64(CountImages(req.Messages)), } diff --git a/internal/usage/caps/service_test.go b/internal/usage/caps/service_test.go index 6a597404..da8899db 100644 --- a/internal/usage/caps/service_test.go +++ b/internal/usage/caps/service_test.go @@ -424,6 +424,7 @@ type fakeUsageCapStore struct { resolveErr error resolveCalls int reserveErr error + enforceTokenCaps bool reserved store.UsageReserveRequest reconciled store.UsageReconcileRequest reconcileCalls int @@ -470,6 +471,13 @@ func (s *fakeUsageCapStore) ReserveUsage(_ context.Context, req store.UsageReser if s.reserveErr != nil { return nil, s.reserveErr } + if s.enforceTokenCaps { + for _, p := range policies { + if p.MaxTokens != nil && req.EstimatedTokens > *p.MaxTokens { + return nil, &store.UsageCapExceededError{PolicyID: p.ID, Reason: "max_tokens"} + } + } + } return &store.UsageReservationResult{ReservationKey: req.ReservationKey, Policies: policies}, nil } func (s *fakeUsageCapStore) ReconcileUsage(ctx context.Context, req store.UsageReconcileRequest) error { @@ -525,3 +533,88 @@ func (s *fakeProviderStore) UpdateProvider(context.Context, uuid.UUID, map[strin func (s *fakeProviderStore) DeleteProvider(context.Context, uuid.UUID) error { return nil } func int64Ptr(v int64) *int64 { return &v } + +// TestPreflightCountsExtraInputTokens proves an out-of-band media charge reaches +// the reservation, not just the context-window guard. Without it the reservation +// under-reserves by the whole payload. +func TestPreflightCountsExtraInputTokens(t *testing.T) { + providerID := uuid.New() + tenantID := uuid.New() + newService := func() (*Service, *fakeUsageCapStore) { + policy := store.UsageCapPolicy{ID: uuid.New(), TenantID: tenantID, MaxTokens: int64Ptr(1_000_000), Enabled: true} + usageStore := &fakeUsageCapStore{policies: []store.UsageCapPolicy{policy}, resolveErr: sql.ErrNoRows} + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: providerID}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + return NewService(usageStore, providerStore), usageStore + } + request := func(extra int) Request { + return Request{ + TenantID: tenantID, ProviderName: "openrouter", ModelID: "some/model", + ReservationKey: "extra-input", Messages: []providers.Message{{Role: "user", Content: "describe this video"}}, + MaxOutputTokens: 10, ExtraInputTokens: extra, + } + } + + svcBase, storeBase := newService() + if _, err := svcBase.Preflight(context.Background(), request(0)); err != nil { + t.Fatalf("Preflight without extra tokens returned error: %v", err) + } + + svcExtra, storeExtra := newService() + if _, err := svcExtra.Preflight(context.Background(), request(1600)); err != nil { + t.Fatalf("Preflight with extra tokens returned error: %v", err) + } + + if got := storeExtra.reserved.EstimatedTokens - storeBase.reserved.EstimatedTokens; got != 1600 { + t.Fatalf("EstimatedTokens grew by %d, want 1600", got) + } +} + +// TestPreflightDeniesWhenMediaChargeExceedsTokenCap proves an out-of-band media +// payload is refused by a tenant token cap the same request passes without it. +// Counting the charge is not enough on its own: what has to hold is that the +// call is denied before transport. +func TestPreflightDeniesWhenMediaChargeExceedsTokenCap(t *testing.T) { + providerID := uuid.New() + tenantID := uuid.New() + policyID := uuid.New() + newService := func() *Service { + policy := store.UsageCapPolicy{ID: policyID, TenantID: tenantID, MaxTokens: int64Ptr(1_000), Enabled: true} + usageStore := &fakeUsageCapStore{ + policies: []store.UsageCapPolicy{policy}, + resolveErr: sql.ErrNoRows, + enforceTokenCaps: true, + } + providerStore := &fakeProviderStore{provider: &store.LLMProviderData{ + BaseModel: store.BaseModel{ID: providerID}, + Name: "openrouter", + ProviderType: store.ProviderOpenRouter, + APIKey: "sk-test", + }} + return NewService(usageStore, providerStore) + } + request := func(extra int) Request { + return Request{ + TenantID: tenantID, ProviderName: "openrouter", ModelID: "some/model", + ReservationKey: "media-cap", Messages: []providers.Message{{Role: "user", Content: "describe this video"}}, + MaxOutputTokens: 10, ExtraInputTokens: extra, + } + } + + if _, err := newService().Preflight(context.Background(), request(0)); err != nil { + t.Fatalf("Preflight without the media charge must pass the 1000-token cap, got %v", err) + } + + reservation, err := newService().Preflight(context.Background(), request(5_000)) + if !errors.Is(err, ErrCapExceeded) { + t.Fatalf("Preflight with the media charge returned %v, want ErrCapExceeded", err) + } + metadata := reservation.TraceMetadata() + if metadata.Decision != store.UsageCapEventBlock { + t.Fatalf("Decision = %q, want block", metadata.Decision) + } +}