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