From ce9e7af3fbee417417fe47c09e788793c8c401bf Mon Sep 17 00:00:00 2001 From: thotam <78630736+thotam@users.noreply.github.com> Date: Sat, 12 Sep 2026 10:31:12 +0700 Subject: [PATCH] fix(tools): price out-of-band media by measured duration, not by its bytes (#1564) read_video called with a url parameter always failed under an agent budget: tool:read_video: cannot verify streamed native media against the agent context budget (no in-memory payload to count); refusing to send read_audio and read_document appeared to work but only for very small files. Both symptoms share one root cause. 5ca433b8 introduced the complete-input invariant and, to keep out-of-band media honest, appended the standard-base64 encoding of the payload as a synthetic guard-only message. That message is counted as text. Measured against the bundled BudgetCounter, base64 costs 0.956 tokens per raw byte, so an 11.6 MB video counted 11,161,058 tokens and exceeded a 200k, a 1M and a 2M window alike. The practical ceiling was roughly 200 KB at a 200k window while videoMaxBytes is 100 MB. The read_video URL transport streams to the File API without buffering, so it had no bytes to count at all and was routed to a helper that refused to send whenever an agent budget was present. Media is now priced by what a provider actually bills for it. Every number below is either a published per-unit rate or a published provider limit; none is derived from a byte size, because a static-image video compresses arbitrarily small and no bitrate floor exists. An earlier revision of this branch tried one and a 120-second, 20,627-byte clip priced at 526 tokens against a true 31,560. internal/mediabudget prices one payload: - Video and audio: ffprobe measures duration. 263 tokens/second for video, 32 for audio, both published for static processing at 1 FPS. - PDF: pdfinfo counts pages at 258 tokens/page. When the page count cannot be read, the charge is the proven ceiling of 1000 pages, which is the most a provider will accept and therefore the most it can bill. - Video and audio that cannot be measured are refused. Upstream already refused unverifiable native media for every URL on the Gemini streamed path; this narrows that refusal from every URL to only what genuinely cannot be measured, rather than removing it. PDF differs because its page count is cheaply measurable and its ceiling is small enough to stay usable. A remote video is measured without downloading it: two ranged GETs, 512 KB from the head and 512 KB from the tail, written into a sparse temp file sized to the declared total and handed to ffprobe. The tail matters because every container that puts its index at the end keeps it there: head-only probing under-reports mpeg by 98% and ogg by 79%, and a plain prefix makes ffprobe under-report a 30-second WAV as 0.74 seconds because it clamps to the bytes it can see. Sizing the temp file to the real total fixes that. Verified end to end on an 11,673,105 byte MP4 served by nginx: 1 MB of ranged reads yielded duration 40.000000, identical to ffprobe reading the whole URL, for a charge of 10,520 tokens. Those requests reuse the existing SSRF-safe path, security.WithPinnedIP plus security.NewSafeClient(0), and no URL is ever handed to an external binary. Beyond the reported bug, two pre-existing gaps let large media reach a provider almost unpriced. ExecuteWithChain treated every callProvider error as a provider failure and advanced to the next entry, and the non-Gemini branches of read_video and read_document reserved without pricing their payload at all. Measured under a 20,000-token window before this change, a 40 MB video and a 1000-page PDF each reached a provider charged about 1,600 tokens. Budget refusals are now terminal in the chain and every media branch prices its payload, so both reach no provider at all. Genuine provider failures still fail over. Known limits, stated rather than discovered: - The /Type /Page scan that guards against a forged /Count is a floor, not a bound. Pages inside a compressed object stream are invisible to it, and a 9,484-byte PDF built that way is charged 258 tokens for 1000 pages. A real pdfinfo reads such files correctly; the scan only ever raises a probed count. - A video URL whose origin does not serve byte ranges is now refused on the non-Gemini path too, and the refusal is terminal. Upstream forwarded such URLs unpriced. A HEAD giving only a size is not enough to price one. - read_video and read_audio require ffprobe. Docker images install it by default except the base variant; bare binaries and the desktop build do not ship it. - A hostile origin can craft a container ffprobe reads as about one second. Reservation.Reconcile overwrites the estimate with the provider's reported usage, so this weakens the gate rather than defeating it. Byte ceilings videoMaxBytes, audioMaxBytes and documentMaxBytes are unchanged. No new module dependency; ffprobe and pdfinfo are optional runtime probes. --- .github/workflows/dev-beta-release.yaml | 3 + .github/workflows/release-beta.yaml | 3 + .github/workflows/release.yaml | 4 + Dockerfile | 11 + README.md | 2 +- docs/14-skills-runtime.md | 15 +- internal/mediabudget/mediabudget.go | 574 +++++++++++++++++++++ internal/mediabudget/mediabudget_test.go | 508 ++++++++++++++++++ internal/tokencount/budget_counter.go | 9 +- internal/tokencount/budget_counter_test.go | 51 ++ internal/tools/media_budget_chain_test.go | 256 +++++++++ internal/tools/media_provider_chain.go | 11 + internal/tools/read_audio.go | 4 + internal/tools/read_audio_resolve.go | 8 +- internal/tools/read_document.go | 4 + internal/tools/read_document_resolve.go | 10 +- internal/tools/read_video.go | 22 +- internal/tools/read_video_resolve.go | 161 +++++- internal/tools/read_video_test.go | 517 ++++++++++++++++++- internal/tools/usage_caps.go | 143 ++--- internal/tools/usage_caps_service_test.go | 127 +++++ internal/tools/usage_caps_test.go | 216 +++++--- internal/usage/caps/context_guard.go | 11 + internal/usage/caps/context_guard_test.go | 29 ++ internal/usage/caps/service.go | 5 +- internal/usage/caps/service_test.go | 93 ++++ 26 files changed, 2621 insertions(+), 176 deletions(-) create mode 100644 internal/mediabudget/mediabudget.go create mode 100644 internal/mediabudget/mediabudget_test.go create mode 100644 internal/tokencount/budget_counter_test.go create mode 100644 internal/tools/media_budget_chain_test.go create mode 100644 internal/tools/usage_caps_service_test.go 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) + } +}