feat(a2asrv): add CostLimiter for budget-aware execution control by kuangmi-bit · Pull Request #378 · a2aproject/a2a-go

kuangmi-bit · 2026-07-18T22:03:37Z

Summary

Add a new a2asrv/costlimiter package that provides cost-based budget enforcement as a CallInterceptor.

Motivation

The existing limiter package provides concurrency control (MaxExecutions), but there is no built-in mechanism for cost-based quotas. This is useful for:

LLM token budgets: limit total tokens consumed by an agent over time
Compute-time quotas: fair scheduling when execution cost varies widely
Monetary credit limits: per-tenant billing enforcement

Design

CostLimiter: a server-side CallInterceptor. Before each request, it estimates the cost via a pluggable CostFunc and reserves budget via a CostStore. If budget is insufficient, it returns an error before the executor runs.
CostStore: an interface with Available/Reserve/Release. Negative availability = unlimited.
InMemoryCostStore: thread-safe in-memory implementation using sync.Mutex.
ScopeFunc: derives a scope (e.g., per-tenant, per-user) from CallContext for independent budget tracking.
Cancelled requests (where no real work was done) release their reserved budget.

Example

store := costlimiter.NewInMemoryCostStore(map[string]int64{"default": 1_000_000})
cl := costlimiter.NewCostLimiter(store,
    func(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) int64 {
        return estimateTokens(callCtx)
    },
)
handler := a2asrv.NewHandler(executor, a2asrv.WithCallInterceptors(cl))

Testing

All 8 tests pass, including:

Zero-cost passthrough
Within-budget / exceeds-budget
Unlimited budget
Release on cancellation
No release on success
Per-scope budgeting via ScopeFunc
Concurrent Store access

No existing tests affected (verified with go test ./a2asrv/...).

gemini-code-assist

Code Review

This pull request introduces the costlimiter package, which implements cost-based budget enforcement for A2A server request handling using a pluggable interceptor and budget store. The review feedback highlights several critical improvements: ensuring budget release operations do not fail on cancelled contexts by using context.WithoutCancel, implementing the documented fallback to limiter.AttachScope when no custom scope function is provided, and removing a redundant Available check in Before to prevent race conditions and double round-trips. Additionally, the feedback points out a bug in InMemoryCostStore.Reserve where storing negative costs for unlimited scopes can incorrectly restrict budgets, and suggests adding corresponding unit tests.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-18T22:06:00Z

+	if resp != nil && resp.Err != nil && errors.Is(resp.Err, context.Canceled) {
+		_ = cl.store.Release(ctx, res.scope, res.cost)
+	}


When a request is cancelled, the context ctx is already cancelled. If cl.store.Release performs any network or database I/O (e.g., in a distributed store implementation), passing the cancelled ctx will cause the release call to fail immediately, preventing the budget from being restored. Use context.WithoutCancel(ctx) to ensure the cleanup/release operation can complete successfully while preserving trace and logging context.

Suggested change

if resp != nil && resp.Err != nil && errors.Is(resp.Err, context.Canceled) {

_ = cl.store.Release(ctx, res.scope, res.cost)

}

if resp != nil && resp.Err != nil && errors.Is(resp.Err, context.Canceled) {

_ = cl.store.Release(context.WithoutCancel(ctx), res.scope, res.cost)

}

Done. After() now uses context.WithoutCancel(ctx) for the Release call when the request was cancelled.

gemini-code-assist · 2026-07-18T22:06:00Z

+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/a2aproject/a2a-go/v2/a2asrv"
+)


Import the limiter package to support retrieving the scope attached to the context via limiter.AttachScope.

Suggested change

import (

"context"

"errors"

"fmt"

"sync"

"github.com/a2aproject/a2a-go/v2/a2asrv"

)

import (

"context"

"errors"

"fmt"

"sync"

"github.com/a2aproject/a2a-go/v2/a2asrv"

"github.com/a2aproject/a2a-go/v2/a2asrv/limiter"

)

Done. limiter package is now imported.

