feat(agents): /v1/agents per-org fail-closed metering + long-running scheduler
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:
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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, ¬null, &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)
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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"}`)
|
||||
|
||||
@@ -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"}`)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user