From a6bcb177fe07908a3130a674c6d2122e9d445ab7 Mon Sep 17 00:00:00 2001
From: mlsmaycon <mlsmaycon@gmail.com>
Date: Sat, 22 Aug 2026 22:00:39 +0000
Subject: [PATCH] [misc] Bound each access-log poll, and assert the streamed
 total
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two review findings on the e2e suite.

lookupAccessLogBySession polled under the caller's context, so one stalled
request could hold the loop open well past the 30s ingest window it exists to
enforce — and the caller would read that delay as a missing row rather than a
slow one. Each poll now expires with the window; the parent's cancellation
still applies, since the request context derives from it.

The streaming test asserted only that the total cost was positive. Input and
output are positive on their own, so a cache-read bucket that was parsed and
then never billed would have passed. Assert the sum of the three buckets: the
gap a dropped cache read leaves is 7e-6, well outside the delta.
---
 e2e/agentnetwork/custom_pricing_test.go | 23 +++++++++++++++++++----
 e2e/agentnetwork/streaming_test.go      |  7 ++++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/e2e/agentnetwork/custom_pricing_test.go b/e2e/agentnetwork/custom_pricing_test.go
index 7c13cbdde..fded58304 100644
--- a/e2e/agentnetwork/custom_pricing_test.go
+++ b/e2e/agentnetwork/custom_pricing_test.go
@@ -181,24 +181,39 @@ const accessLogIngestWindow = 30 * time.Second
 func lookupAccessLogBySession(ctx context.Context, sessionID string, within time.Duration) (api.AgentNetworkAccessLog, bool) {
 	deadline := time.Now().Add(within)
 	for {
-		if logs, lerr := srv.ListAccessLogs(ctx); lerr == nil {
+		// Each poll is bounded by what is left of the window rather than by the
+		// caller's context: a single stalled request would otherwise hold the
+		// loop open long past the ingest window it is meant to enforce, and the
+		// caller would read the delay as a missing row.
+		if logs, lerr := listAccessLogsBy(ctx, deadline); lerr == nil {
 			for _, r := range logs.Data {
 				if r.SessionId != nil && *r.SessionId == sessionID {
 					return r, true
 				}
 			}
 		}
-		if time.Now().After(deadline) {
-			return api.AgentNetworkAccessLog{}, false
-		}
 		select {
 		case <-ctx.Done():
 			return api.AgentNetworkAccessLog{}, false
 		case <-time.After(2 * time.Second):
 		}
+		// Checked after the wait rather than before the request: a poll issued
+		// past the deadline carries no budget and would fail on arrival.
+		if !time.Now().Before(deadline) {
+			return api.AgentNetworkAccessLog{}, false
+		}
 	}
 }
 
+// listAccessLogsBy fetches one access-log page under a context that expires at
+// deadline, so no single call can outlive the window its caller is polling
+// within. The parent's cancellation still applies: the child inherits it.
+func listAccessLogsBy(ctx context.Context, deadline time.Time) (api.AgentNetworkAccessLogsResponse, error) {
+	reqCtx, cancel := context.WithDeadline(ctx, deadline)
+	defer cancel()
+	return srv.ListAccessLogs(reqCtx)
+}
+
 // findAccessLogBySession polls the access-log page for the row carrying
 // sessionID, failing the test if it never lands. Use it for a request whose row
 // must exist; where a missing row is a recoverable race, use
diff --git a/e2e/agentnetwork/streaming_test.go b/e2e/agentnetwork/streaming_test.go
index 2f7e9e3e2..4637da10c 100644
--- a/e2e/agentnetwork/streaming_test.go
+++ b/e2e/agentnetwork/streaming_test.go
@@ -66,9 +66,14 @@ func TestStreamingResponseMetersInputTokens(t *testing.T) {
 	// prices the full input count rather than a remainder.
 	wantInput := float64(harness.VLLMStreamInputTokens) / 1000 * streamInRate
 	wantOutput := float64(harness.VLLMStreamOutputTokens) / 1000 * streamOutRate
+	wantCacheRead := float64(harness.VLLMStreamCacheReadTokens) / 1000 * streamCacheReadRate
 	assert.InDelta(t, wantInput, row.InputCostUsd, 1e-6, "input cost must price the streamed input tokens")
 	assert.InDelta(t, wantOutput, row.OutputCostUsd, 1e-6, "output cost must price the streamed output tokens")
-	assert.Greater(t, row.CostUsd, 0.0, "a streamed request must never record as free")
+	// The total, not merely a positive number: input and output alone are
+	// positive, so a cache bucket parsed and then never billed would pass any
+	// weaker assertion. The gap is 7e-6, well outside the delta.
+	assert.InDelta(t, wantInput+wantOutput+wantCacheRead, row.CostUsd, 1e-6,
+		"the recorded cost must be every bucket the surface bills, cache reads included")
 }
 
 // TestStreamingOnGatewayTypedProvider drives the same streamed Anthropic call