gemini-code-assist · 2026-07-18T22:06:00Z

+func NewCostLimiter(store CostStore, costFn CostFunc, opts ...Option) *CostLimiter {
+	cl := &CostLimiter{
+		store:  store,
+		costFn: costFn,
+		scopeFn: func(_ *a2asrv.CallContext) string {
+			return "default"
+		},
+	}
+	for _, opt := range opts {
+		opt(cl)
+	}
+	return cl
+}


Do not initialize scopeFn to a default function in NewCostLimiter. Leaving it as nil allows us to detect if a custom scope function was provided, and fall back to the scope attached to the context via limiter.AttachScope as documented.

func NewCostLimiter(store CostStore, costFn CostFunc, opts ...Option) *CostLimiter { cl := &CostLimiter{ store: store, costFn: costFn, } for _, opt := range opts { opt(cl) } return cl }

Done. scopeFn is left nil in NewCostLimiter; fallback is handled in resolveScope().

gemini-code-assist · 2026-07-18T22:06:00Z

+func (cl *CostLimiter) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) {
+	scope := cl.scopeFn(callCtx)
+	cost := cl.costFn(ctx, callCtx, req)
+	if cost <= 0 {
+		return ctx, nil, nil
+	}
+
+	available, err := cl.store.Available(ctx, scope)
+	if err != nil {
+		return ctx, nil, fmt.Errorf("costlimiter: checking budget: %w", err)
+	}
+	if available >= 0 && available < cost {
+		return ctx, nil, fmt.Errorf("%w: scope %q: need %d, available %d",
+			ErrBudgetExceeded, scope, cost, available)
+	}
+
+	ok, err := cl.store.Reserve(ctx, scope, cost)
+	if err != nil {
+		return ctx, nil, fmt.Errorf("costlimiter: reserving budget: %w", err)
+	}
+	if !ok {
+		return ctx, nil, fmt.Errorf("%w: scope %q: budget exhausted", ErrBudgetExceeded, scope)
+	}
+
+	return context.WithValue(ctx, reservationKeyType{}, reservation{scope: scope, cost: cost}), nil, nil
+}


This refactoring addresses two key issues:

Fallback to limiter.AttachScope: Implements the documented fallback behavior when scopeFn is not provided, by checking limiter.ScopeFrom(ctx) before defaulting to "default".

Performance Optimization: Removes the redundant Available check before calling Reserve. Since Reserve atomically checks and deducts the budget, calling Available beforehand introduces an unnecessary double round-trip to the store (which is a major performance bottleneck for distributed stores like Redis or SQL) and creates a check-then-act race condition.

func (cl *CostLimiter) Before(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) (context.Context, any, error) { var scope string if cl.scopeFn != nil { scope = cl.scopeFn(callCtx) } else if sc, ok := limiter.ScopeFrom(ctx); ok { scope = sc } else { scope = "default" } cost := cl.costFn(ctx, callCtx, req) if cost <= 0 { return ctx, nil, nil } ok, err := cl.store.Reserve(ctx, scope, cost) if err != nil { return ctx, nil, fmt.Errorf("costlimiter: reserving budget: %w", err) } if !ok { return ctx, nil, fmt.Errorf("%w: scope %q: budget exhausted", ErrBudgetExceeded, scope) } return context.WithValue(ctx, reservationKeyType{}, reservation{scope: scope, cost: cost}), nil, nil }

Done. Added resolveScope() with fallback to limiter.AttachScope. Removed the redundant Available() check — Reserve() now called directly.

gemini-code-assist · 2026-07-18T22:06:00Z

+	if !ok {
+		s.budget[scope] = -cost
+		return true, nil
+	}


If a scope is not present in the map, it defaults to unlimited. Storing -cost in the map is problematic:

If cost is 0, it stores 0, which turns the unlimited scope into a limited scope with 0 budget (exhausted) on subsequent calls.

It causes unnecessary map growth and memory leaks for unlimited scopes.

Since unlimited scopes do not need to track or deduct budget, we should simply return true, nil without modifying the map.

