feat(agents): /v1/agents per-org fail-closed metering + long-running scheduler
release / build-amd64 (push) Failing after 7m11s
release / notify-universe (push) Skipped

Lands the agent-backend metering feature onto the zip->zap-proto-migrated main
(cloud is already fully migrated on main; 8 subsystem pins + MountAll clean).

- /v1/agents/* self-meters a per-run fee to commerce: pre-authorize the org's
  prepaid credit balance fail-closed (402 insufficient_balance), debit on success,
  attributed to product "agent". Added to selfMeteredPrefixes so the edge gate
  never double-bills. Run path (money-moving) requires a VALIDATED principal
  (c.User() non-empty), refusing the no-bearer forge path; scheduled runs carry an
  unforgeable 'scheduler'-prefixed actor.
- ResourceMeter.Gate now forwards costCents as AuthInput.AmountCents so the gate
  enforces available >= fee (not merely > 0) — a 1-cent balance can no longer
  authorize a run that takes the ledger negative. MeterUsage generalizes the
  per-org debit (Actor/Model/token attribution) while Meter keeps its signature.
- Long-running agents: cron scheduler scans once a minute over a partial index
  (ix_agents_scheduled), bounded per-org cap (CLOUD_AGENT_MAX_LONG_RUNNING);
  scheduler.stop drains in-flight runs inside the SIGTERM budget via ShutdownAll.

Deps: commerce/metering v0.1.0 -> v0.1.2 (Actor field + AuthInput.AmountCents).
All other pins inherited from migrated main (ai v1.789.1, authz v1.10.3,
base v1.4.6, commerce v1.42.29, licensing v0.1.1, metrics v0.4.1, o11y v1.3.12,
vfs v0.4.4). hanzoai/zip stays out of the graph; zip == zap-proto/zip v1.2.0.

go mod verify clean; go build ./... EXIT 0; MountAll boot-smoke clean (no
want *zip.App); agents -race + root/ml/provisioning billing tests green.
This commit is contained in:
2026-07-02 12:44:31 -07:00
parent b4322fa7ac
commit fe3a5fd553
20 changed files with 2000 additions and 97 deletions
+37 -3
View File
@@ -1,6 +1,7 @@
package cloud
import (
"context"
"fmt"
"github.com/hanzoai/commerce/metering"
@@ -248,13 +249,20 @@ func pickVaultClient(cfg *Config, log luxlog.Logger) VaultClient {
// calls it.
type MountFunc func(app any, deps Deps) error // app is *zip.App; using any here to avoid an import cycle in pkg/cloud
// ShutdownFunc releases a subsystem's process-lifetime resources (background
// goroutines, open DB handles) on graceful shutdown. It must be idempotent and
// bounded — Serve calls it within the shutdown deadline. ctx carries that
// deadline so a slow teardown is cut off rather than hanging SIGTERM.
type ShutdownFunc func(ctx context.Context) error
// MountSpec describes one subsystem registered for mounting. The Order
// is used when ordering matters for inter-subsystem deps (e.g. iam
// before authz before commerce).
type MountSpec struct {
Name string
Order int
Mount MountFunc
Name string
Order int
Mount MountFunc
Shutdown ShutdownFunc // optional; nil means the subsystem has nothing to tear down.
}
// Registry is the in-process subsystem registry. Subsystems register via
@@ -267,6 +275,32 @@ func Register(name string, order int, mount MountFunc) {
Registry = append(Registry, MountSpec{Name: name, Order: order, Mount: mount})
}
// RegisterWithShutdown adds a subsystem that owns process-lifetime resources: a
// background worker (e.g. the agents scheduler) or a DB handle that must be
// flushed. shutdown is invoked by ShutdownAll on graceful stop. This is the ONE
// way a subsystem gets a teardown — Register stays the zero-teardown default.
func RegisterWithShutdown(name string, order int, mount MountFunc, shutdown ShutdownFunc) {
Registry = append(Registry, MountSpec{Name: name, Order: order, Mount: mount, Shutdown: shutdown})
}
// ShutdownAll tears down every ENABLED subsystem that registered a ShutdownFunc,
// in REVERSE mount order (a dependency is torn down after its dependents), best
// effort: a failure is collected and the rest still run, so one stuck subsystem
// can't strand another's flush. Serve calls this inside the shutdown deadline.
func ShutdownAll(ctx context.Context, cfg *Config) error {
var firstErr error
for i := len(Registry) - 1; i >= 0; i-- {
spec := Registry[i]
if spec.Shutdown == nil || !cfg.Enabled(spec.Name) {
continue
}
if err := spec.Shutdown(ctx); err != nil && firstErr == nil {
firstErr = fmt.Errorf("shutdown %s: %w", spec.Name, err)
}
}
return firstErr
}
// MountAll iterates the registry in order and calls Mount() on each
// enabled subsystem.
func MountAll(app any, cfg *Config, deps Deps) error {
+295 -32
View File
@@ -36,8 +36,9 @@ import (
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/clients/principal"
"github.com/hanzoai/cloud/types"
"github.com/zap-proto/zip"
"github.com/hanzoai/commerce/metering"
luxlog "github.com/luxfi/log"
"github.com/zap-proto/zip"
)
// nameRE is the org-unique handle AND the URL path segment — the traversal
@@ -47,12 +48,50 @@ var nameRE = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$`)
const (
maxInstructions = 32 * 1024 // system prompt cap
maxInput = 128 * 1024
// maxRef bounds the free-text bot-lifecycle references (compute machine id,
// service-account id). They are opaque identifiers, not documents — a
// generous 256 keeps a client from bloating the per-org SQLite with a
// multi-megabyte "id".
maxRef = 256
// agentFeeEnvPrefix is the operator knob for the flat per-run fee. The
// effective fee is cloud.ResourceFeeCents(agentFeeEnvPrefix, meterKind): a
// global CLOUD_AGENT_FEE_CENTS override wins over the $1.00 default; set it
// to 0 to make agent runs free (and therefore un-gated). This is a per-RUN
// fee — the honest, policy-set unit an agent run bills. Token-based pricing
// is intentionally NOT used here: the in-process AIClient returns only the
// completion content (types.ChatResponse{Content}), no token counts, so
// charging per-token would be fabricated. Duration is recorded on the run.
agentFeeEnvPrefix = "CLOUD_AGENT_FEE_CENTS"
// meterKind is the commerce "provider"/attribution label for agent spend —
// the task's product:"agent". One value so every agent run (HTTP or
// scheduled) is attributed identically.
meterKind = "agent"
// schedulerActor is the Actor recorded on a scheduled run that has no IAM
// service account bound. Real service-account identity (the keystone) rides
// in Agent.ServiceAccountID when present.
schedulerActor = "scheduler"
// maxLongRunningPerOrg caps an org's scheduler footprint: how many scheduled
// long-running agents it may create. Each scheduled agent adds recurring
// gate+run+debit load to the shared store, so a per-org bound stops one
// tenant from self-amplifying the once-a-minute scan. Overridable by ops via
// CLOUD_AGENT_MAX_LONG_RUNNING.
maxLongRunningPerOrg = 100
longRunningCapEnv = "CLOUD_AGENT_MAX_LONG_RUNNING"
)
type svc struct {
store *Store
ai types.AIClient
log luxlog.Logger
// bill is the shared per-org gate+meter (reuses deps.Metering, the ONE
// commerce client — the same object ml/provisioning use). Nil/!Enabled()
// makes Gate allow and Meter a no-op, so an unconfigured deployment runs
// agents without billing rather than failing closed on a missing ledger.
bill *cloud.ResourceMeter
// sched is the long-running-agent scheduler; nil until started, stopped on
// Shutdown. It shares svc so it runs agents through the SAME runAgent path.
sched *scheduler
}
var mounted *svc
@@ -60,15 +99,19 @@ var mounted *svc
// ---- HTTP response shapes (the published contract) ----
type agentView struct {
ID string `json:"id"`
Name string `json:"name"`
Model string `json:"model"`
Description string `json:"description,omitempty"`
Tools []string `json:"tools"`
Status string `json:"status"`
Runs int `json:"runs"`
CreatedAt string `json:"createdAt"`
UpdatedAt string `json:"updatedAt"`
ID string `json:"id"`
Name string `json:"name"`
Model string `json:"model"`
Description string `json:"description,omitempty"`
Tools []string `json:"tools"`
Status string `json:"status"`
ExecutionMode string `json:"executionMode"`
Schedule string `json:"schedule,omitempty"`
ComputeRef string `json:"computeRef,omitempty"`
ServiceAccountID string `json:"serviceAccountId,omitempty"`
Runs int `json:"runs"`
CreatedAt string `json:"createdAt"`
UpdatedAt string `json:"updatedAt"`
}
type agentDetail struct {
@@ -98,7 +141,10 @@ func rfc3339(unix int64) string {
func toView(a Agent, runs int) agentView {
return agentView{
ID: a.ID, Name: a.Name, Model: a.Model, Description: a.Description,
Tools: nonNil(a.Tools), Status: a.Status, Runs: runs,
Tools: nonNil(a.Tools), Status: a.Status,
ExecutionMode: a.ExecutionMode, Schedule: a.Schedule,
ComputeRef: a.ComputeRef, ServiceAccountID: a.ServiceAccountID,
Runs: runs,
CreatedAt: rfc3339(a.CreatedAt), UpdatedAt: rfc3339(a.UpdatedAt),
}
}
@@ -138,7 +184,12 @@ func Mount(app *zip.App, deps cloud.Deps) error {
return fmt.Errorf("agents.Mount: open store: %w", err)
}
// deps.AI may be nil when no gateway is configured; run() degrades honestly.
s := &svc{store: store, ai: deps.AI, log: log}
s := &svc{
store: store,
ai: deps.AI,
log: log,
bill: cloud.NewResourceMeter(deps, meterKind),
}
mounted = s
app.Get("/v1/agents", s.list)
@@ -149,28 +200,47 @@ func Mount(app *zip.App, deps cloud.Deps) error {
app.Post("/v1/agents/:name/run", s.run)
app.Get("/v1/agents/:name/runs", s.runs)
log.Info("agents mounted", "ai", s.ai != nil, "brand", deps.Brand)
// Long-running scheduler: invokes each long-running agent's run on its cron
// cadence through the SAME runAgent path as the HTTP handler (one run path,
// one gate, one meter). Only started when inference is wired — with no AI a
// scheduled run could never execute, so there is nothing to schedule.
if s.ai != nil {
s.sched = newScheduler(s, log)
s.sched.start()
}
log.Info("agents mounted", "ai", s.ai != nil, "billing", s.bill.Enabled(),
"scheduler", s.sched != nil, "brand", deps.Brand)
return nil
}
func init() {
cloud.Register("agents", 127, func(app any, deps cloud.Deps) error {
cloud.RegisterWithShutdown("agents", 127, func(app any, deps cloud.Deps) error {
a, ok := app.(*zip.App)
if !ok {
return fmt.Errorf("agents.Mount: app is %T, want *zip.App", app)
}
return Mount(a, deps)
}, func(ctx context.Context) error {
// Graceful teardown: stop the scheduler (drain in-flight runs) and close
// the store. Bounded by the caller's shutdown deadline so a stuck run
// can't hang SIGTERM.
return Shutdown(ctx)
})
}
// ---- handlers ----
type createReq struct {
Name string `json:"name"`
Model string `json:"model"`
Instructions string `json:"instructions"`
Description string `json:"description"`
Tools []string `json:"tools"`
Name string `json:"name"`
Model string `json:"model"`
Instructions string `json:"instructions"`
Description string `json:"description"`
Tools []string `json:"tools"`
ExecutionMode string `json:"executionMode"`
Schedule string `json:"schedule"`
ComputeRef string `json:"computeRef"`
ServiceAccountID string `json:"serviceAccountId"`
}
func (s *svc) create(c *zip.Ctx) error {
@@ -196,6 +266,31 @@ func (s *svc) create(c *zip.Ctx) error {
if len(body.Instructions) > maxInstructions {
return zip.ErrBadRequest("instructions too large")
}
mode, schedule, err := validateLifecycle(body.ExecutionMode, body.Schedule)
if err != nil {
return err
}
computeRef, err := validateRef("computeRef", body.ComputeRef)
if err != nil {
return err
}
serviceAccountID, err := validateRef("serviceAccountId", body.ServiceAccountID)
if err != nil {
return err
}
// Cap the org's scheduler footprint (Red LOW-1): a tenant cannot create an
// unbounded number of scheduled agents that each add recurring load to the
// shared store. Only counts when this create is itself long-running.
if mode == ModeLongRunning {
n, err := s.store.CountLongRunning(c.Context(), org)
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "count: %v", err)
}
if n >= longRunningCap() {
return zip.Errorf(http.StatusConflict,
"long-running agent limit reached for this org (max %d)", longRunningCap())
}
}
id, err := genID("agent")
if err != nil {
return zip.Errorf(http.StatusInternalServerError, "rng: %v", err)
@@ -204,7 +299,9 @@ func (s *svc) create(c *zip.Ctx) error {
a := Agent{
ID: id, Org: org, Name: name, Model: model, Instructions: body.Instructions,
Description: strings.TrimSpace(body.Description), Tools: cleanList(body.Tools),
Status: "ready", CreatedAt: now, UpdatedAt: now,
Status: "ready", ExecutionMode: mode, Schedule: schedule,
ComputeRef: computeRef, ServiceAccountID: serviceAccountID,
CreatedAt: now, UpdatedAt: now,
}
if err := s.store.Create(c.Context(), a); err != nil {
if err == errConflict {
@@ -262,10 +359,14 @@ func (s *svc) get(c *zip.Ctx) error {
}
type updateReq struct {
Model *string `json:"model"`
Instructions *string `json:"instructions"`
Description *string `json:"description"`
Tools *[]string `json:"tools"`
Model *string `json:"model"`
Instructions *string `json:"instructions"`
Description *string `json:"description"`
Tools *[]string `json:"tools"`
ExecutionMode *string `json:"executionMode"`
Schedule *string `json:"schedule"`
ComputeRef *string `json:"computeRef"`
ServiceAccountID *string `json:"serviceAccountId"`
}
func (s *svc) update(c *zip.Ctx) error {
@@ -304,6 +405,45 @@ func (s *svc) update(c *zip.Ctx) error {
if body.Tools != nil {
a.Tools = cleanList(*body.Tools)
}
if body.ComputeRef != nil {
if a.ComputeRef, err = validateRef("computeRef", *body.ComputeRef); err != nil {
return err
}
}
if body.ServiceAccountID != nil {
if a.ServiceAccountID, err = validateRef("serviceAccountId", *body.ServiceAccountID); err != nil {
return err
}
}
// Re-validate the lifecycle from the RESULTING mode+schedule so a partial
// update can't leave a long-running agent without a valid cron (which the
// scheduler would then skip forever). Absent fields keep the stored value.
wasLongRunning := a.ExecutionMode == ModeLongRunning
mode, schedule := a.ExecutionMode, a.Schedule
if body.ExecutionMode != nil {
mode = *body.ExecutionMode
}
if body.Schedule != nil {
schedule = *body.Schedule
}
if a.ExecutionMode, a.Schedule, err = validateLifecycle(mode, schedule); err != nil {
return err
}
// Enforce the per-org scheduler cap on a TRANSITION into long-running, so a
// tenant can't sidestep the create-time cap by making N one-shot agents and
// PATCHing them to long-running (Red LOW-1 follow-up). Only counts when the
// agent was NOT already long-running (a no-op re-save of an existing
// long-running agent must not 409 against its own row).
if a.ExecutionMode == ModeLongRunning && !wasLongRunning {
n, cerr := s.store.CountLongRunning(c.Context(), org)
if cerr != nil {
return zip.Errorf(http.StatusInternalServerError, "count: %v", cerr)
}
if n >= longRunningCap() {
return zip.Errorf(http.StatusConflict,
"long-running agent limit reached for this org (max %d)", longRunningCap())
}
}
a.UpdatedAt = time.Now().Unix()
if err := s.store.Update(c.Context(), a); err != nil {
if err == errNotFound {
@@ -344,6 +484,18 @@ func (s *svc) run(c *zip.Ctx) error {
if !ok {
return zip.ErrForbidden("X-Org-Id required")
}
// A run MOVES MONEY (it debits org's commerce ledger), so it requires a
// VALIDATED principal — not merely a client-supplied X-Org-Id. SanitizeIdentity
// sets X-User-Id (c.User()) ONLY from a JWT it verified, and on the no-bearer
// direct-to-pod path it restores the client's raw X-Org-Id but leaves X-User-Id
// EMPTY. Gating the run on c.User() refuses exactly that anonymous-forge path:
// without it, a direct caller could set X-Org-Id to any tenant and charge a run
// against that victim's balance (Red MEDIUM-2). Read/list/create are the softer
// Phase-1 data path; the money action is held to the higher bar. Same guard the
// s3 / provisioning subsystems use.
if strings.TrimSpace(c.User()) == "" {
return zip.ErrForbidden("a validated principal is required to run an agent")
}
name := nameParam(c)
a, err := s.store.Get(c.Context(), org, name)
if err == errNotFound {
@@ -363,10 +515,14 @@ func (s *svc) run(c *zip.Ctx) error {
return zip.Errorf(http.StatusServiceUnavailable, "inference is not configured on this deployment")
}
r := executeRun(c.Context(), s.ai, org, a, body.Input)
// Record the run regardless of inference outcome — the history is real.
if err := s.store.InsertRun(c.Context(), r); err != nil {
s.log.Warn("record run failed", "org", org, "agent", name, "err", err)
// Pre-authorize the caller's org balance BEFORE any inference (fail-closed).
// The actor is the validated principal (org/sub) when present, else the bare
// org — recorded on the debit for attribution. Gating here means an unfunded
// org gets 402 and NO free inference; an unreachable commerce gets 503.
actor := billingActor(org, c.User())
r, gateErr := s.runAgent(c.Context(), a, body.Input, actor, c.RequestID(), cloud.ClientIP(c))
if gateErr != nil {
return cloud.DenyResource(c, gateErr)
}
if r.Status != "ok" {
// The run is recorded; surface the upstream failure honestly.
@@ -375,6 +531,42 @@ func (s *svc) run(c *zip.Ctx) error {
return c.JSON(http.StatusOK, toRunView(r))
}
// runAgent is the ONE run path — shared by the HTTP handler and the scheduler.
// It (1) pre-authorizes the AGENT's OWN org balance (fail-closed) so no unfunded
// tenant ever gets free inference, (2) executes one real completion, (3) records
// the run regardless of outcome (the history is real), and (4) debits the run
// fee to the agent's org ONLY on success. A non-nil error is a BALANCE-GATE
// denial (out-of-funds / commerce-unknown) that the caller renders (402/503) —
// it means no run happened. A run that executed but the model failed returns a
// recorded error-status Run and a nil error.
func (s *svc) runAgent(ctx context.Context, a Agent, input, actor, requestID, clientIP string) (Run, error) {
fee := cloud.ResourceFeeCents(agentFeeEnvPrefix, meterKind)
// Gate the AGENT's own org — never a caller default, never another tenant.
// fee<=0 or unconfigured billing makes this a no-op (allows).
if err := s.bill.Gate(ctx, a.Org, meterKind, fee); err != nil {
return Run{}, err
}
r := executeRun(ctx, s.ai, a.Org, a, input)
if err := s.store.InsertRun(ctx, r); err != nil {
s.log.Warn("record run failed", "org", a.Org, "agent", a.Name, "err", err)
}
// Bill only a successful run (mirrors the edge gate: failed work is not
// charged). Rich attribution: product=agent (Provider), the agent's model,
// and the actor for the audit trail. Fire-and-forget on a background context.
if r.Status == "ok" {
s.bill.MeterUsage(a.Org, meterKind, metering.Usage{
AmountCents: fee,
Model: a.Model,
Actor: actor,
RequestID: requestID,
ClientIP: clientIP,
})
}
return r, nil
}
// executeRun composes the agent's instructions with the caller input, runs one
// real chat completion through the AI client, and returns the resulting Run —
// status "ok" with output, or "error" with the upstream failure. Pure of HTTP
@@ -442,6 +634,70 @@ func nameParam(c *zip.Ctx) string { return strings.TrimSpace(c.Param("name")) }
// carries an explicit org, so an empty org is a true 403.
func tenant(c *zip.Ctx) (string, bool) { return principal.Tenant(c) }
// validateLifecycle normalizes and validates the execution mode + schedule.
// Empty mode defaults to one-shot. A long-running agent MUST carry a schedule
// that parses as a 5-field cron (else the scheduler would silently never fire
// it); a one-shot agent's schedule is cleared (it is meaningless without the
// scheduler). Returns the normalized (mode, schedule) or a 400.
func validateLifecycle(mode, schedule string) (string, string, error) {
mode = strings.TrimSpace(mode)
if mode == "" {
mode = ModeOneShot
}
schedule = strings.TrimSpace(schedule)
switch mode {
case ModeOneShot:
return ModeOneShot, "", nil // schedule is meaningless one-shot; drop it.
case ModeLongRunning:
if schedule == "" {
return "", "", zip.ErrBadRequest("a long-running agent requires a 'schedule' (5-field cron)")
}
if _, err := parseCron(schedule); err != nil {
return "", "", zip.ErrBadRequest("invalid 'schedule': " + err.Error())
}
return ModeLongRunning, schedule, nil
default:
return "", "", zip.ErrBadRequest("executionMode must be 'one-shot' or 'long-running'")
}
}
// longRunningCap resolves the per-org scheduled-agent limit from the operator
// env override, falling back to the default. A non-positive/invalid override is
// ignored so a typo can never remove the cap.
func longRunningCap() int {
if v := strings.TrimSpace(os.Getenv(longRunningCapEnv)); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
return n
}
}
return maxLongRunningPerOrg
}
// validateRef bounds an opaque lifecycle reference (compute id / service-account
// id). Returns the trimmed value or a 400 when it exceeds maxRef.
func validateRef(field, v string) (string, error) {
v = strings.TrimSpace(v)
if len(v) > maxRef {
return "", zip.ErrBadRequest(field + " too long")
}
return v, nil
}
// billingActor is the "org/sub" identity recorded on a debit for the audit
// trail. It never selects which balance is gated — that is always the org — but
// attributes the spend to a principal. Falls back to the bare org when no
// validated user subject is present (e.g. a service-token caller).
func billingActor(org, sub string) string {
sub = strings.TrimSpace(sub)
if org != "" && sub != "" {
return org + "/" + sub
}
if sub != "" {
return sub
}
return org
}
func cleanList(xs []string) []string {
seen := map[string]bool{}
var out []string
@@ -467,12 +723,19 @@ func genID(prefix string) (string, error) {
return prefix + "_" + hex.EncodeToString(b[:]), nil
}
// Shutdown closes the agents store. Idempotent.
func Shutdown() error {
if mounted == nil || mounted.store == nil {
// Shutdown stops the scheduler (draining in-flight runs, bounded by ctx) and
// closes the agents store. Idempotent — safe to call when nothing is mounted.
func Shutdown(ctx context.Context) error {
if mounted == nil {
return nil
}
err := mounted.store.Close()
if mounted.sched != nil {
mounted.sched.stop(ctx)
}
var err error
if mounted.store != nil {
err = mounted.store.Close()
}
mounted = nil
return err
}
+246
View File
@@ -0,0 +1,246 @@
package agents
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/types"
"github.com/hanzoai/commerce/metering"
luxlog "github.com/luxfi/log"
"github.com/zap-proto/zip"
)
// errTest is the model-failure the "failed run is not billed" case injects.
var errTest = errors.New("model unavailable")
// billServer is a minimal commerce double: it returns a fixed balance and
// records the X-Org-Id header (the tenant the debit lands on) + the usage body
// of every debit. X-Org-Id is the header commerce's service-token auth reads
// (metering >= v0.1.2), so a wrong tenant here would prove a cross-tenant leak.
type billServer struct {
available int64
mu sync.Mutex
usageOrg string
usageBody []byte
usages int32
balances int32
}
func (b *billServer) start(t *testing.T) string {
t.Helper()
mux := http.NewServeMux()
mux.HandleFunc("/v1/billing/balance", func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&b.balances, 1)
_ = json.NewEncoder(w).Encode(map[string]any{"available": b.available})
})
mux.HandleFunc("/v1/billing/usage", func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&b.usages, 1)
body, _ := io.ReadAll(r.Body)
b.mu.Lock()
b.usageOrg, b.usageBody = r.Header.Get("X-Org-Id"), body
b.mu.Unlock()
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `{"transactionId":"tx_1","type":"usage"}`)
})
srv := httptest.NewServer(mux)
t.Cleanup(srv.Close)
return srv.URL
}
func (b *billServer) debits() int32 { return atomic.LoadInt32(&b.usages) }
func (b *billServer) lastDebit() (string, []byte) {
b.mu.Lock()
defer b.mu.Unlock()
return b.usageOrg, b.usageBody
}
// waitForDebit polls a condition briefly — debits are recorded on a detached
// goroutine, so the assertion must wait for the async write.
func waitForDebit(cond func() bool) bool {
for i := 0; i < 200; i++ {
if cond() {
return true
}
time.Sleep(5 * time.Millisecond)
}
return cond()
}
// mountBilled mounts the agents surface with a REAL metering client pointed at
// the fake commerce (default org "hanzo", so every "acme is billed" assertion
// proves the per-call org override scopes the ledger to the CALLER). No
// scheduler is started here (deps.AI is set, but these tests exercise the HTTP
// run path; scheduler tests drive tick() directly).
func mountBilled(t *testing.T, commerceURL string, ai types.AIClient) *zip.App {
t.Helper()
m, err := metering.New(metering.Config{BaseURL: commerceURL, Token: "svc-tok", Org: "hanzo"})
if err != nil {
t.Fatalf("metering.New: %v", err)
}
app := zip.New(zip.Config{Logger: luxlog.New("test")})
deps := cloud.Deps{Logger: luxlog.New("test"), DataDir: t.TempDir(), AI: ai, Metering: m}
if err := Mount(app, deps); err != nil {
t.Fatalf("Mount: %v", err)
}
t.Cleanup(func() { _ = Shutdown(context.Background()) })
return app
}
// TestRunGatesUnfundedOrg: a run for an org with a non-positive balance is
// refused 402 and NO usage is recorded and (fail-closed) no inference output is
// returned — an unfunded tenant gets no free agent run.
func TestRunGatesUnfundedOrg(t *testing.T) {
bs := &billServer{available: 0}
app := mountBilled(t, bs.start(t), &fakeAI{content: "should not run"})
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "gpt-4o-mini", "instructions": "x"}); code != http.StatusCreated {
t.Fatalf("create want 201, got %d", code)
}
code, body := do(t, app, http.MethodPost, "/v1/agents/a/run", "acme", map[string]any{"input": "hi"})
if code != http.StatusPaymentRequired {
t.Fatalf("unfunded run want 402, got %d (%s)", code, body)
}
if bs.debits() != 0 {
t.Fatalf("a refused run must not debit, got %d", bs.debits())
}
}
// TestRunGatesUnderfundedOrg: an org with a POSITIVE balance that is still less
// than the run fee is refused 402 — the gate enforces available >= fee, not
// merely available > 0, so a 1-cent balance can't authorize a $1 run and take
// the ledger negative (Red MEDIUM-1). Default fee is $1.00 (100c).
func TestRunGatesUnderfundedOrg(t *testing.T) {
bs := &billServer{available: 1} // 1 cent, fee is 100 cents
app := mountBilled(t, bs.start(t), &fakeAI{content: "should not run"})
do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "m", "instructions": "x"})
code, body := do(t, app, http.MethodPost, "/v1/agents/a/run", "acme", map[string]any{"input": "hi"})
if code != http.StatusPaymentRequired {
t.Fatalf("underfunded (1c < 100c fee) run want 402, got %d (%s)", code, body)
}
if bs.debits() != 0 {
t.Fatalf("a gate-refused run must not debit, got %d", bs.debits())
}
}
// TestRunDebitsCallerOrg: a funded run returns the output AND debits the CALLER
// org (acme, never the client default 'hanzo'), with product=agent + the agent's
// model on the usage transaction.
func TestRunDebitsCallerOrg(t *testing.T) {
bs := &billServer{available: 100000}
app := mountBilled(t, bs.start(t), &fakeAI{content: "the answer"})
do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "gpt-4o-mini", "instructions": "x"})
code, body := do(t, app, http.MethodPost, "/v1/agents/a/run", "acme", map[string]any{"input": "hi"})
if code != http.StatusOK {
t.Fatalf("funded run want 200, got %d (%s)", code, body)
}
if !waitForDebit(func() bool { return bs.debits() == 1 }) {
t.Fatalf("a successful run must debit once, got %d", bs.debits())
}
org, ubody := bs.lastDebit()
if org != "acme" {
t.Fatalf("debited org %q, want caller %q (never default 'hanzo')", org, "acme")
}
var u struct {
User string `json:"user"`
Amount int64 `json:"amount"`
Model string `json:"model"`
Provider string `json:"provider"`
Actor string `json:"actor"`
}
_ = json.Unmarshal(ubody, &u)
if u.User != "acme" {
t.Fatalf("debit user = %q, want caller org %q", u.User, "acme")
}
if u.Amount != cloud.DefaultResourceFeeCents {
t.Fatalf("debit amount = %d, want default fee %d", u.Amount, cloud.DefaultResourceFeeCents)
}
if u.Provider != meterKind {
t.Fatalf("debit provider = %q, want %q (product:agent)", u.Provider, meterKind)
}
if u.Model != "gpt-4o-mini" {
t.Fatalf("debit model = %q, want the agent's model", u.Model)
}
if u.Actor == "" {
t.Fatalf("debit must carry an actor for the audit trail")
}
}
// TestFailedRunNotBilled: when the model errors, the run is recorded as an error
// but NOT billed — failed work is never charged (mirrors the edge gate).
func TestFailedRunNotBilled(t *testing.T) {
bs := &billServer{available: 100000}
app := mountBilled(t, bs.start(t), &fakeAI{err: errTest})
do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "m", "instructions": "x"})
code, _ := do(t, app, http.MethodPost, "/v1/agents/a/run", "acme", map[string]any{"input": "hi"})
if code != http.StatusBadGateway {
t.Fatalf("errored run want 502, got %d", code)
}
// Give any (erroneous) async debit a chance to land, then assert none did.
if waitForDebit(func() bool { return bs.debits() > 0 }) {
t.Fatalf("a failed run must NOT be billed, got %d debits", bs.debits())
}
}
// TestRunRequiresValidatedPrincipal: a run with only a client X-Org-Id (no
// validated X-User-Id — the direct-to-pod no-bearer path) is refused 403 and
// NEVER debits. A money-moving action can't ride an unauthenticated, forgeable
// org header (Red MEDIUM-2). Read/create still work on the org header alone.
func TestRunRequiresValidatedPrincipal(t *testing.T) {
bs := &billServer{available: 100000}
app := mountBilled(t, bs.start(t), &fakeAI{content: "must not run"})
// create is allowed with X-User-Id (via do()).
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "m", "instructions": "x"}); code != http.StatusCreated {
t.Fatalf("create want 201, got %d", code)
}
// A raw run request carrying ONLY X-Org-Id (no X-User-Id) must be 403.
req := httptest.NewRequest(http.MethodPost, "/v1/agents/a/run", nil)
req.Header.Set("X-Org-Id", "acme") // forged/unvalidated org, no principal
resp, err := app.Fiber().Test(req)
if err != nil {
t.Fatalf("Test: %v", err)
}
_ = resp.Body.Close()
if resp.StatusCode != http.StatusForbidden {
t.Fatalf("run without a validated principal want 403, got %d", resp.StatusCode)
}
if waitForDebit(func() bool { return bs.debits() > 0 }) {
t.Fatalf("an unauthenticated run must never debit, got %d", bs.debits())
}
}
// TestRunAgentGateFailClosedOnUnreachableCommerce: when commerce cannot be
// reached, the gate denies (fail-closed) and no run executes — runAgent returns
// the gate error and the fake AI is never called.
func TestRunAgentGateFailClosedOnUnreachableCommerce(t *testing.T) {
// Point at a dead URL so Authorize errors (unknown balance -> fail-closed).
m, _ := metering.New(metering.Config{BaseURL: "http://127.0.0.1:1", Token: "t", Org: "hanzo", Timeout: 200 * time.Millisecond})
ai := &fakeAI{content: "must not run"}
s := &svc{store: testStore(t), ai: ai, log: luxlog.New("test"), bill: cloud.NewResourceMeter(cloud.Deps{Metering: m, Logger: luxlog.New("test")}, meterKind)}
a := mk("acme", "x")
_, gateErr := s.runAgent(context.Background(), a, "hi", "acme", "", "")
if gateErr == nil {
t.Fatal("unreachable commerce must fail closed (non-nil gate error)")
}
if ai.gotPrompt != "" {
t.Fatalf("no inference must run when the gate denies, got prompt %q", ai.gotPrompt)
}
}
+150
View File
@@ -0,0 +1,150 @@
package agents
// A minimal, dependency-free 5-field cron matcher — the ONE schedule grammar
// long-running agents use. It is deliberately tiny (no seconds field, no
// @-macros, no timezones beyond UTC) because the scheduler ticks once a minute
// and only needs to answer "does this expression fire at this minute?".
//
// Grammar (standard 5-field, all times UTC):
//
// minute hour day-of-month month day-of-week
// 0-59 0-23 1-31 1-12 0-6 (Sun=0)
//
// Each field is a comma list of terms; a term is "*", a number, a range "a-b",
// or a step "*/n" or "a-b/n". Day-of-month and day-of-week combine with OR when
// BOTH are restricted (Vixie-cron semantics), else AND — matching what operators
// expect from "0 9 * * 1" (09:00 on Mondays).
//
// Hand-rolled instead of pulling a cron module: it is ~one screen, fully
// unit-tested, and keeps the dependency surface minimal (no new module in a
// binary that mounts every subsystem).
import (
"fmt"
"strconv"
"strings"
"time"
)
// schedule is a parsed cron expression: one bitset per field (bit i set => the
// field matches value i). domRestricted/dowRestricted record whether the
// day-of-month / day-of-week field was anything other than "*", which selects
// the OR-vs-AND combination rule.
type schedule struct {
min, hour, dom, mon, dow uint64
domRestricted bool
dowRestricted bool
}
// fieldRange bounds each cron field (inclusive).
type fieldRange struct{ min, max int }
var cronRanges = [5]fieldRange{
{0, 59}, // minute
{0, 23}, // hour
{1, 31}, // day of month
{1, 12}, // month
{0, 6}, // day of week (Sunday=0)
}
// parseCron parses a 5-field cron expression or returns an error describing the
// first malformed field. Whitespace between fields is collapsed.
func parseCron(expr string) (schedule, error) {
fields := strings.Fields(strings.TrimSpace(expr))
if len(fields) != 5 {
return schedule{}, fmt.Errorf("cron: want 5 fields, got %d in %q", len(fields), expr)
}
var s schedule
dst := []*uint64{&s.min, &s.hour, &s.dom, &s.mon, &s.dow}
for i, f := range fields {
bits, err := parseField(f, cronRanges[i])
if err != nil {
return schedule{}, fmt.Errorf("cron field %d (%q): %w", i+1, f, err)
}
*dst[i] = bits
}
s.domRestricted = fields[2] != "*"
s.dowRestricted = fields[4] != "*"
return s, nil
}
// parseField parses one comma-separated cron field into a bitset over r.
func parseField(f string, r fieldRange) (uint64, error) {
if f == "" {
return 0, fmt.Errorf("empty field")
}
var bits uint64
for _, term := range strings.Split(f, ",") {
tb, err := parseTerm(term, r)
if err != nil {
return 0, err
}
bits |= tb
}
return bits, nil
}
// parseTerm parses a single term: "*", "n", "a-b", "*/n", or "a-b/n".
func parseTerm(term string, r fieldRange) (uint64, error) {
step := 1
if i := strings.IndexByte(term, '/'); i >= 0 {
n, err := strconv.Atoi(term[i+1:])
if err != nil || n <= 0 {
return 0, fmt.Errorf("bad step %q", term)
}
step = n
term = term[:i]
}
lo, hi := r.min, r.max
switch {
case term == "*":
// full range with the parsed step.
case strings.IndexByte(term, '-') > 0:
i := strings.IndexByte(term, '-')
a, err1 := strconv.Atoi(term[:i])
b, err2 := strconv.Atoi(term[i+1:])
if err1 != nil || err2 != nil {
return 0, fmt.Errorf("bad range %q", term)
}
lo, hi = a, b
default:
n, err := strconv.Atoi(term)
if err != nil {
return 0, fmt.Errorf("bad number %q", term)
}
lo, hi = n, n
}
if lo < r.min || hi > r.max || lo > hi {
return 0, fmt.Errorf("value out of range [%d,%d]", r.min, r.max)
}
var bits uint64
for v := lo; v <= hi; v += step {
bits |= 1 << uint(v)
}
return bits, nil
}
// matches reports whether the schedule fires at t (evaluated in UTC, minute
// granularity). The day-of-month / day-of-week combination follows Vixie cron:
// when BOTH are restricted the day matches if EITHER matches (OR); otherwise the
// unrestricted field is a wildcard and the restricted one is ANDed.
func (s schedule) matches(t time.Time) bool {
t = t.UTC()
if s.min&(1<<uint(t.Minute())) == 0 {
return false
}
if s.hour&(1<<uint(t.Hour())) == 0 {
return false
}
if s.mon&(1<<uint(int(t.Month()))) == 0 {
return false
}
domHit := s.dom&(1<<uint(t.Day())) != 0
dowHit := s.dow&(1<<uint(int(t.Weekday()))) != 0
if s.domRestricted && s.dowRestricted {
return domHit || dowHit
}
return domHit && dowHit
}
+105
View File
@@ -0,0 +1,105 @@
package agents
import (
"testing"
"time"
)
func TestParseCronErrors(t *testing.T) {
bad := []string{
"", // empty
"* * * *", // 4 fields
"* * * * * *", // 6 fields
"60 * * * *", // minute out of range
"* 24 * * *", // hour out of range
"* * 0 * *", // dom below range
"* * 32 * *", // dom above range
"* * * 13 *", // month above range
"* * * * 7", // dow above range
"*/0 * * * *", // zero step
"5-1 * * * *", // inverted range
"abc * * * *", // non-numeric
"1,,2 * * * *", // empty term
}
for _, expr := range bad {
if _, err := parseCron(expr); err == nil {
t.Errorf("parseCron(%q) = nil error, want error", expr)
}
}
}
func TestParseCronValid(t *testing.T) {
for _, expr := range []string{
"* * * * *", "*/5 * * * *", "0 9 * * 1", "0 0 1 * *",
"0,30 * * * *", "0-15 * * * *", "0 9-17/2 * * 1-5", "0 0 * * 0",
} {
if _, err := parseCron(expr); err != nil {
t.Errorf("parseCron(%q) unexpected error: %v", expr, err)
}
}
}
func at(t *testing.T, s string) time.Time {
t.Helper()
tm, err := time.Parse("2006-01-02 15:04 MST", s+" UTC")
if err != nil {
t.Fatalf("bad test time %q: %v", s, err)
}
return tm
}
func TestCronMatches(t *testing.T) {
cases := []struct {
expr string
when string // "YYYY-MM-DD HH:MM"
want bool
}{
{"* * * * *", "2026-07-01 12:34", true},
{"*/5 * * * *", "2026-07-01 12:35", true},
{"*/5 * * * *", "2026-07-01 12:36", false},
{"0 9 * * *", "2026-07-01 09:00", true},
{"0 9 * * *", "2026-07-01 09:01", false},
{"0 9 * * *", "2026-07-01 10:00", false},
// 2026-07-06 is a Monday; 0 9 * * 1 fires 09:00 Mondays.
{"0 9 * * 1", "2026-07-06 09:00", true},
{"0 9 * * 1", "2026-07-07 09:00", false}, // Tuesday
{"0-15 * * * *", "2026-07-01 12:15", true},
{"0-15 * * * *", "2026-07-01 12:16", false},
{"0 0 1 * *", "2026-08-01 00:00", true}, // first of month
{"0 0 1 * *", "2026-08-02 00:00", false}, // second of month
{"0 9-17/2 * * *", "2026-07-01 09:00", true},
{"0 9-17/2 * * *", "2026-07-01 11:00", true},
{"0 9-17/2 * * *", "2026-07-01 10:00", false}, // 10 not in 9,11,13,15,17
}
for _, c := range cases {
s, err := parseCron(c.expr)
if err != nil {
t.Fatalf("parseCron(%q): %v", c.expr, err)
}
if got := s.matches(at(t, c.when)); got != c.want {
t.Errorf("%q matches %q = %v, want %v", c.expr, c.when, got, c.want)
}
}
}
// TestCronDOMDOWOrSemantics: when BOTH day-of-month and day-of-week are
// restricted, Vixie cron fires if EITHER matches. "0 0 13 * 5" fires on the
// 13th OR on any Friday.
func TestCronDOMDOWOrSemantics(t *testing.T) {
s, err := parseCron("0 0 13 * 5") // 5 = Friday
if err != nil {
t.Fatalf("parse: %v", err)
}
// 2026-07-13 is a Monday -> matches via DOM (the 13th).
if !s.matches(at(t, "2026-07-13 00:00")) {
t.Error("should fire on the 13th regardless of weekday")
}
// 2026-07-03 is a Friday -> matches via DOW.
if !s.matches(at(t, "2026-07-03 00:00")) {
t.Error("should fire on a Friday regardless of day-of-month")
}
// 2026-07-06 is a Monday, not the 13th -> no match.
if s.matches(at(t, "2026-07-06 00:00")) {
t.Error("must NOT fire on a non-13th non-Friday")
}
}
+11 -2
View File
@@ -2,6 +2,7 @@ package agents
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
@@ -10,8 +11,8 @@ import (
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/types"
"github.com/zap-proto/zip"
luxlog "github.com/luxfi/log"
"github.com/zap-proto/zip"
)
// mountApp mounts the agents surface with a deterministic fake AI so run() is
@@ -23,6 +24,10 @@ func mountApp(t *testing.T, ai types.AIClient) *zip.App {
if err := Mount(app, cloud.Deps{Logger: luxlog.New("test"), DataDir: t.TempDir(), AI: ai}); err != nil {
t.Fatalf("Mount: %v", err)
}
// Mount starts the scheduler goroutine when AI is non-nil and sets the global
// `mounted` singleton; tear both down at test end so the loop goroutine can't
// leak and clobber a later test's singleton (Red re-review LOW).
t.Cleanup(func() { _ = Shutdown(context.Background()) })
return app
}
@@ -39,7 +44,11 @@ func do(t *testing.T, app *zip.App, method, path, org string, body any) (int, []
}
if org != "" {
req.Header.Set("X-Org-Id", org)
req.Header.Set("X-User-Id", "u_"+org) // validated principal (tenant() gates on it)
// A validated principal: the run path (money-moving) requires a non-empty
// c.User() (X-User-Id). SanitizeIdentity sets this only from a verified
// JWT; the test app has no sanitizer, so we inject it directly, exactly as
// the gateway would. Empty org => no user (the anonymous 403 path).
req.Header.Set("X-User-Id", "u-"+org)
}
resp, err := app.Fiber().Test(req)
if err != nil {
+153
View File
@@ -0,0 +1,153 @@
package agents
import (
"encoding/json"
"net/http"
"strings"
"testing"
)
// TestCreateRejectsOversizedRefs: computeRef/serviceAccountId are opaque ids,
// bounded at the boundary — a multi-KB "id" must be a 400, not persisted.
func TestCreateRejectsOversizedRefs(t *testing.T) {
app := mountApp(t, &fakeAI{content: "x"})
huge := strings.Repeat("a", maxRef+1)
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "big", "model": "m", "computeRef": huge}); code != http.StatusBadRequest {
t.Fatalf("oversized computeRef want 400, got %d", code)
}
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "big2", "model": "m", "serviceAccountId": huge}); code != http.StatusBadRequest {
t.Fatalf("oversized serviceAccountId want 400, got %d", code)
}
}
// TestCreateLongRunningRequiresValidCron: a long-running agent must carry a
// parseable cron; missing/invalid schedule is a 400. A valid one is 201 and the
// mode+schedule round-trip in the view.
func TestCreateLongRunningValidation(t *testing.T) {
app := mountApp(t, &fakeAI{content: "x"})
// long-running without a schedule -> 400.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "a", "model": "m", "executionMode": "long-running"}); code != http.StatusBadRequest {
t.Fatalf("long-running w/o schedule want 400, got %d", code)
}
// long-running with a bad cron -> 400.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "b", "model": "m", "executionMode": "long-running", "schedule": "not a cron"}); code != http.StatusBadRequest {
t.Fatalf("long-running w/ bad cron want 400, got %d", code)
}
// unknown mode -> 400.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "c", "model": "m", "executionMode": "daemon"}); code != http.StatusBadRequest {
t.Fatalf("unknown mode want 400, got %d", code)
}
// valid long-running -> 201, fields echoed.
code, body := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "cron", "model": "m", "executionMode": "long-running",
"schedule": "*/5 * * * *", "computeRef": "vm-1", "serviceAccountId": "acme-cron"})
if code != http.StatusCreated {
t.Fatalf("valid long-running want 201, got %d (%s)", code, body)
}
var v agentView
_ = json.Unmarshal(body, &v)
if v.ExecutionMode != "long-running" || v.Schedule != "*/5 * * * *" ||
v.ComputeRef != "vm-1" || v.ServiceAccountID != "acme-cron" {
t.Fatalf("lifecycle fields not echoed in view: %+v", v)
}
}
// TestCreateOneShotDropsSchedule: a one-shot agent's schedule is meaningless and
// dropped, so the view carries no schedule and the scheduler will never pick it.
func TestCreateOneShotDropsSchedule(t *testing.T) {
app := mountApp(t, &fakeAI{content: "x"})
code, body := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "one", "model": "m", "schedule": "* * * * *"})
if code != http.StatusCreated {
t.Fatalf("create want 201, got %d", code)
}
var v agentView
_ = json.Unmarshal(body, &v)
if v.ExecutionMode != "one-shot" || v.Schedule != "" {
t.Fatalf("one-shot must drop schedule, got mode=%q schedule=%q", v.ExecutionMode, v.Schedule)
}
}
// TestLongRunningPerOrgCap: an org cannot create more than the configured number
// of scheduled long-running agents (Red LOW-1). One-shot agents don't count.
func TestLongRunningPerOrgCap(t *testing.T) {
t.Setenv(longRunningCapEnv, "2")
app := mountApp(t, &fakeAI{content: "x"})
mk := func(name string) map[string]any {
return map[string]any{"name": name, "model": "m", "executionMode": "long-running", "schedule": "* * * * *"}
}
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme", mk("a")); code != http.StatusCreated {
t.Fatalf("1st long-running want 201, got %d", code)
}
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme", mk("b")); code != http.StatusCreated {
t.Fatalf("2nd long-running want 201, got %d", code)
}
// 3rd exceeds the cap of 2 -> 409.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme", mk("c")); code != http.StatusConflict {
t.Fatalf("3rd long-running want 409 (cap), got %d", code)
}
// A one-shot agent is unaffected by the cap.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "one", "model": "m"}); code != http.StatusCreated {
t.Fatalf("one-shot must not be capped, got %d", code)
}
// A DIFFERENT org has its own budget.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "beta", mk("a")); code != http.StatusCreated {
t.Fatalf("other org's 1st long-running want 201, got %d", code)
}
}
// TestLongRunningCapNotBypassedByPatch: the cap can't be dodged by creating
// one-shot agents (uncapped) then PATCHing them to long-running. Transition into
// long-running is capped too; re-saving an already-long-running agent is not.
func TestLongRunningCapNotBypassedByPatch(t *testing.T) {
t.Setenv(longRunningCapEnv, "1")
app := mountApp(t, &fakeAI{content: "x"})
// Fill the cap with one long-running agent.
if code, _ := do(t, app, http.MethodPost, "/v1/agents", "acme",
map[string]any{"name": "lr", "model": "m", "executionMode": "long-running", "schedule": "* * * * *"}); code != http.StatusCreated {
t.Fatalf("seed long-running want 201, got %d", code)
}
// Create a one-shot agent (uncapped), then try to PATCH it to long-running.
do(t, app, http.MethodPost, "/v1/agents", "acme", map[string]any{"name": "sneaky", "model": "m"})
if code, _ := do(t, app, http.MethodPatch, "/v1/agents/sneaky", "acme",
map[string]any{"executionMode": "long-running", "schedule": "* * * * *"}); code != http.StatusConflict {
t.Fatalf("PATCH one-shot->long-running over cap want 409, got %d", code)
}
// Re-saving the EXISTING long-running agent (no transition) must NOT 409.
if code, _ := do(t, app, http.MethodPatch, "/v1/agents/lr", "acme",
map[string]any{"schedule": "*/2 * * * *"}); code != http.StatusOK {
t.Fatalf("no-op re-save of own long-running agent want 200, got %d", code)
}
}
// TestPatchToLongRunningValidates: PATCHing an agent to long-running without a
// schedule is rejected; supplying a valid schedule in the same PATCH succeeds.
func TestPatchToLongRunningValidates(t *testing.T) {
app := mountApp(t, &fakeAI{content: "x"})
do(t, app, http.MethodPost, "/v1/agents", "acme", map[string]any{"name": "a", "model": "m"})
// flip to long-running with no schedule -> 400.
if code, _ := do(t, app, http.MethodPatch, "/v1/agents/a", "acme",
map[string]any{"executionMode": "long-running"}); code != http.StatusBadRequest {
t.Fatalf("patch to long-running w/o schedule want 400, got %d", code)
}
// flip with a schedule -> 200.
code, body := do(t, app, http.MethodPatch, "/v1/agents/a", "acme",
map[string]any{"executionMode": "long-running", "schedule": "0 * * * *"})
if code != http.StatusOK {
t.Fatalf("patch to long-running w/ schedule want 200, got %d (%s)", code, body)
}
var v agentView
_ = json.Unmarshal(body, &v)
if v.ExecutionMode != "long-running" || v.Schedule != "0 * * * *" {
t.Fatalf("patch did not apply lifecycle: %+v", v)
}
}
+290
View File
@@ -0,0 +1,290 @@
package agents
// The long-running-agent scheduler: it invokes each long-running agent's run on
// its cron cadence, through the SAME svc.runAgent path the HTTP handler uses —
// so a scheduled run is gated (fail-closed on the agent's OWN org balance),
// executed, recorded, and billed identically to an interactive one. There is no
// self-HTTP call: the endpoint's BEHAVIOR is the contract, and calling runAgent
// directly keeps ONE run path (no duplicated gate/meter, no re-crossing the
// identity boundary with a synthetic token).
//
// Cadence: one ticker fires every minute (cron's finest granularity). On each
// tick it loads the long-running work set and, for every agent whose schedule
// matches the current minute, launches a run — subject to two safety controls:
//
// - Concurrency cap: at most maxConcurrentPerAgent in-flight runs per agent.
// A slow model must never let cron stack unbounded goroutines for one agent.
// - Exponential backoff: after a failed run (gate denial OR model error) an
// agent is skipped for a growing number of ticks (1,2,4,… up to a cap), so a
// persistently-failing or unfunded agent stops hammering commerce/the model.
// A success resets the backoff.
//
// All state is in-memory and keyed by org/name: the scheduler is a per-process
// singleton owned by the mounted svc, torn down on Shutdown.
import (
"context"
"strings"
"sync"
"time"
luxlog "github.com/luxfi/log"
)
const (
// tickInterval is cron's resolution. Aligned to the top of each minute so a
// "* * * * *" agent fires once per minute, not on process-start phase.
tickInterval = time.Minute
// maxConcurrentPerAgent caps in-flight runs for ONE agent. A cron agent is
// expected to complete within its period; 1 means "never overlap a run with
// itself" (the safe default for periodic work). >1 would allow catch-up.
maxConcurrentPerAgent = 1
// maxBackoffTicks caps the exponential skip so a failing agent still retries
// roughly hourly rather than backing off forever.
maxBackoffTicks = 60
// runTimeout bounds a single scheduled run so one stuck completion cannot pin
// a concurrency slot indefinitely.
runTimeout = 10 * time.Minute
)
// agentState is the per-agent runtime bookkeeping the scheduler keeps between
// ticks: how many runs are in flight, and the backoff countdown after failures.
type agentState struct {
inFlight int
failstreak int // consecutive failures; drives the backoff window.
skipRemain int // ticks still to skip before the next attempt.
parsed schedule
parsedExpr string // the expression `parsed` was compiled from (recompile on change).
}
type scheduler struct {
svc *svc
log luxlog.Logger
cancel context.CancelFunc // cancels the loop + all in-flight run contexts.
mu sync.Mutex
states map[string]*agentState // key: org + "\x00" + name
wg sync.WaitGroup // tracks in-flight run goroutines for clean shutdown.
// now is time.Now, overridable in tests for deterministic cron evaluation.
now func() time.Time
// tick, when non-nil, replaces the internal ticker so tests drive cadence.
tickC <-chan time.Time
}
func newScheduler(s *svc, log luxlog.Logger) *scheduler {
return &scheduler{
svc: s,
log: log.New("component", "scheduler"),
states: map[string]*agentState{},
now: time.Now,
}
}
// start launches the scheduler loop in its own goroutine. Cancelled by stop().
func (sc *scheduler) start() {
ctx, cancel := context.WithCancel(context.Background())
sc.cancel = cancel
sc.wg.Add(1)
go sc.loop(ctx)
}
// stop halts the scheduler and waits for in-flight runs to drain BEFORE the
// caller closes the store — otherwise a run could InsertRun into a closed DB.
//
// Cancelling the loop context also cancels every in-flight run's derived
// context, so a run whose AIClient honors ctx returns promptly. The drain wait
// is bounded by the caller's shutdown ctx: if a run ignores cancellation and
// runs long (up to runTimeout), stop returns at the deadline rather than hanging
// SIGTERM. Idempotent.
func (sc *scheduler) stop(ctx context.Context) {
if sc.cancel == nil {
return
}
sc.cancel()
sc.cancel = nil
done := make(chan struct{})
go func() { sc.wg.Wait(); close(done) }()
select {
case <-done: // clean drain
case <-ctx.Done():
sc.log.Warn("scheduler drain timed out; in-flight runs may not have recorded",
"err", ctx.Err())
}
}
// loop is the cadence driver. It uses the injected tick channel in tests, else a
// real minute ticker. Each tick evaluates the whole long-running work set.
func (sc *scheduler) loop(ctx context.Context) {
defer sc.wg.Done()
tickC := sc.tickC
if tickC == nil {
t := time.NewTicker(tickInterval)
defer t.Stop()
tickC = t.C
}
sc.log.Info("scheduler started", "interval", tickInterval)
for {
select {
case <-ctx.Done():
sc.log.Info("scheduler stopped")
return
case <-tickC:
sc.tick(ctx, sc.now())
}
}
}
// tick evaluates every long-running agent against the wall-clock minute now and
// launches the ones that are due, are not backed off, and have a free
// concurrency slot. It is separated from loop() so tests can invoke it directly.
func (sc *scheduler) tick(ctx context.Context, now time.Time) {
agents, err := sc.svc.store.ListLongRunning(ctx)
if err != nil {
sc.log.Warn("scheduler: list long-running failed", "err", err)
return
}
live := make(map[string]bool, len(agents))
for _, a := range agents {
key := stateKey(a.Org, a.Name)
live[key] = true
if sc.due(a, key, now) {
sc.launch(ctx, a, key)
}
}
sc.pruneDeleted(live)
}
// due decides, under a SINGLE lock acquisition, whether agent a should run this
// tick: it (re)compiles the cron on change, decrements a live backoff window,
// checks the cron against now and the per-agent concurrency slot, and — when it
// returns true — has already reserved the slot (inFlight++). All shared state
// (parsed cron, backoff, inFlight) is touched only while holding sc.mu, so there
// is no data race with the completion goroutine in launch().
func (sc *scheduler) due(a Agent, key string, now time.Time) bool {
sc.mu.Lock()
defer sc.mu.Unlock()
st := sc.stateForLocked(key)
// Recompile the cron only when the expression changed (edits via PATCH).
if st.parsedExpr != a.Schedule {
p, err := parseCron(a.Schedule)
if err != nil {
// Stored schedule is invalid (create/update validate, but a
// hand-edited DB could carry garbage). Skip, don't crash.
sc.log.Warn("scheduler: bad stored schedule, skipping",
"org", a.Org, "agent", a.Name, "schedule", a.Schedule, "err", err)
return false
}
st.parsed, st.parsedExpr = p, a.Schedule
}
if st.skipRemain > 0 { // in a backoff window — consume one tick.
st.skipRemain--
return false
}
if !st.parsed.matches(now) || st.inFlight >= maxConcurrentPerAgent {
return false
}
st.inFlight++ // reserve the slot before launching.
return true
}
// launch runs one scheduled invocation in its own goroutine, updating the
// agent's backoff/concurrency state on completion. The run is empty-input (a
// scheduled agent acts on its own instructions) and attributed to its service
// account when bound, else the synthetic scheduler actor.
func (sc *scheduler) launch(ctx context.Context, a Agent, key string) {
sc.wg.Add(1)
go func() {
defer sc.wg.Done()
runCtx, cancel := context.WithTimeout(ctx, runTimeout)
defer cancel()
// Scheduled runs carry no HTTP request/IP; requestID/clientIP are empty.
r, gateErr := sc.svc.runAgent(runCtx, a, "", scheduledActor(a), "", "")
ok := gateErr == nil && r.Status == "ok"
sc.mu.Lock()
st := sc.stateForLocked(key)
st.inFlight--
if ok {
st.failstreakReset()
} else {
st.failstreakBump()
}
remain := st.skipRemain
streak := st.failstreak
sc.mu.Unlock()
switch {
case gateErr != nil:
sc.log.Warn("scheduled run gated (not executed)",
"org", a.Org, "agent", a.Name, "err", gateErr, "failstreak", streak, "backoffTicks", remain)
case r.Status != "ok":
sc.log.Warn("scheduled run errored",
"org", a.Org, "agent", a.Name, "err", r.Error, "failstreak", streak, "backoffTicks", remain)
default:
sc.log.Info("scheduled run ok", "org", a.Org, "agent", a.Name, "durationMs", r.DurationMs)
}
}()
}
// stateForLocked returns (creating if needed) the runtime state for an agent
// key. The CALLER MUST hold sc.mu — every read/write of agentState fields is
// serialized by that one lock, so the scheduler has no data race between a tick
// deciding to run and a completion goroutine updating backoff/inFlight.
func (sc *scheduler) stateForLocked(key string) *agentState {
st := sc.states[key]
if st == nil {
st = &agentState{}
sc.states[key] = st
}
return st
}
// pruneDeleted drops runtime state for agents that no longer appear in the work
// set (deleted or switched to one-shot), but keeps any with a run still in
// flight so its completion bookkeeping lands on live state.
func (sc *scheduler) pruneDeleted(live map[string]bool) {
sc.mu.Lock()
defer sc.mu.Unlock()
for key, st := range sc.states {
if !live[key] && st.inFlight == 0 {
delete(sc.states, key)
}
}
}
// failstreakReset clears the failure streak and backoff after a success.
func (st *agentState) failstreakReset() { st.failstreak, st.skipRemain = 0, 0 }
// failstreakBump grows the failure streak and sets the next backoff window to
// 2^(streak-1) ticks, capped — 1,2,4,8,… minutes between retries.
func (st *agentState) failstreakBump() {
st.failstreak++
skip := 1 << uint(min(st.failstreak-1, 30)) // guard the shift; 2^30 >> cap.
if skip > maxBackoffTicks {
skip = maxBackoffTicks
}
st.skipRemain = skip
}
// stateKey namespaces runtime state by org+name. The NUL separator can never
// appear in either (nameRE + org validation forbid it), so keys are injective.
func stateKey(org, name string) string { return org + "\x00" + name }
// scheduledActor is the audit-trail Actor for a scheduled run. It is ALWAYS
// prefixed "scheduler" so a scheduled run can never masquerade as a validated
// interactive principal (org/sub). When the agent carries a service-account id
// it is appended as an UNVERIFIED hint (Red LOW-2): the id is client-supplied on
// create and not yet checked against IAM (that is the service-account keystone),
// so it must be clearly non-authoritative, not the bare "principal". Once IAM
// agent service accounts land, this becomes a verified identity.
func scheduledActor(a Agent) string {
if sa := strings.TrimSpace(a.ServiceAccountID); sa != "" {
return schedulerActor + ":" + a.Org + "/" + a.Name + " (sa:" + sa + " unverified)"
}
return schedulerActor + ":" + a.Org + "/" + a.Name
}
+311
View File
@@ -0,0 +1,311 @@
package agents
import (
"context"
"encoding/json"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/hanzoai/cloud"
"github.com/hanzoai/cloud/types"
"github.com/hanzoai/commerce/metering"
luxlog "github.com/luxfi/log"
)
// countingAI records how many completions ran and can be told to fail, so
// scheduler tests can assert on run count + drive backoff deterministically.
type countingAI struct {
mu sync.Mutex
calls int32
fail bool
err error
block chan struct{} // when non-nil, ChatCompletion blocks until closed.
}
func (c *countingAI) ChatCompletion(_ context.Context, _ *types.ChatRequest) (*types.ChatResponse, error) {
atomic.AddInt32(&c.calls, 1)
if c.block != nil {
<-c.block
}
c.mu.Lock()
fail, err := c.fail, c.err
c.mu.Unlock()
if fail {
if err == nil {
err = errTest
}
return nil, err
}
return &types.ChatResponse{Content: "done"}, nil
}
func (c *countingAI) count() int32 { return atomic.LoadInt32(&c.calls) }
// schedSvc builds an svc + scheduler with NO billing (gate allows) and the given
// AI, seeded with the supplied agents. Returns the scheduler for direct tick().
func schedSvc(t *testing.T, ai types.AIClient, seed ...Agent) *scheduler {
t.Helper()
s := &svc{store: testStore(t), ai: ai, log: luxlog.New("test")}
for _, a := range seed {
if err := s.store.Create(context.Background(), a); err != nil {
t.Fatalf("seed %s/%s: %v", a.Org, a.Name, err)
}
}
sc := newScheduler(s, luxlog.New("test"))
return sc
}
func longRunning(org, name, cron string) Agent {
a := mk(org, name)
a.ExecutionMode, a.Schedule = ModeLongRunning, cron
return a
}
// waitFor polls a condition briefly (async run goroutines).
func waitFor(cond func() bool) bool {
for i := 0; i < 200; i++ {
if cond() {
return true
}
time.Sleep(5 * time.Millisecond)
}
return cond()
}
// TestSchedulerFiresDueAgent: a tick at a minute the cron matches launches one
// run; a tick at a non-matching minute launches none.
func TestSchedulerFiresDueAgent(t *testing.T) {
ai := &countingAI{}
sc := schedSvc(t, ai, longRunning("acme", "cron", "*/5 * * * *"))
ctx := context.Background()
sc.tick(ctx, at(t, "2026-07-01 12:36")) // 36 not multiple of 5 -> no fire
time.Sleep(20 * time.Millisecond)
if ai.count() != 0 {
t.Fatalf("non-matching minute must not fire, got %d", ai.count())
}
sc.tick(ctx, at(t, "2026-07-01 12:35")) // matches */5
if !waitFor(func() bool { return ai.count() == 1 }) {
t.Fatalf("matching minute must fire once, got %d", ai.count())
}
}
// TestSchedulerRecordsRun: a scheduled run is persisted to the run history, just
// like an HTTP run — the scheduler shares runAgent.
func TestSchedulerRecordsRun(t *testing.T) {
ai := &countingAI{}
sc := schedSvc(t, ai, longRunning("acme", "cron", "* * * * *"))
ctx := context.Background()
sc.tick(ctx, at(t, "2026-07-01 12:00"))
if !waitFor(func() bool {
runs, _ := sc.svc.store.ListRuns(ctx, "acme", "cron", 10)
return len(runs) == 1 && runs[0].Status == "ok"
}) {
runs, _ := sc.svc.store.ListRuns(ctx, "acme", "cron", 10)
t.Fatalf("scheduled run not recorded: %+v", runs)
}
}
// TestSchedulerBackoffOnFailure: after a failed run the agent is skipped for a
// growing number of ticks, so a broken agent stops hammering. The first failing
// tick fires; the immediately-following matching tick is skipped (backoff=1).
func TestSchedulerBackoffOnFailure(t *testing.T) {
ai := &countingAI{fail: true}
sc := schedSvc(t, ai, longRunning("acme", "cron", "* * * * *"))
ctx := context.Background()
sc.tick(ctx, at(t, "2026-07-01 12:00"))
if !waitFor(func() bool { return ai.count() == 1 }) {
t.Fatalf("first tick should attempt the run, got %d", ai.count())
}
// Wait for the failure to register the backoff window.
if !waitFor(func() bool {
sc.mu.Lock()
defer sc.mu.Unlock()
st := sc.states[stateKey("acme", "cron")]
return st != nil && st.failstreak == 1 && st.skipRemain == 1
}) {
t.Fatal("failure should set failstreak=1, skipRemain=1")
}
// Next matching tick is consumed by backoff -> no new run.
sc.tick(ctx, at(t, "2026-07-01 12:01"))
time.Sleep(20 * time.Millisecond)
if ai.count() != 1 {
t.Fatalf("backoff tick must not fire, got %d", ai.count())
}
// The tick after that (skip exhausted) fires again.
sc.tick(ctx, at(t, "2026-07-01 12:02"))
if !waitFor(func() bool { return ai.count() == 2 }) {
t.Fatalf("post-backoff tick should fire, got %d", ai.count())
}
}
// TestSchedulerConcurrencyCap: a slow run holds the single per-agent slot, so a
// second matching tick while it is in flight does NOT start a second run.
func TestSchedulerConcurrencyCap(t *testing.T) {
ai := &countingAI{block: make(chan struct{})}
sc := schedSvc(t, ai, longRunning("acme", "cron", "* * * * *"))
ctx := context.Background()
sc.tick(ctx, at(t, "2026-07-01 12:00")) // starts run #1, which blocks
if !waitFor(func() bool { return ai.count() == 1 }) {
t.Fatalf("first run should start, got %d", ai.count())
}
sc.tick(ctx, at(t, "2026-07-01 12:01")) // slot busy -> no second run
time.Sleep(20 * time.Millisecond)
if ai.count() != 1 {
t.Fatalf("concurrency cap breached: %d runs in flight", ai.count())
}
close(ai.block) // let run #1 finish
if !waitFor(func() bool {
sc.mu.Lock()
defer sc.mu.Unlock()
st := sc.states[stateKey("acme", "cron")]
return st != nil && st.inFlight == 0
}) {
t.Fatal("in-flight count should drain to 0 after completion")
}
}
// TestSchedulerBillsScheduledRun: a scheduled tick goes through the SAME gate +
// meter as an HTTP run — a funded agent's scheduled run debits its OWN org via
// commerce (product=agent), proving the billing path is live on the cron path,
// not just the HTTP handler (Red INFO-2).
func TestSchedulerBillsScheduledRun(t *testing.T) {
bs := &billServer{available: 100000}
m, err := metering.New(metering.Config{BaseURL: bs.start(t), Token: "svc-tok", Org: "hanzo"})
if err != nil {
t.Fatalf("metering.New: %v", err)
}
s := &svc{
store: testStore(t),
ai: &countingAI{},
log: luxlog.New("test"),
bill: cloud.NewResourceMeter(cloud.Deps{Metering: m, Logger: luxlog.New("test")}, meterKind),
}
if err := s.store.Create(context.Background(), longRunning("acme", "cron", "* * * * *")); err != nil {
t.Fatalf("seed: %v", err)
}
sc := newScheduler(s, luxlog.New("test"))
sc.tick(context.Background(), at(t, "2026-07-01 12:00"))
if !waitForDebit(func() bool { return bs.debits() == 1 }) {
t.Fatalf("a scheduled run on a funded org must debit once, got %d", bs.debits())
}
org, ubody := bs.lastDebit()
if org != "acme" {
t.Fatalf("scheduled debit org = %q, want the agent's own org acme", org)
}
var u struct {
User string `json:"user"`
Provider string `json:"provider"`
Actor string `json:"actor"`
}
_ = json.Unmarshal(ubody, &u)
if u.User != "acme" || u.Provider != meterKind {
t.Fatalf("scheduled debit user/provider = %q/%q, want acme/%s", u.User, u.Provider, meterKind)
}
// The actor MUST be the "scheduler:" namespace, never a bare "org/sub" that
// could be mistaken for a validated interactive principal (Red LOW-2).
if !strings.HasPrefix(u.Actor, schedulerActor+":") {
t.Fatalf("scheduled actor = %q, want a %q-prefixed (non-principal) actor", u.Actor, schedulerActor)
}
}
// TestSchedulerGatesUnfundedRun: a scheduled run on an unfunded org is gated
// (fail-closed) so the model NEVER runs and nothing is debited — an unfunded
// long-running agent can't burn free inference every minute.
func TestSchedulerGatesUnfundedRun(t *testing.T) {
bs := &billServer{available: 0}
m, _ := metering.New(metering.Config{BaseURL: bs.start(t), Token: "t", Org: "hanzo"})
ai := &countingAI{}
s := &svc{
store: testStore(t),
ai: ai,
log: luxlog.New("test"),
bill: cloud.NewResourceMeter(cloud.Deps{Metering: m, Logger: luxlog.New("test")}, meterKind),
}
if err := s.store.Create(context.Background(), longRunning("acme", "cron", "* * * * *")); err != nil {
t.Fatalf("seed: %v", err)
}
sc := newScheduler(s, luxlog.New("test"))
sc.tick(context.Background(), at(t, "2026-07-01 12:00"))
time.Sleep(40 * time.Millisecond)
if ai.count() != 0 {
t.Fatalf("unfunded scheduled run must NOT invoke the model, got %d", ai.count())
}
if bs.debits() != 0 {
t.Fatalf("unfunded scheduled run must not debit, got %d", bs.debits())
}
}
// TestSchedulerStopDrainsCleanly: stop() with an un-expired ctx cancels the loop
// and waits for the (fast) in-flight run to finish before returning — the drain
// path that lets Shutdown close the store safely.
func TestSchedulerStopDrainsCleanly(t *testing.T) {
ai := &countingAI{}
sc := schedSvc(t, ai, longRunning("acme", "cron", "* * * * *"))
sc.start()
// Fire one run via a direct tick, then stop — stop must return after drain.
sc.tick(context.Background(), at(t, "2026-07-01 12:00"))
done := make(chan struct{})
go func() { sc.stop(context.Background()); close(done) }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("stop() did not return — drain hung")
}
// After a clean stop, no run goroutine is left holding a slot.
sc.mu.Lock()
st := sc.states[stateKey("acme", "cron")]
inFlight := 0
if st != nil {
inFlight = st.inFlight
}
sc.mu.Unlock()
if inFlight != 0 {
t.Fatalf("after drain inFlight=%d, want 0", inFlight)
}
}
// TestSchedulerStopHonorsDeadline: a run that IGNORES cancellation (blocks) must
// not hang stop() past the caller's deadline — stop returns at the ctx deadline
// rather than waiting the full runTimeout.
func TestSchedulerStopHonorsDeadline(t *testing.T) {
ai := &countingAI{block: make(chan struct{})}
sc := schedSvc(t, ai, longRunning("acme", "cron", "* * * * *"))
sc.start()
sc.tick(context.Background(), at(t, "2026-07-01 12:00"))
if !waitFor(func() bool { return ai.count() == 1 }) {
t.Fatal("run should have started")
}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
start := time.Now()
sc.stop(ctx) // the run is blocked and ignores ctx; stop must return at deadline
if elapsed := time.Since(start); elapsed > 2*time.Second {
t.Fatalf("stop() waited %v — did not honor the 100ms deadline", elapsed)
}
// Release the stuck run and let it fully drain BEFORE the test's store
// cleanup, so the late InsertRun can't race a closed DB.
close(ai.block)
sc.wg.Wait()
}
// TestSchedulerOnlyLongRunning: a one-shot agent is never fired by the
// scheduler even if its (dropped) schedule would have matched.
func TestSchedulerOnlyLongRunning(t *testing.T) {
ai := &countingAI{}
one := mk("acme", "one") // one-shot default, no schedule
sc := schedSvc(t, ai, one)
sc.tick(context.Background(), at(t, "2026-07-01 12:00"))
time.Sleep(20 * time.Millisecond)
if ai.count() != 0 {
t.Fatalf("one-shot agent must never be scheduled, got %d", ai.count())
}
}
+175 -26
View File
@@ -22,19 +22,44 @@ var (
// prompt (instructions), and a set of tool names it may call. Tenant isolation
// is the org column, enforced on every query. It never stores a secret — tool
// credentials live in KMS and are referenced by name at run time.
//
// The bot-lifecycle fields promote an agent from a one-shot callable into a
// long-running bot (per hanzo-agent-bot-architecture: "Bot = Agent + compute +
// long-running"):
//
// - ExecutionMode: "one-shot" (default; runs only when POSTed) or
// "long-running" (the scheduler invokes it on Schedule).
// - Schedule: a 5-field cron expression; required when long-running, ignored
// otherwise. The scheduler evaluates it once a minute.
// - ComputeRef: an optional visor machine id the bot is bound to. It is an
// opaque reference here; binding/lifecycle is owned elsewhere.
// - ServiceAccountID: an optional IAM agent service-account (<org>-<agent>).
// When set it is the Actor recorded on scheduled-run billing so an
// autonomous run is attributable to a principal, not just the org.
type Agent struct {
ID string
Org string
Name string
Model string
Instructions string
Description string
Tools []string
Status string
CreatedAt int64
UpdatedAt int64
ID string
Org string
Name string
Model string
Instructions string
Description string
Tools []string
Status string
ExecutionMode string
Schedule string
ComputeRef string
ServiceAccountID string
CreatedAt int64
UpdatedAt int64
}
// Execution modes. One-shot agents run only on an explicit POST; long-running
// agents are additionally invoked by the scheduler on their Schedule.
const (
ModeOneShot = "one-shot"
ModeLongRunning = "long-running"
)
// Run is one execution of an agent: the input, the produced output (or error),
// which model served it, and how long it took. Real history — every row is a
// call that actually happened.
@@ -84,16 +109,20 @@ func openStore(path string) (*Store, error) {
func (s *Store) migrate() error {
const ddl = `
CREATE TABLE IF NOT EXISTS agents (
id TEXT PRIMARY KEY,
org TEXT NOT NULL,
name TEXT NOT NULL,
model TEXT NOT NULL DEFAULT '',
instructions TEXT NOT NULL DEFAULT '',
description TEXT NOT NULL DEFAULT '',
tools TEXT NOT NULL DEFAULT '[]',
status TEXT NOT NULL DEFAULT 'ready',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
id TEXT PRIMARY KEY,
org TEXT NOT NULL,
name TEXT NOT NULL,
model TEXT NOT NULL DEFAULT '',
instructions TEXT NOT NULL DEFAULT '',
description TEXT NOT NULL DEFAULT '',
tools TEXT NOT NULL DEFAULT '[]',
status TEXT NOT NULL DEFAULT 'ready',
execution_mode TEXT NOT NULL DEFAULT 'one-shot',
schedule TEXT NOT NULL DEFAULT '',
compute_ref TEXT NOT NULL DEFAULT '',
service_account_id TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE UNIQUE INDEX IF NOT EXISTS ux_agents_org_name ON agents(org, name);
CREATE INDEX IF NOT EXISTS ix_agents_org_updated ON agents(org, updated_at);
@@ -115,9 +144,73 @@ CREATE INDEX IF NOT EXISTS ix_runs_org_agent_created ON agent_runs(org, agent_na
if _, err := s.db.Exec(ddl); err != nil {
return fmt.Errorf("migrate: %w", err)
}
// Forward, idempotent migration for databases created before the
// bot-lifecycle columns existed. Each ADD COLUMN is guarded by a live
// column-existence check (PRAGMA table_info), so re-running migrate() on an
// already-upgraded DB is a no-op and never errors — the DDL above handles
// fresh DBs, this handles pre-existing ones. It touches no storage-backend
// knob (driverName/DSN), so the SQLite-only storage lockdown is unaffected.
if err := s.addColumns("agents", map[string]string{
"execution_mode": "TEXT NOT NULL DEFAULT 'one-shot'",
"schedule": "TEXT NOT NULL DEFAULT ''",
"compute_ref": "TEXT NOT NULL DEFAULT ''",
"service_account_id": "TEXT NOT NULL DEFAULT ''",
}); err != nil {
return err
}
// Partial index for the once-a-minute scheduler scan — created AFTER the
// lifecycle columns exist (a legacy DB gains them just above), so it selects
// only the (typically few) scheduled long-running agents instead of
// full-scanning every org's agents on the single shared SQLite connection.
if _, err := s.db.Exec(`CREATE INDEX IF NOT EXISTS ix_agents_scheduled
ON agents(org, name) WHERE execution_mode='long-running' AND schedule<>''`); err != nil {
return fmt.Errorf("migrate: scheduled index: %w", err)
}
return nil
}
// addColumns adds each missing column to table, idempotently. A column already
// present is skipped; a fresh install (all present from the CREATE) is a no-op.
func (s *Store) addColumns(table string, cols map[string]string) error {
have, err := s.columns(table)
if err != nil {
return err
}
for name, def := range cols {
if have[name] {
continue
}
// name/def are package-internal literals, never user input — no
// injection surface. SQLite forbids parameterizing DDL identifiers.
if _, err := s.db.Exec(`ALTER TABLE ` + table + ` ADD COLUMN ` + name + ` ` + def); err != nil {
return fmt.Errorf("migrate: add %s.%s: %w", table, name, err)
}
}
return nil
}
// columns returns the set of column names on table via PRAGMA table_info.
func (s *Store) columns(table string) (map[string]bool, error) {
rows, err := s.db.Query(`PRAGMA table_info(` + table + `)`)
if err != nil {
return nil, fmt.Errorf("migrate: table_info %s: %w", table, err)
}
defer func() { _ = rows.Close() }()
have := map[string]bool{}
for rows.Next() {
var (
cid, notnull, pk int
name, ctype string
dflt sql.NullString
)
if err := rows.Scan(&cid, &name, &ctype, &notnull, &dflt, &pk); err != nil {
return nil, fmt.Errorf("migrate: scan table_info: %w", err)
}
have[name] = true
}
return have, rows.Err()
}
func (s *Store) Close() error { return s.db.Close() }
func encodeList(xs []string) string {
@@ -142,23 +235,37 @@ func decodeList(s string) []string {
return xs
}
const agentCols = `id,org,name,model,instructions,description,tools,status,created_at,updated_at`
const agentCols = `id,org,name,model,instructions,description,tools,status,execution_mode,schedule,compute_ref,service_account_id,created_at,updated_at`
func scanAgent(sc interface{ Scan(...any) error }) (Agent, error) {
var a Agent
var tools string
err := sc.Scan(&a.ID, &a.Org, &a.Name, &a.Model, &a.Instructions, &a.Description,
&tools, &a.Status, &a.CreatedAt, &a.UpdatedAt)
&tools, &a.Status, &a.ExecutionMode, &a.Schedule, &a.ComputeRef, &a.ServiceAccountID,
&a.CreatedAt, &a.UpdatedAt)
a.Tools = decodeList(tools)
return a, err
}
// normalizeMode is the lowest-layer fail-safe default: an empty execution_mode
// is stored as one-shot so NO path (handler, scheduler, or a direct store call)
// can persist an agent the scheduler would treat ambiguously. The HTTP handler
// also defaults+validates, but this makes the invariant hold at the store.
func normalizeMode(m string) string {
if strings.TrimSpace(m) == "" {
return ModeOneShot
}
return m
}
// Create inserts one agent. A UNIQUE(org,name) violation surfaces as errConflict.
func (s *Store) Create(ctx context.Context, a Agent) error {
a.ExecutionMode = normalizeMode(a.ExecutionMode)
_, err := s.db.ExecContext(ctx,
`INSERT INTO agents (`+agentCols+`) VALUES (?,?,?,?,?,?,?,?,?,?)`,
`INSERT INTO agents (`+agentCols+`) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
a.ID, a.Org, a.Name, a.Model, a.Instructions, a.Description,
encodeList(a.Tools), a.Status, a.CreatedAt, a.UpdatedAt)
encodeList(a.Tools), a.Status, a.ExecutionMode, a.Schedule, a.ComputeRef,
a.ServiceAccountID, a.CreatedAt, a.UpdatedAt)
if err != nil {
if strings.Contains(err.Error(), "UNIQUE constraint failed") {
return errConflict
@@ -202,9 +309,13 @@ func (s *Store) List(ctx context.Context, org string) ([]Agent, error) {
// Update overwrites the mutable fields of an existing agent.
func (s *Store) Update(ctx context.Context, a Agent) error {
a.ExecutionMode = normalizeMode(a.ExecutionMode)
res, err := s.db.ExecContext(ctx,
`UPDATE agents SET model=?,instructions=?,description=?,tools=?,status=?,updated_at=? WHERE org=? AND name=?`,
a.Model, a.Instructions, a.Description, encodeList(a.Tools), a.Status, a.UpdatedAt, a.Org, a.Name)
`UPDATE agents SET model=?,instructions=?,description=?,tools=?,status=?,
execution_mode=?,schedule=?,compute_ref=?,service_account_id=?,updated_at=?
WHERE org=? AND name=?`,
a.Model, a.Instructions, a.Description, encodeList(a.Tools), a.Status,
a.ExecutionMode, a.Schedule, a.ComputeRef, a.ServiceAccountID, a.UpdatedAt, a.Org, a.Name)
if err != nil {
return fmt.Errorf("update agent: %w", err)
}
@@ -215,6 +326,44 @@ func (s *Store) Update(ctx context.Context, a Agent) error {
return nil
}
// ListLongRunning returns every agent across ALL orgs whose execution_mode is
// long-running and that carries a non-empty schedule — the scheduler's work
// set. It is the ONE cross-org query in this store; the scheduler is a trusted
// in-process subsystem (not a tenant request), and each returned agent carries
// its own Org so every downstream action (run, gate, meter) stays scoped to the
// agent's own tenant.
func (s *Store) ListLongRunning(ctx context.Context) ([]Agent, error) {
rows, err := s.db.QueryContext(ctx,
`SELECT `+agentCols+` FROM agents
WHERE execution_mode=? AND schedule<>'' ORDER BY org, name`, ModeLongRunning)
if err != nil {
return nil, fmt.Errorf("list long-running: %w", err)
}
defer func() { _ = rows.Close() }()
var out []Agent
for rows.Next() {
a, err := scanAgent(rows)
if err != nil {
return nil, fmt.Errorf("scan agent: %w", err)
}
out = append(out, a)
}
return out, rows.Err()
}
// CountLongRunning returns how many scheduled long-running agents an org has —
// used to cap an org's scheduler footprint at create time.
func (s *Store) CountLongRunning(ctx context.Context, org string) (int, error) {
var n int
err := s.db.QueryRowContext(ctx,
`SELECT COUNT(*) FROM agents WHERE org=? AND execution_mode=? AND schedule<>''`,
org, ModeLongRunning).Scan(&n)
if err != nil {
return 0, fmt.Errorf("count long-running: %w", err)
}
return n, nil
}
// Delete removes an agent and its run history. Reports whether a row went.
func (s *Store) Delete(ctx context.Context, org, name string) (bool, error) {
tx, err := s.db.BeginTx(ctx, nil)
+146
View File
@@ -0,0 +1,146 @@
package agents
import (
"context"
"database/sql"
"path/filepath"
"testing"
"time"
)
// TestLifecycleFieldsRoundTrip: the four bot-lifecycle columns persist and read
// back through Create/Get/Update.
func TestLifecycleFieldsRoundTrip(t *testing.T) {
s := testStore(t)
ctx := context.Background()
a := mk("acme", "sweeper")
a.ExecutionMode = ModeLongRunning
a.Schedule = "*/5 * * * *"
a.ComputeRef = "vm-123"
a.ServiceAccountID = "acme-sweeper"
if err := s.Create(ctx, a); err != nil {
t.Fatalf("create: %v", err)
}
got, err := s.Get(ctx, "acme", "sweeper")
if err != nil {
t.Fatalf("get: %v", err)
}
if got.ExecutionMode != ModeLongRunning || got.Schedule != "*/5 * * * *" ||
got.ComputeRef != "vm-123" || got.ServiceAccountID != "acme-sweeper" {
t.Fatalf("lifecycle fields not persisted: %+v", got)
}
got.Schedule = "0 9 * * 1"
got.ComputeRef = "vm-456"
got.UpdatedAt = time.Now().Unix()
if err := s.Update(ctx, got); err != nil {
t.Fatalf("update: %v", err)
}
got2, _ := s.Get(ctx, "acme", "sweeper")
if got2.Schedule != "0 9 * * 1" || got2.ComputeRef != "vm-456" {
t.Fatalf("update did not persist lifecycle edits: %+v", got2)
}
}
// TestDefaultExecutionMode: a fresh agent created without a mode reads back as
// one-shot (the DEFAULT that the DDL + migration guarantee), never empty.
func TestDefaultExecutionMode(t *testing.T) {
s := testStore(t)
ctx := context.Background()
if err := s.Create(ctx, mk("acme", "plain")); err != nil {
t.Fatalf("create: %v", err)
}
got, _ := s.Get(ctx, "acme", "plain")
if got.ExecutionMode != ModeOneShot {
t.Fatalf("default execution_mode = %q, want %q", got.ExecutionMode, ModeOneShot)
}
}
// TestListLongRunning: returns only long-running agents WITH a schedule, across
// orgs, and each carries its own org (the scheduler scopes actions per agent).
func TestListLongRunning(t *testing.T) {
s := testStore(t)
ctx := context.Background()
oneShot := mk("acme", "oneshot") // default one-shot
lr := mk("acme", "cron")
lr.ExecutionMode, lr.Schedule = ModeLongRunning, "* * * * *"
lrNoSched := mk("beta", "cron")
lrNoSched.ExecutionMode, lrNoSched.Schedule = ModeLongRunning, "" // no schedule -> excluded
lrOther := mk("beta", "nightly")
lrOther.ExecutionMode, lrOther.Schedule = ModeLongRunning, "0 0 * * *"
for _, a := range []Agent{oneShot, lr, lrNoSched, lrOther} {
if err := s.Create(ctx, a); err != nil {
t.Fatalf("seed %s/%s: %v", a.Org, a.Name, err)
}
}
got, err := s.ListLongRunning(ctx)
if err != nil {
t.Fatalf("list long-running: %v", err)
}
if len(got) != 2 {
t.Fatalf("want 2 scheduled agents (acme/cron, beta/nightly), got %d: %+v", len(got), got)
}
seen := map[string]string{}
for _, a := range got {
seen[a.Org+"/"+a.Name] = a.Schedule
}
if seen["acme/cron"] != "* * * * *" || seen["beta/nightly"] != "0 0 * * *" {
t.Fatalf("wrong scheduled set: %v", seen)
}
if _, bad := seen["beta/cron"]; bad {
t.Fatalf("long-running agent WITHOUT a schedule must be excluded")
}
}
// TestMigrationIdempotentOnLegacyDB: a DB created with the PRE-lifecycle schema
// (no new columns) is migrated forward on open, existing rows survive with the
// column defaults, and re-opening (re-running migrate) is a clean no-op.
func TestMigrationIdempotentOnLegacyDB(t *testing.T) {
path := filepath.Join(t.TempDir(), "legacy.db")
// Hand-build the legacy schema + a legacy row, exactly as the pre-lifecycle
// migrate() would have, then close.
legacy, err := sql.Open("sqlite", path)
if err != nil {
t.Fatalf("open legacy: %v", err)
}
const legacyDDL = `
CREATE TABLE agents (
id TEXT PRIMARY KEY, org TEXT NOT NULL, name TEXT NOT NULL,
model TEXT NOT NULL DEFAULT '', instructions TEXT NOT NULL DEFAULT '',
description TEXT NOT NULL DEFAULT '', tools TEXT NOT NULL DEFAULT '[]',
status TEXT NOT NULL DEFAULT 'ready', created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL);`
if _, err := legacy.Exec(legacyDDL); err != nil {
t.Fatalf("legacy ddl: %v", err)
}
now := time.Now().Unix()
if _, err := legacy.Exec(
`INSERT INTO agents (id,org,name,model,instructions,description,tools,status,created_at,updated_at)
VALUES ('old-id','acme','legacy','m','i','d','[]','ready',?,?)`, now, now); err != nil {
t.Fatalf("legacy insert: %v", err)
}
_ = legacy.Close()
// Open through the real store TWICE — the first migrates, the second proves
// idempotency (no error re-adding existing columns).
for i := 0; i < 2; i++ {
st, err := openStore(path)
if err != nil {
t.Fatalf("open #%d migrate failed: %v", i, err)
}
got, err := st.Get(context.Background(), "acme", "legacy")
if err != nil {
t.Fatalf("open #%d: legacy row lost: %v", i, err)
}
if got.ExecutionMode != ModeOneShot {
t.Fatalf("open #%d: migrated row default mode = %q, want %q", i, got.ExecutionMode, ModeOneShot)
}
if got.Schedule != "" || got.ComputeRef != "" || got.ServiceAccountID != "" {
t.Fatalf("open #%d: migrated defaults not empty: %+v", i, got)
}
_ = st.Close()
}
}
+1 -1
View File
@@ -43,7 +43,7 @@ func (b *billDouble) start(t *testing.T) string {
atomic.AddInt32(&b.usages, 1)
body, _ := io.ReadAll(r.Body)
b.mu.Lock()
b.usageOrg, b.usageBody = r.Header.Get("X-IAM-Org-Id"), body
b.usageOrg, b.usageBody = r.Header.Get("X-Org-Id"), body
b.mu.Unlock()
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `{"transactionId":"tx_1","type":"usage"}`)
+5 -2
View File
@@ -25,7 +25,10 @@ import (
)
// billServer is a minimal commerce double: it returns a fixed balance and
// records the X-IAM-Org-Id header + body of any usage debit.
// records the X-Org-Id header + body of any usage debit. X-Org-Id is the
// header commerce's service-token auth actually reads (metering >= v0.1.2);
// the tenant namespace resolves from it, so a stale name would silently debit
// the default org.
type billServer struct {
available int64
@@ -45,7 +48,7 @@ func (b *billServer) start(t *testing.T) string {
atomic.AddInt32(&b.usages, 1)
body, _ := io.ReadAll(r.Body)
b.mu.Lock()
b.usageOrg, b.usageBody = r.Header.Get("X-IAM-Org-Id"), body
b.usageOrg, b.usageBody = r.Header.Get("X-Org-Id"), body
b.mu.Unlock()
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `{"transactionId":"tx_1","type":"usage"}`)
+1 -1
View File
@@ -8,7 +8,7 @@ require (
github.com/ClickHouse/clickhouse-go/v2 v2.40.1
github.com/coder/websocket v1.8.12
github.com/dop251/goja v0.0.0-20260627200808-0b76000cabdb
github.com/hanzoai/commerce/metering v0.1.0
github.com/hanzoai/commerce/metering v0.1.2
github.com/hanzoai/goa v0.0.0-20260629022130-202bcdc87340
github.com/hanzoai/iam v1.19.8-0.20260622075908-c2fd4872545b
github.com/hanzoai/kms/sdk/go v1.1.1
+2 -2
View File
@@ -1129,8 +1129,8 @@ github.com/hanzoai/builder v0.3.13 h1:tAOJ+0Q0xrrovk7lkvaZxuKZ4lqENIB6tE0Rr9+6Bo
github.com/hanzoai/builder v0.3.13/go.mod h1:TWZaiP0Y9tCMwtLH2EvQqBAeT1f3aJI5Y0XPM8S0wcE=
github.com/hanzoai/commerce v1.42.29 h1:7uvorI2/o8OO5QwL+D9TxAELprBvfj4j07dmsPO+xwE=
github.com/hanzoai/commerce v1.42.29/go.mod h1:M/a2EobHgkTWcaWd69z7YOg5xccSyj1oSmyA0N48C2w=
github.com/hanzoai/commerce/metering v0.1.0 h1:qKK5eqHbZwiWVZLzxvQh+rWI/h+yIt6BJlCWTEE7feI=
github.com/hanzoai/commerce/metering v0.1.0/go.mod h1:LLFOtgJM5OczHb0D/4GRntL0a/jl80VkmyvTYCZN3yE=
github.com/hanzoai/commerce/metering v0.1.2 h1:r0PnKVeseb7J/TAiL1RUEWD91oBHP3yOeC7JiJkylR4=
github.com/hanzoai/commerce/metering v0.1.2/go.mod h1:LLFOtgJM5OczHb0D/4GRntL0a/jl80VkmyvTYCZN3yE=
github.com/hanzoai/common v0.67.7 h1:6LAzDF4MOPVE6TvtA7gtYRvWESCTI5g0Hsqhe00XTd4=
github.com/hanzoai/common v0.67.7/go.mod h1:5bHRjMM3zyieuv0xi4S4HRlC2drcyfPGVJRQUT8moQo=
github.com/hanzoai/dashscope-go-sdk v0.0.2 h1:L/FlStjXeehrNSBqU8y/ecSG/MnaAJqfyfjmLYmY344=
+8 -2
View File
@@ -164,6 +164,9 @@ func billingEnabled(m *metering.Client) bool { return m != nil && m.Enabled() }
// - health/liveness probes (every /v1/<svc>/health and bare /health),
// - /v1/ai/* : the ai subsystem self-meters its own LLM token costs to
// commerce; charging again here would DOUBLE-BILL,
// - /v1/agents/* : the canonical agents subsystem now gates + meters its OWN
// per-run fee to commerce (clients/agents runAgent); the edge must stay 0 or
// every agent run is billed twice,
// - other subsystems that already self-meter their units (commerce billing
// itself, o11y telemetry, mcp tool dispatch),
//
@@ -188,8 +191,10 @@ func DefaultPrice(c *zip.Ctx) int64 {
}
}
// Generic agent/compute edge with no finer meter: a flat per-request charge.
if strings.HasPrefix(path, "/v1/agent/") || strings.HasPrefix(path, "/v1/agents/") {
// Legacy singular /v1/agent/* edge (the bot reverse-proxy) has no finer
// meter of its own: a flat per-request charge. The canonical plural
// /v1/agents/* is self-metered (above) and never reaches here.
if strings.HasPrefix(path, "/v1/agent/") {
return cloudEdgePriceCents
}
@@ -209,6 +214,7 @@ const cloudEdgePriceCents int64 = 1
// product label, so an additional flat edge charge would double-bill.
var selfMeteredPrefixes = []string{
"/v1/ai/", // LLM token costs metered by the ai subsystem.
"/v1/agents/", // per-run agent fee metered by the agents subsystem.
"/v1/commerce/", // billing itself; not metered as usage.
"/v1/o11y/", // telemetry ingest; not user-billable here.
"/v1/mcp/", // tool dispatch meters per-tool downstream.
+3 -2
View File
@@ -233,8 +233,9 @@ func TestDefaultPrice(t *testing.T) {
{"/healthz", true, "liveness probe"},
{"/v1/iam/health", true, "subsystem health suffix"},
{"/v1/base/health", true, "subsystem health suffix"},
{"/v1/agent/run", false, "generic agent edge has no finer meter"},
{"/v1/agents/list", false, "generic agents edge"},
{"/v1/agent/run", false, "legacy singular agent edge has no finer meter"},
{"/v1/agents/list", true, "agents subsystem self-meters per-run fee — must not double-bill"},
{"/v1/agents/x/run", true, "agent run self-metered by the agents subsystem"},
{"/v1/unknown/thing", true, "unpriced path defaults to 0 (opt-in metering)"},
}
for _, tc := range cases {
+46 -17
View File
@@ -38,8 +38,8 @@ import (
"strings"
"github.com/hanzoai/commerce/metering"
"github.com/zap-proto/zip"
luxlog "github.com/luxfi/log"
"github.com/zap-proto/zip"
)
// DefaultResourceFeeCents is the fallback flat provision/create fee (in cents)
@@ -87,14 +87,20 @@ func (rm *ResourceMeter) Enabled() bool { return rm != nil && rm.m != nil && rm.
//
// costCents<=0 means the kind is free → no gate (mirrors BillingGate's price==0
// short-circuit). org MUST be the caller's resolved slug; it is sent as the
// commerce user AND X-IAM-Org-Id so the CALLER's ledger is checked, overriding
// commerce user AND X-Org-Id so the CALLER's ledger is checked, overriding
// the client default org — the anti-cross-tenant property. The balance check
// honors ctx (a client disconnect/timeout cancels it).
//
// costCents is forwarded as AuthInput.AmountCents so the gate enforces
// available >= costCents, not merely available > 0 — otherwise a 1-cent balance
// would authorize an arbitrarily expensive charge (the debit still lands, taking
// the ledger negative). This mirrors what a prepaid gate must do: refuse a
// request the balance cannot cover BEFORE the work runs.
func (rm *ResourceMeter) Gate(ctx context.Context, org, kind string, costCents int64) error {
if !rm.Enabled() || costCents <= 0 {
return nil
}
return rm.m.Authorize(ctx, metering.AuthInput{User: org, Org: org})
return rm.m.Authorize(ctx, metering.AuthInput{User: org, Org: org, AmountCents: costCents})
}
// Meter records a successful charge to the caller's org ledger. It is the ONE
@@ -108,26 +114,49 @@ func (rm *ResourceMeter) Gate(ctx context.Context, org, kind string, costCents i
// received, and a request-context cancellation must not cancel the debit (mirror
// of BillingGate). A debit failure is logged for reconciliation, not swallowed.
func (rm *ResourceMeter) Meter(org, kind string, amountCents int64, requestID, clientIP string) {
if !rm.Enabled() || amountCents <= 0 {
rm.MeterUsage(org, kind, metering.Usage{
AmountCents: amountCents,
RequestID: requestID,
ClientIP: clientIP,
})
}
// MeterUsage is the general-purpose per-org debit: it records the caller-built
// usage event after forcing the per-org billing invariants that make the debit
// land on the CALLER's ledger and never another tenant's:
//
// - u.User and u.Org are OVERWRITTEN to the caller's org slug (the per-org
// prepaid billing key + the X-Org-Id namespace) — a caller can never bill
// someone else, and a surface can't accidentally leave them unset (which
// would debit the client-default org).
// - Provider defaults to the meter's provider; Status defaults to "success";
// Currency defaults to "usd".
//
// Everything else the caller supplies (AmountCents, Model, Actor, RequestID,
// token counts, ClientIP) flows through so a metered surface can attribute spend
// richly. Like Meter it is fire-and-forget on a background context and a no-op
// when billing is unconfigured or AmountCents<=0. kind is for the failure log.
func (rm *ResourceMeter) MeterUsage(org, kind string, u metering.Usage) {
if !rm.Enabled() || u.AmountCents <= 0 {
return
}
usage := metering.Usage{
User: org, // per-ORG billing: ledger keyed on the org slug.
Org: org, // X-IAM-Org-Id -> caller's namespace (overrides client default).
Currency: "usd",
AmountCents: amountCents,
Provider: rm.provider, // the product/surface that metered (e.g. "functions", "s3", "provisioning").
Model: kind, // the billed unit within the product (e.g. "sql", "invoke", "op") — per-item attribution.
RequestID: requestID,
Status: "success",
ClientIP: clientIP,
u.User = org // per-ORG billing: ledger keyed on the org slug.
u.Org = org // X-Org-Id -> caller's namespace (overrides client default).
if u.Provider == "" {
u.Provider = rm.provider
}
if u.Status == "" {
u.Status = "success"
}
if u.Currency == "" {
u.Currency = "usd"
}
m, log, env := rm.m, rm.log, rm.env
go func() {
if _, err := m.Record(context.Background(), usage); err != nil && log != nil {
if _, err := m.Record(context.Background(), u); err != nil && log != nil {
log.Error("resource debit failed (resource created, not billed)",
"org", org, "kind", kind, "provider", usage.Provider,
"cents", amountCents, "env", env, "err", err)
"org", org, "kind", kind, "provider", u.Provider,
"cents", u.AmountCents, "env", env, "err", err)
}
}()
}
+7 -7
View File
@@ -2,7 +2,7 @@ package cloud
// Tests for the shared per-org resource gate+meter (ResourceMeter). They drive
// the real metering client against a fake commerce server that RECORDS the
// X-IAM-Org-Id tenant header and request bodies, so the multitenancy contract is
// X-Org-Id tenant header and request bodies, so the multitenancy contract is
// proven end-to-end over HTTP — no mock of the metering client itself. The fake
// is built with a metering DEFAULT org of "hanzo"; every assertion that the
// caller "acme" is billed (not "hanzo") proves the per-call org override is what
@@ -25,16 +25,16 @@ import (
)
// recCommerce answers the metering client's balance + usage calls and records,
// per endpoint, the X-IAM-Org-Id tenant header it saw plus the usage body — the
// per endpoint, the X-Org-Id tenant header it saw plus the usage body — the
// evidence for the per-org / cross-tenant assertions.
type recCommerce struct {
balanceAvailable int64 // returned as {"available":N} on GET /v1/billing/balance
balanceStatus int // 0 => 200
mu sync.Mutex
balanceOrg string // last X-IAM-Org-Id on a balance call
balanceOrg string // last X-Org-Id on a balance call
balanceCalls int32
usageOrg string // last X-IAM-Org-Id on a usage call
usageOrg string // last X-Org-Id on a usage call
usageCount int32
usageBody []byte
}
@@ -45,7 +45,7 @@ func (f *recCommerce) server(t *testing.T) *httptest.Server {
mux.HandleFunc("/v1/billing/balance", func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&f.balanceCalls, 1)
f.mu.Lock()
f.balanceOrg = r.Header.Get("X-IAM-Org-Id")
f.balanceOrg = r.Header.Get("X-Org-Id")
f.mu.Unlock()
status := f.balanceStatus
if status == 0 {
@@ -58,7 +58,7 @@ func (f *recCommerce) server(t *testing.T) *httptest.Server {
atomic.AddInt32(&f.usageCount, 1)
body, _ := io.ReadAll(r.Body)
f.mu.Lock()
f.usageOrg = r.Header.Get("X-IAM-Org-Id")
f.usageOrg = r.Header.Get("X-Org-Id")
f.usageBody = body
f.mu.Unlock()
w.WriteHeader(http.StatusOK)
@@ -157,7 +157,7 @@ func TestResourceMeter_GateFailOpenOnCommerceError(t *testing.T) {
}
}
// Meter debits the CALLER org: usage POST fires once, carries X-IAM-Org-Id:acme
// Meter debits the CALLER org: usage POST fires once, carries X-Org-Id:acme
// (not the default 'hanzo'), body user=="acme", amount==cost. This is the
// per-org / anti-cross-tenant debit proof.
func TestResourceMeter_MeterDebitsCallerOrg(t *testing.T) {
+8
View File
@@ -180,6 +180,14 @@ func Serve(enable []string) error {
shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
_ = healthSrv.Shutdown(shutdownCtx)
// Tear down subsystems that own process-lifetime resources (background
// workers, DB handles) BEFORE the HTTP server closes, within the deadline —
// e.g. the agents scheduler drains its in-flight runs (so a scheduled run's
// InsertRun + debit land) and closes its store. Best-effort: a teardown error
// is logged, not fatal, so one subsystem can't strand shutdown.
if err := ShutdownAll(shutdownCtx, cfg); err != nil {
deps.Logger.Warn("subsystem shutdown", "err", err)
}
// Close the audit store last so any in-flight append has drained through the
// serialized writer and the SQLite file is flushed cleanly.
if auditRec != nil {