if !ok { return true, nil }

Done. Reserve() returns true, nil for unlimited scopes without storing in the map. Added TestInMemoryCostStore_UnlimitedScopeNotStored to verify.

gemini-code-assist · 2026-07-18T22:06:00Z

+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/a2aproject/a2a-go/v2/a2asrv"
+)


Import the limiter package to support testing the fallback to limiter.AttachScope.

Suggested change

import (

"context"

"errors"

"testing"

"github.com/a2aproject/a2a-go/v2/a2asrv"

)

import (

"context"

"errors"

"testing"

"github.com/a2aproject/a2a-go/v2/a2asrv"

"github.com/a2aproject/a2a-go/v2/a2asrv/limiter"

)

Done. limiter package imported in test file.

gemini-code-assist · 2026-07-18T22:06:00Z

+	if !errors.Is(err, ErrBudgetExceeded) {
+		t.Fatalf("agent-a should be out of budget, got: %v", err)
+	}
+}


Add a unit test to verify that the CostLimiter correctly falls back to the scope attached to the context via limiter.AttachScope when no custom ScopeFunc is provided.

if !errors.Is(err, ErrBudgetExceeded) { t.Fatalf("agent-a should be out of budget, got: %v", err) } } func TestCostLimiter_FallbackToLimiterScope(t *testing.T) { store := NewInMemoryCostStore(map[string]int64{"custom-scope": 10}) cl := NewCostLimiter(store, func(ctx context.Context, callCtx *a2asrv.CallContext, req *a2asrv.Request) int64 { return 10 }) ctx := context.Background() ctx = limiter.AttachScope(ctx, "custom-scope") ctx, callCtx = a2asrv.NewCallContext(ctx, nil) ctx, _, err := cl.Before(ctx, callCtx, &a2asrv.Request{}) if err != nil { t.Fatalf("should have successfully reserved from custom-scope, got: %v", err) } // Second call should fail as budget is exhausted _, _, err = cl.Before(ctx, callCtx, &a2asrv.Request{}) if !errors.Is(err, ErrBudgetExceeded) { t.Fatalf("expected ErrBudgetExceeded, got: %v", err) } }

Done. Added TestCostLimiter_FallbackToLimiterScope verifying the AttachScope fallback.

Add a new a2asrv/costlimiter package that provides cost-based budget enforcement as a CallInterceptor. This complements the existing ConcurrencyConfig (which limits in-flight executions) with cost-aware budgeting — useful for LLM token budgets, compute-time quotas, or monetary credit limits. Key components: - CostLimiter: server-side CallInterceptor with pluggable CostFunc/CostStore - CostStore interface: Available/Reserve/Release with atomic semantics - InMemoryCostStore: thread-safe, unknown scopes treated as unlimited - ScopeFunc: per-agent/per-tenant scoping, falls back to limiter.AttachScope - Cancelled requests release budget via context.WithoutCancel - Removed redundant Available check before Reserve (race condition fix) Addresses review feedback: context.WithoutCancel on release, limiter.AttachScope fallback, InMemoryCostStore negative budget fix. 8 tests pass.

…ored tests - Import limiter package in test file - Add TestCostLimiter_FallbackToLimiterScope: verifies CostLimiter falls back to limiter.AttachScope when no custom ScopeFunc is provided - Add TestInMemoryCostStore_UnlimitedScopeNotStored: verifies unlimited scopes are not stored in the budget map

gemini-code-assist Bot reviewed Jul 18, 2026

View reviewed changes

kuangmi-bit force-pushed the feat/cost-limiter branch from 3edd3a0 to 82b30ff Compare July 19, 2026 12:20

邝谧 and others added 3 commits July 20, 2026 08:13

fix: remove ineffectual ctx assignments in costlimiter_test.go

fefbdb4

chore: re-trigger CI after GitHub outage

ceb7421

Uh oh!

Conversation

kuangmi-bit commented Jul 18, 2026

Summary

Motivation

Design

Example

Testing

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

kuangmi-bit Jul 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants